diff --git a/Cargo.lock b/Cargo.lock index 2c0b23e67..403fbc8ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10689,6 +10689,7 @@ dependencies = [ "temps-ai", "temps-auth", "temps-core", + "temps-database", "temps-entities", "thiserror 2.0.18", "tokio", diff --git a/crates/temps-ai-gateway/Cargo.toml b/crates/temps-ai-gateway/Cargo.toml index 100e7ffcb..e1432a828 100644 --- a/crates/temps-ai-gateway/Cargo.toml +++ b/crates/temps-ai-gateway/Cargo.toml @@ -40,3 +40,4 @@ tower = { workspace = true } http-body-util = { workspace = true } http = { workspace = true } hyper = { workspace = true } +temps-database = { path = "../temps-database" } diff --git a/crates/temps-ai-gateway/src/error.rs b/crates/temps-ai-gateway/src/error.rs index 8b0391d7e..e97d64402 100644 --- a/crates/temps-ai-gateway/src/error.rs +++ b/crates/temps-ai-gateway/src/error.rs @@ -14,6 +14,48 @@ pub enum AiGatewayError { #[error("Model '{model}' is not allowed in scope '{scope}'")] ModelNotAllowed { model: String, scope: String }, + #[error( + "AI gateway rate limit exceeded for scope '{scope}': maximum {limit_per_minute} requests per minute; retry in {retry_after_seconds}s" + )] + RateLimitExceeded { + scope: String, + limit_per_minute: i64, + retry_after_seconds: u64, + }, + + #[error( + "AI gateway monthly budget exceeded for scope '{scope}': spent {spent_microcents} of {limit_microcents} microcents" + )] + MonthlyBudgetExceeded { + scope: String, + spent_microcents: i64, + limit_microcents: i64, + }, + + #[error("No pricing is configured for model '{model}' required by budget scope '{scope}'")] + PricingUnavailable { model: String, scope: String }, + + #[error( + "AI gateway budget scope '{scope}' requires max_tokens so spend can be reserved safely" + )] + BudgetRequiresMaxTokens { scope: String }, + + #[error("AI gateway budget scope '{scope}' cannot safely project remote image token cost")] + BudgetProjectionUnavailable { scope: String }, + + #[error("Invalid AI gateway governance config for scope '{scope}': {field}={value} must be non-negative")] + InvalidGovernanceConfig { + scope: String, + field: &'static str, + value: i64, + }, + + #[error("Invalid AI gateway governance scope '{scope}'; expected instance, project:, environment:, or token:")] + InvalidGovernanceScope { scope: String }, + + #[error("AI gateway governance config for scope '{scope}' was not found")] + GovernanceConfigNotFound { scope: String }, + #[error("Upstream provider error for model '{model}': {status} {message}")] UpstreamError { model: String, diff --git a/crates/temps-ai-gateway/src/handlers/gateway.rs b/crates/temps-ai-gateway/src/handlers/gateway.rs index dd2807054..364dece9d 100644 --- a/crates/temps-ai-gateway/src/handlers/gateway.rs +++ b/crates/temps-ai-gateway/src/handlers/gateway.rs @@ -9,8 +9,7 @@ use axum::{ use bytes::Bytes; use std::sync::Arc; use std::time::Instant; -use temps_auth::permission_guard; -use temps_auth::RequireAuth; +use temps_auth::{permission_guard, AuthContext, RequireAuth}; use temps_core::problemdetails::Problem; use tracing::{debug, error, info}; use utoipa::OpenApi; @@ -18,8 +17,8 @@ use utoipa::OpenApi; use crate::error::AiGatewayError; use crate::handlers::types::AiGatewayAppState; use crate::services::gateway_service::{ByokOverride, CredentialType}; -use crate::services::usage_service::AiRequestContext; -use crate::services::UsageService; +use crate::services::usage_service::{AiRequestContext, AiUsageAttribution}; +use crate::services::{GovernanceReservation, UsageService}; use crate::types::*; /// Extract BYOK overrides from request headers. @@ -67,6 +66,22 @@ fn extract_ai_context(headers: &HeaderMap) -> AiRequestContext { } } +fn usage_attribution(auth: &AuthContext) -> AiUsageAttribution { + match auth.deployment_token_info() { + Some(token) => AiUsageAttribution { + user_id: None, + project_id: Some(token.project_id), + environment_id: token.environment_id, + deployment_id: token.deployment_id, + deployment_token_id: Some(token.token_id), + }, + None => AiUsageAttribution { + user_id: auth.user_id_opt(), + ..Default::default() + }, + } +} + fn credential_type_str(ct: CredentialType) -> &'static str { match ct { CredentialType::System => "system", @@ -74,6 +89,63 @@ fn credential_type_str(ct: CredentialType) -> &'static str { } } +fn normalize_output_limit( + request: &mut ChatCompletionRequest, +) -> Result, AiGatewayError> { + let max_completion_tokens = request + .extra + .as_mut() + .and_then(|extra| extra.remove("max_completion_tokens")) + .map(|value| { + value.as_i64().ok_or_else(|| AiGatewayError::Validation { + message: "max_completion_tokens must be a positive integer".to_string(), + }) + }) + .transpose()?; + + if let (Some(max_tokens), Some(max_completion_tokens)) = + (request.max_tokens, max_completion_tokens) + { + if max_tokens != max_completion_tokens { + return Err(AiGatewayError::Validation { + message: "max_tokens and max_completion_tokens must match when both are provided" + .to_string(), + }); + } + } + + let effective = request.max_tokens.or(max_completion_tokens); + if effective.is_some_and(|value| value <= 0) { + return Err(AiGatewayError::Validation { + message: "max_tokens must be greater than zero".to_string(), + }); + } + // Keep one canonical cap. Provider adapters may translate this field but + // cannot preserve a conflicting caller-controlled alias afterward. + request.max_tokens = effective; + Ok(effective) +} + +fn projected_chat_input_tokens(request: &ChatCompletionRequest) -> Option { + let contains_image = request.messages.iter().any(|message| { + matches!( + message.content.as_ref(), + Some(MessageContent::Parts(parts)) if parts.iter().any(|part| part.image_url.is_some()) + ) + }); + if contains_image { + // A short remote URL can expand into a high-resolution provider-side + // image charge. Governance rejects it only when an operator-funded + // budget needs a hard upper bound; unbudgeted and BYOK calls proceed. + return None; + } + Some( + serde_json::to_vec(request) + .map(|body| i64::try_from(body.len()).unwrap_or(i64::MAX)) + .unwrap_or(i64::MAX), + ) +} + // ============================================================================ // Streaming usage extraction // ============================================================================ @@ -112,12 +184,13 @@ fn wrap_stream_with_usage_tracking( Box> + Send>, >, usage_service: Arc, - user_id: Option, + attribution: AiUsageAttribution, provider: String, model: String, start: Instant, is_byok: bool, ai_context: AiRequestContext, + reservation: GovernanceReservation, ) -> std::pin::Pin> + Send>> { use tokio_stream::StreamExt; @@ -162,20 +235,27 @@ fn wrap_stream_with_usage_tracking( let latency_ms = start.elapsed().as_millis() as i32; if input > 0 || output > 0 { + let estimated_cost = + crate::handlers::pricing::estimate_cost_microcents(&model, input, output) + .unwrap_or_else(|| { + tracing::warn!(model, "No pricing found for streaming AI usage record"); + 0 + }); tokio::spawn(async move { if let Err(e) = usage_service - .log_usage_with_context( - user_id, + .log_usage_with_context_and_reservation( + &attribution, &provider, &model, input, output, latency_ms, - 0, + estimated_cost, 200, true, // streaming is_byok, &ai_context, + Some(&reservation), ) .await { @@ -305,6 +385,30 @@ fn error_to_response(error: AiGatewayError) -> impl IntoResponse { "model_not_allowed", ), ), + AiGatewayError::RateLimitExceeded { .. } => ( + StatusCode::TOO_MANY_REQUESTS, + OpenAiErrorResponse::rate_limit(error.to_string()), + ), + AiGatewayError::MonthlyBudgetExceeded { .. } => ( + StatusCode::PAYMENT_REQUIRED, + OpenAiErrorResponse::new( + error.to_string(), + "insufficient_quota", + Some("budget_exceeded"), + ), + ), + AiGatewayError::PricingUnavailable { .. } => ( + StatusCode::SERVICE_UNAVAILABLE, + OpenAiErrorResponse::server_error(error.to_string(), "pricing_unavailable"), + ), + AiGatewayError::BudgetRequiresMaxTokens { .. } => ( + StatusCode::BAD_REQUEST, + OpenAiErrorResponse::invalid_request(error.to_string(), "max_tokens_required"), + ), + AiGatewayError::BudgetProjectionUnavailable { .. } => ( + StatusCode::BAD_REQUEST, + OpenAiErrorResponse::invalid_request(error.to_string(), "unsupported_budget_input"), + ), AiGatewayError::Validation { message } => ( StatusCode::BAD_REQUEST, OpenAiErrorResponse::invalid_request(message, "invalid_request"), @@ -325,10 +429,25 @@ fn error_to_response(error: AiGatewayError) -> impl IntoResponse { "invalid_provider_url", ), ), - _ => ( - StatusCode::INTERNAL_SERVER_ERROR, - OpenAiErrorResponse::server_error(error.to_string(), "internal_error"), - ), + AiGatewayError::InvalidGovernanceConfig { .. } + | AiGatewayError::InvalidGovernanceScope { .. } + | AiGatewayError::GovernanceConfigNotFound { .. } + | AiGatewayError::ProviderKeyNotFound { .. } + | AiGatewayError::TranslationError { .. } + | AiGatewayError::StreamError { .. } + | AiGatewayError::Encryption(_) + | AiGatewayError::HttpClient(_) + | AiGatewayError::Internal { .. } + | AiGatewayError::Database(_) => { + error!(error = %error, "Internal AI gateway error"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + OpenAiErrorResponse::server_error( + "AI gateway request failed due to an internal error.", + "internal_error", + ), + ) + } }; (status, Json(body)) @@ -356,7 +475,7 @@ async fn chat_completions( RequireAuth(auth): RequireAuth, State(app_state): State>, headers: HeaderMap, - Json(request): Json, + Json(mut request): Json, ) -> Result { permission_guard!(auth, AiGatewayExecute); @@ -399,14 +518,44 @@ async fn chat_completions( } } + let effective_max_tokens = match normalize_output_limit(&mut request) { + Ok(limit) => limit, + Err(error) => return Ok(error_to_response(error).into_response()), + }; + if request.n.is_some_and(|value| value <= 0) { + return Ok(error_to_response(AiGatewayError::Validation { + message: "n must be greater than zero".to_string(), + }) + .into_response()); + } + let byok = extract_byok(&headers); let ai_context = extract_ai_context(&headers); let start = Instant::now(); let model = request.model.clone(); let is_streaming = request.stream; - // None for deployment tokens (machine callers) so usage rows store NULL - // instead of falsely attributing all deployed-app traffic to user id 0. - let user_id = auth.user_id_opt(); + let attribution = usage_attribution(&auth); + // UTF-8 bytes are a conservative upper bound for tokenizer input tokens. + let projected_input_tokens = projected_chat_input_tokens(&request); + let projected_output_tokens = + effective_max_tokens.map(|tokens| tokens.saturating_mul(i64::from(request.n.unwrap_or(1)))); + let reservation = match app_state + .governance_service + .check_request( + &attribution, + &model, + byok.api_key.is_some(), + projected_input_tokens, + projected_output_tokens, + ) + .await + { + Ok(reservation) => reservation, + Err(error) => { + debug!(error = %error, "AI gateway governance rejected request"); + return Ok(error_to_response(error).into_response()); + } + }; if is_streaming { match app_state @@ -420,7 +569,7 @@ async fn chat_completions( info!( model = model, - user_id = ?user_id, + user_id = ?attribution.user_id, streaming = true, credential_type = credential_type_str(cred_type), "AI gateway streaming request started" @@ -440,12 +589,13 @@ async fn chat_completions( let wrapped = wrap_stream_with_usage_tracking( stream, app_state.usage_service.clone(), - user_id, + attribution.clone(), provider_id.to_string(), model.clone(), start, cred_type == CredentialType::Byok, ai_context.clone(), + reservation.clone(), ); let body = Body::from_stream(wrapped); @@ -469,6 +619,13 @@ async fn chat_completions( } } Err(e) => { + if let Err(release_error) = app_state + .governance_service + .release_cost_reservation(&reservation) + .await + { + error!(error = %release_error, "Failed to release AI gateway cost reservation"); + } error!(model = model, error = %e, "AI gateway streaming request failed"); Ok(error_to_response(e).into_response()) } @@ -484,7 +641,8 @@ async fn chat_completions( let provider_id = crate::providers::route_model_to_provider(&model).unwrap_or("unknown"); - // Log usage asynchronously (don't block the response) + // Persist before responding so the next request observes this + // spend when checking the monthly budget. if let Some(ref usage) = response.usage { let usage_service = app_state.usage_service.clone(); let model_clone = model.clone(); @@ -494,31 +652,37 @@ async fn chat_completions( let latency_ms = latency.as_millis() as i32; let is_byok = cred_type == CredentialType::Byok; let ctx = ai_context.clone(); - tokio::spawn(async move { - if let Err(e) = usage_service - .log_usage_with_context( - user_id, - &provider_clone, - &model_clone, - input, - output, - latency_ms, - 0, - 200, - false, // non-streaming path - is_byok, - &ctx, - ) - .await - { - error!(error = %e, "Failed to log AI usage"); - } - }); + let usage_attribution = attribution.clone(); + let estimated_cost = + crate::handlers::pricing::estimate_cost_microcents(&model, input, output) + .unwrap_or_else(|| { + tracing::warn!(model, "No pricing found for AI usage record"); + 0 + }); + if let Err(e) = usage_service + .log_usage_with_context_and_reservation( + &usage_attribution, + &provider_clone, + &model_clone, + input, + output, + latency_ms, + estimated_cost, + 200, + false, // non-streaming path + is_byok, + &ctx, + Some(&reservation), + ) + .await + { + error!(error = %e, "Failed to log AI usage"); + } } info!( model = model, - user_id = ?user_id, + user_id = ?attribution.user_id, latency_ms = latency.as_millis() as u64, credential_type = credential_type_str(cred_type), "AI gateway request completed" @@ -573,6 +737,13 @@ async fn chat_completions( } } Err(e) => { + if let Err(release_error) = app_state + .governance_service + .release_cost_reservation(&reservation) + .await + { + error!(error = %release_error, "Failed to release AI gateway cost reservation"); + } let latency = start.elapsed(); error!( model = model, @@ -661,7 +832,25 @@ async fn embeddings( let byok = extract_byok(&headers); let ai_context = extract_ai_context(&headers); let start = Instant::now(); - let user_id = auth.user_id_opt(); + let attribution = usage_attribution(&auth); + let projected_input_tokens = serde_json::to_vec(&request.input) + .map(|body| i64::try_from(body.len()).unwrap_or(i64::MAX)) + .unwrap_or(i64::MAX) + .into(); + let reservation = match app_state + .governance_service + .check_request( + &attribution, + &request.model, + byok.api_key.is_some(), + projected_input_tokens, + Some(0), + ) + .await + { + Ok(reservation) => reservation, + Err(error) => return Ok(error_to_response(error).into_response()), + }; match app_state.gateway_service.embeddings(&request, &byok).await { Ok((response, cred_type)) => { @@ -669,8 +858,8 @@ async fn embeddings( let provider_id = crate::providers::route_model_to_provider(&request.model).unwrap_or("unknown"); - // Log usage asynchronously (don't block the response). Embeddings - // only consume prompt tokens; there is no completion output. + // Persist before responding for quota consistency. Embeddings only + // consume prompt tokens; there is no completion output. { let usage_service = app_state.usage_service.clone(); let model = request.model.clone(); @@ -679,31 +868,37 @@ async fn embeddings( let latency_ms = latency.as_millis() as i32; let is_byok = cred_type == CredentialType::Byok; let ctx = ai_context.clone(); - tokio::spawn(async move { - if let Err(e) = usage_service - .log_usage_with_context( - user_id, - &provider, - &model, - input_tokens, - 0, - latency_ms, - 0, - 200, - false, - is_byok, - &ctx, - ) - .await - { - error!(error = %e, "Failed to log AI embedding usage"); - } - }); + let usage_attribution = attribution.clone(); + let estimated_cost = + crate::handlers::pricing::estimate_cost_microcents(&model, input_tokens, 0) + .unwrap_or_else(|| { + tracing::warn!(model, "No pricing found for AI embedding usage record"); + 0 + }); + if let Err(e) = usage_service + .log_usage_with_context_and_reservation( + &usage_attribution, + &provider, + &model, + input_tokens, + 0, + latency_ms, + estimated_cost, + 200, + false, + is_byok, + &ctx, + Some(&reservation), + ) + .await + { + error!(error = %e, "Failed to log AI embedding usage"); + } } info!( model = request.model, - user_id = ?user_id, + user_id = ?attribution.user_id, latency_ms = latency.as_millis() as u64, credential_type = credential_type_str(cred_type), "AI gateway embedding request completed" @@ -735,13 +930,23 @@ async fn embeddings( } } } - Err(e) => Ok(error_to_response(e).into_response()), + Err(e) => { + if let Err(release_error) = app_state + .governance_service + .release_cost_reservation(&reservation) + .await + { + error!(error = %release_error, "Failed to release AI gateway cost reservation"); + } + Ok(error_to_response(e).into_response()) + } } } #[cfg(test)] mod tests { use super::*; + use temps_entities::deployment_tokens::DeploymentTokenPermission; // Helper to build a chat completion request fn sample_chat_request() -> ChatCompletionRequest { @@ -771,6 +976,55 @@ mod tests { } } + #[test] + fn conflicting_output_caps_cannot_under_reserve_budget() { + let mut request = sample_chat_request(); + request.max_tokens = Some(1); + request.extra = Some(serde_json::Map::from_iter([( + "max_completion_tokens".to_string(), + serde_json::json!(100_000), + )])); + + assert!(matches!( + normalize_output_limit(&mut request), + Err(AiGatewayError::Validation { ref message }) + if message.contains("must match") + )); + } + + #[test] + fn max_completion_tokens_is_normalized_to_canonical_cap() { + let mut request = sample_chat_request(); + request.extra = Some(serde_json::Map::from_iter([( + "max_completion_tokens".to_string(), + serde_json::json!(512), + )])); + + assert_eq!( + normalize_output_limit(&mut request).expect("valid alias should normalize"), + Some(512) + ); + assert_eq!(request.max_tokens, Some(512)); + assert!(!request + .extra + .as_ref() + .is_some_and(|extra| extra.contains_key("max_completion_tokens"))); + } + + #[test] + fn remote_image_input_has_no_safe_budget_projection() { + let mut request = sample_chat_request(); + request.messages[0].content = Some(MessageContent::Parts(vec![ContentPart { + r#type: "image_url".to_string(), + text: None, + image_url: Some(serde_json::json!({ + "url": "https://example.test/high-resolution.png" + })), + }])); + + assert_eq!(projected_chat_input_tokens(&request), None); + } + #[test] fn test_error_to_response_model_not_found() { let err = AiGatewayError::ModelNotFound { @@ -1003,6 +1257,29 @@ mod tests { assert_eq!(credential_type_str(CredentialType::Byok), "byok"); } + #[test] + fn deployment_token_attribution_uses_trusted_auth_scope() { + let auth = AuthContext::new_deployment_token( + 7, + Some(11), + Some(13), + 17, + "production".to_string(), + vec![DeploymentTokenPermission::AiGatewayExecute], + ); + + assert_eq!( + usage_attribution(&auth), + AiUsageAttribution { + user_id: None, + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), + } + ); + } + #[test] fn test_extract_usage_from_sse_openai_final_chunk() { let line = r#"data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":42,"completion_tokens":18,"total_tokens":60}}"#; diff --git a/crates/temps-ai-gateway/src/handlers/governance.rs b/crates/temps-ai-gateway/src/handlers/governance.rs new file mode 100644 index 000000000..dde9d1e6a --- /dev/null +++ b/crates/temps-ai-gateway/src/handlers/governance.rs @@ -0,0 +1,258 @@ +use std::sync::Arc; + +use axum::{ + extract::{Path, State}, + http::StatusCode, + response::IntoResponse, + routing::{get, put}, + Extension, Json, Router, +}; +use serde::{Deserialize, Serialize}; +use temps_auth::{permission_guard, RequireAuth}; +use temps_core::audit::{AuditContext, AuditOperation}; +use temps_core::problemdetails::{Problem, ProblemDetails}; +use temps_core::RequestMetadata; +use utoipa::{OpenApi, ToSchema}; + +use crate::handlers::types::AiGatewayAppState; + +#[derive(OpenApi)] +#[openapi( + paths(list_governance_configs, upsert_governance_config, delete_governance_config), + components(schemas(GovernanceConfigResponse, UpsertGovernanceConfigRequest)), + info( + title = "AI Gateway Governance API", + description = "Configure instance, project, environment, and deployment-token AI gateway limits", + version = "1.0.0" + ), + tags((name = "AI Gateway Governance", description = "AI model, rate, and budget policy")) +)] +pub struct AiGatewayGovernanceApiDoc; + +pub fn configure_governance_routes() -> Router> { + Router::new() + .route("/ai/governance", get(list_governance_configs)) + .route( + "/ai/governance/{scope}", + put(upsert_governance_config).delete(delete_governance_config), + ) +} + +#[derive(Debug, Serialize, ToSchema)] +pub struct GovernanceConfigResponse { + pub id: i32, + pub scope: String, + pub allowed_models: Option>, + pub max_requests_per_minute: Option, + pub max_cost_per_month_microcents: Option, + pub created_at: String, + pub updated_at: String, +} + +impl From for GovernanceConfigResponse { + fn from(model: temps_entities::ai_gateway_config::Model) -> Self { + Self { + id: model.id, + scope: model.scope, + allowed_models: model.allowed_models.and_then(|value| { + value.as_array().map(|models| { + models + .iter() + .filter_map(|model| model.as_str().map(String::from)) + .collect() + }) + }), + max_requests_per_minute: model.max_requests_per_minute, + max_cost_per_month_microcents: model.max_cost_per_month_microcents, + created_at: model.created_at.to_rfc3339(), + updated_at: model.updated_at.to_rfc3339(), + } + } +} + +#[derive(Debug, Deserialize, ToSchema)] +pub struct UpsertGovernanceConfigRequest { + /// NULL allows every model; an empty array denies every model. + pub allowed_models: Option>, + /// NULL disables the request-rate limit for this scope. + pub max_requests_per_minute: Option, + /// NULL disables the operator-funded monthly budget for this scope. + pub max_cost_per_month_microcents: Option, +} + +#[derive(Debug, Serialize)] +struct GovernanceConfigAudit { + context: AuditContext, + action: String, + scope: String, +} + +impl AuditOperation for GovernanceConfigAudit { + fn operation_type(&self) -> String { + format!("AI_GATEWAY_GOVERNANCE_{}", self.action) + } + + fn user_id(&self) -> Option { + Some(self.context.user_id) + } + + fn ip_address(&self) -> Option { + self.context.ip_address.clone() + } + + fn user_agent(&self) -> &str { + &self.context.user_agent + } + + fn serialize(&self) -> temps_core::anyhow::Result { + serde_json::to_string(self).map_err(|error| { + temps_core::anyhow::anyhow!( + "Failed to serialize AI gateway governance audit for scope '{}': {}", + self.scope, + error + ) + }) + } +} + +#[utoipa::path( + tag = "AI Gateway Governance", + get, + path = "/ai/governance", + responses( + (status = 200, body = Vec), + (status = 401, body = ProblemDetails), + (status = 403, body = ProblemDetails), + ), + security(("bearer_auth" = [])) +)] +async fn list_governance_configs( + RequireAuth(auth): RequireAuth, + State(app_state): State>, +) -> Result { + // Governance policies expose budget and allowlist details and therefore + // use the same operator-only permission as mutations. + permission_guard!(auth, AiGatewayWrite); + let configs = app_state.governance_service.list_configs().await?; + Ok(Json( + configs + .into_iter() + .map(GovernanceConfigResponse::from) + .collect::>(), + )) +} + +#[utoipa::path( + tag = "AI Gateway Governance", + put, + path = "/ai/governance/{scope}", + params(("scope" = String, Path, description = "instance, project:, environment:, or token:")), + request_body = UpsertGovernanceConfigRequest, + responses( + (status = 200, body = GovernanceConfigResponse), + (status = 400, body = ProblemDetails), + (status = 401, body = ProblemDetails), + (status = 403, body = ProblemDetails), + ), + security(("bearer_auth" = [])) +)] +async fn upsert_governance_config( + RequireAuth(auth): RequireAuth, + State(app_state): State>, + Extension(metadata): Extension, + Path(scope): Path, + Json(request): Json, +) -> Result { + permission_guard!(auth, AiGatewayWrite); + + let allowed_models = request + .allowed_models + .map(|models| serde_json::json!(models)); + let config = app_state + .governance_service + .upsert_config( + &scope, + allowed_models, + request.max_requests_per_minute, + request.max_cost_per_month_microcents, + ) + .await?; + + audit_change(&app_state, &auth, &metadata, "UPSERTED", &scope).await; + Ok(Json(GovernanceConfigResponse::from(config))) +} + +#[utoipa::path( + tag = "AI Gateway Governance", + delete, + path = "/ai/governance/{scope}", + params(("scope" = String, Path, description = "instance, project:, environment:, or token:")), + responses( + (status = 204), + (status = 401, body = ProblemDetails), + (status = 403, body = ProblemDetails), + (status = 404, body = ProblemDetails), + ), + security(("bearer_auth" = [])) +)] +async fn delete_governance_config( + RequireAuth(auth): RequireAuth, + State(app_state): State>, + Extension(metadata): Extension, + Path(scope): Path, +) -> Result { + permission_guard!(auth, AiGatewayWrite); + app_state.governance_service.delete_config(&scope).await?; + audit_change(&app_state, &auth, &metadata, "DELETED", &scope).await; + Ok(StatusCode::NO_CONTENT) +} + +async fn audit_change( + app_state: &AiGatewayAppState, + auth: &temps_auth::AuthContext, + metadata: &RequestMetadata, + action: &str, + scope: &str, +) { + let audit = GovernanceConfigAudit { + context: AuditContext { + user_id: auth.user_id(), + ip_address: Some(metadata.ip_address.clone()), + user_agent: metadata.user_agent.clone(), + }, + action: action.to_string(), + scope: scope.to_string(), + }; + if let Err(error) = app_state.audit_service.create_audit_log(&audit).await { + tracing::error!( + error = %error, + scope, + action, + "Failed to create AI gateway governance audit log" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_response_preserves_scope_and_limits() { + let response = GovernanceConfigResponse::from(temps_entities::ai_gateway_config::Model { + id: 1, + scope: "project:7".to_string(), + allowed_models: Some(serde_json::json!(["gpt-5-mini"])), + max_requests_per_minute: Some(60), + max_cost_per_month_microcents: Some(1_000_000), + created_at: chrono::Utc::now(), + updated_at: chrono::Utc::now(), + }); + assert_eq!(response.scope, "project:7"); + assert_eq!( + response.allowed_models, + Some(vec!["gpt-5-mini".to_string()]) + ); + assert_eq!(response.max_requests_per_minute, Some(60)); + } +} diff --git a/crates/temps-ai-gateway/src/handlers/mod.rs b/crates/temps-ai-gateway/src/handlers/mod.rs index 29fb71872..a86814ac4 100644 --- a/crates/temps-ai-gateway/src/handlers/mod.rs +++ b/crates/temps-ai-gateway/src/handlers/mod.rs @@ -1,10 +1,12 @@ pub mod gateway; +pub mod governance; pub mod pricing; pub mod providers; pub mod types; pub mod usage; pub use gateway::configure_gateway_routes; +pub use governance::configure_governance_routes; pub use pricing::configure_pricing_routes; pub use providers::configure_admin_routes; pub use types::{create_ai_gateway_app_state, AiGatewayAppState}; diff --git a/crates/temps-ai-gateway/src/handlers/pricing.rs b/crates/temps-ai-gateway/src/handlers/pricing.rs index 1ad2294c7..6f80cefc6 100644 --- a/crates/temps-ai-gateway/src/handlers/pricing.rs +++ b/crates/temps-ai-gateway/src/handlers/pricing.rs @@ -42,7 +42,7 @@ pub struct PricingResponse { /// Pricing for a single model, all values in USD per 1M tokens. /// Fields are optional because not every provider supports every pricing tier. -#[derive(Debug, Serialize, ToSchema)] +#[derive(Debug, Clone, Serialize, ToSchema)] pub struct ModelPricing { /// Model identifier (e.g. "gpt-5.4", "claude-sonnet-4-6") pub model: String, @@ -145,7 +145,7 @@ impl PricingBuilder { } } -fn build_pricing() -> Vec { +pub(crate) fn build_pricing() -> Vec { vec![ // ── Anthropic ─────────────────────────────────────────────────── PricingBuilder::new("anthropic", "claude-opus-4-6", "Claude Opus 4.6", 5.0, 25.0) @@ -223,6 +223,22 @@ fn build_pricing() -> Vec { .cache(0.0, 0.0, 0.075) .batch(0.075, 0.30) .build(), + PricingBuilder::new( + "openai", + "text-embedding-3-small", + "Text Embedding 3 Small", + 0.02, + 0.0, + ) + .build(), + PricingBuilder::new( + "openai", + "text-embedding-3-large", + "Text Embedding 3 Large", + 0.13, + 0.0, + ) + .build(), // ── xAI ───────────────────────────────────────────────────────── PricingBuilder::new( "xai", @@ -306,6 +322,34 @@ fn build_pricing() -> Vec { ] } +/// Estimate provider cost in microcents (1/1,000,000 cent) from token usage. +/// +/// A price expressed as USD per million tokens becomes 100 microcents per +/// token because 1 USD is 100,000,000 microcents. Returning `None` for an +/// unknown model is intentional; callers must not claim such traffic was free. +pub(crate) fn estimate_cost_microcents( + model: &str, + input_tokens: i64, + output_tokens: i64, +) -> Option { + if input_tokens < 0 || output_tokens < 0 { + return None; + } + + let pricing = build_pricing() + .into_iter() + .find(|pricing| pricing.model == model)?; + let cost = (input_tokens as f64 * pricing.input_per_million + + output_tokens as f64 * pricing.output_per_million) + * 100.0; + if !cost.is_finite() || cost > i64::MAX as f64 { + return None; + } + // Always round positive fractional microcents upward. Rounding down would + // let sufficiently small requests consume a zero-cost budget reservation. + Some(cost.ceil() as i64) +} + // ============================================================================ // Handlers // ============================================================================ @@ -360,13 +404,35 @@ mod tests { model.model ); assert!( - model.output_per_million > 0.0, - "Output price for {} must be positive", + model.output_per_million >= 0.0, + "Output price for {} must not be negative", model.model ); } } + #[test] + fn estimate_cost_uses_input_and_output_prices() { + assert_eq!( + estimate_cost_microcents("gpt-5-mini", 1_000, 500), + Some(120_000) + ); + } + + #[test] + fn estimate_embedding_cost_uses_input_only() { + assert_eq!( + estimate_cost_microcents("text-embedding-3-small", 10_000, 0), + Some(20_000) + ); + } + + #[test] + fn estimate_cost_rejects_unknown_model_and_negative_usage() { + assert_eq!(estimate_cost_microcents("custom-model", 10, 10), None); + assert_eq!(estimate_cost_microcents("gpt-5-mini", -1, 10), None); + } + #[test] fn test_anthropic_has_cache_pricing() { let pricing = build_pricing(); diff --git a/crates/temps-ai-gateway/src/handlers/providers.rs b/crates/temps-ai-gateway/src/handlers/providers.rs index 5d6cc9f4f..556f7dc9c 100644 --- a/crates/temps-ai-gateway/src/handlers/providers.rs +++ b/crates/temps-ai-gateway/src/handlers/providers.rs @@ -49,6 +49,46 @@ impl From for Problem { AiGatewayError::ModelNotAllowed { .. } => problemdetails::new(StatusCode::FORBIDDEN) .with_title("Model Not Allowed") .with_detail(error.to_string()), + AiGatewayError::RateLimitExceeded { .. } => { + problemdetails::new(StatusCode::TOO_MANY_REQUESTS) + .with_title("AI Gateway Rate Limit Exceeded") + .with_detail(error.to_string()) + } + AiGatewayError::MonthlyBudgetExceeded { .. } => { + problemdetails::new(StatusCode::PAYMENT_REQUIRED) + .with_title("AI Gateway Budget Exceeded") + .with_detail(error.to_string()) + } + AiGatewayError::PricingUnavailable { .. } => { + problemdetails::new(StatusCode::SERVICE_UNAVAILABLE) + .with_title("AI Gateway Pricing Unavailable") + .with_detail(error.to_string()) + } + AiGatewayError::BudgetRequiresMaxTokens { .. } => { + problemdetails::new(StatusCode::BAD_REQUEST) + .with_title("AI Gateway max_tokens Required") + .with_detail(error.to_string()) + } + AiGatewayError::BudgetProjectionUnavailable { .. } => { + problemdetails::new(StatusCode::BAD_REQUEST) + .with_title("Unsupported Budgeted AI Input") + .with_detail(error.to_string()) + } + AiGatewayError::InvalidGovernanceConfig { .. } => { + problemdetails::new(StatusCode::INTERNAL_SERVER_ERROR) + .with_title("Invalid AI Gateway Configuration") + .with_detail(error.to_string()) + } + AiGatewayError::InvalidGovernanceScope { .. } => { + problemdetails::new(StatusCode::BAD_REQUEST) + .with_title("Invalid AI Gateway Scope") + .with_detail(error.to_string()) + } + AiGatewayError::GovernanceConfigNotFound { .. } => { + problemdetails::new(StatusCode::NOT_FOUND) + .with_title("AI Gateway Configuration Not Found") + .with_detail(error.to_string()) + } AiGatewayError::UpstreamError { status, .. } => { let http_status = StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); problemdetails::new(http_status) diff --git a/crates/temps-ai-gateway/src/handlers/types.rs b/crates/temps-ai-gateway/src/handlers/types.rs index 105db2d5e..af5c1558f 100644 --- a/crates/temps-ai-gateway/src/handlers/types.rs +++ b/crates/temps-ai-gateway/src/handlers/types.rs @@ -1,12 +1,13 @@ use std::sync::Arc; use temps_core::AuditLogger; -use crate::services::{GatewayService, ProviderKeyService, UsageService}; +use crate::services::{GatewayService, GovernanceService, ProviderKeyService, UsageService}; pub struct AiGatewayAppState { pub gateway_service: Arc, pub provider_key_service: Arc, pub usage_service: Arc, + pub governance_service: Arc, pub audit_service: Arc, pub telemetry: Arc, } @@ -15,6 +16,7 @@ pub async fn create_ai_gateway_app_state( gateway_service: Arc, provider_key_service: Arc, usage_service: Arc, + governance_service: Arc, audit_service: Arc, telemetry: Arc, ) -> Arc { @@ -22,6 +24,7 @@ pub async fn create_ai_gateway_app_state( gateway_service, provider_key_service, usage_service, + governance_service, audit_service, telemetry, }) diff --git a/crates/temps-ai-gateway/src/handlers/usage.rs b/crates/temps-ai-gateway/src/handlers/usage.rs index 3dc8b04fd..d6c736478 100644 --- a/crates/temps-ai-gateway/src/handlers/usage.rs +++ b/crates/temps-ai-gateway/src/handlers/usage.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use temps_auth::permission_guard; use temps_auth::RequireAuth; use temps_core::problemdetails::{Problem, ProblemDetails}; -use utoipa::{OpenApi, ToSchema}; +use utoipa::{IntoParams, OpenApi, ToSchema}; use crate::error::AiGatewayError; use crate::handlers::types::AiGatewayAppState; @@ -78,7 +78,8 @@ pub fn configure_usage_routes() -> Router> { // Query param structs // ============================================================================ -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Deserialize, IntoParams, ToSchema)] +#[into_params(parameter_in = Query)] pub struct UsageQueryParams { /// ISO 8601 start time (defaults to 24h ago) pub from: Option, @@ -86,6 +87,14 @@ pub struct UsageQueryParams { pub to: Option, /// Filter by user ID pub user_id: Option, + /// Filter by deployment-token project ID + pub project_id: Option, + /// Filter by deployment-token environment ID + pub environment_id: Option, + /// Filter by deployment ID + pub deployment_id: Option, + /// Filter by deployment-token ID + pub deployment_token_id: Option, /// Filter by conversation ID pub conversation_id: Option, /// Filter by tags (comma-separated, AND logic) @@ -100,6 +109,10 @@ impl UsageQueryParams { fn to_filter(&self) -> UsageFilter { UsageFilter { user_id: self.user_id, + project_id: self.project_id, + environment_id: self.environment_id, + deployment_id: self.deployment_id, + deployment_token_id: self.deployment_token_id, conversation_id: self.conversation_id.clone(), tags: self.tags.clone(), model: self.model.clone(), @@ -109,7 +122,8 @@ impl UsageQueryParams { } } -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Deserialize, IntoParams, ToSchema)] +#[into_params(parameter_in = Query)] pub struct TimeseriesQueryParams { /// ISO 8601 start time (defaults to 24h ago) pub from: Option, @@ -119,6 +133,10 @@ pub struct TimeseriesQueryParams { pub bucket: Option, /// Filter by user ID pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, + pub deployment_id: Option, + pub deployment_token_id: Option, /// Filter by conversation ID pub conversation_id: Option, /// Filter by tags (comma-separated, AND logic) @@ -133,6 +151,10 @@ impl TimeseriesQueryParams { fn to_filter(&self) -> UsageFilter { UsageFilter { user_id: self.user_id, + project_id: self.project_id, + environment_id: self.environment_id, + deployment_id: self.deployment_id, + deployment_token_id: self.deployment_token_id, conversation_id: self.conversation_id.clone(), tags: self.tags.clone(), model: self.model.clone(), @@ -142,7 +164,8 @@ impl TimeseriesQueryParams { } } -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Deserialize, IntoParams, ToSchema)] +#[into_params(parameter_in = Query)] pub struct TopModelsQueryParams { /// ISO 8601 start time (defaults to 24h ago) pub from: Option, @@ -152,6 +175,8 @@ pub struct TopModelsQueryParams { pub limit: Option, /// Filter by user ID pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, /// Filter by tags (comma-separated, AND logic) pub tags: Option, } @@ -160,6 +185,8 @@ impl TopModelsQueryParams { fn to_filter(&self) -> UsageFilter { UsageFilter { user_id: self.user_id, + project_id: self.project_id, + environment_id: self.environment_id, tags: self.tags.clone(), ..Default::default() } @@ -170,7 +197,8 @@ impl TopModelsQueryParams { pub const RECENT_DEFAULT_LIMIT: u64 = 20; pub const RECENT_MAX_LIMIT: u64 = 50; -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Deserialize, IntoParams, ToSchema)] +#[into_params(parameter_in = Query)] pub struct RecentQueryParams { /// Page size (defaults to 20, max 50) pub limit: Option, @@ -178,6 +206,10 @@ pub struct RecentQueryParams { pub offset: Option, /// Filter by user ID pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, + pub deployment_id: Option, + pub deployment_token_id: Option, /// Filter by conversation ID pub conversation_id: Option, /// Filter by tags (comma-separated, AND logic) @@ -210,6 +242,10 @@ impl RecentQueryParams { fn to_filter(&self) -> UsageFilter { UsageFilter { user_id: self.user_id, + project_id: self.project_id, + environment_id: self.environment_id, + deployment_id: self.deployment_id, + deployment_token_id: self.deployment_token_id, conversation_id: self.conversation_id.clone(), tags: self.tags.clone(), model: self.model.clone(), @@ -234,7 +270,8 @@ impl RecentQueryParams { } } -#[derive(Debug, Deserialize, ToSchema)] +#[derive(Debug, Deserialize, IntoParams, ToSchema)] +#[into_params(parameter_in = Query)] pub struct ConversationsQueryParams { /// ISO 8601 start time (defaults to 24h ago) pub from: Option, @@ -244,6 +281,8 @@ pub struct ConversationsQueryParams { pub limit: Option, /// Filter by user ID pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, /// Filter by tags (comma-separated, AND logic) pub tags: Option, /// Filter by model name @@ -254,6 +293,8 @@ impl ConversationsQueryParams { fn to_filter(&self) -> UsageFilter { UsageFilter { user_id: self.user_id, + project_id: self.project_id, + environment_id: self.environment_id, tags: self.tags.clone(), model: self.model.clone(), ..Default::default() @@ -303,10 +344,7 @@ fn parse_time_range( tag = "AI Gateway Usage", get, path = "/ai/usage/summary", - params( - ("from" = Option, Query, description = "ISO 8601 start time (defaults to 24h ago)"), - ("to" = Option, Query, description = "ISO 8601 end time (defaults to now)"), - ), + params(UsageQueryParams), responses( (status = 200, description = "Usage summary for the time range", body = UsageSummary), (status = 400, description = "Invalid query parameters", body = ProblemDetails), @@ -336,10 +374,7 @@ async fn get_usage_summary( tag = "AI Gateway Usage", get, path = "/ai/usage/by-provider", - params( - ("from" = Option, Query, description = "ISO 8601 start time (defaults to 24h ago)"), - ("to" = Option, Query, description = "ISO 8601 end time (defaults to now)"), - ), + params(UsageQueryParams), responses( (status = 200, description = "Usage broken down by provider", body = Vec), (status = 400, description = "Invalid query parameters", body = ProblemDetails), @@ -369,11 +404,7 @@ async fn get_usage_by_provider( tag = "AI Gateway Usage", get, path = "/ai/usage/timeseries", - params( - ("from" = Option, Query, description = "ISO 8601 start time (defaults to 24h ago)"), - ("to" = Option, Query, description = "ISO 8601 end time (defaults to now)"), - ("bucket" = Option, Query, description = "Bucket size: hour, day, week (defaults to day)"), - ), + params(TimeseriesQueryParams), responses( (status = 200, description = "Time-series usage data", body = Vec), (status = 400, description = "Invalid query parameters", body = ProblemDetails), @@ -404,11 +435,7 @@ async fn get_usage_timeseries( tag = "AI Gateway Usage", get, path = "/ai/usage/top-models", - params( - ("from" = Option, Query, description = "ISO 8601 start time (defaults to 24h ago)"), - ("to" = Option, Query, description = "ISO 8601 end time (defaults to now)"), - ("limit" = Option, Query, description = "Max results (defaults to 10)"), - ), + params(TopModelsQueryParams), responses( (status = 200, description = "Top models by request count", body = Vec), (status = 400, description = "Invalid query parameters", body = ProblemDetails), @@ -439,21 +466,7 @@ async fn get_usage_top_models( tag = "AI Gateway Usage", get, path = "/ai/usage/recent", - params( - ("limit" = Option, Query, description = "Page size (defaults to 20, max 50)"), - ("offset" = Option, Query, description = "Number of results to skip for pagination (defaults to 0)"), - ("provider" = Option, Query, description = "Filter by provider name"), - ("model" = Option, Query, description = "Filter by model name"), - ("status" = Option, Query, description = "Filter by HTTP status code (exact match)"), - ("cost_gte" = Option, Query, description = "Cost greater-than-or-equal, in microcents"), - ("cost_gt" = Option, Query, description = "Cost strictly greater-than, in microcents"), - ("cost_lte" = Option, Query, description = "Cost less-than-or-equal, in microcents"), - ("cost_lt" = Option, Query, description = "Cost strictly less-than, in microcents"), - ("tokens_gte" = Option, Query, description = "Total tokens greater-than-or-equal"), - ("tokens_gt" = Option, Query, description = "Total tokens strictly greater-than"), - ("tokens_lte" = Option, Query, description = "Total tokens less-than-or-equal"), - ("tokens_lt" = Option, Query, description = "Total tokens strictly less-than"), - ), + params(RecentQueryParams), responses( (status = 200, description = "Page of recent usage log entries", body = UsageLogPage), (status = 401, description = "Unauthorized", body = ProblemDetails), @@ -483,14 +496,7 @@ async fn get_usage_recent( tag = "AI Gateway Usage", get, path = "/ai/usage/conversations", - params( - ("from" = Option, Query, description = "ISO 8601 start time (defaults to 24h ago)"), - ("to" = Option, Query, description = "ISO 8601 end time (defaults to now)"), - ("limit" = Option, Query, description = "Max results (defaults to 50, max 100)"), - ("user_id" = Option, Query, description = "Filter by user ID"), - ("tags" = Option, Query, description = "Filter by tags (comma-separated)"), - ("model" = Option, Query, description = "Filter by model name"), - ), + params(ConversationsQueryParams), responses( (status = 200, description = "Conversation summaries", body = Vec), (status = 400, description = "Invalid query parameters", body = ProblemDetails), diff --git a/crates/temps-ai-gateway/src/plugin.rs b/crates/temps-ai-gateway/src/plugin.rs index 765c6d5e9..cd5dee69f 100644 --- a/crates/temps-ai-gateway/src/plugin.rs +++ b/crates/temps-ai-gateway/src/plugin.rs @@ -11,7 +11,7 @@ use utoipa::OpenApi as OpenApiTrait; use crate::{ handlers::{self, create_ai_gateway_app_state, AiGatewayAppState}, - services::{GatewayService, ProviderKeyService, UsageService}, + services::{GatewayService, GovernanceService, ProviderKeyService, UsageService}, }; pub struct AiGatewayPlugin; @@ -57,9 +57,12 @@ impl TempsPlugin for AiGatewayPlugin { )); context.register_service(ai_service as Arc); - let usage_service = Arc::new(UsageService::new(db)); + let usage_service = Arc::new(UsageService::new(db.clone())); context.register_service(usage_service.clone()); + let governance_service = Arc::new(GovernanceService::new(db)); + context.register_service(governance_service.clone()); + let audit_service = context.require_service::(); let telemetry = context @@ -72,6 +75,7 @@ impl TempsPlugin for AiGatewayPlugin { gateway_service, provider_key_service, usage_service, + governance_service, audit_service, telemetry, ) @@ -93,6 +97,7 @@ impl TempsPlugin for AiGatewayPlugin { let routes = handlers::configure_admin_routes() .merge(handlers::configure_usage_routes()) .merge(handlers::configure_pricing_routes()) + .merge(handlers::configure_governance_routes()) .merge(handlers::configure_gateway_routes()) .with_state(app_state); @@ -107,6 +112,9 @@ impl TempsPlugin for AiGatewayPlugin { schema.merge(usage_schema); let pricing_schema = ::openapi(); schema.merge(pricing_schema); + let governance_schema = + ::openapi(); + schema.merge(governance_schema); Some(schema) } } diff --git a/crates/temps-ai-gateway/src/providers/openai_compat.rs b/crates/temps-ai-gateway/src/providers/openai_compat.rs index 6d8c504d9..c1b3137aa 100644 --- a/crates/temps-ai-gateway/src/providers/openai_compat.rs +++ b/crates/temps-ai-gateway/src/providers/openai_compat.rs @@ -202,10 +202,7 @@ impl AiProvider for OpenAiCompatProvider { } // Inject stream_options.include_usage so the final chunk includes token counts - let extra = req.extra.get_or_insert_with(Default::default); - extra - .entry("stream_options") - .or_insert_with(|| serde_json::json!({"include_usage": true})); + force_stream_usage(&mut req); let response = self .client @@ -294,6 +291,16 @@ impl AiProvider for OpenAiCompatProvider { } } +fn force_stream_usage(request: &mut ChatCompletionRequest) { + let extra = request.extra.get_or_insert_with(Default::default); + // This must overwrite caller input: final usage is needed to convert a + // conservative budget reservation into the provider's actual token cost. + extra.insert( + "stream_options".to_string(), + serde_json::json!({"include_usage": true}), + ); +} + /// Returns true if the model is an OpenAI o-series reasoning model. fn is_o_series_model(model: &str) -> bool { let m = model.to_lowercase(); @@ -526,4 +533,23 @@ mod tests { assert!(!is_o_series_model("gpt-4o")); assert!(!is_o_series_model("grok-3")); } + + #[test] + fn force_stream_usage_overrides_caller_opt_out() { + let mut request = test_request("gpt-4o"); + request.extra = Some(serde_json::Map::from_iter([( + "stream_options".to_string(), + serde_json::json!({"include_usage": false}), + )])); + + force_stream_usage(&mut request); + + assert_eq!( + request + .extra + .as_ref() + .and_then(|extra| extra.get("stream_options")), + Some(&serde_json::json!({"include_usage": true})) + ); + } } diff --git a/crates/temps-ai-gateway/src/services/ai_service.rs b/crates/temps-ai-gateway/src/services/ai_service.rs index 35e1b895f..1231cfd2d 100644 --- a/crates/temps-ai-gateway/src/services/ai_service.rs +++ b/crates/temps-ai-gateway/src/services/ai_service.rs @@ -1,10 +1,12 @@ //! ADR-022: the gateway-backed implementation of the general [`AiService`] //! foundation. //! -//! Wraps [`GatewayService`] so every internal AI call inherits provider-key -//! resolution, model routing, and per-scope rate/cost governance. Structured -//! output rides the gateway's existing `response_format` plumbing. Best-effort: -//! returns [`AiError`] rather than panicking; callers add the timeout. +//! Wraps [`GatewayService`] so internal AI calls inherit provider-key resolution +//! and model routing. Deployment-token rate/cost governance and usage +//! attribution are enforced by the public gateway HTTP handlers; internal +//! control-plane calls are intentionally outside that accounting boundary. +//! Structured output rides the gateway's existing `response_format` plumbing. +//! Best-effort: returns [`AiError`] rather than panicking; callers add the timeout. use std::sync::Arc; diff --git a/crates/temps-ai-gateway/src/services/governance_service.rs b/crates/temps-ai-gateway/src/services/governance_service.rs new file mode 100644 index 000000000..bf2b26393 --- /dev/null +++ b/crates/temps-ai-gateway/src/services/governance_service.rs @@ -0,0 +1,756 @@ +use std::sync::Arc; + +use chrono::{Datelike, TimeZone, Utc}; +use sea_orm::{ + ActiveModelTrait, ActiveValue::Set, ColumnTrait, ConnectionTrait, DatabaseBackend, + DatabaseConnection, DatabaseTransaction, EntityTrait, FromQueryResult, QueryFilter, QueryOrder, + Statement, TransactionTrait, +}; + +use crate::error::AiGatewayError; +use crate::handlers::pricing::estimate_cost_microcents; + +use super::AiUsageAttribution; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GovernanceReservation { + request_id: String, + billing_period: chrono::NaiveDate, +} + +impl GovernanceReservation { + pub(crate) fn request_id(&self) -> &str { + &self.request_id + } + + pub(crate) fn billing_period(&self) -> chrono::NaiveDate { + self.billing_period + } +} + +#[derive(Debug, FromQueryResult)] +struct CountRow { + count: Option, + retry_after_seconds: Option, +} + +#[derive(Debug, FromQueryResult)] +struct CostRow { + cost: Option, +} + +#[derive(Debug)] +struct AppliedConfig { + scope: String, + max_requests_per_minute: Option, + max_cost_per_month_microcents: Option, +} + +/// Enforces persisted instance/project/environment/token AI gateway policy. +/// +/// Rate events and conservative cost reservations live in PostgreSQL. Advisory +/// transaction locks serialize checks for each scope, so limits remain valid +/// when multiple console processes serve requests concurrently. +pub struct GovernanceService { + db: Arc, +} + +impl GovernanceService { + pub fn new(db: Arc) -> Self { + Self { db } + } + + #[allow(clippy::too_many_arguments)] + pub async fn check_request( + &self, + attribution: &AiUsageAttribution, + model: &str, + is_byok: bool, + projected_input_tokens: Option, + max_output_tokens: Option, + ) -> Result { + let scope_names = applicable_scope_names(attribution); + let rows = temps_entities::ai_gateway_config::Entity::find() + .filter(temps_entities::ai_gateway_config::Column::Scope.is_in(scope_names.clone())) + .all(self.db.as_ref()) + .await?; + + let mut configs = Vec::new(); + for scope in scope_names { + let Some(config) = rows.iter().find(|row| row.scope == scope) else { + continue; + }; + + if let Some(allowed_models) = config.allowed_models.as_ref() { + let allowed = allowed_models + .as_array() + .is_some_and(|models| models.iter().any(|entry| entry.as_str() == Some(model))); + if !allowed { + return Err(AiGatewayError::ModelNotAllowed { + model: model.to_string(), + scope: scope.clone(), + }); + } + } + + validate_nonnegative( + &scope, + "max_requests_per_minute", + config.max_requests_per_minute, + )?; + validate_nonnegative( + &scope, + "max_cost_per_month_microcents", + config.max_cost_per_month_microcents, + )?; + + configs.push(AppliedConfig { + scope, + max_requests_per_minute: config.max_requests_per_minute, + max_cost_per_month_microcents: config.max_cost_per_month_microcents, + }); + } + + let budget_scope = configs + .iter() + .find(|config| config.max_cost_per_month_microcents.is_some()) + .map(|config| config.scope.clone()); + let projected_cost = if is_byok || budget_scope.is_none() { + None + } else { + let output_tokens = + max_output_tokens.ok_or_else(|| AiGatewayError::BudgetRequiresMaxTokens { + scope: budget_scope + .clone() + .unwrap_or_else(|| "instance".to_string()), + })?; + let input_tokens = projected_input_tokens.ok_or_else(|| { + AiGatewayError::BudgetProjectionUnavailable { + scope: budget_scope + .clone() + .unwrap_or_else(|| "instance".to_string()), + } + })?; + estimate_cost_microcents(model, input_tokens, output_tokens) + .ok_or_else(|| AiGatewayError::PricingUnavailable { + model: model.to_string(), + scope: budget_scope.unwrap_or_else(|| "instance".to_string()), + })? + .into() + }; + + let request_id = uuid::Uuid::new_v4().to_string(); + let billing_period = current_month_start()?.date_naive(); + let has_rate_limit = configs + .iter() + .any(|config| config.max_requests_per_minute.is_some()); + if configs.is_empty() || (!has_rate_limit && projected_cost.is_none()) { + return Ok(GovernanceReservation { + request_id, + billing_period, + }); + } + + // Settle stale reservations in their own transaction so a subsequent + // budget rejection cannot roll the conservative debit back. + let cleanup_txn = self.db.begin().await?; + self.cleanup_expired_state(&cleanup_txn).await?; + cleanup_txn.commit().await?; + + let txn = self.db.begin().await?; + self.lock_scopes(&txn, &configs).await?; + self.check_rates_and_record(&txn, &configs, &request_id) + .await?; + if let Some(projected_cost) = projected_cost { + self.check_budgets_and_reserve( + &txn, + attribution, + &configs, + &request_id, + projected_cost, + billing_period, + ) + .await?; + } + txn.commit().await?; + + Ok(GovernanceReservation { + request_id, + billing_period, + }) + } + + pub async fn release_cost_reservation( + &self, + reservation: &GovernanceReservation, + ) -> Result<(), AiGatewayError> { + self.db + .execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "DELETE FROM ai_gateway_cost_reservations WHERE request_id = $1::uuid", + [reservation.request_id.clone().into()], + )) + .await?; + Ok(()) + } + + pub async fn list_configs( + &self, + ) -> Result, AiGatewayError> { + Ok(temps_entities::ai_gateway_config::Entity::find() + .order_by_asc(temps_entities::ai_gateway_config::Column::Scope) + .all(self.db.as_ref()) + .await?) + } + + pub async fn upsert_config( + &self, + scope: &str, + allowed_models: Option, + max_requests_per_minute: Option, + max_cost_per_month_microcents: Option, + ) -> Result { + validate_scope(scope)?; + validate_nonnegative(scope, "max_requests_per_minute", max_requests_per_minute)?; + validate_nonnegative( + scope, + "max_cost_per_month_microcents", + max_cost_per_month_microcents, + )?; + validate_allowed_models(scope, allowed_models.as_ref())?; + + let existing = temps_entities::ai_gateway_config::Entity::find() + .filter(temps_entities::ai_gateway_config::Column::Scope.eq(scope)) + .one(self.db.as_ref()) + .await?; + + let model = match existing { + Some(existing) => { + let mut active: temps_entities::ai_gateway_config::ActiveModel = existing.into(); + active.allowed_models = Set(allowed_models); + active.max_requests_per_minute = Set(max_requests_per_minute); + active.max_cost_per_month_microcents = Set(max_cost_per_month_microcents); + active.update(self.db.as_ref()).await? + } + None => { + temps_entities::ai_gateway_config::ActiveModel { + scope: Set(scope.to_string()), + allowed_models: Set(allowed_models), + max_requests_per_minute: Set(max_requests_per_minute), + max_cost_per_month_microcents: Set(max_cost_per_month_microcents), + ..Default::default() + } + .insert(self.db.as_ref()) + .await? + } + }; + + Ok(model) + } + + pub async fn delete_config(&self, scope: &str) -> Result<(), AiGatewayError> { + validate_scope(scope)?; + let result = temps_entities::ai_gateway_config::Entity::delete_many() + .filter(temps_entities::ai_gateway_config::Column::Scope.eq(scope)) + .exec(self.db.as_ref()) + .await?; + if result.rows_affected == 0 { + return Err(AiGatewayError::GovernanceConfigNotFound { + scope: scope.to_string(), + }); + } + Ok(()) + } + + async fn lock_scopes( + &self, + txn: &DatabaseTransaction, + configs: &[AppliedConfig], + ) -> Result<(), AiGatewayError> { + let mut scopes = configs + .iter() + .filter(|config| { + config.max_requests_per_minute.is_some() + || config.max_cost_per_month_microcents.is_some() + }) + .map(|config| config.scope.as_str()) + .collect::>(); + scopes.sort_unstable(); + for scope in scopes { + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "SELECT pg_advisory_xact_lock(hashtextextended($1, 908245731))", + [scope.into()], + )) + .await?; + } + Ok(()) + } + + async fn cleanup_expired_state(&self, txn: &DatabaseTransaction) -> Result<(), AiGatewayError> { + let month_start = current_month_start()?; + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "DELETE FROM ai_gateway_rate_events WHERE occurred_at < NOW() - INTERVAL '1 hour'", + [], + )) + .await?; + // A reservation that outlives the provider timeout may represent a streamed + // response whose final usage could not be persisted. Convert it to a durable + // debit instead of releasing it, otherwise disconnecting clients could bypass + // the monthly budget. Explicit upstream failures release their reservations. + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "UPDATE ai_gateway_cost_reservations SET is_conservative_debit = TRUE WHERE billing_period = $1 AND expires_at <= NOW() AND is_conservative_debit = FALSE", + [month_start.date_naive().into()], + )) + .await?; + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "DELETE FROM ai_gateway_cost_reservations WHERE billing_period < $1", + [month_start.date_naive().into()], + )) + .await?; + Ok(()) + } + + async fn check_rates_and_record( + &self, + txn: &DatabaseTransaction, + configs: &[AppliedConfig], + request_id: &str, + ) -> Result<(), AiGatewayError> { + for config in configs { + let Some(limit) = config.max_requests_per_minute else { + continue; + }; + let row = CountRow::find_by_statement(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + r#"SELECT + COUNT(*)::INT8 AS count, + GREATEST( + 1, + CEIL(EXTRACT(EPOCH FROM (MIN(occurred_at) + INTERVAL '60 seconds' - NOW()))) + )::INT8 AS retry_after_seconds + FROM ai_gateway_rate_events + WHERE scope = $1 AND occurred_at > NOW() - INTERVAL '60 seconds'"#, + [config.scope.clone().into()], + )) + .one(txn) + .await? + .unwrap_or(CountRow { + count: Some(0), + retry_after_seconds: Some(1), + }); + if row.count.unwrap_or(0) >= limit { + return Err(AiGatewayError::RateLimitExceeded { + scope: config.scope.clone(), + limit_per_minute: limit, + retry_after_seconds: row + .retry_after_seconds + .and_then(|seconds| u64::try_from(seconds).ok()) + .unwrap_or(1), + }); + } + } + + for config in configs { + if config.max_requests_per_minute.is_none() { + continue; + } + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "INSERT INTO ai_gateway_rate_events (request_id, scope) VALUES ($1::uuid, $2)", + [request_id.into(), config.scope.clone().into()], + )) + .await?; + } + Ok(()) + } + + async fn check_budgets_and_reserve( + &self, + txn: &DatabaseTransaction, + attribution: &AiUsageAttribution, + configs: &[AppliedConfig], + request_id: &str, + projected_cost: i64, + billing_period: chrono::NaiveDate, + ) -> Result<(), AiGatewayError> { + for config in configs { + let Some(limit) = config.max_cost_per_month_microcents else { + continue; + }; + let row = CostRow::find_by_statement(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + r#"SELECT ( + SELECT COALESCE(SUM(estimated_cost_microcents), 0)::INT8 + FROM ai_usage_logs + WHERE is_byok = FALSE + AND ((billing_period = $1 + AND timestamp >= $1::date + AND timestamp < $1::date + INTERVAL '1 month 6 minutes') + OR (billing_period IS NULL + AND timestamp >= $1::date + AND timestamp < $1::date + INTERVAL '1 month')) + AND ($2 = 'instance' + OR ($3::INT IS NOT NULL AND $2 LIKE 'project:%' AND project_id = $3) + OR ($4::INT IS NOT NULL AND $2 LIKE 'environment:%' AND environment_id = $4) + OR ($5::INT IS NOT NULL AND $2 LIKE 'token:%' AND deployment_token_id = $5)) + ) + ( + SELECT COALESCE(SUM(reserved_microcents), 0)::INT8 + FROM ai_gateway_cost_reservations + WHERE scope = $2 AND billing_period = $1 + ) AS cost"#, + [ + billing_period.into(), + config.scope.clone().into(), + attribution.project_id.into(), + attribution.environment_id.into(), + attribution.deployment_token_id.into(), + ], + )) + .one(txn) + .await? + .unwrap_or(CostRow { cost: Some(0) }); + let spent = row.cost.unwrap_or(0); + if spent.saturating_add(projected_cost) > limit { + return Err(AiGatewayError::MonthlyBudgetExceeded { + scope: config.scope.clone(), + spent_microcents: spent, + limit_microcents: limit, + }); + } + } + + for config in configs { + if config.max_cost_per_month_microcents.is_none() { + continue; + } + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "INSERT INTO ai_gateway_cost_reservations (request_id, scope, reserved_microcents, billing_period) VALUES ($1::uuid, $2, $3, $4)", + [ + request_id.into(), + config.scope.clone().into(), + projected_cost.into(), + billing_period.into(), + ], + )) + .await?; + } + Ok(()) + } +} + +fn current_month_start() -> Result, AiGatewayError> { + let now = Utc::now(); + Utc.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0) + .single() + .ok_or_else(|| AiGatewayError::Internal { + message: format!( + "Failed to calculate current billing month for {}-{}", + now.year(), + now.month() + ), + }) +} + +fn applicable_scope_names(attribution: &AiUsageAttribution) -> Vec { + let mut scopes = vec!["instance".to_string()]; + if let Some(project_id) = attribution.project_id { + scopes.push(format!("project:{project_id}")); + } + if let Some(environment_id) = attribution.environment_id { + scopes.push(format!("environment:{environment_id}")); + } + if let Some(token_id) = attribution.deployment_token_id { + scopes.push(format!("token:{token_id}")); + } + scopes +} + +fn validate_nonnegative( + scope: &str, + field: &'static str, + value: Option, +) -> Result<(), AiGatewayError> { + if let Some(value) = value { + if value < 0 { + return Err(AiGatewayError::InvalidGovernanceConfig { + scope: scope.to_string(), + field, + value, + }); + } + } + Ok(()) +} + +fn validate_scope(scope: &str) -> Result<(), AiGatewayError> { + if scope == "instance" { + return Ok(()); + } + let valid = ["project:", "environment:", "token:"].iter().any(|prefix| { + scope + .strip_prefix(prefix) + .and_then(|raw_id| raw_id.parse::().ok().map(|id| (raw_id, id))) + .is_some_and(|(raw_id, id)| id > 0 && raw_id == id.to_string()) + }); + if valid { + Ok(()) + } else { + Err(AiGatewayError::InvalidGovernanceScope { + scope: scope.to_string(), + }) + } +} + +fn validate_allowed_models( + scope: &str, + allowed_models: Option<&serde_json::Value>, +) -> Result<(), AiGatewayError> { + let Some(value) = allowed_models else { + return Ok(()); + }; + let valid = value.as_array().is_some_and(|models| { + models + .iter() + .all(|model| model.as_str().is_some_and(|model| !model.trim().is_empty())) + }); + if valid { + Ok(()) + } else { + Err(AiGatewayError::Validation { + message: format!( + "allowed_models for AI gateway scope '{}' must be an array of non-empty model IDs", + scope + ), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sea_orm::{DbErr, MockDatabase, MockExecResult}; + + fn config( + scope: &str, + models: Option, + rpm: Option, + budget: Option, + ) -> temps_entities::ai_gateway_config::Model { + temps_entities::ai_gateway_config::Model { + id: 1, + scope: scope.to_string(), + allowed_models: models, + max_requests_per_minute: rpm, + max_cost_per_month_microcents: budget, + created_at: Utc::now(), + updated_at: Utc::now(), + } + } + + #[test] + fn scope_names_are_ordered_from_broadest_to_most_specific() { + let attribution = AiUsageAttribution { + project_id: Some(7), + environment_id: Some(11), + deployment_token_id: Some(17), + ..Default::default() + }; + assert_eq!( + applicable_scope_names(&attribution), + vec!["instance", "project:7", "environment:11", "token:17"] + ); + } + + #[test] + fn scope_validation_requires_canonical_positive_ids() { + assert!(validate_scope("instance").is_ok()); + assert!(validate_scope("project:7").is_ok()); + assert!(validate_scope("environment:11").is_ok()); + assert!(validate_scope("token:17").is_ok()); + for scope in [ + "token:0", + "token:017", + "project:0", + "project:007", + "project:+7", + ] { + assert!(matches!( + validate_scope(scope), + Err(AiGatewayError::InvalidGovernanceScope { .. }) + )); + } + } + + #[test] + fn invalid_negative_limit_has_context() { + assert!(matches!( + validate_nonnegative("project:7", "max_requests_per_minute", Some(-1)), + Err(AiGatewayError::InvalidGovernanceConfig { + ref scope, + field: "max_requests_per_minute", + value: -1, + }) if scope == "project:7" + )); + } + + #[test] + fn allowed_model_validation_rejects_malformed_values() { + assert!(validate_allowed_models("instance", None).is_ok()); + assert!(validate_allowed_models("instance", Some(&serde_json::json!([]))).is_ok()); + assert!( + validate_allowed_models("project:7", Some(&serde_json::json!(["gpt-5-mini"]))).is_ok() + ); + assert!(validate_allowed_models("project:7", Some(&serde_json::json!([""]))).is_err()); + assert!(validate_allowed_models("project:7", Some(&serde_json::json!({}))).is_err()); + } + + #[tokio::test] + async fn model_allowlist_rejects_before_opening_transaction() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![config( + "project:7", + Some(serde_json::json!(["gpt-5-mini"])), + None, + None, + )]]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + let attribution = AiUsageAttribution { + project_id: Some(7), + ..Default::default() + }; + + assert!(matches!( + service + .check_request(&attribution, "claude-sonnet-4-6", false, Some(10), Some(10)) + .await, + Err(AiGatewayError::ModelNotAllowed { ref scope, .. }) if scope == "project:7" + )); + } + + #[tokio::test] + async fn budget_requires_bounded_output_before_opening_transaction() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![config("project:7", None, None, Some(100))]]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + + assert!(matches!( + service + .check_request( + &AiUsageAttribution { project_id: Some(7), ..Default::default() }, + "gpt-5-mini", + false, + Some(10), + None, + ) + .await, + Err(AiGatewayError::BudgetRequiresMaxTokens { ref scope }) if scope == "project:7" + )); + } + + #[tokio::test] + async fn budget_rejects_input_without_safe_cost_projection() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![config("project:7", None, None, Some(100))]]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + + assert!(matches!( + service + .check_request( + &AiUsageAttribution { + project_id: Some(7), + ..Default::default() + }, + "gpt-5-mini", + false, + None, + Some(10), + ) + .await, + Err(AiGatewayError::BudgetProjectionUnavailable { ref scope }) + if scope == "project:7" + )); + } + + #[tokio::test] + async fn byok_skips_operator_budget_reservation() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results(vec![vec![config("project:7", None, None, Some(1))]]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + let result = service + .check_request( + &AiUsageAttribution { + project_id: Some(7), + ..Default::default() + }, + "custom-model", + true, + Some(10), + None, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn upsert_config_creates_valid_scope() { + let expected = config("project:7", None, Some(30), Some(500)); + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results([ + Vec::::new(), + vec![expected.clone()], + ]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + + let created = service + .upsert_config("project:7", None, Some(30), Some(500)) + .await + .expect("valid project config should be created"); + assert_eq!(created.scope, expected.scope); + } + + #[tokio::test] + async fn delete_config_returns_typed_not_found() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_exec_results([MockExecResult { + last_insert_id: 0, + rows_affected: 0, + }]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + assert!(matches!( + service.delete_config("project:7").await, + Err(AiGatewayError::GovernanceConfigNotFound { ref scope }) if scope == "project:7" + )); + } + + #[tokio::test] + async fn check_request_propagates_database_error() { + let db = MockDatabase::new(DatabaseBackend::Postgres) + .append_query_errors([DbErr::Custom("governance database unavailable".to_string())]) + .into_connection(); + let service = GovernanceService::new(Arc::new(db)); + assert!(matches!( + service + .check_request( + &AiUsageAttribution::default(), + "gpt-5-mini", + false, + Some(10), + Some(10), + ) + .await, + Err(AiGatewayError::Database(DbErr::Custom(ref message))) + if message == "governance database unavailable" + )); + } +} diff --git a/crates/temps-ai-gateway/src/services/mod.rs b/crates/temps-ai-gateway/src/services/mod.rs index 1c7af9c16..897d9201f 100644 --- a/crates/temps-ai-gateway/src/services/mod.rs +++ b/crates/temps-ai-gateway/src/services/mod.rs @@ -1,9 +1,11 @@ pub mod ai_service; pub mod gateway_service; +pub mod governance_service; pub mod provider_key_service; pub mod usage_service; pub use ai_service::GatewayAiService; pub use gateway_service::{ByokOverride, CredentialType, GatewayService}; +pub use governance_service::{GovernanceReservation, GovernanceService}; pub use provider_key_service::ProviderKeyService; -pub use usage_service::{AiRequestContext, UsageFilter, UsageService}; +pub use usage_service::{AiRequestContext, AiUsageAttribution, UsageFilter, UsageService}; diff --git a/crates/temps-ai-gateway/src/services/usage_service.rs b/crates/temps-ai-gateway/src/services/usage_service.rs index f56187f89..aa605c9de 100644 --- a/crates/temps-ai-gateway/src/services/usage_service.rs +++ b/crates/temps-ai-gateway/src/services/usage_service.rs @@ -1,12 +1,14 @@ use chrono::{DateTime, Utc}; use sea_orm::{ - ActiveModelTrait, DatabaseBackend, DatabaseConnection, FromQueryResult, Set, Statement, + ActiveModelTrait, ConnectionTrait, DatabaseBackend, DatabaseConnection, FromQueryResult, Set, + Statement, TransactionTrait, }; use serde::{Deserialize, Serialize}; use std::sync::Arc; use temps_entities::ai_usage_logs; use utoipa::ToSchema; +use super::GovernanceReservation; use crate::error::AiGatewayError; // ============================================================================ @@ -71,6 +73,16 @@ pub struct UsageLogEntry { pub is_streaming: bool, pub is_byok: bool, #[serde(skip_serializing_if = "Option::is_none")] + pub user_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub project_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub environment_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub deployment_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub deployment_token_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub conversation_id: Option, pub tags: Vec, #[serde(skip_serializing_if = "Option::is_none")] @@ -112,6 +124,17 @@ pub struct AiRequestContext { pub trace_id: Option, } +/// Stable ownership dimensions attached by trusted authentication context. +/// These values never come from caller-controlled request headers. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AiUsageAttribution { + pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, + pub deployment_id: Option, + pub deployment_token_id: Option, +} + /// Filters for querying AI usage data. /// /// Cost bounds are expressed in microcents (the unit stored in @@ -121,6 +144,10 @@ pub struct AiRequestContext { #[derive(Debug, Clone, Default, Deserialize, ToSchema)] pub struct UsageFilter { pub user_id: Option, + pub project_id: Option, + pub environment_id: Option, + pub deployment_id: Option, + pub deployment_token_id: Option, pub conversation_id: Option, /// Comma-separated tags to filter by (AND logic). pub tags: Option, @@ -206,6 +233,11 @@ struct UsageLogRow { status: Option, is_streaming: Option, is_byok: Option, + user_id: Option, + project_id: Option, + environment_id: Option, + deployment_id: Option, + deployment_token_id: Option, conversation_id: Option, tags: Option, request_id: Option, @@ -259,7 +291,10 @@ impl UsageService { is_byok: bool, ) -> Result<(), AiGatewayError> { self.log_usage_with_context( - user_id, + &AiUsageAttribution { + user_id, + ..Default::default() + }, provider, model, input_tokens, @@ -277,7 +312,42 @@ impl UsageService { #[allow(clippy::too_many_arguments)] pub async fn log_usage_with_context( &self, - user_id: Option, + attribution: &AiUsageAttribution, + provider: &str, + model: &str, + input_tokens: i64, + output_tokens: i64, + latency_ms: i32, + estimated_cost_microcents: i64, + status: i16, + is_streaming: bool, + is_byok: bool, + context: &AiRequestContext, + ) -> Result<(), AiGatewayError> { + self.log_usage_with_context_and_reservation( + attribution, + provider, + model, + input_tokens, + output_tokens, + latency_ms, + estimated_cost_microcents, + status, + is_streaming, + is_byok, + context, + None, + ) + .await + } + + /// Atomically converts a conservative governance reservation into actual + /// usage. If persistence fails, the reservation remains in place so a + /// transient database error cannot open a quota bypass. + #[allow(clippy::too_many_arguments)] + pub async fn log_usage_with_context_and_reservation( + &self, + attribution: &AiUsageAttribution, provider: &str, model: &str, input_tokens: i64, @@ -288,10 +358,20 @@ impl UsageService { is_streaming: bool, is_byok: bool, context: &AiRequestContext, + reservation: Option<&GovernanceReservation>, ) -> Result<(), AiGatewayError> { let record = ai_usage_logs::ActiveModel { timestamp: Set(chrono::Utc::now()), - user_id: Set(user_id), + user_id: Set(attribution.user_id), + project_id: Set(attribution.project_id), + environment_id: Set(attribution.environment_id), + deployment_id: Set(attribution.deployment_id), + deployment_token_id: Set(attribution.deployment_token_id), + billing_period: Set(Some( + reservation + .map(GovernanceReservation::billing_period) + .unwrap_or_else(current_billing_period), + )), provider: Set(provider.to_string()), model: Set(model.to_string()), input_tokens: Set(input_tokens), @@ -308,7 +388,19 @@ impl UsageService { ..Default::default() }; - record.insert(self.db.as_ref()).await?; + if let Some(reservation) = reservation { + let txn = self.db.begin().await?; + record.insert(&txn).await?; + txn.execute(Statement::from_sql_and_values( + DatabaseBackend::Postgres, + "DELETE FROM ai_gateway_cost_reservations WHERE request_id = $1::uuid", + [reservation.request_id().into()], + )) + .await?; + txn.commit().await?; + } else { + record.insert(self.db.as_ref()).await?; + } Ok(()) } @@ -630,6 +722,11 @@ impl UsageService { status, is_streaming, is_byok, + user_id, + project_id, + environment_id, + deployment_id, + deployment_token_id, conversation_id, array_to_string(tags, ',') as tags, request_id, @@ -738,6 +835,11 @@ impl UsageService { status, is_streaming, is_byok, + user_id, + project_id, + environment_id, + deployment_id, + deployment_token_id, conversation_id, array_to_string(tags, ',') as tags, request_id, @@ -772,6 +874,30 @@ impl UsageService { param_idx += 1; } + if let Some(project_id) = filter.project_id { + conditions.push(format!("project_id = ${}", param_idx)); + values.push(project_id.into()); + param_idx += 1; + } + + if let Some(environment_id) = filter.environment_id { + conditions.push(format!("environment_id = ${}", param_idx)); + values.push(environment_id.into()); + param_idx += 1; + } + + if let Some(deployment_id) = filter.deployment_id { + conditions.push(format!("deployment_id = ${}", param_idx)); + values.push(deployment_id.into()); + param_idx += 1; + } + + if let Some(deployment_token_id) = filter.deployment_token_id { + conditions.push(format!("deployment_token_id = ${}", param_idx)); + values.push(deployment_token_id.into()); + param_idx += 1; + } + if let Some(ref conv_id) = filter.conversation_id { conditions.push(format!("conversation_id = ${}", param_idx)); values.push(conv_id.clone().into()); @@ -862,6 +988,13 @@ impl UsageService { } } +fn current_billing_period() -> chrono::NaiveDate { + use chrono::Datelike; + + let today = Utc::now().date_naive(); + chrono::NaiveDate::from_ymd_opt(today.year(), today.month(), 1).unwrap_or(today) +} + /// Shift all `$N` parameter placeholders in a SQL fragment by `offset`. fn shift_params(sql: &str, offset: usize) -> String { let mut result = String::with_capacity(sql.len()); @@ -904,6 +1037,11 @@ fn usage_log_from_row(r: UsageLogRow) -> UsageLogEntry { status: r.status.unwrap_or(0), is_streaming: r.is_streaming.unwrap_or(false), is_byok: r.is_byok.unwrap_or(false), + user_id: r.user_id, + project_id: r.project_id, + environment_id: r.environment_id, + deployment_id: r.deployment_id, + deployment_token_id: r.deployment_token_id, conversation_id: r.conversation_id, tags: r .tags @@ -920,6 +1058,7 @@ fn usage_log_from_row(r: UsageLogRow) -> UsageLogEntry { #[cfg(test)] mod tests { use super::*; + use sea_orm::{DatabaseBackend, MockDatabase}; #[test] fn test_shift_params_basic() { @@ -957,6 +1096,74 @@ mod tests { assert_eq!(values.len(), 2); } + #[tokio::test] + async fn log_usage_persists_deployment_attribution() { + let db = Arc::new( + MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results([vec![ai_usage_logs::Model { + id: 1, + timestamp: Utc::now(), + user_id: None, + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), + billing_period: Some( + chrono::NaiveDate::from_ymd_opt(2026, 8, 1) + .expect("valid test billing period"), + ), + provider: "openai".to_string(), + model: "gpt-5-mini".to_string(), + input_tokens: 100, + output_tokens: 50, + latency_ms: 25, + estimated_cost_microcents: 120, + status: 200, + is_streaming: false, + is_byok: false, + conversation_id: None, + tags: Vec::new(), + request_id: None, + trace_id: None, + }]]) + .into_connection(), + ); + let service = UsageService::new(db.clone()); + let attribution = AiUsageAttribution { + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), + ..Default::default() + }; + + service + .log_usage_with_context( + &attribution, + "openai", + "gpt-5-mini", + 100, + 50, + 25, + 120, + 200, + false, + false, + &AiRequestContext::default(), + ) + .await + .expect("scoped usage should be inserted"); + + drop(service); + let db = Arc::try_unwrap(db).expect("usage service should release database"); + let log = db.into_transaction_log(); + let sql = log[0].statements()[0].sql.to_lowercase(); + assert!(sql.contains("project_id")); + assert!(sql.contains("environment_id")); + assert!(sql.contains("deployment_id")); + assert!(sql.contains("deployment_token_id")); + } + #[test] fn test_build_filter_clause_with_user_id() { let db = sea_orm::DatabaseConnection::Disconnected; @@ -973,6 +1180,27 @@ mod tests { assert_eq!(values.len(), 3); } + #[test] + fn test_build_filter_clause_with_deployment_scope() { + let service = UsageService::new(Arc::new(sea_orm::DatabaseConnection::Disconnected)); + let from = Utc::now() - chrono::Duration::hours(1); + let to = Utc::now(); + let filter = UsageFilter { + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), + ..Default::default() + }; + + let (clause, values) = service.build_filter_clause(from, to, &filter); + assert!(clause.contains("project_id = $3")); + assert!(clause.contains("environment_id = $4")); + assert!(clause.contains("deployment_id = $5")); + assert!(clause.contains("deployment_token_id = $6")); + assert_eq!(values.len(), 6); + } + #[test] fn test_build_filter_clause_with_tags() { let db = sea_orm::DatabaseConnection::Disconnected; @@ -1069,6 +1297,11 @@ mod tests { status: Some(200), is_streaming: Some(false), is_byok: Some(false), + user_id: None, + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), conversation_id: Some("conv_123".to_string()), tags: Some("agent:support,env:prod".to_string()), request_id: Some("req_abc".to_string()), @@ -1080,6 +1313,10 @@ mod tests { assert_eq!(entry.tags, vec!["agent:support", "env:prod"]); assert_eq!(entry.request_id, Some("req_abc".to_string())); assert_eq!(entry.trace_id, Some("trace_xyz".to_string())); + assert_eq!(entry.project_id, Some(7)); + assert_eq!(entry.environment_id, Some(11)); + assert_eq!(entry.deployment_id, Some(13)); + assert_eq!(entry.deployment_token_id, Some(17)); } #[test] @@ -1096,6 +1333,11 @@ mod tests { status: Some(200), is_streaming: Some(true), is_byok: Some(true), + user_id: Some(3), + project_id: None, + environment_id: None, + deployment_id: None, + deployment_token_id: None, conversation_id: None, tags: Some("".to_string()), request_id: None, @@ -1122,6 +1364,10 @@ mod tests { fn test_usage_filter_default() { let filter = UsageFilter::default(); assert!(filter.user_id.is_none()); + assert!(filter.project_id.is_none()); + assert!(filter.environment_id.is_none()); + assert!(filter.deployment_id.is_none()); + assert!(filter.deployment_token_id.is_none()); assert!(filter.conversation_id.is_none()); assert!(filter.tags.is_none()); assert!(filter.model.is_none()); @@ -1142,6 +1388,11 @@ mod tests { status: 200, is_streaming: false, is_byok: false, + user_id: None, + project_id: None, + environment_id: None, + deployment_id: None, + deployment_token_id: None, conversation_id: None, tags: vec![], request_id: None, @@ -1168,6 +1419,11 @@ mod tests { status: 200, is_streaming: false, is_byok: false, + user_id: None, + project_id: Some(7), + environment_id: Some(11), + deployment_id: Some(13), + deployment_token_id: Some(17), conversation_id: Some("conv_abc".to_string()), tags: vec!["agent:support".to_string()], request_id: Some("req_123".to_string()), @@ -1179,6 +1435,8 @@ mod tests { assert!(json.contains("agent:support")); assert!(json.contains("req_123")); assert!(json.contains("trace_456")); + assert!(json.contains("\"project_id\":7")); + assert!(json.contains("\"environment_id\":11")); } #[test] diff --git a/crates/temps-ai-gateway/tests/governance_integration.rs b/crates/temps-ai-gateway/tests/governance_integration.rs new file mode 100644 index 000000000..c59926301 --- /dev/null +++ b/crates/temps-ai-gateway/tests/governance_integration.rs @@ -0,0 +1,251 @@ +use std::sync::Arc; + +use sea_orm::{ConnectionTrait, DatabaseBackend, Statement}; +use temps_ai_gateway::error::AiGatewayError; +use temps_ai_gateway::services::{ + AiRequestContext, AiUsageAttribution, GovernanceService, UsageService, +}; +use temps_database::test_utils::TestDatabase; + +async fn setup() -> Option<(TestDatabase, Arc)> { + let test_db = match TestDatabase::with_migrations().await { + Ok(test_db) => test_db, + Err(error) + if temps_database::test_utils::is_container_runtime_unavailable(&error.to_string()) => + { + eprintln!( + "Docker/Postgres unavailable, skipping AI governance integration test: {error}" + ); + return None; + } + Err(error) => panic!("AI governance test database setup failed: {error}"), + }; + let service = Arc::new(GovernanceService::new(test_db.connection_arc())); + Some((test_db, service)) +} + +async fn scalar_count(db: &sea_orm::DatabaseConnection, table: &str) -> i64 { + let statement = Statement::from_string( + DatabaseBackend::Postgres, + format!("SELECT COUNT(*)::INT8 AS count FROM {table}"), + ); + db.query_one(statement) + .await + .expect("count query should succeed") + .expect("count query should return a row") + .try_get("", "count") + .expect("count should be an i64") +} + +fn project_attribution() -> AiUsageAttribution { + AiUsageAttribution { + project_id: Some(7), + environment_id: Some(11), + deployment_token_id: Some(17), + ..Default::default() + } +} + +#[tokio::test] +async fn concurrent_project_budget_checks_allow_only_one_reservation() { + let Some((_test_db, service)) = setup().await else { + return; + }; + service + .upsert_config("project:7", None, None, Some(2_000)) + .await + .expect("project budget should be created"); + + let first_service = service.clone(); + let second_service = service.clone(); + let first = tokio::spawn(async move { + first_service + .check_request( + &project_attribution(), + "gpt-5-mini", + false, + Some(10), + Some(10), + ) + .await + }); + let second = tokio::spawn(async move { + second_service + .check_request( + &project_attribution(), + "gpt-5-mini", + false, + Some(10), + Some(10), + ) + .await + }); + let (first, second) = tokio::join!(first, second); + let results = [ + first.expect("first budget task should complete"), + second.expect("second budget task should complete"), + ]; + + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); + assert_eq!( + results + .iter() + .filter(|result| matches!(result, Err(AiGatewayError::MonthlyBudgetExceeded { .. }))) + .count(), + 1 + ); +} + +#[tokio::test] +async fn token_rate_limit_is_shared_between_service_instances() { + let Some((test_db, first_service)) = setup().await else { + return; + }; + first_service + .upsert_config("token:17", None, Some(1), None) + .await + .expect("deployment-token rate limit should be created"); + let second_service = Arc::new(GovernanceService::new(test_db.connection_arc())); + + let first = tokio::spawn(async move { + first_service + .check_request(&project_attribution(), "custom-model", true, Some(10), None) + .await + }); + let second = tokio::spawn(async move { + second_service + .check_request(&project_attribution(), "custom-model", true, Some(10), None) + .await + }); + let (first, second) = tokio::join!(first, second); + let results = [ + first.expect("first rate task should complete"), + second.expect("second rate task should complete"), + ]; + + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); + assert_eq!( + results + .iter() + .filter(|result| matches!(result, Err(AiGatewayError::RateLimitExceeded { .. }))) + .count(), + 1 + ); +} + +#[tokio::test] +async fn finalized_usage_atomically_replaces_cost_reservation() { + let Some((test_db, governance)) = setup().await else { + return; + }; + governance + .upsert_config("project:7", None, None, Some(10_000)) + .await + .expect("project budget should be created"); + let attribution = project_attribution(); + let reservation = governance + .check_request(&attribution, "gpt-5-mini", false, Some(10), Some(10)) + .await + .expect("request should reserve project budget"); + assert_eq!( + scalar_count(test_db.db.as_ref(), "ai_gateway_cost_reservations").await, + 1 + ); + + let usage = UsageService::new(test_db.connection_arc()); + usage + .log_usage_with_context_and_reservation( + &attribution, + "openai", + "gpt-5-mini", + 10, + 5, + 20, + 1_200, + 200, + false, + false, + &AiRequestContext::default(), + Some(&reservation), + ) + .await + .expect("usage should atomically replace its reservation"); + + assert_eq!( + scalar_count(test_db.db.as_ref(), "ai_gateway_cost_reservations").await, + 0 + ); + let usage_row = test_db + .db + .query_one(Statement::from_string( + DatabaseBackend::Postgres, + "SELECT estimated_cost_microcents, billing_period FROM ai_usage_logs".to_string(), + )) + .await + .expect("usage query should succeed") + .expect("finalized usage row should exist"); + assert_eq!( + usage_row + .try_get::("", "estimated_cost_microcents") + .expect("stored cost should be an i64"), + 1_200 + ); + assert!(usage_row + .try_get::("", "billing_period") + .is_ok()); +} + +#[tokio::test] +async fn abandoned_cost_reservation_becomes_a_durable_conservative_debit() { + let Some((test_db, governance)) = setup().await else { + return; + }; + governance + .upsert_config("project:7", None, None, Some(2_000)) + .await + .expect("project budget should be created"); + governance + .check_request( + &project_attribution(), + "gpt-5-mini", + false, + Some(10), + Some(10), + ) + .await + .expect("first request should reserve the complete budget"); + test_db + .db + .execute(Statement::from_string( + DatabaseBackend::Postgres, + "UPDATE ai_gateway_cost_reservations SET expires_at = NOW() - INTERVAL '1 second'" + .to_string(), + )) + .await + .expect("test should expire the abandoned reservation"); + + assert!(matches!( + governance + .check_request( + &project_attribution(), + "gpt-5-mini", + false, + Some(10), + Some(10), + ) + .await, + Err(AiGatewayError::MonthlyBudgetExceeded { .. }) + )); + let reservation = test_db + .db + .query_one(Statement::from_string( + DatabaseBackend::Postgres, + "SELECT is_conservative_debit FROM ai_gateway_cost_reservations".to_string(), + )) + .await + .expect("reservation query should succeed") + .expect("conservative debit should remain for the billing period"); + assert!(reservation + .try_get::("", "is_conservative_debit") + .expect("debit marker should be a boolean")); +} diff --git a/crates/temps-entities/src/ai_gateway_config.rs b/crates/temps-entities/src/ai_gateway_config.rs index 64cd98c0d..cec4e3a18 100644 --- a/crates/temps-entities/src/ai_gateway_config.rs +++ b/crates/temps-entities/src/ai_gateway_config.rs @@ -9,7 +9,7 @@ use temps_core::DBDateTime; pub struct Model { #[sea_orm(primary_key)] pub id: i32, - /// Scope: "instance", "project:{id}", "environment:{id}" + /// Scope: "instance", "project:{id}", "environment:{id}", "token:{id}" pub scope: String, /// JSON array of allowed model IDs, NULL means all models allowed pub allowed_models: Option, diff --git a/crates/temps-entities/src/ai_usage_logs.rs b/crates/temps-entities/src/ai_usage_logs.rs index 90dfc4117..2d8b615f4 100644 --- a/crates/temps-entities/src/ai_usage_logs.rs +++ b/crates/temps-entities/src/ai_usage_logs.rs @@ -9,12 +9,22 @@ pub struct Model { pub id: i64, pub timestamp: DBDateTime, pub user_id: Option, + /// Project charged for a deployment-token request. + pub project_id: Option, + /// Optional environment scope carried by the deployment token. + pub environment_id: Option, + /// Optional deployment scope carried by the deployment token. + pub deployment_id: Option, + /// Deployment-token row used for the request. Kept after token deletion. + pub deployment_token_id: Option, + /// UTC calendar month in which governance reserved this request's cost. + pub billing_period: Option, pub provider: String, pub model: String, pub input_tokens: i64, pub output_tokens: i64, pub latency_ms: i32, - /// Estimated cost in microcents (1/10000 of a cent) + /// Estimated cost in microcents (1/1,000,000 of a cent) pub estimated_cost_microcents: i64, /// HTTP status code returned to the caller pub status: i16, diff --git a/crates/temps-migrations/src/migration/m20260803_000001_add_ai_usage_scope.rs b/crates/temps-migrations/src/migration/m20260803_000001_add_ai_usage_scope.rs new file mode 100644 index 000000000..9db56606a --- /dev/null +++ b/crates/temps-migrations/src/migration/m20260803_000001_add_ai_usage_scope.rs @@ -0,0 +1,94 @@ +use sea_orm_migration::prelude::*; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .alter_table( + Table::alter() + .table(AiUsageLogs::Table) + .add_column(ColumnDef::new(AiUsageLogs::ProjectId).integer()) + .add_column(ColumnDef::new(AiUsageLogs::EnvironmentId).integer()) + .add_column(ColumnDef::new(AiUsageLogs::DeploymentId).integer()) + .add_column(ColumnDef::new(AiUsageLogs::DeploymentTokenId).integer()) + .add_column(ColumnDef::new(AiUsageLogs::BillingPeriod).date()) + .to_owned(), + ) + .await?; + + // Keep identifiers denormalized so deleting a deployment token or + // deployment does not erase historical cost attribution. The existing + // timestamp index bounds usage and governance queries. Avoid building + // new indexes on this populated hot table during startup migration; + // those can be added later through the decoupled migration path. + let db = manager.get_connection(); + db.execute_unprepared( + r#"CREATE TABLE ai_gateway_rate_events ( + request_id UUID NOT NULL, + scope VARCHAR(255) NOT NULL, + occurred_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (request_id, scope) + )"#, + ) + .await?; + db.execute_unprepared( + "CREATE INDEX idx_ai_gateway_rate_events_scope_time ON ai_gateway_rate_events (scope, occurred_at DESC)", + ) + .await?; + + db.execute_unprepared( + r#"CREATE TABLE ai_gateway_cost_reservations ( + request_id UUID NOT NULL, + scope VARCHAR(255) NOT NULL, + reserved_microcents BIGINT NOT NULL CHECK (reserved_microcents >= 0), + billing_period DATE NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + INTERVAL '6 minutes', + is_conservative_debit BOOLEAN NOT NULL DEFAULT FALSE, + PRIMARY KEY (request_id, scope) + )"#, + ) + .await?; + db.execute_unprepared( + "CREATE INDEX idx_ai_gateway_cost_reservations_scope_period_expiry ON ai_gateway_cost_reservations (scope, billing_period, expires_at)", + ) + .await?; + + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + let db = manager.get_connection(); + db.execute_unprepared("DROP TABLE IF EXISTS ai_gateway_cost_reservations") + .await?; + db.execute_unprepared("DROP TABLE IF EXISTS ai_gateway_rate_events") + .await?; + manager + .alter_table( + Table::alter() + .table(AiUsageLogs::Table) + .drop_column(AiUsageLogs::DeploymentTokenId) + .drop_column(AiUsageLogs::DeploymentId) + .drop_column(AiUsageLogs::EnvironmentId) + .drop_column(AiUsageLogs::ProjectId) + .drop_column(AiUsageLogs::BillingPeriod) + .to_owned(), + ) + .await?; + + Ok(()) + } +} + +#[derive(DeriveIden)] +enum AiUsageLogs { + Table, + ProjectId, + EnvironmentId, + DeploymentId, + DeploymentTokenId, + BillingPeriod, +} diff --git a/crates/temps-migrations/src/migration/mod.rs b/crates/temps-migrations/src/migration/mod.rs index 96602e424..271029007 100644 --- a/crates/temps-migrations/src/migration/mod.rs +++ b/crates/temps-migrations/src/migration/mod.rs @@ -166,6 +166,7 @@ mod m20260725_000001_sandboxes_agent_run_link; mod m20260728_000001_add_environment_id_to_metric_alert_rules; mod m20260730_000001_add_architecture_to_nodes; mod m20260802_000001_add_environment_force_https; +mod m20260803_000001_add_ai_usage_scope; pub struct Migrator; @@ -339,6 +340,7 @@ impl MigratorTrait for Migrator { ), Box::new(m20260730_000001_add_architecture_to_nodes::Migration), Box::new(m20260802_000001_add_environment_force_https::Migration), + Box::new(m20260803_000001_add_ai_usage_scope::Migration), ] } }