diff --git a/crates/api/src/database_encryption.rs b/crates/api/src/database_encryption.rs index bde2f587b..4f5384b9a 100644 --- a/crates/api/src/database_encryption.rs +++ b/crates/api/src/database_encryption.rs @@ -171,6 +171,7 @@ const APPROVED: &[ApprovedGroup] = &[ ApprovedGroup { table: "organization_invitations", columns: &["email", "role", "status", "token", "email_status", "email_last_error", "email_message_id"], reason: "Invitation workflow data required for indexed acceptance and delivery operations; tokens are short-lived" }, ApprovedGroup { table: "organization_limits_history", columns: &["changed_by", "change_reason", "changed_by_user_email", "credit_type", "source", "currency"], reason: "Restricted billing and limits audit history" }, ApprovedGroup { table: "organization_member_role_audit_log", columns: &["previous_role", "new_role"], reason: "Restricted administrator member-role audit history" }, + ApprovedGroup { table: "organization_model_access_log", columns: &["action"], reason: "Restricted administrator model-access audit history: a grant-or-revoke enum next to UUIDs, no customer content" }, ApprovedGroup { table: "organization_members", columns: &["role"], reason: "Queryable authorization role" }, ApprovedGroup { table: "organization_credit_consumption", columns: &["credit_type"], reason: "Restricted current billing consumption counter" }, ApprovedGroup { table: "organization_reporting_tokens", columns: &["name", "token_hash", "token_prefix"], reason: "Reporting credentials are one-way hashed; name and prefix are display metadata" }, diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 8c2727a89..cbb2dbfe8 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -406,12 +406,21 @@ pub async fn init_domain_services_with_pool( database::repositories::OrganizationUsageRepository::new(database.pool().clone()), ) as Arc; + // Create models service. It owns the model-resolve cache (models and + // access decisions on unlisted models) that the attestation and completion + // services below share, so its invalidation reaches all three. + let models_service = Arc::new(services::models::ModelsServiceImpl::new( + inference_provider_pool.clone(), + models_repo.clone(), + )); + // Create attestation service let attestation_service = Arc::new( services::attestation::AttestationService::init( attestation_repo, inference_provider_pool.clone(), models_repo.clone(), + models_service.model_resolve_cache(), metrics_service.clone(), usage_repository_for_attestation, config.ita.clone(), @@ -420,12 +429,6 @@ pub async fn init_domain_services_with_pool( .unwrap(), ); - // Create models service - let models_service = Arc::new(services::models::ModelsServiceImpl::new( - inference_provider_pool.clone(), - models_repo.clone(), - )); - // Prepare repositories for usage service (will be created after workspace service) let usage_repository = Arc::new( database::repositories::OrganizationUsageRepository::with_accounting_config( @@ -1861,11 +1864,14 @@ pub fn build_conversation_routes( /// Route classification (nearai/infra#193) — see `routes/attestation.rs` for /// the full table: /// - `GET /v1/attestation/report` and `GET /v1/signature/{chat_id}` require an -/// API key (`auth_middleware_with_api_key`). The middleware validates the key -/// and confirms its workspace and organization remain active (rejecting -/// missing/invalid/expired/deactivated credentials with 401); like -/// the signature route, report retrieval is non-billable — no usage or -/// billing records are created. +/// API key (`auth_middleware_with_workspace_context`). The middleware +/// validates the key and confirms its workspace and organization remain +/// active (rejecting missing/invalid/expired/deactivated credentials with +/// 401); like the signature route, report retrieval is non-billable — no +/// usage or billing records are created. It is the same check as +/// `auth_middleware_with_api_key`, and it also hands the handler the key's +/// organization: the report route needs it to decide whether the caller +/// may ask about an unlisted model. /// - `GET /v1/attestation/ita-token` and `GET /v1/attestation/pins` are /// deliberately public; the rationale is documented on /// `build_public_attestation_routes`. @@ -1881,7 +1887,7 @@ pub fn build_attestation_routes( .with_state(attestation_route_state.clone()) .layer(from_fn_with_state( auth_state_middleware.clone(), - auth_middleware_with_api_key, + middleware::auth::auth_middleware_with_workspace_context, )); // One section per attested backup provider; add new providers here. @@ -2085,8 +2091,10 @@ pub fn build_gateway_routes( /// `CLOUD_API_USAGE_TOKEN` service secret rather than the standard `sk-…` /// API-key middleware. Mounted under `/v1/internal/*` after the global /// `/v1` nest. Routes: `POST /internal/usage` (inference-proxy's -/// service-token usage reporter) and `GET /internal/models` (the model -/// catalog including unlisted models, for a gateway lane). Further internal +/// service-token usage reporter), `GET /internal/models` (the model +/// catalog including unlisted models, for a gateway lane) and +/// `GET /internal/model-access` (which organizations may use each unlisted +/// model, for a gateway lane that enforces access itself). Further internal /// endpoints can be added here without touching the sk-auth stack. pub fn build_internal_routes(app_state: AppState) -> Router { Router::new() @@ -2098,6 +2106,10 @@ pub fn build_internal_routes(app_state: AppState) -> Router { "/internal/models", get(crate::routes::completions::models_internal), ) + .route( + "/internal/model-access", + get(crate::routes::completions::model_access_internal), + ) .with_state(app_state) } @@ -2355,6 +2367,19 @@ fn build_admin_routes_with_options( "/admin/models/{model_name}/deprecation/confirm", axum::routing::post(confirm_model_deprecation), ) + .route( + "/admin/models/{model_name}/access", + axum::routing::get(crate::routes::admin_model_access::get_model_access), + ) + .route( + "/admin/models/{model_name}/access/{org_id}", + axum::routing::put(crate::routes::admin_model_access::grant_model_access) + .delete(crate::routes::admin_model_access::revoke_model_access), + ) + .route( + "/admin/organizations/{org_id}/model-access", + axum::routing::get(crate::routes::admin_model_access::list_organization_model_access), + ) .route("/admin/services", axum::routing::post(create_service)) .route("/admin/services/{id}", axum::routing::patch(update_service)) .route( diff --git a/crates/api/src/middleware/admin_policy.rs b/crates/api/src/middleware/admin_policy.rs index f4adf2001..4b1dab93d 100644 --- a/crates/api/src/middleware/admin_policy.rs +++ b/crates/api/src/middleware/admin_policy.rs @@ -38,6 +38,8 @@ pub(crate) fn admin_operation(method: &Method, matched_path: &str) -> AdminOpera "/admin/models" | "/admin/models/pricing-changes" | "/admin/models/{model_name}/history" + | "/admin/models/{model_name}/access" + | "/admin/organizations/{org_id}/model-access" | "/admin/organizations/{org_id}/limits/history" | "/admin/organizations/{org_id}/usage/balance" | "/admin/organizations/{org_id}/usage-discount" @@ -140,6 +142,24 @@ mod tests { } } + #[test] + fn model_access_listings_are_reads_and_grant_and_revoke_are_writes() { + for prefix in ["/v1", ""] { + for listing in [ + "/admin/models/{model_name}/access", + "/admin/organizations/{org_id}/model-access", + ] { + let path = format!("{prefix}{listing}"); + assert_eq!(admin_operation(&Method::GET, &path), AdminOperation::Read); + assert_eq!(admin_operation(&Method::PUT, &path), AdminOperation::Write); + } + let grant = format!("{prefix}/admin/models/{{model_name}}/access/{{org_id}}"); + for method in [Method::PUT, Method::DELETE, Method::GET] { + assert_eq!(admin_operation(&method, &grant), AdminOperation::Write); + } + } + } + #[test] fn api_keys_get_is_read() { for path in ["/v1/admin/api-keys", "/admin/api-keys"] { diff --git a/crates/api/src/middleware/model_deprecation.rs b/crates/api/src/middleware/model_deprecation.rs index 612d63a67..b36017757 100644 --- a/crates/api/src/middleware/model_deprecation.rs +++ b/crates/api/src/middleware/model_deprecation.rs @@ -4,7 +4,8 @@ //! //! - `x-model-deprecation-date`: the planned deprecation date, in the same //! form as `deprecation_date` on `GET /v1/models`. -//! - `x-model-successor`: the recommended replacement model, when one is set. +//! - `x-model-successor`: the recommended replacement model, when one is set +//! and the caller may use it. //! - `x-model-pricing-change-date`: the earliest pending pricing change, //! independently of whether the model has a planned deprecation. //! @@ -30,7 +31,9 @@ use axum::{ response::Response, }; use chrono::{DateTime, Utc}; -use services::models::{ModelWithPricing, ModelsServiceTrait}; +use services::models::{ModelCaller, ModelWithPricing, ModelsServiceTrait}; + +use super::auth::AuthenticatedApiKey; pub const HEADER_MODEL_DEPRECATION_DATE: &str = "x-model-deprecation-date"; pub const HEADER_MODEL_SUCCESSOR: &str = "x-model-successor"; @@ -119,6 +122,14 @@ pub async fn model_deprecation_middleware( ) -> Response { let requested = RequestedModel::default(); request.extensions_mut().insert(requested.clone()); + // Who is asking, read before the request moves on. The auth layer runs + // outside this one and has stored the API key's organization. + let caller = request + .extensions() + .get::() + .map_or(ModelCaller::Anonymous, |api_key| { + ModelCaller::Organization(api_key.organization.id.0) + }); let mut response = next.run(request).await; @@ -132,7 +143,35 @@ pub async fn model_deprecation_middleware( // ago. An alias resolves to its canonical model. Advisory only: a failed // or slow lookup (cache miss while the database is struggling) leaves the // response as it is instead of holding it back. - let lookup = state.models_service.resolve_and_get_model(model_name); + // + // Resolved for the request's own caller, never as a trusted one. A + // successful status does not mean the handler resolved the model: a + // streaming Responses request answers 200 first and fails inside the + // stream. The headers of a model the caller may not use would tell it + // that the model exists. + // + // The successor is another model, possibly an unlisted one, so it is + // named only if it resolves for the same caller. Both lookups share the + // time limit. + let lookup = async { + let mut model = state + .models_service + .resolve_and_get_model(model_name, caller) + .await?; + if model.deprecation_date.is_some() { + if let Some(successor) = model.successor_model_name.take() { + if state + .models_service + .resolve_and_get_model(&successor, caller) + .await + .is_ok() + { + model.successor_model_name = Some(successor); + } + } + } + Ok::<_, services::models::ModelsError>(model) + }; let Ok(Ok(model)) = tokio::time::timeout(LOOKUP_TIMEOUT, lookup).await else { return response; }; diff --git a/crates/api/src/models.rs b/crates/api/src/models.rs index 48979566a..97d7e984b 100644 --- a/crates/api/src/models.rs +++ b/crates/api/src/models.rs @@ -1155,6 +1155,26 @@ pub struct ModelsResponse { pub data: Vec, } +/// `GET /v1/internal/model-access`: which organizations may use each +/// restricted model. A model that is absent from `models` is unrestricted. +#[derive(Debug, Serialize, Deserialize)] +pub struct ModelAccessDocument { + /// Always `"model_access"`. + pub object: String, + /// One entry per unlisted model, active or not, ordered by model name. + pub models: Vec, +} + +/// One restricted model of [`ModelAccessDocument`]. +#[derive(Debug, Serialize, Deserialize)] +pub struct ModelAccessDocumentEntry { + /// Canonical model name, the `id` of `GET /v1/internal/models`. + pub id: String, + /// The active organizations that hold a grant. Empty when nobody may use + /// the model. + pub organization_ids: Vec, +} + /// Model pricing information. /// /// `input` / `output` are the legacy HuggingFace-style fields (USD per million @@ -1286,7 +1306,8 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub deprecation_date: Option, /// Recommended replacement for a model with a planned deprecation - /// (canonical model id). Omitted when none is set. + /// (canonical model id). Omitted when none is set. The public catalogs + /// name only a listed successor. #[serde(skip_serializing_if = "Option::is_none")] pub successor_model_id: Option, /// Human-readable description (OpenRouter `description`). @@ -4121,7 +4142,8 @@ pub struct ModelMetadata { #[serde(rename = "deprecationDate", skip_serializing_if = "Option::is_none")] pub deprecation_date: Option, /// Recommended replacement for a model with a planned deprecation - /// (canonical model id). Omitted when none is set. + /// (canonical model id). Omitted when none is set. The public catalogs + /// name only a listed successor. #[serde(rename = "successorModelId", skip_serializing_if = "Option::is_none")] pub successor_model_id: Option, /// OpenRouter `openrouter.slug` override (lowercase `author/slug`). Omitted @@ -4129,8 +4151,9 @@ pub struct ModelMetadata { /// `openrouter: { slug }` object; the admin view exposes the raw value. #[serde(rename = "openrouterSlug", skip_serializing_if = "Option::is_none")] pub openrouter_slug: Option, - /// Whether the model appears in the public model catalogs. Set on admin - /// views only; the public catalog endpoints omit it. + /// Whether the model is listed. An unlisted model is left out of the + /// public model catalogs and only organizations that hold a grant may + /// use it. Set on admin views only; the public catalog endpoints omit it. #[serde(rename = "isListed", skip_serializing_if = "Option::is_none")] pub is_listed: Option, } @@ -4185,11 +4208,14 @@ pub struct UpdateModelApiRequest { /// Required to be set explicitly when activating a zero-price model. #[serde(rename = "allowFree", skip_serializing_if = "Option::is_none")] pub allow_free: Option, - /// `false` keeps the model out of the public model catalogs - /// (`GET /v1/models`, `GET /v1/model/list`, `GET /v1/model/{name}`) while - /// it stays active and billable; `true` lists it again. Omitted = leave - /// unchanged (new models are listed). It is a catalog switch only: a - /// model that has an inference URL can still be called by name. + /// `false` unlists the model. It stays active and billable, but it + /// leaves the public model catalogs (`GET /v1/models`, + /// `GET /v1/model/list`, `GET /v1/model/{name}`) and becomes restricted: + /// only organizations that hold a grant + /// (`PUT /v1/admin/models/{model_name}/access/{org_id}`) may use it, and + /// everyone else is answered as if the model did not exist. `true` lists + /// it again and opens it to every caller. Omitted = leave unchanged (new + /// models are listed). #[serde(rename = "isListed", skip_serializing_if = "Option::is_none")] pub is_listed: Option, pub aliases: Option>, @@ -4664,7 +4690,8 @@ pub struct ModelHistoryEntry { /// If true, this model was allowed to serve without pricing at this point in time. #[serde(rename = "allowFree")] pub allow_free: bool, - /// Whether the model was in the public model catalogs at this point in time. + /// Whether the model was listed (in the public model catalogs and open to + /// every caller) at this point in time. #[serde(rename = "isListed")] pub is_listed: bool, } diff --git a/crates/api/src/openapi.rs b/crates/api/src/openapi.rs index 66739fa15..a14bbb35c 100644 --- a/crates/api/src/openapi.rs +++ b/crates/api/src/openapi.rs @@ -167,6 +167,10 @@ use utoipa::{Modify, OpenApi}; crate::routes::admin::list_model_pricing_changes, crate::routes::admin::cancel_model_pricing_change, crate::routes::admin::get_model_history, + crate::routes::admin_model_access::get_model_access, + crate::routes::admin_model_access::grant_model_access, + crate::routes::admin_model_access::revoke_model_access, + crate::routes::admin_model_access::list_organization_model_access, crate::routes::admin::get_admin_organization_balance, crate::routes::admin::update_organization_limits, crate::routes::admin::get_organization_limits_history, @@ -256,6 +260,10 @@ use utoipa::{Modify, OpenApi}; UpdateOrganizationRequest, CreateApiKeyRequest, ApiKeyResponse, UpdateApiKeySpendLimitRequest, UpdateApiKeyRequest, // usage_hourly repair models + crate::routes::admin_model_access::ModelAccessResponse, + crate::routes::admin_model_access::ModelAccessGrantResponse, + crate::routes::admin_model_access::OrganizationModelAccessResponse, + crate::routes::admin_model_access::OrganizationModelAccessEntry, crate::routes::admin_usage_discount::UsageDiscountRequest, crate::routes::admin_credit_reconciliation::CreditReconciliationRequest, crate::routes::admin_credit_reconciliation::CreditReconciliationResponse, diff --git a/crates/api/src/routes/admin_model_access.rs b/crates/api/src/routes/admin_model_access.rs new file mode 100644 index 000000000..ddaeae34d --- /dev/null +++ b/crates/api/src/routes/admin_model_access.rs @@ -0,0 +1,344 @@ +//! Administrator-managed access to unlisted models. +//! +//! An unlisted model (`isListed: false`) is restricted: only an organization +//! that holds a grant may use it, and everyone else is answered as if the +//! model did not exist. A grant is an (organization, model) pair. Grants on a +//! listed model have no effect, so access can be prepared before the model is +//! unlisted. A grant counts only while the organization and the model are +//! both active. +//! +//! A grant or a revocation applies at once on the instance that served the +//! request. Other instances follow within the model-resolve cache TTL (30 +//! seconds). Each one that changes something is recorded in +//! `organization_model_access_log`. +use crate::{middleware::AdminUser, models::ErrorResponse, routes::admin::AdminAppState}; +use axum::{ + extract::{Extension, Path, State}, + http::StatusCode, + Json, +}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use services::admin::{AdminError, ModelAccessGrant, OrganizationModelAccessGrant}; +use tracing::{debug, error}; +use utoipa::ToSchema; +use uuid::Uuid; + +/// A model and the organizations that hold a grant for it. +#[derive(Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct ModelAccessResponse { + /// The model's id. + pub id: Uuid, + /// The model's canonical name. + pub model_name: String, + /// `false`: the model is restricted to the organizations below. `true`: + /// every organization may use it and the grants have no effect. + pub is_listed: bool, + /// Grants count only while the model is active. + pub is_active: bool, + /// The active organizations that hold a grant, oldest grant first. + pub organizations: Vec, +} + +/// One organization's grant for a model. +#[derive(Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct ModelAccessGrantResponse { + pub organization_id: Uuid, + pub organization_name: String, + pub granted_at: DateTime, + /// The admin who made the grant. + pub granted_by_user_id: Option, + pub granted_by_user_email: Option, +} + +impl From for ModelAccessGrantResponse { + fn from(grant: ModelAccessGrant) -> Self { + Self { + organization_id: grant.organization_id, + organization_name: grant.organization_name, + granted_at: grant.granted_at, + granted_by_user_id: grant.granted_by_user_id, + granted_by_user_email: grant.granted_by_user_email, + } + } +} + +/// The models an organization holds a grant for. +#[derive(Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct OrganizationModelAccessResponse { + pub organization_id: Uuid, + /// Ordered by model name. + pub models: Vec, +} + +/// One model an organization holds a grant for. +#[derive(Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct OrganizationModelAccessEntry { + /// The model's id. + pub id: Uuid, + /// The model's canonical name. + pub model_name: String, + /// `true`: every organization may use the model, so this grant has no + /// effect until the model is unlisted. + pub is_listed: bool, + /// The grant counts only while the model is active. + pub is_active: bool, + pub granted_at: DateTime, + /// The admin who made the grant. + pub granted_by_user_id: Option, + pub granted_by_user_email: Option, +} + +impl From for OrganizationModelAccessEntry { + fn from(grant: OrganizationModelAccessGrant) -> Self { + Self { + id: grant.model.id, + model_name: grant.model.model_name, + is_listed: grant.model.is_listed, + is_active: grant.model.is_active, + granted_at: grant.granted_at, + granted_by_user_id: grant.granted_by_user_id, + granted_by_user_email: grant.granted_by_user_email, + } + } +} + +type ApiError = (StatusCode, Json); + +fn parse_organization_id(org_id: &str) -> Result { + Uuid::parse_str(org_id).map_err(|_| { + ( + StatusCode::BAD_REQUEST, + Json(ErrorResponse::new( + "Invalid organization ID format".to_string(), + "invalid_id".to_string(), + )), + ) + }) +} + +fn api_error(error: AdminError) -> ApiError { + match error { + AdminError::ModelNotFound(message) => ( + StatusCode::NOT_FOUND, + Json(ErrorResponse::new(message, "model_not_found".to_string())), + ), + AdminError::OrganizationNotFound(message) => ( + StatusCode::NOT_FOUND, + Json(ErrorResponse::new( + message, + "organization_not_found".to_string(), + )), + ), + other => { + // Discriminant only: the message can carry database error text. + error!( + error_kind = ?std::mem::discriminant(&other), + "Model access operation failed" + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(ErrorResponse::new( + "Failed to read or change model access".to_string(), + "internal_server_error".to_string(), + )), + ) + } + } +} + +/// List the organizations that may use a model (Admin only) +/// +/// Returns the model's id, name and `isListed`, and the active organizations +/// that hold a grant for it. The grants restrict the model only while it is +/// unlisted; on a listed model they are kept but have no effect. +/// +/// **Note:** Model names containing forward slashes must be URL-encoded +/// (e.g. "vendor%2Fmodel"). +#[utoipa::path( + get, + path = "/v1/admin/models/{model_name}/access", + tag = "Admin", + params( + ("model_name" = String, Path, description = "Canonical model name (URL-encode if it contains slashes)") + ), + responses( + (status = 200, description = "The model and its grants", body = ModelAccessResponse), + (status = 401, description = "Unauthorized", body = ErrorResponse), + (status = 403, description = "Forbidden: not an admin", body = ErrorResponse), + (status = 404, description = "Model not found", body = ErrorResponse), + (status = 500, description = "Internal server error", body = ErrorResponse) + ), + security( + ("session_token" = []) + ) +)] +pub async fn get_model_access( + State(app_state): State, + Path(model_name): Path, + Extension(_admin_user): Extension, +) -> Result, ApiError> { + debug!("Get model access request for: {}", model_name); + + let access = app_state + .admin_service + .get_model_access(&model_name) + .await + .map_err(api_error)?; + + Ok(Json(ModelAccessResponse { + id: access.model.id, + model_name: access.model.model_name, + is_listed: access.model.is_listed, + is_active: access.model.is_active, + organizations: access.grants.into_iter().map(Into::into).collect(), + })) +} + +/// Grant an organization access to a model (Admin only) +/// +/// Idempotent: granting again returns the existing grant unchanged. The grant +/// takes effect when the model is unlisted; while the model is listed every +/// organization may use it anyway. It applies at once on the instance that +/// served this request and within 30 seconds on the others. +/// +/// **Note:** Model names containing forward slashes must be URL-encoded +/// (e.g. "vendor%2Fmodel"). +#[utoipa::path( + put, + path = "/v1/admin/models/{model_name}/access/{org_id}", + tag = "Admin", + params( + ("model_name" = String, Path, description = "Canonical model name (URL-encode if it contains slashes)"), + ("org_id" = String, Path, description = "The organization's ID (as a UUID)") + ), + responses( + (status = 200, description = "The grant", body = ModelAccessGrantResponse), + (status = 400, description = "Invalid organization ID", body = ErrorResponse), + (status = 401, description = "Unauthorized", body = ErrorResponse), + (status = 403, description = "Forbidden: not an admin, or a read-only admin access token", body = ErrorResponse), + (status = 404, description = "Model or organization not found", body = ErrorResponse), + (status = 500, description = "Internal server error", body = ErrorResponse) + ), + security( + ("session_token" = []) + ) +)] +pub async fn grant_model_access( + State(app_state): State, + Path((model_name, org_id)): Path<(String, String)>, + Extension(admin_user): Extension, +) -> Result, ApiError> { + debug!( + "Grant model access request: model={}, org_id={}", + model_name, org_id + ); + let organization_id = parse_organization_id(&org_id)?; + + let grant = app_state + .admin_service + .grant_model_access(&model_name, organization_id, Some(admin_user.0.id)) + .await + .map_err(api_error)?; + + Ok(Json(grant.into())) +} + +/// Revoke an organization's access to a model (Admin only) +/// +/// Idempotent: revoking a grant that does not exist succeeds. If the model is +/// unlisted the organization loses access at once on the instance that served +/// this request and within 30 seconds on the others. +/// +/// **Note:** Model names containing forward slashes must be URL-encoded +/// (e.g. "vendor%2Fmodel"). +#[utoipa::path( + delete, + path = "/v1/admin/models/{model_name}/access/{org_id}", + tag = "Admin", + params( + ("model_name" = String, Path, description = "Canonical model name (URL-encode if it contains slashes)"), + ("org_id" = String, Path, description = "The organization's ID (as a UUID)") + ), + responses( + (status = 204, description = "The organization holds no grant for the model"), + (status = 400, description = "Invalid organization ID", body = ErrorResponse), + (status = 401, description = "Unauthorized", body = ErrorResponse), + (status = 403, description = "Forbidden: not an admin, or a read-only admin access token", body = ErrorResponse), + (status = 404, description = "Model or organization not found", body = ErrorResponse), + (status = 500, description = "Internal server error", body = ErrorResponse) + ), + security( + ("session_token" = []) + ) +)] +pub async fn revoke_model_access( + State(app_state): State, + Path((model_name, org_id)): Path<(String, String)>, + Extension(admin_user): Extension, +) -> Result { + debug!( + "Revoke model access request: model={}, org_id={}", + model_name, org_id + ); + let organization_id = parse_organization_id(&org_id)?; + + app_state + .admin_service + .revoke_model_access(&model_name, organization_id, Some(admin_user.0.id)) + .await + .map_err(api_error)?; + + Ok(StatusCode::NO_CONTENT) +} + +/// List the models an organization holds grants for (Admin only) +/// +/// Each entry carries the model's `isListed` and `isActive`: a grant restricts +/// nothing while its model is listed and counts for nothing while its model is +/// inactive. +#[utoipa::path( + get, + path = "/v1/admin/organizations/{org_id}/model-access", + tag = "Admin", + params( + ("org_id" = String, Path, description = "The organization's ID (as a UUID)") + ), + responses( + (status = 200, description = "The organization's grants", body = OrganizationModelAccessResponse), + (status = 400, description = "Invalid organization ID", body = ErrorResponse), + (status = 401, description = "Unauthorized", body = ErrorResponse), + (status = 403, description = "Forbidden: not an admin", body = ErrorResponse), + (status = 404, description = "Organization not found", body = ErrorResponse), + (status = 500, description = "Internal server error", body = ErrorResponse) + ), + security( + ("session_token" = []) + ) +)] +pub async fn list_organization_model_access( + State(app_state): State, + Path(org_id): Path, + Extension(_admin_user): Extension, +) -> Result, ApiError> { + debug!( + "List organization model access request for org_id: {}", + org_id + ); + let organization_id = parse_organization_id(&org_id)?; + + let grants = app_state + .admin_service + .list_organization_model_access(organization_id) + .await + .map_err(api_error)?; + + Ok(Json(OrganizationModelAccessResponse { + organization_id, + models: grants.into_iter().map(Into::into).collect(), + })) +} diff --git a/crates/api/src/routes/anthropic.rs b/crates/api/src/routes/anthropic.rs index 132de3204..534244cc9 100644 --- a/crates/api/src/routes/anthropic.rs +++ b/crates/api/src/routes/anthropic.rs @@ -32,7 +32,7 @@ use inference_providers::{ }; use services::completions::hash_inference_id_to_uuid; use services::completions::ports::{CompletionError, ConcurrentRequestGuard}; -use services::models::{ModelWithPricing, ModelsError, ModelsServiceTrait}; +use services::models::{ModelCaller, ModelWithPricing, ModelsError, ModelsServiceTrait}; use services::usage::{ five_minute_cache_write_rate, CacheWriteBilling, InferenceType, ProviderAttribution, RecordUsageServiceRequest, StopReason, UsageServiceTrait, @@ -466,6 +466,7 @@ async fn handle_request( let (model, alias_from) = match resolve_model( app_state.models_service.as_ref(), + ModelCaller::Organization(api_key.organization.id.0), &headers, &prepared.requested_model, ) @@ -916,6 +917,7 @@ fn normalized_beta_header( async fn resolve_model( models_service: &dyn ModelsServiceTrait, + caller: ModelCaller, headers: &HeaderMap, requested: &str, ) -> Result<(ModelWithPricing, Option), Response> { @@ -925,10 +927,13 @@ async fn resolve_model( format!("anthropic/{requested}") }; - let model = match models_service.resolve_and_get_model(&normalized).await { + let model = match models_service + .resolve_and_get_model(&normalized, caller) + .await + { Ok(model) => model, Err(ModelsError::NotFound(_)) if normalized != requested => models_service - .resolve_and_get_model(requested) + .resolve_and_get_model(requested, caller) .await .map_err(model_resolution_error)?, Err(error) => return Err(model_resolution_error(error)), diff --git a/crates/api/src/routes/attestation.rs b/crates/api/src/routes/attestation.rs index 10a8a2669..5a95f1db0 100644 --- a/crates/api/src/routes/attestation.rs +++ b/crates/api/src/routes/attestation.rs @@ -6,9 +6,9 @@ //! //! | Route | Auth | Rationale | //! |---------------------------------|---------|-----------| -//! | `GET /v1/attestation/report` | API key | Data-plane endpoint, documented as key-protected. Key validation only — retrieval is not billed and never creates usage records. | +//! | `GET /v1/attestation/report` | API key | Data-plane endpoint, documented as key-protected. Key validation only — retrieval is not billed and never creates usage records. The key's organization decides whether an unlisted model resolves. | //! | `GET /v1/signature/{chat_id}` | API key | Returns per-completion signatures; completions are key-scoped, so lookups are too. | -//! | `GET /v1/attestation/ita-token` | Public | Deliberate exception — see `build_public_attestation_routes`. | +//! | `GET /v1/attestation/ita-token` | Public | Deliberate exception — see `build_public_attestation_routes`. Listed models only. | //! | `GET /v1/attestation/pins` | Public | Deliberate exception — see `build_public_attestation_routes`. | use crate::{ohttp_gateway::OhttpAttestation, routes::api::AppState}; diff --git a/crates/api/src/routes/attestation/alias.rs b/crates/api/src/routes/attestation/alias.rs index 39cc35ae4..c7890c934 100644 --- a/crates/api/src/routes/attestation/alias.rs +++ b/crates/api/src/routes/attestation/alias.rs @@ -3,19 +3,27 @@ use axum::{ http::{HeaderMap, StatusCode}, response::{Json as ResponseJson, Response}, }; -use services::models::ModelsServiceTrait; +use services::models::{ModelCaller, ModelsServiceTrait}; use std::sync::Arc; +/// The `(requested, canonical)` pair when `model` is an alias `caller` may +/// use. A model the caller may not use is `NotFound` here, so its canonical +/// name is never disclosed and the attestation service answers the request +/// like one for an unknown model. pub(super) async fn resolve_attestation_alias( model: Option<&str>, models_service: &Arc, + caller: ModelCaller, headers: &HeaderMap, ) -> Result, (StatusCode, ResponseJson)> { let Some(requested) = model else { return Ok(None); }; - match models_service.resolve_and_get_model(requested).await { + match models_service + .resolve_and_get_model(requested, caller) + .await + { Ok(m) if m.model_name != requested => { if crate::routes::common::no_aliasing_requested(headers) { return Err(( diff --git a/crates/api/src/routes/attestation/ita_token.rs b/crates/api/src/routes/attestation/ita_token.rs index 1e0c2f6dc..bd305ebbe 100644 --- a/crates/api/src/routes/attestation/ita_token.rs +++ b/crates/api/src/routes/attestation/ita_token.rs @@ -13,6 +13,7 @@ use axum::{ use services::attestation::ita::{ ItaGatewaySigningAlg as ServiceItaGatewaySigningAlg, ItaTokenQuery as ServiceItaTokenQuery, }; +use services::models::ModelCaller; /// Get Intel Trust Authority attestation token /// @@ -45,10 +46,15 @@ pub async fn get_ita_token( State(state): State, headers: HeaderMap, ) -> Result { - let alias_resolved = - resolve_attestation_alias(params.model.as_deref(), &state.models_service, &headers) - .await - .map_err(error_tuple_into_response)?; + // Public route: no organization is known, so only listed models resolve. + let alias_resolved = resolve_attestation_alias( + params.model.as_deref(), + &state.models_service, + ModelCaller::Anonymous, + &headers, + ) + .await + .map_err(error_tuple_into_response)?; let query = params .into_service_query() .map_err(error_tuple_into_response)?; diff --git a/crates/api/src/routes/attestation/report.rs b/crates/api/src/routes/attestation/report.rs index 003dd683e..5153e4932 100644 --- a/crates/api/src/routes/attestation/report.rs +++ b/crates/api/src/routes/attestation/report.rs @@ -3,14 +3,18 @@ use super::{ errors::{attestation_report_error_response, error_response}, AttestationRouteState, }; -use crate::{models::ErrorResponse, ohttp_gateway::OhttpAttestation}; +use crate::{ + middleware::auth::AuthenticatedApiKey, models::ErrorResponse, ohttp_gateway::OhttpAttestation, +}; use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::Json as ResponseJson, + Extension, }; use inference_providers::{ProviderFilter, ProviderSource}; use serde::{Deserialize, Serialize}; +use services::models::ModelCaller; use utoipa::{IntoParams, ToSchema}; /// Query parameters for attestation report @@ -155,7 +159,7 @@ impl From for AttestationRespo ), responses( (status = 200, description = "Attestation report retrieved", body = AttestationResponse), - (status = 400, description = "Invalid nonce, signing_algo or provider, or unknown model", body = ErrorResponse), + (status = 400, description = "Invalid nonce, signing_algo or provider, or unknown model (a model the key's organization may not use is reported as unknown)", body = ErrorResponse), (status = 401, description = "Missing or invalid API key", body = ErrorResponse), (status = 503, description = "Service unavailable", body = ErrorResponse) ), @@ -167,10 +171,18 @@ impl From for AttestationRespo pub async fn get_attestation_report( Query(params): Query, State(state): State, + Extension(api_key): Extension, headers: HeaderMap, ) -> Result)> { - let alias_resolved = - resolve_attestation_alias(params.model.as_deref(), &state.models_service, &headers).await?; + // An unlisted model has a report only for an organization that may use it. + let caller = ModelCaller::Organization(api_key.organization.id.0); + let alias_resolved = resolve_attestation_alias( + params.model.as_deref(), + &state.models_service, + caller, + &headers, + ) + .await?; let provider_filter = match params .provider .as_deref() @@ -202,6 +214,7 @@ pub async fn get_attestation_report( params.signing_address, params.include_tls_fingerprint.unwrap_or(false), provider_filter, + caller, ) .await .map_err(attestation_report_error_response)?; diff --git a/crates/api/src/routes/attestation/tests/ita_token_route_test_support.rs b/crates/api/src/routes/attestation/tests/ita_token_route_test_support.rs index 7f9499719..cbf565445 100644 --- a/crates/api/src/routes/attestation/tests/ita_token_route_test_support.rs +++ b/crates/api/src/routes/attestation/tests/ita_token_route_test_support.rs @@ -13,7 +13,7 @@ use services::{ }, AttestationError, SignatureLookupResult, }, - models::{ModelInfo, ModelWithPricing, ModelsError, ModelsServiceTrait}, + models::{ModelCaller, ModelInfo, ModelWithPricing, ModelsError, ModelsServiceTrait}, }; use std::sync::{Arc, Mutex}; use uuid::Uuid; @@ -99,6 +99,7 @@ impl services::attestation::ports::AttestationServiceTrait for RecordingItaAttes _signing_address: Option, _include_tls_fingerprint: bool, _provider_filter: Option, + _caller: ModelCaller, ) -> Result { Err(AttestationError::InternalError("unused".to_string())) } @@ -127,14 +128,33 @@ impl services::attestation::ports::AttestationServiceTrait for RecordingItaAttes #[derive(Clone, Default)] pub(super) struct TestModelsService { alias: Option<(String, String)>, + /// When set, the alias target is an unlisted model. + unlisted: bool, + callers: Arc>>, } impl TestModelsService { pub(super) fn with_alias(alias: &str, canonical: &str) -> Self { Self { alias: Some((alias.to_string(), canonical.to_string())), + ..Self::default() } } + + pub(super) fn with_unlisted_alias(alias: &str, canonical: &str) -> Self { + Self { + unlisted: true, + ..Self::with_alias(alias, canonical) + } + } + + /// The callers the route resolved models for, in order. + pub(super) fn callers(&self) -> Vec { + self.callers + .lock() + .unwrap_or_else(|error| panic!("caller lock poisoned: {error}")) + .clone() + } } #[async_trait] @@ -147,24 +167,41 @@ impl ModelsServiceTrait for TestModelsService { Ok(Vec::new()) } - async fn get_model_by_name(&self, model_name: &str) -> Result { + async fn get_model_by_name( + &self, + model_name: &str, + _caller: ModelCaller, + ) -> Result { Ok(model_with_name(model_name)) } async fn resolve_and_get_model( &self, identifier: &str, + caller: ModelCaller, ) -> Result { - let canonical = self + self.callers + .lock() + .unwrap_or_else(|error| panic!("caller lock poisoned: {error}")) + .push(caller); + let alias_target = self .alias .as_ref() .filter(|(alias, _)| alias == identifier) - .map(|(_, canonical)| canonical.as_str()) - .unwrap_or(identifier); - Ok(model_with_name(canonical)) + .map(|(_, canonical)| canonical.as_str()); + if self.unlisted && alias_target.is_some() && caller != ModelCaller::Trusted { + return Err(ModelsError::NotFound(format!( + "Model '{identifier}' not found" + ))); + } + Ok(model_with_name(alias_target.unwrap_or(identifier))) } - async fn resolve_alias_cached(&self, _identifier: &str) -> Option { + async fn resolve_alias_cached( + &self, + _identifier: &str, + _caller: ModelCaller, + ) -> Option { None } diff --git a/crates/api/src/routes/attestation/tests/ita_token_route_tests.rs b/crates/api/src/routes/attestation/tests/ita_token_route_tests.rs index 9bbbf8e80..922de988f 100644 --- a/crates/api/src/routes/attestation/tests/ita_token_route_tests.rs +++ b/crates/api/src/routes/attestation/tests/ita_token_route_tests.rs @@ -150,6 +150,39 @@ async fn ita_token_route_rejects_alias_when_no_aliasing_requested() { assert!(attestation.queries().is_empty()); } +/// The route is public: it must ask for the model as an anonymous caller, and +/// must not name the canonical model behind an alias that caller cannot use. +#[tokio::test] +async fn ita_token_route_resolves_the_model_as_an_anonymous_caller() { + let attestation = RecordingItaAttestationService::err(AttestationError::UnknownModel( + "alias-model".to_string(), + )); + let models = TestModelsService::with_unlisted_alias("alias-model", "canonical-model"); + let server = public_ita_server(attestation.clone(), models.clone()); + + let response = server + .get("/attestation/ita-token") + .add_query_param("model", "alias-model") + .add_header(crate::routes::common::HEADER_NO_ALIASING, "true") + .await; + + assert_eq!( + models.callers(), + vec![services::models::ModelCaller::Anonymous] + ); + // Not the alias rejection, which would name the canonical model: the + // request reaches the service and gets its unknown-model answer. + assert_eq!(response.status_code(), StatusCode::BAD_REQUEST); + assert!(response.maybe_header(HEADER_MODEL_ALIAS_RESOLVED).is_none()); + let body = response.text(); + assert!(!body.contains("canonical-model"), "{body}"); + assert!(body.contains("Model 'alias-model' not found"), "{body}"); + assert_eq!( + attestation.only_query().model.as_deref(), + Some("alias-model") + ); +} + #[tokio::test] async fn ita_token_route_maps_service_errors_to_http_contract() { for scenario in [ diff --git a/crates/api/src/routes/completions/audio.rs b/crates/api/src/routes/completions/audio.rs index 0214ee60d..8c39ec585 100644 --- a/crates/api/src/routes/completions/audio.rs +++ b/crates/api/src/routes/completions/audio.rs @@ -154,7 +154,10 @@ pub async fn audio_transcriptions( // Resolve model to get UUID for usage tracking (handles aliases like chat_completions) let model = match app_state .models_service - .resolve_and_get_model(&request.model) + .resolve_and_get_model( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/chat.rs b/crates/api/src/routes/completions/chat.rs index 77a675311..95cf888a5 100644 --- a/crates/api/src/routes/completions/chat.rs +++ b/crates/api/src/routes/completions/chat.rs @@ -217,9 +217,17 @@ pub(super) async fn chat_completions_inner( let e2ee_active = e2ee_requested(&encryption_headers); let include_stream_usage_in_response = chat_stream_include_usage_requested(&request); + // The catalog reads below answer for the requested model before the + // completion service has resolved it, so each one states who is asking. + // For an unlisted model this organization may not use they find nothing, + // and the request reaches the completion service like one for an unknown + // model and gets its error. + let caller = ModelCaller::Organization(api_key.organization.id.0); + // Strict alias mode: refuse to serve through an alias before any // inference happens (issue #573). - if let Err(resp) = reject_if_aliased(&app_state.models_service, &headers, &request.model).await + if let Err(resp) = + reject_if_aliased(&app_state.models_service, caller, &headers, &request.model).await { return resp; } @@ -234,7 +242,7 @@ pub(super) async fn chat_completions_inner( // path); advisory only — strict mode above stays authoritative. let alias_canonical = app_state .models_service - .resolve_alias_cached(&request.model) + .resolve_alias_cached(&request.model, caller) .await; let resolved_model_name = alias_canonical.as_deref().unwrap_or(&request.model); let (model_attestation_supported, model_input_modalities) = @@ -242,21 +250,41 @@ pub(super) async fn chat_completions_inner( // Exact catalog name first, then the alias target: a name that is // both a model and another model's alias must read its own // capabilities. - Ok(models) => models - .iter() - .find(|model| model.model_name == request.model) - .or_else(|| { - models - .iter() - .find(|model| model.model_name == resolved_model_name) - }) - .map(|model| { - ( - Some(model.attestation_supported), - model.input_modalities.clone(), - ) - }) - .unwrap_or((None, None)), + Ok(models) => { + let model = models + .iter() + .find(|model| model.model_name == request.model) + .or_else(|| { + models + .iter() + .find(|model| model.model_name == resolved_model_name) + }); + // The list holds every active model. Read the row only if + // this organization may use the model (no cost for a listed + // one); a failed check hides it too. Otherwise the modality + // check below would answer for a model the caller must not + // learn about. + let model = match model { + Some(model) + if app_state + .models_service + .caller_can_use_model(caller, model) + .await + .unwrap_or(false) => + { + Some(model) + } + _ => None, + }; + model + .map(|model| { + ( + Some(model.attestation_supported), + model.input_modalities.clone(), + ) + }) + .unwrap_or((None, None)) + } Err(error) => { tracing::warn!( model = %request.model, diff --git a/crates/api/src/routes/completions/chat_auto_redact.rs b/crates/api/src/routes/completions/chat_auto_redact.rs index 36dd98c90..760ecdbde 100644 --- a/crates/api/src/routes/completions/chat_auto_redact.rs +++ b/crates/api/src/routes/completions/chat_auto_redact.rs @@ -113,9 +113,11 @@ pub(super) async fn bill_auto_redact_classify( }, ); + // A trusted lookup: the platform chose this model and has already run the + // classify pass, so the usage is billed whoever the organization is. let model = match app_state .models_service - .get_model_by_name(auto_redact::DEFAULT_PII_MODEL) + .get_model_by_name(auto_redact::DEFAULT_PII_MODEL, ModelCaller::Trusted) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/embeddings.rs b/crates/api/src/routes/completions/embeddings.rs index b34416d55..467eb2902 100644 --- a/crates/api/src/routes/completions/embeddings.rs +++ b/crates/api/src/routes/completions/embeddings.rs @@ -136,7 +136,10 @@ pub async fn embeddings( // Resolve model to get UUID for usage tracking let model = match app_state .models_service - .get_model_by_name(&model_name) + .get_model_by_name( + &model_name, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/images.rs b/crates/api/src/routes/completions/images.rs index 9951fe932..ce9a16b32 100644 --- a/crates/api/src/routes/completions/images.rs +++ b/crates/api/src/routes/completions/images.rs @@ -59,7 +59,10 @@ pub async fn image_generations( // Resolve model to get UUID for usage tracking let model = match app_state .models_service - .get_model_by_name(&request.model) + .get_model_by_name( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, @@ -457,7 +460,10 @@ pub async fn image_edits( // Resolve model to get UUID for usage tracking let model = match app_state .models_service - .get_model_by_name(&request.model) + .get_model_by_name( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/mod.rs b/crates/api/src/routes/completions/mod.rs index 1d09175c1..a29539d16 100644 --- a/crates/api/src/routes/completions/mod.rs +++ b/crates/api/src/routes/completions/mod.rs @@ -26,6 +26,7 @@ use services::completions::{ ports::{CompletionMessage, CompletionRequest as ServiceCompletionRequest}, MAX_PROVIDER_DONE_EVENT_BYTES, }; +use services::models::ModelCaller; use sha2::{Digest, Sha256}; use std::convert::Infallible; use std::sync::Arc; @@ -155,16 +156,22 @@ fn e2ee_requested(encryption_headers: &crate::routes::common::EncryptionHeaders) /// clients, before a payload is bound to a different model TD's signing key. /// /// Unknown models fall through (`Ok`) so the completion service returns its -/// canonical "not a valid model name or alias" error. +/// canonical "not a valid model name or alias" error. So does an alias of a +/// model `caller` may not use: naming its canonical model here would disclose +/// it. async fn reject_if_aliased( models_service: &Arc, + caller: ModelCaller, headers: &header::HeaderMap, requested_model: &str, ) -> Result<(), Response> { if !no_aliasing_requested(headers) { return Ok(()); } - match models_service.resolve_and_get_model(requested_model).await { + match models_service + .resolve_and_get_model(requested_model, caller) + .await + { Ok(m) if m.model_name != requested_model => Err(( StatusCode::BAD_REQUEST, ResponseJson(ErrorResponse::new( diff --git a/crates/api/src/routes/completions/models.rs b/crates/api/src/routes/completions/models.rs index 244cff362..bf76dce30 100644 --- a/crates/api/src/routes/completions/models.rs +++ b/crates/api/src/routes/completions/models.rs @@ -45,8 +45,9 @@ pub async fn models( /// /// The document has the shape of `GET /v1/models` and adds the models whose /// `is_listed` is false. A gateway lane reads its catalog here, so it can -/// serve and advertise a model that this API keeps out of its own public -/// catalog. +/// serve a model that this API keeps out of its own public catalog. Which +/// organizations may use an unlisted model is in +/// `GET /v1/internal/model-access`. /// /// Auth: `Authorization: Bearer `, the service secret /// of `POST /v1/internal/usage`. The endpoint answers 503 until that token is @@ -91,6 +92,71 @@ pub async fn models_internal( )) } +/// List the restricted models and who may use them, for trusted services. +/// +/// An unlisted model is restricted: only organizations that hold a grant may +/// use it. This document lists every unlisted model with the ids of the +/// active organizations that hold a grant for it (an empty list when there +/// are none). A model that is absent from it is unrestricted, so a model +/// that is inactive here is listed too: a lane may still be serving it. A +/// gateway lane polls the document and enforces access itself for the +/// requests it serves. It is read from the database on every call, so a +/// grant or a revocation shows up on the next poll of any instance. +/// +/// Each entry names a model by its canonical name only. A gateway that +/// accepts an alias must check the model the alias stands for (the aliases +/// are in `GET /v1/internal/models`); checking the alias itself would find +/// no entry and read the model as unrestricted. +/// +/// Auth: `Authorization: Bearer `, the service secret +/// of `POST /v1/internal/usage`. The endpoint answers 503 until that token is +/// configured. Not part of the public OpenAPI document. The answer is marked +/// `no-store`: it must not sit in a shared cache. +pub async fn model_access_internal( + State(app_state): State, + headers: header::HeaderMap, +) -> Result< + ( + [(header::HeaderName, &'static str); 1], + ResponseJson, + ), + (StatusCode, ResponseJson), +> { + crate::routes::usage::verify_internal_usage_token( + &headers, + app_state.config.internal_usage_token.as_deref(), + )?; + + let restricted = app_state + .models_service + .get_restricted_model_access() + .await + .map_err(|e| { + tracing::error!(error = %e, "Failed to list model access for the internal document"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + ResponseJson(ErrorResponse::new( + "Failed to retrieve model access".to_string(), + "internal_server_error".to_string(), + )), + ) + })?; + + Ok(( + [(header::CACHE_CONTROL, "no-store")], + ResponseJson(ModelAccessDocument { + object: "model_access".to_string(), + models: restricted + .into_iter() + .map(|model| ModelAccessDocumentEntry { + id: model.model_name, + organization_ids: model.organization_ids, + }) + .collect(), + }), + )) +} + pub(super) fn model_with_pricing_to_info(model: services::models::ModelWithPricing) -> ModelInfo { let pricing = CatalogPricingAmounts { input_cost_per_token: model.input_cost_per_token, diff --git a/crates/api/src/routes/completions/privacy.rs b/crates/api/src/routes/completions/privacy.rs index bc30b6cc1..f146dc297 100644 --- a/crates/api/src/routes/completions/privacy.rs +++ b/crates/api/src/routes/completions/privacy.rs @@ -325,7 +325,10 @@ pub async fn privacy_classify( let model = match app_state .models_service - .get_model_by_name(&model_name) + .get_model_by_name( + &model_name, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, @@ -568,7 +571,10 @@ pub async fn privacy_redact( let model = match app_state .models_service - .get_model_by_name(&model_name) + .get_model_by_name( + &model_name, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/rerank.rs b/crates/api/src/routes/completions/rerank.rs index 0895056f8..1a412f753 100644 --- a/crates/api/src/routes/completions/rerank.rs +++ b/crates/api/src/routes/completions/rerank.rs @@ -175,7 +175,10 @@ pub async fn rerank( // Models must be registered in the database to be available for use let model = match app_state .models_service - .get_model_by_name(&request.model) + .get_model_by_name( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/score.rs b/crates/api/src/routes/completions/score.rs index 0562f5f82..895a55968 100644 --- a/crates/api/src/routes/completions/score.rs +++ b/crates/api/src/routes/completions/score.rs @@ -52,7 +52,10 @@ pub async fn score( // Models must be registered in the database to be available for use let model = match app_state .models_service - .get_model_by_name(&request.model) + .get_model_by_name( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/src/routes/completions/text.rs b/crates/api/src/routes/completions/text.rs index 8a837507c..fc0b1b484 100644 --- a/crates/api/src/routes/completions/text.rs +++ b/crates/api/src/routes/completions/text.rs @@ -257,13 +257,15 @@ pub(super) async fn completions_inner( // service resolves aliases for this endpoint exactly like chat, so it // gets the same contract: x-no-aliasing rejects, and aliased responses // carry the warning + x-model-alias-resolved header. - if let Err(resp) = reject_if_aliased(&app_state.models_service, &headers, &request.model).await + let caller = ModelCaller::Organization(api_key.organization.id.0); + if let Err(resp) = + reject_if_aliased(&app_state.models_service, caller, &headers, &request.model).await { return resp; } let alias_canonical = app_state .models_service - .resolve_alias_cached(&request.model) + .resolve_alias_cached(&request.model, caller) .await; let resolved_model_name = alias_canonical.as_deref().unwrap_or(&request.model); diff --git a/crates/api/src/routes/mod.rs b/crates/api/src/routes/mod.rs index 3cdbcf7ba..24eee2254 100644 --- a/crates/api/src/routes/mod.rs +++ b/crates/api/src/routes/mod.rs @@ -1,5 +1,6 @@ pub mod admin; pub mod admin_credit_reconciliation; +pub mod admin_model_access; pub mod admin_usage_hourly; pub mod anthropic; pub mod api; diff --git a/crates/api/src/routes/models.rs b/crates/api/src/routes/models.rs index 8ea88291b..f36563f55 100644 --- a/crates/api/src/routes/models.rs +++ b/crates/api/src/routes/models.rs @@ -365,6 +365,7 @@ mod tests { async fn get_model_by_name( &self, model_name: &str, + _caller: services::models::ModelCaller, ) -> Result { Err(ModelsError::NotFound(format!( "Model '{model_name}' not found" @@ -374,6 +375,7 @@ mod tests { async fn resolve_and_get_model( &self, _identifier: &str, + _caller: services::models::ModelCaller, ) -> Result { self.db_calls.fetch_add(1, Ordering::SeqCst); Ok(service_model("db-backed/detail", Some(512))) @@ -388,7 +390,11 @@ mod tests { Ok(service_model("public/enriched-detail", Some(4_096))) } - async fn resolve_alias_cached(&self, _identifier: &str) -> Option { + async fn resolve_alias_cached( + &self, + _identifier: &str, + _caller: services::models::ModelCaller, + ) -> Option { None } diff --git a/crates/api/src/routes/responses.rs b/crates/api/src/routes/responses.rs index ed5fcdc32..eca48a227 100644 --- a/crates/api/src/routes/responses.rs +++ b/crates/api/src/routes/responses.rs @@ -268,7 +268,13 @@ pub async fn create_response( requested_model.set(model); } let native = if let Some(body) = body.as_ref() { - state.native_service.selected_model(body).await + state + .native_service + .selected_model( + body, + services::models::ModelCaller::Organization(api_key.organization.id.0), + ) + .await } else { None }; diff --git a/crates/api/src/routes/systemone.rs b/crates/api/src/routes/systemone.rs index ca6f0ca77..e6701aec6 100644 --- a/crates/api/src/routes/systemone.rs +++ b/crates/api/src/routes/systemone.rs @@ -19,7 +19,7 @@ use inference_providers::{ use services::{ attestation::SignatureKind, completions::hash_inference_id_to_uuid, - models::ModelsError, + models::{ModelCaller, ModelsError}, usage::{InferenceType, ProviderAttribution, RecordUsageServiceRequest, StopReason}, }; use sha2::{Digest, Sha256}; @@ -72,7 +72,10 @@ pub async fn systemone( } let model = match state .models_service - .resolve_and_get_model(&request.model) + .resolve_and_get_model( + &request.model, + ModelCaller::Organization(api_key.organization.id.0), + ) .await { Ok(model) => model, diff --git a/crates/api/tests/e2e_all/admin_pricing_changes.rs b/crates/api/tests/e2e_all/admin_pricing_changes.rs index 6b3231a4b..81f908c3f 100644 --- a/crates/api/tests/e2e_all/admin_pricing_changes.rs +++ b/crates/api/tests/e2e_all/admin_pricing_changes.rs @@ -124,16 +124,22 @@ impl services::models::ModelsServiceTrait for NoopModelsService { async fn get_model_by_name( &self, _model_name: &str, + _caller: services::models::ModelCaller, ) -> Result { unimplemented!() } async fn resolve_and_get_model( &self, _identifier: &str, + _caller: services::models::ModelCaller, ) -> Result { unimplemented!() } - async fn resolve_alias_cached(&self, _identifier: &str) -> Option { + async fn resolve_alias_cached( + &self, + _identifier: &str, + _caller: services::models::ModelCaller, + ) -> Option { None } async fn get_configured_model_names( diff --git a/crates/api/tests/e2e_all/main.rs b/crates/api/tests/e2e_all/main.rs index fc451847e..fb9cf365f 100644 --- a/crates/api/tests/e2e_all/main.rs +++ b/crates/api/tests/e2e_all/main.rs @@ -86,6 +86,7 @@ mod openrouter_params; mod org_system_prompt; mod organization_deletion; mod organization_membership_order; +mod organization_model_access; mod organization_priority; mod organization_usage_discount; mod pagination_validation; diff --git a/crates/api/tests/e2e_all/model_pricing_notices.rs b/crates/api/tests/e2e_all/model_pricing_notices.rs index ec2e6f3e1..fb28aaa0f 100644 --- a/crates/api/tests/e2e_all/model_pricing_notices.rs +++ b/crates/api/tests/e2e_all/model_pricing_notices.rs @@ -2,7 +2,7 @@ use crate::common::*; use chrono::{Duration, Timelike, Utc}; use futures::FutureExt; use serde_json::{json, Value}; -use services::models::ModelsServiceTrait; +use services::models::{ModelCaller, ModelsServiceTrait}; use std::{panic::AssertUnwindSafe, sync::Arc}; const HEADER: &str = "x-model-pricing-change-date"; @@ -342,7 +342,7 @@ async fn serial_scheduler_replaces_the_notice_with_the_announced_tiered_prices() } let resolved = fixture .models - .resolve_and_get_model(&fixture.alias) + .resolve_and_get_model(&fixture.alias, ModelCaller::Trusted) .await .unwrap(); assert!(resolved.upcoming_pricing.is_none()); @@ -366,7 +366,7 @@ async fn assert_invalid_pending_profile_is_advisory(profile: Value) { // Exercise the canonical-only repository path as well as alias resolution. let model = fixture .models - .get_model_by_name(&fixture.model) + .get_model_by_name(&fixture.model, ModelCaller::Trusted) .await .unwrap(); assert!(model.upcoming_pricing.is_none()); diff --git a/crates/api/tests/e2e_all/organization_model_access.rs b/crates/api/tests/e2e_all/organization_model_access.rs new file mode 100644 index 000000000..d1162c865 --- /dev/null +++ b/crates/api/tests/e2e_all/organization_model_access.rs @@ -0,0 +1,2034 @@ +// E2E tests for per-organization access to unlisted models. +// +// An unlisted model (`isListed: false`) is restricted: an organization may use +// it through this API only while it holds a grant. Everyone else must get +// exactly the answer a model that does not exist gets. Grants are managed by +// admins (`/v1/admin/models/{model_name}/access`) and published to trusted +// services (`GET /v1/internal/model-access`). + +use crate::common::{fake_ita::FakeIta, fake_ita::FakeItaMode, *}; +use api::models::BatchUpdateModelApiRequest; +use axum::http::Method; +use bytes::Bytes; +use inference_providers::{ + mock::MockProvider, responses_raw::ResponsesRawResponse, SystemOneResponseWithBytes, +}; +use serde_json::{json, Value}; +use std::collections::BTreeMap; +use std::sync::Arc; +use uuid::Uuid; + +const INTERNAL_TOKEN: &str = "organization-model-access-internal-secret"; +const NONCE: &str = "0000000000000000000000000000000000000000000000000000000000000001"; +/// Every fixture model has a planned deprecation, so a served request carries +/// the deprecation headers and a refused one must not. `2030-01-01` defaults +/// to 13:00 UTC. +const DEPRECATION_DATE: &str = "2030-01-01"; +const DEPRECATION_DATE_HEADER: &str = "2030-01-01T13:00:00Z"; +const HEADER_DEPRECATION_DATE: &str = "x-model-deprecation-date"; +const HEADER_SUCCESSOR: &str = "x-model-successor"; + +/// The fixture models, one per kind of route. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ModelKind { + /// A text model, with an alias. + Text, + /// A System One model (output modality `decisions`). + Decisions, + /// A model the native Anthropic Messages routes accept. + Anthropic, + /// A model on the native Responses allowlist. + NativeResponses, +} + +struct Fixture { + server: axum_test::TestServer, + database: Arc, + pool: Arc, + /// Serves every fixture model and records the chat requests it receives. + provider: Arc, + text: String, + text_alias: String, + decisions: String, + anthropic: String, + native: String, + /// The listed model every fixture model names as its successor. + successor: String, +} + +impl Fixture { + fn model(&self, kind: ModelKind) -> &str { + match kind { + ModelKind::Text => &self.text, + ModelKind::Decisions => &self.decisions, + ModelKind::Anthropic => &self.anthropic, + ModelKind::NativeResponses => &self.native, + } + } +} + +/// A name with the shape and length of `model` that no model has. The last 36 +/// characters of every fixture name are a UUID. +fn unknown_twin(model: &str) -> String { + format!("{}{}", &model[..model.len() - 36], Uuid::new_v4()) +} + +fn model_fields(extra: Value) -> Value { + let mut fields = json!({ + "inputCostPerToken": { "amount": 1_000_000, "currency": "USD" }, + "outputCostPerToken": { "amount": 2_000_000, "currency": "USD" }, + "modelDisplayName": "Model Access Test Model", + "modelDescription": "A model used by the model access e2e tests", + "contextLength": 4096, + "maxOutputLength": 1024, + "verifiable": false, + "isActive": true, + "isListed": false, + "inputModalities": ["text"], + "outputModalities": ["text"] + }); + for (key, value) in extra.as_object().expect("extra is an object") { + fields[key] = value.clone(); + } + fields +} + +async fn upsert_models(server: &axum_test::TestServer, models: Vec<(&str, Value)>) { + let mut batch = BatchUpdateModelApiRequest::new(); + for (name, fields) in models { + batch.insert( + name.to_string(), + serde_json::from_value(fields).expect("valid model update"), + ); + } + admin_batch_upsert_models(server, batch, get_session_id()).await; +} + +/// A server with one unlisted model of each kind, all served by one mock +/// provider. Each has a planned deprecation and names a listed successor. +async fn setup() -> Fixture { + let suffix = Uuid::new_v4(); + let text = format!("access-text/model-{suffix}"); + let text_alias = format!("access-text-alias-{suffix}"); + let decisions = format!("access-decisions/model-{suffix}"); + let anthropic = format!("anthropic/access-model-{suffix}"); + let native = format!("access-native/model-{suffix}"); + let successor = format!("access-successor/model-{suffix}"); + + let native_allowlist = vec![native.clone()]; + let (server, pool, _shared_mock, database) = + setup_test_server_with_pool_and_config(move |config| { + config.internal_usage_token = Some(INTERNAL_TOKEN.to_string()); + config.external_providers.enable_anthropic_messages = true; + config.native_responses_models = native_allowlist; + }) + .await; + + // The successor first: a model can only name one that exists. + upsert_models( + &server, + vec![(&successor, model_fields(json!({ "isListed": true })))], + ) + .await; + let deprecated = |extra: Value| { + let mut fields = model_fields(extra); + fields["deprecationDate"] = json!(DEPRECATION_DATE); + fields["successorModelId"] = json!(successor); + fields + }; + upsert_models( + &server, + vec![ + (&text, deprecated(json!({ "aliases": [text_alias] }))), + ( + &decisions, + deprecated(json!({ "outputModalities": ["decisions"] })), + ), + ( + &anthropic, + deprecated(json!({ + "providerType": "external", + "attestationSupported": false, + "providerConfig": { + "backend": "anthropic", + "base_url": "https://anthropic.invalid/v1" + } + })), + ), + (&native, deprecated(json!({ "ownedBy": "openai" }))), + ], + ) + .await; + + let provider = Arc::new( + MockProvider::new_accept_all() + .with_systemone_handler(|request| { + let raw = serde_json::to_vec(&json!({ + "id": format!("decision-{}", Uuid::new_v4()), + "model": "decision-model", + "answers": { "billing": { "type": "noul", "noul": 0.95 } }, + "usage": { "input_tokens": 10, "output_tokens": 3 } + })) + .expect("fixture serializes"); + SystemOneResponseWithBytes::parse(raw, &request) + }) + .with_responses_handler(|_body: Value| { + let native = json!({ + "id": format!("resp_{}", Uuid::new_v4()), + "object": "response", + "status": "completed", + "output": [], + "usage": { "input_tokens": 10, "output_tokens": 5, "total_tokens": 15 }, + "service_tier": "default", + "store": false + }); + ResponsesRawResponse { + status: axum::http::StatusCode::OK, + headers: [( + axum::http::header::CONTENT_TYPE, + axum::http::HeaderValue::from_static("application/json"), + )] + .into_iter() + .collect(), + body: Box::pin(futures_util::stream::iter(vec![Ok(Bytes::from( + native.to_string(), + ))])), + } + }), + ); + for model in [&text, &decisions, &anthropic, &native] { + pool.register_provider(model.clone(), provider.clone()) + .await; + } + + Fixture { + server, + database, + pool, + provider, + text, + text_alias, + decisions, + anthropic, + native, + successor, + } +} + +struct Organization { + id: String, + api_key: String, +} + +async fn organization(server: &axum_test::TestServer) -> Organization { + let org = setup_org_with_credits(server, 10_000_000_000i64).await; + let api_key = get_api_key_for_org(server, org.id.clone()).await; + Organization { + id: org.id, + api_key, + } +} + +fn access_path(model: &str) -> String { + format!("/v1/admin/models/{}/access", urlencoding::encode(model)) +} + +fn grant_path(model: &str, organization_id: &str) -> String { + format!("{}/{organization_id}", access_path(model)) +} + +async fn admin_call( + server: &axum_test::TestServer, + method: Method, + path: &str, + token: &str, +) -> axum_test::TestResponse { + server + .method(method, path) + .add_header("Authorization", format!("Bearer {token}")) + .add_header("User-Agent", MOCK_USER_AGENT) + .await +} + +async fn grant(server: &axum_test::TestServer, model: &str, organization_id: &str) -> Value { + let response = admin_call( + server, + Method::PUT, + &grant_path(model, organization_id), + &get_session_id(), + ) + .await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + response.json() +} + +async fn revoke(server: &axum_test::TestServer, model: &str, organization_id: &str) { + let response = admin_call( + server, + Method::DELETE, + &grant_path(model, organization_id), + &get_session_id(), + ) + .await; + assert_eq!(response.status_code(), 204, "{}", response.text()); +} + +/// `PATCH /v1/admin/models` for one model. +async fn set_listed(server: &axum_test::TestServer, model: &str, is_listed: bool) { + upsert_models(server, vec![(model, json!({ "isListed": is_listed }))]).await; +} + +/// Every route that takes a model name from the caller and answers for it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Route { + ChatCompletions, + ChatCompletionsStream, + Completions, + CompletionsStream, + Responses, + ResponsesStream, + ResponsesNative, + Messages, + MessagesCountTokens, + Embeddings, + ImageGenerations, + ImageEdits, + AudioTranscriptions, + Rerank, + Score, + PrivacyClassify, + PrivacyRedact, + SystemOne, + AttestationReport, + AttestationReportWithNonce, +} + +const ROUTES: &[Route] = &[ + Route::ChatCompletions, + Route::ChatCompletionsStream, + Route::Completions, + Route::CompletionsStream, + Route::Responses, + Route::ResponsesStream, + Route::ResponsesNative, + Route::Messages, + Route::MessagesCountTokens, + Route::Embeddings, + Route::ImageGenerations, + Route::ImageEdits, + Route::AudioTranscriptions, + Route::Rerank, + Route::Score, + Route::PrivacyClassify, + Route::PrivacyRedact, + Route::SystemOne, + Route::AttestationReport, + Route::AttestationReportWithNonce, +]; + +impl Route { + /// The route template as it is registered in the router. + fn path(self) -> &'static str { + match self { + Route::ChatCompletions | Route::ChatCompletionsStream => "/chat/completions", + Route::Completions | Route::CompletionsStream => "/completions", + Route::Responses | Route::ResponsesStream | Route::ResponsesNative => "/responses", + Route::Messages => "/messages", + Route::MessagesCountTokens => "/messages/count_tokens", + Route::Embeddings => "/embeddings", + Route::ImageGenerations => "/images/generations", + Route::ImageEdits => "/images/edits", + Route::AudioTranscriptions => "/audio/transcriptions", + Route::Rerank => "/rerank", + Route::Score => "/score", + Route::PrivacyClassify => "/privacy/classify", + Route::PrivacyRedact => "/privacy/redact", + Route::SystemOne => "/systemone", + Route::AttestationReport | Route::AttestationReportWithNonce => "/attestation/report", + } + } + + fn model_kind(self) -> ModelKind { + match self { + Route::SystemOne => ModelKind::Decisions, + Route::Messages | Route::MessagesCountTokens => ModelKind::Anthropic, + Route::ResponsesNative => ModelKind::NativeResponses, + _ => ModelKind::Text, + } + } + + /// The method the route is registered with. + fn method(self) -> &'static str { + match self { + Route::AttestationReport | Route::AttestationReportWithNonce => "GET", + _ => "POST", + } + } + + /// Whether the mock provider can serve the route, so an organization + /// that holds a grant gets a 200. The native Anthropic routes need a + /// provider that speaks that protocol, which the mock does not: there a + /// grant only moves the request past the unknown-model answer. + fn mock_serves_it(self) -> bool { + !matches!(self, Route::Messages | Route::MessagesCountTokens) + } + + /// Whether a served request on this route carries the deprecation + /// headers of its model. The attestation report is not an inference + /// route and has no such headers; the Anthropic routes are not served by + /// the mock. + fn announces_deprecation(self) -> bool { + self.mock_serves_it() + && !matches!( + self, + Route::AttestationReport | Route::AttestationReportWithNonce + ) + } + + /// Whether the route accepts an alias in place of the model's name. The + /// others look the model up by its canonical name only. + fn resolves_aliases(self) -> bool { + matches!( + self, + Route::ChatCompletions + | Route::ChatCompletionsStream + | Route::Completions + | Route::CompletionsStream + | Route::Responses + | Route::ResponsesStream + | Route::AudioTranscriptions + | Route::AttestationReport + | Route::AttestationReportWithNonce + ) + } +} + +/// The deprecation headers of a response: (date, successor). +fn deprecation_headers(response: &Observed) -> (Option<&str>, Option<&str>) { + let header = |name: &str| response.headers.get(name).map(String::as_str); + (header(HEADER_DEPRECATION_DATE), header(HEADER_SUCCESSOR)) +} + +/// A minimal PNG image (1x1 pixel). +fn png() -> Vec { + vec![ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, + 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, + 0x15, 0xC4, 0x89, 0x00, 0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0x00, + 0x01, 0x00, 0x00, 0x05, 0x00, 0x01, 0x0D, 0x0A, 0x2D, 0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, + 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82, + ] +} + +/// Send the route's request for `model` with `api_key`. +async fn send( + server: &axum_test::TestServer, + route: Route, + api_key: &str, + model: &str, +) -> axum_test::TestResponse { + let auth = format!("Bearer {api_key}"); + let post = |path: &str, body: Value| { + server + .post(path) + .add_header("Authorization", auth.clone()) + .json(&body) + }; + let chat_body = |stream: bool| { + json!({ + "model": model, + "messages": [{ "role": "user", "content": "Hello" }], + "stream": stream, + "max_tokens": 16 + }) + }; + match route { + Route::ChatCompletions => post("/v1/chat/completions", chat_body(false)).await, + Route::ChatCompletionsStream => post("/v1/chat/completions", chat_body(true)).await, + Route::Completions | Route::CompletionsStream => { + post( + "/v1/completions", + json!({ + "model": model, + "prompt": "Hello", + "stream": route == Route::CompletionsStream, + "max_tokens": 16 + }), + ) + .await + } + Route::Responses | Route::ResponsesStream => { + post( + "/v1/responses", + json!({ + "model": model, + "input": "Hello", + "stream": route == Route::ResponsesStream + }), + ) + .await + } + Route::ResponsesNative => { + post( + "/v1/responses", + json!({ "model": model, "store": false, "input": "Hello" }), + ) + .await + } + Route::Messages | Route::MessagesCountTokens => { + let path = if route == Route::Messages { + "/v1/messages" + } else { + "/v1/messages/count_tokens" + }; + let mut body = json!({ + "model": model, + "messages": [{ "role": "user", "content": "Hello" }] + }); + if route == Route::Messages { + body["max_tokens"] = json!(16); + } + post(path, body) + .add_header("anthropic-version", "2023-06-01") + .await + } + Route::Embeddings => { + post( + "/v1/embeddings", + json!({ "model": model, "input": "Hello world" }), + ) + .await + } + Route::ImageGenerations => { + post( + "/v1/images/generations", + json!({ "model": model, "prompt": "A lighthouse at dusk", "n": 1 }), + ) + .await + } + Route::ImageEdits => { + server + .post("/v1/images/edits") + .add_header("Authorization", auth.clone()) + .multipart( + axum_test::multipart::MultipartForm::new() + .add_part( + "image", + axum_test::multipart::Part::bytes(png()) + .file_name("image.png") + .mime_type("image/png"), + ) + .add_text("model", model) + .add_text("prompt", "Make the sky clear"), + ) + .await + } + Route::AudioTranscriptions => { + server + .post("/v1/audio/transcriptions") + .add_header("Authorization", auth.clone()) + .multipart( + axum_test::multipart::MultipartForm::new() + .add_part( + "file", + axum_test::multipart::Part::bytes(vec![0u8; 1024]) + .file_name("test.mp3") + .mime_type("audio/mpeg"), + ) + .add_text("model", model), + ) + .await + } + Route::Rerank => { + post( + "/v1/rerank", + json!({ + "model": model, + "query": "What is the capital of France?", + "documents": ["Paris is the capital of France.", "Cows are animals."] + }), + ) + .await + } + Route::Score => { + post( + "/v1/score", + json!({ "model": model, "text_1": "Hello world", "text_2": "Hi there" }), + ) + .await + } + Route::PrivacyClassify => { + post( + "/v1/privacy/classify", + json!({ "model": model, "input": "My name is Alice" }), + ) + .await + } + Route::PrivacyRedact => { + post( + "/v1/privacy/redact", + json!({ "model": model, "input": "My name is Alice" }), + ) + .await + } + Route::SystemOne => { + post( + "/v1/systemone", + json!({ + "model": model, + "state": { "message": "Please refund this charge" }, + "questions": { + "billing": { "type": "noul", "instructions": "Is this about billing?" } + } + }), + ) + .await + } + Route::AttestationReport | Route::AttestationReportWithNonce => { + let model = urlencoding::encode(model); + let nonce = if route == Route::AttestationReportWithNonce { + format!("&nonce={NONCE}") + } else { + String::new() + }; + server + .get(&format!("/v1/attestation/report?model={model}{nonce}")) + .add_header("Authorization", auth.clone()) + .await + } + } +} + +/// Everything a client can observe of a response, with the model name it +/// asked for blanked out and the values that differ on every request removed. +#[derive(Debug, PartialEq, Eq)] +struct Observed { + status: u16, + headers: BTreeMap, + body: String, +} + +fn observe(response: &axum_test::TestResponse, model: &str) -> Observed { + // Per-request values a client cannot tell two models apart by. + const PER_REQUEST_HEADERS: &[&str] = &["date", "x-request-id"]; + let headers = response + .headers() + .iter() + .filter(|(name, _)| !PER_REQUEST_HEADERS.contains(&name.as_str())) + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().unwrap_or("").to_string(), + ) + }) + .collect(); + Observed { + status: response.status_code().as_u16(), + headers, + body: blank_generated_ids(&response.text().replace(model, "")), + } +} + +/// Blank the ids and timestamps the Responses API generates for every +/// request (`resp_…`, `msg_…`, `"created_at":…`), so two answers to the same +/// request compare equal. +fn blank_generated_ids(body: &str) -> String { + let mut out = String::with_capacity(body.len()); + let mut rest = body; + while let Some(start) = ["resp_", "msg_", "\"created_at\":"] + .iter() + .filter_map(|marker| rest.find(marker).map(|at| (at, *marker))) + .min() + { + let (at, marker) = start; + out.push_str(&rest[..at + marker.len()]); + let tail = &rest[at + marker.len()..]; + let generated = tail + .find(|c: char| !c.is_ascii_alphanumeric() && c != '-') + .unwrap_or(tail.len()); + out.push('#'); + rest = &tail[generated..]; + } + out.push_str(rest); + out +} + +/// Every `.route("", (..))` call in a router source file, as +/// (method, path) pairs. Comment lines and the unit-test module are left out. +fn registered_routes(source: &str) -> Vec<(String, String)> { + const METHODS: &[&str] = &[ + "get", "post", "put", "patch", "delete", "head", "options", "trace", "any", + ]; + let production = source + .split("\n#[cfg(test)]\nmod tests") + .next() + .expect("split yields at least one part"); + let code = production + .lines() + .filter(|line| !line.trim_start().starts_with("//")) + .collect::>() + .join("\n"); + + let mut routes = Vec::new(); + let mut rest = code.as_str(); + while let Some(at) = rest.find(".route(") { + rest = &rest[at + ".route(".len()..]; + // The arguments of the call: up to the parenthesis that closes it. + let mut depth = 1usize; + let end = rest + .char_indices() + .find_map(|(index, c)| { + match c { + '(' => depth += 1, + ')' => depth -= 1, + _ => {} + } + (depth == 0).then_some(index) + }) + .expect("a `.route(` call that is never closed"); + let call = &rest[..end]; + let quoted = call.trim_start(); + assert!( + quoted.starts_with('"'), + "a route path that is not a string literal: {call}" + ); + let (path, handlers) = quoted[1..] + .split_once('"') + .expect("the path's closing quote"); + + let mut methods = Vec::new(); + for method in METHODS { + let token = format!("{method}("); + let mut from = 0; + while let Some(found) = handlers[from..].find(&token) { + let at = from + found; + // `get(` is a method; `budget(` or `my_get(` is not. + let starts_a_word = handlers[..at] + .chars() + .next_back() + .is_none_or(|c| !(c.is_ascii_alphanumeric() || c == '_')); + if starts_a_word { + methods.push(method.to_uppercase()); + } + from = at + token.len(); + } + } + assert!( + !methods.is_empty(), + "no method found for route {path}: {handlers}" + ); + routes.extend(methods.into_iter().map(|method| (method, path.to_string()))); + rest = &rest[end..]; + } + routes +} + +/// The access tests cover the routes in `ROUTES`. This one fails when a route +/// is registered in any router builder of `src/lib.rs` (or in the public +/// attestation, credit transfer and organization credential routers it +/// merges) without a decision on record here: either +/// the route is in `ROUTES`, and so in the table test, or it is listed below +/// with the reason it is not. A new route cannot ship without that decision. +/// +/// Routes are matched by method and path, as they are registered. +#[test] +fn every_registered_route_has_a_model_access_decision() { + /// Routes that deal in model names without being inference for an + /// organization. Each has its own rule, tested elsewhere in this file or + /// in `unlisted_models.rs`. + const OTHER_MODEL_ROUTES: &[(&str, &str)] = &[ + // Public catalogs: listed models only. + ("GET", "/models"), + ("GET", "/model/list"), + ("GET", "/model/{model_name}"), + // Public: resolves as an anonymous caller. + ("GET", "/attestation/ita-token"), + // Public: the attested providers' trust pins, published for + // verifiers. It takes no model name and reads no catalog row. + ("GET", "/attestation/pins"), + // Service token: usage is recorded for any model, and the two + // documents are for a gateway that enforces access itself. + ("POST", "/internal/usage"), + ("GET", "/internal/models"), + ("GET", "/internal/model-access"), + ]; + /// Routes that take no model name from the caller. + const NO_MODEL_NAME: &[(&str, &str)] = &[ + ("GET", "/health"), + ("POST", "/ohttp"), + ("GET", "/.well-known/ohttp-gateway"), + ("GET", "/ohttp/config"), + ("GET", "/docs"), + ("GET", "/api-docs/openapi.json"), + // Sign-in and invitations. + ("POST", "/auth/vpc/login"), + ("POST", "/near"), + ("GET", "/login"), + ("GET", "/github"), + ("GET", "/google"), + ("GET", "/callback"), + ("GET", "/user"), + ("POST", "/logout"), + ("GET", "/{token}"), + ("POST", "/{token}/accept"), + // Stored responses, by id. + ("GET", "/responses/{response_id}"), + ("DELETE", "/responses/{response_id}"), + ("POST", "/responses/{response_id}/cancel"), + ("GET", "/responses/{response_id}/input_items"), + // The MCP endpoint serves web search only. + ("POST", "/mcp"), + ("POST", "/conversations"), + ("POST", "/conversations/batch"), + ("GET", "/conversations/{conversation_id}"), + ("POST", "/conversations/{conversation_id}"), + ("DELETE", "/conversations/{conversation_id}"), + ("POST", "/conversations/{conversation_id}/pin"), + ("DELETE", "/conversations/{conversation_id}/pin"), + ("POST", "/conversations/{conversation_id}/archive"), + ("DELETE", "/conversations/{conversation_id}/archive"), + ("POST", "/conversations/{conversation_id}/clone"), + ("GET", "/conversations/{conversation_id}/items"), + ("POST", "/conversations/{conversation_id}/items"), + // Looked up by chat id. The `model` query parameter is not read. + ("GET", "/signature/{chat_id}"), + ("GET", "/organizations/{org_id}/workspaces"), + ("POST", "/organizations/{org_id}/workspaces"), + ("GET", "/workspaces/{workspace_id}"), + ("PUT", "/workspaces/{workspace_id}"), + ("DELETE", "/workspaces/{workspace_id}"), + ("GET", "/workspaces/{workspace_id}/api-keys"), + ("POST", "/workspaces/{workspace_id}/api-keys"), + ("DELETE", "/workspaces/{workspace_id}/api-keys/{key_id}"), + ("PATCH", "/workspaces/{workspace_id}/api-keys/{key_id}"), + ( + "PATCH", + "/workspaces/{workspace_id}/api-keys/{key_id}/spend-limit", + ), + ( + "GET", + "/workspaces/{workspace_id}/api-keys/{key_id}/usage/history", + ), + ("POST", "/files"), + ("GET", "/files"), + ("GET", "/files/{file_id}"), + ("DELETE", "/files/{file_id}"), + ("GET", "/files/{file_id}/content"), + ("POST", "/feature-requests"), + // An organization's own usage and costs. + ("POST", "/billing/costs"), + ("GET", "/organizations/{org_id}/usage/export"), + ("GET", "/organizations/{org_id}/usage/summary"), + ( + "GET", + "/organizations/{org_id}/usage/reporting-token-auth-probe", + ), + ("POST", "/check_api_key"), + ("GET", "/services"), + ("GET", "/services/{service_name}"), + // Credit transfers and the scoped credentials that authorize them. + ("POST", "/organizations/{org_id}/credit/transfer"), + ("GET", "/organizations/{org_id}/credit/transfers"), + ("GET", "/organizations/{org_id}/credit/transfers/{id}"), + ("GET", "/organizations/{org_id}/credit/transfer/capacity"), + ("POST", "/organizations/{org_id}/credentials"), + ("GET", "/organizations/{org_id}/credentials"), + ("DELETE", "/organizations/{org_id}/credentials/{id}"), + ("PUT", "/organizations/{org_id}/credentials/{id}"), + ("POST", "/organizations/{org_id}/credentials/{id}/rotate"), + ]; + + let mut registered = registered_routes(include_str!("../../src/lib.rs")); + // Router builders that `src/lib.rs` calls from other modules. + for source in [ + include_str!("../../src/routes/attestation.rs"), + include_str!("../../src/routes/credit_transfers.rs"), + include_str!("../../src/routes/organization_credentials.rs"), + ] { + registered.extend(registered_routes(source)); + } + assert!( + registered.len() >= 100, + "found too few routes, the scan is broken: {registered:?}" + ); + let is_registered = + |method: &str, path: &str| registered.iter().any(|(m, p)| m == method && p == path); + + for (method, path) in ®istered { + // Admin routes act as the platform: their model names are not + // subject to grants. Their own access rule is the admin permission + // policy. + let admin = path.starts_with("/admin/"); + let in_table = ROUTES + .iter() + .any(|route| route.method() == method && route.path() == path); + let listed = [OTHER_MODEL_ROUTES, NO_MODEL_NAME] + .concat() + .contains(&(method.as_str(), path.as_str())); + assert!( + admin || in_table || listed, + "{method} {path} has no model access decision. If it takes a model name from \ + the caller, add it to ROUTES (and `send`) so the table test covers it. If not, \ + add it to NO_MODEL_NAME or OTHER_MODEL_ROUTES." + ); + } + + // The lists must not rot: every entry is still a registered route, and + // the scan saw every route of the table. + for (method, path) in OTHER_MODEL_ROUTES.iter().chain(NO_MODEL_NAME) { + assert!( + is_registered(method, path), + "{method} {path} is listed here but no longer registered" + ); + } + for route in ROUTES { + assert!( + is_registered(route.method(), route.path()), + "{route:?} is in the table but {} {} is not registered", + route.method(), + route.path() + ); + } +} + +/// The scanner itself: methods chained on one path, a layered handler, and +/// names that only look like methods. +#[test] +fn registered_routes_reads_methods_and_paths() { + let source = r#" + Router::new() + .route("/a", get(get_a).delete(delete_a)) + .route( + "/b/{id}", + post(budget(my_get)).layer(DefaultBodyLimit::max(LIMIT)), + ) + // .route("/commented-out", get(gone)) + .route("/c", axum::routing::patch(update_c)) + +#[cfg(test)] +mod tests { + fn app() -> Router { Router::new().route("/test-only", get(ok)) } +} +"#; + let routes = registered_routes(source); + let routes: Vec<(&str, &str)> = routes + .iter() + .map(|(method, path)| (method.as_str(), path.as_str())) + .collect(); + assert_eq!( + routes, + vec![ + ("GET", "/a"), + ("DELETE", "/a"), + ("POST", "/b/{id}"), + ("PATCH", "/c"), + ] + ); +} + +/// (a) For an unlisted model that has a provider, an organization without a +/// grant gets on every route exactly what an unknown model gets. With a grant +/// the request passes the check. +#[tokio::test] +async fn unlisted_model_is_unknown_without_a_grant_on_every_inference_route() { + let fixture = setup().await; + let server = &fixture.server; + let outsider = organization(server).await; + let member = organization(server).await; + for kind in [ + ModelKind::Text, + ModelKind::Decisions, + ModelKind::Anthropic, + ModelKind::NativeResponses, + ] { + grant(server, fixture.model(kind), &member.id).await; + } + + for &route in ROUTES { + let model = fixture.model(route.model_kind()); + let unknown = unknown_twin(model); + + let missing = send(server, route, &outsider.api_key, &unknown).await; + let refused = send(server, route, &outsider.api_key, model).await; + let missing = observe(&missing, &unknown); + let refused = observe(&refused, model); + // A streaming Responses request answers 200 and fails in the stream. + assert!( + (400..500).contains(&missing.status) || missing.body.contains("event: response.failed"), + "{route:?}: an unknown model must be an error: {missing:?}" + ); + assert_eq!( + refused, missing, + "{route:?}: an unlisted model without a grant must look unknown" + ); + + let granted = send(server, route, &member.api_key, model).await; + if route.mock_serves_it() { + assert_eq!( + granted.status_code(), + 200, + "{route:?}: the granted organization must be served: {}", + granted.text() + ); + } + let granted = observe(&granted, model); + assert_ne!( + granted, refused, + "{route:?}: a grant must move the request past the unknown-model answer" + ); + assert!( + !granted.body.contains("event: response.failed"), + "{route:?}: the granted stream must not fail: {granted:?}" + ); + + // The model has a planned deprecation. Its headers are part of what + // was compared above; spell the two cases out. + assert_eq!(deprecation_headers(&refused), (None, None), "{route:?}"); + if route.announces_deprecation() { + assert_eq!( + deprecation_headers(&granted), + ( + Some(DEPRECATION_DATE_HEADER), + Some(fixture.successor.as_str()) + ), + "{route:?}" + ); + } + + // The same through the model's alias. A route that does not resolve + // aliases answers "unknown" to everyone. + if route.model_kind() != ModelKind::Text { + continue; + } + let alias = &fixture.text_alias; + let unknown_alias = unknown_twin(alias); + let missing = observe( + &send(server, route, &outsider.api_key, &unknown_alias).await, + &unknown_alias, + ); + let refused = observe(&send(server, route, &outsider.api_key, alias).await, alias); + assert_eq!( + refused, missing, + "{route:?}: the alias of an unlisted model without a grant must look unknown" + ); + assert!( + !refused.body.contains(model), + "{route:?}: the alias answer names the canonical model: {refused:?}" + ); + let granted = observe(&send(server, route, &member.api_key, alias).await, alias); + if route.resolves_aliases() { + assert_ne!(granted, refused, "{route:?} by alias"); + // The audio route resolves the alias but then hands the alias + // itself to the provider pool, which does not know it. That is + // how it behaves for a listed model too; the request is past + // the access check either way. + if route != Route::AudioTranscriptions { + assert_eq!(granted.status, 200, "{route:?} by alias: {granted:?}"); + } + } else { + assert_eq!(granted, refused, "{route:?} does not resolve aliases"); + } + } +} + +/// A planned deprecation is announced in response headers. For an unlisted +/// model they would tell a caller without a grant that the model exists, on +/// every route, and above all on the streaming Responses route, which +/// answers 200 before the model is resolved. +#[tokio::test] +async fn deprecation_headers_are_only_added_for_a_caller_that_may_use_the_model() { + let fixture = setup().await; + let server = &fixture.server; + let outsider = organization(server).await; + let member = organization(server).await; + for kind in [ + ModelKind::Text, + ModelKind::Decisions, + ModelKind::Anthropic, + ModelKind::NativeResponses, + ] { + grant(server, fixture.model(kind), &member.id).await; + } + let announced = ( + Some(DEPRECATION_DATE_HEADER), + Some(fixture.successor.as_str()), + ); + + for &route in ROUTES { + let model = fixture.model(route.model_kind()); + + let refused = send(server, route, &outsider.api_key, model).await; + let refused = observe(&refused, model); + assert_eq!( + deprecation_headers(&refused), + (None, None), + "{route:?}: headers for a model the caller may not use: {refused:?}" + ); + // The successor is not named anywhere else in the answer either. + assert!( + !refused.body.contains(&fixture.successor), + "{route:?}: {refused:?}" + ); + + let granted = observe(&send(server, route, &member.api_key, model).await, model); + if route.announces_deprecation() { + assert_eq!(deprecation_headers(&granted), announced, "{route:?}"); + } + } + + // The streaming Responses route is the one that answers 200 either way. + let refused = send( + server, + Route::ResponsesStream, + &outsider.api_key, + &fixture.text, + ) + .await; + assert_eq!(refused.status_code(), 200); + assert!(refused.text().contains("event: response.failed")); + assert!(refused.maybe_header(HEADER_DEPRECATION_DATE).is_none()); + assert!(refused.maybe_header(HEADER_SUCCESSOR).is_none()); + + // By alias: nothing without a grant, the canonical model's headers with + // one. + for (org, expected) in [(&outsider, (None, None)), (&member, announced)] { + let response = send( + server, + Route::ResponsesStream, + &org.api_key, + &fixture.text_alias, + ) + .await; + let response = observe(&response, &fixture.text_alias); + assert_eq!(deprecation_headers(&response), expected); + } + + // Revoking takes the headers away with the access. + revoke(server, &fixture.text, &member.id).await; + let revoked = send( + server, + Route::ResponsesStream, + &member.api_key, + &fixture.text, + ) + .await; + assert!(revoked.maybe_header(HEADER_DEPRECATION_DATE).is_none()); + assert!(revoked.maybe_header(HEADER_SUCCESSOR).is_none()); +} + +/// A listed model may name an unlisted successor. Its name reaches only a +/// caller that may use it: the deprecation header carries it for an +/// organization that holds a grant for the successor, and the public +/// catalogs leave it out until the successor is listed. +#[tokio::test] +async fn an_unlisted_successor_is_named_only_to_a_caller_that_may_use_it() { + let fixture = setup().await; + let server = &fixture.server; + let suffix = Uuid::new_v4(); + let successor = format!("access-restricted-successor/model-{suffix}"); + let public = format!("access-public/model-{suffix}"); + upsert_models(server, vec![(&successor, model_fields(json!({})))]).await; + let mut fields = model_fields(json!({ "isListed": true })); + fields["deprecationDate"] = json!(DEPRECATION_DATE); + fields["successorModelId"] = json!(successor); + upsert_models(server, vec![(&public, fields)]).await; + fixture + .pool + .register_provider(public.clone(), fixture.provider.clone()) + .await; + let outsider = organization(server).await; + let member = organization(server).await; + grant(server, &successor, &member.id).await; + + let headers = |org: &Organization| { + let api_key = org.api_key.clone(); + let public = public.clone(); + async move { + let response = send(server, Route::ChatCompletions, &api_key, &public).await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + let observed = observe(&response, &public); + ( + deprecation_headers(&observed).0.map(str::to_string), + deprecation_headers(&observed).1.map(str::to_string), + observed.body, + ) + } + }; + let catalogs = || async { + let mut bodies = Vec::new(); + for path in [ + "/v1/models".to_string(), + "/v1/model/list?limit=100000".to_string(), + format!("/v1/model/{}", urlencoding::encode(&public)), + ] { + let response = server.get(&path).await; + assert_eq!(response.status_code(), 200, "{path}"); + assert!(response.text().contains(&public), "{path}"); + bodies.push((path, response.text())); + } + bodies + }; + + let (date, named, body) = headers(&outsider).await; + assert_eq!(date.as_deref(), Some(DEPRECATION_DATE_HEADER)); + assert_eq!(named, None); + assert!(!body.contains(&successor), "{body}"); + let (_, named, _) = headers(&member).await; + assert_eq!(named.as_deref(), Some(successor.as_str())); + for (path, body) in catalogs().await { + assert!(!body.contains(&successor), "{path}: {body}"); + } + + // Once the successor is listed, everyone may use it and see it. + set_listed(server, &successor, true).await; + let (_, named, _) = headers(&outsider).await; + assert_eq!(named.as_deref(), Some(successor.as_str())); + for (path, body) in catalogs().await { + assert!(body.contains(&successor), "{path}: {body}"); + } +} + +/// The pre-dispatch checks of the chat routes answer for the requested model +/// before the completion service resolves it. None of them may answer for an +/// unlisted model the organization holds no grant for. +#[tokio::test] +async fn pre_dispatch_checks_do_not_disclose_an_unlisted_model() { + let fixture = setup().await; + let server = &fixture.server; + let outsider = organization(server).await; + let unknown = unknown_twin(&fixture.text); + let unknown_alias = unknown_twin(&fixture.text_alias); + + // A video part is refused with the model's declared modalities. + let video = |model: &str| { + server + .post("/v1/chat/completions") + .add_header("Authorization", format!("Bearer {}", outsider.api_key)) + .json(&json!({ + "model": model, + "messages": [{ "role": "user", "content": [ + { "type": "text", "text": "Describe this" }, + { "type": "video_url", "video_url": { "url": "https://example.com/clip.mp4" } } + ]}], + "max_tokens": 16 + })) + }; + let refused = video(&fixture.text).await; + assert!( + !refused.text().contains("does not support"), + "the modality check answered for a model the caller may not use: {}", + refused.text() + ); + assert_eq!( + observe(&refused, &fixture.text), + observe(&video(&unknown).await, &unknown) + ); + + // Strict alias mode names the canonical model of an alias. + for path in ["/v1/chat/completions", "/v1/completions"] { + let strict = |model: &str| { + let body = if path == "/v1/completions" { + json!({ "model": model, "prompt": "Hello", "max_tokens": 16 }) + } else { + json!({ + "model": model, + "messages": [{ "role": "user", "content": "Hello" }], + "max_tokens": 16 + }) + }; + server + .post(path) + .add_header("Authorization", format!("Bearer {}", outsider.api_key)) + .add_header("x-no-aliasing", "true") + .json(&body) + }; + let refused = strict(&fixture.text_alias).await; + assert!( + !refused.text().contains(&fixture.text), + "{path}: the alias check named the canonical model: {}", + refused.text() + ); + assert_eq!( + observe(&refused, &fixture.text_alias), + observe(&strict(&unknown_alias).await, &unknown_alias), + "{path}" + ); + } + + // The attestation report resolves aliases too, and announces them. + let report = |model: &str, strict: bool| { + let mut request = server + .get(&format!( + "/v1/attestation/report?model={}", + urlencoding::encode(model) + )) + .add_header("Authorization", format!("Bearer {}", outsider.api_key)); + if strict { + request = request.add_header("x-no-aliasing", "true"); + } + request + }; + for strict in [false, true] { + let refused = report(&fixture.text_alias, strict).await; + assert!( + !refused.text().contains(&fixture.text), + "the report route named the canonical model: {}", + refused.text() + ); + assert_eq!( + observe(&refused, &fixture.text_alias), + observe(&report(&unknown_alias, strict).await, &unknown_alias) + ); + } + + // None of this reached the provider. + assert!(fixture.provider.last_chat_params().await.is_none()); +} + +/// (b) Granting and revoking take effect at once on the instance that served +/// the admin request, also when the answer for the organization was cached. +#[tokio::test] +async fn grant_and_revoke_take_effect_at_once() { + let fixture = setup().await; + let server = &fixture.server; + let org = organization(server).await; + let other = organization(server).await; + let chat = || send(server, Route::ChatCompletions, &org.api_key, &fixture.text); + let embed = || send(server, Route::Embeddings, &org.api_key, &fixture.text); + let unknown = unknown_twin(&fixture.text); + let unknown_chat = observe( + &send(server, Route::ChatCompletions, &org.api_key, &unknown).await, + &unknown, + ); + + // Refused, twice: the second answer comes from the cached decision. + for _ in 0..2 { + assert_eq!(observe(&chat().await, &fixture.text), unknown_chat); + assert_eq!(embed().await.status_code(), 404); + } + + grant(server, &fixture.text, &org.id).await; + for _ in 0..2 { + assert_eq!(chat().await.status_code(), 200); + assert_eq!(embed().await.status_code(), 200); + // The alias resolves for the organization as well. + let aliased = send( + server, + Route::ChatCompletions, + &org.api_key, + &fixture.text_alias, + ) + .await; + assert_eq!(aliased.status_code(), 200, "{}", aliased.text()); + } + // The grant is the organization's, nobody else's. + assert_eq!( + send( + server, + Route::ChatCompletions, + &other.api_key, + &fixture.text + ) + .await + .status_code(), + unknown_chat.status + ); + + revoke(server, &fixture.text, &org.id).await; + for _ in 0..2 { + assert_eq!(observe(&chat().await, &fixture.text), unknown_chat); + assert_eq!(embed().await.status_code(), 404); + } +} + +/// (g) Listing decides. A listed model is open to everyone and grants on it +/// have no effect; unlisting it restricts it at once, and the grants prepared +/// while it was listed apply from that moment. +#[tokio::test] +async fn listing_a_model_opens_it_and_unlisting_restricts_it() { + let fixture = setup().await; + let server = &fixture.server; + let prepared = organization(server).await; + let outsider = organization(server).await; + let chat = |org: &Organization| { + let api_key = org.api_key.clone(); + let model = fixture.text.clone(); + async move { + send(server, Route::ChatCompletions, &api_key, &model) + .await + .status_code() + .as_u16() + } + }; + + set_listed(server, &fixture.text, true).await; + assert_eq!(chat(&outsider).await, 200); + assert_eq!(chat(&prepared).await, 200); + + // Prepare access while the model is still open to everyone. + grant(server, &fixture.text, &prepared.id).await; + assert_eq!(chat(&outsider).await, 200); + + set_listed(server, &fixture.text, false).await; + assert_eq!(chat(&prepared).await, 200); + assert_ne!(chat(&outsider).await, 200); + + set_listed(server, &fixture.text, true).await; + assert_eq!(chat(&outsider).await, 200); +} + +/// A grant counts only while the model is active. +#[tokio::test] +async fn a_grant_on_an_inactive_model_has_no_effect() { + let fixture = setup().await; + let server = &fixture.server; + let org = organization(server).await; + grant(server, &fixture.text, &org.id).await; + assert_eq!( + send(server, Route::Embeddings, &org.api_key, &fixture.text) + .await + .status_code(), + 200 + ); + + upsert_models(server, vec![(&fixture.text, json!({ "isActive": false }))]).await; + assert_eq!( + send(server, Route::Embeddings, &org.api_key, &fixture.text) + .await + .status_code(), + 404 + ); + // The grant is still there for the admin, with the reason it is idle. + let access: Value = admin_call( + server, + Method::GET, + &access_path(&fixture.text), + &get_session_id(), + ) + .await + .json(); + assert_eq!(access["isActive"], false); + assert_eq!(access["organizations"][0]["organizationId"], org.id); + // The gateway document keeps the model, with its grants. A lane may + // still be serving it, and a model that is missing from the document is + // read as unrestricted. + assert_eq!( + internal_entry(server, &fixture.text).await, + Some(json!({ "id": fixture.text, "organization_ids": [org.id] })) + ); + + // Grants on an inactive model can still be managed, so it can be + // prepared before it is activated and cleaned up after. + let other = organization(server).await; + grant(server, &fixture.text, &other.id).await; + revoke(server, &fixture.text, &org.id).await; + upsert_models(server, vec![(&fixture.text, json!({ "isActive": true }))]).await; + // Deactivating a model unregisters its provider. + fixture + .pool + .register_provider(fixture.text.clone(), fixture.provider.clone()) + .await; + assert_eq!( + send(server, Route::Embeddings, &other.api_key, &fixture.text) + .await + .status_code(), + 200 + ); + assert_eq!( + send(server, Route::Embeddings, &org.api_key, &fixture.text) + .await + .status_code(), + 404 + ); +} + +/// The audit trail of one organization's access to one model, oldest first: +/// (action, acting user). +async fn audit_trail( + fixture: &Fixture, + model: &str, + organization_id: &str, +) -> Vec<(String, Option)> { + fixture + .database + .pool() + .get() + .await + .expect("database connection") + .query( + "SELECT l.action, l.acting_user_id \ + FROM organization_model_access_log l \ + JOIN models m ON m.id = l.model_id \ + WHERE m.model_name = $1 AND l.organization_id = $2 \ + ORDER BY l.created_at, l.action", + &[ + &model, + &Uuid::parse_str(organization_id).expect("organization id"), + ], + ) + .await + .expect("read the audit trail") + .iter() + .map(|row| (row.get("action"), row.get("acting_user_id"))) + .collect() +} + +/// Every grant and every revocation leaves one audit row, with the admin who +/// made it. A call that changes nothing leaves none. +#[tokio::test] +async fn grants_and_revocations_are_audited() { + let fixture = setup().await; + let server = &fixture.server; + let org = organization(server).await; + let other = organization(server).await; + let model = &fixture.text; + let admin = Some(Uuid::parse_str(MOCK_USER_ID).expect("mock user id")); + let granted = ("grant".to_string(), admin); + let revoked = ("revoke".to_string(), admin); + + assert_eq!(audit_trail(&fixture, model, &org.id).await, vec![]); + + // Revoking a grant that does not exist changes nothing. + revoke(server, model, &org.id).await; + assert_eq!(audit_trail(&fixture, model, &org.id).await, vec![]); + + // One row for the grant, none for granting it again. + grant(server, model, &org.id).await; + grant(server, model, &org.id).await; + assert_eq!( + audit_trail(&fixture, model, &org.id).await, + vec![granted.clone()] + ); + + // One row for the revocation, none for revoking again. The grant row + // is gone from the grants table; the trail still tells the story. + revoke(server, model, &org.id).await; + revoke(server, model, &org.id).await; + assert_eq!( + audit_trail(&fixture, model, &org.id).await, + vec![granted.clone(), revoked.clone()] + ); + + // Granting again is a new event. + grant(server, model, &org.id).await; + assert_eq!( + audit_trail(&fixture, model, &org.id).await, + vec![granted.clone(), revoked, granted.clone()] + ); + + // The trail is per organization and per model. + assert_eq!(audit_trail(&fixture, model, &other.id).await, vec![]); + assert_eq!( + audit_trail(&fixture, &fixture.decisions, &org.id).await, + vec![] + ); + + // A request that is refused writes nothing. + let unknown_org = Uuid::new_v4().to_string(); + let refused = admin_call( + server, + Method::PUT, + &grant_path(model, &unknown_org), + &get_session_id(), + ) + .await; + assert_eq!(refused.status_code(), 404); + assert_eq!(audit_trail(&fixture, model, &unknown_org).await, vec![]); + + // Deleting the organization removes its grant: a revocation by the + // platform, with no acting user. + let deleted = server + .delete(&format!("/v1/organizations/{}", org.id)) + .add_header("Authorization", format!("Bearer {}", get_session_id())) + .add_header("User-Agent", MOCK_USER_AGENT) + .await; + assert!( + deleted.status_code().is_success(), + "deleting the organization failed: {}", + deleted.text() + ); + let trail = audit_trail(&fixture, model, &org.id).await; + assert_eq!(trail.len(), 4, "{trail:?}"); + assert_eq!(trail[3], ("revoke".to_string(), None)); +} + +/// (c) The admin routes: happy paths, idempotency and 404s. +#[tokio::test] +async fn admin_manages_grants() { + let fixture = setup().await; + let server = &fixture.server; + let session = get_session_id(); + let first = organization(server).await; + let second = organization(server).await; + let model = &fixture.text; + + // Nobody holds a grant yet. + let access = admin_call(server, Method::GET, &access_path(model), &session).await; + assert_eq!(access.status_code(), 200, "{}", access.text()); + let access: Value = access.json(); + assert_eq!(access["modelName"], model.as_str()); + assert_eq!(access["isListed"], false); + assert_eq!(access["isActive"], true); + assert!(Uuid::parse_str(access["id"].as_str().expect("model id")).is_ok()); + assert_eq!(access["organizations"], json!([])); + + // Grant: the answer is the stored grant, with who made it. + let granted = grant(server, model, &first.id).await; + assert_eq!(granted["organizationId"], first.id); + assert!(granted["organizationName"].is_string()); + assert_eq!(granted["grantedByUserId"], MOCK_USER_ID); + assert_eq!(granted["grantedByUserEmail"], "admin@test.com"); + let granted_at = granted["grantedAt"] + .as_str() + .expect("grantedAt") + .to_string(); + + // Idempotent: a second PUT returns the same grant, not a newer one. + let again = grant(server, model, &first.id).await; + assert_eq!(again, granted); + + grant(server, model, &second.id).await; + let access: Value = admin_call(server, Method::GET, &access_path(model), &session) + .await + .json(); + let holders: Vec<&str> = access["organizations"] + .as_array() + .expect("organizations") + .iter() + .map(|grant| grant["organizationId"].as_str().expect("organizationId")) + .collect(); + assert_eq!(holders, vec![first.id.as_str(), second.id.as_str()]); + assert_eq!(access["organizations"][0]["grantedAt"], granted_at.as_str()); + + // The organization's side of the same grants. + let path = format!("/v1/admin/organizations/{}/model-access", first.id); + let of_first = admin_call(server, Method::GET, &path, &session).await; + assert_eq!(of_first.status_code(), 200, "{}", of_first.text()); + let of_first: Value = of_first.json(); + assert_eq!(of_first["organizationId"], first.id); + assert_eq!( + of_first["models"], + json!([{ + "id": access["id"], + "modelName": model, + "isListed": false, + "isActive": true, + "grantedAt": granted_at, + "grantedByUserId": MOCK_USER_ID, + "grantedByUserEmail": "admin@test.com" + }]) + ); + + // Revoke, and revoke again: both succeed. + revoke(server, model, &first.id).await; + revoke(server, model, &first.id).await; + let access: Value = admin_call(server, Method::GET, &access_path(model), &session) + .await + .json(); + assert_eq!(access["organizations"].as_array().unwrap().len(), 1); + assert_eq!(access["organizations"][0]["organizationId"], second.id); + let of_first: Value = admin_call(server, Method::GET, &path, &session) + .await + .json(); + assert_eq!(of_first["models"], json!([])); + + // Unknown model or organization: 404, and nothing is written. + let missing_model = unknown_twin(model); + let missing_org = Uuid::new_v4().to_string(); + for (method, path, error_type) in [ + (Method::GET, access_path(&missing_model), "model_not_found"), + ( + Method::PUT, + grant_path(&missing_model, &first.id), + "model_not_found", + ), + ( + Method::DELETE, + grant_path(&missing_model, &first.id), + "model_not_found", + ), + ( + Method::PUT, + grant_path(model, &missing_org), + "organization_not_found", + ), + ( + Method::DELETE, + grant_path(model, &missing_org), + "organization_not_found", + ), + ( + Method::GET, + format!("/v1/admin/organizations/{missing_org}/model-access"), + "organization_not_found", + ), + ] { + let response = admin_call(server, method.clone(), &path, &session).await; + assert_eq!(response.status_code(), 404, "{method} {path}"); + assert_eq!( + response.json::()["error"]["type"], + error_type, + "{method} {path}" + ); + } + let malformed = admin_call( + server, + Method::PUT, + &grant_path(model, "not-a-uuid"), + &session, + ) + .await; + assert_eq!(malformed.status_code(), 400); + assert_eq!(malformed.json::()["error"]["type"], "invalid_id"); + + // A model name with a slash that is not URL-encoded is not a route. + let unencoded = admin_call( + server, + Method::GET, + &format!("/v1/admin/models/{model}/access"), + &session, + ) + .await; + assert_eq!(unencoded.status_code(), 404); +} + +/// (c) Permissions: the routes need an admin, a read-only admin token can +/// read grants but not change them. +#[tokio::test] +async fn admin_routes_follow_the_admin_token_permissions() { + let suffix = Uuid::new_v4(); + let model = format!("access-permissions/model-{suffix}"); + let (server, database) = setup_test_server_with_config_and_database(|config| { + config.auth.admin_read_only_tokens_enabled = true; + }) + .await; + upsert_models(&server, vec![(&model, model_fields(json!({})))]).await; + let org = organization(&server).await; + let listing = access_path(&model); + let of_org = format!("/v1/admin/organizations/{}/model-access", org.id); + let one_grant = grant_path(&model, &org.id); + + let repo = database::repositories::AdminAccessTokenRepository::new(database.pool().clone()); + let mut tokens = Vec::new(); + for permission in [ + database::models::AdminAccessTokenPermission::ReadOnly, + database::models::AdminAccessTokenPermission::ReadWrite, + ] { + let (_, token) = repo + .create( + Uuid::parse_str(MOCK_USER_ID).unwrap(), + format!("model-access-{suffix}"), + "model access permission test".into(), + chrono::Utc::now() + chrono::Duration::hours(1), + Some(MOCK_USER_AGENT.into()), + permission, + ) + .await + .unwrap(); + tokens.push(token); + } + let (read_only, read_write) = (&tokens[0], &tokens[1]); + + // No credentials, no answer. + for (method, path) in [ + (Method::GET, &listing), + (Method::GET, &of_org), + (Method::PUT, &one_grant), + (Method::DELETE, &one_grant), + ] { + let anonymous = server.method(method.clone(), path).await; + assert_eq!(anonymous.status_code(), 401, "{method} {path}"); + } + + // Read-only: both listings, no writes. + for path in [&listing, &of_org] { + let response = admin_call(&server, Method::GET, path, read_only).await; + assert_eq!(response.status_code(), 200, "{path}: {}", response.text()); + } + for method in [Method::PUT, Method::DELETE] { + let response = admin_call(&server, method.clone(), &one_grant, read_only).await; + assert_eq!(response.status_code(), 403, "{method}"); + assert_eq!( + response.json::()["error"]["type"], + "insufficient_permissions" + ); + } + let access: Value = admin_call(&server, Method::GET, &listing, read_only) + .await + .json(); + assert_eq!(access["organizations"], json!([])); + + // Read-write: grant and revoke. + let granted = admin_call(&server, Method::PUT, &one_grant, read_write).await; + assert_eq!(granted.status_code(), 200, "{}", granted.text()); + assert_eq!(granted.json::()["grantedByUserId"], MOCK_USER_ID); + let revoked = admin_call(&server, Method::DELETE, &one_grant, read_write).await; + assert_eq!(revoked.status_code(), 204); +} + +async fn internal_document(server: &axum_test::TestServer) -> Value { + let response = server + .get("/v1/internal/model-access") + .add_header("Authorization", format!("Bearer {INTERNAL_TOKEN}")) + .await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + assert_eq!(response.header("cache-control"), "no-store"); + let document: Value = response.json(); + assert_eq!(document["object"], "model_access"); + document +} + +/// The entry of `model` in `GET /v1/internal/model-access`. +async fn internal_entry(server: &axum_test::TestServer, model: &str) -> Option { + internal_document(server).await["models"] + .as_array() + .expect("models is an array") + .iter() + .find(|entry| entry["id"] == model) + .cloned() +} + +fn sorted(mut ids: Vec<&str>) -> Vec<&str> { + ids.sort_unstable(); + ids +} + +/// (d) `GET /v1/internal/model-access`: every unlisted model, active or not, +/// with the active organizations that hold a grant. +#[tokio::test] +async fn internal_document_lists_restricted_models_and_their_organizations() { + let fixture = setup().await; + let server = &fixture.server; + let first = organization(server).await; + let second = organization(server).await; + let model = &fixture.text; + + // Restricted, and nobody may use it: listed with no organizations. + assert_eq!( + internal_entry(server, model).await, + Some(json!({ "id": model, "organization_ids": [] })) + ); + + grant(server, model, &first.id).await; + grant(server, model, &second.id).await; + let entry = internal_entry(server, model).await.expect("restricted"); + assert_eq!( + entry.as_object().expect("entry is an object").len(), + 2, + "the entry carries id and organization_ids only: {entry}" + ); + let ids: Vec<&str> = entry["organization_ids"] + .as_array() + .expect("organization_ids") + .iter() + .map(|id| id.as_str().expect("an organization id")) + .collect(); + assert_eq!(ids, sorted(vec![first.id.as_str(), second.id.as_str()])); + + revoke(server, model, &first.id).await; + assert_eq!( + internal_entry(server, model).await, + Some(json!({ "id": model, "organization_ids": [second.id] })) + ); + + // A deleted organization no longer counts, and its grant is gone. + let deleted = server + .delete(&format!("/v1/organizations/{}", second.id)) + .add_header("Authorization", format!("Bearer {}", get_session_id())) + .add_header("User-Agent", MOCK_USER_AGENT) + .await; + assert!( + deleted.status_code().is_success(), + "deleting the organization failed: {}", + deleted.text() + ); + assert_eq!( + internal_entry(server, model).await, + Some(json!({ "id": model, "organization_ids": [] })) + ); + let access: Value = admin_call(server, Method::GET, &access_path(model), &get_session_id()) + .await + .json(); + assert_eq!(access["organizations"], json!([])); + let rows: i64 = fixture + .database + .pool() + .get() + .await + .expect("database connection") + .query_one( + "SELECT COUNT(*) FROM organization_model_access WHERE organization_id = $1", + &[&Uuid::parse_str(&second.id).expect("organization id")], + ) + .await + .expect("count grants") + .get(0); + assert_eq!(rows, 0, "deleting an organization removes its grants"); + + // A listed model is unrestricted: it is not in the document, even + // though a grant on it exists. + grant(server, model, &first.id).await; + set_listed(server, model, true).await; + assert_eq!(internal_entry(server, model).await, None); +} + +/// (d) Authentication of the internal document: the service token only. +#[tokio::test] +async fn internal_document_requires_the_service_token() { + let fixture = setup().await; + let server = &fixture.server; + + let missing = server.get("/v1/internal/model-access").await; + assert_eq!(missing.status_code(), 401, "{}", missing.text()); + + let org = organization(server).await; + for token in ["not-the-token", &org.api_key, &get_session_id()] { + let response = server + .get("/v1/internal/model-access") + .add_header("Authorization", format!("Bearer {token}")) + .add_header("User-Agent", MOCK_USER_AGENT) + .await; + assert_eq!(response.status_code(), 401, "{}", response.text()); + } + + // Disabled until the token is configured. + let unconfigured = setup_test_server().await; + let response = unconfigured + .get("/v1/internal/model-access") + .add_header("Authorization", format!("Bearer {INTERNAL_TOKEN}")) + .await; + assert_eq!(response.status_code(), 503, "{}", response.text()); + assert_eq!( + response.json::()["error"]["type"], + "endpoint_disabled" + ); +} + +/// (e) The ITA token route is public, so no organization is known: an +/// unlisted model is unknown there, whoever holds a grant. +#[tokio::test] +async fn public_ita_token_route_treats_an_unlisted_model_as_unknown() { + let fake_ita = FakeIta::start(FakeItaMode::Success).await; + let server = setup_ita_server(&fake_ita, ItaServerMode::Enabled { max_retries: 0 }).await; + let suffix = Uuid::new_v4(); + let model = format!("access-ita/model-{suffix}"); + let alias = format!("access-ita-alias-{suffix}"); + upsert_models( + &server, + vec![(&model, model_fields(json!({ "aliases": [alias] })))], + ) + .await; + // A grant changes nothing for an anonymous caller. + let org = organization(&server).await; + grant(&server, &model, &org.id).await; + + let token = |model: &str, strict: bool| { + let mut request = server.get(&format!( + "/v1/attestation/ita-token?model={}&nonce={NONCE}", + urlencoding::encode(model) + )); + if strict { + request = request.add_header("x-no-aliasing", "true"); + } + request + }; + for (name, strict) in [(&model, false), (&alias, false), (&alias, true)] { + let unknown = unknown_twin(name); + let refused = token(name, strict).await; + assert_eq!(refused.status_code(), 400, "{}", refused.text()); + assert_eq!( + refused.json::()["error"], + json!({ + "message": format!("Model '{name}' not found. It's not a valid model name or alias."), + "type": "invalid_request_error", + "param": "model", + "code": null, + }) + ); + assert!(!refused.text().contains(&format!("'{model}'")) || name == &model); + assert_eq!( + observe(&refused, name), + observe(&token(&unknown, strict).await, &unknown) + ); + } + // No model evidence was appraised: only verifier nonces were fetched. + assert!( + fake_ita + .paths() + .iter() + .all(|path| *path == "/appraisal/v2/nonce"), + "{:?}", + fake_ita.paths() + ); +} + +/// The public catalog detail route keeps answering 404 for an unlisted model. +#[tokio::test] +async fn public_model_detail_does_not_show_an_unlisted_model() { + let fixture = setup().await; + let server = &fixture.server; + for name in [&fixture.text, &fixture.text_alias] { + let unknown = unknown_twin(name); + let detail = |model: &str| server.get(&format!("/v1/model/{}", urlencoding::encode(model))); + assert_eq!( + observe(&detail(name).await, name), + observe(&detail(&unknown).await, &unknown) + ); + } +} + +/// (f) Usage that was served is billed: `POST /v1/internal/usage` records +/// usage on an unlisted model for an organization that holds no grant. +#[tokio::test] +async fn usage_is_recorded_for_an_organization_without_a_grant() { + let fixture = setup().await; + let server = &fixture.server; + let org = setup_org_with_credits(server, 10_000_000_000i64).await; + let workspace_id = list_workspaces(server, org.id.clone()) + .await + .first() + .expect("org has a default workspace") + .id + .clone(); + let key = create_api_key_in_workspace(server, workspace_id.clone(), "usage".to_string()).await; + + // The organization holds no grant: it cannot call the model here... + let access: Value = admin_call( + server, + Method::GET, + &format!("/v1/admin/organizations/{}/model-access", org.id), + &get_session_id(), + ) + .await + .json(); + assert_eq!(access["models"], json!([])); + let refused = send( + server, + Route::Embeddings, + key.key.as_deref().expect("a new key carries its secret"), + &fixture.text, + ) + .await; + assert_eq!(refused.status_code(), 404); + + // ...yet usage a trusted service reports for it is recorded and priced. + let response = server + .post("/v1/internal/usage") + .add_header("Authorization", format!("Bearer {INTERNAL_TOKEN}")) + .json(&json!({ + "organization_id": org.id, + "workspace_id": workspace_id, + "api_key_id": key.id, + "type": "chat_completion", + "model": fixture.text, + "input_tokens": 1000, + "output_tokens": 200, + "id": format!("model-access-usage-{}", Uuid::new_v4()) + })) + .await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + let body: Value = response.json(); + assert_eq!(body["model"], fixture.text.as_str()); + assert_eq!(body["input_cost"], 1_000_000_000i64); + assert_eq!(body["output_cost"], 400_000_000i64); +} diff --git a/crates/database/src/migrations/sql/V0097__add_organization_model_access.sql b/crates/database/src/migrations/sql/V0097__add_organization_model_access.sql new file mode 100644 index 000000000..c919f7efb --- /dev/null +++ b/crates/database/src/migrations/sql/V0097__add_organization_model_access.sql @@ -0,0 +1,61 @@ +-- Per-organization access to unlisted models. +-- +-- V0091 added `models.is_listed`. It only kept a model out of the public +-- catalogs: anyone who knew the name of an unlisted model could still call it. +-- From this migration on an unlisted model is a restricted model. An +-- organization may use it through this API only while it holds a row here (a +-- grant). Everyone else gets the answer they would get for a model that does +-- not exist. +-- +-- Rules, enforced where a model is resolved for a caller: +-- - A listed model is open to everyone. Grants on it have no effect, so access +-- can be prepared before the model is unlisted. +-- - A grant counts only while both the organization and the model are active. +-- +-- A grant is a relationship between two rows, written by admins, so it gets +-- its own table. No existing column can carry it: `organizations.settings` is +-- customer-editable and must never hold access state. +CREATE TABLE organization_model_access ( + organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE, + model_id UUID NOT NULL REFERENCES models(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + -- The admin who granted access. Organizations, models and users are soft + -- deleted, so the cascades above and this SET NULL are for a manual hard + -- delete only (V0058, V0086 precedent). + created_by_user_id UUID REFERENCES users(id) ON DELETE SET NULL, + PRIMARY KEY (organization_id, model_id) +); + +-- The primary key serves "does this organization hold a grant for this +-- model" and "which models may this organization use". This index serves the +-- other direction: "which organizations may use this model". +CREATE INDEX idx_organization_model_access_model_id + ON organization_model_access (model_id); + +-- The audit trail of the table above. A revocation deletes its grant row, so +-- the table alone cannot say who had access when, or who took it away. +-- +-- Append-only: one row per grant and one per revocation, written in the same +-- statement as the change, and only when the change happened (granting a +-- grant that exists, or revoking one that does not, writes nothing). +-- +-- The foreign keys carry no ON DELETE action, like the role audit log +-- (V0077): an audit row keeps the rows it names from being hard deleted. User +-- erasure rewrites the `users` row in place and keeps its id, and this table +-- holds ids only, so there is nothing for it to scrub here. +CREATE TABLE organization_model_access_log ( + id UUID PRIMARY KEY DEFAULT uuid_generate_v4(), + organization_id UUID NOT NULL REFERENCES organizations(id), + model_id UUID NOT NULL REFERENCES models(id), + action VARCHAR(10) NOT NULL CHECK (action IN ('grant', 'revoke')), + -- The admin who made the change. NULL for a revocation made by the + -- platform: deleting an organization removes its grants. + acting_user_id UUID REFERENCES users(id), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_organization_model_access_log_model_created_at + ON organization_model_access_log (model_id, created_at DESC); + +CREATE INDEX idx_organization_model_access_log_org_created_at + ON organization_model_access_log (organization_id, created_at DESC); diff --git a/crates/database/src/models.rs b/crates/database/src/models.rs index 095cfcbf5..89fe08e43 100644 --- a/crates/database/src/models.rs +++ b/crates/database/src/models.rs @@ -535,10 +535,11 @@ pub struct Model { /// If true, this model may be activated even when both cost fields are 0. /// Intended for intentionally-free models (e.g. community previews). pub allow_free: bool, - /// Whether the model appears in the public model catalogs - /// (`GET /v1/models`, `GET /v1/model/list`, `GET /v1/model/{name}`). - /// `false` keeps an active, billable row out of them. It is a catalog - /// switch only: it does not change routing or who may call the model. + /// Whether the model is listed. A listed model appears in the public + /// model catalogs (`GET /v1/models`, `GET /v1/model/list`, + /// `GET /v1/model/{name}`) and every caller may use it. An unlisted model + /// stays active and billable but is left out of them, and only an + /// organization with a row in `organization_model_access` may use it. pub is_listed: bool, pub owned_by: String, pub created_at: DateTime, @@ -567,9 +568,10 @@ pub struct UpdateModelPricingRequest { pub is_active: Option, /// If true, allows activation even with zero pricing. pub allow_free: Option, - /// `Some(false)` keeps the model out of the public model catalogs, - /// `Some(true)` lists it again, `None` leaves the column unchanged - /// (new rows default to listed). + /// `Some(false)` unlists the model: it leaves the public model catalogs + /// and only organizations that hold a grant may use it. `Some(true)` + /// lists it again, `None` leaves the column unchanged (new rows default + /// to listed). pub is_listed: Option, pub aliases: Option>, pub owned_by: Option, @@ -679,7 +681,8 @@ pub struct ModelHistory { /// If true, this model was allowed to serve without pricing at this point in time. pub allow_free: bool, - /// Whether the model was in the public model catalogs at this point in time. + /// Whether the model was listed (in the public model catalogs and open to + /// every caller) at this point in time. pub is_listed: bool, } diff --git a/crates/database/src/repositories/admin_composite.rs b/crates/database/src/repositories/admin_composite.rs index 74f19c9ed..34d5ed2f4 100644 --- a/crates/database/src/repositories/admin_composite.rs +++ b/crates/database/src/repositories/admin_composite.rs @@ -9,14 +9,14 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use services::admin::{ AdminApiKeyFilters, AdminApiKeyInfo, AdminModelInfo, AdminOrganizationInfo, - AdminOrganizationMemberInfo, AdminRepository, DeprecateModelOutcome, - ModelDeprecationDeliveryRecord, ModelDeprecationEmailStatus, ModelDeprecationModel, - ModelDeprecationRecipient, ModelHistoryEntry, ModelPricing, ModelPricingSnapshot, - ModelValidationState, OrganizationLifecycleFilter, OrganizationLimits, - OrganizationLimitsHistoryEntry, OrganizationLimitsUpdate, PlatformServiceInfo, - PricingChangeDeliveryRecord, PricingChangeOpenConflictError, PricingChangeRecipientRow, - ScheduledPricingChange, ScheduledPricingChangeInsert, ScheduledPricingChangeStatus, - UpdateModelAdminRequest, UserInfo, UserOrganizationInfo, + AdminOrganizationMemberInfo, AdminRepository, DeprecateModelOutcome, ModelAccessGrant, + ModelAccessModel, ModelDeprecationDeliveryRecord, ModelDeprecationEmailStatus, + ModelDeprecationModel, ModelDeprecationRecipient, ModelHistoryEntry, ModelPricing, + ModelPricingSnapshot, ModelValidationState, OrganizationLifecycleFilter, OrganizationLimits, + OrganizationLimitsHistoryEntry, OrganizationLimitsUpdate, OrganizationModelAccessGrant, + PlatformServiceInfo, PricingChangeDeliveryRecord, PricingChangeOpenConflictError, + PricingChangeRecipientRow, ScheduledPricingChange, ScheduledPricingChangeInsert, + ScheduledPricingChangeStatus, UpdateModelAdminRequest, UserInfo, UserOrganizationInfo, }; use services::lifecycle::OrganizationLifecycle; use services::service_usage::ports::ServiceUnit; @@ -161,6 +161,41 @@ fn row_to_admin_api_key_info(row: &tokio_postgres::Row) -> AdminApiKeyInfo { } } +/// One grant with its organization and the admin who made it. Only active +/// organizations: a grant held by an inactive one has no effect. Callers add +/// the `WHERE` clause. +const MODEL_ACCESS_GRANT_SELECT: &str = r#" + SELECT + a.organization_id, + o.name AS organization_name, + a.created_at, + a.created_by_user_id, + u.email AS created_by_user_email + FROM organization_model_access a + JOIN organizations o ON o.id = a.organization_id AND o.is_active = true + LEFT JOIN users u ON u.id = a.created_by_user_id +"#; + +fn row_to_model_access_grant(row: &tokio_postgres::Row) -> ModelAccessGrant { + ModelAccessGrant { + organization_id: row.get("organization_id"), + organization_name: row.get("organization_name"), + granted_at: row.get("created_at"), + granted_by_user_id: row.get("created_by_user_id"), + granted_by_user_email: row.get("created_by_user_email"), + } +} + +/// The row must select `id, model_name, is_listed, is_active` of `models`. +fn row_to_model_access_model(row: &tokio_postgres::Row) -> ModelAccessModel { + ModelAccessModel { + id: row.get("id"), + model_name: row.get("model_name"), + is_listed: row.get("is_listed"), + is_active: row.get("is_active"), + } +} + fn service_to_info(s: &crate::models::Service) -> Result { let unit = ServiceUnit::try_from(s.unit.as_str()).map_err(|e| anyhow::anyhow!("{}", e))?; Ok(PlatformServiceInfo { @@ -1908,6 +1943,155 @@ impl AdminRepository for AdminCompositeRepository { Ok(row.get::<_, bool>("exists")) } + async fn get_model_for_access(&self, model_name: &str) -> Result> { + let client = self.pool.get().await?; + let row = client + .query_opt( + r#" + SELECT id, model_name, is_listed, is_active + FROM models + WHERE model_name = $1 + "#, + &[&model_name], + ) + .await?; + + Ok(row.as_ref().map(row_to_model_access_model)) + } + + async fn list_model_access_grants(&self, model_id: Uuid) -> Result> { + let client = self.pool.get().await?; + let rows = client + .query( + &format!("{MODEL_ACCESS_GRANT_SELECT} WHERE a.model_id = $1 ORDER BY a.created_at, a.organization_id"), + &[&model_id], + ) + .await?; + + Ok(rows.iter().map(row_to_model_access_grant).collect()) + } + + async fn grant_model_access( + &self, + model_id: Uuid, + organization_id: Uuid, + granted_by_user_id: Option, + ) -> Result> { + let mut client = self.pool.get().await?; + let transaction = client.transaction().await?; + // Hold the organization row until the grant is written. Deleting an + // organization locks the same row to deactivate it and then removes + // its grants in that transaction, so the two cannot interleave: either + // the deletion waits and then removes this grant too, or this waits + // for the deletion and finds the organization inactive. Without the + // lock a grant could land after the deletion had removed the others. + let active = transaction + .query_opt( + "SELECT 1 FROM organizations WHERE id = $1 AND is_active = true FOR SHARE", + &[&organization_id], + ) + .await? + .is_some(); + if !active { + return Ok(None); + } + // One statement, so the grant and its audit row are written together + // or not at all. An existing grant keeps its date and its author, and + // then nothing is written to the audit log either. + let created = transaction + .execute( + r#" + WITH granted AS ( + INSERT INTO organization_model_access (organization_id, model_id, created_by_user_id) + VALUES ($1, $2, $3) + ON CONFLICT (organization_id, model_id) DO NOTHING + RETURNING organization_id, model_id, created_by_user_id + ) + INSERT INTO organization_model_access_log (organization_id, model_id, action, acting_user_id) + SELECT organization_id, model_id, 'grant', created_by_user_id FROM granted + "#, + &[&organization_id, &model_id, &granted_by_user_id], + ) + .await? + > 0; + let row = transaction + .query_opt( + &format!( + "{MODEL_ACCESS_GRANT_SELECT} WHERE a.model_id = $1 AND a.organization_id = $2" + ), + &[&model_id, &organization_id], + ) + .await? + .context("the grant is gone: revoked meanwhile")?; + transaction.commit().await?; + + Ok(Some((row_to_model_access_grant(&row), created))) + } + + async fn revoke_model_access( + &self, + model_id: Uuid, + organization_id: Uuid, + revoked_by_user_id: Option, + ) -> Result { + let client = self.pool.get().await?; + // One statement, so the grant goes and its audit row is written + // together or not at all. No grant, no audit row. + let revoked = client + .execute( + r#" + WITH revoked AS ( + DELETE FROM organization_model_access + WHERE organization_id = $1 AND model_id = $2 + RETURNING organization_id, model_id + ) + INSERT INTO organization_model_access_log (organization_id, model_id, action, acting_user_id) + SELECT organization_id, model_id, 'revoke', $3::UUID FROM revoked + "#, + &[&organization_id, &model_id, &revoked_by_user_id], + ) + .await?; + + Ok(revoked > 0) + } + + async fn list_organization_model_access( + &self, + organization_id: Uuid, + ) -> Result> { + let client = self.pool.get().await?; + let rows = client + .query( + r#" + SELECT + m.id, + m.model_name, + m.is_listed, + m.is_active, + a.created_at, + a.created_by_user_id, + u.email AS created_by_user_email + FROM organization_model_access a + JOIN models m ON m.id = a.model_id + LEFT JOIN users u ON u.id = a.created_by_user_id + WHERE a.organization_id = $1 + ORDER BY m.model_name + "#, + &[&organization_id], + ) + .await?; + + Ok(rows + .iter() + .map(|row| OrganizationModelAccessGrant { + model: row_to_model_access_model(row), + granted_at: row.get("created_at"), + granted_by_user_id: row.get("created_by_user_id"), + granted_by_user_email: row.get("created_by_user_email"), + }) + .collect()) + } + async fn list_services( &self, include_inactive: bool, diff --git a/crates/database/src/repositories/model.rs b/crates/database/src/repositories/model.rs index 358d661c5..0841989bc 100644 --- a/crates/database/src/repositories/model.rs +++ b/crates/database/src/repositories/model.rs @@ -1408,6 +1408,95 @@ impl ModelRepository { Ok(row) } + /// Whether the organization holds a grant for the model that is in + /// effect. A grant held by an inactive organization, or on an inactive + /// model, has no effect. + pub async fn organization_has_model_access( + &self, + organization_id: uuid::Uuid, + model_id: uuid::Uuid, + ) -> Result { + let row = retry_db!("organization_has_model_access", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + // One-shot typed query: one round trip, see repositories/mod.rs. + client + .query_typed_opt( + r#" + SELECT 1 + FROM organization_model_access a + JOIN organizations o ON o.id = a.organization_id + JOIN models m ON m.id = a.model_id + WHERE a.organization_id = $1 + AND a.model_id = $2 + AND o.is_active = true + AND m.is_active = true + "#, + &[(&organization_id, Type::UUID), (&model_id, Type::UUID)], + ) + .await + .map_err(map_db_error) + })?; + + Ok(row.is_some()) + } + + /// Every unlisted model, active or not, with the active organizations + /// that hold a grant for it, ordered by model name. A model nobody may + /// use is listed with no organizations. + /// + /// Inactive models are included on purpose. The reader takes a model that + /// is absent from the list to be unrestricted, and it may still be + /// serving a model that was deactivated here. + pub async fn list_restricted_model_access( + &self, + ) -> Result> { + let rows = retry_db!("list_restricted_model_access", { + let client = self + .pool + .get() + .await + .context("Failed to get database connection") + .map_err(RepositoryError::PoolError)?; + + client + .query( + r#" + SELECT + m.model_name, + COALESCE( + array_agg(o.id ORDER BY o.id) FILTER (WHERE o.id IS NOT NULL), + '{}' + ) AS organization_ids + FROM models m + LEFT JOIN organization_model_access a ON a.model_id = m.id + LEFT JOIN organizations o + ON o.id = a.organization_id + AND o.is_active = true + WHERE m.is_listed = false + GROUP BY m.id + ORDER BY m.model_name + "#, + &[], + ) + .await + .map_err(map_db_error) + })?; + + Ok(rows + .iter() + .map(|row| services::models::RestrictedModelAccess { + model_name: row.get("model_name"), + organization_ids: row.get("organization_ids"), + }) + .collect()) + } + /// Helper method to convert database row to Model fn row_to_model(&self, row: &Row) -> Model { let model_name: String = row.get("model_name"); @@ -1739,4 +1828,19 @@ impl services::models::ModelsRepository for ModelRepository { async fn get_configured_model_names(&self) -> Result> { self.get_configured_model_names().await } + + async fn organization_has_model_access( + &self, + organization_id: uuid::Uuid, + model_id: uuid::Uuid, + ) -> Result { + self.organization_has_model_access(organization_id, model_id) + .await + } + + async fn list_restricted_model_access( + &self, + ) -> Result> { + self.list_restricted_model_access().await + } } diff --git a/crates/database/src/repositories/organization.rs b/crates/database/src/repositories/organization.rs index f084abfbf..96d1812c7 100644 --- a/crates/database/src/repositories/organization.rs +++ b/crates/database/src/repositories/organization.rs @@ -717,6 +717,27 @@ where .await .map_err(map_db_error)?; + // A grant on an unlisted model counts only for an active organization, + // but remove the rows as well, for the reason given for reporting + // tokens above: a deleted org cannot regain access through a future path. + // Each removed grant gets its audit row; no acting user, as above: this + // is a system revocation. + client + .execute( + r#" + WITH revoked AS ( + DELETE FROM organization_model_access + WHERE organization_id = $1 + RETURNING organization_id, model_id + ) + INSERT INTO organization_model_access_log (organization_id, model_id, action, acting_user_id) + SELECT organization_id, model_id, 'revoke', NULL::UUID FROM revoked + "#, + &[&organization_id], + ) + .await + .map_err(map_db_error)?; + Ok(()) } diff --git a/crates/database/tests/hot_path_typed_queries.rs b/crates/database/tests/hot_path_typed_queries.rs index 15e9a5549..59d4f4ec1 100644 --- a/crates/database/tests/hot_path_typed_queries.rs +++ b/crates/database/tests/hot_path_typed_queries.rs @@ -208,6 +208,67 @@ async fn typed_hot_path_queries_repeat_on_one_pool() -> anyhow::Result<()> { .await? .is_none()); + // model::organization_has_model_access: no grant, then a grant, which + // counts only while the organization and the model are both active + for _ in 0..3 { + assert!( + !models + .organization_has_model_access(org_id, model_id) + .await? + ); + } + { + let client = pool.get().await?; + client + .execute( + "INSERT INTO organization_model_access (organization_id, model_id, created_by_user_id) VALUES ($1, $2, $3)", + &[&org_id, &model_id, &user_id], + ) + .await?; + for _ in 0..3 { + assert!( + models + .organization_has_model_access(org_id, model_id) + .await? + ); + } + assert!( + !models + .organization_has_model_access(Uuid::new_v4(), model_id) + .await? + ); + assert!( + !models + .organization_has_model_access(org_id, Uuid::new_v4()) + .await? + ); + for (table, id) in [("models", model_id), ("organizations", org_id)] { + client + .execute( + &format!("UPDATE {table} SET is_active = false WHERE id = $1"), + &[&id], + ) + .await?; + assert!( + !models + .organization_has_model_access(org_id, model_id) + .await?, + "a grant must not count while its row in {table} is inactive" + ); + client + .execute( + &format!("UPDATE {table} SET is_active = true WHERE id = $1"), + &[&id], + ) + .await?; + } + assert!( + models + .organization_has_model_access(org_id, model_id) + .await? + ); + } + // organization concurrent limit let orgs = PgOrganizationRepository::new(pool.clone()); for _ in 0..3 { diff --git a/crates/services/src/admin/email_loop_tests.rs b/crates/services/src/admin/email_loop_tests.rs index ed17f43f4..4b86a4543 100644 --- a/crates/services/src/admin/email_loop_tests.rs +++ b/crates/services/src/admin/email_loop_tests.rs @@ -1,5 +1,10 @@ -//! Unit tests for the email batch loops in `AdminServiceImpl`: the recipient -//! activity lookup must only run immediately before a real send. +//! Unit tests for `AdminServiceImpl` over a counting repository. +//! +//! Email batch loops: the recipient activity lookup must only run immediately +//! before a real send. +//! +//! Model access: a grant or a revocation drops one cached access decision, +//! and only when it changed something. use super::*; use crate::admin::{AdminService, AdminServiceImpl}; @@ -55,6 +60,15 @@ struct CountingRepo { deprecation_recipients: Vec, pricing_recipients: Vec, scheduled_rows: Vec, + /// The model the access operations find, if any. + access_model: Option, + /// Whether a grant creates a row (false: the grant already existed). + grant_creates: bool, + /// Whether the grant finds the organization deleted, after the service + /// checked that it was active. + grant_finds_organization_deleted: bool, + /// Whether a revocation finds a grant to delete. + revoke_finds: bool, } #[allow(unused_variables)] @@ -370,7 +384,7 @@ impl AdminRepository for CountingRepo { &self, organization_id: uuid::Uuid, ) -> Result { - unimplemented!("not used by the email loop tests") + Ok(true) } async fn list_services( @@ -410,6 +424,57 @@ impl AdminRepository for CountingRepo { ) -> Result, anyhow::Error> { unimplemented!("not used by the email loop tests") } + + async fn get_model_for_access( + &self, + model_name: &str, + ) -> Result, anyhow::Error> { + Ok(self.access_model.clone()) + } + + async fn list_model_access_grants( + &self, + model_id: uuid::Uuid, + ) -> Result, anyhow::Error> { + unimplemented!("not used by the email loop tests") + } + + async fn grant_model_access( + &self, + model_id: uuid::Uuid, + organization_id: uuid::Uuid, + granted_by_user_id: Option, + ) -> Result, anyhow::Error> { + if self.grant_finds_organization_deleted { + return Ok(None); + } + Ok(Some(( + ModelAccessGrant { + organization_id, + organization_name: "org".to_string(), + granted_at: chrono::Utc::now(), + granted_by_user_id, + granted_by_user_email: None, + }, + self.grant_creates, + ))) + } + + async fn revoke_model_access( + &self, + model_id: uuid::Uuid, + organization_id: uuid::Uuid, + revoked_by_user_id: Option, + ) -> Result { + Ok(self.revoke_finds) + } + + async fn list_organization_model_access( + &self, + organization_id: uuid::Uuid, + ) -> Result, anyhow::Error> { + unimplemented!("not used by the email loop tests") + } } struct CountingEmailSender { @@ -651,3 +716,153 @@ async fn pricing_loop_looks_up_activity_once_before_a_real_send() { assert_eq!(repo.active_calls.load(Ordering::SeqCst), 1); assert_eq!(sender.sends.load(Ordering::SeqCst), 1); } + +/// A service whose models service accepts exactly `access_invalidations` +/// calls to drop the access decision of (`organization_id`, the model), and +/// no other invalidation. The mock checks the count when it is dropped. +fn access_service( + repo: CountingRepo, + organization_id: Uuid, + access_invalidations: usize, +) -> AdminServiceImpl { + let model_id = Uuid::new_v4(); + let repo = CountingRepo { + access_model: Some(ModelAccessModel { + id: model_id, + model_name: "restricted/model".to_string(), + is_listed: false, + is_active: true, + }), + ..repo + }; + let mut models = MockModelsServiceTrait::new(); + // A grant changes neither the model list nor a resolved model, and + // reloading the list asks every backend. + models.expect_invalidate_models_cache().times(0); + models + .expect_invalidate_model_access() + .withf(move |organization, model| *organization == organization_id && *model == model_id) + .times(access_invalidations) + .returning(|_, _| ()); + AdminServiceImpl::new( + Arc::new(repo), + Arc::new(models), + Arc::new(MockCompletionServiceTrait::new()), + Arc::new(CountingEmailSender { + sends: AtomicUsize::new(0), + }), + ) +} + +#[tokio::test] +async fn a_new_grant_drops_only_that_access_decision() { + let organization_id = Uuid::new_v4(); + let repo = CountingRepo { + grant_creates: true, + ..Default::default() + }; + let service = access_service(repo, organization_id, 1); + + let grant = service + .grant_model_access("restricted/model", organization_id, Some(Uuid::new_v4())) + .await + .unwrap(); + + assert_eq!(grant.organization_id, organization_id); +} + +#[tokio::test] +async fn a_grant_that_already_exists_invalidates_nothing() { + let organization_id = Uuid::new_v4(); + let repo = CountingRepo { + grant_creates: false, + ..Default::default() + }; + let service = access_service(repo, organization_id, 0); + + service + .grant_model_access("restricted/model", organization_id, Some(Uuid::new_v4())) + .await + .unwrap(); +} + +#[tokio::test] +async fn a_grant_for_an_organization_deleted_meanwhile_is_not_found() { + let organization_id = Uuid::new_v4(); + let repo = CountingRepo { + grant_creates: true, + grant_finds_organization_deleted: true, + ..Default::default() + }; + let service = access_service(repo, organization_id, 0); + + assert!(matches!( + service + .grant_model_access("restricted/model", organization_id, Some(Uuid::new_v4())) + .await, + Err(AdminError::OrganizationNotFound(_)) + )); +} + +#[tokio::test] +async fn a_revocation_drops_only_that_access_decision() { + let organization_id = Uuid::new_v4(); + let repo = CountingRepo { + revoke_finds: true, + ..Default::default() + }; + let service = access_service(repo, organization_id, 1); + + service + .revoke_model_access("restricted/model", organization_id, Some(Uuid::new_v4())) + .await + .unwrap(); +} + +#[tokio::test] +async fn a_revocation_of_nothing_invalidates_nothing() { + let organization_id = Uuid::new_v4(); + let repo = CountingRepo { + revoke_finds: false, + ..Default::default() + }; + let service = access_service(repo, organization_id, 0); + + service + .revoke_model_access("restricted/model", organization_id, Some(Uuid::new_v4())) + .await + .unwrap(); +} + +#[tokio::test] +async fn access_operations_on_an_unknown_model_change_nothing() { + let organization_id = Uuid::new_v4(); + let mut models = MockModelsServiceTrait::new(); + models.expect_invalidate_models_cache().times(0); + models.expect_invalidate_model_access().times(0); + let service = AdminServiceImpl::new( + Arc::new(CountingRepo { + grant_creates: true, + revoke_finds: true, + ..Default::default() + }), + Arc::new(models), + Arc::new(MockCompletionServiceTrait::new()), + Arc::new(CountingEmailSender { + sends: AtomicUsize::new(0), + }), + ); + + assert!(matches!( + service + .grant_model_access("no/such-model", organization_id, None) + .await, + Err(AdminError::ModelNotFound(_)) + )); + assert!(matches!( + service + .revoke_model_access("no/such-model", organization_id, None) + .await, + Err(AdminError::ModelNotFound(_)) + )); +} diff --git a/crates/services/src/admin/mod.rs b/crates/services/src/admin/mod.rs index f8a5e7f86..4da898be7 100644 --- a/crates/services/src/admin/mod.rs +++ b/crates/services/src/admin/mod.rs @@ -1497,9 +1497,138 @@ impl AdminService for AdminServiceImpl { .map_err(|e| AdminError::InternalError(e.to_string()))? .ok_or_else(|| AdminError::ServiceNotFound(format!("Service {id} not found"))) } + + async fn get_model_access(&self, model_name: &str) -> Result { + let model = self.model_for_access(model_name).await?; + let grants = self + .repository + .list_model_access_grants(model.id) + .await + .map_err(|e| AdminError::InternalError(e.to_string()))?; + + Ok(ModelAccess { model, grants }) + } + + async fn grant_model_access( + &self, + model_name: &str, + organization_id: uuid::Uuid, + granted_by_user_id: Option, + ) -> Result { + let model = self.model_for_access(model_name).await?; + self.require_active_organization(organization_id).await?; + + // The repository checks the organization again, under a lock: it can + // be deleted after the check above. + let (grant, created) = self + .repository + .grant_model_access(model.id, organization_id, granted_by_user_id) + .await + .map_err(|e| AdminError::InternalError(e.to_string()))? + .ok_or_else(|| { + AdminError::OrganizationNotFound(format!( + "Organization '{organization_id}' not found" + )) + })?; + + // A grant that already existed changed nothing: nothing to drop. + if created { + // Drop the one cached access decision, so the grant applies to + // this organization's next request on this instance. The model + // list and the resolved models stay cached: a grant changes + // neither, and reloading the list asks every backend. Other + // instances follow within the cache TTL. + self.models_service + .invalidate_model_access(organization_id, model.id) + .await; + + tracing::info!( + model_id = %model.id, + organization_id = %organization_id, + admin_user_id = ?granted_by_user_id, + "Model access granted" + ); + } + + Ok(grant) + } + + async fn revoke_model_access( + &self, + model_name: &str, + organization_id: uuid::Uuid, + revoked_by_user_id: Option, + ) -> Result<(), AdminError> { + let model = self.model_for_access(model_name).await?; + self.require_active_organization(organization_id).await?; + + let revoked = self + .repository + .revoke_model_access(model.id, organization_id, revoked_by_user_id) + .await + .map_err(|e| AdminError::InternalError(e.to_string()))?; + + // No grant existed: nothing changed, nothing to drop. + if revoked { + // Drop the one cached access decision, so the revocation applies + // to this organization's next request on this instance. Other + // instances follow within the cache TTL. + self.models_service + .invalidate_model_access(organization_id, model.id) + .await; + + tracing::info!( + model_id = %model.id, + organization_id = %organization_id, + admin_user_id = ?revoked_by_user_id, + "Model access revoked" + ); + } + + Ok(()) + } + + async fn list_organization_model_access( + &self, + organization_id: uuid::Uuid, + ) -> Result, AdminError> { + self.require_active_organization(organization_id).await?; + + self.repository + .list_organization_model_access(organization_id) + .await + .map_err(|e| AdminError::InternalError(e.to_string())) + } } impl AdminServiceImpl { + /// The model whose access is being managed, or `ModelNotFound`. + async fn model_for_access(&self, model_name: &str) -> Result { + self.repository + .get_model_for_access(model_name) + .await + .map_err(|e| AdminError::InternalError(e.to_string()))? + .ok_or_else(|| AdminError::ModelNotFound(format!("Model '{model_name}' not found"))) + } + + async fn require_active_organization( + &self, + organization_id: uuid::Uuid, + ) -> Result<(), AdminError> { + let exists = self + .repository + .organization_exists(organization_id) + .await + .map_err(|e| AdminError::InternalError(e.to_string()))?; + if exists { + Ok(()) + } else { + Err(AdminError::OrganizationNotFound(format!( + "Organization '{organization_id}' not found" + ))) + } + } + async fn validate_model_request( model_name: &str, request: &UpdateModelAdminRequest, diff --git a/crates/services/src/admin/ports.rs b/crates/services/src/admin/ports.rs index 5867b5f9a..1d4a80e26 100644 --- a/crates/services/src/admin/ports.rs +++ b/crates/services/src/admin/ports.rs @@ -23,8 +23,9 @@ pub struct UpdateModelAdminRequest { pub is_active: Option, /// If true, allows activation even with zero pricing. pub allow_free: Option, - /// `Some(false)` keeps the model out of the public model catalogs, - /// `Some(true)` lists it again, `None` leaves it unchanged. + /// `Some(false)` unlists the model: it leaves the public model catalogs + /// and only organizations that hold a grant may use it. `Some(true)` + /// lists it again, `None` leaves it unchanged. pub is_listed: Option, pub aliases: Option>, pub owned_by: Option, @@ -131,8 +132,9 @@ pub struct ModelPricing { pub successor_model_name: Option, /// OpenRouter `openrouter.slug` override. NULL = unset. pub openrouter_slug: Option, - /// Whether the model appears in the public model catalogs. `false` keeps - /// an active, billable row out of them; it does not change routing. + /// Whether the model is listed. An unlisted model is left out of the + /// public model catalogs and only organizations that hold a grant may + /// use it. pub is_listed: bool, } @@ -188,7 +190,8 @@ pub struct ModelHistoryEntry { pub openrouter_slug: Option, /// If true, this model was allowed to serve without pricing at this point in time. pub allow_free: bool, - /// Whether the model was in the public model catalogs at this point in time. + /// Whether the model was listed (in the public model catalogs and open to + /// every caller) at this point in time. pub is_listed: bool, pub effective_from: chrono::DateTime, pub effective_until: Option>, @@ -322,8 +325,9 @@ pub struct AdminModelInfo { pub successor_model_name: Option, /// OpenRouter `openrouter.slug` override. NULL = unset. pub openrouter_slug: Option, - /// Whether the model appears in the public model catalogs. `false` keeps - /// an active, billable row out of them; it does not change routing. + /// Whether the model is listed. An unlisted model is left out of the + /// public model catalogs and only organizations that hold a grant may + /// use it. pub is_listed: bool, } @@ -658,6 +662,45 @@ pub struct AdminOrganizationMemberInfo { pub user: UserInfo, } +/// A model as the access-management operations see it. Looked up by +/// canonical name, active or not, so a grant can be reviewed or revoked while +/// its model is inactive. +#[derive(Debug, Clone)] +pub struct ModelAccessModel { + pub id: uuid::Uuid, + pub model_name: String, + pub is_listed: bool, + pub is_active: bool, +} + +/// A grant: one organization that may use one unlisted model. +#[derive(Debug, Clone)] +pub struct ModelAccessGrant { + pub organization_id: uuid::Uuid, + pub organization_name: String, + pub granted_at: chrono::DateTime, + /// The admin who made the grant. `None` if that user row is gone. + pub granted_by_user_id: Option, + pub granted_by_user_email: Option, +} + +/// A model and the active organizations that hold a grant for it. +#[derive(Debug, Clone)] +pub struct ModelAccess { + pub model: ModelAccessModel, + pub grants: Vec, +} + +/// One model an organization holds a grant for. +#[derive(Debug, Clone)] +pub struct OrganizationModelAccessGrant { + pub model: ModelAccessModel, + pub granted_at: chrono::DateTime, + /// The admin who made the grant. `None` if that user row is gone. + pub granted_by_user_id: Option, + pub granted_by_user_email: Option, +} + #[derive(Debug, thiserror::Error)] pub enum AdminError { #[error("Model not found: {0}")] @@ -1023,6 +1066,49 @@ pub trait AdminRepository: Send + Sync { async fn organization_exists(&self, organization_id: uuid::Uuid) -> Result; + /// Look a model up by canonical name for access management, active or + /// not. `None` if no such model exists. + async fn get_model_for_access( + &self, + model_name: &str, + ) -> Result, anyhow::Error>; + + /// The active organizations that hold a grant for the model, oldest + /// grant first. + async fn list_model_access_grants( + &self, + model_id: uuid::Uuid, + ) -> Result, anyhow::Error>; + + /// Grant the organization access to the model. Returns the stored grant + /// and whether this call created it: a grant that already exists is + /// returned unchanged. A grant that is created gets its audit row in the + /// same statement. `None` if the organization is not active: the check + /// holds the organization row until the grant is written, so a + /// concurrent deletion cannot leave a grant behind. + async fn grant_model_access( + &self, + model_id: uuid::Uuid, + organization_id: uuid::Uuid, + granted_by_user_id: Option, + ) -> Result, anyhow::Error>; + + /// Revoke the organization's grant for the model. Returns whether a + /// grant existed. A grant that is revoked gets its audit row in the same + /// statement. + async fn revoke_model_access( + &self, + model_id: uuid::Uuid, + organization_id: uuid::Uuid, + revoked_by_user_id: Option, + ) -> Result; + + /// The models the organization holds a grant for, by model name. + async fn list_organization_model_access( + &self, + organization_id: uuid::Uuid, + ) -> Result, anyhow::Error>; + /// List platform services with pagination (admin only) async fn list_services( &self, @@ -1289,6 +1375,43 @@ pub trait AdminService: Send + Sync { cost_per_unit: Option, is_active: Option, ) -> Result; + + /// A model and the organizations that hold a grant for it (admin only). + /// Returns `ModelNotFound` if no model has this canonical name. + async fn get_model_access(&self, model_name: &str) -> Result; + + /// Grant an organization access to a model (admin only). Idempotent: a + /// grant that already exists is returned unchanged. Returns + /// `ModelNotFound` or `OrganizationNotFound` if either is missing; the + /// organization must be active. A grant that is created is written to + /// the audit log with `granted_by_user_id`. + async fn grant_model_access( + &self, + model_name: &str, + organization_id: uuid::Uuid, + granted_by_user_id: Option, + ) -> Result; + + /// Revoke an organization's access to a model (admin only). Idempotent: + /// revoking a grant that does not exist succeeds. Returns `ModelNotFound` + /// or `OrganizationNotFound` if either is missing; the organization must + /// be active, like for a grant. A deleted organization has no grants + /// left to revoke: its deletion removed them. A grant that is revoked is + /// written to the audit log with `revoked_by_user_id`. + async fn revoke_model_access( + &self, + model_name: &str, + organization_id: uuid::Uuid, + revoked_by_user_id: Option, + ) -> Result<(), AdminError>; + + /// The models an organization holds a grant for (admin only). Returns + /// `OrganizationNotFound` if the organization does not exist or is + /// inactive. + async fn list_organization_model_access( + &self, + organization_id: uuid::Uuid, + ) -> Result, AdminError>; } #[cfg(test)] diff --git a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs index c4c38eafc..c505a4f9e 100644 --- a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs +++ b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs @@ -286,6 +286,7 @@ pub(super) fn lifecycle_service( repository, inference_provider_pool: pool.clone(), models_repository: Arc::new(EmptyModelsRepository), + model_resolve_cache: crate::models::new_model_resolve_cache(), metrics_service: Arc::new(NoopMetricsService), usage_repository: Arc::new(NoopUsageRepository), vpc_info: None, diff --git a/crates/services/src/attestation/ita/service.rs b/crates/services/src/attestation/ita/service.rs index 7f7c9d632..382d44acd 100644 --- a/crates/services/src/attestation/ita/service.rs +++ b/crates/services/src/attestation/ita/service.rs @@ -17,6 +17,7 @@ use crate::{ GatewayQuoteInput, }, inference_provider_pool::InferenceProviderPool, + models::{resolve_model_cached, ModelCaller}, }; pub struct ModelAttestationInput { @@ -212,12 +213,17 @@ impl AttestationService { let Some(requested_model) = query.model.as_ref() else { return Ok(None); }; - let resolved_model = self - .models_repository - .resolve_and_get_model(requested_model) - .await - .map_err(|e| AttestationError::ProviderError(format!("Failed to resolve model: {e}")))? - .ok_or_else(|| AttestationError::UnknownModel(requested_model.clone()))?; + // This token is served to the public, so the model is resolved for an + // anonymous caller: an unlisted model is unknown here. + let resolved_model = resolve_model_cached( + &self.model_resolve_cache, + self.models_repository.as_ref(), + requested_model, + ModelCaller::Anonymous, + ) + .await + .map_err(|e| AttestationError::ProviderError(format!("Failed to resolve model: {e}")))? + .ok_or_else(|| AttestationError::UnknownModel(requested_model.clone()))?; let canonical_model = resolved_model.model_name; let model_alias_resolved = (canonical_model != *requested_model).then(|| ItaModelAliasResolved { diff --git a/crates/services/src/attestation/ita/service_test_support.rs b/crates/services/src/attestation/ita/service_test_support.rs index 995220818..adab2548e 100644 --- a/crates/services/src/attestation/ita/service_test_support.rs +++ b/crates/services/src/attestation/ita/service_test_support.rs @@ -82,6 +82,7 @@ fn service_with_collectors( ExternalProvidersConfig::default(), )), models_repository: Arc::new(FakeModelsRepository), + model_resolve_cache: crate::models::new_model_resolve_cache(), metrics_service: Arc::new(NoopMetricsService), usage_repository: Arc::new(NoopUsageRepository), vpc_info: None, @@ -252,6 +253,13 @@ impl ModelsRepository for FakeModelsRepository { &self, identifier: &str, ) -> anyhow::Result> { + if identifier == "unlisted-model" { + let mut model = canonical_model(); + model.model_name = identifier.to_string(); + model.aliases = Vec::new(); + model.is_listed = false; + return Ok(Some(model)); + } Ok((identifier == "alias-model" || identifier == "canonical-model").then(canonical_model)) } diff --git a/crates/services/src/attestation/ita/service_tests.rs b/crates/services/src/attestation/ita/service_tests.rs index 299a6395c..cd41ed406 100644 --- a/crates/services/src/attestation/ita/service_tests.rs +++ b/crates/services/src/attestation/ita/service_tests.rs @@ -125,6 +125,47 @@ async fn model_flow_uses_canonical_model_and_returns_alias_metadata() -> TestRes Ok(()) } +#[tokio::test] +async fn an_unlisted_model_is_unknown_to_the_public_token_flow() -> TestResult { + // Given: the catalog holds an unlisted model, and the token is served to + // the public, so no organization is known. + let nonce_step = || { + FakeStep::json( + 200, + r#"{"val":"dmVyaWZpZXItdmFsdWU=","iat":"aWF0LWJ5dGVz","signature":"dmVyaWZpZXItc2lnbmF0dXJl"}"#, + ) + }; + let mut errors = Vec::new(); + for model in ["unlisted-model", "no-such-model"] { + let server = FakeIta::start(vec![nonce_step()]).await?; + let model_collector = RecordingModelCollector::compatible(); + let service = service_for_fake_ita( + &server, + 0, + RecordingGatewayQuoteCollector::default(), + Some(model_collector.clone()), + )?; + + // When: the request names it. + let error = service + .get_ita_attestation_token(query(Some(model))) + .await + .expect_err("the model must not resolve"); + + // Then: no provider evidence is collected for it. + assert!( + matches!(&error, AttestationError::UnknownModel(name) if name == model), + "unexpected error for {model}: {error:?}" + ); + assert!(model_collector.calls().is_empty()); + errors.push(error.to_string().replace(model, "")); + } + + // And: it is reported exactly like a model that does not exist. + assert_eq!(errors[0], errors[1]); + Ok(()) +} + #[tokio::test] async fn model_flow_fails_closed_on_incompatible_provider_evidence() -> TestResult { // Given: a model provider returns evidence without ITA-compatible GPU nonce binding. diff --git a/crates/services/src/attestation/lifecycle.rs b/crates/services/src/attestation/lifecycle.rs index f5d0b1cb3..8db6cb823 100644 --- a/crates/services/src/attestation/lifecycle.rs +++ b/crates/services/src/attestation/lifecycle.rs @@ -6,7 +6,8 @@ use super::{ environment::{load_tls_cert_fingerprint, load_vpc_info, load_vpc_shared_secret}, ita::{ita_client_error_class, ItaClient, ProviderPoolModelAttestationCollector}, AttestationError, AttestationRepository, AttestationService, DstackGatewayQuoteCollector, - InferenceProviderPool, MetricsServiceTrait, ModelsRepository, UsageRepository, + InferenceProviderPool, MetricsServiceTrait, ModelResolveCache, ModelsRepository, + UsageRepository, }; /// Default TTL (seconds) for the no-nonce attestation-report cache. Reports for @@ -20,6 +21,7 @@ impl AttestationService { repository: Arc, inference_provider_pool: Arc, models_repository: Arc, + model_resolve_cache: ModelResolveCache, metrics_service: Arc, usage_repository: Arc, ita_config: ItaAttestationConfig, @@ -128,6 +130,7 @@ impl AttestationService { repository, inference_provider_pool, models_repository, + model_resolve_cache, metrics_service, usage_repository, vpc_info, diff --git a/crates/services/src/attestation/mod.rs b/crates/services/src/attestation/mod.rs index d4a816926..a8826db52 100644 --- a/crates/services/src/attestation/mod.rs +++ b/crates/services/src/attestation/mod.rs @@ -57,7 +57,7 @@ use crate::{ attestation::{ita::ItaClient, models::VpcInfo, ports::AttestationRepository}, inference_provider_pool::InferenceProviderPool, metrics::MetricsServiceTrait, - models::ModelsRepository, + models::{ModelResolveCache, ModelsRepository}, usage::UsageRepository, }; @@ -65,6 +65,10 @@ pub struct AttestationService { pub repository: Arc, pub inference_provider_pool: Arc, pub models_repository: Arc, + /// The cache owned by `ModelsServiceImpl`, so a model is resolved here + /// with the access decisions (and the admin invalidation) the rest of + /// the request path uses. + pub model_resolve_cache: ModelResolveCache, pub metrics_service: Arc, pub usage_repository: Arc, pub vpc_info: Option, diff --git a/crates/services/src/attestation/ports.rs b/crates/services/src/attestation/ports.rs index 64eb48164..760afa49f 100644 --- a/crates/services/src/attestation/ports.rs +++ b/crates/services/src/attestation/ports.rs @@ -2,6 +2,7 @@ use crate::attestation::ita::{ItaTokenQuery, ItaTokenResponse}; use crate::attestation::models::{ AttestationError, AttestationReport, ChatSignature, SignatureLookupResult, }; +use crate::models::ModelCaller; use async_trait::async_trait; use inference_providers::ProviderFilter; @@ -81,6 +82,10 @@ pub trait AttestationServiceTrait: Send + Sync { /// /// `provider_filter`: when `Some`, only the matching provider (NEAR fleet or 3P source) is queried. /// `None` keeps the existing behaviour (first successful provider wins). + /// + /// `caller`: who asks. A `model` the caller may not use is reported as + /// `UnknownModel`, like a name that does not exist. + #[allow(clippy::too_many_arguments)] async fn get_attestation_report( &self, model: Option, @@ -89,8 +94,12 @@ pub trait AttestationServiceTrait: Send + Sync { signing_address: Option, include_tls_fingerprint: bool, provider_filter: Option, + caller: ModelCaller, ) -> Result; + /// Intel Trust Authority tokens for the gateway and, when the query names + /// one, a model. Served to the public: the model is resolved for an + /// anonymous caller, so an unlisted model is reported as `UnknownModel`. async fn get_ita_attestation_token( &self, query: ItaTokenQuery, diff --git a/crates/services/src/attestation/report.rs b/crates/services/src/attestation/report.rs index b208cf748..6d3877e81 100644 --- a/crates/services/src/attestation/report.rs +++ b/crates/services/src/attestation/report.rs @@ -6,6 +6,7 @@ use super::{models::AttestationReport, AttestationError, AttestationService, Gat use crate::metrics::consts::{ get_environment, METRIC_ATTESTATION_REPORT_CACHE, TAG_ENVIRONMENT, TAG_RESULT, }; +use crate::models::{resolve_model_cached, ModelCaller}; use inference_providers::ProviderFilter; pub(in crate::attestation) fn generate_nonce_hex() -> String { @@ -87,6 +88,7 @@ fn normalize_signing_algo(signing_algo: Option<&str>) -> Result, } impl AttestationService { + #[allow(clippy::too_many_arguments)] pub(in crate::attestation) async fn get_attestation_report_impl( &self, model: Option, @@ -95,10 +97,38 @@ impl AttestationService { signing_address: Option, include_tls_fingerprint: bool, provider_filter: Option, + caller: ModelCaller, ) -> Result { let env_tag = format!("{TAG_ENVIRONMENT}:{}", get_environment()); let signing_algo = normalize_signing_algo(signing_algo.as_deref())?; + // Resolve the model for the caller before the cache lookup below. A + // cache hit skips the build, so a check made inside the build would + // not run when another caller had already cached the report. A model + // the caller may not use fails here exactly like an unknown one. + let resolved_canonical = if let Some(ref m) = model { + let resolved_model = resolve_model_cached( + &self.model_resolve_cache, + self.models_repository.as_ref(), + m, + caller, + ) + .await + .map_err(|e| AttestationError::ProviderError(format!("Failed to resolve model: {e}")))? + .ok_or_else(|| AttestationError::UnknownModel(m.clone()))?; + let canonical = resolved_model.model_name; + if &canonical != m { + tracing::debug!( + requested_model = %m, + canonical_model = %canonical, + "Resolved alias to canonical model name for attestation report" + ); + } + Some(canonical) + } else { + None + }; + // Precompute the no-nonce cache key BEFORE the params are moved into the // build closure below. let cache_key = format!( @@ -145,30 +175,6 @@ impl AttestationService { .unwrap_or(DEFAULT_GATEWAY_SIGNING_ALGO) .to_string(); - // Resolve model alias synchronously (fast DB lookup ~10 ms) before - // spawning the parallel futures below. - let resolved_canonical = if let Some(ref m) = model { - let resolved_model = self - .models_repository - .resolve_and_get_model(m) - .await - .map_err(|e| { - AttestationError::ProviderError(format!("Failed to resolve model: {e}")) - })? - .ok_or_else(|| AttestationError::UnknownModel(m.clone()))?; - let canonical = resolved_model.model_name.clone(); - if &canonical != m { - tracing::debug!( - requested_model = %m, - canonical_model = %canonical, - "Resolved alias to canonical model name for attestation report" - ); - } - Some(canonical) - } else { - None - }; - let signing_address_to_use = self.get_signing_address_hex(&algo)?; let signing_address_clean = signing_address_to_use .strip_prefix("0x") diff --git a/crates/services/src/attestation/report_tests.rs b/crates/services/src/attestation/report_tests.rs index 9fbbd164e..1e0399aba 100644 --- a/crates/services/src/attestation/report_tests.rs +++ b/crates/services/src/attestation/report_tests.rs @@ -19,7 +19,7 @@ use super::{ }; use crate::{ inference_provider_pool::InferenceProviderPool, - models::{ModelWithPricing, ModelsRepository}, + models::{ModelCaller, ModelWithPricing, ModelsRepository}, }; const MODEL: &str = "test-org/attested-model"; @@ -117,6 +117,52 @@ fn catalog_model() -> ModelWithPricing { } } +/// Catalog holding `MODEL` as an unlisted model that one organization holds a +/// grant for. +struct RestrictedModelRepository { + granted_organization: Uuid, +} + +#[async_trait] +impl ModelsRepository for RestrictedModelRepository { + async fn get_all_active_models(&self) -> anyhow::Result> { + Ok(vec![unlisted_catalog_model()]) + } + + async fn get_model_by_name( + &self, + model_name: &str, + ) -> anyhow::Result> { + Ok((model_name == MODEL).then(unlisted_catalog_model)) + } + + async fn resolve_and_get_model( + &self, + identifier: &str, + ) -> anyhow::Result> { + Ok((identifier == MODEL).then(unlisted_catalog_model)) + } + + async fn get_configured_model_names(&self) -> anyhow::Result> { + Ok(vec![MODEL.to_string()]) + } + + async fn organization_has_model_access( + &self, + organization_id: Uuid, + _model_id: Uuid, + ) -> anyhow::Result { + Ok(organization_id == self.granted_organization) + } +} + +fn unlisted_catalog_model() -> ModelWithPricing { + ModelWithPricing { + is_listed: false, + ..catalog_model() + } +} + /// Service whose pool serves `MODEL` from a recording `MockProvider`. Returns /// the number of attestation calls pool registration already made, so tests /// can look only at the calls their own requests caused. @@ -153,6 +199,7 @@ async fn mixed_case_signing_algo_is_forwarded_lowercased() { None, false, None, + ModelCaller::Anonymous, ) .await .unwrap_or_else(|e| panic!("{requested} must be accepted: {e}")); @@ -178,6 +225,7 @@ async fn omitted_signing_algo_leaves_the_backend_default() { None, false, None, + ModelCaller::Anonymous, ) .await .expect("omitted signing_algo must succeed"); @@ -199,6 +247,7 @@ async fn unknown_model_is_a_client_error_and_skips_the_backend() { None, false, None, + ModelCaller::Anonymous, ) .await { @@ -230,6 +279,7 @@ async fn unsupported_signing_algo_is_rejected_before_the_backend() { None, false, None, + ModelCaller::Anonymous, ) .await; @@ -254,6 +304,7 @@ async fn no_nonce_cache_is_keyed_on_the_normalized_algo() { None, false, None, + ModelCaller::Anonymous, ) }; @@ -305,7 +356,15 @@ async fn no_nonce_cache_does_not_serve_a_source_disabled_after_caching() { .build(), ); let fetch = |filter: Option| { - service.get_attestation_report(Some(MODEL.to_string()), None, None, None, false, filter) + service.get_attestation_report( + Some(MODEL.to_string()), + None, + None, + None, + false, + filter, + ModelCaller::Anonymous, + ) }; let chutes_filter = Some(ProviderFilter::Source(ProviderSource::Chutes)); @@ -328,3 +387,82 @@ async fn no_nonce_cache_does_not_serve_a_source_disabled_after_caching() { "filtered request must not be served a cached report from a disabled source" ); } + +/// An unlisted model has a report only for a caller that may use it. The +/// check must not depend on the no-nonce cache: a report cached for an +/// organization that holds a grant must not be served to anyone else. +#[tokio::test] +async fn unlisted_model_report_is_refused_like_an_unknown_model_even_when_cached() { + let (mut service, provider, baseline) = report_service().await; + let granted = Uuid::new_v4(); + service.models_repository = Arc::new(RestrictedModelRepository { + granted_organization: granted, + }); + service.report_cache = Some( + moka::future::Cache::builder() + .max_capacity(16) + .time_to_live(Duration::from_secs(60)) + .build(), + ); + let fetch = |model: &str, nonce: Option<&str>, caller: ModelCaller| { + service.get_attestation_report( + Some(model.to_string()), + None, + nonce.map(str::to_string), + None, + false, + None, + caller, + ) + }; + + // The granted organization gets the report, and the nonce-less one is + // now cached. + let granted_caller = ModelCaller::Organization(granted); + fetch(MODEL, None, granted_caller) + .await + .expect("no-nonce report for the granted organization"); + fetch(MODEL, Some(NONCE), granted_caller) + .await + .expect("nonce report for the granted organization"); + let backend_calls = backend_algos_since(&provider, baseline).len(); + assert_eq!(backend_calls, 2); + + for caller in [ + ModelCaller::Organization(Uuid::new_v4()), + ModelCaller::Anonymous, + ] { + for nonce in [None, Some(NONCE)] { + let refused = match fetch(MODEL, nonce, caller).await { + Ok(_) => panic!("{caller:?} must not get a report for the unlisted model"), + Err(error) => error, + }; + let unknown = match fetch("test-org/no-such-model", nonce, caller).await { + Ok(_) => panic!("unknown model must be rejected"), + Err(error) => error, + }; + assert!(matches!(refused, AttestationError::UnknownModel(_))); + assert!(matches!(unknown, AttestationError::UnknownModel(_))); + assert_eq!( + refused.to_string().replace(MODEL, ""), + unknown + .to_string() + .replace("test-org/no-such-model", "") + ); + } + } + // The refusals never reached the backend. + assert_eq!( + backend_algos_since(&provider, baseline).len(), + backend_calls + ); + + // The cached report still serves the organization that may read it. + fetch(MODEL, None, granted_caller) + .await + .expect("cached report for the granted organization"); + assert_eq!( + backend_algos_since(&provider, baseline).len(), + backend_calls + ); +} diff --git a/crates/services/src/attestation/service_trait.rs b/crates/services/src/attestation/service_trait.rs index a4fa6b1c3..35ef1891c 100644 --- a/crates/services/src/attestation/service_trait.rs +++ b/crates/services/src/attestation/service_trait.rs @@ -6,6 +6,7 @@ use super::{ models::AttestationReport, ports, AttestationError, AttestationService, SignatureLookupResult, }; +use crate::models::ModelCaller; use inference_providers::ProviderFilter; #[async_trait] @@ -80,6 +81,7 @@ impl ports::AttestationServiceTrait for AttestationService { signing_address: Option, include_tls_fingerprint: bool, provider_filter: Option, + caller: ModelCaller, ) -> Result { self.get_attestation_report_impl( model, @@ -88,6 +90,7 @@ impl ports::AttestationServiceTrait for AttestationService { signing_address, include_tls_fingerprint, provider_filter, + caller, ) .await } diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index be80b9211..87e9cb66d 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -790,8 +790,9 @@ pub struct CompletionServiceImpl { concurrent_limit: u32, /// Cache for per-organization concurrent limits (5-minute TTL) org_concurrent_limits: Cache, - /// Positive-only cache of model resolution (30 s TTL), shared with - /// `ModelsServiceImpl` so admin invalidation clears it. + /// Cache of model resolution and of access decisions on unlisted models + /// (30 s TTL), shared with `ModelsServiceImpl` so admin invalidation + /// clears it. model_resolve_cache: crate::models::ModelResolveCache, /// Repository for fetching organization concurrent limits organization_limit_repository: Arc, @@ -934,14 +935,18 @@ impl CompletionServiceImpl { self } + /// Resolve a model for `caller`. `None` covers both an unknown name and an + /// unlisted model the caller may not use, so both get the same error. async fn resolve_model_cached( &self, identifier: &str, + caller: crate::models::ModelCaller, ) -> Result, anyhow::Error> { crate::models::resolve_model_cached( &self.model_resolve_cache, self.models_repository.as_ref(), identifier, + caller, ) .await } @@ -1845,8 +1850,15 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { chat_params.placement = Default::default(); // Resolve model name (could be an alias) and get model details in a single DB call - // This also validates that the model exists and is active - let model = match self.resolve_model_cached(&request.model).await { + // This also validates that the model exists, is active, and may be + // used by the request's organization + let model = match self + .resolve_model_cached( + &request.model, + crate::models::ModelCaller::Organization(request.organization_id), + ) + .await + { Ok(Some(m)) => m, Ok(None) => { let err = ports::CompletionError::InvalidModel(format!( @@ -1982,10 +1994,23 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { Ok(event_stream) } - async fn create_chat_completion( + async fn run_chat_completion( &self, - request: ports::CompletionRequest, + mut request: ports::CompletionRequest, + platform_model: Option, ) -> Result { + // Who the model is resolved for. A platform model replaces the + // request's model and is resolved as a trusted caller: the platform + // chose it, so it must not depend on the grants of the organization + // the work is done for. Its name comes from `PlatformModel`, never + // from the request. + let model_caller = match platform_model { + Some(platform_model) => { + request.model = platform_model.name().to_string(); + crate::models::ModelCaller::Trusted + } + None => crate::models::ModelCaller::Organization(request.organization_id), + }; let service_start_time = Instant::now(); let organization_id = request.organization_id; let workspace_id = request.workspace_id; @@ -2042,8 +2067,12 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { chat_params.placement = Default::default(); // Resolve model name (could be an alias) and get model details in a single DB call - // This also validates that the model exists and is active - let model = match self.resolve_model_cached(&request.model).await { + // This also validates that the model exists, is active, and may be + // used by whoever chose it (see `model_caller` above) + let model = match self + .resolve_model_cached(&request.model, model_caller) + .await + { Ok(Some(m)) => m, Ok(None) => { let err = ports::CompletionError::InvalidModel(format!( @@ -2560,8 +2589,9 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { async fn get_model( &self, model_name: &str, + caller: crate::models::ModelCaller, ) -> Result, anyhow::Error> { - self.resolve_model_cached(model_name).await + self.resolve_model_cached(model_name, caller).await } fn get_inference_provider_pool( diff --git a/crates/services/src/completions/model_resolve_cache_tests.rs b/crates/services/src/completions/model_resolve_cache_tests.rs index 88e7fc097..2023a3248 100644 --- a/crates/services/src/completions/model_resolve_cache_tests.rs +++ b/crates/services/src/completions/model_resolve_cache_tests.rs @@ -1,33 +1,35 @@ -use super::provider_attribution_tests::{test_model, StaticOrganizationLimitRepository}; +use super::provider_attribution_tests::{ + completion_request, test_model, StaticOrganizationLimitRepository, +}; use super::*; use crate::inference_provider_pool::InferenceProviderPool; use crate::metrics::capturing::CapturingMetricsService; use crate::models::{ - model_resolve_cache_with_ttl, ModelWithPricing, ModelsServiceImpl, ModelsServiceTrait, + model_resolve_cache_with_ttl, ModelCaller, ModelWithPricing, ModelsServiceImpl, + ModelsServiceTrait, PlatformModel, }; use crate::test_utils::{CapturingUsageService, MockAttestationService}; use config::ExternalProvidersConfig; +use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; /// Repository whose resolved model can be swapped (or removed) between calls -/// and which counts `resolve_and_get_model` reads. +/// and which counts `resolve_and_get_model` reads. It also holds the grants +/// on that model and counts `organization_has_model_access` reads. struct CountingModelsRepository { model: Mutex>, resolves: AtomicUsize, delay: Duration, fail: std::sync::atomic::AtomicBool, + granted_organizations: Mutex>, + access_lookups: AtomicUsize, } impl CountingModelsRepository { fn new(model: Option) -> Arc { - Arc::new(Self { - model: Mutex::new(model), - resolves: AtomicUsize::new(0), - delay: Duration::ZERO, - fail: std::sync::atomic::AtomicBool::new(false), - }) + Self::slow(model, Duration::ZERO) } fn slow(model: Option, delay: Duration) -> Arc { @@ -36,12 +38,39 @@ impl CountingModelsRepository { resolves: AtomicUsize::new(0), delay, fail: std::sync::atomic::AtomicBool::new(false), + granted_organizations: Mutex::new(HashSet::new()), + access_lookups: AtomicUsize::new(0), }) } + /// A repository holding one unlisted model named `m`. + fn unlisted() -> Arc { + let mut model = test_model("m"); + model.is_listed = false; + Self::new(Some(model)) + } + fn resolves(&self) -> usize { self.resolves.load(Ordering::SeqCst) } + + fn access_lookups(&self) -> usize { + self.access_lookups.load(Ordering::SeqCst) + } + + fn grant(&self, organization_id: Uuid) { + self.granted_organizations + .lock() + .unwrap() + .insert(organization_id); + } + + fn revoke(&self, organization_id: Uuid) { + self.granted_organizations + .lock() + .unwrap() + .remove(&organization_id); + } } #[async_trait::async_trait] @@ -71,6 +100,19 @@ impl ModelsRepository for CountingModelsRepository { async fn get_configured_model_names(&self) -> Result, anyhow::Error> { Ok(vec![]) } + + async fn organization_has_model_access( + &self, + organization_id: Uuid, + _model_id: Uuid, + ) -> Result { + self.access_lookups.fetch_add(1, Ordering::SeqCst); + Ok(self + .granted_organizations + .lock() + .unwrap() + .contains(&organization_id)) + } } fn pool() -> Arc { @@ -96,12 +138,20 @@ async fn hit_does_not_call_repository_and_miss_does() { let repo = CountingModelsRepository::new(Some(test_model("m"))); let service = service(repo.clone()); - assert!(service.resolve_model_cached("m").await.unwrap().is_some()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_some()); assert_eq!(repo.resolves(), 1); - assert!(service.resolve_model_cached("m").await.unwrap().is_some()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_some()); assert_eq!(repo.resolves(), 1); assert!(service - .resolve_model_cached("other") + .resolve_model_cached("other", ModelCaller::Trusted) .await .unwrap() .is_some()); @@ -113,13 +163,25 @@ async fn none_is_never_cached() { let repo = CountingModelsRepository::new(None); let service = service(repo.clone()); - assert!(service.resolve_model_cached("m").await.unwrap().is_none()); - assert!(service.resolve_model_cached("m").await.unwrap().is_none()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_none()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_none()); assert_eq!(repo.resolves(), 2); // A newly activated model is visible immediately. *repo.model.lock().unwrap() = Some(test_model("m")); - assert!(service.resolve_model_cached("m").await.unwrap().is_some()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_some()); } #[tokio::test] @@ -129,14 +191,24 @@ async fn models_service_invalidation_clears_the_shared_cache() { let service = service(repo.clone()).with_model_resolve_cache(models_service.model_resolve_cache()); - service.resolve_model_cached("m").await.unwrap(); - service.resolve_model_cached("m").await.unwrap(); + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap(); + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap(); assert_eq!(repo.resolves(), 1); // Admin deactivation: invalidate, then the next request must re-read. *repo.model.lock().unwrap() = None; models_service.invalidate_models_cache().await; - assert!(service.resolve_model_cached("m").await.unwrap().is_none()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_none()); assert_eq!(repo.resolves(), 2); } @@ -146,9 +218,15 @@ async fn entries_expire_after_ttl() { let service = service(repo.clone()) .with_model_resolve_cache(model_resolve_cache_with_ttl(Duration::from_millis(50))); - service.resolve_model_cached("m").await.unwrap(); + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(150)).await; - service.resolve_model_cached("m").await.unwrap(); + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap(); assert_eq!(repo.resolves(), 2); } @@ -160,7 +238,11 @@ async fn concurrent_misses_for_one_key_read_the_repository_once() { let tasks: Vec<_> = (0..16) .map(|_| { let service = service.clone(); - tokio::spawn(async move { service.resolve_model_cached("m").await }) + tokio::spawn(async move { + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + }) }) .collect(); for task in tasks { @@ -177,7 +259,11 @@ async fn concurrent_none_misses_are_coalesced_but_not_cached() { let tasks: Vec<_> = (0..8) .map(|_| { let service = service.clone(); - tokio::spawn(async move { service.resolve_model_cached("m").await }) + tokio::spawn(async move { + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + }) }) .collect(); for task in tasks { @@ -187,7 +273,11 @@ async fn concurrent_none_misses_are_coalesced_but_not_cached() { // Not cached: the next request re-reads and sees a newly activated model. *repo.model.lock().unwrap() = Some(test_model("m")); - assert!(service.resolve_model_cached("m").await.unwrap().is_some()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_some()); assert_eq!(repo.resolves(), 2); } @@ -200,7 +290,11 @@ async fn repository_errors_propagate_to_coalesced_callers_and_are_not_cached() { let tasks: Vec<_> = (0..4) .map(|_| { let service = service.clone(); - tokio::spawn(async move { service.resolve_model_cached("m").await }) + tokio::spawn(async move { + service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + }) }) .collect(); for task in tasks { @@ -209,5 +303,284 @@ async fn repository_errors_propagate_to_coalesced_callers_and_are_not_cached() { assert_eq!(repo.resolves(), 1); repo.fail.store(false, Ordering::SeqCst); - assert!(service.resolve_model_cached("m").await.unwrap().is_some()); + assert!(service + .resolve_model_cached("m", ModelCaller::Trusted) + .await + .unwrap() + .is_some()); +} + +#[tokio::test] +async fn a_listed_model_resolves_for_any_organization_without_a_grant_lookup() { + let repo = CountingModelsRepository::new(Some(test_model("m"))); + let service = service(repo.clone()); + + for _ in 0..3 { + let caller = ModelCaller::Organization(Uuid::new_v4()); + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_some()); + } + assert_eq!(repo.resolves(), 1); + assert_eq!(repo.access_lookups(), 0); +} + +#[tokio::test] +async fn an_unlisted_model_resolves_only_for_an_organization_that_holds_a_grant() { + let repo = CountingModelsRepository::unlisted(); + let service = service(repo.clone()); + let granted = Uuid::new_v4(); + repo.grant(granted); + + assert!(service + .resolve_model_cached("m", ModelCaller::Organization(granted)) + .await + .unwrap() + .is_some()); + for caller in [ + ModelCaller::Organization(Uuid::new_v4()), + ModelCaller::Anonymous, + ] { + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_none()); + // The lookup the Responses API makes before dispatch sees the same. + assert!(service.get_model("m", caller).await.unwrap().is_none()); + } + // One read of the model serves every caller; only organizations reach + // the grants. + assert_eq!(repo.resolves(), 1); + assert_eq!(repo.access_lookups(), 2); +} + +#[tokio::test] +async fn models_service_invalidation_clears_the_shared_access_decisions() { + let repo = CountingModelsRepository::unlisted(); + let models_service = ModelsServiceImpl::new(pool(), repo.clone()); + let service = + service(repo.clone()).with_model_resolve_cache(models_service.model_resolve_cache()); + let organization = Uuid::new_v4(); + let caller = ModelCaller::Organization(organization); + + // A refusal is cached. + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_none()); + repo.grant(organization); + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_none()); + assert_eq!(repo.access_lookups(), 1); + + // Granting through the admin API invalidates: the grant applies at once. + models_service.invalidate_models_cache().await; + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_some()); + + // So does revoking. + repo.revoke(organization); + models_service.invalidate_models_cache().await; + assert!(service + .resolve_model_cached("m", caller) + .await + .unwrap() + .is_none()); + assert_eq!(repo.access_lookups(), 3); +} + +/// The completion service must answer a model the organization may not use +/// with the error it gives for a name that does not exist. +#[tokio::test] +async fn a_refused_model_gets_the_completion_error_of_an_unknown_one() { + let repo = CountingModelsRepository::unlisted(); + let service = service(repo.clone()); + + let error_for = |model: &'static str| { + let service = &service; + async move { + let Err(blocking) = service + .create_chat_completion(completion_request(model)) + .await + else { + panic!("the request must be refused"); + }; + let Err(streaming) = service + .create_chat_completion_stream(completion_request(model)) + .await + else { + panic!("the streaming request must be refused"); + }; + (blocking, streaming) + } + }; + + let (refused, refused_stream) = error_for("m").await; + // The repository holds one model, so remove it to make `m` unknown. + *repo.model.lock().unwrap() = None; + let (unknown, unknown_stream) = error_for("m").await; + + for (refused, unknown) in [(refused, unknown), (refused_stream, unknown_stream)] { + assert!(matches!(refused, CompletionError::InvalidModel(_))); + assert!(matches!(unknown, CompletionError::InvalidModel(_))); + assert_eq!(refused.to_string(), unknown.to_string()); + } +} + +#[tokio::test] +async fn a_granted_organization_is_served_an_unlisted_model() { + let repo = CountingModelsRepository::unlisted(); + let pool = pool(); + pool.register_provider( + "m".to_string(), + Arc::new(inference_providers::mock::MockProvider::new_accept_all()), + ) + .await; + let service = CompletionServiceImpl::new( + pool, + Arc::new(MockAttestationService), + Arc::new(CapturingUsageService::new()), + Arc::new(CapturingMetricsService::new()), + repo.clone(), + Arc::new(StaticOrganizationLimitRepository), + ); + + // The provider is there, yet an organization without a grant is refused. + let Err(refused) = service + .create_chat_completion(completion_request("m")) + .await + else { + panic!("an organization without a grant must be refused"); + }; + assert_eq!( + refused.to_string(), + "Invalid model: Model 'm' not found. It's not a valid model name or alias." + ); + + let request = completion_request("m"); + repo.grant(request.organization_id); + service + .create_chat_completion(request) + .await + .expect("the granted organization is served"); +} + +/// Unlisted models by name, and no grants (the trait default refuses). +struct UnlistedModelsRepository(Vec); + +impl UnlistedModelsRepository { + fn named(names: &[&str]) -> Arc { + Arc::new(Self( + names + .iter() + .map(|name| { + let mut model = test_model(name); + model.is_listed = false; + model + }) + .collect(), + )) + } + + fn find(&self, name: &str) -> Option { + self.0 + .iter() + .find(|model| model.model_name == name) + .cloned() + } +} + +#[async_trait::async_trait] +impl ModelsRepository for UnlistedModelsRepository { + async fn get_all_active_models(&self) -> Result, anyhow::Error> { + Ok(self.0.clone()) + } + + async fn get_model_by_name( + &self, + name: &str, + ) -> Result, anyhow::Error> { + Ok(self.find(name)) + } + + async fn resolve_and_get_model( + &self, + name: &str, + ) -> Result, anyhow::Error> { + Ok(self.find(name)) + } + + async fn get_configured_model_names(&self) -> Result, anyhow::Error> { + Ok(self + .0 + .iter() + .map(|model| model.model_name.clone()) + .collect()) + } +} + +/// A model the platform chose for its own work (the title model) must not +/// depend on the grants of the organization the work is done for. And a +/// model name from a request must not be able to take that path. +#[tokio::test] +async fn a_platform_model_is_resolved_as_a_trusted_caller() { + let platform_model = PlatformModel::title_generation(); + let title_model = platform_model.name().to_string(); + let request_model = "request/unlisted"; + // Two unlisted models nobody holds a grant for: the platform's, and one + // a request could name. + let repo = UnlistedModelsRepository::named(&[&title_model, request_model]); + let provider = Arc::new(inference_providers::mock::MockProvider::new_accept_all()); + let pool = pool(); + for model in [title_model.as_str(), request_model] { + pool.register_provider(model.to_string(), provider.clone()) + .await; + } + let service = CompletionServiceImpl::new( + pool, + Arc::new(MockAttestationService), + Arc::new(CapturingUsageService::new()), + Arc::new(CapturingMetricsService::new()), + repo, + Arc::new(StaticOrganizationLimitRepository), + ); + + // Named by a request, both are refused: the organization holds no grant. + for model in [title_model.as_str(), request_model] { + let Err(refused) = service + .create_chat_completion(completion_request(model)) + .await + else { + panic!("{model} must be refused as a request's model"); + }; + assert!(matches!(refused, CompletionError::InvalidModel(_))); + } + assert!(provider.last_chat_params().await.is_none()); + + // Chosen by the platform, the title model is served. + service + .create_platform_chat_completion(platform_model.clone(), completion_request(&title_model)) + .await + .expect("the platform's own model is served without a grant"); + let served = provider.last_chat_params().await.expect("a served request"); + assert_eq!(served.model, title_model); + + // The platform model replaces the model the request names: that name + // is not resolved as trusted, and the other unlisted model is not served. + service + .create_platform_chat_completion(platform_model, completion_request(request_model)) + .await + .expect("served on the platform model"); + let served = provider.last_chat_params().await.expect("a served request"); + assert_eq!(served.model, title_model); } diff --git a/crates/services/src/completions/ports.rs b/crates/services/src/completions/ports.rs index 340622d1a..d476a7e09 100644 --- a/crates/services/src/completions/ports.rs +++ b/crates/services/src/completions/ports.rs @@ -193,9 +193,38 @@ pub trait CompletionServiceTrait: Send + Sync { request: CompletionRequest, ) -> Result; + /// Create a completion on the model the request names. The model is + /// resolved for the request's organization. async fn create_chat_completion( &self, request: CompletionRequest, + ) -> Result { + self.run_chat_completion(request, None).await + } + + /// Create a completion on a model the platform chose for its own work, + /// such as a conversation title. `model` replaces `request.model` and is + /// resolved as a trusted caller, so the work does not depend on the + /// grants of the organization it is done (and billed) for. + /// + /// A model name from a request cannot take this path: a + /// [`PlatformModel`](crate::models::PlatformModel) cannot be built from a + /// string. + async fn create_platform_chat_completion( + &self, + model: crate::models::PlatformModel, + request: CompletionRequest, + ) -> Result { + self.run_chat_completion(request, Some(model)).await + } + + /// The completion behind [`Self::create_chat_completion`] and + /// [`Self::create_platform_chat_completion`]. Implement this one; call + /// one of those two. With a platform model, `request.model` is ignored. + async fn run_chat_completion( + &self, + request: CompletionRequest, + platform_model: Option, ) -> Result; /// Execute an audio transcription request with concurrent request limiting @@ -257,10 +286,13 @@ pub trait CompletionServiceTrait: Send + Sync { params: inference_providers::ScoreParams, ) -> Result; - /// Resolve an active model by canonical name or alias (for capability checks, etc.). + /// Resolve an active model by canonical name or alias, for `caller` (for + /// capability checks, etc.). `None` if the name is unknown or names an + /// unlisted model the caller may not use. async fn get_model( &self, model_name: &str, + caller: crate::models::ModelCaller, ) -> Result, anyhow::Error>; /// Get the inference provider pool for direct access diff --git a/crates/services/src/completions/provider_attribution_tests.rs b/crates/services/src/completions/provider_attribution_tests.rs index 787f97dce..df4dbdef8 100644 --- a/crates/services/src/completions/provider_attribution_tests.rs +++ b/crates/services/src/completions/provider_attribution_tests.rs @@ -83,7 +83,7 @@ pub(super) fn test_model(model_name: &str) -> ModelWithPricing { } } -fn completion_request(model: &str) -> ports::CompletionRequest { +pub(super) fn completion_request(model: &str) -> ports::CompletionRequest { ports::CompletionRequest { request_priority: 0, session_hint: None, diff --git a/crates/services/src/models/mod.rs b/crates/services/src/models/mod.rs index 31d56211f..85882b311 100644 --- a/crates/services/src/models/mod.rs +++ b/crates/services/src/models/mod.rs @@ -11,8 +11,12 @@ use std::time::Duration; use async_trait::async_trait; use moka::future::Cache; -pub use ports::{ModelInfo, ModelWithPricing, ModelsError, ModelsRepository, ModelsServiceTrait}; +pub use ports::{ + ModelCaller, ModelInfo, ModelWithPricing, ModelsError, ModelsRepository, ModelsServiceTrait, + PlatformModel, RestrictedModelAccess, +}; use tracing::warn; +use uuid::Uuid; use crate::inference_provider_pool::{BackendModelMetadata, InferenceProviderPool}; @@ -43,21 +47,59 @@ const MODELS_LIST_CACHE_KEY: &str = "all"; const MODEL_RESOLVE_CACHE_TTL_SECS: u64 = 30; const MODEL_RESOLVE_CACHE_CAPACITY: u64 = 1_000; -/// Positive-only cache of `ModelsRepository::resolve_and_get_model`, keyed by the -/// requested model string (alias or canonical name). Cloning shares the cache. -/// Owned by `ModelsServiceImpl` (which clears it on invalidation) and handed to -/// the completion service. -pub type ModelResolveCache = Cache; +/// Capacity of the grant cache: one small entry per (organization, unlisted +/// model) pair that was asked for within the TTL. +const MODEL_ACCESS_CACHE_CAPACITY: u64 = 10_000; + +/// The caches behind `resolve_model_cached`. Cloning shares them. Owned by +/// `ModelsServiceImpl` (which clears them on invalidation) and handed to the +/// other services that resolve a model for a request. +/// +/// - `models`: positive-only cache of `ModelsRepository::resolve_and_get_model`, +/// keyed by the requested model string (alias or canonical name). It holds +/// the model itself, for any caller. +/// - `access`: cache of `ModelsRepository::organization_has_model_access`, +/// keyed by (organization id, model id). It holds both answers, so an +/// organization that keeps asking for a model it may not use does not reach +/// the database every time. Only unlisted models are ever looked up here. +/// +/// Both share one TTL, so a grant or a revocation made through another +/// instance takes effect here within that time. +#[derive(Clone)] +pub struct ModelResolveCache { + models: Cache, + access: Cache<(Uuid, Uuid), bool>, +} + +impl ModelResolveCache { + /// Drop every cached model and every cached access decision. + pub fn invalidate_all(&self) { + self.models.invalidate_all(); + self.access.invalidate_all(); + } + + /// Drop the cached access decision for one organization and one model. + /// The cached models are kept: a grant does not change a model. + pub async fn invalidate_access(&self, organization_id: Uuid, model_id: Uuid) { + self.access.invalidate(&(organization_id, model_id)).await; + } +} pub fn new_model_resolve_cache() -> ModelResolveCache { model_resolve_cache_with_ttl(Duration::from_secs(MODEL_RESOLVE_CACHE_TTL_SECS)) } pub fn model_resolve_cache_with_ttl(ttl: Duration) -> ModelResolveCache { - Cache::builder() - .max_capacity(MODEL_RESOLVE_CACHE_CAPACITY) - .time_to_live(ttl) - .build() + ModelResolveCache { + models: Cache::builder() + .max_capacity(MODEL_RESOLVE_CACHE_CAPACITY) + .time_to_live(ttl) + .build(), + access: Cache::builder() + .max_capacity(MODEL_ACCESS_CACHE_CAPACITY) + .time_to_live(ttl) + .build(), + } } /// Outcome of a failed single-flight load: not-found is not cached, errors @@ -67,22 +109,32 @@ enum ResolveMiss { Failed(anyhow::Error), } -/// `ModelsRepository::resolve_and_get_model` through the positive-only cache. -/// Concurrent misses for the same identifier are coalesced into one repository -/// read (moka single-flight). Only found models are cached; `None` is surfaced -/// as an `Err` to moka so it is never stored and always re-reads, so a newly -/// activated model works immediately. +/// Resolve `identifier` (alias or canonical name) to an active model that +/// `caller` may use. This is the one place a model is resolved for a request. +/// +/// `None` means "no such model for this caller": the name is unknown, or it +/// names an unlisted model the caller holds no grant for. The two cases are +/// deliberately not told apart, so callers cannot report them differently. +/// +/// The model comes from `ModelsRepository::resolve_and_get_model` through the +/// positive-only cache. Concurrent misses for the same identifier are coalesced +/// into one repository read (moka single-flight). Only found models are cached; +/// `None` is surfaced as an `Err` to moka so it is never stored and always +/// re-reads, so a newly activated model works immediately. /// /// Invalidation: `invalidate_all` only discards entries inserted before it, so a /// load already in flight when an admin write invalidates can still insert its /// pre-write value afterwards. That staleness is bounded by the TTL; moka has no -/// built-in way to make an in-flight load respect invalidation. +/// built-in way to make an in-flight load respect invalidation. The same holds +/// for a grant lookup in flight while a grant is made or revoked. pub async fn resolve_model_cached( cache: &ModelResolveCache, repository: &dyn ModelsRepository, identifier: &str, + caller: ModelCaller, ) -> Result, anyhow::Error> { let loaded = cache + .models .try_get_with_by_ref(identifier, async { match repository.resolve_and_get_model(identifier).await { Ok(Some(model)) => Ok(model), @@ -91,12 +143,47 @@ pub async fn resolve_model_cached( } }) .await; - match loaded { - Ok(model) => Ok(Some(model)), + let model = match loaded { + Ok(model) => model, Err(miss) => match &*miss { - ResolveMiss::NotFound => Ok(None), - ResolveMiss::Failed(e) => Err(anyhow::anyhow!("{e:#}")), + ResolveMiss::NotFound => return Ok(None), + ResolveMiss::Failed(e) => return Err(anyhow::anyhow!("{e:#}")), }, + }; + if caller_can_use_model(cache, repository, caller, &model).await? { + Ok(Some(model)) + } else { + Ok(None) + } +} + +/// The access rule. A listed model is open to every caller and returns before +/// any lookup, so it costs nothing here. An unlisted model is open to a +/// trusted caller, closed to an anonymous one, and open to an organization +/// only if it holds a grant (cached, both answers). +/// +/// A failed grant lookup is an error, never an allow. +async fn caller_can_use_model( + cache: &ModelResolveCache, + repository: &dyn ModelsRepository, + caller: ModelCaller, + model: &ModelWithPricing, +) -> Result { + if model.is_listed { + return Ok(true); + } + match caller { + ModelCaller::Trusted => Ok(true), + ModelCaller::Anonymous => Ok(false), + ModelCaller::Organization(organization_id) => cache + .access + .try_get_with((organization_id, model.id), async { + repository + .organization_has_model_access(organization_id, model.id) + .await + }) + .await + .map_err(|e: Arc| anyhow::anyhow!("{e:#}")), } } @@ -139,7 +226,8 @@ pub struct ModelsServiceImpl { /// sentinel since pagination has been dropped — there is only ever one /// list to serve. models_list_cache: Cache<&'static str, Arc>>, - /// Shared with the completion service; cleared by `invalidate_models_cache`. + /// Shared with the completion and attestation services; cleared by + /// `invalidate_models_cache`. model_resolve_cache: ModelResolveCache, } @@ -160,8 +248,9 @@ impl ModelsServiceImpl { } } - /// Handle to the resolve cache for the completion service. Invalidation on - /// this service clears it, so admin writes take effect immediately here. + /// Handle to the resolve cache for the other services that resolve a + /// model for a request. Invalidation on this service clears it, so admin + /// writes take effect immediately here. pub fn model_resolve_cache(&self) -> ModelResolveCache { self.model_resolve_cache.clone() } @@ -219,50 +308,103 @@ impl ModelsServiceTrait for ModelsServiceImpl { Ok((*arc).clone()) } - async fn get_model_by_name(&self, model_name: &str) -> Result { - self.models_repository + async fn get_model_by_name( + &self, + model_name: &str, + caller: ModelCaller, + ) -> Result { + let not_found = || ModelsError::NotFound(format!("Model '{model_name}' not found")); + let model = self + .models_repository .get_model_by_name(model_name) .await .map_err(|e| ModelsError::InternalError(e.to_string()))? - .ok_or_else(|| ModelsError::NotFound(format!("Model '{model_name}' not found"))) + .ok_or_else(not_found)?; + // A model the caller may not use gets the error of an unknown name. + if self.caller_can_use_model(caller, &model).await? { + Ok(model) + } else { + Err(not_found()) + } } async fn resolve_and_get_model( &self, identifier: &str, + caller: ModelCaller, ) -> Result { resolve_model_cached( &self.model_resolve_cache, self.models_repository.as_ref(), identifier, + caller, ) .await .map_err(|e| ModelsError::InternalError(e.to_string()))? .ok_or_else(|| ModelsError::NotFound(format!("Model '{identifier}' not found"))) } + async fn caller_can_use_model( + &self, + caller: ModelCaller, + model: &ModelWithPricing, + ) -> Result { + caller_can_use_model( + &self.model_resolve_cache, + self.models_repository.as_ref(), + caller, + model, + ) + .await + .map_err(|e| ModelsError::InternalError(e.to_string())) + } + async fn resolve_public_model( &self, identifier: &str, ) -> Result { let models = self.cached_models().await?; // An unlisted model is answered like a model that does not exist. - let mut listed = models.iter().filter(|model| model.is_listed); - if let Some(model) = listed.clone().find(|model| model.model_name == identifier) { - return Ok(model.clone()); - } - listed - .find(|model| model.aliases.iter().any(|alias| alias == identifier)) + let listed = models.iter().filter(|model| model.is_listed); + let mut model = listed + .clone() + .find(|model| model.model_name == identifier) + .or_else(|| { + listed + .clone() + .find(|model| model.aliases.iter().any(|alias| alias == identifier)) + }) .cloned() - .ok_or_else(|| ModelsError::NotFound(format!("Model '{identifier}' not found"))) + .ok_or_else(|| ModelsError::NotFound(format!("Model '{identifier}' not found")))?; + // So is an unlisted successor: a listed model may name one. + if let Some(successor) = &model.successor_model_name { + if !listed.clone().any(|listed| &listed.model_name == successor) { + model.successor_model_name = None; + } + } + Ok(model) } - async fn resolve_alias_cached(&self, identifier: &str) -> Option { + async fn resolve_alias_cached(&self, identifier: &str, caller: ModelCaller) -> Option { let models = self.cached_models().await.ok()?; - models + let model = models .iter() - .find(|m| m.aliases.iter().any(|a| a == identifier)) - .map(|m| m.model_name.clone()) + .find(|m| m.aliases.iter().any(|a| a == identifier))?; + // The canonical name of a model the caller may not use stays hidden. + // A failed check hides it too, and is logged: otherwise a database + // outage on this advisory path would look like an unknown alias. + match self.caller_can_use_model(caller, model).await { + Ok(true) => Some(model.model_name.clone()), + Ok(false) => None, + Err(error) => { + warn!( + model_id = %model.id, + error = %error, + "Model access check failed during alias resolution; the alias stays unresolved" + ); + None + } + } } async fn get_configured_model_names(&self) -> Result, ModelsError> { @@ -272,10 +414,23 @@ impl ModelsServiceTrait for ModelsServiceImpl { .map_err(|e| ModelsError::InternalError(e.to_string())) } + async fn get_restricted_model_access(&self) -> Result, ModelsError> { + self.models_repository + .list_restricted_model_access() + .await + .map_err(|e| ModelsError::InternalError(e.to_string())) + } + async fn invalidate_models_cache(&self) { self.models_list_cache.invalidate_all(); self.model_resolve_cache.invalidate_all(); } + + async fn invalidate_model_access(&self, organization_id: Uuid, model_id: Uuid) { + self.model_resolve_cache + .invalidate_access(organization_id, model_id) + .await; + } } #[cfg(test)] @@ -283,12 +438,24 @@ mod tests { use super::*; use crate::inference_provider_pool::BackendModelMetadata; use config::ExternalProvidersConfig; + use std::collections::HashSet; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::Mutex; use uuid::Uuid; struct StaticModelsRepository { active_models: Vec, models_by_name: HashMap, resolved_models: HashMap, + /// (organization id, model id) pairs that hold a grant. + grants: Mutex>, + /// Number of `organization_has_model_access` reads. + access_lookups: AtomicUsize, + fail_access_lookups: AtomicBool, + /// Number of `get_all_active_models` reads. + list_reads: AtomicUsize, + /// Number of `resolve_and_get_model` reads. + resolves: AtomicUsize, } impl StaticModelsRepository { @@ -305,6 +472,11 @@ mod tests { active_models, models_by_name, resolved_models, + grants: Mutex::new(HashSet::new()), + access_lookups: AtomicUsize::new(0), + fail_access_lookups: AtomicBool::new(false), + list_reads: AtomicUsize::new(0), + resolves: AtomicUsize::new(0), } } @@ -312,11 +484,30 @@ mod tests { self.resolved_models.insert(identifier.to_string(), model); self } + + fn grant(&self, organization_id: Uuid, model_id: Uuid) { + self.grants + .lock() + .unwrap() + .insert((organization_id, model_id)); + } + + fn revoke(&self, organization_id: Uuid, model_id: Uuid) { + self.grants + .lock() + .unwrap() + .remove(&(organization_id, model_id)); + } + + fn access_lookups(&self) -> usize { + self.access_lookups.load(Ordering::SeqCst) + } } #[async_trait] impl ModelsRepository for StaticModelsRepository { async fn get_all_active_models(&self) -> Result, anyhow::Error> { + self.list_reads.fetch_add(1, Ordering::SeqCst); Ok(self.active_models.clone()) } @@ -331,12 +522,108 @@ mod tests { &self, identifier: &str, ) -> Result, anyhow::Error> { + self.resolves.fetch_add(1, Ordering::SeqCst); Ok(self.resolved_models.get(identifier).cloned()) } async fn get_configured_model_names(&self) -> Result, anyhow::Error> { Ok(self.models_by_name.keys().cloned().collect()) } + + async fn organization_has_model_access( + &self, + organization_id: Uuid, + model_id: Uuid, + ) -> Result { + self.access_lookups.fetch_add(1, Ordering::SeqCst); + if self.fail_access_lookups.load(Ordering::SeqCst) { + anyhow::bail!("db down"); + } + Ok(self + .grants + .lock() + .unwrap() + .contains(&(organization_id, model_id))) + } + + async fn list_restricted_model_access( + &self, + ) -> Result, anyhow::Error> { + let grants = self.grants.lock().unwrap(); + Ok(self + .active_models + .iter() + .filter(|model| !model.is_listed) + .map(|model| RestrictedModelAccess { + model_name: model.model_name.clone(), + organization_ids: grants + .iter() + .filter(|(_, model_id)| *model_id == model.id) + .map(|(organization_id, _)| *organization_id) + .collect(), + }) + .collect()) + } + } + + /// A repository that keeps the trait defaults for grants. + struct NoGrantsRepository(ModelWithPricing); + + #[async_trait] + impl ModelsRepository for NoGrantsRepository { + async fn get_all_active_models(&self) -> Result, anyhow::Error> { + Ok(vec![self.0.clone()]) + } + + async fn get_model_by_name( + &self, + model_name: &str, + ) -> Result, anyhow::Error> { + Ok((model_name == self.0.model_name).then(|| self.0.clone())) + } + + async fn resolve_and_get_model( + &self, + identifier: &str, + ) -> Result, anyhow::Error> { + Ok((identifier == self.0.model_name).then(|| self.0.clone())) + } + + async fn get_configured_model_names(&self) -> Result, anyhow::Error> { + Ok(vec![self.0.model_name.clone()]) + } + } + + /// A service over `repository` with no backends, for the access tests. + fn service_over(repository: Arc) -> ModelsServiceImpl { + ModelsServiceImpl::new( + Arc::new(InferenceProviderPool::new( + None, + ExternalProvidersConfig::default(), + )), + repository, + ) + } + + const LISTED: &str = "test/listed"; + const UNLISTED: &str = "test/unlisted"; + const UNLISTED_ALIAS: &str = "unlisted-alias"; + + /// One listed and one unlisted model (with an alias), and their service. + fn access_fixture() -> (Arc, ModelsServiceImpl, Uuid) { + let mut unlisted = test_catalog_model(UNLISTED); + unlisted.is_listed = false; + unlisted.aliases = vec![UNLISTED_ALIAS.to_string()]; + let unlisted_id = unlisted.id; + let repository = Arc::new( + StaticModelsRepository::with_active_models(vec![ + test_catalog_model(LISTED), + unlisted.clone(), + ]) + .with_resolved_model(UNLISTED_ALIAS, unlisted), + ); + let service = service_over(repository.clone()); + (repository, service, unlisted_id) } fn test_catalog_model(model_name: &str) -> ModelWithPricing { @@ -736,15 +1023,395 @@ mod tests { )); } - // Resolution for serving and billing does not look at the switch. - let resolved = service.resolve_and_get_model(unlisted_name).await.unwrap(); + // A trusted caller (usage recording, admin flows) still resolves it. + let resolved = service + .resolve_and_get_model(unlisted_name, ModelCaller::Trusted) + .await + .unwrap(); assert_eq!(resolved.model_name, unlisted_name); assert_eq!( - service.resolve_alias_cached("unlisted-alias").await, + service + .resolve_alias_cached("unlisted-alias", ModelCaller::Trusted) + .await, + Some(unlisted_name.to_string()) + ); + } + + #[tokio::test] + async fn the_public_views_name_only_a_listed_successor() { + let to_unlisted = "test/to-unlisted"; + let to_listed = "test/to-listed"; + let unlisted_name = "test/unlisted-successor"; + let listed_name = "test/listed-successor"; + let mut unlisted = test_catalog_model(unlisted_name); + unlisted.is_listed = false; + let mut names_unlisted = test_catalog_model(to_unlisted); + names_unlisted.aliases = vec!["to-unlisted-alias".to_string()]; + names_unlisted.successor_model_name = Some(unlisted_name.to_string()); + let mut names_listed = test_catalog_model(to_listed); + names_listed.successor_model_name = Some(listed_name.to_string()); + let service = service_with_backend_models( + StaticModelsRepository::with_active_models(vec![ + names_unlisted, + names_listed, + unlisted, + test_catalog_model(listed_name), + ]), + to_listed, + vec![provider_model(to_listed, Some(32_768), Some(4_096))], + ) + .await; + + let successor_in = |models: &[ModelWithPricing], name: &str| { + models + .iter() + .find(|model| model.model_name == name) + .unwrap() + .successor_model_name + .clone() + }; + let public = service.get_public_models_with_pricing().await.unwrap(); + assert_eq!(successor_in(&public, to_unlisted), None); + assert_eq!( + successor_in(&public, to_listed), + Some(listed_name.to_string()) + ); + for identifier in [to_unlisted, "to-unlisted-alias"] { + let model = service.resolve_public_model(identifier).await.unwrap(); + assert_eq!(model.successor_model_name, None, "{identifier}"); + } + assert_eq!( + service + .resolve_public_model(to_listed) + .await + .unwrap() + .successor_model_name, + Some(listed_name.to_string()) + ); + + // The trusted view keeps the successor. + let all = service.get_models_with_pricing().await.unwrap(); + assert_eq!( + successor_in(&all, to_unlisted), Some(unlisted_name.to_string()) ); } + #[tokio::test] + async fn a_listed_model_is_open_to_every_caller_and_never_consults_grants() { + let (repository, service, _) = access_fixture(); + + for caller in [ + ModelCaller::Anonymous, + ModelCaller::Organization(Uuid::new_v4()), + ModelCaller::Trusted, + ] { + let resolved = service.resolve_and_get_model(LISTED, caller).await.unwrap(); + assert_eq!(resolved.model_name, LISTED); + let by_name = service.get_model_by_name(LISTED, caller).await.unwrap(); + assert_eq!(by_name.model_name, LISTED); + assert!(service + .caller_can_use_model(caller, &by_name) + .await + .unwrap()); + } + assert_eq!(repository.access_lookups(), 0); + } + + #[tokio::test] + async fn an_unlisted_model_resolves_only_for_a_trusted_caller_or_a_granted_organization() { + let (repository, service, unlisted_id) = access_fixture(); + let granted = Uuid::new_v4(); + let other = Uuid::new_v4(); + repository.grant(granted, unlisted_id); + // A grant on another model opens nothing here. + repository.grant(other, Uuid::new_v4()); + + for caller in [ModelCaller::Trusted, ModelCaller::Organization(granted)] { + for identifier in [UNLISTED, UNLISTED_ALIAS] { + let resolved = service + .resolve_and_get_model(identifier, caller) + .await + .unwrap(); + assert_eq!(resolved.model_name, UNLISTED, "{caller:?} {identifier}"); + } + let by_name = service.get_model_by_name(UNLISTED, caller).await.unwrap(); + assert!(service + .caller_can_use_model(caller, &by_name) + .await + .unwrap()); + assert_eq!( + service.resolve_alias_cached(UNLISTED_ALIAS, caller).await, + Some(UNLISTED.to_string()) + ); + } + + let unlisted = service + .get_model_by_name(UNLISTED, ModelCaller::Trusted) + .await + .unwrap(); + for caller in [ModelCaller::Anonymous, ModelCaller::Organization(other)] { + for identifier in [UNLISTED, UNLISTED_ALIAS] { + assert!( + matches!( + service.resolve_and_get_model(identifier, caller).await, + Err(ModelsError::NotFound(_)) + ), + "{caller:?} {identifier}" + ); + } + assert!(matches!( + service.get_model_by_name(UNLISTED, caller).await, + Err(ModelsError::NotFound(_)) + )); + assert!(!service + .caller_can_use_model(caller, &unlisted) + .await + .unwrap()); + // The alias does not disclose the canonical name either. + assert_eq!( + service.resolve_alias_cached(UNLISTED_ALIAS, caller).await, + None + ); + } + } + + #[tokio::test] + async fn a_refused_model_gets_the_error_of_an_unknown_one() { + let (_, service, _) = access_fixture(); + let caller = ModelCaller::Organization(Uuid::new_v4()); + let unknown = "test/no-such-model"; + + let refused = service + .resolve_and_get_model(UNLISTED, caller) + .await + .unwrap_err(); + let missing = service + .resolve_and_get_model(unknown, caller) + .await + .unwrap_err(); + assert!(matches!(refused, ModelsError::NotFound(_))); + assert!(matches!(missing, ModelsError::NotFound(_))); + assert_eq!( + refused.to_string().replace(UNLISTED, unknown), + missing.to_string() + ); + + let refused = service + .get_model_by_name(UNLISTED, caller) + .await + .unwrap_err(); + let missing = service + .get_model_by_name(unknown, caller) + .await + .unwrap_err(); + assert!(matches!(refused, ModelsError::NotFound(_))); + assert!(matches!(missing, ModelsError::NotFound(_))); + assert_eq!( + refused.to_string().replace(UNLISTED, unknown), + missing.to_string() + ); + } + + #[tokio::test] + async fn access_decisions_are_cached_both_ways_and_cleared_by_invalidation() { + let (repository, service, unlisted_id) = access_fixture(); + let organization = Uuid::new_v4(); + let caller = ModelCaller::Organization(organization); + let resolves = || async { + service + .resolve_and_get_model(UNLISTED, caller) + .await + .is_ok() + }; + + // Refused, and the refusal is cached: one grant lookup for three asks. + for _ in 0..3 { + assert!(!resolves().await); + } + assert_eq!(repository.access_lookups(), 1); + + // A grant written behind the cache's back is not seen... + repository.grant(organization, unlisted_id); + assert!(!resolves().await); + assert_eq!(repository.access_lookups(), 1); + // ...until the cache is invalidated, as granting through the admin + // API does. + service.invalidate_models_cache().await; + for _ in 0..3 { + assert!(resolves().await); + } + assert_eq!(repository.access_lookups(), 2); + + // The same for a revocation. + repository.revoke(organization, unlisted_id); + assert!(resolves().await); + service.invalidate_models_cache().await; + assert!(!resolves().await); + assert_eq!(repository.access_lookups(), 3); + + // Decisions are per organization. + assert!(service + .resolve_and_get_model(UNLISTED, ModelCaller::Organization(Uuid::new_v4())) + .await + .is_err()); + assert_eq!(repository.access_lookups(), 4); + } + + /// A grant or a revocation changes one answer. Dropping it must not make + /// the service reload the model list (which asks every backend) or + /// resolve the model again, and must leave other organizations' answers + /// cached. + #[tokio::test] + async fn access_only_invalidation_drops_one_decision_and_keeps_the_model_caches() { + let (repository, service, unlisted_id) = access_fixture(); + let organization = Uuid::new_v4(); + let other = Uuid::new_v4(); + let resolves_for = |organization: Uuid| { + let service = &service; + async move { + service + .resolve_and_get_model(UNLISTED, ModelCaller::Organization(organization)) + .await + .is_ok() + } + }; + + // Warm every cache: the list, the resolved model, both decisions. + service.get_models_with_pricing().await.unwrap(); + assert!(!resolves_for(organization).await); + assert!(!resolves_for(other).await); + let reads = |repository: &StaticModelsRepository| { + ( + repository.list_reads.load(Ordering::SeqCst), + repository.resolves.load(Ordering::SeqCst), + repository.access_lookups(), + ) + }; + assert_eq!(reads(&repository), (1, 1, 2)); + + // Both organizations get a grant, one decision is dropped. + repository.grant(organization, unlisted_id); + repository.grant(other, unlisted_id); + service + .invalidate_model_access(organization, unlisted_id) + .await; + + // That organization's answer is read again... + assert!(resolves_for(organization).await); + // ...the other one's is still the cached refusal... + assert!(!resolves_for(other).await); + // ...and neither the list nor the model was read again. + service.get_models_with_pricing().await.unwrap(); + assert_eq!(reads(&repository), (1, 1, 3)); + + // Dropping a decision that is not cached is harmless. + service + .invalidate_model_access(Uuid::new_v4(), unlisted_id) + .await; + assert!(resolves_for(organization).await); + assert_eq!(reads(&repository), (1, 1, 3)); + } + + #[tokio::test] + async fn access_decisions_expire_with_the_resolve_cache_ttl() { + let (repository, _, unlisted_id) = access_fixture(); + let cache = model_resolve_cache_with_ttl(Duration::from_millis(50)); + let organization = Uuid::new_v4(); + let caller = ModelCaller::Organization(organization); + + assert!( + resolve_model_cached(&cache, repository.as_ref(), UNLISTED, caller) + .await + .unwrap() + .is_none() + ); + repository.grant(organization, unlisted_id); + tokio::time::sleep(Duration::from_millis(150)).await; + // No invalidation: this is how another instance learns of the grant. + assert!( + resolve_model_cached(&cache, repository.as_ref(), UNLISTED, caller) + .await + .unwrap() + .is_some() + ); + assert_eq!(repository.access_lookups(), 2); + } + + #[tokio::test] + async fn a_failed_grant_lookup_is_an_error_and_is_not_cached() { + let (repository, service, unlisted_id) = access_fixture(); + let organization = Uuid::new_v4(); + let caller = ModelCaller::Organization(organization); + repository.grant(organization, unlisted_id); + + repository.fail_access_lookups.store(true, Ordering::SeqCst); + assert!(matches!( + service.resolve_and_get_model(UNLISTED, caller).await, + Err(ModelsError::InternalError(_)) + )); + assert!(matches!( + service.get_model_by_name(UNLISTED, caller).await, + Err(ModelsError::InternalError(_)) + )); + // The advisory alias lookup hides the model rather than failing. + assert_eq!( + service.resolve_alias_cached(UNLISTED_ALIAS, caller).await, + None + ); + + repository + .fail_access_lookups + .store(false, Ordering::SeqCst); + assert!(service + .resolve_and_get_model(UNLISTED, caller) + .await + .is_ok()); + } + + #[tokio::test] + async fn a_repository_without_grants_opens_an_unlisted_model_to_trusted_callers_only() { + let mut unlisted = test_catalog_model(UNLISTED); + unlisted.is_listed = false; + let service = service_over(Arc::new(NoGrantsRepository(unlisted))); + + assert!(service + .resolve_and_get_model(UNLISTED, ModelCaller::Trusted) + .await + .is_ok()); + for caller in [ + ModelCaller::Anonymous, + ModelCaller::Organization(Uuid::new_v4()), + ] { + assert!(matches!( + service.resolve_and_get_model(UNLISTED, caller).await, + Err(ModelsError::NotFound(_)) + )); + } + // It cannot say which models are restricted, so it must not claim + // that none is. + assert!(service.get_restricted_model_access().await.is_err()); + } + + #[tokio::test] + async fn restricted_model_access_lists_unlisted_models_with_their_organizations() { + let (repository, service, unlisted_id) = access_fixture(); + + let access = service.get_restricted_model_access().await.unwrap(); + assert_eq!( + access, + vec![RestrictedModelAccess { + model_name: UNLISTED.to_string(), + organization_ids: Vec::new(), + }] + ); + + let organization = Uuid::new_v4(); + repository.grant(organization, unlisted_id); + let access = service.get_restricted_model_access().await.unwrap(); + assert_eq!(access[0].organization_ids, vec![organization]); + } + #[tokio::test] async fn get_models_with_pricing_public_resolver_exact_model_name_wins_over_alias() { let aliased_model_name = "test/aliased"; @@ -809,9 +1476,13 @@ mod tests { ) .await; - let by_name = service.get_model_by_name(model_name).await.unwrap(); - let by_alias = service.get_model_by_name("friendly").await; - let resolved_alias = service.resolve_and_get_model("friendly").await.unwrap(); + let caller = ModelCaller::Anonymous; + let by_name = service.get_model_by_name(model_name, caller).await.unwrap(); + let by_alias = service.get_model_by_name("friendly", caller).await; + let resolved_alias = service + .resolve_and_get_model("friendly", caller) + .await + .unwrap(); assert_eq!(by_name.context_length, 4096); assert_eq!(by_name.max_output_length, Some(1024)); diff --git a/crates/services/src/models/ports.rs b/crates/services/src/models/ports.rs index 8dc1f4f3d..0bcf83924 100644 --- a/crates/services/src/models/ports.rs +++ b/crates/services/src/models/ports.rs @@ -91,15 +91,83 @@ pub struct ModelWithPricing { /// OpenRouter `openrouter.slug` override. When set, the public API emits a /// nested `openrouter: { slug: }` object; NULL = unset (omitted). pub openrouter_slug: Option, - /// Whether the model appears in the public model catalogs - /// (`GET /v1/models`, `GET /v1/model/list`, `GET /v1/model/{name}`). - /// `false` keeps an active, billable row out of them; everything else - /// (resolution, routing, usage recording) treats the model as before. + /// Whether the model is listed. A listed model appears in the public + /// model catalogs (`GET /v1/models`, `GET /v1/model/list`, + /// `GET /v1/model/{name}`) and every caller may use it. An unlisted model + /// is left out of them and is restricted: only an organization that holds + /// a grant may use it (see [`ModelCaller`]). Usage recording does not + /// look at this field. pub is_listed: bool, /// When the model row was created — used as OpenRouter's `created` unix timestamp. pub created_at: chrono::DateTime, } +/// Who is asking for a model. Every lookup of a model for a request states +/// its caller, and the answer depends on it: +/// +/// - a listed model resolves for every caller; +/// - an unlisted model resolves only for [`ModelCaller::Trusted`] and for an +/// [`ModelCaller::Organization`] that holds a grant for it. +/// +/// A model that does not resolve for the caller is reported exactly like a +/// model that does not exist, so the caller cannot learn that it is there. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ModelCaller { + /// No organization is known, as on an unauthenticated route. Listed + /// models only. + Anonymous, + /// A request made on behalf of this organization. Listed models, plus the + /// unlisted models the organization holds a grant for. + Organization(Uuid), + /// The platform itself: admin flows, usage recording, background jobs and + /// work on a model the platform chose (see [`PlatformModel`]). Every + /// model. Never for a model name taken from a request: a response can be + /// sent before its model is resolved, so "the request was served" proves + /// nothing about the caller's access. + Trusted, +} + +/// A model the platform chose for its own work, such as writing the title of +/// a conversation. It is not a name from a request. +/// +/// A completion that runs a `PlatformModel` resolves it as +/// [`ModelCaller::Trusted`], so the work does not depend on the grants of the +/// organization it is done for. That is safe only because a `PlatformModel` +/// cannot be built from a string: the constructors below are the whole list, +/// and each reads the platform's own configuration. Do not add one that takes +/// a name from a caller. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlatformModel(String); + +impl PlatformModel { + /// Model used when `TITLE_GENERATION_MODEL` is not set. + const DEFAULT_TITLE_GENERATION_MODEL: &'static str = "Qwen/Qwen3-30B-A3B-Instruct-2507"; + + /// The model that writes conversation titles: `TITLE_GENERATION_MODEL`, + /// or the default. + pub fn title_generation() -> Self { + Self( + std::env::var("TITLE_GENERATION_MODEL") + .unwrap_or_else(|_| Self::DEFAULT_TITLE_GENERATION_MODEL.to_string()), + ) + } + + /// The model's name (canonical name or alias). + pub fn name(&self) -> &str { + &self.0 + } +} + +/// One unlisted model, active or not, and the organizations that may use it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RestrictedModelAccess { + /// Canonical model name. + pub model_name: String, + /// Active organizations that hold a grant for the model. Empty when + /// nobody may use it. + pub organization_ids: Vec, +} + #[derive(Debug, thiserror::Error)] pub enum ModelsError { #[error("Internal error: {0}")] @@ -111,6 +179,11 @@ pub enum ModelsError { } /// Repository trait for accessing model data +/// +/// The lookups here return every active model, unlisted ones included, and +/// apply no access rule. Code that looks a model up for a request must go +/// through [`ModelsServiceTrait`] or `resolve_model_cached`, which take the +/// [`ModelCaller`]. #[async_trait] pub trait ModelsRepository: Send + Sync { /// Get all active models with pricing @@ -132,6 +205,34 @@ pub trait ModelsRepository: Send + Sync { /// Get list of configured model names (canonical names) from database /// Returns only active models that have been configured with pricing async fn get_configured_model_names(&self) -> Result, anyhow::Error>; + + /// Whether the organization holds a grant for the model that is in + /// effect: the organization and the model must both be active. + /// + /// The default knows of no grants and refuses, so an implementor that + /// does not override it cannot open an unlisted model to anyone. + async fn organization_has_model_access( + &self, + _organization_id: Uuid, + _model_id: Uuid, + ) -> Result { + Ok(false) + } + + /// Every unlisted model, active or not, with the active organizations + /// that hold a grant for it, ordered by model name. + /// + /// A reader takes a model that is absent from the list to be + /// unrestricted. So an inactive model stays on it (the reader may still + /// be serving it), and the default fails instead of returning an empty + /// list. + async fn list_restricted_model_access( + &self, + ) -> Result, anyhow::Error> { + Err(anyhow::anyhow!( + "this models repository does not list model access" + )) + } } #[cfg_attr(test, mockall::automock)] @@ -142,6 +243,11 @@ pub trait ModelsServiceTrait: Send + Sync { /// Get models with pricing and metadata (from database). /// + /// Every active model, unlisted ones included: this is the trusted view. + /// Before answering a request from one of these rows, check it with + /// [`Self::caller_can_use_model`]; public catalogs read + /// [`Self::get_public_models_with_pricing`] instead. + /// /// Implementations are expected to cache the result with a short TTL since /// this endpoint is publicly accessible and called frequently. async fn get_models_with_pricing(&self) -> Result, ModelsError>; @@ -149,24 +255,67 @@ pub trait ModelsServiceTrait: Send + Sync { /// The models of [`Self::get_models_with_pricing`] that are listed in the /// public catalogs (`is_listed`). Every unauthenticated catalog endpoint /// must read this one, so an unlisted model never reaches them. + /// + /// A successor is kept only if it is one of these models: a listed model + /// may name an unlisted successor, and the public catalogs must not. async fn get_public_models_with_pricing(&self) -> Result, ModelsError> { let mut models = self.get_models_with_pricing().await?; models.retain(|model| model.is_listed); + let public: std::collections::HashSet = models + .iter() + .map(|model| model.model_name.clone()) + .collect(); + for model in &mut models { + if model + .successor_model_name + .as_ref() + .is_some_and(|successor| !public.contains(successor)) + { + model.successor_model_name = None; + } + } Ok(models) } - /// Get a specific model by name - async fn get_model_by_name(&self, model_name: &str) -> Result; + /// Get a specific model by name, for `caller`. + /// + /// A model the caller may not use is `NotFound`, with the error an + /// unknown name gets. + async fn get_model_by_name( + &self, + model_name: &str, + caller: ModelCaller, + ) -> Result; - /// Resolve a model identifier (alias or canonical name) and return the full model details - /// Returns an error if the model is not found or not active + /// Resolve a model identifier (alias or canonical name) for `caller` and + /// return the full model details. + /// + /// Returns an error if the model is not found or not active. A model the + /// caller may not use is `NotFound`, with the error an unknown name gets. async fn resolve_and_get_model( &self, identifier: &str, + caller: ModelCaller, ) -> Result; + /// Whether `caller` may use `model`, a row this service returned. + /// + /// A listed model costs nothing to check. Only an unlisted model asked + /// for by an organization consults the grants. + /// + /// The default knows of no grants: it opens an unlisted model to + /// [`ModelCaller::Trusted`] only. + async fn caller_can_use_model( + &self, + caller: ModelCaller, + model: &ModelWithPricing, + ) -> Result { + Ok(model.is_listed || caller == ModelCaller::Trusted) + } + /// Resolve a model identifier for public model-detail reads. An unlisted - /// model is answered like a model that does not exist. + /// model is answered like a model that does not exist, and an unlisted + /// successor is left out like in [`Self::get_public_models_with_pricing`]. /// /// Production implementations should use the cached active public model /// list, preferring exact canonical model-name matches before alias matches. @@ -176,8 +325,19 @@ pub trait ModelsServiceTrait: Send + Sync { &self, identifier: &str, ) -> Result { - let model = self.resolve_and_get_model(identifier).await?; + let mut model = self + .resolve_and_get_model(identifier, ModelCaller::Anonymous) + .await?; if model.is_listed { + if let Some(successor) = &model.successor_model_name { + if self + .resolve_and_get_model(successor, ModelCaller::Anonymous) + .await + .is_err() + { + model.successor_model_name = None; + } + } Ok(model) } else { Err(ModelsError::NotFound(format!( @@ -188,20 +348,47 @@ pub trait ModelsServiceTrait: Send + Sync { /// Resolve `identifier` against the cached active-models alias map. /// Returns `Some(canonical_name)` when `identifier` is a registered - /// alias of an active model, `None` when it is a canonical name, - /// unknown, or the lookup fails. + /// alias of an active model that `caller` may use, `None` when it is a + /// canonical name, unknown, an alias of a model the caller may not use, + /// or the lookup fails. /// /// Unlike `resolve_and_get_model`, this never hits the DB on the hot - /// path (cache-backed, short TTL) and is intended for advisory - /// annotations such as alias-substitution warnings — authoritative - /// checks must use `resolve_and_get_model`. - async fn resolve_alias_cached(&self, identifier: &str) -> Option; + /// path for a listed model (cache-backed, short TTL) and is intended for + /// advisory annotations such as alias-substitution warnings — + /// authoritative checks must use `resolve_and_get_model`. + async fn resolve_alias_cached(&self, identifier: &str, caller: ModelCaller) -> Option; /// Get list of configured model names (canonical names) from database /// Returns only active models that have been configured with pricing async fn get_configured_model_names(&self) -> Result, ModelsError>; - /// Invalidate any cached model list / metadata. Called by admin writes - /// that mutate the `models` or `model_aliases` tables. + /// Every unlisted model, active or not, with the active organizations + /// that hold a grant for it. Trusted services (a gateway that enforces + /// access itself) read this; a model that is absent from it is + /// unrestricted, which is why a model that was deactivated here stays on + /// it. + /// + /// The default fails instead of returning an empty list, which a reader + /// would take to mean that no model is restricted. + async fn get_restricted_model_access(&self) -> Result, ModelsError> { + Err(ModelsError::InternalError( + "this models service does not list model access".to_string(), + )) + } + + /// Invalidate any cached model list / metadata and cached access + /// decisions. Called by admin writes that mutate the `models` or + /// `model_aliases` tables. async fn invalidate_models_cache(&self); + + /// Forget the cached answer to "may this organization use this model", + /// and nothing else. Called when a grant is made or revoked: the model + /// list and the resolved models are not affected by a grant, and + /// reloading the list asks every backend for its metadata. + /// + /// The default clears everything, which is correct for an implementor + /// that does not keep access decisions apart. + async fn invalidate_model_access(&self, _organization_id: Uuid, _model_id: Uuid) { + self.invalidate_models_cache().await; + } } diff --git a/crates/services/src/responses/native.rs b/crates/services/src/responses/native.rs index bcb86679f..370c27139 100644 --- a/crates/services/src/responses/native.rs +++ b/crates/services/src/responses/native.rs @@ -65,14 +65,20 @@ pub fn selected(models: &[String], canonical_model: &str, body: &Value) -> bool impl NativeResponsesService { /// Resolve aliases before checking the canonical allowlist. Non-selected - /// requests retain legacy validation and error handling. - pub async fn selected_model(&self, body: &Value) -> Option { + /// requests retain legacy validation and error handling. A model the + /// caller may not use is never selected: the request takes the legacy + /// path and gets its unknown-model answer. + pub async fn selected_model( + &self, + body: &Value, + caller: crate::models::ModelCaller, + ) -> Option { if self.models.is_empty() || !inference_providers::responses_raw::is_stateless(body) { return None; } let model = self .models_service - .resolve_and_get_model(body["model"].as_str()?) + .resolve_and_get_model(body["model"].as_str()?, caller) .await .ok()?; selected(&self.models, &model.model_name, body).then_some(model) diff --git a/crates/services/src/responses/service.rs b/crates/services/src/responses/service.rs index 5f2d0924c..76e879c32 100644 --- a/crates/services/src/responses/service.rs +++ b/crates/services/src/responses/service.rs @@ -178,9 +178,15 @@ impl ports::ResponseServiceTrait for ResponseServiceImpl { // Reuse the lookup needed for image dispatch, but perform it before // starting the stream or tools so endpoint errors remain HTTP 400s. + // A model this organization may not use is `None`, like an unknown + // one: no endpoint check and no image dispatch for it, and the + // completion service answers the request as an unknown model. let model = self .completion_service - .get_model(&request.model) + .get_model( + &request.model, + crate::models::ModelCaller::Organization(organization_id), + ) .await .ok() .flatten(); @@ -3153,15 +3159,17 @@ impl ResponseServiceImpl { ); // Generate title using completion service (names not included - tracked via database) - let title_model = std::env::var("TITLE_GENERATION_MODEL") - .unwrap_or_else(|_| "Qwen/Qwen3-30B-A3B-Instruct-2507".to_string()); + // The platform chose this model, not the request: it is resolved as a + // trusted caller, so a title does not depend on the organization's + // grants. + let title_model = crate::models::PlatformModel::title_generation(); // `session_hint: None` here too: this is the internal title-generation // request, issued by the deprecated /responses API, which does not forward // `x-session-id`. See the comment on the other CompletionRequest above. let completion_request = crate::completions::ports::CompletionRequest { request_id, session_hint: None, - model: title_model, + model: title_model.name().to_string(), messages: vec![crate::completions::ports::CompletionMessage { reasoning_content: None, role: "user".to_string(), @@ -3196,7 +3204,7 @@ impl ResponseServiceImpl { // Call completion service to generate title let completion_result = completion_service - .create_chat_completion(completion_request) + .create_platform_chat_completion(title_model, completion_request) .await .map_err(|e| { errors::ResponseError::InternalError(format!("Failed to generate title: {e}")) diff --git a/crates/services/src/test_utils.rs b/crates/services/src/test_utils.rs index fc8adfceb..bedbe1396 100644 --- a/crates/services/src/test_utils.rs +++ b/crates/services/src/test_utils.rs @@ -64,6 +64,7 @@ impl AttestationServiceTrait for MockAttestationService { _signing_address: Option, _include_tls_fingerprint: bool, _provider_filter: Option, + _caller: crate::models::ModelCaller, ) -> Result { Err(AttestationError::InternalError( "Not implemented".to_string(),