diff --git a/README.md b/README.md index ec0952536..e15c65047 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ [![GHCR](https://img.shields.io/badge/GHCR-codex--proxy--rs-2496ED?logo=docker&logoColor=white&style=flat-square)](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs) [![MIT](https://img.shields.io/badge/License-MIT-blue.svg?style=flat-square)](https://opensource.org/license/mit) -[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署文档](deploy/README.md) · [架构](docs/architecture.md) +[快速开始](#快速开始) · [客户端接入](#客户端接入) · [运维](#运维) · [部署文档](deploy/README.md) · [架构](docs/app-layer-refactor-architecture.md) @@ -21,14 +21,14 @@ ## 能力 -| 领域 | 实现 | -| -------- | ------------------------------------------------ | -| 协议 | OpenAI Responses JSON/SSE 与模型目录 | -| Provider | Codex OAuth、xAI/Grok OAuth session | -| 路由 | Provider instance、能力过滤、显式 fallback 与多 credential | -| 延续 | Provider 原生 continuation 与加密 portable transcript | -| 管理 | Client Key、Provider、模型、Route、Target、价格与 OAuth | -| 计量 | Request/Attempt usage、版本价格、预算、聚合与恢复 | +| 领域 | 实现 | +| -------- | ------------------------------------------------------------ | +| 协议 | OpenAI Responses JSON、SSE、WebSocket 与模型目录 | +| Provider | OpenAI OAuth、xAI OAuth，多账号隔离与自动刷新 | +| 路由 | Client Key 平台绑定、Provider instance、模型映射与安全 fallback | +| 延续 | 客户端携带历史；Provider 原生 continuation 固定账号与连接 | +| 管理 | Client Key、账号、模型目录、设置、观测、系统与 OAuth | +| 计量 | 模型请求 Token、费用、延迟、账号与 Provider 归因 | ## 快速开始 @@ -47,19 +47,21 @@ sudo chown "$(id -u):10001" deploy/config.yaml chmod 0640 deploy/config.yaml ``` -分别执行三次： +生成并导出两个基础设施密码： ```bash -openssl rand -hex 24 +export CPR_DATABASE_PASSWORD="$(openssl rand -hex 24)" +export CPR_REDIS_PASSWORD="$(openssl rand -hex 24)" ``` -将结果写入 `deploy/config.yaml`： +再生成一次管理员初始密码并写入 `deploy/config.yaml`： -| 配置 | 约束 | -| ------------------------------ | ------------------------ | -| `x-cpr.database.password` | 48 位十六进制 | -| `x-cpr.redis.password` | 48 位十六进制 | -| `x-cpr.admin.default_password` | 至少 12 位，不能包含 `$` | +| 配置 | 约束 | +| ------------------------ | ------------------------ | +| `admin.default_password` | 至少 12 位，不能包含 `$` | + +`CPR_DATABASE_PASSWORD` 与 `CPR_REDIS_PASSWORD` 必须是 48 位十六进制，并在每次 Compose +命令的 shell 中可用。 Linux 需要允许容器组写入运行目录： @@ -83,9 +85,9 @@ curl -i http://127.0.0.1:8080/healthz ### 3. 初始化 1. 使用 `admin@cpr.local` 与初始密码登录。 -2. 创建 Provider instance，并通过 OAuth 导入 Codex 或 xAI credential。 -3. 配置 Provider model、对外 Route/Target；严格预算场景同时发布价格版本。 -4. 创建 `sk_...` 客户端 Key，并设置模型 allowlist、速率、并发和预算策略。 +2. 创建 Provider instance，并通过 OAuth 或 JSON 导入 OpenAI、xAI 账号。 +3. 按平台配置客户端模型到上游模型的映射；未配置的模型名原样透传。 +4. 创建 `sk_...` 客户端 Key，绑定平台并设置模型范围、速率与并发限制。 > [!IMPORTANT] > xAI 使用 OAuth session，不支持把 xAI API Key 作为上游 credential。 @@ -123,6 +125,8 @@ curl http://127.0.0.1:8080/v1/responses \ | 路由 | 用途 | | --------------------------- | ------------------- | | `POST /v1/responses` | JSON 或 SSE 透明代理 | +| `GET /v1/responses` | WebSocket 透明代理 | +| `POST /v1/responses/review` | OpenAI review 请求 | | `GET /v1/models` | 启用的公开模型列表 | | `GET /v1/models/{model_id}` | 公开模型详情 | @@ -146,7 +150,7 @@ docker compose -f deploy/compose.yaml up -d ``` > [!IMPORTANT] -> `.runtime/` 保存数据库、Redis、身份密钥和日志。删除该目录会永久清除运行状态。 +> `.runtime/` 保存 PostgreSQL、Redis、自更新状态和日志。删除该目录会永久清除运行状态。 Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理，转发 WebSocket upgrade 与真实客户端 IP；不要暴露 PostgreSQL 或 Redis。 @@ -154,8 +158,8 @@ upgrade 与真实客户端 IP；不要暴露 PostgreSQL 或 Redis。 ## 文档 - [部署](deploy/README.md) -- [架构](docs/architecture.md) -- [多 Provider 目标架构](docs/multi-provider-architecture.md) -- [终态数据模型](docs/multi-provider-database.md) +- [终态代码架构](docs/app-layer-refactor-architecture.md) +- [终态数据模型](docs/multi-provider-final-database.md) +- [重构执行报告](docs/app-layer-refactor-execution-report.md) - [Release](https://github.com/zyycn/codex-proxy-rs/releases) - [容器镜像](https://github.com/zyycn/codex-proxy-rs/pkgs/container/codex-proxy-rs) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 4fb1f3f45..d8b3f21e6 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -200,7 +200,6 @@ dependencies = [ "tower", "tower-layer", "tower-service", - "tracing", ] [[package]] @@ -219,7 +218,6 @@ dependencies = [ "sync_wrapper", "tower-layer", "tower-service", - "tracing", ] [[package]] @@ -383,52 +381,19 @@ checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" name = "codex-proxy-rs" version = "0.0.0" dependencies = [ - "argon2", - "async-trait", - "axum", - "base64", - "bytes", - "chrono", "config", - "flate2", - "futures", + "gateway-admin", "gateway-api", "gateway-core", - "gateway-protocol", + "gateway-host", "gateway-store", - "hex", - "hmac 0.13.0", - "ipnet", "provider-openai", "provider-xai", - "rand 0.10.2", - "rand_core 0.6.4", - "redis", - "reqwest", - "ring", - "secrecy", - "semver", "serde", - "serde_json", - "sha2 0.11.0", - "sqlx", - "subtle", "syn", - "tar", "tempfile", "thiserror", "tokio", - "tokio-tungstenite 0.29.0", - "tokio-util", - "tower", - "tower-http", - "tracing", - "tracing-appender", - "tracing-subscriber", - "url", - "uuid", - "wiremock", - "zeroize", ] [[package]] @@ -1077,26 +1042,47 @@ dependencies = [ "slab", ] +[[package]] +name = "gateway-admin" +version = "0.0.0" +dependencies = [ + "argon2", + "async-trait", + "base64", + "chrono", + "futures", + "gateway-core", + "hex", + "rand_core 0.6.4", + "secrecy", + "serde", + "subtle", + "thiserror", + "tokio", + "uuid", +] + [[package]] name = "gateway-api" version = "0.1.0" dependencies = [ - "argon2", "async-trait", "axum", "base64", "chrono", "futures", + "gateway-admin", "gateway-core", "gateway-protocol", - "rand_core 0.6.4", "serde", "serde_json", - "subtle", "thiserror", "tokio", "tower", + "tower-http", "tracing", + "url", + "uuid", ] [[package]] @@ -1113,6 +1099,37 @@ dependencies = [ "uuid", ] +[[package]] +name = "gateway-host" +version = "0.0.0" +dependencies = [ + "async-trait", + "axum", + "chrono", + "config", + "filetime", + "flate2", + "futures", + "gateway-admin", + "gateway-core", + "hex", + "reqwest", + "semver", + "serde", + "serde_json", + "sha2 0.11.0", + "tar", + "tempfile", + "thiserror", + "tokio", + "tower", + "tracing", + "tracing-appender", + "tracing-subscriber", + "uuid", + "wiremock", +] + [[package]] name = "gateway-protocol" version = "0.0.0" @@ -1130,6 +1147,7 @@ dependencies = [ "async-trait", "chrono", "futures", + "gateway-admin", "gateway-core", "hex", "redis", @@ -1139,6 +1157,8 @@ dependencies = [ "sqlx", "thiserror", "tokio", + "tracing", + "url", "uuid", ] @@ -2029,17 +2049,20 @@ dependencies = [ "chrono", "cookie", "futures", + "gateway-admin", "gateway-core", "gateway-protocol", "getrandom 0.3.4", "hex", "jsonwebtoken", "reqwest", + "roxmltree", "rsa", "rustls", "rustls-native-certs", "rustls-pki-types", "secrecy", + "semver", "serde", "serde_json", "sha2 0.11.0", @@ -2066,6 +2089,7 @@ dependencies = [ "base64", "chrono", "futures", + "gateway-admin", "gateway-core", "gateway-protocol", "getrandom 0.3.4", @@ -2395,6 +2419,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "roxmltree" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1964b10c76125c36f8afe190065a4bf9a87bf324842c05701330bba9f1cacbb" +dependencies = [ + "memchr", +] + [[package]] name = "rsa" version = "0.9.10" @@ -3470,7 +3503,6 @@ checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a" dependencies = [ "getrandom 0.4.3", "js-sys", - "serde_core", "wasm-bindgen", ] diff --git a/backend/Cargo.toml b/backend/Cargo.toml index d35692af2..6cd9f6c33 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -1,8 +1,10 @@ [workspace] members = [ "apps/gateway", + "crates/gateway-admin", "crates/gateway-api", "crates/gateway-core", + "crates/gateway-host", "crates/gateway-protocol", "crates/gateway-store", "crates/providers/openai", diff --git a/backend/apps/gateway/Cargo.toml b/backend/apps/gateway/Cargo.toml index 252f64bbc..954c9d40b 100644 --- a/backend/apps/gateway/Cargo.toml +++ b/backend/apps/gateway/Cargo.toml @@ -7,66 +7,26 @@ publish.workspace = true rust-version.workspace = true autotests = false -[lib] -path = "src/bootstrap.rs" - [[test]] name = "main" path = "tests/main.rs" [dependencies] -argon2 = "0.5.3" -async-trait = "0.1.89" -axum = { version = "0.8.9", default-features = false, features = ["http1", "json", "query", "tokio", "tracing", "ws"] } -base64 = "0.22.1" -bytes = "1.11.1" -chrono = { version = "0.4.45", features = ["serde", "clock"] } -config = { version = "0.15.23", default-features = false, features = ["yaml"] } -flate2 = "1.1.9" -futures = "0.3.32" +gateway-admin = { path = "../../crates/gateway-admin" } gateway-api = { path = "../../crates/gateway-api" } gateway-core = { path = "../../crates/gateway-core" } -gateway-protocol = { path = "../../crates/gateway-protocol" } +gateway-host = { path = "../../crates/gateway-host" } gateway-store = { path = "../../crates/gateway-store" } provider-openai = { path = "../../crates/providers/openai" } provider-xai = { path = "../../crates/providers/xai" } -hex = "0.4.3" -hmac = "0.13.0" -ipnet = "2.12.0" -rand = "0.10.1" -# Enables password_hash::rand_core::OsRng for argon2 0.5. -rand_core = { version = "0.6.4", features = ["getrandom"] } -redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] } -ring = "0.17.14" -# Pinned to the Codex Desktop 26.707.41301 HTTP stack; re-audit a newer Desktop before upgrading. -reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] } -secrecy = { version = "0.10.3", features = ["serde"] } -semver = "1.0.27" serde = { version = "1.0.228", features = ["derive"] } -serde_json = { version = "1.0.150", features = ["preserve_order"] } -sha2 = "0.11.0" -sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "json"] } -subtle = "2.6.1" -tar = "0.4.44" thiserror = "2.0.18" -tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net"] } -tokio-util = { version = "0.7.18", features = ["rt"] } -tower = { version = "0.5.3", features = ["util"] } -tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout", "fs"] } -tracing = "0.1.44" -tracing-appender = "0.2.5" -tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] } -url = "2.5.8" -uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] } -zeroize = { version = "1.8.2", features = ["derive"] } +tokio = { version = "1.52.3", default-features = false, features = ["rt-multi-thread"] } [dev-dependencies] -sqlx = { version = "0.9.0", default-features = false, features = ["derive"] } -syn = { version = "2.0.118", features = ["full", "visit"] } +config = { version = "0.15.23", default-features = false, features = ["yaml"] } +syn = { version = "2.0.118", features = ["full"] } tempfile = "3.27.0" -tokio = { version = "1.52.3", default-features = false, features = ["test-util"] } -tokio-tungstenite = "0.29.0" -wiremock = "0.6.5" [lints] workspace = true diff --git a/backend/apps/gateway/src/bootstrap.rs b/backend/apps/gateway/src/bootstrap.rs index e06779260..f907e07ed 100644 --- a/backend/apps/gateway/src/bootstrap.rs +++ b/backend/apps/gateway/src/bootstrap.rs @@ -1,11399 +1,118 @@ -//! 网关唯一组合根：启动配置、基础设施装配、Provider/Admin 投影和 HTTP 生命周期。 +//! 网关唯一组合根：加载包级配置并连接各 Bundle。 -pub mod workers; - -use std::{ - env, - path::{Path, PathBuf}, -}; - -use chrono::{DateTime, Utc}; -use futures::StreamExt as _; -use secrecy::{ExposeSecret, SecretBox}; -use serde::{Deserialize, de::IgnoredAny}; -use tracing_subscriber::EnvFilter; -use url::Url; +use gateway_core::engine::provider::{ProviderRegistry, RegistryError}; +use gateway_host::{ConfigError, HostConfig, LoadableConfig}; +use serde::Deserialize; const CONFIG_SCHEMA_VERSION: u32 = 1; -const CONFIG_RELATIVE_PATH: &str = "deploy/config.yaml"; -const SERVER_HOST_ENV: &str = "CPR_SERVER_HOST"; -const SERVER_PORT_ENV: &str = "CPR_SERVER_PORT"; -const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL"; -const REDIS_URL_ENV: &str = "CPR_REDIS_URL"; -const SERVICE_PASSWORD_HEX_LENGTH: usize = 48; - -const WEAK_ADMIN_PASSWORDS: &[&str] = &[ - "", - "admin", - "123456", - "password", - "changeme", - "change-me", - "replace-me", - "codex-proxy-rs", -]; - -type SecretValue = SecretBox; - -/// 已加载并完成校验的启动配置。 -/// -/// 数据库、Redis 和管理员初始化密码只在启动阶段保留，并在释放时清零。 -#[derive(Debug)] -pub struct BootstrapConfig { - app: AppConfig, - database_url: SecretValue, - redis_url: SecretValue, - admin_default_password: SecretValue, -} - -impl BootstrapConfig { - /// 从当前目录或父目录中的 `deploy/config.yaml` 加载配置。 - pub fn load() -> Result { - let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?; - let config_path = discover_config_path(¤t_directory)?; - let overrides = TopologyOverrides::from_environment()?; - Self::load_from_path_with_overrides(config_path, overrides) - } - - /// 从指定 YAML 文件加载配置，不应用容器拓扑覆盖。 - pub fn load_from_path(config_path: impl AsRef) -> Result { - Self::load_from_path_with_overrides(config_path, TopologyOverrides::default()) - } - - /// 从指定 YAML 文件加载配置，并应用明确的容器拓扑覆盖。 - pub fn load_from_path_with_overrides( - config_path: impl AsRef, - overrides: TopologyOverrides, - ) -> Result { - let config_path = absolute_path(config_path.as_ref())?; - let config_directory = config_path.parent().ok_or(ConfigError::InvalidConfigPath)?; - - let document: ConfigDocument = ::config::Config::builder() - .add_source(::config::File::from(config_path.as_path()).required(true)) - .build() - .and_then(::config::Config::try_deserialize) - .map_err(|_| ConfigError::InvalidDocument { - path: config_path.clone(), - })?; - - let ConfigDocument { - cpr: mut file_config, - _services: _, - } = document; - if file_config.schema_version != CONFIG_SCHEMA_VERSION { - return Err(ConfigError::UnsupportedSchemaVersion); - } - - overrides.apply(&mut file_config); - validate_file_config(&file_config)?; - - resolve_relative_path(config_directory, &mut file_config.logging.file.directory); - - let database_url = connection_url_with_password( - &file_config.database.url, - &file_config.database.password, - "database.url", - )?; - let redis_url = connection_url_with_password( - &file_config.redis.url, - &file_config.redis.password, - "redis.url", - )?; - - let FileAppConfig { - schema_version: _, - server, - database: _, - redis: _, - wire_profile, - admin, - logging, - } = file_config; - - let app = AppConfig { - server, - wire_profile, - admin: AdminConfig { - session_ttl_minutes: admin.session_ttl_minutes, - default_username: admin.default_username, - }, - logging, - }; - - Ok(Self { - app, - database_url, - redis_url, - admin_default_password: admin.default_password, - }) - } - - /// 返回不包含启动密码的应用运行配置。 - pub fn app(&self) -> &AppConfig { - &self.app - } - - /// 返回已安全注入密码的 PostgreSQL 连接地址。 - pub fn database_url(&self) -> &str { - self.database_url.expose_secret() - } - - /// 返回已安全注入密码的 Redis 连接地址。 - pub fn redis_url(&self) -> &str { - self.redis_url.expose_secret() - } - - pub(crate) fn into_parts(self) -> BootstrapConfigParts { - BootstrapConfigParts { - app: self.app, - database_url: self.database_url, - redis_url: self.redis_url, - admin_default_password: self.admin_default_password, - } - } -} - -pub(crate) struct BootstrapConfigParts { - pub(crate) app: AppConfig, - pub(crate) database_url: SecretValue, - pub(crate) redis_url: SecretValue, - pub(crate) admin_default_password: SecretValue, -} - -/// Docker Compose 对容器内部网络拓扑的固定覆盖。 -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct TopologyOverrides { - /// 容器内监听地址。 - pub server_host: Option, - /// 容器内监听端口。 - pub server_port: Option, - /// 容器内 PostgreSQL 地址，不包含密码。 - pub database_url: Option, - /// 容器内 Redis 地址，不包含密码。 - pub redis_url: Option, -} - -impl TopologyOverrides { - fn from_environment() -> Result { - let server_port = optional_environment_value(SERVER_PORT_ENV)? - .map(|value| { - value - .parse::() - .ok() - .filter(|port| *port > 0) - .ok_or(ConfigError::InvalidTopologyOverride(SERVER_PORT_ENV)) - }) - .transpose()?; - - Ok(Self { - server_host: optional_environment_value(SERVER_HOST_ENV)?, - server_port, - database_url: optional_environment_value(DATABASE_URL_ENV)?, - redis_url: optional_environment_value(REDIS_URL_ENV)?, - }) - } - - fn apply(self, config: &mut FileAppConfig) { - if let Some(host) = self.server_host { - config.server.host = host; - } - if let Some(port) = self.server_port { - config.server.port = port; - } - if let Some(url) = self.database_url { - config.database.url = url; - } - if let Some(url) = self.redis_url { - config.redis.url = url; - } - } -} - -/// 应用运行总配置，不包含任何启动密码。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AppConfig { - /// HTTP 服务配置。 - pub server: ServerConfig, - /// 经审计固定的 Codex Desktop 上游请求画像。 - pub wire_profile: WireProfileConfig, - /// 管理员初始化配置，不包含密码。 - pub admin: AdminConfig, - /// 日志配置。 - pub logging: LoggingConfig, -} - -/// HTTP 监听配置。 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(deny_unknown_fields)] -pub struct ServerConfig { - /// 监听主机。 - pub host: String, - /// 监听端口。 - pub port: u16, -} - -/// 经审计固定的 Codex Desktop 上游请求画像配置。 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(deny_unknown_fields)] -pub struct WireProfileConfig { - /// `originator` 请求头及 User-Agent 产品名。 - pub originator: String, - /// bundled Codex Core 版本。 - pub codex_version: String, - /// Desktop 应用版本。 - pub desktop_version: String, - /// Desktop 制品构建号。 - pub desktop_build: String, - /// Codex Core UA 中的目标操作系统类型。 - pub os_type: String, - /// Codex Core UA 中的目标操作系统版本。 - pub os_version: String, - /// Codex Core UA 中的目标架构。 - pub arch: String, - /// Codex Core UA 中的终端标记。 - pub terminal: String, - /// 此画像最后一次经制品与源码核验的时间。 - pub verified_at: DateTime, -} - -impl WireProfileConfig { - #[must_use] - pub fn user_agent(&self) -> String { - format!( - "{}/{} ({} {}; {}) {} ({}; {})", - self.originator, - self.codex_version, - self.os_type, - self.os_version, - self.arch, - self.terminal, - self.originator, - self.desktop_version, - ) - } - - #[must_use] - pub fn dashboard_view(&self) -> gateway_api::admin::observability::DashboardWireProfileView { - gateway_api::admin::observability::DashboardWireProfileView { - originator: self.originator.clone(), - codex_version: self.codex_version.clone(), - desktop_version: self.desktop_version.clone(), - desktop_build: self.desktop_build.clone(), - target: gateway_api::admin::observability::DashboardWireTargetView { - os_type: self.os_type.clone(), - os_version: self.os_version.clone(), - arch: self.arch.clone(), - terminal: self.terminal.clone(), - }, - user_agent: self.user_agent(), - verified_at: self.verified_at, - release: gateway_api::admin::observability::DashboardDesktopReleaseView { - status: "unchecked".to_owned(), - checked_at: None, - latest_version: None, - latest_build: None, - published_at: None, - minimum_system_version: None, - hardware_requirements: None, - download_url: None, - download_size: None, - signature_present: None, - error: None, - }, - } - } -} - -/// 管理员初始化与会话配置，不包含初始化密码。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminConfig { - /// 会话有效期（分钟）。 - pub session_ttl_minutes: u64, - /// 默认管理员用户名（首次启动时创建）。 - pub default_username: String, -} - -/// 应用结构化日志配置。 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(deny_unknown_fields)] -pub struct LoggingConfig { - /// 默认日志级别；进程 `RUST_LOG` 可作临时覆盖。 - pub level: String, - /// 是否写入标准输出。 - pub stdout: bool, - /// 文件日志配置。 - pub file: FileLoggingConfig, -} - -/// 应用结构化文件日志配置。 -#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] -#[serde(deny_unknown_fields)] -pub struct FileLoggingConfig { - /// 是否写入轮转文件。 - pub enabled: bool, - /// 日志目录。 - pub directory: PathBuf, - /// 按中国自然日计算的保留天数。 - pub retention_days: usize, - /// 单文件大小上限，单位 MiB。 - pub max_file_size_mb: u64, - /// 文件总数上限。 - pub max_files: usize, -} +/// 顶层配置只组合各包拥有的配置段，不解释任何业务字段。 #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] -struct ConfigDocument { - #[serde(rename = "x-cpr")] - cpr: FileAppConfig, - #[serde(rename = "services")] - _services: IgnoredAny, -} - -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct FileAppConfig { +pub struct GatewayConfig { schema_version: u32, - server: ServerConfig, - database: FileDatabaseConfig, - redis: FileRedisConfig, - wire_profile: WireProfileConfig, - admin: FileAdminConfig, - logging: LoggingConfig, + host: HostConfig, + store: gateway_store::StoreConfig, + admin: gateway_admin::AdminConfig, + api: gateway_api::ApiConfig, + openai: provider_openai::OpenAiConfig, + xai: provider_xai::XaiConfig, } -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct FileDatabaseConfig { - url: String, - password: SecretValue, -} - -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct FileRedisConfig { - url: String, - password: SecretValue, -} - -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct FileAdminConfig { - session_ttl_minutes: u64, - default_username: String, - default_password: SecretValue, -} - -fn discover_config_path(start: &Path) -> Result { - start - .ancestors() - .map(|directory| directory.join(CONFIG_RELATIVE_PATH)) - .find(|candidate| candidate.is_file()) - .ok_or(ConfigError::ConfigFileNotFound) -} - -fn absolute_path(path: &Path) -> Result { - if path.is_absolute() { - return Ok(path.to_path_buf()); - } - let current_directory = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?; - Ok(current_directory.join(path)) -} - -fn resolve_relative_path(base: &Path, path: &mut PathBuf) { - if path.is_relative() { - *path = base.join(&*path); - } -} - -fn optional_environment_value(name: &'static str) -> Result, ConfigError> { - match env::var(name) { - Ok(value) if value.trim().is_empty() => Err(ConfigError::InvalidTopologyOverride(name)), - Ok(value) => Ok(Some(value)), - Err(env::VarError::NotPresent) => Ok(None), - Err(env::VarError::NotUnicode(_)) => Err(ConfigError::InvalidTopologyOverride(name)), - } -} - -fn validate_file_config(config: &FileAppConfig) -> Result<(), ConfigError> { - validate_nonempty("server.host", &config.server.host)?; - if config.server.port == 0 { - return Err(ConfigError::InvalidField("server.port")); - } - validate_url( - "database.url", - &config.database.url, - &["postgres", "postgresql"], - true, - )?; - validate_url("redis.url", &config.redis.url, &["redis", "rediss"], true)?; - validate_service_password("database.password", &config.database.password)?; - validate_service_password("redis.password", &config.redis.password)?; - validate_admin_password(config.admin.default_password.expose_secret())?; - validate_nonempty("admin.default_username", &config.admin.default_username)?; - if config.admin.session_ttl_minutes == 0 { - return Err(ConfigError::InvalidField("admin.session_ttl_minutes")); - } - validate_wire_profile(&config.wire_profile)?; - validate_logging(&config.logging)?; - Ok(()) -} - -use chrono::{Duration as ChronoDuration, FixedOffset, TimeZone, Timelike}; -use gateway_api::admin::{ - WireValidationError, - observability::{ - AttemptMetricsView, BillingView, CostCoverageView, CostView, CursorWire, - DashboardAccountUsageView, DashboardCacheCardView, DashboardCapacityInfoView, - DashboardCardsView, DashboardCredentialUsageView, DashboardCredentialsCardView, - DashboardDataView, DashboardPoolSummaryView, DashboardQuery, DashboardTokensCardView, - DashboardTrafficCardView, DetailQuery, DiagnosticDimension as WireDiagnosticDimension, - DiagnosticItemView, DiagnosticsQuery, DiagnosticsView, HealthTimelinePointView, - HealthTimelineView, OpsErrorMetadataView, OpsErrorView, OpsQuery, OverviewCostPointView, - OverviewCostView, OverviewHealthPointView, OverviewHealthView, - OverviewPerformancePointView, OverviewPerformanceView, PageData, PageMeta, - ProviderOverviewView, RequestMetricsView, TokenDetailsView, TrendData, TrendKind, - TrendPointView, TrendSummaryView, UsageAttemptView, UsageInsightsOverviewView, UsageQuery, - UsageRecordDetailView, UsageRecordMetadataView, UsageRecordView, UsageSummaryView, - parse_attempt_index as parse_wire_attempt_index, parse_datetime as parse_wire_datetime, - parse_status as parse_wire_status, - }, -}; -use gateway_store::DecimalAmount; -use gateway_store::postgres::{ - AttemptMetrics, CurrencyCostTotal, DashboardObservation, - DiagnosticDimension as StoreDiagnosticDimension, DiagnosticObservation, ObservabilityCursor, - ObservabilityPageNumber, ObservabilityPageSize, ObservabilityRange, ObservabilityRepository, - OpsErrorFilter, OpsErrorPage, OpsErrorQuery, OpsErrorRecord, RequestMetricPoint, - RequestMetrics, UsageAttemptObservation, UsageOverview, UsageRecord, UsageRecordDetail, - UsageRecordFilter, UsageRecordPage, UsageRecordQuery, -}; -use serde::{Serialize, de::DeserializeOwned}; - -const MAX_CURSOR_BYTES: usize = 1024; -const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15; -const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4; -const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10; -const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3; -const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 99.0; - -#[derive(Debug, Clone, Copy, Default)] -struct HealthWindow { - success_requests: u64, - failed_requests: u64, - cancelled_requests: u64, - incomplete_requests: u64, - caller_error_requests: u64, -} - -fn format_number(value: u64) -> String { - let text = value.to_string(); - let mut output = String::with_capacity(text.len() + text.len() / 3); - for (index, character) in text.chars().rev().enumerate() { - if index > 0 && index % 3 == 0 { - output.push(','); +impl LoadableConfig for GatewayConfig { + fn resolve_and_validate(&mut self, source_dir: &std::path::Path) -> Result<(), ConfigError> { + if self.schema_version != CONFIG_SCHEMA_VERSION { + return Err(ConfigError::InvalidField("schema_version")); } - output.push(character); - } - output.chars().rev().collect() -} - -fn format_compact_number(value: u64) -> String { - if value < 1_000 { - return format_number(value); - } - for (suffix, threshold) in [ - ("P", 1_000_000_000_000_000_u64), - ("T", 1_000_000_000_000_u64), - ("B", 1_000_000_000_u64), - ("M", 1_000_000_u64), - ("K", 1_000_u64), - ] { - if value >= threshold { - let scaled = value as f64 / threshold as f64; - return format!("{scaled:.1}{suffix}").replace(".0", ""); - } - } - format_number(value) -} - -fn format_tokens(value: u64) -> String { - format_compact_number(value) -} - -fn format_billing_amount(value: &DecimalAmount) -> String { - format!("${}", value.as_str()) -} - -fn format_duration_ms(value: Option) -> String { - let Some(value) = value.filter(|value| *value >= 0) else { - return "—".to_owned(); - }; - if value < 1_000 { - format!("{value} ms") - } else if value < 60_000 { - format!("{:.2} s", value as f64 / 1_000.0) - } else if value < 3_600_000 { - format!("{:.1} min", value as f64 / 60_000.0) - } else { - format!("{:.1} h", value as f64 / 3_600_000.0) - } -} - -pub struct ObservabilityAdminAdapter { - repository: Arc, - accounts: Arc, - runtime_settings: Arc, - wire_profile: gateway_api::admin::observability::DashboardWireProfileView, -} - -impl ObservabilityAdminAdapter { - #[must_use] - pub fn new( - repository: Arc, - accounts: Arc, - runtime_settings: Arc, - wire_profile: gateway_api::admin::observability::DashboardWireProfileView, - ) -> Self { - Self { - repository, - accounts, - runtime_settings, - wire_profile, - } - } -} - -#[async_trait] -impl ObservabilityAdminService for ObservabilityAdminAdapter { - async fn dashboard_summary( - &self, - query: DashboardQuery, - ) -> Result { - let kind = query.trend_kind().map_err(map_wire_error)?; - let range = dashboard_range(query.start_time, query.end_time)?; - let (observation, accounts, runtime_settings) = tokio::try_join!( - self.repository.dashboard_summary(range), - self.accounts.list_provider_accounts(None, true), - self.runtime_settings.load_runtime_settings(), - ) - .map_err(|error| map_store_error(error, "Dashboard"))?; - Ok(dashboard_view( - observation, - accounts, - &runtime_settings, - kind, - self.wire_profile.clone(), - )) - } - - async fn dashboard_trend(&self, query: DashboardQuery) -> Result { - let kind = query.trend_kind().map_err(map_wire_error)?; - let range = dashboard_today_range(query.start_time, query.end_time)?; - let points = self - .repository - .dashboard_trend(range) - .await - .map_err(|error| map_store_error(error, "Dashboard trend"))?; - Ok(trend_view(kind, &points)) - } - - async fn usage_records( - &self, - query: UsageQuery, - ) -> Result, AdminServiceError> { - let (store_query, page, page_size) = usage_store_query(&query)?; - let result = self - .repository - .list_usage_records(store_query) - .await - .map_err(|error| map_store_error(error, "Usage records"))?; - usage_page_view(result, page, page_size) - } - - async fn usage_record_detail( - &self, - query: DetailQuery, - ) -> Result { - query.validate().map_err(map_wire_error)?; - let detail = self - .repository - .usage_record_detail(query.id.trim()) - .await - .map_err(|error| map_store_error(error, "Usage record"))?; - Ok(usage_detail_view(detail)) - } - - async fn usage_records_summary( - &self, - query: UsageQuery, - ) -> Result { - let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref())?; - let filter = usage_filter(&query)?; - let overview = self - .repository - .usage_summary(range, filter) - .await - .map_err(|error| map_store_error(error, "Usage summary"))?; - let average_latency = average( - overview.requests.latency_sum, - overview.requests.latency_count, - ); - Ok(UsageSummaryView { - total_requests: format_compact_number(overview.requests.request_count), - input_tokens: format_tokens(overview.requests.input_tokens), - output_tokens: format_tokens(overview.requests.output_tokens), - cached_tokens: format_tokens(overview.requests.cached_tokens), - cache_write_tokens: format_tokens(overview.requests.cache_write_tokens), - total_tokens: format_tokens(overview.requests.total_tokens), - average_latency_ms: display_duration(average_latency), - logical_requests: request_metrics_view(&overview.requests), - attempts: attempt_metrics_view(&overview.attempts), - }) - } - - async fn usage_insights_overview( - &self, - query: UsageQuery, - ) -> Result { - let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref())?; - let filter = usage_filter(&query)?; - let (overview, points) = tokio::try_join!( - self.repository.usage_summary(range, filter.clone()), - self.repository.usage_trend(range, filter), - ) - .map_err(|error| map_store_error(error, "Usage insights"))?; - Ok(overview_view(overview, &points)) - } - - async fn usage_insights_diagnostics( - &self, - query: DiagnosticsQuery, - ) -> Result { - let wire_dimension = query.dimension().map_err(map_wire_error)?; - let (dimension, dimension_name) = map_diagnostic_dimension(wire_dimension); - let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref())?; - let status_code = parse_wire_status(query.status_code).map_err(map_wire_error)?; - let filter = UsageRecordFilter { - provider_kind: non_empty(query.provider), - model: non_empty(query.model), - status_code, - search: non_empty(query.search), - ..UsageRecordFilter::default() - }; - let items = self - .repository - .usage_diagnostics(range, filter, dimension) - .await - .map_err(|error| map_store_error(error, "Usage diagnostics"))?; - let total = items.iter().map(|item| item.request_count).sum(); - Ok(DiagnosticsView { - dimension: dimension_name.to_owned(), - items: items - .into_iter() - .map(|item| diagnostic_item_view(item, total)) - .collect(), - }) - } - - async fn ops_errors( - &self, - query: OpsQuery, - ) -> Result, AdminServiceError> { - let (page_number, page_size_number) = query.validate_page().map_err(map_wire_error)?; - query.validate_cursor().map_err(map_wire_error)?; - let page = ObservabilityPageNumber::new(page_number).map_err(map_invalid_store)?; - let page_size = ObservabilityPageSize::new(page_size_number).map_err(map_invalid_store)?; - let cursor = decode_observability_cursor(query.cursor.as_deref())?; - let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref())?; - let status_code = parse_wire_status( - query - .upstream_status_code - .or(query.client_status_code) - .or(query.status_code), - ) - .map_err(map_wire_error)?; - let filter = OpsErrorFilter { - client_api_key_ref: non_empty(query.client_api_key_id), - request_id: non_empty(query.request_id), - provider_account_ref: non_empty(query.account_id), - provider_kind: non_empty(query.provider), - operation: non_empty(query.route).or_else(|| non_empty(query.kind)), - model: non_empty(query.model), - transport: non_empty(query.transport), - attempt_index: parse_wire_attempt_index(query.attempt_index).map_err(map_wire_error)?, - response_id: non_empty(query.response_id), - upstream_request_id: non_empty(query.upstream_request_id), - failure_kind: non_empty(query.failure_class), - status_code, - search: non_empty(query.search), - }; - let result = self - .repository - .list_ops_errors(OpsErrorQuery { - range, - filter, - cursor, - page, - page_size, - }) - .await - .map_err(|error| map_store_error(error, "Ops errors"))?; - ops_page_view(result, page.get(), page_size.get()) - } -} - -fn dashboard_view( - observation: DashboardObservation, - accounts: Vec, - runtime_settings: &gateway_store::postgres::RuntimeSettings, - kind: TrendKind, - wire_profile: gateway_api::admin::observability::DashboardWireProfileView, -) -> DashboardDataView { - let today_start = china_day_start(observation.range.end); - let yesterday_start = today_start - ChronoDuration::days(1); - let today = sum_points(&observation.trend, today_start, observation.range.end); - let yesterday = sum_points(&observation.trend, yesterday_start, today_start); - let costs = cost_views(&observation.attempts.costs); - let usd = currency_amount(&observation.attempts.costs, "USD"); - let trend = trend_view(kind, &observation.trend); - let health_timeline = health_timeline_view(&observation.trend); - let account_usage = observation - .account_usage - .iter() - .map(|credential| DashboardAccountUsageView { - id: credential.account_id.clone(), - email: credential - .email - .clone() - .unwrap_or_else(|| credential.name.clone()), - plan_type: credential.plan_type.clone(), - tokens: credential - .total_tokens - .map_or_else(|| "—".to_owned(), format_tokens), - quota_used_percent: None, - last_used: relative_time(credential.last_used_at, observation.range.end), - }) - .collect::>(); - let credential_usage = observation - .account_usage - .into_iter() - .map(|credential| DashboardCredentialUsageView { - id: credential.account_id, - display_name: credential.email.unwrap_or(credential.name), - plan_type: credential.plan_type, - tokens: credential - .total_tokens - .map_or_else(|| "-".to_owned(), format_tokens), - tokens_value: credential.total_tokens, - last_used: credential - .last_used_at - .map_or_else(|| "-".to_owned(), |value| china_datetime(&value)), - provider: credential.provider_kind, - availability: credential.availability, - request_count: credential.request_count, - }) - .collect::>(); - let now = observation.range.end; - let total = u64::try_from(accounts.len()).unwrap_or(u64::MAX); - let active = u64::try_from( - accounts - .iter() - .filter(|account| { - account.enabled - && account.availability == ProviderAccountAvailability::Ready - && account.access_token_expires_at > now - }) - .count(), - ) - .unwrap_or(u64::MAX); - let pool_summary = DashboardPoolSummaryView { - total, - active, - expired: u64::try_from( - accounts - .iter() - .filter(|account| { - account.availability == ProviderAccountAvailability::Expired - || account.access_token_expires_at <= now - }) - .count(), - ) - .unwrap_or(u64::MAX), - quota_exhausted: u64::try_from( - accounts - .iter() - .filter(|account| { - account.availability == ProviderAccountAvailability::QuotaExhausted - }) - .count(), - ) - .unwrap_or(u64::MAX), - refreshing: None, - disabled: u64::try_from(accounts.iter().filter(|account| !account.enabled).count()) - .unwrap_or(u64::MAX), - banned: u64::try_from( - accounts - .iter() - .filter(|account| account.availability == ProviderAccountAvailability::Banned) - .count(), - ) - .unwrap_or(u64::MAX), - }; - let max_concurrent = u64::from(runtime_settings.max_concurrent_per_account); - let total_slots = active.saturating_mul(max_concurrent); - let usage_records = observation - .recent_requests - .into_iter() - .map(usage_record_view) - .collect::>(); - - DashboardDataView { - cards: DashboardCardsView { - credentials: DashboardCredentialsCardView { - total: format_compact_number(observation.provider_accounts.total), - total_value: observation.provider_accounts.total, - enabled: format_compact_number(observation.provider_accounts.enabled), - enabled_value: observation.provider_accounts.enabled, - unavailable: format_compact_number(observation.provider_accounts.unavailable), - unavailable_value: observation.provider_accounts.unavailable, - }, - traffic: DashboardTrafficCardView { - today_requests: format_compact_number(today.request_count), - today_requests_value: today.request_count, - yesterday_requests_value: yesterday.request_count, - total_requests: format_compact_number(observation.requests.request_count), - }, - tokens: DashboardTokensCardView { - today_tokens: format_tokens(today.total_tokens), - today_tokens_value: today.total_tokens, - yesterday_tokens_value: yesterday.total_tokens, - total_tokens: format_tokens(observation.requests.total_tokens), - total_billing_amount_usd: usd.map_or_else(|| "—".to_owned(), format_billing_amount), - }, - cache: DashboardCacheCardView { - today_hit_rate: display_rate(rate(today.cached_tokens, today.input_tokens)), - today_hit_rate_value: optional_rate(today.cached_tokens, today.input_tokens), - yesterday_hit_rate_value: optional_rate( - yesterday.cached_tokens, - yesterday.input_tokens, - ), - total_hit_rate: display_rate(rate( - observation.requests.cached_tokens, - observation.requests.input_tokens, - )), - total_cached_tokens: format_tokens(observation.requests.cached_tokens), - average_first_token_latency_ms: display_duration(average( - observation.requests.first_token_latency_sum, - observation.requests.first_token_latency_count, - )), - }, - }, - trend, - health_timeline, - wire_profile, - account_usage, - credential_usage, - usage_records, - pool_summary, - capacity_info: DashboardCapacityInfoView { - max_concurrent_per_account: max_concurrent, - total_slots, - used_slots: None, - available_slots: None, - }, - rotation_strategy: runtime_settings.rotation_strategy.clone(), - logical_requests: request_metrics_view(&observation.requests), - attempts: attempt_metrics_view(&observation.attempts), - costs, - } -} - -fn relative_time(value: Option>, now: DateTime) -> String { - let Some(value) = value else { - return "从未使用".to_owned(); - }; - let elapsed = now.signed_duration_since(value); - if elapsed.num_seconds() < 0 { - return china_datetime(&value); - } - if elapsed.num_seconds() < 60 { - return "刚刚".to_owned(); - } - if elapsed.num_minutes() < 60 { - return format!("{} 分钟前", elapsed.num_minutes()); - } - if elapsed.num_hours() < 24 { - return format!("{} 小时前", elapsed.num_hours()); - } - format!("{} 天前", elapsed.num_days()) -} - -fn usage_page_view( - page: UsageRecordPage, - page_number: u32, - page_size: u16, -) -> Result, AdminServiceError> { - Ok(PageData { - items: page.items.into_iter().map(usage_record_view).collect(), - page: page_meta(page_number, page_size, page.total), - next_cursor: page - .next_cursor - .as_ref() - .map(encode_observability_cursor) - .transpose()?, - }) -} - -fn ops_page_view( - page: OpsErrorPage, - page_number: u32, - page_size: u16, -) -> Result, AdminServiceError> { - Ok(PageData { - items: page.items.into_iter().map(ops_error_view).collect(), - page: page_meta(page_number, page_size, page.total), - next_cursor: page - .next_cursor - .as_ref() - .map(encode_observability_cursor) - .transpose()?, - }) -} - -fn page_meta(page: u32, page_size: u16, total: u64) -> PageMeta { - let page_size_u64 = u64::from(page_size); - let total_pages_u64 = total.saturating_add(page_size_u64 - 1) / page_size_u64; - PageMeta { - page, - page_size, - total, - total_pages: u32::try_from(total_pages_u64).unwrap_or(u32::MAX), - } -} - -fn usage_record_view(record: UsageRecord) -> UsageRecordView { - let tokens = token_details(&record); - let billing = usage_billing_view(&record); - let costs = record - .cost_amount - .as_ref() - .zip(record.cost_currency.as_ref()) - .map(|(amount, currency)| { - vec![CostView { - currency: currency.clone(), - estimated_amount: amount.to_string(), - }] - }) - .unwrap_or_default(); - let status_code = record - .client_status_code - .or(record.upstream_status_code) - .map(i64::from); - let message = record - .error_message - .clone() - .unwrap_or_else(|| record.outcome.clone()); - let first_token_display = display_duration(record.first_token_ms); - let latency_display = display_duration(record.latency_ms); - let cost_coverage = match record.cost_source.as_str() { - "provider_reported" | "calculated" => CostCoverageView { - known: 1, - partial: 0, - unknown: 0, - not_billable: 0, - }, - _ => CostCoverageView { - known: 0, - partial: 0, - unknown: 1, - not_billable: 0, - }, - }; - UsageRecordView { - id: record.id.clone(), - request_id: record.id, - client_api_key_id: Some(record.client_api_key_ref), - kind: record.operation.clone(), - provider: record.provider_kind, - account_id: record.provider_account_ref, - account_email: record.provider_account_email, - route: None, - model: record - .upstream_model_id - .clone() - .unwrap_or_else(|| record.requested_model_id.clone()), - requested_model: Some(record.requested_model_id), - upstream_model: record.upstream_model_id, - service_tier: None, - status_code, - transport: record.upstream_transport.or(Some(record.client_transport)), - attempt_index: None, - attempt_count: u64::from(record.attempt_count), - response_id: record.client_response_id, - upstream_request_id: record.upstream_request_id, - latency_ms: record.latency_ms, - first_token_ms: record.first_token_ms, - input_tokens: record.input_tokens, - output_tokens: record.output_tokens, - cached_tokens: record.cached_tokens, - cache_write_tokens: record.cache_write_tokens, - reasoning_tokens: record.reasoning_tokens, - message, - metadata: UsageRecordMetadataView { - protocol: record.protocol, - logical_outcome: record.outcome.clone(), - attempt_count: u64::from(record.attempt_count), - }, - created_at: record.started_at, - created_at_display: china_datetime(&record.started_at), - client_ip: record.client_ip, - user_agent: record.user_agent, - reasoning_effort: record.reasoning_effort, - reasoning_preset: record.reasoning_preset, - compact: Some(record.compact), - request_kind: record.request_kind, - subagent_kind: record.subagent_kind, - token_details: tokens, - billing, - costs, - cost_coverage, - first_token_latency_ms: record.first_token_ms, - first_token_latency_ms_display: first_token_display, - latency_ms_display: latency_display, - logical_outcome: record.outcome, - } -} - -fn usage_billing_view(record: &UsageRecord) -> Option { - let total = record.cost_amount.as_ref()?; - let currency = record.cost_currency.as_deref()?; - if record.cost_source != "calculated" { - return Some(total_only_billing_view(total, currency)); - } - let Some(breakdown) = calculated_billing_breakdown(record) else { - return Some(total_only_billing_view(total, currency)); - }; - let Ok(persisted_total) = total.as_str().parse::() else { - return Some(total_only_billing_view(total, currency)); - }; - if breakdown.total_amount().amount() != persisted_total - || breakdown.total_amount().currency().as_str() != currency - { - return Some(total_only_billing_view(total, currency)); - } - - Some(BillingView { - input_amount_display: format_money(breakdown.input_amount()), - output_amount_display: format_money(breakdown.output_amount()), - cache_read_amount_display: format_money(breakdown.cache_read_amount()), - cache_write_amount_display: format_money(breakdown.cache_write_amount()), - standard_amount_display: format_money(breakdown.standard_amount()), - total_amount_display: format_decimal_currency(total, currency), - input_price_display: format_token_price(breakdown.input_price_per_million()), - output_price_display: format_token_price(breakdown.output_price_per_million()), - cache_read_price_display: format_token_price(breakdown.cache_read_price_per_million()), - cache_write_price_display: format_token_price(breakdown.cache_write_price_per_million()), - service_tier_display: format_service_tier(breakdown.service_tier()), - multiplier_display: format!("{:.2}x", f64::from(breakdown.multiplier_percent()) / 100.0), - }) -} - -fn calculated_billing_breakdown( - record: &UsageRecord, -) -> Option { - let model = record.upstream_model_id.as_deref()?; - let input = record.input_tokens?; - let output = record.output_tokens?; - let cached = record.cached_tokens?; - match record.provider_kind.as_deref()? { - "xai" => provider_xai::grok_billing_breakdown(model, input, output, cached), - "openai" => { - let cache_write = record.cache_write_tokens?; - [None, Some("fast"), Some("flex")] - .into_iter() - .filter_map(|tier| { - provider_openai::openai_billing_breakdown( - model, - input, - output, - cached, - cache_write, - tier, - ) - }) - .find(|breakdown| { - record.cost_amount.as_ref().is_some_and(|total| { - total - .as_str() - .parse::() - .is_ok_and(|persisted| breakdown.total_amount().amount() == persisted) - }) - }) - } - _ => None, - } -} - -fn total_only_billing_view(total: &DecimalAmount, currency: &str) -> BillingView { - BillingView { - input_amount_display: "—".to_owned(), - output_amount_display: "—".to_owned(), - cache_read_amount_display: "—".to_owned(), - cache_write_amount_display: "—".to_owned(), - standard_amount_display: "—".to_owned(), - total_amount_display: format_decimal_currency(total, currency), - input_price_display: "—".to_owned(), - output_price_display: "—".to_owned(), - cache_read_price_display: "—".to_owned(), - cache_write_price_display: "—".to_owned(), - service_tier_display: "—".to_owned(), - multiplier_display: "—".to_owned(), - } -} - -fn format_money(money: gateway_core::accounting::Money) -> String { - let currency = money.currency(); - let amount = money.amount().to_string(); - format_decimal_currency_value(&amount, currency.as_str()) -} - -fn format_decimal_currency(amount: &DecimalAmount, currency: &str) -> String { - format_decimal_currency_value(amount.as_str(), currency) -} - -fn format_decimal_currency_value(amount: &str, currency: &str) -> String { - if currency == "USD" { - let value = amount.parse::().unwrap_or_default(); - let precision = if value != 0.0 && value.abs() < 0.01 { - 4 - } else { - 2 - }; - format!("${value:.precision$}") - } else { - format!("{currency} {amount}") - } -} - -fn format_token_price(price: gateway_core::accounting::Money) -> String { - if price.currency().as_str() != "USD" { - return format!("{} {} / 1M Token", price.currency(), price.amount()); - } - let value = price.amount().scaled() as f64 / 10_000_000_000.0; - format!("${value:.4} / 1M Token") -} - -fn format_service_tier(service_tier: Option<&str>) -> String { - match service_tier { - Some("priority" | "fast") => "Fast".to_owned(), - Some("flex") => "Flex".to_owned(), - Some("default") | None => "Default".to_owned(), - Some(other) => other.to_owned(), - } -} - -fn usage_detail_view(detail: UsageRecordDetail) -> UsageRecordDetailView { - UsageRecordDetailView { - request: usage_record_view(detail.request), - attempts: detail.attempts.into_iter().map(attempt_view).collect(), - } -} - -fn attempt_view(attempt: UsageAttemptObservation) -> UsageAttemptView { - let occurred_at = attempt.occurred_at; - UsageAttemptView { - id: attempt.id, - attempt_index: attempt.attempt_index, - trigger: attempt.source, - provider: attempt - .provider_kind - .unwrap_or_else(|| "unknown".to_owned()), - provider_instance_id: attempt - .provider_instance_id - .unwrap_or_else(|| "unknown".to_owned()), - model: attempt - .upstream_model_id - .unwrap_or_else(|| "unknown".to_owned()), - transport: attempt - .upstream_transport - .unwrap_or_else(|| "unknown".to_owned()), - send_state: attempt - .upstream_send_state - .unwrap_or_else(|| "unknown".to_owned()), - outcome: attempt.outcome, - downstream_committed: attempt.downstream_committed, - status_code: attempt.status_code, - provider_error_code: attempt.provider_error_code, - failure_class: attempt.failure_kind, - cost_estimate_status: attempt - .cost_source - .clone() - .unwrap_or_else(|| "unavailable".to_owned()), - estimated_cost_amount: attempt.cost_amount.map(|amount| amount.to_string()), - estimated_cost_currency: attempt.cost_currency, - input_tokens: attempt.input_tokens, - output_tokens: attempt.output_tokens, - cached_tokens: attempt.cached_tokens, - total_tokens: attempt.total_tokens, - first_token_ms: None, - latency_ms: attempt.latency_ms, - credential_name: attempt.provider_account_ref, - account_email: None, - started_at: occurred_at, - completed_at: Some(occurred_at), - } -} - -fn token_details(record: &UsageRecord) -> TokenDetailsView { - let input = record.input_tokens; - let output = record.output_tokens; - let cached = record.cached_tokens; - let cache_write = record.cache_write_tokens; - let reasoning = record.reasoning_tokens; - let total = record.total_tokens; - TokenDetailsView { - input_tokens: input, - output_tokens: output, - cached_tokens: cached, - cache_write_tokens: cache_write, - reasoning_tokens: reasoning, - total_tokens: total, - input_tokens_display: optional_number(input), - output_tokens_display: optional_number(output), - cached_tokens_display: optional_compact_number(cached), - cache_write_tokens_display: optional_compact_number(cache_write), - reasoning_tokens_display: optional_number(reasoning), - total_tokens_display: optional_number(total), - } -} - -fn optional_number(value: Option) -> String { - value.map_or_else(|| "-".to_owned(), format_number) -} - -fn optional_compact_number(value: Option) -> String { - value.map_or_else(|| "-".to_owned(), format_compact_number) -} - -fn trend_view(kind: TrendKind, points: &[RequestMetricPoint]) -> TrendData { - let summary = trend_summary(kind, points); - let points = points.iter().map(trend_point_view).collect::>(); - TrendData { - kind, - points, - summary, - } -} - -fn trend_point_view(point: &RequestMetricPoint) -> TrendPointView { - let metrics = &point.metrics; - let latency = average(metrics.latency_sum, metrics.latency_count); - let first_token_latency = average( - metrics.first_token_latency_sum, - metrics.first_token_latency_count, - ); - let errors = service_failure_count(metrics); - let success_rate = (metrics.request_count > 0) - .then(|| metrics.success_count as f64 / metrics.request_count as f64 * 100.0); - let local = point.bucket_start.with_timezone(&china_offset()); - let label = local.format("%m-%d %H:%M").to_string(); - TrendPointView { - time: local.format("%H:%M").to_string(), - bucket: point.bucket_start, - label, - requests: format_compact_number(metrics.request_count), - requests_value: metrics.request_count, - input_tokens: format_tokens(metrics.input_tokens), - input_tokens_value: metrics.input_tokens, - output_tokens: format_tokens(metrics.output_tokens), - output_tokens_value: metrics.output_tokens, - cached_tokens: format_tokens(metrics.cached_tokens), - cached_tokens_value: metrics.cached_tokens, - cache_hit_rate_value: rate(metrics.cached_tokens, metrics.input_tokens), - tokens_value: metrics.total_tokens, - errors: format_compact_number(errors), - errors_value: errors, - latency: display_duration(latency), - latency_value: latency, - first_token_latency: display_duration(first_token_latency), - first_token_latency_value: first_token_latency, - max_latency: display_duration(metrics.max_latency_ms), - max_latency_value: metrics.max_latency_ms, - min_latency: display_duration(metrics.min_latency_ms), - min_latency_value: metrics.min_latency_ms, - success_rate: success_rate.map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), - success_rate_value: success_rate, - } -} - -fn trend_summary(kind: TrendKind, points: &[RequestMetricPoint]) -> Vec { - match kind { - TrendKind::Usage => vec![ - summary( - "输入", - points.iter().map(|point| point.metrics.input_tokens).sum(), - None, - ), - summary( - "输出", - points.iter().map(|point| point.metrics.output_tokens).sum(), - None, - ), - summary( - "缓存", - points.iter().map(|point| point.metrics.cached_tokens).sum(), - None, - ), - ], - TrendKind::Latency => { - let latency_sum = points.iter().map(|point| point.metrics.latency_sum).sum(); - let latency_count = points.iter().map(|point| point.metrics.latency_count).sum(); - let maximum = points - .iter() - .filter_map(|point| point.metrics.max_latency_ms) - .max(); - let minimum = points - .iter() - .filter_map(|point| point.metrics.min_latency_ms) - .min(); - vec![ - summary_duration("平均", average(latency_sum, latency_count)), - summary_duration("最高", maximum), - summary_duration("最低", minimum), - ] - } - TrendKind::Errors => { - let errors = points - .iter() - .map(|point| service_failure_count(&point.metrics)) - .sum::(); - let successes = points - .iter() - .map(|point| point.metrics.success_count) - .sum::(); - let requests = points - .iter() - .map(|point| point.metrics.request_count) - .sum::(); - let health_requests = successes.saturating_add(errors); - let success = - (health_requests > 0).then(|| successes as f64 / health_requests as f64 * 100.0); - vec![ - summary("错误数", errors, None), - TrendSummaryView { - label: "成功率".to_owned(), - value: "—".to_owned(), - ratio: success.map(|value| format!("{value:.1}%")), - }, - summary("总请求", requests, None), - ] - } - } -} - -fn summary(label: &str, value: u64, ratio: Option) -> TrendSummaryView { - TrendSummaryView { - label: label.to_owned(), - value: format_compact_number(value), - ratio, - } -} - -fn summary_duration(label: &str, value: Option) -> TrendSummaryView { - TrendSummaryView { - label: label.to_owned(), - value: display_duration(value), - ratio: None, - } -} - -fn health_timeline_view(points: &[RequestMetricPoint]) -> HealthTimelineView { - health_timeline_view_at(points, Utc::now()) -} - -/// 按指定时刻投影固定的中国自然日健康时间线。 -#[must_use] -pub fn health_timeline_view_at( - records: &[RequestMetricPoint], - now: DateTime, -) -> HealthTimelineView { - let current_slot = china_quarter_hour_start(now); - let start = china_day_start(now); - let mut buckets = (0..HEALTH_TIMELINE_SLOTS) - .map(|index| { - ( - start + ChronoDuration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index), - HealthWindow::default(), - ) - }) - .collect::>(); - for record in records { - if record.bucket_start < start || record.bucket_start > now { - continue; - } - let record_slot = china_quarter_hour_start(record.bucket_start); - if let Some((_, bucket)) = buckets - .iter_mut() - .find(|(bucket_start, _)| *bucket_start == record_slot) - { - bucket.success_requests = bucket - .success_requests - .saturating_add(record.metrics.success_count); - bucket.failed_requests = bucket - .failed_requests - .saturating_add(service_failure_count(&record.metrics)); - bucket.cancelled_requests = bucket - .cancelled_requests - .saturating_add(record.metrics.cancelled_count); - bucket.incomplete_requests = bucket - .incomplete_requests - .saturating_add(record.metrics.incomplete_count); - bucket.caller_error_requests = bucket - .caller_error_requests - .saturating_add(record.metrics.caller_error_count); - } - } - - let totals = buckets - .iter() - .filter(|(bucket_start, _)| *bucket_start <= current_slot) - .fold(HealthWindow::default(), |mut totals, (_, bucket)| { - totals.success_requests = totals - .success_requests - .saturating_add(bucket.success_requests); - totals.failed_requests = totals - .failed_requests - .saturating_add(bucket.failed_requests); - totals.cancelled_requests = totals - .cancelled_requests - .saturating_add(bucket.cancelled_requests); - totals.incomplete_requests = totals - .incomplete_requests - .saturating_add(bucket.incomplete_requests); - totals.caller_error_requests = totals - .caller_error_requests - .saturating_add(bucket.caller_error_requests); - totals - }); - HealthTimelineView { - title: "请求健康时间线".to_owned(), - description: "有效请求可用性".to_owned(), - reliability_display: format_health_reliability(totals), - status: health_status(totals, false).to_owned(), - success_requests: totals.success_requests, - failed_requests: totals.failed_requests, - cancelled_requests: totals.cancelled_requests, - incomplete_requests: totals.incomplete_requests, - caller_error_requests: totals.caller_error_requests, - points: buckets - .into_iter() - .enumerate() - .map(|(index, (bucket_start, bucket))| { - let elapsed_minutes = index as i64 * HEALTH_TIMELINE_SLOT_MINUTES; - HealthTimelinePointView { - time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60), - status: health_status(bucket, bucket_start > current_slot).to_owned(), - reliability_display: format_health_reliability(bucket), - success_requests: bucket.success_requests, - failed_requests: bucket.failed_requests, - cancelled_requests: bucket.cancelled_requests, - incomplete_requests: bucket.incomplete_requests, - caller_error_requests: bucket.caller_error_requests, - } - }) - .collect(), - } -} - -fn health_status(bucket: HealthWindow, is_future: bool) -> &'static str { - let eligible_requests = bucket - .success_requests - .saturating_add(bucket.failed_requests); - if is_future { - "future" - } else if eligible_requests == 0 { - "no_data" - } else if bucket.success_requests == 0 - && bucket.failed_requests >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD - { - "unavailable" - } else if eligible_requests < HEALTH_TIMELINE_MIN_SAMPLE_SIZE { - "low_sample" - } else if health_reliability(bucket) - .is_some_and(|reliability| reliability < HEALTH_TIMELINE_STABLE_RELIABILITY) - { - "unstable" - } else { - "stable" - } -} - -fn health_reliability(bucket: HealthWindow) -> Option { - let eligible_requests = bucket - .success_requests - .saturating_add(bucket.failed_requests); - (eligible_requests > 0) - .then(|| bucket.success_requests as f64 / eligible_requests as f64 * 100.0) -} - -fn format_health_reliability(bucket: HealthWindow) -> String { - health_reliability(bucket) - .map(|value| format!("{value:.1}%")) - .unwrap_or_else(|| "-".to_owned()) -} - -fn service_failure_count(metrics: &RequestMetrics) -> u64 { - metrics - .failure_count - .saturating_sub(metrics.caller_error_count) -} - -fn overview_view( - overview: UsageOverview, - points: &[RequestMetricPoint], -) -> UsageInsightsOverviewView { - let requests = &overview.requests; - let estimated_cost = currency_amount(&overview.attempts.costs, "USD"); - let failed_requests = service_failure_count(requests); - let health_requests = requests.success_count.saturating_add(failed_requests); - let health_points = points - .iter() - .map(|point| { - let failed = service_failure_count(&point.metrics); - let health_requests = point.metrics.success_count.saturating_add(failed); - OverviewHealthPointView { - bucket: point.bucket_start, - label: point - .bucket_start - .with_timezone(&china_offset()) - .format("%m-%d %H:%M") - .to_string(), - success_requests: point.metrics.success_count, - failed_requests: failed, - cancelled_requests: point.metrics.cancelled_count, - incomplete_requests: point.metrics.incomplete_count, - caller_error_requests: point.metrics.caller_error_count, - error_rate: rate(failed, health_requests), - } - }) - .collect(); - let performance_points = points - .iter() - .map(|point| OverviewPerformancePointView { - bucket: point.bucket_start, - label: point - .bucket_start - .with_timezone(&china_offset()) - .format("%m-%d %H:%M") - .to_string(), - latency_p50_ms: point - .metrics - .latency_percentiles - .p50_ms - .map(|value| value.as_f64()), - latency_p95_ms: point - .metrics - .latency_percentiles - .p95_ms - .map(|value| value.as_f64()), - latency_p99_ms: point - .metrics - .latency_percentiles - .p99_ms - .map(|value| value.as_f64()), - first_token_p50_ms: point - .metrics - .first_token_latency_percentiles - .p50_ms - .map(|value| value.as_f64()), - first_token_p95_ms: point - .metrics - .first_token_latency_percentiles - .p95_ms - .map(|value| value.as_f64()), - first_token_p99_ms: point - .metrics - .first_token_latency_percentiles - .p99_ms - .map(|value| value.as_f64()), - }) - .collect(); - let cost_points = points - .iter() - .map(|point| OverviewCostPointView { - bucket: point.bucket_start, - label: point - .bucket_start - .with_timezone(&china_offset()) - .format("%m-%d %H:%M") - .to_string(), - input_tokens: point.metrics.input_tokens, - output_tokens: point.metrics.output_tokens, - cached_tokens: point.metrics.cached_tokens, - total_tokens: point.metrics.total_tokens, - estimated_cost: None, - standard_cost: None, - cached_token_rate: rate(point.metrics.cached_tokens, point.metrics.input_tokens), - cache_hit_request_rate: point.metrics.cache_hit_request_rate(), - }) - .collect(); - - UsageInsightsOverviewView { - granularity: "15m".to_owned(), - health: OverviewHealthView { - total_requests: health_requests, - success_requests: requests.success_count, - failed_requests, - cancelled_requests: requests.cancelled_count, - incomplete_requests: requests.incomplete_count, - caller_error_requests: requests.caller_error_count, - success_rate: rate(requests.success_count, health_requests), - request_change_rate: None, - success_rate_change: None, - points: health_points, - }, - performance: OverviewPerformanceView { - latency_p50_ms: requests - .latency_percentiles - .p50_ms - .map(|value| value.as_f64()), - latency_p95_ms: requests - .latency_percentiles - .p95_ms - .map(|value| value.as_f64()), - latency_p99_ms: requests - .latency_percentiles - .p99_ms - .map(|value| value.as_f64()), - first_token_p50_ms: requests - .first_token_latency_percentiles - .p50_ms - .map(|value| value.as_f64()), - first_token_p95_ms: requests - .first_token_latency_percentiles - .p95_ms - .map(|value| value.as_f64()), - first_token_p99_ms: requests - .first_token_latency_percentiles - .p99_ms - .map(|value| value.as_f64()), - latency_coverage: rate(requests.latency_count, requests.request_count), - first_token_coverage: rate(requests.first_token_latency_count, requests.request_count), - points: performance_points, - }, - cost: OverviewCostView { - estimated_cost: estimated_cost.map(ToString::to_string), - standard_cost: None, - cost_per_request: None, - tokens_per_request: if requests.request_count > 0 { - requests.total_tokens as f64 / requests.request_count as f64 - } else { - 0.0 - }, - cached_token_rate: rate(requests.cached_tokens, requests.input_tokens), - cache_hit_request_rate: requests.cache_hit_request_rate(), - input_tokens: requests.input_tokens, - output_tokens: requests.output_tokens, - cached_tokens: requests.cached_tokens, - total_tokens: requests.total_tokens, - points: cost_points, - costs: cost_views(&overview.attempts.costs), - coverage: cost_coverage_view(&overview.attempts), - }, - attempts: attempt_metrics_view(&overview.attempts), - providers: overview - .providers - .into_iter() - .map(|provider| ProviderOverviewView { - provider: provider.provider_kind, - request_count: provider.request_count, - attempt_count: provider.attempt_count, - failure_count: provider.failure_count, - total_tokens: provider.total_tokens, - }) - .collect(), - } -} - -fn diagnostic_item_view(item: DiagnosticObservation, total: u64) -> DiagnosticItemView { - let error_count = item.failure_count; - DiagnosticItemView { - name: display_dimension_name(&item.name), - request_count: item.request_count, - success_count: item.request_count.saturating_sub(error_count), - error_count, - error_rate: rate(error_count, item.request_count), - request_share: rate(item.request_count, total), - average_latency_ms: item.average_latency_ms, - estimated_cost: None, - attempt_count: item.attempt_count, - total_tokens: item.total_tokens, - } -} - -fn ops_error_view(error: OpsErrorRecord) -> OpsErrorView { - let status = error.status_code.map(i64::from); - OpsErrorView { - id: error.event_id, - request_id: error.request_id, - client_api_key_id: error.client_api_key_ref, - kind: error.operation.clone(), - provider: error.provider_kind, - account_id: error.provider_account_ref, - route: error.operation, - model: error.upstream_model_id, - status_code: status, - client_status_code: None, - upstream_status_code: None, - transport: error.upstream_transport, - attempt_index: error.attempt_index, - failure_class: error.failure_kind, - response_id: error.client_response_id, - upstream_request_id: error.upstream_request_id, - latency_ms: error.latency_ms, - message: error.message, - metadata: OpsErrorMetadataView { - source: error.source, - component: error.component, - attempt_id: None, - provider_instance_id: error.provider_instance_id, - account_label: None, - }, - created_at: error.occurred_at, - created_at_display: china_datetime(&error.occurred_at), - } -} - -fn request_metrics_view(metrics: &RequestMetrics) -> RequestMetricsView { - RequestMetricsView { - request_count: metrics.request_count, - success_count: metrics.success_count, - failure_count: metrics.failure_count, - cancelled_count: metrics.cancelled_count, - incomplete_count: metrics.incomplete_count, - caller_error_count: metrics.caller_error_count, - input_tokens: metrics.input_tokens, - output_tokens: metrics.output_tokens, - cached_tokens: metrics.cached_tokens, - cache_write_tokens: metrics.cache_write_tokens, - reasoning_tokens: metrics.reasoning_tokens, - total_tokens: metrics.total_tokens, - } -} - -fn attempt_metrics_view(metrics: &AttemptMetrics) -> AttemptMetricsView { - AttemptMetricsView { - attempt_count: metrics.attempt_count, - success_count: metrics.success_count, - failure_count: metrics.failure_count, - cancelled_count: metrics.cancelled_count, - incomplete_count: metrics.incomplete_count, - rate_limited_count: metrics.rate_limited_count, - auth_failure_count: metrics.auth_failure_count, - provider5xx_count: metrics.provider_5xx_count, - cost_coverage: cost_coverage_view(metrics), - costs: cost_views(&metrics.costs), - } -} - -fn cost_coverage_view(metrics: &AttemptMetrics) -> CostCoverageView { - CostCoverageView { - known: metrics - .cost_coverage - .provider_reported_count - .saturating_add(metrics.cost_coverage.calculated_count), - partial: 0, - unknown: metrics.cost_coverage.unavailable_count, - not_billable: 0, - } -} - -fn cost_views(costs: &[CurrencyCostTotal]) -> Vec { - costs - .iter() - .map(|cost| CostView { - currency: cost.currency.clone(), - estimated_amount: cost.amount.to_string(), - }) - .collect() -} - -fn currency_amount<'a>( - costs: &'a [CurrencyCostTotal], - currency: &str, -) -> Option<&'a DecimalAmount> { - costs - .iter() - .find(|cost| cost.currency.eq_ignore_ascii_case(currency)) - .map(|cost| &cost.amount) -} - -fn usage_store_query( - query: &UsageQuery, -) -> Result<(UsageRecordQuery, u32, u16), AdminServiceError> { - let (page_number, page_size_number) = query.validate_page().map_err(map_wire_error)?; - query.validate_cursor().map_err(map_wire_error)?; - let page = ObservabilityPageNumber::new(page_number).map_err(map_invalid_store)?; - let page_size = ObservabilityPageSize::new(page_size_number).map_err(map_invalid_store)?; - let cursor = decode_observability_cursor(query.cursor.as_deref())?; - let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref())?; - let filter = usage_filter(query)?; - Ok(( - UsageRecordQuery { - range, - filter, - cursor, - page, - page_size, - }, - page.get(), - page_size.get(), - )) -} - -fn usage_filter(query: &UsageQuery) -> Result { - let kind = non_empty(query.outcome.clone()).or_else(|| { - non_empty(query.kind.clone()).filter(|value| { - matches!( - value.as_str(), - "running" | "succeeded" | "failed" | "cancelled" | "incomplete" - ) - }) - }); - Ok(UsageRecordFilter { - client_api_key_ref: non_empty(query.client_api_key_id.clone()), - request_id: non_empty(query.request_id.clone()), - provider_account_ref: non_empty(query.account_id.clone()), - operation: non_empty(query.route.clone()), - provider_kind: non_empty(query.provider.clone()), - model: non_empty(query.model.clone()), - outcome: kind, - status_code: parse_wire_status(query.status_code).map_err(map_wire_error)?, - transport: non_empty(query.transport.clone()), - attempt_index: parse_wire_attempt_index(query.attempt_index).map_err(map_wire_error)?, - response_id: non_empty(query.response_id.clone()), - upstream_request_id: non_empty(query.upstream_request_id.clone()), - search: non_empty(query.search.clone()), - }) -} - -fn decode_cursor(encoded: &str) -> Result { - if encoded.is_empty() || encoded.len() > MAX_CURSOR_BYTES { - return Err(AdminServiceError::invalid("Invalid observability cursor")); - } - let bytes = URL_SAFE_NO_PAD - .decode(encoded) - .map_err(|_| AdminServiceError::invalid("Invalid observability cursor"))?; - serde_json::from_slice(&bytes) - .map_err(|_| AdminServiceError::invalid("Invalid observability cursor")) -} - -fn encode_cursor(cursor: &T) -> Result { - let bytes = serde_json::to_vec(cursor) - .map_err(|_| AdminServiceError::internal("Failed to encode observability cursor"))?; - Ok(URL_SAFE_NO_PAD.encode(bytes)) -} - -fn decode_observability_cursor( - value: Option<&str>, -) -> Result, AdminServiceError> { - value - .map(|encoded| { - let wire: CursorWire = decode_cursor(encoded)?; - ObservabilityCursor::new(wire.observed_at, wire.stable_id).map_err(map_invalid_store) - }) - .transpose() -} - -fn encode_observability_cursor(cursor: &ObservabilityCursor) -> Result { - encode_cursor(&CursorWire { - observed_at: cursor.observed_at, - stable_id: cursor.stable_id.clone(), - }) -} - -fn usage_range( - start: Option<&str>, - end: Option<&str>, -) -> Result { - let end = parse_wire_datetime(end) - .map_err(map_wire_error)? - .unwrap_or_else(Utc::now); - let start = parse_wire_datetime(start) - .map_err(map_wire_error)? - .unwrap_or(end - ChronoDuration::days(7)); - external_observability_range(start, end) -} - -fn dashboard_range( - start: Option, - end: Option, -) -> Result { - let end = parse_wire_datetime(end.as_deref()) - .map_err(map_wire_error)? - .unwrap_or_else(Utc::now); - let start = parse_wire_datetime(start.as_deref()) - .map_err(map_wire_error)? - .unwrap_or_else(|| china_day_start(end) - ChronoDuration::days(1)); - external_observability_range(start, end) -} - -fn dashboard_today_range( - start: Option, - end: Option, -) -> Result { - let end = parse_wire_datetime(end.as_deref()) - .map_err(map_wire_error)? - .unwrap_or_else(Utc::now); - let start = parse_wire_datetime(start.as_deref()) - .map_err(map_wire_error)? - .unwrap_or_else(|| china_day_start(end)); - external_observability_range(start, end) -} - -/// 外部管理查询最多覆盖 366 天；内部账号留存投影不经过此边界。 -pub fn external_observability_range( - start: DateTime, - end: DateTime, -) -> Result { - let duration = end.signed_duration_since(start); - if duration <= ChronoDuration::zero() || duration > ChronoDuration::days(366) { - return Err(AdminServiceError::invalid( - "Invalid observability time range", - )); - } - ObservabilityRange::new(start, end).map_err(map_invalid_store) -} - -fn map_invalid_store(error: StoreError) -> AdminServiceError { - map_store_error(error, "Observability query") -} - -fn map_wire_error(error: WireValidationError) -> AdminServiceError { - match error.field() { - "timeRange" => AdminServiceError::invalid("Invalid time range"), - "statusCode" => AdminServiceError::invalid("Status code must be between 100 and 599"), - "attemptIndex" => AdminServiceError::invalid("Attempt index is out of range"), - "kind" => AdminServiceError::invalid("Invalid dashboard trend kind"), - "dimension" => AdminServiceError::invalid("Invalid diagnostics dimension"), - "id" => AdminServiceError::invalid("Usage record ID is required"), - "page" | "pageSize" => AdminServiceError::invalid("Invalid Observability query"), - "cursor" => AdminServiceError::invalid("Invalid observability cursor"), - _ => AdminServiceError::invalid("Invalid observability query"), - } -} - -fn map_diagnostic_dimension( - dimension: WireDiagnosticDimension, -) -> (StoreDiagnosticDimension, &'static str) { - match dimension { - WireDiagnosticDimension::Model => (StoreDiagnosticDimension::Model, "model"), - WireDiagnosticDimension::Account => (StoreDiagnosticDimension::Account, "account"), - WireDiagnosticDimension::ApiKey => (StoreDiagnosticDimension::ApiKey, "apiKey"), - WireDiagnosticDimension::Provider => (StoreDiagnosticDimension::Provider, "provider"), - WireDiagnosticDimension::Transport => (StoreDiagnosticDimension::Transport, "transport"), - WireDiagnosticDimension::Failure => (StoreDiagnosticDimension::Failure, "failureClass"), - WireDiagnosticDimension::Status => (StoreDiagnosticDimension::Status, "status"), - } -} - -fn map_store_error(error: StoreError, resource: &'static str) -> AdminServiceError { - match error { - StoreError::NotFound { .. } => { - AdminServiceError::not_found(format!("{resource} not found")) - } - StoreError::InvalidData { .. } => AdminServiceError::invalid(format!("Invalid {resource}")), - StoreError::Conflict { .. } => AdminServiceError::conflict(format!("{resource} conflict")), - StoreError::Unavailable { .. } => { - tracing::error!(error = %error, resource, "Observability repository unavailable"); - AdminServiceError::unavailable("Observability repository unavailable") - } - } -} - -fn non_empty(value: Option) -> Option { - value - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -fn average(sum: u64, count: u64) -> Option { - (count > 0).then(|| sum / count) -} - -fn rate(numerator: u64, denominator: u64) -> f64 { - if denominator == 0 { - 0.0 - } else { - numerator as f64 / denominator as f64 - } -} - -fn optional_rate(numerator: u64, denominator: u64) -> Option { - (denominator > 0).then(|| rate(numerator, denominator)) -} - -fn display_rate(value: f64) -> String { - if value.is_finite() { - format!("{:.1}%", value * 100.0) - } else { - "—".to_owned() - } -} - -fn display_duration(value: Option) -> String { - format_duration_ms(value.and_then(|value| i64::try_from(value).ok())) -} - -fn china_offset() -> FixedOffset { - FixedOffset::east_opt(8 * 60 * 60).expect("China offset is valid") -} - -fn china_rfc3339(value: &DateTime) -> String { - value.with_timezone(&china_offset()).to_rfc3339() -} - -fn china_datetime(value: &DateTime) -> String { - value - .with_timezone(&china_offset()) - .format("%Y-%m-%d %H:%M:%S") - .to_string() -} - -fn china_day_start(value: DateTime) -> DateTime { - let offset = china_offset(); - let local = value.with_timezone(&offset); - let local_start = local - .date_naive() - .and_hms_opt(0, 0, 0) - .expect("midnight is valid"); - offset - .from_local_datetime(&local_start) - .single() - .expect("fixed offset has one local time") - .with_timezone(&Utc) -} - -fn china_quarter_hour_start(value: DateTime) -> DateTime { - let offset = china_offset(); - let local = value.with_timezone(&offset); - let local_start = local - .date_naive() - .and_hms_opt(local.hour(), local.minute() / 15 * 15, 0) - .expect("quarter hour is valid"); - offset - .from_local_datetime(&local_start) - .single() - .expect("fixed offset has one local time") - .with_timezone(&Utc) -} - -fn sum_points( - points: &[RequestMetricPoint], - start: DateTime, - end: DateTime, -) -> RequestMetrics { - let mut total = RequestMetrics::default(); - for metrics in points - .iter() - .filter(|point| point.bucket_start >= start && point.bucket_start < end) - .map(|point| &point.metrics) - { - total.request_count = total.request_count.saturating_add(metrics.request_count); - total.success_count = total.success_count.saturating_add(metrics.success_count); - total.failure_count = total.failure_count.saturating_add(metrics.failure_count); - total.input_tokens = total.input_tokens.saturating_add(metrics.input_tokens); - total.output_tokens = total.output_tokens.saturating_add(metrics.output_tokens); - total.cached_tokens = total.cached_tokens.saturating_add(metrics.cached_tokens); - total.total_tokens = total.total_tokens.saturating_add(metrics.total_tokens); - } - total -} - -fn display_dimension_name(value: &str) -> String { - match value { - "__none__" => "未知".to_owned(), - value => value.to_owned(), - } -} - -fn validate_url( - field: &'static str, - raw_url: &str, - allowed_schemes: &[&str], - reject_password: bool, -) -> Result<(), ConfigError> { - let url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?; - if !allowed_schemes.contains(&url.scheme()) || url.host_str().is_none() { - return Err(ConfigError::InvalidField(field)); - } - if reject_password && url.password().is_some() { - return Err(ConfigError::PasswordInUrl(field)); - } - Ok(()) -} - -fn validate_service_password( - field: &'static str, - password: &SecretValue, -) -> Result<(), ConfigError> { - let password = password.expose_secret(); - if password.len() != SERVICE_PASSWORD_HEX_LENGTH - || !password.bytes().all(|byte| byte.is_ascii_hexdigit()) - { - return Err(ConfigError::InvalidServicePassword(field)); - } - Ok(()) -} - -fn validate_admin_password(password: &str) -> Result<(), ConfigError> { - let password = password.trim(); - if password.len() < 12 - || password.contains('$') - || WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str()) - { - return Err(ConfigError::WeakAdminPassword); - } - Ok(()) -} - -fn validate_wire_profile(config: &WireProfileConfig) -> Result<(), ConfigError> { - for (field, value) in [ - ("wire_profile.originator", config.originator.as_str()), - ("wire_profile.codex_version", config.codex_version.as_str()), - ( - "wire_profile.desktop_version", - config.desktop_version.as_str(), - ), - ("wire_profile.desktop_build", config.desktop_build.as_str()), - ("wire_profile.os_type", config.os_type.as_str()), - ("wire_profile.os_version", config.os_version.as_str()), - ("wire_profile.arch", config.arch.as_str()), - ("wire_profile.terminal", config.terminal.as_str()), - ] { - validate_nonempty(field, value)?; - } - - if semver::Version::parse(&config.codex_version).is_err() { - return Err(ConfigError::InvalidField("wire_profile.codex_version")); - } - if !is_numeric_dotted_version(&config.desktop_version) { - return Err(ConfigError::InvalidField("wire_profile.desktop_version")); - } - if !config - .desktop_build - .bytes() - .all(|byte| byte.is_ascii_digit()) - { - return Err(ConfigError::InvalidField("wire_profile.desktop_build")); - } - Ok(()) -} - -fn is_numeric_dotted_version(value: &str) -> bool { - let mut parts = value.split('.'); - let valid_parts = parts - .by_ref() - .filter(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit())) - .count(); - valid_parts >= 2 && valid_parts == value.split('.').count() -} - -fn validate_logging(config: &LoggingConfig) -> Result<(), ConfigError> { - EnvFilter::try_new(&config.level).map_err(|_| ConfigError::InvalidField("logging.level"))?; - if !config.stdout && !config.file.enabled { - return Err(ConfigError::InvalidField("logging")); - } - if config.file.directory.as_os_str().is_empty() { - return Err(ConfigError::InvalidField("logging.file.directory")); - } - if config.file.retention_days == 0 { - return Err(ConfigError::InvalidField("logging.file.retention_days")); - } - if config.file.max_file_size_mb == 0 { - return Err(ConfigError::InvalidField("logging.file.max_file_size_mb")); - } - if config.file.max_files == 0 { - return Err(ConfigError::InvalidField("logging.file.max_files")); - } - Ok(()) -} - -fn validate_nonempty(field: &'static str, value: &str) -> Result<(), ConfigError> { - if value.trim().is_empty() { - return Err(ConfigError::InvalidField(field)); - } - Ok(()) -} - -fn connection_url_with_password( - raw_url: &str, - password: &SecretValue, - field: &'static str, -) -> Result { - let mut url = Url::parse(raw_url).map_err(|_| ConfigError::InvalidField(field))?; - url.set_password(Some(password.expose_secret())) - .map_err(|()| ConfigError::InvalidField(field))?; - Ok(SecretBox::new(Box::new(url.into()))) -} - -/// 启动配置加载错误。错误文本不会包含配置值或连接地址。 -#[derive(Debug, thiserror::Error, PartialEq, Eq)] -pub enum ConfigError { - /// 无法确定当前工作目录。 - #[error("failed to determine the current working directory")] - CurrentDirectory, - /// 从当前目录及父目录找不到配置文件。 - #[error("deploy/config.yaml was not found in the current directory or its parents")] - ConfigFileNotFound, - /// 配置路径没有父目录。 - #[error("configuration file path is invalid")] - InvalidConfigPath, - /// YAML 无法读取、解析或不符合结构；不透传解析器文本以避免泄露密码。 - #[error("failed to load configuration file `{path}`; check its permissions and YAML structure")] - InvalidDocument { - /// 配置文件路径。 - path: PathBuf, - }, - /// 配置格式版本不受支持。 - #[error("x-cpr.schema_version must be 1")] - UnsupportedSchemaVersion, - /// 字段值非法。 - #[error("configuration field `{0}` is invalid")] - InvalidField(&'static str), - /// 连接 URL 不允许内嵌密码。 - #[error("configuration field `{0}` must not contain a password; use its password field")] - PasswordInUrl(&'static str), - /// PostgreSQL 或 Redis 密码格式非法。 - #[error("configuration field `{0}` must be exactly 48 hexadecimal characters")] - InvalidServicePassword(&'static str), - /// 管理员初始化密码非法。 - #[error("admin.default_password must be strong, at least 12 characters, and contain no `$`")] - WeakAdminPassword, - /// Compose 拓扑覆盖非法。 - #[error("container topology override `{0}` is invalid")] - InvalidTopologyOverride(&'static str), -} - -use std::sync::Arc; - -use async_trait::async_trait; -use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; -use gateway_api::admin::{ - AdminServiceError, - auth::{ - AdminAuthAuditEvent, AdminAuthBackend, AdminAuthService, AdminBackendSession, - AdminRequestContext, AdminSessionResolver, DefaultAdminAuthService, - }, -}; -use gateway_store::postgres::{ - AdminAuditActorKind, AdminAuditEvent, AdminSecurityAuditRepository, RuntimeSettingsRepository, -}; -use gateway_store::{ConflictKind, Revision, StoreError}; -use subtle::ConstantTimeEq as _; -use uuid::Uuid; - -/// Admin auth domain 与两个 Store owner 之间的组合 adapter。 -pub struct StoreAdminAuthBackend { - security: Arc, - runtime_settings: Arc, - runtime_auth: Arc, -} - -impl StoreAdminAuthBackend { - #[must_use] - pub const fn new( - security: Arc, - runtime_settings: Arc, - runtime_auth: Arc, - ) -> Self { - Self { - security, - runtime_settings, - runtime_auth, - } - } -} - -#[async_trait] -impl AdminAuthBackend for StoreAdminAuthBackend { - async fn password_hash( - &self, - admin_user_id: &str, - ) -> Result, AdminServiceError> { - self.security - .password_hash(admin_user_id) - .await - .map_err(|_| AdminServiceError::unavailable("Admin security store unavailable")) - } - - async fn store_password_hash( - &self, - admin_user_id: &str, - password_hash: &str, - ) -> Result<(), AdminServiceError> { - self.security - .replace_password_hash(admin_user_id, password_hash) - .await - .map_err(|_| AdminServiceError::unavailable("Admin security store unavailable")) - } - - async fn admin_api_key(&self) -> Result, AdminServiceError> { - self.runtime_settings - .load_runtime_settings() - .await - .map(|settings| settings.admin_api_key) - .map_err(|_| AdminServiceError::unavailable("Runtime settings unavailable")) - } - - async fn load_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError> { - self.runtime_auth - .load_admin_session(session_id) - .await - .map(|session| { - session.map(|session| AdminBackendSession { - admin_user_id: session.admin_user_id, - expires_at: session.expires_at, - }) - }) - .map_err(|_| AdminServiceError::unavailable("Admin session store unavailable")) - } - - async fn store_admin_session( - &self, - session_id: &str, - session: &AdminBackendSession, - ) -> Result<(), AdminServiceError> { - self.runtime_auth - .store_admin_session( - session_id, - &gateway_store::redis::AdminSessionRecord { - admin_user_id: session.admin_user_id.clone(), - expires_at: session.expires_at, - }, - ) - .await - .map_err(|_| AdminServiceError::unavailable("Admin session store unavailable")) - } - - async fn delete_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError> { - self.runtime_auth - .delete_admin_session(session_id) - .await - .map(|session| { - session.map(|session| AdminBackendSession { - admin_user_id: session.admin_user_id, - expires_at: session.expires_at, - }) - }) - .map_err(|_| AdminServiceError::unavailable("Admin session store unavailable")) - } - - async fn login_source_is_throttled( - &self, - source: &str, - failure_limit: u32, - window_seconds: u64, - ) -> Result { - self.runtime_auth - .login_source_is_throttled(source, failure_limit, window_seconds) - .await - .map_err(|_| AdminServiceError::unavailable("Admin login state unavailable")) - } - - async fn record_login_failure( - &self, - source: &str, - failure_limit: u32, - window_seconds: u64, - ) -> Result { - self.runtime_auth - .record_login_failure(source, failure_limit, window_seconds) - .await - .map_err(|_| AdminServiceError::unavailable("Admin login state unavailable")) - } - - async fn clear_login_failures(&self, source: &str) -> Result<(), AdminServiceError> { - self.runtime_auth - .clear_login_failures(source) - .await - .map_err(|_| AdminServiceError::unavailable("Admin login state unavailable")) - } - - async fn append_auth_audit(&self, event: AdminAuthAuditEvent) -> Result<(), AdminServiceError> { - self.security - .append_admin_audit_event(AdminAuditEvent { - id: format!("audit_{}", Uuid::now_v7().simple()), - actor_kind: AdminAuditActorKind::AdminSession, - actor_admin_user_id: Some(event.admin_user_id.clone()), - actor_ref: format!("admin:{}", event.admin_user_id), - admin_request_id: None, - action: event.action.to_owned(), - entity_kind: "admin_session".to_owned(), - entity_ref: event.admin_user_id, - config_revision: None, - changed_fields: Vec::new(), - created_at: event.occurred_at, - }) - .await - .map_err(|_| AdminServiceError::unavailable("Admin audit store unavailable")) - } -} - -pub type AdminSessionService = DefaultAdminAuthService; -fn audit_actor(context: &AdminRequestContext) -> (AdminAuditActorKind, Option, String) { - if let Some(admin_user_id) = context.admin_user_id() { - ( - AdminAuditActorKind::AdminSession, - Some(admin_user_id.to_owned()), - format!("admin:{admin_user_id}"), - ) - } else { - ( - AdminAuditActorKind::AdminApiKey, - None, - "admin_api_key".to_owned(), - ) - } -} - -fn admin_audit_event( - context: &AdminRequestContext, - action: &'static str, - entity_kind: &'static str, - entity_ref: String, - changed_fields: &[&str], -) -> AdminAuditEvent { - let (actor_kind, actor_admin_user_id, actor_ref) = audit_actor(context); - AdminAuditEvent { - id: format!("audit_{}", Uuid::now_v7().simple()), - actor_kind, - actor_admin_user_id, - actor_ref, - admin_request_id: Some(context.request_id().to_owned()), - action: action.to_owned(), - entity_kind: entity_kind.to_owned(), - entity_ref, - config_revision: None, - changed_fields: changed_fields - .iter() - .map(|field| (*field).to_owned()) - .collect(), - created_at: Utc::now(), - } -} - -fn revision(value: u64) -> Result { - Revision::new(value) - .map_err(|_| AdminServiceError::invalid("expectedConfigRevision must be positive")) -} - -fn map_admin_store_error(error: StoreError, entity: &'static str) -> AdminServiceError { - match error { - StoreError::NotFound { .. } => { - AdminServiceError::not_found(format!("{entity} was not found")) - } - StoreError::Conflict { - kind: ConflictKind::StaleRevision, - .. - } => AdminServiceError::conflict("Configuration revision is stale; reload and retry"), - StoreError::Conflict { .. } => { - AdminServiceError::conflict(format!("{entity} state does not allow this operation")) - } - StoreError::InvalidData { .. } => { - AdminServiceError::invalid(format!("Invalid {entity} configuration")) - } - StoreError::Unavailable { .. } => { - AdminServiceError::unavailable("Configuration repository unavailable") - } - } -} - -use gateway_api::admin::{ - AdminSessionState, - accounts::{ - AccountAdminService, AccountAdminState, AccountConnectionTestEvent, - AccountConnectionTestEventStream, AccountExportData, AccountModelView, AccountModelsData, - AccountPageData, AccountQuotaData, AccountQuotaView, AccountQuotaWindowView, - AccountRefreshData, AccountRefreshRequest, AccountSort, AccountStatus, AccountSummaryView, - AccountUsageView, AccountView, CurrencyCostView, ModelUsageView, ProviderFilter, - SortDirection, SortField, ValidatedListQuery, - }, - auth::AdminAuthState, - catalog::{CatalogAdminService, CatalogAdminState}, - client_keys::{ClientKeyAdminService, ClientKeyAdminState}, - observability::{ObservabilityAdminService, ObservabilityAdminState}, - openai::{CodexAdminService, CodexAdminState}, - settings::{AdminSettingsService, AdminSettingsState}, - system::{SystemAdminService, SystemAdminState}, - xai::{XaiAdminService, XaiAdminState}, -}; -use gateway_api::openai::{OpenAiApiState, OpenAiClientService}; - -/// 旧账号页需要的统一只读目录；Provider 专属动作由后续同一 adapter 分派给 owner。 -pub struct AccountAdminAdapter { - accounts: Arc, - admin_accounts: Arc, - core_store: Arc, - control_plane: Arc, - instances: Arc, - observability: Arc, - security: Arc, - codex_owner: Arc, - codex_quota: Arc, - codex_catalog: Arc, - xai_repository: provider_xai::GrokCredentialRepository, - xai_refresh: Arc, - xai_quota: Arc, - xai_catalog: Arc, - connection_test: Arc, - publisher: RuntimeSnapshotPublisher, -} - -/// [`AccountAdminAdapter`] 的显式 owner 组合，避免位置参数错接。 -pub struct AccountAdminPorts { - pub accounts: Arc, - pub admin_accounts: Arc, - pub core_store: Arc, - pub control_plane: Arc, - pub instances: Arc, - pub observability: Arc, - pub security: Arc, - pub codex_owner: Arc, - pub codex_quota: Arc, - pub codex_catalog: Arc, - pub xai_repository: provider_xai::GrokCredentialRepository, - pub xai_refresh: Arc, - pub xai_quota: Arc, - pub xai_catalog: Arc, - pub connection_test: Arc, - pub publisher: RuntimeSnapshotPublisher, -} - -struct AccountListPresentationInput { - id: String, - account: gateway_store::postgres::ProviderAccountSummary, - instance: gateway_store::postgres::ProviderInstanceRecord, - rolling_usage: Option, -} - -impl AccountAdminAdapter { - #[must_use] - pub fn new(ports: AccountAdminPorts) -> Self { - Self { - accounts: ports.accounts, - admin_accounts: ports.admin_accounts, - core_store: ports.core_store, - control_plane: ports.control_plane, - instances: ports.instances, - observability: ports.observability, - security: ports.security, - codex_owner: ports.codex_owner, - codex_quota: ports.codex_quota, - codex_catalog: ports.codex_catalog, - xai_repository: ports.xai_repository, - xai_refresh: ports.xai_refresh, - xai_quota: ports.xai_quota, - xai_catalog: ports.xai_catalog, - connection_test: ports.connection_test, - publisher: ports.publisher, - } - } - - async fn load_account( - &self, - id: &str, - ) -> Result { - self.accounts - .load_provider_account(id) - .await - .map_err(|error| map_admin_store_error(error, "provider account"))? - .ok_or_else(|| AdminServiceError::not_found("Provider account was not found")) - } - - async fn load_instance( - &self, - id: &str, - ) -> Result { - self.instances - .get_provider_instance(id) - .await - .map_err(|error| map_admin_store_error(error, "Provider instance"))? - .ok_or_else(|| AdminServiceError::not_found("Provider instance was not found")) - } - - async fn load_account_usage_by_id( - &self, - range: ObservabilityRange, - account_ids: &[String], - ) -> Result< - BTreeMap, - AdminServiceError, - > { - let mut usage_by_account = BTreeMap::new(); - for account_ids in account_ids.chunks(200) { - let query = gateway_store::postgres::ProviderAccountUsageQuery::for_accounts( - range, - account_ids.to_vec(), - ) - .map_err(|error| map_admin_store_error(error, "account usage"))?; - let observations = self - .observability - .provider_account_usage(query) - .await - .map_err(|error| map_admin_store_error(error, "account usage"))?; - usage_by_account.extend( - observations - .into_iter() - .map(|observation| (observation.account_id.clone(), observation)), - ); - } - Ok(usage_by_account) - } - - async fn load_account_view( - &self, - id: &str, - refresh_quota: bool, - ) -> Result { - let record = self.load_account(id).await?; - let instance = self - .load_instance(&record.summary.provider_instance_id) - .await?; - let settings = self - .control_plane - .load_control_plane() - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))? - .settings; - let now = Utc::now(); - let usage_range = ObservabilityRange::new( - now - ChronoDuration::days(i64::from(settings.usage_retention_days)), - now, - ) - .map_err(|error| map_admin_store_error(error, "account usage"))?; - let rolling_range = ObservabilityRange::new(now - ChronoDuration::hours(24), now) - .map_err(|error| map_admin_store_error(error, "account usage"))?; - let account_ids = vec![id.to_owned()]; - let (mut usage_by_account, mut rolling_usage_by_account) = tokio::try_join!( - self.load_account_usage_by_id(usage_range, &account_ids), - self.load_account_usage_by_id(rolling_range, &account_ids), - )?; - let usage = usage_by_account.remove(id); - let rolling_usage = rolling_usage_by_account.remove(id); - let quota = self - .provider_quota_view( - &record.summary, - &instance, - refresh_quota, - rolling_usage.as_ref(), - ) - .await?; - let refresh_token_expires_at = if record.summary.provider_kind == "xai" { - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(id.to_owned()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - self.xai_repository - .read_lifecycle(&account_id) - .await - .map_err(map_xai_repository_error)? - .refresh_token_expires_at() - .map(china_rfc3339) - } else { - None - }; - let mut view = account_view(record.summary, instance.name, usage, now); - view.quota = quota; - view.refresh_token_expires_at = refresh_token_expires_at; - Ok(view) - } - - async fn provider_quota_view( - &self, - account: &gateway_store::postgres::ProviderAccountSummary, - instance: &gateway_store::postgres::ProviderInstanceRecord, - refresh: bool, - rolling_usage: Option<&gateway_store::postgres::ProviderAccountUsageObservation>, - ) -> Result { - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(account.id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - match account.provider_kind.as_str() { - "openai" => { - let snapshot = if refresh { - Some( - self.codex_quota - .refresh_account(&core_provider_instance(instance)?, &account_id) - .await - .map_err(map_codex_quota_error)?, - ) - } else { - self.codex_quota - .read_account(&account_id) - .await - .map_err(map_codex_quota_error)? - }; - Ok(snapshot.map_or_else(empty_quota_view, codex_quota_view)) - } - "xai" => { - let snapshot = if refresh { - Some( - self.xai_quota - .refresh_account(&account_id) - .await - .map_err(map_xai_quota_error)?, - ) - } else { - self.xai_quota - .read_account(&account_id) - .await - .map_err(map_xai_quota_error)? - }; - Ok(snapshot.map_or_else(empty_quota_view, |snapshot| { - xai_quota_view(snapshot, rolling_usage) - })) - } - _ => Ok(empty_quota_view()), - } - } - - async fn account_models( - &self, - id: &str, - refresh: bool, - ) -> Result { - let record = self.load_account(id).await?; - let instance = self - .load_instance(&record.summary.provider_instance_id) - .await?; - let account_id = gateway_core::engine::credential::ProviderAccountId::new(id.to_owned()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - let revision = CoreCredentialRevision::new(record.summary.credential_revision.get()) - .map_err(|_| AdminServiceError::invalid("Invalid credential revision"))?; - let models = match record.summary.provider_kind.as_str() { - "openai" if refresh => self - .codex_catalog - .synchronize_account(&core_provider_instance(&instance)?, &account_id) - .await - .map_err(map_codex_catalog_error)?, - "openai" => self - .codex_catalog - .cached_account_models( - &gateway_core::routing::ProviderInstanceId::new(instance.id) - .map_err(|_| AdminServiceError::invalid("Invalid Provider instance"))?, - &account_id, - ) - .map_err(map_codex_catalog_error)? - .unwrap_or_default(), - "xai" if refresh => self - .xai_catalog - .refresh_account_catalog(&account_id, provider_xai::transport::GROK_CLIENT_VERSION) - .await - .map_err(map_xai_catalog_error)? - .seed() - .models() - .to_vec(), - "xai" => self - .xai_catalog - .read_account_catalog(&account_id, revision) - .await - .map_err(map_xai_catalog_error)? - .map(|catalog| catalog.seed().models().to_vec()) - .unwrap_or_default(), - _ => Vec::new(), - }; - Ok(AccountModelsData { - models: models - .into_iter() - .map(|id| AccountModelView { - label: id.clone(), - id, - }) - .collect(), - }) - } - - async fn list_presentation( - &self, - input: AccountListPresentationInput, - ) -> Result<(String, AccountQuotaView, Option), AdminServiceError> { - let quota = self - .provider_quota_view( - &input.account, - &input.instance, - false, - input.rolling_usage.as_ref(), - ) - .await?; - let refresh_token_expires_at = if input.account.provider_kind == "xai" { - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(input.account.id) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - self.xai_repository - .read_lifecycle(&account_id) - .await - .map_err(map_xai_repository_error)? - .refresh_token_expires_at() - .map(china_rfc3339) - } else { - None - }; - Ok((input.id, quota, refresh_token_expires_at)) - } -} - -enum AccountRefreshGuard { - Codex(provider_openai::credential::PreparedCodexCredentialRotationGuard), - Xai(provider_xai::PreparedGrokCredentialRotationGuard), -} - -fn core_provider_instance( - record: &gateway_store::postgres::ProviderInstanceRecord, -) -> Result { - Ok(gateway_core::routing::ProviderInstance::new( - gateway_core::routing::ProviderInstanceId::new(record.id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid Provider instance ID"))?, - gateway_core::routing::ProviderKind::new(record.provider_kind.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid Provider kind"))?, - record.base_url.clone(), - record.enabled, - gateway_core::routing::InstanceHealth::Healthy, - )) -} - -fn empty_quota_view() -> AccountQuotaView { - AccountQuotaView { - refreshed_at_display: "—".to_owned(), - windows: Vec::new(), - } -} - -fn codex_quota_view( - snapshot: provider_openai::credential::CodexAccountQuotaSnapshot, -) -> AccountQuotaView { - let now = Utc::now(); - let observed_at = DateTime::::from(snapshot.observed_at()); - AccountQuotaView { - refreshed_at_display: relative_time(Some(observed_at), now), - windows: snapshot - .windows() - .iter() - .map(|window| { - let label = codex_quota_window_label(window); - let used = window.used_percent(); - AccountQuotaWindowView { - key: window.key().to_owned(), - group: codex_quota_window_group(window.kind()).to_owned(), - window_seconds: window.window_seconds(), - label_display: label, - used_percent: used, - used_percent_display: used - .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), - local_usage: None, - reset_at_display: window - .reset_at() - .map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), - } - }) - .collect(), - } -} - -fn codex_quota_window_group( - kind: provider_openai::credential::CodexQuotaWindowKind, -) -> &'static str { - match kind { - provider_openai::credential::CodexQuotaWindowKind::Monthly => "monthly", - provider_openai::credential::CodexQuotaWindowKind::ShortTerm - | provider_openai::credential::CodexQuotaWindowKind::Weekly => "shortTerm", - provider_openai::credential::CodexQuotaWindowKind::Other => "other", - } -} - -fn codex_quota_window_label(window: &provider_openai::credential::CodexQuotaWindow) -> String { - let base = match window.kind() { - provider_openai::credential::CodexQuotaWindowKind::ShortTerm => "5小时限额".to_owned(), - provider_openai::credential::CodexQuotaWindowKind::Weekly => "周限额".to_owned(), - provider_openai::credential::CodexQuotaWindowKind::Monthly => "月限额".to_owned(), - provider_openai::credential::CodexQuotaWindowKind::Other => { - custom_quota_window_label(window.window_seconds()) - } - }; - let source = window.source(); - if matches!(source, "core" | "spend_control" | "monthly_limit") { - return base; - } - let source = if is_codex_review_limit(source) { - "代码审查" - } else { - source - }; - format!("{source} · {base}") -} - -fn custom_quota_window_label(window_seconds: Option) -> String { - let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else { - return "额度".to_owned(); - }; - if seconds % 86_400 == 0 { - format!("{}天限额", seconds / 86_400) - } else if seconds % 3_600 == 0 { - format!("{}小时限额", seconds / 3_600) - } else { - format!("{}分钟限额", seconds.div_ceil(60)) - } -} - -fn is_codex_review_limit(value: &str) -> bool { - let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_"); - matches!( - normalized.as_str(), - "review" | "code_review" | "codex_review" | "codex_code_review" - ) || normalized.contains("code_review") - || normalized.contains("codex_review") -} - -fn xai_quota_view( - snapshot: provider_xai::GrokQuotaSnapshot, - rolling_usage: Option<&gateway_store::postgres::ProviderAccountUsageObservation>, -) -> AccountQuotaView { - let billing = snapshot.billing(); - if !billing.has_authoritative_quota() { - return xai_free_quota_view(snapshot.observed_at(), rolling_usage); - } - let used = billing.used_percent(); - let period = billing.period_type().unwrap_or("billing"); - let window_seconds = xai_quota_window_seconds(billing.period_start(), billing.period_end()); - AccountQuotaView { - refreshed_at_display: relative_time(Some(snapshot.observed_at()), Utc::now()), - windows: vec![AccountQuotaWindowView { - key: period.to_owned(), - group: xai_quota_window_group(billing.period_kind()).to_owned(), - window_seconds, - label_display: xai_quota_window_label(billing.period_kind(), window_seconds), - used_percent: used, - used_percent_display: used - .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), - local_usage: Some(serde_json::json!({ - "periodStart": billing.period_start(), - "periodEnd": billing.period_end(), - "monthlyLimitCents": billing.monthly_limit_cents(), - "includedUsedCents": billing.included_used_cents(), - "onDemandCapCents": billing.on_demand_cap_cents(), - "onDemandUsedCents": billing.on_demand_used_cents(), - "prepaidBalanceCents": billing.prepaid_balance_cents(), - })), - reset_at_display: billing - .period_end() - .and_then(parse_rfc3339_utc) - .map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), - }], - } -} - -fn xai_free_quota_view( - observed_at: DateTime, - rolling_usage: Option<&gateway_store::postgres::ProviderAccountUsageObservation>, -) -> AccountQuotaView { - let local_usage = rolling_usage.map(|usage| { - let total_tokens = usage.total_tokens.unwrap_or(0); - serde_json::json!({ - "requestCount": usage.request_count, - "requestCountDisplay": format_number(usage.request_count), - "inputTokens": usage.input_tokens.unwrap_or(0), - "inputTokensDisplay": display_optional_tokens(usage.input_tokens), - "outputTokens": usage.output_tokens.unwrap_or(0), - "outputTokensDisplay": display_optional_tokens(usage.output_tokens), - "cachedTokens": usage.cached_tokens.unwrap_or(0), - "cachedTokensDisplay": display_optional_tokens(usage.cached_tokens), - "totalTokens": total_tokens, - "totalTokensDisplay": format_tokens(total_tokens), - }) - }); - AccountQuotaView { - refreshed_at_display: relative_time(Some(observed_at), Utc::now()), - windows: vec![AccountQuotaWindowView { - key: "free-rolling-24h".to_owned(), - group: "shortTerm".to_owned(), - window_seconds: Some(provider_xai::GROK_FREE_ROLLING_WINDOW_SECONDS), - label_display: "24小时额度".to_owned(), - used_percent: None, - used_percent_display: "—".to_owned(), - local_usage, - reset_at_display: "—".to_owned(), - }], - } -} - -fn xai_quota_window_group(kind: provider_xai::GrokQuotaPeriodKind) -> &'static str { - match kind { - provider_xai::GrokQuotaPeriodKind::Weekly => "shortTerm", - provider_xai::GrokQuotaPeriodKind::Monthly => "monthly", - provider_xai::GrokQuotaPeriodKind::Other => "other", - } -} - -fn xai_quota_window_label( - kind: provider_xai::GrokQuotaPeriodKind, - window_seconds: Option, -) -> String { - match kind { - provider_xai::GrokQuotaPeriodKind::Weekly => "周限额".to_owned(), - provider_xai::GrokQuotaPeriodKind::Monthly => "月限额".to_owned(), - provider_xai::GrokQuotaPeriodKind::Other => custom_quota_window_label(window_seconds), - } -} - -fn xai_quota_window_seconds(start: Option<&str>, end: Option<&str>) -> Option { - let start = start.and_then(parse_rfc3339_utc)?; - let end = end.and_then(parse_rfc3339_utc)?; - end.signed_duration_since(start) - .num_seconds() - .try_into() - .ok() -} - -fn parse_rfc3339_utc(value: &str) -> Option> { - DateTime::parse_from_rfc3339(value) - .ok() - .map(|value| value.with_timezone(&Utc)) -} - -fn map_codex_quota_error( - error: provider_openai::credential::CodexCredentialQuotaError, -) -> AdminServiceError { - use provider_openai::credential::CodexCredentialQuotaError as Error; - match error { - Error::NotFound => AdminServiceError::not_found("Codex account was not found"), - Error::RevisionConflict => AdminServiceError::conflict("Codex quota snapshot is stale"), - Error::InvalidInstance | Error::InvalidCredentialData => { - AdminServiceError::invalid("Codex quota data is invalid") - } - Error::TransportInitialization | Error::Repository(_) | Error::Store | Error::Upstream => { - AdminServiceError::unavailable("Codex quota service is unavailable") - } - } -} - -fn map_xai_quota_error(error: provider_xai::GrokQuotaError) -> AdminServiceError { - use provider_xai::GrokQuotaError as Error; - match error { - Error::AccountUnavailable => AdminServiceError::not_found("xAI account is unavailable"), - Error::StaleCredentialSnapshot => { - AdminServiceError::conflict("xAI quota snapshot is stale") - } - Error::InvalidData => AdminServiceError::invalid("xAI quota data is invalid"), - Error::Upstream | Error::Store => { - AdminServiceError::unavailable("xAI quota service is unavailable") - } - } -} - -fn map_codex_catalog_error( - error: provider_openai::credential::CodexCredentialCatalogError, -) -> AdminServiceError { - use provider_openai::credential::CodexCredentialCatalogError as Error; - match error { - Error::NoEligibleCredential => { - AdminServiceError::not_found("Codex account has no available model catalog") - } - Error::InvalidInstance | Error::InvalidCredentialData | Error::ConflictingModelFacts => { - AdminServiceError::invalid("Codex model catalog is invalid") - } - Error::Upstream | Error::Cache => { - AdminServiceError::unavailable("Codex model catalog is unavailable") - } - } -} - -fn map_xai_catalog_error(error: provider_xai::GrokCredentialCatalogError) -> AdminServiceError { - use provider_xai::GrokCredentialCatalogError as Error; - match error { - Error::NoEligibleCredential => { - AdminServiceError::not_found("xAI account has no available model catalog") - } - Error::StaleCredentialSnapshot => { - AdminServiceError::conflict("xAI model catalog snapshot is stale") - } - Error::InvalidInstance | Error::InvalidCredentialData | Error::ConflictingModelFacts => { - AdminServiceError::invalid("xAI model catalog is invalid") - } - Error::Upstream | Error::Cache | Error::Store => { - AdminServiceError::unavailable("xAI model catalog is unavailable") - } - } -} - -fn map_xai_refresh_error(error: provider_xai::GrokCredentialRefreshError) -> AdminServiceError { - use provider_xai::GrokCredentialRefreshError as Error; - match error { - Error::InvalidConfiguration | Error::InvalidRefreshResponse => { - AdminServiceError::invalid("xAI credential refresh data is invalid") - } - Error::ManualFailure(provider_xai::GrokRefreshFailure::InvalidGrant) => { - AdminServiceError::invalid("xAI refresh token is expired") - } - Error::Repository(provider_xai::GrokCredentialRepositoryError::CredentialNotFound) => { - AdminServiceError::not_found("xAI account was not found") - } - Error::Repository(provider_xai::GrokCredentialRepositoryError::StaleCredentialRevision) => { - AdminServiceError::conflict("xAI credential revision is stale") - } - _ => AdminServiceError::unavailable("xAI credential refresh is unavailable"), - } -} - -#[async_trait] -impl AccountAdminService for AccountAdminAdapter { - async fn list(&self, query: ValidatedListQuery) -> Result { - let (control_plane, accounts, instances) = tokio::try_join!( - self.control_plane.load_control_plane(), - self.accounts.list_provider_accounts(None, true), - self.instances.list_provider_instances(true), - ) - .map_err(|error| map_admin_store_error(error, "account catalog"))?; - let instances_by_id = instances - .into_iter() - .map(|instance| (instance.id.clone(), instance)) - .collect::>(); - let now = Utc::now(); - let usage_range = ObservabilityRange::new( - now - ChronoDuration::days(i64::from(control_plane.settings.usage_retention_days)), - now, - ) - .map_err(|error| map_admin_store_error(error, "account usage"))?; - let rolling_range = ObservabilityRange::new(now - ChronoDuration::hours(24), now) - .map_err(|error| map_admin_store_error(error, "account usage"))?; - let account_ids = accounts - .iter() - .map(|account| account.id.clone()) - .collect::>(); - let (mut usage_by_account, mut rolling_usage_by_account) = tokio::try_join!( - self.load_account_usage_by_id(usage_range, &account_ids), - self.load_account_usage_by_id(rolling_range, &account_ids), - )?; - let accounts_by_id = accounts - .iter() - .cloned() - .map(|account| (account.id.clone(), account)) - .collect::>(); - let summary = account_summary(&accounts, now); - let mut items = accounts - .into_iter() - .filter(|account| account_matches_query(account, &query, now)) - .map(|account| { - let instance_name = instances_by_id - .get(&account.provider_instance_id) - .map(|instance| instance.name.clone()) - .unwrap_or_else(|| account.provider_instance_id.clone()); - let usage = usage_by_account.remove(&account.id); - account_view(account, instance_name, usage, now) - }) - .collect::>(); - sort_account_views(&mut items, query.sort); - let total = u64::try_from(items.len()).unwrap_or(u64::MAX); - let page_size = usize::try_from(query.page_size).unwrap_or(usize::MAX); - let offset = usize::try_from(query.page.saturating_sub(1)) - .unwrap_or(usize::MAX) - .saturating_mul(page_size); - let mut items = items - .into_iter() - .skip(offset) - .take(page_size) - .collect::>(); - let presentation_inputs = items - .iter() - .map(|view| { - let account = accounts_by_id.get(&view.id).cloned().ok_or_else(|| { - AdminServiceError::invalid("Provider account presentation is missing") - })?; - let instance = instances_by_id - .get(&account.provider_instance_id) - .cloned() - .ok_or_else(|| { - AdminServiceError::invalid("Provider account instance is missing") - })?; - Ok(AccountListPresentationInput { - id: view.id.clone(), - account, - instance, - rolling_usage: rolling_usage_by_account.remove(&view.id), - }) - }) - .collect::, AdminServiceError>>()?; - let presentations = futures::stream::iter( - presentation_inputs - .into_iter() - .map(|input| self.list_presentation(input)), - ) - .buffer_unordered(8) - .collect::>() - .await; - let mut presentations_by_id = BTreeMap::new(); - for presentation in presentations { - let (id, quota, refresh_token_expires_at) = presentation?; - presentations_by_id.insert(id, (quota, refresh_token_expires_at)); - } - for item in &mut items { - if let Some((quota, refresh_token_expires_at)) = presentations_by_id.remove(&item.id) { - item.quota = quota; - item.refresh_token_expires_at = refresh_token_expires_at; - } - } - let total_pages = if total == 0 { - 0 - } else { - u32::try_from(total.div_ceil(u64::from(query.page_size))).unwrap_or(u32::MAX) - }; - Ok(AccountPageData { - config_revision: control_plane.settings.config_revision.get(), - items, - page: gateway_api::admin::PageMeta::new( - query.page, - query.page_size, - total, - total_pages, - ), - summary, - }) - } - - async fn export( - &self, - context: &AdminRequestContext, - ids: Vec, - ) -> Result { - let mut codex = Vec::new(); - let mut xai = Vec::new(); - for id in &ids { - let record = self - .accounts - .load_provider_account(id) - .await - .map_err(|error| map_admin_store_error(error, "provider account"))? - .ok_or_else(|| AdminServiceError::not_found("Provider account was not found"))?; - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - let credential_revision = - CoreCredentialRevision::new(record.summary.credential_revision.get()) - .map_err(|_| AdminServiceError::invalid("Invalid credential revision"))?; - let loaded = self - .core_store - .load_credential(&account_id, credential_revision) - .await - .map_err(|_| AdminServiceError::conflict("Credential revision is stale"))?; - match record.summary.provider_kind.as_str() { - "openai" => codex.push(provider_openai::credential::ExportManagedCodexCredential { - current: loaded, - added_at: record.summary.created_at, - updated_at: record.summary.updated_at, - }), - "xai" => xai.push(loaded), - _ => { - return Err(AdminServiceError::invalid( - "Provider does not support account export", - )); - } - } - } - let mut documents = Vec::new(); - if !codex.is_empty() { - let document = provider_openai::credential::CodexCredentialAdmin - .format_cpr_export(codex) - .and_then(provider_openai::credential::CodexCprExportDocument::into_json) - .map_err(map_codex_admin_error)?; - documents.push(serde_json::json!({"provider": "openai", "document": document})); - } - if !xai.is_empty() { - let document = provider_xai::GrokCredentialAdmin - .export_oauth_bundle(&xai, Utc::now()) - .map_err(map_xai_repository_error)? - .into_value(); - documents.push(serde_json::json!({"provider": "xai", "document": document})); - } - self.security - .append_admin_audit_event(admin_audit_event( - context, - "export_sensitive", - "provider_account", - format!("{} accounts", ids.len()), - &[], - )) - .await - .map_err(|error| map_admin_store_error(error, "account export audit"))?; - Ok(AccountExportData::new(serde_json::json!({ - "exportedAt": Utc::now().to_rfc3339(), - "documents": documents, - }))) - } - - async fn refresh( - &self, - context: &AdminRequestContext, - request: AccountRefreshRequest, - ) -> Result { - let record = self.load_account(&request.id).await?; - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(request.id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - let credential_revision = - CoreCredentialRevision::new(record.summary.credential_revision.get()) - .map_err(|_| AdminServiceError::invalid("Invalid credential revision"))?; - let (profile, credential, guard) = match record.summary.provider_kind.as_str() { - "openai" => { - let (profile, credential, guard) = self - .codex_owner - .manual_refresh(account_id, credential_revision) - .await - .map_err(map_codex_admin_error)? - .into_parts(); - (profile, credential, AccountRefreshGuard::Codex(guard)) - } - "xai" => { - let (profile, credential, guard) = self - .xai_refresh - .prepare_manual_refresh(&account_id, credential_revision) - .await - .map_err(map_xai_refresh_error)? - .into_parts(); - (profile, credential, AccountRefreshGuard::Xai(guard)) - } - _ => { - return Err(AdminServiceError::invalid( - "Provider does not support OAuth refresh", - )); - } - }; - let rotation = self - .admin_accounts - .rotate_provider_account( - revision(request.expected_config_revision)?, - RotateProviderAccount { - scope: provider_scope( - record.summary.provider_kind, - record.summary.provider_instance_id, - ), - profile: store_profile(profile), - credential: store_credential_update(credential)?, - audit: admin_audit_event( - context, - "refresh", - "provider_account", - request.id.clone(), - &["provider_credentials_json", "credential_revision"], - ), - }, - ) - .await - .map_err(|error| map_admin_store_error(error, "provider account"))?; - match guard { - AccountRefreshGuard::Codex(guard) => drop(guard), - AccountRefreshGuard::Xai(guard) => drop(guard), - } - self.publisher - .publish_committed(rotation.config_revision) - .await; - Ok(AccountRefreshData { - config_revision: rotation.config_revision.get(), - account: self.load_account_view(&request.id, false).await?, - }) - } - - async fn quota(&self, id: String) -> Result { - Ok(AccountQuotaData { - account: self.load_account_view(&id, false).await?, - }) - } - - async fn refresh_quota(&self, id: String) -> Result { - Ok(AccountQuotaData { - account: self.load_account_view(&id, true).await?, - }) - } - - async fn models(&self, id: String) -> Result { - self.account_models(&id, false).await - } - - async fn refresh_models(&self, id: String) -> Result { - self.account_models(&id, true).await - } - - async fn test_connection( - &self, - id: String, - model_id: String, - ) -> Result { - let record = self.load_account(&id).await?; - let initial_account_status = - account_status_text(account_status(&record.summary, Utc::now())) - .0 - .to_owned(); - let provider_instance_id = - gateway_core::routing::ProviderInstanceId::new(record.summary.provider_instance_id) - .map_err(|_| AdminServiceError::invalid("Invalid Provider instance ID"))?; - let upstream_model = UpstreamModelId::new(model_id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid upstream model ID"))?; - let account_id = gateway_core::engine::credential::ProviderAccountId::new(id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid provider account ID"))?; - let initial_events = vec![ - AccountConnectionTestEvent::started(model_id.clone()), - AccountConnectionTestEvent::request(serde_json::json!({ - "model": model_id, - "input": [{ - "role": "user", - "content": [{ "type": "input_text", "text": "Reply with exactly OK." }] - }], - "stream": true, - "store": false - })), - ]; - let connection_test = Arc::clone(&self.connection_test); - let accounts = Arc::clone(&self.accounts); - let terminal_events = futures::stream::once(async move { - let result = connection_test - .test_account(account_id, provider_instance_id, upstream_model) - .await; - let account_status = accounts - .load_provider_account(&id) - .await - .ok() - .flatten() - .map(|account| { - account_status_text(account_status(&account.summary, Utc::now())) - .0 - .to_owned() - }) - .unwrap_or(initial_account_status); - let mut events = Vec::new(); - match result { - Ok(content) => { - events.extend(content.into_iter().map(AccountConnectionTestEvent::content)); - events.push(AccountConnectionTestEvent::completed(account_status)); - } - Err(error) => { - let message = match error.kind() { - GatewayErrorKind::InvalidRequest - | GatewayErrorKind::Unsupported - | GatewayErrorKind::ModelNotFound => { - AdminServiceError::invalid(error.to_string()) - } - _ => AdminServiceError::unavailable("Provider connection test failed"), - }; - events.push(AccountConnectionTestEvent::failed( - message.to_string(), - account_status, - )); - } - } - events - }) - .flat_map(futures::stream::iter); - Ok(Box::pin( - futures::stream::iter(initial_events).chain(terminal_events), - )) - } -} - -fn account_matches_query( - account: &gateway_store::postgres::ProviderAccountSummary, - query: &ValidatedListQuery, - now: DateTime, -) -> bool { - let provider_matches = match &query.provider { - ProviderFilter::All => true, - ProviderFilter::Provider(provider) => account.provider_kind == provider.as_str(), - }; - let search_matches = query.search.as_ref().is_none_or(|search| { - let search = search.to_lowercase(); - [ - Some(account.id.as_str()), - Some(account.name.as_str()), - account.email.as_deref(), - account.upstream_account_id.as_deref(), - Some(account.upstream_user_id.as_str()), - ] - .into_iter() - .flatten() - .any(|value| value.to_lowercase().contains(&search)) - }); - let status_matches = query - .status - .is_none_or(|status| account_status(account, now) == status); - provider_matches && search_matches && status_matches -} - -fn account_status( - account: &gateway_store::postgres::ProviderAccountSummary, - now: DateTime, -) -> AccountStatus { - if !account.enabled { - AccountStatus::Disabled - } else if account.availability == ProviderAccountAvailability::Banned { - AccountStatus::Banned - } else if account.availability == ProviderAccountAvailability::QuotaExhausted { - AccountStatus::QuotaExhausted - } else if account.availability == ProviderAccountAvailability::Expired - || account.access_token_expires_at <= now - { - AccountStatus::Expired - } else { - AccountStatus::Active - } -} - -fn account_status_text(status: AccountStatus) -> (&'static str, &'static str) { - match status { - AccountStatus::Active => ("active", "正常"), - AccountStatus::Expired => ("expired", "已过期"), - AccountStatus::QuotaExhausted => ("quota_exhausted", "额度耗尽"), - AccountStatus::Disabled => ("disabled", "已禁用"), - AccountStatus::Banned => ("banned", "已封禁"), - } -} - -fn account_summary( - accounts: &[gateway_store::postgres::ProviderAccountSummary], - now: DateTime, -) -> AccountSummaryView { - let total = u64::try_from(accounts.len()).unwrap_or(u64::MAX); - let active = u64::try_from( - accounts - .iter() - .filter(|account| account_status(account, now) == AccountStatus::Active) - .count(), - ) - .unwrap_or(u64::MAX); - let quota_exhausted = u64::try_from( - accounts - .iter() - .filter(|account| account_status(account, now) == AccountStatus::QuotaExhausted) - .count(), - ) - .unwrap_or(u64::MAX); - AccountSummaryView { - total, - active, - quota_exhausted, - attention: total.saturating_sub(active), - } -} - -fn account_view( - account: gateway_store::postgres::ProviderAccountSummary, - provider_instance_name: String, - usage: Option, - now: DateTime, -) -> AccountView { - let status = account_status(&account, now); - let (status, _) = account_status_text(status); - let expires_at = china_rfc3339(&account.access_token_expires_at); - let added_at = china_rfc3339(&account.created_at); - let updated_at = china_rfc3339(&account.updated_at); - AccountView { - id: account.id.clone(), - name: account.name, - provider: account.provider_kind, - provider_instance_id: account.provider_instance_id, - provider_instance_name, - resource_ref: account.id, - email: account.email, - account_id: account.upstream_account_id, - user_id: Some(account.upstream_user_id), - label: None, - plan_type: account.plan_type, - has_refresh_token: account.has_refresh_token, - status: status.to_owned(), - display_status: status.to_owned(), - token_refreshing: false, - availability: account.availability.as_str().to_owned(), - enabled: account.enabled, - credential_revision: account.credential_revision.get(), - state_revision: None, - access_token_expires_at: Some(expires_at), - access_token_expires_at_display: Some(china_datetime(&account.access_token_expires_at)), - refresh_token_expires_at: None, - next_refresh_at: account.next_refresh_at.map(|value| china_rfc3339(&value)), - added_at, - added_at_display: china_datetime(&account.created_at), - updated_at, - updated_at_display: china_datetime(&account.updated_at), - quota: AccountQuotaView { - refreshed_at_display: account - .quota_observed_at - .map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), - windows: Vec::new(), - }, - usage: account_usage_view(usage), - } -} - -fn account_usage_view( - usage: Option, -) -> AccountUsageView { - let Some(usage) = usage else { - return empty_account_usage(); - }; - let coverage = &usage.cost_coverage; - let known_count = coverage - .provider_reported_count - .saturating_add(coverage.calculated_count); - let cost_estimate_status = if known_count == 0 { - "unknown" - } else if coverage.unavailable_count > 0 { - "partial" - } else { - "known" - }; - let costs = usage.costs.iter().map(account_currency_cost_view).collect(); - let models = usage - .models - .into_iter() - .map(account_model_usage_view) - .collect(); - AccountUsageView { - request_count: Some(usage.request_count), - request_count_display: format_number(usage.request_count), - input_tokens: usage.input_tokens, - input_tokens_display: display_optional_tokens(usage.input_tokens), - output_tokens: usage.output_tokens, - output_tokens_display: display_optional_tokens(usage.output_tokens), - cached_tokens: usage.cached_tokens, - cached_tokens_display: display_optional_tokens(usage.cached_tokens), - total_tokens: usage.total_tokens, - total_tokens_display: display_optional_tokens(usage.total_tokens), - created_tokens: usage.cache_write_tokens, - created_tokens_display: display_optional_tokens(usage.cache_write_tokens), - read_tokens: usage.cached_tokens, - read_tokens_display: display_optional_tokens(usage.cached_tokens), - last_used_at: usage.last_used_at.map(|value| china_rfc3339(&value)), - last_used_at_display: usage.last_used_at.map_or_else( - || "—".to_owned(), - |value| relative_time(Some(value), Utc::now()), - ), - cost_estimate_status: cost_estimate_status.to_owned(), - known_cost_count: Some(known_count), - partial_cost_count: Some(if cost_estimate_status == "partial" { - 1 - } else { - 0 - }), - unknown_cost_count: Some(coverage.unavailable_count), - costs, - models, - } -} - -fn account_model_usage_view( - usage: gateway_store::postgres::ProviderAccountModelUsageObservation, -) -> ModelUsageView { - let known_count = usage - .cost_coverage - .provider_reported_count - .saturating_add(usage.cost_coverage.calculated_count); - let cost_estimate_status = if known_count == 0 { - "unknown" - } else if usage.cost_coverage.unavailable_count > 0 { - "partial" - } else { - "known" - }; - let usd = usage - .costs - .iter() - .find(|cost| cost.currency.eq_ignore_ascii_case("USD")); - ModelUsageView { - model: usage.model, - request_count: usage.request_count, - request_count_display: format_number(usage.request_count), - success_rate: (usage.request_count > 0) - .then(|| usage.success_count as f64 * 100.0 / usage.request_count as f64), - success_rate_display: if usage.request_count == 0 { - "—".to_owned() - } else { - format!( - "{:.1}%", - usage.success_count as f64 * 100.0 / usage.request_count as f64 - ) - }, - input_tokens: usage.input_tokens, - input_tokens_display: display_optional_tokens(usage.input_tokens), - output_tokens: usage.output_tokens, - output_tokens_display: display_optional_tokens(usage.output_tokens), - cached_tokens: usage.cached_tokens, - cached_tokens_display: display_optional_tokens(usage.cached_tokens), - total_tokens: usage.total_tokens, - total_tokens_display: display_optional_tokens(usage.total_tokens), - billing_amount_usd: usd.map(|cost| cost.amount.as_str().to_owned()), - billing_amount_usd_display: usd.map_or_else( - || "—".to_owned(), - |cost| format_billing_amount(&cost.amount), - ), - cost_estimate_status: cost_estimate_status.to_owned(), - known_cost_count: known_count, - partial_cost_count: u64::from(cost_estimate_status == "partial"), - unknown_cost_count: usage.cost_coverage.unavailable_count, - costs: usage.costs.iter().map(account_currency_cost_view).collect(), - last_used_at: china_rfc3339(&usage.last_used_at), - last_used_at_display: relative_time(Some(usage.last_used_at), Utc::now()), - } -} - -fn account_currency_cost_view(cost: &CurrencyCostTotal) -> CurrencyCostView { - CurrencyCostView { - currency: cost.currency.clone(), - estimated_amount: cost.amount.as_str().to_owned(), - estimated_amount_display: format!("{} {}", cost.currency, cost.amount.as_str()), - } -} - -fn display_optional_tokens(value: Option) -> String { - value.map_or_else(|| "—".to_owned(), format_tokens) -} - -fn empty_account_usage() -> AccountUsageView { - AccountUsageView { - request_count: None, - request_count_display: "—".to_owned(), - input_tokens: None, - input_tokens_display: "—".to_owned(), - output_tokens: None, - output_tokens_display: "—".to_owned(), - cached_tokens: None, - cached_tokens_display: "—".to_owned(), - total_tokens: None, - total_tokens_display: "—".to_owned(), - created_tokens: None, - created_tokens_display: "—".to_owned(), - read_tokens: None, - read_tokens_display: "—".to_owned(), - last_used_at: None, - last_used_at_display: "—".to_owned(), - cost_estimate_status: "unavailable".to_owned(), - known_cost_count: None, - partial_cost_count: None, - unknown_cost_count: None, - costs: Vec::new(), - models: Vec::new(), - } -} - -fn sort_account_views(items: &mut [AccountView], sort: Option) { - let Some(sort) = sort else { - items.sort_by(|left, right| left.id.cmp(&right.id)); - return; - }; - items.sort_by(|left, right| { - let ordering = match sort.field { - SortField::Email => left.email.cmp(&right.email), - SortField::Status => left.status.cmp(&right.status), - SortField::PlanType => left.plan_type.cmp(&right.plan_type), - SortField::Usage => left.usage.total_tokens.cmp(&right.usage.total_tokens), - SortField::LastUsedAt => left.usage.last_used_at.cmp(&right.usage.last_used_at), - SortField::ExpiresAt => left - .access_token_expires_at - .cmp(&right.access_token_expires_at), - } - .then_with(|| left.id.cmp(&right.id)); - match sort.direction { - SortDirection::Asc => ordering, - SortDirection::Desc => ordering.reverse(), - } - }); -} - -/// 管理端各领域端口的只读组合。 -pub struct AdminServices { - sessions: Arc, - accounts: Arc, - catalog: Arc, - client_keys: Arc, - codex: Arc, - observability: Arc, - settings: Arc, - system: Arc, - xai: Arc, -} - -/// [`AdminServices`] 的显式构造参数，避免位置参数错接领域端口。 -pub struct AdminServicePorts { - pub sessions: Arc, - pub accounts: Arc, - pub catalog: Arc, - pub client_keys: Arc, - pub codex: Arc, - pub observability: Arc, - pub settings: Arc, - pub system: Arc, - pub xai: Arc, -} - -impl AdminServices { - #[must_use] - pub fn new(ports: AdminServicePorts) -> Self { - Self { - sessions: ports.sessions, - accounts: ports.accounts, - catalog: ports.catalog, - client_keys: ports.client_keys, - codex: ports.codex, - observability: ports.observability, - settings: ports.settings, - system: ports.system, - xai: ports.xai, - } - } -} - -use gateway_api::admin::catalog::{ - CatalogListQuery, CatalogMutationData, CatalogRevisionRequest, CreateProviderInstanceRequest, - ProviderInstanceDetailData, ProviderInstanceListData, ProviderInstanceView, - UpdateProviderInstanceRequest, -}; -use gateway_store::postgres::{ - ConfigCatalogRepository, ControlPlaneReplacement, ControlPlaneRepository, ControlPlaneSnapshot, - NewProviderInstance, RuntimeSettings, RuntimeSettingsUpdate, UpdateProviderInstanceDetails, -}; - -/// Provider instance 的 PostgreSQL 管理投影。 -pub struct CatalogAdminAdapter { - control_plane: Arc, - instances: Arc, - publisher: RuntimeSnapshotPublisher, -} - -impl CatalogAdminAdapter { - #[must_use] - pub const fn new( - control_plane: Arc, - instances: Arc, - publisher: RuntimeSnapshotPublisher, - ) -> Self { - Self { - control_plane, - instances, - publisher, - } - } - - async fn revision(&self) -> Result { - self.control_plane - .load_control_plane() - .await - .map(|snapshot| snapshot.settings.config_revision.get()) - .map_err(|error| map_admin_store_error(error, "catalog")) - } - - async fn provider_instance_mutation( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - enabled: Option, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let id = request.id; - let (action, changed_fields) = match enabled { - Some(true) => ("enable", &["enabled"][..]), - Some(false) => ("disable", &["enabled"][..]), - None => ("delete", &[][..]), - }; - let audit = admin_audit_event( - context, - action, - "provider_instance", - id.clone(), - changed_fields, - ); - let committed = match enabled { - Some(enabled) => { - self.control_plane - .set_provider_instance_enabled(expected, &id, enabled, audit) - .await - } - None => { - self.control_plane - .delete_provider_instance(expected, &id, audit) - .await - } - } - .map_err(|error| map_admin_store_error(error, "provider instance"))?; - self.publisher.publish_committed(committed).await; - Ok(CatalogMutationData { - config_revision: committed.get(), - id, - }) - } -} - -#[async_trait] -impl CatalogAdminService for CatalogAdminAdapter { - async fn list_provider_instances( - &self, - query: CatalogListQuery, - ) -> Result { - let revision = self.revision().await?; - let records = self - .instances - .list_provider_instances(true) - .await - .map_err(|error| map_admin_store_error(error, "provider instance"))?; - let (records, next_cursor) = paginate_by_id( - records, - query.cursor.as_deref(), - usize::from(query.page_size()), - |record| &record.id, - )?; - Ok(ProviderInstanceListData { - config_revision: revision, - items: records.into_iter().map(provider_instance_view).collect(), - next_cursor, - }) - } - - async fn provider_instance( - &self, - id: String, - ) -> Result { - let revision = self.revision().await?; - let item = self - .instances - .get_provider_instance(&id) - .await - .map_err(|error| map_admin_store_error(error, "provider instance"))? - .ok_or_else(|| AdminServiceError::not_found("Provider instance was not found"))?; - Ok(ProviderInstanceDetailData { - config_revision: revision, - item: provider_instance_view(item), - }) - } - - async fn create_provider_instance( - &self, - context: &AdminRequestContext, - request: CreateProviderInstanceRequest, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let id = request.id; - let committed = self - .control_plane - .create_provider_instance( - expected, - NewProviderInstance { - id: id.clone(), - provider_kind: request.provider_kind, - name: request.name, - base_url: request.base_url, - }, - admin_audit_event( - context, - "create", - "provider_instance", - id.clone(), - &["provider_kind", "name", "base_url", "enabled"], - ), - ) - .await - .map_err(|error| map_admin_store_error(error, "provider instance"))?; - self.publisher.publish_committed(committed).await; - Ok(CatalogMutationData { - config_revision: committed.get(), - id, - }) - } - - async fn update_provider_instance( - &self, - context: &AdminRequestContext, - request: UpdateProviderInstanceRequest, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let id = request.id; - let committed = self - .control_plane - .update_provider_instance( - expected, - UpdateProviderInstanceDetails { - id: id.clone(), - name: request.name, - base_url: request.base_url, - }, - admin_audit_event( - context, - "update", - "provider_instance", - id.clone(), - &["name", "base_url"], - ), - ) - .await - .map_err(|error| map_admin_store_error(error, "provider instance"))?; - self.publisher.publish_committed(committed).await; - Ok(CatalogMutationData { - config_revision: committed.get(), - id, - }) - } - - async fn enable_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result { - self.provider_instance_mutation(context, request, Some(true)) - .await - } - - async fn disable_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result { - self.provider_instance_mutation(context, request, Some(false)) - .await - } - - async fn delete_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result { - self.provider_instance_mutation(context, request, None) - .await - } -} - -fn runtime_settings_update(settings: &RuntimeSettings) -> RuntimeSettingsUpdate { - RuntimeSettingsUpdate { - admin_api_key: settings.admin_api_key.clone(), - refresh_margin_seconds: settings.refresh_margin_seconds, - refresh_concurrency: settings.refresh_concurrency, - max_concurrent_per_account: settings.max_concurrent_per_account, - request_interval_ms: settings.request_interval_ms, - rotation_strategy: settings.rotation_strategy.clone(), - provider_model_mappings: settings.provider_model_mappings.clone(), - usage_retention_days: settings.usage_retention_days, - ops_event_retention_days: settings.ops_event_retention_days, - audit_retention_days: settings.audit_retention_days, - } -} - -fn provider_instance_view( - record: gateway_store::postgres::ProviderInstanceRecord, -) -> ProviderInstanceView { - ProviderInstanceView { - id: record.id, - provider_kind: record.provider_kind, - name: record.name, - base_url: record.base_url, - enabled: record.enabled, - created_at: record.created_at, - updated_at: record.updated_at, - } -} - -fn paginate_by_id( - items: Vec, - cursor: Option<&str>, - limit: usize, - id: F, -) -> Result<(Vec, Option), AdminServiceError> -where - F: Fn(&T) -> &str, -{ - let start = match cursor { - Some(cursor) => items - .iter() - .position(|item| id(item) == cursor) - .map(|index| index + 1) - .ok_or_else(|| AdminServiceError::invalid("Invalid catalog cursor"))?, - None => 0, - }; - let has_more = items.len().saturating_sub(start) > limit; - let page = items - .into_iter() - .skip(start) - .take(limit) - .collect::>(); - let next_cursor = has_more - .then(|| page.last().map(|item| id(item).to_owned())) - .flatten(); - Ok((page, next_cursor)) -} - -use gateway_api::admin::settings::{RuntimeSettingsView, UpdateRuntimeSettingsRequest}; -use rand_core::{OsRng, RngCore as _}; - -/// Runtime settings 与旧设置页只读 route 投影的组合 adapter。 -pub struct RuntimeSettingsAdminAdapter { - control_plane: Arc, - publisher: RuntimeSnapshotPublisher, -} - -impl RuntimeSettingsAdminAdapter { - #[must_use] - pub const fn new( - control_plane: Arc, - publisher: RuntimeSnapshotPublisher, - ) -> Self { - Self { - control_plane, - publisher, - } - } - - async fn replace_snapshot( - &self, - context: &AdminRequestContext, - _snapshot: ControlPlaneSnapshot, - expected_revision: Revision, - settings: RuntimeSettingsUpdate, - action: &'static str, - changed_fields: &[&str], - ) -> Result { - let replacement = ControlPlaneReplacement { - settings, - audit: admin_audit_event( - context, - action, - "runtime_settings", - "1".to_owned(), - changed_fields, - ), - }; - let committed = self - .control_plane - .replace_control_plane(expected_revision, replacement) - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))?; - self.publisher - .publish_committed(committed.settings.config_revision) - .await; - Ok(committed) - } -} - -#[async_trait] -impl AdminSettingsService for RuntimeSettingsAdminAdapter { - async fn load(&self) -> Result { - self.control_plane - .load_control_plane() - .await - .map(settings_view) - .map_err(|error| map_admin_store_error(error, "runtime settings")) - } - - async fn replace( - &self, - context: &AdminRequestContext, - request: UpdateRuntimeSettingsRequest, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let snapshot = self - .control_plane - .load_control_plane() - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))?; - let settings = RuntimeSettingsUpdate { - admin_api_key: snapshot.settings.admin_api_key.clone(), - refresh_margin_seconds: request.refresh_margin_seconds, - refresh_concurrency: to_u32_admin(request.refresh_concurrency, "refreshConcurrency")?, - max_concurrent_per_account: to_u32_admin( - request.max_concurrent_per_account, - "maxConcurrentPerAccount", - )?, - request_interval_ms: request.request_interval_ms, - rotation_strategy: request.rotation_strategy, - provider_model_mappings: request.provider_model_mappings, - usage_retention_days: to_u32_admin(request.usage_retention_days, "usageRetentionDays")?, - ops_event_retention_days: to_u32_admin( - request.ops_event_retention_days, - "opsEventRetentionDays", - )?, - audit_retention_days: to_u32_admin(request.audit_retention_days, "auditRetentionDays")?, - }; - self.replace_snapshot( - context, - snapshot, - expected, - settings, - "update", - &[ - "refresh_margin_seconds", - "refresh_concurrency", - "max_concurrent_per_account", - "request_interval_ms", - "rotation_strategy", - "provider_model_mappings", - "usage_retention_days", - "ops_event_retention_days", - "audit_retention_days", - ], - ) - .await - .map(settings_view) - } - - async fn admin_api_key_exists(&self) -> Result { - self.control_plane - .load_control_plane() - .await - .map(|snapshot| snapshot.settings.admin_api_key.is_some()) - .map_err(|error| map_admin_store_error(error, "runtime settings")) - } - - async fn regenerate_admin_api_key( - &self, - context: &AdminRequestContext, - ) -> Result { - let snapshot = self - .control_plane - .load_control_plane() - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))?; - let expected = snapshot.settings.config_revision; - let mut bytes = [0_u8; 32]; - OsRng.fill_bytes(&mut bytes); - let key = format!("admin-{}", hex::encode(bytes)); - let mut settings = runtime_settings_update(&snapshot.settings); - settings.admin_api_key = Some(key.clone()); - self.replace_snapshot( - context, - snapshot, - expected, - settings, - "regenerate_admin_api_key", - &["admin_api_key"], - ) - .await?; - Ok(key) - } - - async fn delete_admin_api_key( - &self, - context: &AdminRequestContext, - ) -> Result<(), AdminServiceError> { - let snapshot = self - .control_plane - .load_control_plane() - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))?; - if snapshot.settings.admin_api_key.is_none() { - return Ok(()); - } - let expected = snapshot.settings.config_revision; - let mut settings = runtime_settings_update(&snapshot.settings); - settings.admin_api_key = None; - self.replace_snapshot( - context, - snapshot, - expected, - settings, - "delete_admin_api_key", - &["admin_api_key"], - ) - .await?; + self.host.resolve_and_validate(source_dir)?; + self.store + .resolve_and_validate(source_dir) + .map_err(|_| ConfigError::InvalidField("store"))?; + self.admin + .resolve_and_validate(source_dir) + .map_err(|_| ConfigError::InvalidField("admin"))?; + self.api + .resolve_and_validate(source_dir) + .map_err(|_| ConfigError::InvalidField("api"))?; + self.openai + .resolve_and_validate(source_dir) + .map_err(|_| ConfigError::InvalidField("openai"))?; + self.xai + .resolve_and_validate(source_dir) + .map_err(|_| ConfigError::InvalidField("xai"))?; Ok(()) } } -fn settings_view(snapshot: ControlPlaneSnapshot) -> RuntimeSettingsView { - RuntimeSettingsView { - config_revision: snapshot.settings.config_revision.get(), - provider_model_mappings: snapshot.settings.provider_model_mappings, - refresh_margin_seconds: snapshot.settings.refresh_margin_seconds, - refresh_concurrency: u64::from(snapshot.settings.refresh_concurrency), - max_concurrent_per_account: u64::from(snapshot.settings.max_concurrent_per_account), - request_interval_ms: snapshot.settings.request_interval_ms, - rotation_strategy: snapshot.settings.rotation_strategy, - usage_retention_days: u64::from(snapshot.settings.usage_retention_days), - ops_event_retention_days: u64::from(snapshot.settings.ops_event_retention_days), - audit_retention_days: u64::from(snapshot.settings.audit_retention_days), - updated_at: snapshot.settings.updated_at, - } -} - -fn to_u32_admin(value: u64, field: &'static str) -> Result { - u32::try_from(value).map_err(|_| AdminServiceError::invalid(format!("Invalid {field}"))) -} - -use gateway_api::admin::client_keys::{ - ClientKeyCursorData, ClientKeyCursorValue, ClientKeyListData, ClientKeySort, - ClientKeySortDirection, ClientKeySortField, ClientKeyView, ClientKeyViewFields, - CreateClientKeyFields, CreatedClientKeyData, ListClientKeysFields, MutatedClientKeyData, - RevealedClientKeyData, UpdateClientKeyFields, decode_client_key_cursor, - encode_client_key_cursor, -}; -use gateway_store::postgres::{ - ClientApiKeyCursor as StoreClientKeyCursor, - ClientApiKeyCursorValue as StoreClientKeyCursorValue, ClientApiKeyListQuery, - ClientApiKeyRecord, ClientApiKeyRepository, ClientApiKeySort as StoreClientKeySort, - ClientApiKeySortDirection as StoreClientKeySortDirection, - ClientApiKeySortField as StoreClientKeySortField, NewClientApiKey, UpdateClientApiKeyDetails, -}; - -const DEFAULT_CLIENT_KEY_PAGE_SIZE: u16 = 50; - -/// 明文 Client API Key 的安全列表、CAS 管理与显式 reveal adapter。 -pub struct ClientKeyAdminAdapter { - control_plane: Arc, - repository: Arc, - publisher: RuntimeSnapshotPublisher, -} - -impl ClientKeyAdminAdapter { - #[must_use] - pub const fn new( - control_plane: Arc, - repository: Arc, - publisher: RuntimeSnapshotPublisher, - ) -> Self { - Self { - control_plane, - repository, - publisher, - } - } - - async fn mutate_enabled( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - enabled: Option, - ) -> Result { - let expected = revision(expected_config_revision)?; - let (action, fields) = match enabled { - Some(true) => ("enable", &["enabled"][..]), - Some(false) => ("disable", &["enabled"][..]), - None => ("delete", &[][..]), - }; - let audit = admin_audit_event(context, action, "client_api_key", id.clone(), fields); - let committed = match enabled { - Some(enabled) => { - self.control_plane - .set_client_api_key_enabled(expected, &id, enabled, audit) - .await - } - None => { - self.control_plane - .delete_client_api_key(expected, &id, audit) - .await - } - } - .map_err(|error| map_admin_store_error(error, "client API key"))?; - self.publisher.publish_committed(committed).await; - Ok(MutatedClientKeyData::new(committed.get(), id)) - } -} - -#[async_trait] -impl ClientKeyAdminService for ClientKeyAdminAdapter { - async fn list( - &self, - query: ListClientKeysFields, - ) -> Result { - let sort = store_client_key_sort(query.sort); - let cursor = query - .cursor - .as_deref() - .map(decode_client_key_cursor) - .transpose() - .map_err(|_| AdminServiceError::invalid("Invalid client key cursor"))? - .map(|cursor| store_client_key_cursor(cursor, sort)) - .transpose() - .map_err(|error| map_admin_store_error(error, "client API key cursor"))?; - let page = self - .repository - .list_client_api_keys(ClientApiKeyListQuery { - cursor, - page_size: query.limit.unwrap_or(DEFAULT_CLIENT_KEY_PAGE_SIZE), - search: query.search, - sort, - }) - .await - .map_err(|error| map_admin_store_error(error, "client API key"))?; - let next_cursor = page - .next_cursor - .map(|cursor| encode_client_key_cursor(&wire_client_key_cursor(cursor))) - .transpose() - .map_err(|_| AdminServiceError::internal("Failed to encode client key cursor"))?; - let config_revision = self - .control_plane - .load_control_plane() - .await - .map_err(|error| map_admin_store_error(error, "runtime settings"))? - .settings - .config_revision - .get(); - Ok(ClientKeyListData::new( - config_revision, - page.items.into_iter().map(client_key_view).collect(), - next_cursor, - page.total, - )) - } - - async fn create( - &self, - context: &AdminRequestContext, - fields: CreateClientKeyFields, - ) -> Result { - let expected = revision(fields.expected_config_revision)?; - let id = format!("key_{}", Uuid::now_v7().simple()); - let mut bytes = [0_u8; 32]; - OsRng.fill_bytes(&mut bytes); - let plaintext_key = format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes)); - let prefix = plaintext_key.chars().take(10).collect::(); - let committed = self - .control_plane - .create_client_api_key( - expected, - NewClientApiKey { - id: id.clone(), - name: fields.name, - label: fields.label, - provider_kind: fields.provider_kind, - key: plaintext_key.clone(), - max_concurrency: fields.max_concurrency, - requests_per_minute: fields.requests_per_minute, - tokens_per_minute: fields.tokens_per_minute, - }, - admin_audit_event( - context, - "create", - "client_api_key", - id.clone(), - &[ - "name", - "label", - "provider_kind", - "key", - "enabled", - "max_concurrency", - "requests_per_minute", - "tokens_per_minute", - ], - ), - ) - .await - .map_err(|error| map_admin_store_error(error, "client API key"))?; - self.publisher.publish_committed(committed).await; - Ok(CreatedClientKeyData::new( - committed.get(), - id, - prefix, - plaintext_key, - )) - } - - async fn reveal(&self, id: String) -> Result { - let secret = self - .repository - .reveal_client_api_key(&id) - .await - .map_err(|error| map_admin_store_error(error, "client API key"))? - .ok_or_else(|| AdminServiceError::not_found("Client API key was not found"))?; - Ok(RevealedClientKeyData::new(secret.id, secret.key)) - } - - async fn update( - &self, - context: &AdminRequestContext, - fields: UpdateClientKeyFields, - ) -> Result { - let expected = revision(fields.expected_config_revision)?; - let id = fields.id; - let committed = self - .control_plane - .update_client_api_key( - expected, - UpdateClientApiKeyDetails { - id: id.clone(), - name: fields.name, - label: fields.label, - provider_kind: fields.provider_kind, - max_concurrency: fields.max_concurrency, - requests_per_minute: fields.requests_per_minute, - tokens_per_minute: fields.tokens_per_minute, - }, - admin_audit_event( - context, - "update", - "client_api_key", - id.clone(), - &[ - "name", - "label", - "provider_kind", - "max_concurrency", - "requests_per_minute", - "tokens_per_minute", - ], - ), - ) - .await - .map_err(|error| map_admin_store_error(error, "client API key"))?; - self.publisher.publish_committed(committed).await; - Ok(MutatedClientKeyData::new(committed.get(), id)) - } - - async fn disable( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result { - self.mutate_enabled(context, id, expected_config_revision, Some(false)) - .await - } - - async fn enable( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result { - self.mutate_enabled(context, id, expected_config_revision, Some(true)) - .await - } - - async fn delete( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result { - self.mutate_enabled(context, id, expected_config_revision, None) - .await - } -} - -const fn store_client_key_sort(sort: ClientKeySort) -> StoreClientKeySort { - StoreClientKeySort { - field: match sort.field { - ClientKeySortField::Name => StoreClientKeySortField::Name, - ClientKeySortField::Enabled => StoreClientKeySortField::Enabled, - ClientKeySortField::CreatedAt => StoreClientKeySortField::CreatedAt, - ClientKeySortField::LastUsedAt => StoreClientKeySortField::LastUsedAt, - }, - direction: match sort.direction { - ClientKeySortDirection::Asc => StoreClientKeySortDirection::Asc, - ClientKeySortDirection::Desc => StoreClientKeySortDirection::Desc, - }, - } -} - -fn store_client_key_cursor( - cursor: ClientKeyCursorData, - expected_sort: StoreClientKeySort, -) -> gateway_store::StoreResult { - let sort = store_client_key_sort(cursor.sort); - if sort != expected_sort { - return Err(gateway_store::StoreError::InvalidData { - entity: "client API key cursor", - message: "cursor sort does not match the requested sort".to_owned(), - }); - } - let value = match cursor.value { - ClientKeyCursorValue::Name(value) => StoreClientKeyCursorValue::Name(value), - ClientKeyCursorValue::Enabled(value) => StoreClientKeyCursorValue::Enabled(value), - ClientKeyCursorValue::CreatedAt(value) => StoreClientKeyCursorValue::CreatedAt(value), - ClientKeyCursorValue::LastUsedAt(value) => StoreClientKeyCursorValue::LastUsedAt(value), - }; - StoreClientKeyCursor::new(sort, value, cursor.id) -} - -fn wire_client_key_cursor(cursor: StoreClientKeyCursor) -> ClientKeyCursorData { - let sort = ClientKeySort { - field: match cursor.sort.field { - StoreClientKeySortField::Name => ClientKeySortField::Name, - StoreClientKeySortField::Enabled => ClientKeySortField::Enabled, - StoreClientKeySortField::CreatedAt => ClientKeySortField::CreatedAt, - StoreClientKeySortField::LastUsedAt => ClientKeySortField::LastUsedAt, - }, - direction: match cursor.sort.direction { - StoreClientKeySortDirection::Asc => ClientKeySortDirection::Asc, - StoreClientKeySortDirection::Desc => ClientKeySortDirection::Desc, - }, - }; - let value = match cursor.value { - StoreClientKeyCursorValue::Name(value) => ClientKeyCursorValue::Name(value), - StoreClientKeyCursorValue::Enabled(value) => ClientKeyCursorValue::Enabled(value), - StoreClientKeyCursorValue::CreatedAt(value) => ClientKeyCursorValue::CreatedAt(value), - StoreClientKeyCursorValue::LastUsedAt(value) => ClientKeyCursorValue::LastUsedAt(value), - }; - ClientKeyCursorData { - sort, - value, - id: cursor.id, - } -} - -fn client_key_view(record: ClientApiKeyRecord) -> ClientKeyView { - ClientKeyView::new(ClientKeyViewFields { - id: record.id, - name: record.name, - label: record.label, - provider_kind: record.provider_kind, - prefix: record.prefix, - enabled: record.enabled, - max_concurrency: record.max_concurrency, - requests_per_minute: record.requests_per_minute, - tokens_per_minute: record.tokens_per_minute, - created_at: record.created_at, - updated_at: record.updated_at, - last_used_at: record.last_used_at, - }) -} - -use gateway_api::admin::openai::{ - CodexCredentialDetailsData, CodexCredentialListData, CodexCredentialMutationData, - CodexCredentialRotationData, CodexCredentialView, CodexCredentialsDocumentImportData, - CodexOAuthAuthorizationStartedData, CompleteOAuthAuthorizationRequest, - CredentialCursorWire as CodexCredentialCursor, CredentialMutationRequest as CodexMutation, - ImportCredentialsDocumentRequest as CodexImportDocument, - ListCredentialsQuery as CodexListQuery, RotateCredentialRequest as CodexRotate, - StartOAuthAuthorizationRequest as CodexStartAuthorization, -}; -use gateway_core::engine::credential::{ - CredentialCasUpdate as CoreCredentialUpdate, CredentialRevision as CoreCredentialRevision, - NewProviderAccount as CoreNewProviderAccount, ProviderAccountStore, -}; -use gateway_store::{ - JsonObject, - postgres::{ - DeleteProviderAccounts, ImportProviderAccounts, - NewProviderAccount as StoreNewProviderAccount, ProviderAccountAdminRepository, - ProviderAccountAdminScope, ProviderAccountAvailability, ProviderAccountRepository, - ProviderAccountSummary, ProviderCredentialUpdate, RotateProviderAccount, - SetProviderAccountEnabled, UpdateProviderAccount, - }, -}; -use provider_openai::credential::{ - CodexCredentialAdmin, CodexCredentialAdminError, CodexCredentialAdminService, CodexOAuthAdmin, - CodexOAuthAdminError, CodexOAuthFlowBinding, CodexOAuthPendingStore, - CodexOAuthPendingStoreError, CodexOAuthSecret, CodexPendingAuthorization, - CodexTokenIdentityVerifier, CompleteCodexOAuthAuthorization, RotateManagedCodexCredential, - StartCodexOAuthAuthorization, StartCodexOAuthReauthorization, StoredCodexPendingAuthorization, -}; -use secrecy::SecretString; -use tokio::sync::Mutex as AsyncMutex; - -const PROVIDER_CREDENTIAL_JSON_LIMIT: usize = 256 * 1024; - -#[derive(Debug, Clone)] -struct PendingCodexAdminFlow { - owner_ref: String, - expected_config_revision: Revision, - kind: PendingCodexAdminFlowKind, -} - -#[derive(Debug, Clone)] -enum PendingCodexAdminFlowKind { - Create { - provider_instance_id: String, - }, - Reauthorize { - account_id: String, - provider_instance_id: String, - }, -} - -/// 单进程 OAuth flow 的一次性 server-only 状态。 -pub struct InMemoryCodexOAuthPendingStore { - flows: AsyncMutex>, -} - -impl InMemoryCodexOAuthPendingStore { - #[must_use] - pub fn new() -> Self { - Self { - flows: AsyncMutex::new(BTreeMap::new()), - } - } -} - -impl Default for InMemoryCodexOAuthPendingStore { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl CodexOAuthPendingStore for InMemoryCodexOAuthPendingStore { - async fn create( - &self, - pending: &CodexPendingAuthorization, - ) -> Result<(), CodexOAuthPendingStoreError> { - let duplicate = CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization { - flow_id: pending.flow_id().to_owned(), - owner_ref: pending.owner_ref().to_owned(), - started_request_ref: pending.started_request_ref().to_owned(), - provider_instance_id: pending.provider_instance_id().to_owned(), - name: pending.name().to_owned(), - expires_at: pending.expires_at(), - state: pending.state().clone(), - nonce: pending.nonce().clone(), - code_verifier: pending.code_verifier().clone(), - reauthorization_account_id: pending - .reauthorization() - .map(|target| target.account_id().to_string()), - reauthorization_credential_revision: pending - .reauthorization() - .map(|target| target.credential_revision().get()), - })?; - let mut flows = self.flows.lock().await; - if flows - .insert(pending.flow_id().to_owned(), duplicate) - .is_some() - { - return Err(CodexOAuthPendingStoreError::Conflict); - } - Ok(()) - } - - async fn take( - &self, - owner_ref: &str, - flow_id: &str, - ) -> Result, CodexOAuthPendingStoreError> { - let mut flows = self.flows.lock().await; - if flows - .get(flow_id) - .is_some_and(|flow| flow.owner_ref() == owner_ref) - { - Ok(flows.remove(flow_id)) - } else { - Ok(None) - } - } -} - -/// Codex Provider 验证结果到 PostgreSQL 原子 revision + audit 的组合 adapter。 -pub struct CodexAdminAdapter { - accounts: Arc, - admin_accounts: Arc, - core_store: Arc, - control_plane: Arc, - verifier: Arc, - oauth: Arc, - owner: Arc, - preparer: CodexCredentialAdmin, - publisher: RuntimeSnapshotPublisher, - pending_flows: AsyncMutex>, -} - -/// [`CodexAdminAdapter`] 的显式 owner 组合。 -pub struct CodexAdminPorts { - pub accounts: Arc, - pub admin_accounts: Arc, - pub core_store: Arc, - pub control_plane: Arc, - pub verifier: Arc, - pub oauth: Arc, - pub owner: Arc, - pub publisher: RuntimeSnapshotPublisher, -} - -impl CodexAdminAdapter { - #[must_use] - pub fn new(ports: CodexAdminPorts) -> Self { - Self { - accounts: ports.accounts, - admin_accounts: ports.admin_accounts, - core_store: ports.core_store, - control_plane: ports.control_plane, - verifier: ports.verifier, - oauth: ports.oauth, - owner: ports.owner, - preparer: CodexCredentialAdmin, - publisher: ports.publisher, - pending_flows: AsyncMutex::new(BTreeMap::new()), - } - } - - async fn current_revision(&self) -> Result { - self.control_plane - .load_control_plane() - .await - .map(|snapshot| snapshot.settings.config_revision) - .map_err(|error| map_admin_store_error(error, "runtime settings")) - } - - async fn require_revision(&self, expected: u64) -> Result { - let expected = revision(expected)?; - let current = self.current_revision().await?; - if current != expected { - return Err(AdminServiceError::conflict( - "Configuration revision is stale", - )); - } - Ok(expected) - } - - async fn import_prepared( - &self, - context: &AdminRequestContext, - expected: Revision, - provider_instance_id: String, - prepared: Vec, - action: &'static str, - ) -> Result<(Revision, Vec), AdminServiceError> { - let mut accounts = Vec::with_capacity(prepared.len()); - for account in prepared { - accounts.push(store_new_provider_account(account)?); - } - let accounts = resolve_admin_import_accounts( - self.accounts.as_ref(), - "openai", - &provider_instance_id, - accounts, - ) - .await?; - let ids: Vec = accounts.iter().map(|account| account.id.clone()).collect(); - let entity_ref = if ids.len() == 1 { - ids[0].clone() - } else { - format!("{} credentials", ids.len()) - }; - let committed = self - .admin_accounts - .import_provider_accounts( - expected, - ImportProviderAccounts { - scope: provider_scope("openai", provider_instance_id), - accounts, - audit: admin_audit_event( - context, - action, - "provider_account", - entity_ref, - &["provider_credentials_json", "credential_revision"], - ), - }, - ) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))?; - self.publisher.publish_committed(committed).await; - Ok((committed, ids)) - } - - async fn mutate_enabled( - &self, - context: &AdminRequestContext, - request: CodexMutation, - enabled: Option, - ) -> Result { - let record = self - .accounts - .load_provider_account(&request.credential_id) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))? - .ok_or_else(|| AdminServiceError::not_found("Codex credential was not found"))?; - require_provider(&record.summary, "openai")?; - let expected = revision(request.expected_config_revision)?; - let id = record.summary.id.clone(); - let scope = provider_scope("openai", record.summary.provider_instance_id); - let committed = match enabled { - Some(enabled) => { - self.admin_accounts - .set_provider_account_enabled_admin( - expected, - SetProviderAccountEnabled { - scope, - account_id: id.clone(), - enabled, - audit: admin_audit_event( - context, - if enabled { "enable" } else { "disable" }, - "provider_account", - id.clone(), - &["enabled"], - ), - }, - ) - .await - } - None => { - self.admin_accounts - .delete_provider_accounts_admin( - expected, - DeleteProviderAccounts { - scope, - account_ids: vec![id.clone()], - audit: admin_audit_event( - context, - "delete", - "provider_account", - id.clone(), - &[], - ), - }, - ) - .await - } - } - .map_err(|error| map_admin_store_error(error, "Codex credential"))?; - self.publisher.publish_committed(committed).await; - Ok(CodexCredentialMutationData { - config_revision: committed.get(), - credential_id: id, - }) - } -} - -#[async_trait] -impl CodexAdminService for CodexAdminAdapter { - async fn list( - &self, - query: CodexListQuery, - ) -> Result { - let cursor = query - .cursor - .as_deref() - .map(decode_cursor::) - .transpose()?; - let mut accounts = self - .accounts - .list_provider_accounts(query.provider_instance_id.as_deref(), true) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))? - .into_iter() - .filter(|account| account.provider_kind == "openai") - .filter(|account| { - query - .enabled - .is_none_or(|enabled| account.enabled == enabled) - }) - .filter(|account| { - query.availability.as_deref().is_none_or(|availability| { - codex_availability(account.availability) == availability - }) - }) - .collect::>(); - accounts.sort_by(|left, right| { - (left.created_at, left.id.as_str()).cmp(&(right.created_at, right.id.as_str())) - }); - if let Some(cursor) = cursor.as_ref() { - accounts.retain(|account| { - (account.created_at, account.id.as_str()) - > (cursor.created_at, cursor.credential_id.as_str()) - }); - } - let limit = usize::from( - query - .limit - .unwrap_or(gateway_api::admin::openai::DEFAULT_PAGE_SIZE), - ); - let has_more = accounts.len() > limit; - accounts.truncate(limit); - let next_cursor = has_more - .then(|| accounts.last()) - .flatten() - .map(|account| { - encode_cursor(&CodexCredentialCursor { - created_at: account.created_at, - credential_id: account.id.clone(), - }) - }) - .transpose()?; - Ok(CodexCredentialListData { - config_revision: self.current_revision().await?.get(), - items: accounts.into_iter().map(codex_credential_view).collect(), - next_cursor, - }) - } - - async fn details( - &self, - credential_id: String, - ) -> Result { - let record = self - .accounts - .load_provider_account(&credential_id) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))? - .ok_or_else(|| AdminServiceError::not_found("Codex credential was not found"))?; - require_provider(&record.summary, "openai")?; - Ok(CodexCredentialDetailsData { - config_revision: self.current_revision().await?.get(), - credential: codex_credential_view(record.summary), - }) - } - - async fn import_document( - &self, - context: &AdminRequestContext, - request: CodexImportDocument, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let instance = - gateway_core::routing::ProviderInstanceId::new(request.provider_instance_id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid Codex Provider instance"))?; - let profile = self - .owner - .prepare_import_document(instance, request.document) - .await - .map_err(map_codex_admin_error)?; - let (committed, ids) = self - .import_prepared( - context, - expected, - request.provider_instance_id, - profile.into_accounts(), - "import_document", - ) - .await?; - Ok(CodexCredentialsDocumentImportData { - config_revision: committed.get(), - credential_ids: ids, - }) - } - - async fn start_authorization( - &self, - context: &AdminRequestContext, - request: CodexStartAuthorization, - ) -> Result { - let expected = self - .require_revision(request.expected_config_revision) - .await?; - let owner_ref = admin_flow_owner(context); - let binding = CodexOAuthFlowBinding::new(owner_ref.clone(), context.request_id()) - .map_err(map_codex_oauth_error)?; - let (started, kind) = match (request.credential_id, request.expected_credential_revision) { - (None, None) => { - let provider_instance_id = request.provider_instance_id; - let started = self - .oauth - .start_authorization(StartCodexOAuthAuthorization { - binding, - provider_instance_id: provider_instance_id.clone(), - name: request.name, - }) - .await - .map_err(map_codex_oauth_error)?; - ( - started, - PendingCodexAdminFlowKind::Create { - provider_instance_id, - }, - ) - } - (Some(account_id), Some(expected_credential_revision)) => { - let record = self - .accounts - .load_provider_account(&account_id) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))? - .ok_or_else(|| { - AdminServiceError::not_found("Codex credential was not found") - })?; - require_provider(&record.summary, "openai")?; - if record.summary.provider_instance_id != request.provider_instance_id { - return Err(AdminServiceError::conflict( - "Codex credential does not belong to the selected Provider instance", - )); - } - let account = - gateway_core::engine::credential::ProviderAccountId::new(account_id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid Codex credential ID"))?; - let credential_revision = CoreCredentialRevision::new(expected_credential_revision) - .map_err(|_| AdminServiceError::invalid("Invalid credential revision"))?; - let started = self - .oauth - .start_reauthorization(StartCodexOAuthReauthorization { - binding, - account_id: account, - expected_credential_revision: credential_revision, - }) - .await - .map_err(map_codex_oauth_error)?; - ( - started, - PendingCodexAdminFlowKind::Reauthorize { - account_id, - provider_instance_id: record.summary.provider_instance_id, - }, - ) - } - _ => { - return Err(AdminServiceError::invalid( - "Codex reauthorization target is incomplete", - )); - } - }; - self.pending_flows.lock().await.insert( - started.flow_id.clone(), - PendingCodexAdminFlow { - owner_ref, - expected_config_revision: expected, - kind, - }, - ); - Ok(CodexOAuthAuthorizationStartedData { - flow_id: started.flow_id, - authorization_url: started.authorization_url, - expires_at: started.expires_at, - }) - } - - async fn complete_authorization( - &self, - context: &AdminRequestContext, - request: CompleteOAuthAuthorizationRequest, - ) -> Result { - let owner_ref = admin_flow_owner(context); - let pending = { - let mut flows = self.pending_flows.lock().await; - if !flows - .get(&request.flow_id) - .is_some_and(|pending| pending.owner_ref == owner_ref) - { - return Err(AdminServiceError::not_found( - "Codex OAuth flow was not found", - )); - } - flows - .remove(&request.flow_id) - .expect("checked Codex OAuth flow exists") - }; - let command = CompleteCodexOAuthAuthorization { - owner_ref, - flow_id: request.flow_id, - callback_url: SecretString::from(request.callback_url), - }; - match pending.kind { - PendingCodexAdminFlowKind::Create { - provider_instance_id, - } => { - let prepared = self - .oauth - .complete_authorization(command) - .await - .map_err(map_codex_oauth_error)?; - let (committed, ids) = self - .import_prepared( - context, - pending.expected_config_revision, - provider_instance_id, - vec![prepared], - "oauth_complete", - ) - .await?; - Ok(CodexCredentialMutationData { - config_revision: committed.get(), - credential_id: ids.into_iter().next().expect("single OAuth account"), - }) - } - PendingCodexAdminFlowKind::Reauthorize { - account_id, - provider_instance_id, - } => { - let prepared = self - .oauth - .complete_reauthorization(command) - .await - .map_err(map_codex_oauth_error)?; - let (profile, credential, guard) = prepared.into_parts(); - let rotation = self - .admin_accounts - .rotate_provider_account( - pending.expected_config_revision, - RotateProviderAccount { - scope: provider_scope("openai", provider_instance_id), - profile: store_profile(profile), - credential: store_credential_update(credential)?, - audit: admin_audit_event( - context, - "oauth_reauthorize", - "provider_account", - account_id.clone(), - &["provider_credentials_json", "credential_revision"], - ), - }, - ) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))?; - drop(guard); - self.publisher - .publish_committed(rotation.config_revision) - .await; - Ok(CodexCredentialMutationData { - config_revision: rotation.config_revision.get(), - credential_id: account_id, - }) - } - } - } - - async fn rotate( - &self, - context: &AdminRequestContext, - request: CodexRotate, - ) -> Result { - let account_id = - gateway_core::engine::credential::ProviderAccountId::new(request.credential_id.clone()) - .map_err(|_| AdminServiceError::invalid("Invalid Codex credential ID"))?; - let expected_credential = CoreCredentialRevision::new(request.expected_credential_revision) - .map_err(|_| AdminServiceError::invalid("Invalid credential revision"))?; - let current = self - .core_store - .load_credential(&account_id, expected_credential) - .await - .map_err(|_| AdminServiceError::conflict("Credential revision is stale"))?; - let secret = CodexOAuthSecret { - access_token: SecretString::from(request.access_token), - refresh_token: request.refresh_token.map(SecretString::from), - id_token: None, - }; - let profile = self - .verifier - .verify(&secret) - .await - .map_err(map_codex_identity_error)?; - let prepared = self - .preparer - .prepare_rotation(RotateManagedCodexCredential { - current, - secret, - verified_account: profile, - }) - .map_err(map_codex_admin_error)?; - let scope = provider_scope( - "openai", - prepared_instance(&self.accounts, &request.credential_id).await?, - ); - let rotation = self - .admin_accounts - .rotate_provider_account( - revision(request.expected_config_revision)?, - RotateProviderAccount { - scope, - profile: store_profile(prepared.profile), - credential: store_credential_update(prepared.credential)?, - audit: admin_audit_event( - context, - "rotate", - "provider_account", - request.credential_id.clone(), - &["provider_credentials_json", "credential_revision"], - ), - }, - ) - .await - .map_err(|error| map_admin_store_error(error, "Codex credential"))?; - self.publisher - .publish_committed(rotation.config_revision) - .await; - Ok(CodexCredentialRotationData { - credential_id: request.credential_id, - credential_revision: i64::try_from(rotation.credential_revision.get()) - .map_err(|_| AdminServiceError::internal("Credential revision overflow"))?, - }) - } - - async fn enable( - &self, - context: &AdminRequestContext, - request: CodexMutation, - ) -> Result { - self.mutate_enabled(context, request, Some(true)).await - } - - async fn disable( - &self, - context: &AdminRequestContext, - request: CodexMutation, - ) -> Result { - self.mutate_enabled(context, request, Some(false)).await - } - - async fn delete( - &self, - context: &AdminRequestContext, - request: CodexMutation, - ) -> Result { - self.mutate_enabled(context, request, None).await - } -} - -fn provider_scope( - provider_kind: impl Into, - provider_instance_id: impl Into, -) -> ProviderAccountAdminScope { - ProviderAccountAdminScope { - provider_kind: provider_kind.into(), - provider_instance_id: provider_instance_id.into(), - } -} - -async fn resolve_admin_import_accounts( - repository: &dyn ProviderAccountRepository, - provider_kind: &str, - provider_instance_id: &str, - mut accounts: Vec, -) -> Result, AdminServiceError> { - let existing = repository - .list_provider_accounts(None, true) - .await - .map_err(|error| map_admin_store_error(error, "Provider credential"))?; - let mut resolved_ids = BTreeSet::new(); - let mut batch_identities = BTreeSet::new(); - - for account in &mut accounts { - if account.provider_kind != provider_kind - || account.provider_instance_id != provider_instance_id - { - return Err(AdminServiceError::invalid( - "Imported credential is outside the selected Provider instance", - )); - } - let existing_by_id = existing.iter().find(|current| current.id == account.id); - let existing_by_identity = existing.iter().find(|current| { - current.provider_kind == provider_kind - && current.upstream_user_id == account.upstream_user_id - && current.upstream_account_id == account.upstream_account_id - }); - - if let Some(current) = existing_by_id - && (current.provider_kind != provider_kind - || current.provider_instance_id != provider_instance_id - || current.upstream_user_id != account.upstream_user_id - || current.upstream_account_id != account.upstream_account_id) - { - return Err(AdminServiceError::conflict( - "Provider credential ID is already bound to another identity", - )); - } - if let Some(current) = existing_by_identity { - if current.provider_instance_id != provider_instance_id - || existing_by_id.is_some_and(|by_id| by_id.id != current.id) - { - return Err(AdminServiceError::conflict( - "Provider credential identity is already bound to another instance", - )); - } - account.id.clone_from(¤t.id); - } - - if !resolved_ids.insert(account.id.clone()) - || !batch_identities.insert(( - account.upstream_user_id.clone(), - account.upstream_account_id.clone(), - )) - { - return Err(AdminServiceError::invalid( - "Import document contains duplicate Provider identities", - )); - } - } - Ok(accounts) -} - -fn require_provider( - account: &ProviderAccountSummary, - provider: &str, -) -> Result<(), AdminServiceError> { - if account.provider_kind == provider { - Ok(()) - } else { - Err(AdminServiceError::not_found( - "Provider credential was not found", - )) - } -} - -fn store_new_provider_account( - prepared: CoreNewProviderAccount, -) -> Result { - let account = prepared.account; - let credential = JsonObject::try_from_value( - "provider credential", - serde_json::Value::Object(prepared.credential.into_inner()), - PROVIDER_CREDENTIAL_JSON_LIMIT, - ) - .map_err(|error| map_admin_store_error(error, "Provider credential"))?; - Ok(StoreNewProviderAccount { - id: account.id().as_str().to_owned(), - provider_instance_id: account.instance().as_str().to_owned(), - provider_kind: account.provider().as_str().to_owned(), - name: account.name().to_owned(), - email: account.email().map(str::to_owned), - upstream_user_id: account.upstream_user_id().to_owned(), - upstream_account_id: account.upstream_account_id().map(str::to_owned), - plan_type: account.plan_type().map(str::to_owned), - provider_credentials_json: credential, - has_refresh_token: account.has_refresh_token(), - access_token_expires_at: DateTime::::from(account.access_token_expires_at()), - next_refresh_at: account.next_refresh_at().map(DateTime::::from), - enabled: account.enabled(), - availability: store_account_availability(account.availability()), - cooldown_until: account.cooldown_until().map(DateTime::::from), - availability_observed_at: Utc::now(), - }) -} - -const fn store_account_availability( - availability: gateway_core::engine::credential::AccountAvailability, -) -> ProviderAccountAvailability { - match availability { - gateway_core::engine::credential::AccountAvailability::Unknown => { - ProviderAccountAvailability::Unknown - } - gateway_core::engine::credential::AccountAvailability::Ready => { - ProviderAccountAvailability::Ready - } - gateway_core::engine::credential::AccountAvailability::Cooldown => { - ProviderAccountAvailability::Cooldown - } - gateway_core::engine::credential::AccountAvailability::QuotaExhausted => { - ProviderAccountAvailability::QuotaExhausted - } - gateway_core::engine::credential::AccountAvailability::Expired => { - ProviderAccountAvailability::Expired - } - gateway_core::engine::credential::AccountAvailability::Banned => { - ProviderAccountAvailability::Banned - } - gateway_core::engine::credential::AccountAvailability::Invalid => { - ProviderAccountAvailability::Invalid - } - } -} - -fn store_profile( - profile: gateway_core::engine::credential::ProviderAccountUpdate, -) -> UpdateProviderAccount { - UpdateProviderAccount { - id: profile.account_id.as_str().to_owned(), - name: profile.name, - email: profile.email, - plan_type: profile.plan_type, - } -} - -fn store_credential_update( - update: CoreCredentialUpdate, -) -> Result { - let ( - account_id, - expected_revision, - _profile, - credential, - has_refresh_token, - access_token_expires_at, - next_refresh_at, - ) = update.into_parts(); - Ok(ProviderCredentialUpdate { - account_id: account_id.as_str().to_owned(), - expected_revision: revision(expected_revision.get())?, - provider_credentials_json: JsonObject::try_from_value( - "provider credential", - serde_json::Value::Object(credential.into_inner()), - PROVIDER_CREDENTIAL_JSON_LIMIT, - ) - .map_err(|error| map_admin_store_error(error, "Provider credential"))?, - has_refresh_token, - access_token_expires_at: DateTime::::from(access_token_expires_at), - next_refresh_at: next_refresh_at.map(DateTime::::from), - }) -} - -async fn prepared_instance( - accounts: &Arc, - account_id: &str, -) -> Result { - accounts - .load_provider_account(account_id) - .await - .map_err(|error| map_admin_store_error(error, "Provider credential"))? - .map(|record| record.summary.provider_instance_id) - .ok_or_else(|| AdminServiceError::not_found("Provider credential was not found")) -} - -fn codex_credential_view(account: ProviderAccountSummary) -> CodexCredentialView { - CodexCredentialView { - id: account.id, - provider_instance_id: account.provider_instance_id, - name: account.name, - email: account.email, - upstream_user_id: account.upstream_user_id, - upstream_account_id: account.upstream_account_id, - enabled: account.enabled, - credential_revision: i64::try_from(account.credential_revision.get()).unwrap_or(i64::MAX), - has_refresh_token: account.has_refresh_token, - availability: codex_availability(account.availability).to_owned(), - availability_reason: account.availability_reason, - plan_type: account.plan_type, - access_token_expires_at: account.access_token_expires_at, - next_refresh_at: account.next_refresh_at, - cooldown_until: account.cooldown_until, - created_at: account.created_at, - updated_at: account.updated_at, - } -} - -const fn codex_availability(value: ProviderAccountAvailability) -> &'static str { - match value { - ProviderAccountAvailability::Unknown => "unknown", - ProviderAccountAvailability::Ready => "ready", - ProviderAccountAvailability::Cooldown => "cooldown", - ProviderAccountAvailability::QuotaExhausted => "exhausted", - ProviderAccountAvailability::Expired - | ProviderAccountAvailability::Banned - | ProviderAccountAvailability::Invalid => "invalid", - } -} - -fn admin_flow_owner(context: &AdminRequestContext) -> String { - context - .admin_user_id() - .map(|id| format!("admin:{id}")) - .unwrap_or_else(|| "admin:api-key".to_owned()) -} - -fn map_codex_identity_error( - error: provider_openai::credential::CodexIdentityVerificationError, -) -> AdminServiceError { - match error { - provider_openai::credential::CodexIdentityVerificationError::Rejected => { - AdminServiceError::invalid("Codex OAuth identity was rejected") - } - provider_openai::credential::CodexIdentityVerificationError::Unavailable => { - AdminServiceError::unavailable("Codex identity verification is unavailable") - } - } -} - -fn map_codex_admin_error(error: CodexCredentialAdminError) -> AdminServiceError { - match error { - CodexCredentialAdminError::InvalidInput | CodexCredentialAdminError::InvalidCredential => { - AdminServiceError::invalid("Codex credential is invalid") - } - CodexCredentialAdminError::IdentityMismatch => { - AdminServiceError::conflict("Codex credential identity cannot be rebound") - } - CodexCredentialAdminError::NotFound => { - AdminServiceError::not_found("Codex account was not found") - } - CodexCredentialAdminError::RevisionConflict - | CodexCredentialAdminError::RefreshLeaseUnavailable - | CodexCredentialAdminError::RefreshAmbiguous => { - AdminServiceError::conflict("Codex account changed while the credential was refreshed") - } - CodexCredentialAdminError::MissingRefreshToken => { - AdminServiceError::invalid("Codex account has no refresh token") - } - CodexCredentialAdminError::RefreshRejected - | CodexCredentialAdminError::AccountBanned - | CodexCredentialAdminError::IdentityRejected => { - AdminServiceError::invalid("Codex refreshed credential was rejected") - } - CodexCredentialAdminError::StoreUnavailable - | CodexCredentialAdminError::RefreshUnavailable - | CodexCredentialAdminError::IdentityUnavailable => { - AdminServiceError::unavailable("Codex credential refresh is unavailable") - } - } -} - -fn map_codex_oauth_error(error: CodexOAuthAdminError) -> AdminServiceError { - match error { - CodexOAuthAdminError::InvalidInput | CodexOAuthAdminError::UpstreamRejected => { - AdminServiceError::invalid("Codex OAuth operation was rejected") - } - CodexOAuthAdminError::NotFound | CodexOAuthAdminError::FlowExpired => { - AdminServiceError::not_found("Codex OAuth flow was not found") - } - CodexOAuthAdminError::Conflict | CodexOAuthAdminError::Ambiguous => { - AdminServiceError::conflict("Codex OAuth operation conflicts with current state") - } - CodexOAuthAdminError::UpstreamUnavailable | CodexOAuthAdminError::StorageUnavailable => { - AdminServiceError::unavailable("Codex OAuth service is unavailable") - } - CodexOAuthAdminError::Credential => { - AdminServiceError::internal("Codex credential preparation failed") - } - } -} - -use gateway_api::admin::xai::{ - AuthorizationStartResult, CompleteAuthorizationRequest as XaiCompleteAuthorization, - CredentialMutationRequest as XaiMutation, ListCredentialsQuery as XaiListQuery, - StartAuthorizationRequest as XaiStartAuthorization, XaiCredentialImportData, - XaiCredentialImportDocumentRequest, XaiCredentialListData, XaiCredentialMutationData, - XaiCredentialViewData, -}; -use provider_xai::{ - AuthorizationCallback, DiscoveryDocument, GrokCredentialAdmin, GrokOAuthClient, - GrokOAuthImportDocument, OFFICIAL_REDIRECT_URI, PendingAuthorization, RedirectUriAllowlist, - VerifiedGrokAccount, -}; - -struct PendingXaiAuthorizationFlow { - owner_ref: String, - expected_config_revision: Revision, - provider_instance_id: String, - name: String, - expires_at: DateTime, - discovery: DiscoveryDocument, - pending: PendingAuthorization, -} - -/// xAI 官方 OAuth/OIDC 流到 Store 原子账号事务的组合 adapter。 -pub struct XaiAdminAdapter { - accounts: Arc, - admin_accounts: Arc, - control_plane: Arc, - runtime_settings: Arc, - oauth: Arc, - preparer: GrokCredentialAdmin, - publisher: RuntimeSnapshotPublisher, - authorization_flows: AsyncMutex>, -} - -impl XaiAdminAdapter { - #[must_use] - pub fn new( - accounts: Arc, - admin_accounts: Arc, - control_plane: Arc, - runtime_settings: Arc, - oauth: Arc, - publisher: RuntimeSnapshotPublisher, - ) -> Self { - Self { - accounts, - admin_accounts, - control_plane, - runtime_settings, - oauth, - preparer: GrokCredentialAdmin, - publisher, - authorization_flows: AsyncMutex::new(BTreeMap::new()), - } - } - - async fn current_revision(&self) -> Result { - self.control_plane - .load_control_plane() - .await - .map(|snapshot| snapshot.settings.config_revision) - .map_err(|error| map_admin_store_error(error, "runtime settings")) - } - - async fn require_revision(&self, expected: u64) -> Result { - let expected = revision(expected)?; - if self.current_revision().await? != expected { - return Err(AdminServiceError::conflict( - "Configuration revision is stale", - )); - } - Ok(expected) - } - - async fn refresh_margin(&self) -> Result { - self.runtime_settings - .load_runtime_settings() - .await - .map(|settings| Duration::from_secs(settings.refresh_margin_seconds)) - .map_err(|error| map_admin_store_error(error, "runtime settings")) - } - - async fn prepare_verified( - &self, - provider_instance_id: String, - name: String, - email: Option, - tokens: provider_xai::VerifiedTokenSet, - ) -> Result { - let account_id = gateway_core::engine::credential::ProviderAccountId::new(format!( - "acct_{}", - Uuid::now_v7().simple() - )) - .map_err(|_| AdminServiceError::internal("Failed to allocate xAI account ID"))?; - let provider_instance_id = - gateway_core::routing::ProviderInstanceId::new(provider_instance_id) - .map_err(|_| AdminServiceError::invalid("Invalid xAI Provider instance"))?; - self.preparer - .prepare_verified_account(&VerifiedGrokAccount { - account_id, - provider_instance_id, - name, - email, - upstream_account_id: None, - plan_type: None, - tokens, - enabled: true, - refresh_margin: self.refresh_margin().await?, - }) - .map_err(map_xai_repository_error) - } - - async fn commit_import( - &self, - context: &AdminRequestContext, - expected: Revision, - provider_instance_id: String, - prepared: Vec, - action: &'static str, - ) -> Result<(Revision, Vec), AdminServiceError> { - let mut accounts = Vec::with_capacity(prepared.len()); - for account in prepared { - accounts.push(store_new_provider_account(account)?); - } - let accounts = resolve_admin_import_accounts( - self.accounts.as_ref(), - "xai", - &provider_instance_id, - accounts, - ) - .await?; - let ids: Vec = accounts.iter().map(|account| account.id.clone()).collect(); - let committed = self - .admin_accounts - .import_provider_accounts( - expected, - ImportProviderAccounts { - scope: provider_scope("xai", provider_instance_id), - accounts, - audit: admin_audit_event( - context, - action, - "provider_account", - if ids.len() == 1 { - ids[0].clone() - } else { - format!("{} credentials", ids.len()) - }, - &["provider_credentials_json", "credential_revision"], - ), - }, - ) - .await - .map_err(|error| map_admin_store_error(error, "xAI credential"))?; - self.publisher.publish_committed(committed).await; - Ok((committed, ids)) - } - - async fn mutate_xai_enabled( - &self, - context: &AdminRequestContext, - request: XaiMutation, - enabled: Option, - ) -> Result { - let record = self - .accounts - .load_provider_account(&request.credential_id) - .await - .map_err(|error| map_admin_store_error(error, "xAI credential"))? - .ok_or_else(|| AdminServiceError::not_found("xAI credential was not found"))?; - require_provider(&record.summary, "xai")?; - let id = record.summary.id.clone(); - let scope = provider_scope("xai", record.summary.provider_instance_id); - let expected = revision(request.expected_config_revision)?; - let committed = match enabled { - Some(enabled) => { - self.admin_accounts - .set_provider_account_enabled_admin( - expected, - SetProviderAccountEnabled { - scope, - account_id: id.clone(), - enabled, - audit: admin_audit_event( - context, - if enabled { "enable" } else { "disable" }, - "provider_account", - id.clone(), - &["enabled"], - ), - }, - ) - .await - } - None => { - self.admin_accounts - .delete_provider_accounts_admin( - expected, - DeleteProviderAccounts { - scope, - account_ids: vec![id.clone()], - audit: admin_audit_event( - context, - "delete", - "provider_account", - id.clone(), - &[], - ), - }, - ) - .await - } - } - .map_err(|error| map_admin_store_error(error, "xAI credential"))?; - self.publisher.publish_committed(committed).await; - Ok(XaiCredentialMutationData { - config_revision: committed.get(), - credential_id: id, - }) - } -} - -#[async_trait] -impl XaiAdminService for XaiAdminAdapter { - async fn list(&self, query: XaiListQuery) -> Result { - let items = self - .accounts - .list_provider_accounts(query.provider_instance_id.as_deref(), true) - .await - .map_err(|error| map_admin_store_error(error, "xAI credential"))? - .into_iter() - .filter(|account| account.provider_kind == "xai") - .map(xai_credential_view) - .collect(); - Ok(XaiCredentialListData { - config_revision: self.current_revision().await?.get(), - items, - }) - } - - async fn import_document( - &self, - context: &AdminRequestContext, - request: XaiCredentialImportDocumentRequest, - ) -> Result { - let expected = revision(request.expected_config_revision)?; - let encoded = serde_json::to_vec(&request.document) - .map_err(|_| AdminServiceError::invalid("xAI import document is invalid"))?; - let document = - GrokOAuthImportDocument::parse_json(&encoded).map_err(map_xai_import_error)?; - let discovery = self.oauth.discover().await.map_err(map_xai_oauth_error)?; - let entries = document.into_entries(); - let mut prepared = Vec::with_capacity(entries.len()); - for entry in entries { - let name = entry.name().to_owned(); - let email = entry.email().map(str::to_owned); - let tokens = self - .oauth - .verify_imported_credential(&discovery, entry.into_candidate()) - .await - .map_err(map_xai_import_error)?; - prepared.push( - self.prepare_verified(request.provider_instance_id.clone(), name, email, tokens) - .await?, - ); - } - let (committed, ids) = self - .commit_import( - context, - expected, - request.provider_instance_id, - prepared, - "import_document", - ) - .await?; - Ok(XaiCredentialImportData::new(committed.get(), ids)) - } - - async fn start_authorization( - &self, - context: &AdminRequestContext, - request: XaiStartAuthorization, - ) -> Result { - let expected = self - .require_revision(request.credential.expected_config_revision) - .await?; - let discovery = self.oauth.discover().await.map_err(map_xai_oauth_error)?; - let allowlist = RedirectUriAllowlist::new([OFFICIAL_REDIRECT_URI]) - .map_err(|_| AdminServiceError::invalid("xAI redirect URI is invalid"))?; - let redirect_uri = allowlist - .authorize(OFFICIAL_REDIRECT_URI) - .map_err(|_| AdminServiceError::invalid("xAI redirect URI is not allowed"))?; - let pending = self - .oauth - .start_authorization_code(&discovery, redirect_uri, None) - .map_err(map_xai_oauth_error)?; - let authorization_url = pending.authorization_url().to_string(); - let flow_id = format!("flow_{}", Uuid::now_v7().simple()); - let expires_at = Utc::now() + ChronoDuration::minutes(30); - self.authorization_flows.lock().await.insert( - flow_id.clone(), - PendingXaiAuthorizationFlow { - owner_ref: admin_flow_owner(context), - expected_config_revision: expected, - provider_instance_id: request.credential.provider_instance_id, - name: request.credential.name, - expires_at, - discovery, - pending, - }, - ); - Ok(AuthorizationStartResult { - flow_id, - authorization_url, - expires_at, - }) - } - - async fn complete_authorization( - &self, - context: &AdminRequestContext, - request: XaiCompleteAuthorization, - ) -> Result { - let owner_ref = admin_flow_owner(context); - let flow = { - let mut flows = self.authorization_flows.lock().await; - if !flows - .get(&request.flow_id) - .is_some_and(|flow| flow.owner_ref == owner_ref) - { - return Err(AdminServiceError::not_found( - "xAI authorization flow was not found", - )); - } - flows - .remove(&request.flow_id) - .expect("checked xAI authorization flow exists") - }; - if flow.expires_at <= Utc::now() { - return Err(AdminServiceError::not_found( - "xAI authorization flow expired", - )); - } - let callback_url = url::Url::parse(request.callback_url.trim()) - .map_err(|_| AdminServiceError::invalid("xAI OAuth callback URL is invalid"))?; - let expected_callback = url::Url::parse(OFFICIAL_REDIRECT_URI) - .map_err(|_| AdminServiceError::internal("xAI OAuth callback is invalid"))?; - if callback_url.scheme() != expected_callback.scheme() - || callback_url.host_str() != expected_callback.host_str() - || callback_url.port_or_known_default() != expected_callback.port_or_known_default() - || callback_url.path() != expected_callback.path() - || !callback_url.username().is_empty() - || callback_url.password().is_some() - || callback_url.fragment().is_some() - { - return Err(AdminServiceError::invalid( - "xAI OAuth callback URL is not allowed", - )); - } - let callback = AuthorizationCallback::parse(callback_url.query().unwrap_or_default()) - .map_err(|_| AdminServiceError::invalid("xAI OAuth callback was rejected"))?; - let grant = flow - .pending - .accept_callback(callback) - .map_err(map_xai_oauth_error)?; - let tokens = self - .oauth - .exchange_authorization_code(&flow.discovery, grant) - .await - .map_err(map_xai_oauth_error)?; - let prepared = self - .prepare_verified(flow.provider_instance_id.clone(), flow.name, None, tokens) - .await?; - let (committed, ids) = self - .commit_import( - context, - flow.expected_config_revision, - flow.provider_instance_id, - vec![prepared], - "authorization_oauth_complete", - ) - .await?; - Ok(XaiCredentialMutationData { - config_revision: committed.get(), - credential_id: ids - .into_iter() - .next() - .expect("single authorization account"), - }) - } - - async fn disable( - &self, - context: &AdminRequestContext, - request: XaiMutation, - ) -> Result { - self.mutate_xai_enabled(context, request, Some(false)).await - } - - async fn enable( - &self, - context: &AdminRequestContext, - request: XaiMutation, - ) -> Result { - self.mutate_xai_enabled(context, request, Some(true)).await - } - - async fn delete( - &self, - context: &AdminRequestContext, - request: XaiMutation, - ) -> Result { - self.mutate_xai_enabled(context, request, None).await - } -} - -fn xai_credential_view(account: ProviderAccountSummary) -> XaiCredentialViewData { - XaiCredentialViewData { - id: account.id, - provider_instance_id: account.provider_instance_id, - name: account.name, - email: account.email, - upstream_user_id: account.upstream_user_id, - upstream_account_id: account.upstream_account_id, - plan_type: account.plan_type, - enabled: account.enabled, - credential_revision: i64::try_from(account.credential_revision.get()).unwrap_or(i64::MAX), - has_refresh_token: account.has_refresh_token, - availability: account.availability.as_str().to_owned(), - availability_reason: account.availability_reason, - access_token_expires_at: account.access_token_expires_at, - next_refresh_at: account.next_refresh_at, - cooldown_until: account.cooldown_until, - created_at: account.created_at, - updated_at: account.updated_at, - } -} - -fn map_xai_oauth_error(error: provider_xai::OAuthError) -> AdminServiceError { - match error.class() { - provider_xai::FailureClass::Transient => { - AdminServiceError::unavailable("xAI OAuth service is unavailable") - } - provider_xai::FailureClass::Ambiguous => { - AdminServiceError::conflict("xAI OAuth send state is ambiguous") - } - provider_xai::FailureClass::CredentialPermanent - | provider_xai::FailureClass::ConfigurationPermanent - | provider_xai::FailureClass::UserActionRequired - | provider_xai::FailureClass::Security - | provider_xai::FailureClass::Unsupported => { - AdminServiceError::invalid("xAI OAuth operation was rejected") - } - } -} - -fn map_xai_import_error(error: provider_xai::GrokOAuthImportError) -> AdminServiceError { - match error { - provider_xai::GrokOAuthImportError::OAuth(error) => map_xai_oauth_error(error), - provider_xai::GrokOAuthImportError::InvalidField(_) => { - AdminServiceError::invalid("xAI OAuth import is invalid") - } - } -} - -fn map_xai_repository_error( - error: provider_xai::GrokCredentialRepositoryError, -) -> AdminServiceError { - use provider_xai::GrokCredentialRepositoryError as Error; - match error { - Error::InvalidInput(_) | Error::InvalidCredentialData | Error::WrongProviderKind => { - AdminServiceError::invalid("xAI credential is invalid") - } - Error::CredentialNotFound => AdminServiceError::not_found("xAI credential was not found"), - Error::StaleCredentialRevision | Error::IdentityRebind | Error::Conflict => { - AdminServiceError::conflict("xAI credential conflicts with current state") - } - Error::RevisionOverflow | Error::Store => { - AdminServiceError::unavailable("xAI credential store is unavailable") - } - } -} - -use gateway_store::redis::CredentialLeaseRepository as _; - -const PROVIDER_ACCOUNT_LEASE_TTL: Duration = Duration::from_secs(10 * 60); -const OAUTH_REFRESH_LEASE_TTL: Duration = Duration::from_secs(5 * 60); -const XAI_CATALOG_CACHE_TTL_SECONDS: u64 = 5 * 60; - -/// Provider 调度和 OAuth refresh 共用的 Redis lease 组合器。 -pub struct ProviderLeaseAdapter { - repository: gateway_store::redis::RedisCredentialLeaseRepository, - process_id: String, - sequence: std::sync::atomic::AtomicU64, - scheduling_cursors: std::sync::Mutex>, -} - -impl ProviderLeaseAdapter { - #[must_use] - pub fn new(repository: gateway_store::redis::RedisCredentialLeaseRepository) -> Self { - Self { - repository, - process_id: format!("gateway_{}", Uuid::now_v7().simple()), - sequence: std::sync::atomic::AtomicU64::new(0), - scheduling_cursors: std::sync::Mutex::new(BTreeMap::new()), - } - } - - fn owner_id(&self, operation: &str) -> String { - let sequence = self - .sequence - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - format!("{}:{operation}:{sequence}", self.process_id) - } - - fn next_scheduling_cursor(&self, provider_instance_id: &ProviderInstanceId) -> Option { - let mut cursors = self.scheduling_cursors.lock().ok()?; - let cursor = cursors.entry(provider_instance_id.clone()).or_default(); - let current = *cursor; - *cursor = cursor.wrapping_add(1); - Some(current) - } - - async fn scheduling_signals( - &self, - accounts: &[gateway_core::engine::credential::ProviderAccountId], - ) -> Result< - BTreeMap< - gateway_core::engine::credential::ProviderAccountId, - gateway_core::engine::credential::AccountRuntimeSignals, - >, - gateway_store::StoreError, - > { - let ids = accounts - .iter() - .map(|account| account.as_str().to_owned()) - .collect::>(); - let signals = self.repository.credential_runtime_signals(&ids).await?; - signals - .into_iter() - .map(|signal| { - let account = - gateway_core::engine::credential::ProviderAccountId::new(signal.resource_id) - .map_err(|_| gateway_store::StoreError::InvalidData { - entity: "credential runtime signal", - message: "Redis returned an invalid account ID".to_owned(), - })?; - Ok(( - account, - gateway_core::engine::credential::AccountRuntimeSignals { - in_flight: signal.in_flight, - last_started_at: signal.last_started_at.map(Into::into), - quota_reset_at: None, - quota_remaining_rank: None, - }, - )) - }) - .collect() - } - - async fn scheduling_lease( - &self, - account_id: &gateway_core::engine::credential::ProviderAccountId, - max_concurrent: u32, - request_interval: Duration, - ttl: Duration, - ) -> Result - { - self.repository - .try_acquire_scheduling_lease(&gateway_store::redis::CredentialSchedulingLeaseRequest { - resource_id: account_id.as_str().to_owned(), - owner_id: self.owner_id("request"), - max_concurrent, - request_interval, - ttl, - }) - .await - } - - async fn refresh_lease( - &self, - account_id: &gateway_core::engine::credential::ProviderAccountId, - ) -> Result, gateway_store::StoreError> { - let request = gateway_store::redis::CredentialLeaseRequest { - scope: gateway_store::redis::CredentialLeaseScope::OAuthRefresh, - resource_id: account_id.as_str().to_owned(), - owner_id: self.owner_id("refresh"), - ttl: OAUTH_REFRESH_LEASE_TTL, - }; - self.repository.try_acquire_guard(request).await - } -} - -fn request_lease_ttl(deadline: SystemTime) -> Option { - deadline - .duration_since(SystemTime::now()) - .ok() - .filter(|remaining| !remaining.is_zero()) - .map(|remaining| remaining.min(PROVIDER_ACCOUNT_LEASE_TTL)) -} - -#[async_trait] -impl provider_openai::credential::CredentialLeaseCoordinator for ProviderLeaseAdapter { - async fn runtime_signals( - &self, - accounts: &[gateway_core::engine::credential::ProviderAccountId], - ) -> Result< - BTreeMap< - gateway_core::engine::credential::ProviderAccountId, - gateway_core::engine::credential::AccountRuntimeSignals, - >, - provider_openai::credential::CredentialLeaseCoordinatorError, - > { - self.scheduling_signals(accounts) - .await - .map_err(|_| provider_openai::credential::CredentialLeaseCoordinatorError::Unavailable) - } - - fn next_round_robin_cursor( - &self, - provider_instance_id: &ProviderInstanceId, - ) -> Result { - self.next_scheduling_cursor(provider_instance_id) - .ok_or(provider_openai::credential::CredentialLeaseCoordinatorError::Unavailable) - } - - async fn try_acquire( - &self, - request: provider_openai::credential::CredentialLeaseRequest, - ) -> Result< - provider_openai::credential::LeaseAcquisition, - provider_openai::credential::CredentialLeaseCoordinatorError, - > { - let ttl = request_lease_ttl(request.deadline) - .ok_or(provider_openai::credential::CredentialLeaseCoordinatorError::Unavailable)?; - match self - .scheduling_lease( - &request.account_id, - request.max_concurrent, - request.request_interval, - ttl, - ) - .await - .map_err(|_| { - provider_openai::credential::CredentialLeaseCoordinatorError::Unavailable - })? { - gateway_store::redis::CredentialSchedulingLeaseAcquisition::Acquired(guard) => Ok( - provider_openai::credential::LeaseAcquisition::Acquired(Box::new(guard)), - ), - gateway_store::redis::CredentialSchedulingLeaseAcquisition::Busy { retry_after } => { - Ok(provider_openai::credential::LeaseAcquisition::Busy { retry_after }) - } - } - } -} - -#[async_trait] -impl provider_xai::GrokCredentialLeaseCoordinator for ProviderLeaseAdapter { - async fn load_scheduling_state( - &self, - provider_instance_id: &gateway_core::routing::ProviderInstanceId, - accounts: &[gateway_core::engine::credential::ProviderAccountId], - ) -> Result< - provider_xai::GrokAccountSchedulingState, - provider_xai::GrokCredentialLeaseCoordinatorError, - > { - let signals = self - .scheduling_signals(accounts) - .await - .map_err(|_| provider_xai::GrokCredentialLeaseCoordinatorError::Unavailable)?; - Ok(provider_xai::GrokAccountSchedulingState { - signals, - sticky_account: None, - round_robin_cursor: self - .next_scheduling_cursor(provider_instance_id) - .ok_or(provider_xai::GrokCredentialLeaseCoordinatorError::Unavailable)?, - }) - } - - async fn try_acquire( - &self, - request: &provider_xai::GrokCredentialLeaseRequest, - ) -> Result< - provider_xai::GrokCredentialLeaseAcquisition, - provider_xai::GrokCredentialLeaseCoordinatorError, - > { - match self - .scheduling_lease( - &request.account_id, - request.max_concurrent_per_account, - request.request_interval, - PROVIDER_ACCOUNT_LEASE_TTL, - ) - .await - .map_err(|_| provider_xai::GrokCredentialLeaseCoordinatorError::Unavailable)? - { - gateway_store::redis::CredentialSchedulingLeaseAcquisition::Acquired(guard) => Ok( - provider_xai::GrokCredentialLeaseAcquisition::Acquired(Box::new(guard)), - ), - gateway_store::redis::CredentialSchedulingLeaseAcquisition::Busy { retry_after } => { - Ok(provider_xai::GrokCredentialLeaseAcquisition::Unavailable { retry_after }) - } - } - } -} - -#[async_trait] -impl provider_openai::credential::CodexRefreshLeaseCoordinator for ProviderLeaseAdapter { - async fn try_acquire( - &self, - request: &provider_openai::credential::CodexRefreshLeaseRequest, - ) -> Result< - provider_openai::credential::CodexRefreshLeaseAcquisition, - provider_openai::credential::CodexRefreshLeaseError, - > { - self.refresh_lease(&request.account_id) - .await - .map(|guard| match guard { - Some(guard) => provider_openai::credential::CodexRefreshLeaseAcquisition::Acquired( - Box::new(guard), - ), - None => provider_openai::credential::CodexRefreshLeaseAcquisition::Unavailable, - }) - .map_err(|_| provider_openai::credential::CodexRefreshLeaseError::Unavailable) - } -} - -#[async_trait] -impl provider_xai::GrokRefreshLeaseCoordinator for ProviderLeaseAdapter { - async fn try_acquire( - &self, - request: &provider_xai::GrokRefreshLeaseRequest, - ) -> Result - { - self.refresh_lease(&request.account_id) - .await - .map(|guard| match guard { - Some(guard) => provider_xai::GrokRefreshLeaseAcquisition::Acquired(Box::new(guard)), - None => provider_xai::GrokRefreshLeaseAcquisition::Unavailable, - }) - .map_err(|_| provider_xai::GrokRefreshLeaseError::Unavailable) - } -} - -/// xAI Provider-owned catalog 与 Redis opaque cache 的唯一转换边界。 -pub struct XaiCatalogCacheAdapter { - repository: Arc, -} - -impl XaiCatalogCacheAdapter { - #[must_use] - pub const fn new( - repository: Arc, - ) -> Self { - Self { repository } - } - - fn key( - account_id: &gateway_core::engine::credential::ProviderAccountId, - revision: gateway_core::engine::credential::CredentialRevision, - ) -> Result< - gateway_store::redis::ProviderAccountCatalogCacheKey, - provider_xai::GrokCatalogCacheError, - > { - Ok(gateway_store::redis::ProviderAccountCatalogCacheKey { - provider_kind: "xai".to_owned(), - provider_account_id: account_id.as_str().to_owned(), - credential_revision: Revision::new(revision.get()) - .map_err(|_| provider_xai::GrokCatalogCacheError::InvalidData)?, - }) - } -} - -#[async_trait] -impl provider_xai::GrokCredentialCatalogCache for XaiCatalogCacheAdapter { - async fn replace( - &self, - catalog: provider_xai::GrokAccountCatalog, - ) -> Result<(), provider_xai::GrokCatalogCacheError> { - let mut document = serde_json::Map::new(); - document.insert("version".to_owned(), serde_json::Value::from(1)); - document.insert( - "observedAt".to_owned(), - serde_json::Value::String(catalog.observed_at().to_rfc3339()), - ); - if let Some(etag) = catalog.seed().etag() { - document.insert( - "etag".to_owned(), - serde_json::Value::String(etag.to_owned()), - ); - } - document.insert( - "models".to_owned(), - serde_json::Value::Array( - catalog - .seed() - .models() - .iter() - .cloned() - .map(serde_json::Value::String) - .collect(), - ), - ); - self.repository - .replace_provider_account_catalog( - &Self::key(catalog.account_id(), catalog.revision())?, - &gateway_core::engine::credential::OpaqueProviderData::new(document), - XAI_CATALOG_CACHE_TTL_SECONDS, - ) - .await - .map_err(|_| provider_xai::GrokCatalogCacheError::Unavailable) - } - - async fn read( - &self, - account_id: &gateway_core::engine::credential::ProviderAccountId, - revision: gateway_core::engine::credential::CredentialRevision, - ) -> Result, provider_xai::GrokCatalogCacheError> { - let Some(document) = self - .repository - .get_provider_account_catalog(&Self::key(account_id, revision)?) - .await - .map_err(|_| provider_xai::GrokCatalogCacheError::Unavailable)? - else { - return Ok(None); - }; - let mut fields = document.into_inner(); - if fields.remove("version").and_then(|value| value.as_u64()) != Some(1) { - return Err(provider_xai::GrokCatalogCacheError::InvalidData); - } - let observed_at = fields - .remove("observedAt") - .and_then(|value| value.as_str().map(ToOwned::to_owned)) - .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) - .map(|value| value.with_timezone(&Utc)) - .ok_or(provider_xai::GrokCatalogCacheError::InvalidData)?; - let etag = match fields.remove("etag") { - None => None, - Some(serde_json::Value::String(value)) => Some(value), - Some(_) => return Err(provider_xai::GrokCatalogCacheError::InvalidData), - }; - let models = fields - .remove("models") - .and_then(|value| value.as_array().cloned()) - .ok_or(provider_xai::GrokCatalogCacheError::InvalidData)? - .into_iter() - .map(|value| { - value - .as_str() - .map(ToOwned::to_owned) - .ok_or(provider_xai::GrokCatalogCacheError::InvalidData) - }) - .collect::, _>>()?; - if !fields.is_empty() { - return Err(provider_xai::GrokCatalogCacheError::InvalidData); - } - let seed = provider_xai::GrokCredentialCatalogSeed::new(models, etag) - .map_err(|_| provider_xai::GrokCatalogCacheError::InvalidData)?; - Ok(Some(provider_xai::GrokAccountCatalog::new( - account_id.clone(), - revision, - observed_at, - seed, - ))) - } - - async fn permits( - &self, - account_id: &gateway_core::engine::credential::ProviderAccountId, - revision: gateway_core::engine::credential::CredentialRevision, - model: &str, - ) -> Result { - Ok(self - .read(account_id, revision) - .await? - .is_some_and(|catalog| catalog.seed().permits(model))) - } -} - -use gateway_api::admin::system::{ - SystemAdminError, SystemAdminErrorKind, SystemUpdateEvent, SystemUpdateEventStream, -}; -use tokio::sync::broadcast; -use tokio_util::sync::CancellationToken as TokioCancellationToken; - -const SYSTEM_UPDATE_APP_BINARY_NAME: &str = "codex-proxy-rs"; -const SYSTEM_UPDATE_DEFAULT_WEB_DIST_DIR: &str = "/app/web/dist"; -const SYSTEM_UPDATE_GITHUB_API_BASE: &str = "https://api.github.com/repos"; -const SYSTEM_UPDATE_CACHE_TTL: Duration = Duration::from_secs(20 * 60); -const SYSTEM_UPDATE_MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024; -const SYSTEM_UPDATE_MAX_CHECKSUM_SIZE: u64 = 1024 * 1024; -const SYSTEM_UPDATE_MAX_EXTRACTED_SIZE: u64 = 1024 * 1024 * 1024; -const SYSTEM_UPDATE_MAX_ARCHIVE_FILES: usize = 20_000; -const SYSTEM_UPDATE_RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS"; -const SYSTEM_UPDATE_REPLACEMENT_START_DELAY_MS: &str = "1200"; - -/// 系统更新运行参数；生产从既有 `CPR_UPDATE_*` 环境变量装配。 -#[derive(Debug, Clone)] -pub struct SystemUpdateConfig { - pub version: String, - pub git_sha: String, - pub build_time: String, - pub deployment_mode: String, - pub build_type: String, - pub update_channel: String, - pub update_repository: Option, - pub github_api_base: String, - pub executable_path: Option, - pub web_dist_dir: PathBuf, - pub update_state_file: PathBuf, - pub update_lock_file: PathBuf, - pub update_temp_dir: PathBuf, - pub self_restart_enabled: bool, -} - -impl SystemUpdateConfig { - fn from_environment() -> Self { - let update_repository = system_update_environment_value("CPR_UPDATE_REPOSITORY"); - let deployment_mode = system_update_environment_value("CPR_DEPLOYMENT_MODE") - .unwrap_or_else(|| "source".to_owned()); - let executable_path = system_update_environment_value("CPR_UPDATE_EXE_PATH") - .map(PathBuf::from) - .or_else(|| { - (deployment_mode == "docker") - .then(|| PathBuf::from("/app/bin").join(SYSTEM_UPDATE_APP_BINARY_NAME)) - }); - let update_state_file = system_update_environment_value("CPR_UPDATE_STATE_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from("/app/.runtime/data/update-state.json")); - let update_lock_file = system_update_environment_value("CPR_UPDATE_LOCK_FILE") - .map(PathBuf::from) - .unwrap_or_else(|| update_state_file.with_extension("lock")); - let update_temp_dir = system_update_environment_value("CPR_UPDATE_TEMP_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| system_update_default_temp_dir(&update_state_file)); - Self { - version: option_env!("CPR_VERSION") - .unwrap_or(env!("CARGO_PKG_VERSION")) - .to_owned(), - git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), - build_time: option_env!("CPR_BUILD_TIME") - .unwrap_or("unknown") - .to_owned(), - deployment_mode, - build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), - update_channel: system_update_environment_value("CPR_UPDATE_CHANNEL") - .unwrap_or_else(|| "stable".to_owned()), - update_repository, - github_api_base: system_update_environment_value("CPR_GITHUB_API_BASE") - .unwrap_or_else(|| SYSTEM_UPDATE_GITHUB_API_BASE.to_owned()), - executable_path, - web_dist_dir: system_update_environment_value("CPR_WEB_DIST_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(SYSTEM_UPDATE_DEFAULT_WEB_DIST_DIR)), - update_state_file, - update_lock_file, - update_temp_dir, - self_restart_enabled: system_update_environment_value("CPR_ENABLE_SELF_RESTART") - .as_deref() - == Some("true"), - } - } - - fn update_support_error(&self) -> Option { - if self.build_type != "release" { - return Some("一键更新需要正式构建包".to_owned()); - } - let Some(repository) = self.update_repository.as_deref() else { - return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); - }; - if let Err(error) = system_update_validate_repository(repository) { - return Some(error.to_string()); - } - if let Err(error) = system_update_validate_github_api_base(&self.github_api_base) { - return Some(error); - } - None - } - - fn release_cache_key(&self) -> String { - format!( - "{}|{}|{}|{}|{}|{}", - self.update_repository.as_deref().unwrap_or_default(), - self.github_api_base, - self.version, - self.deployment_mode, - self.build_type, - self.update_channel, - ) - } - - fn executable_path(&self) -> Result { - if let Some(path) = &self.executable_path { - return Ok(path.clone()); - } - env::current_exe() - .and_then(fs::canonicalize) - .map_err(system_update_internal_with("Failed to resolve executable")) - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct SystemUpdateInfo { - current_version: String, - latest_version: String, - has_update: bool, - deployment_mode: String, - deployment_mode_label: String, - build_type: String, - build_type_label: String, - release_url: Option, - notes: Option, - cached: bool, - update_supported: bool, - unsupported_reason: Option, - warning: Option, -} - -#[derive(Debug, Clone)] -struct CachedSystemUpdateInfo { - key: String, - info: SystemUpdateInfo, - cached_at: std::time::Instant, -} - -struct SystemReleaseCache { - entry: AsyncMutex>, -} - -impl Default for SystemReleaseCache { - fn default() -> Self { - Self { - entry: AsyncMutex::const_new(None), - } - } -} - -#[derive(Debug, Clone, Deserialize)] -struct SystemGitHubRelease { - tag_name: String, - name: Option, - body: Option, - html_url: Option, - prerelease: bool, - #[serde(default)] - assets: Vec, -} - -#[derive(Debug, Clone, Deserialize)] -struct SystemGitHubAsset { - name: String, - browser_download_url: String, - size: u64, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct SystemVersionData { - version: String, - git_sha: String, - build_time: String, - deployment_mode: String, - deployment_mode_label: String, - update_channel: String, - latest_version: String, - has_update: bool, - update_cached: bool, - update_warning: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct SystemUpdateStartedData { - operation_id: String, - deployment_mode: String, - message: String, - need_restart: bool, - target_version: String, -} - -#[derive(Debug, Deserialize, Serialize, Default)] -#[serde(rename_all = "camelCase")] -struct SystemUpdateStatusData { - previous_version: Option, - current_version: Option, - #[serde(default)] - operation: SystemOperationState, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -enum SystemOperationKind { - Update, - Rollback, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "camelCase")] -enum SystemOperationStatus { - #[default] - Idle, - Running, - Succeeded, - Failed, -} - -#[derive(Debug, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -struct SystemOperationState { - operation_id: Option, - kind: Option, - status: SystemOperationStatus, - target_version: Option, - message: Option, - error: Option, - started_at: Option, - finished_at: Option, -} - -impl Default for SystemOperationState { - fn default() -> Self { - Self { - operation_id: None, - kind: None, - status: SystemOperationStatus::Idle, - target_version: None, - message: None, - error: None, - started_at: None, - finished_at: None, - } - } -} - -#[derive(Debug, Clone, Copy, Serialize)] -#[serde(rename_all = "lowercase")] -enum SystemUpdateLogLevel { - Info, - Success, - Warning, - Error, -} - -struct SystemUpdateEventSender { - sender: broadcast::Sender<(String, serde_json::Value)>, - sequence: AtomicU64, -} - -impl Default for SystemUpdateEventSender { - fn default() -> Self { - let (sender, _receiver) = broadcast::channel(256); - Self { - sender, - sequence: AtomicU64::new(0), - } - } -} - -impl SystemUpdateEventSender { - fn subscribe(&self) -> broadcast::Receiver<(String, serde_json::Value)> { - self.sender.subscribe() - } - - fn emit( - &self, - level: SystemUpdateLogLevel, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit_with_terminal(level, operation_id, step, message, false); - } - - fn emit_terminal( - &self, - level: SystemUpdateLogLevel, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - ) { - self.emit_with_terminal(level, operation_id, step, message, true); - } - - fn emit_with_terminal( - &self, - level: SystemUpdateLogLevel, - operation_id: Option<&str>, - step: Option<&str>, - message: impl Into, - terminal: bool, - ) { - let now = Utc::now(); - let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); - let id = format!( - "update-log-{}-{sequence}", - now.timestamp_nanos_opt() - .unwrap_or_else(|| now.timestamp_millis()) - ); - let data = serde_json::json!({ - "operationId": operation_id, - "level": level, - "step": step, - "message": message.into(), - "terminal": terminal, - "at": now.to_rfc3339(), - }); - let _ = self.sender.send((id, data)); - } -} - -/// 当前进程版本、更新、回滚与真实 shutdown 信号的系统端口。 -pub struct ProcessSystemAdminService { - shutdown: TokioCancellationToken, - events: SystemUpdateEventSender, - config: SystemUpdateConfig, - operation_lock: AsyncMutex<()>, - release_cache: SystemReleaseCache, -} - -impl ProcessSystemAdminService { - #[must_use] - pub fn new(shutdown: TokioCancellationToken) -> Self { - Self::with_config(shutdown, SystemUpdateConfig::from_environment()) - } - - /// 使用显式配置装配系统更新端口，供组合根和黑盒边界测试使用。 - #[must_use] - pub fn with_config(shutdown: TokioCancellationToken, config: SystemUpdateConfig) -> Self { - Self { - shutdown, - events: SystemUpdateEventSender::default(), - config, - operation_lock: AsyncMutex::const_new(()), - release_cache: SystemReleaseCache::default(), - } - } - - async fn version_value(&self) -> Result { - let info = - system_update_check_latest_release(&self.release_cache, &self.config, false).await; - system_update_json_value(SystemVersionData { - version: self.config.version.clone(), - git_sha: self.config.git_sha.clone(), - build_time: self.config.build_time.clone(), - deployment_mode: self.config.deployment_mode.clone(), - deployment_mode_label: deployment_mode_label(&self.config.deployment_mode).to_owned(), - update_channel: self.config.update_channel.clone(), - latest_version: info.latest_version, - has_update: info.has_update, - update_cached: info.cached, - update_warning: info.warning, - }) - } - - async fn perform_update_inner( - &self, - target_version: Option, - ) -> Result { - let _operation_guard = self - .operation_lock - .try_lock() - .map_err(|_| system_conflict("System update already running"))?; - if let Some(reason) = self.config.update_support_error() { - self.events.emit_terminal( - SystemUpdateLogLevel::Error, - None, - Some("preflight"), - reason.clone(), - ); - return Err(system_conflict(reason)); - } - let confirmed_target = system_update_confirmed_target(target_version)?; - let repository = self - .config - .update_repository - .as_deref() - .ok_or_else(|| system_conflict("检查更新需要配置 CPR_UPDATE_REPOSITORY"))?; - - self.events.emit( - SystemUpdateLogLevel::Info, - None, - Some("release"), - "正在获取最新 Release 信息", - ); - let release = system_update_fetch_latest_release(&self.config.github_api_base, repository) - .await - .inspect_err(|error| { - self.events.emit_terminal( - SystemUpdateLogLevel::Error, - None, - Some("release"), - error.to_string(), - ); - })?; - let info = system_update_info_from_release(&self.config, release.clone()); - if info.latest_version != confirmed_target { - let message = format!( - "远端最新版本已变更为 v{}，请重新检查并确认", - info.latest_version - ); - self.events.emit_terminal( - SystemUpdateLogLevel::Warning, - None, - Some("release"), - message.clone(), - ); - return Err(system_conflict(message)); - } - if !info.has_update { - self.events.emit_terminal( - SystemUpdateLogLevel::Warning, - None, - Some("release"), - "当前版本已是最新", - ); - return Err(system_conflict("Already up to date")); - } - - let target_version = info.latest_version; - let operation_id = system_update_operation_id("update"); - let file_lock = SystemUpdateFileLock::acquire(&self.config.update_lock_file)?; - system_update_set_operation_running( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Update, - Some(&target_version), - &self.config.version, - )?; - self.events.emit( - SystemUpdateLogLevel::Info, - Some(&operation_id), - Some("prepare"), - format!("准备更新到 v{target_version}"), - ); - let result = self - .perform_release_update(&release, &target_version, &operation_id) - .await; - match &result { - Ok(()) => self.events.emit_terminal( - SystemUpdateLogLevel::Success, - Some(&operation_id), - Some("done"), - "更新文件已替换，等待服务重启生效", - ), - Err(error) => self.events.emit_terminal( - SystemUpdateLogLevel::Error, - Some(&operation_id), - Some("failed"), - error.to_string(), - ), - } - system_update_finish_operation( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Update, - result.as_ref().ok().map(|()| target_version.clone()), - result.as_ref().err().map(ToString::to_string), - ); - drop(file_lock); - result?; - - Ok(SystemUpdateStartedData { - operation_id, - deployment_mode: self.config.deployment_mode.clone(), - message: "更新完成，请重启服务。".to_owned(), - need_restart: true, - target_version, - }) - } - - async fn perform_release_update( - &self, - release: &SystemGitHubRelease, - version: &str, - operation_id: &str, - ) -> Result<(), SystemAdminError> { - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("asset"), - "正在选择匹配当前平台的更新包", - ); - let archive = system_update_select_release_archive(release, version)?; - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("asset"), - format!( - "已选择更新包 {} ({})", - archive.name, - system_update_format_bytes(archive.size) - ), - ); - system_update_validate_download_url( - &archive.browser_download_url, - &self.config.github_api_base, - )?; - if archive.size == 0 || archive.size > SYSTEM_UPDATE_MAX_DOWNLOAD_SIZE { - return Err(system_invalid("Release archive has an invalid size")); - } - let checksum = release - .assets - .iter() - .find(|asset| asset.name == "checksums.txt") - .ok_or_else(|| system_bad_gateway("Release checksums.txt is required"))?; - if checksum.size > SYSTEM_UPDATE_MAX_CHECKSUM_SIZE { - return Err(system_invalid("Release checksums.txt is too large")); - } - system_update_validate_download_url( - &checksum.browser_download_url, - &self.config.github_api_base, - )?; - - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("prepare"), - "正在创建临时更新目录", - ); - fs::create_dir_all(&self.config.update_temp_dir).map_err(system_update_internal_with( - "Failed to prepare update temp dir", - ))?; - let update_temp_root = fs::canonicalize(&self.config.update_temp_dir).map_err( - system_update_internal_with("Failed to resolve update temp dir"), - )?; - let temp_dir = system_update_unique_temp_dir(&update_temp_root)?; - let temp_guard = SystemUpdateTempDir::new(temp_dir.clone()); - let archive_path = temp_dir.join(&archive.name); - - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("download"), - "开始下载更新包", - ); - system_update_download_file( - &archive.browser_download_url, - &archive_path, - archive.size, - &self.config.github_api_base, - Some(SystemDownloadProgress { - operation_id, - events: &self.events, - }), - ) - .await?; - self.events.emit( - SystemUpdateLogLevel::Success, - Some(operation_id), - Some("download"), - "更新包下载完成", - ); - - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("checksum"), - "正在校验 checksum", - ); - system_update_verify_checksum( - &archive_path, - &archive.name, - &checksum.browser_download_url, - &self.config.github_api_base, - ) - .await?; - self.events.emit( - SystemUpdateLogLevel::Success, - Some(operation_id), - Some("checksum"), - "checksum 校验通过", - ); - - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("extract"), - "正在解压更新包", - ); - let extracted = system_update_extract_release_archive(&archive_path, &temp_dir)?; - self.events.emit( - SystemUpdateLogLevel::Success, - Some(operation_id), - Some("extract"), - "更新包解压完成", - ); - - self.events.emit( - SystemUpdateLogLevel::Info, - Some(operation_id), - Some("replace"), - "正在替换应用文件", - ); - let executable_path = self.config.executable_path()?; - system_update_replace_release_files( - &executable_path, - &self.config.web_dist_dir, - extracted, - )?; - self.events.emit( - SystemUpdateLogLevel::Success, - Some(operation_id), - Some("replace"), - "应用文件替换完成", - ); - drop(temp_guard); - Ok(()) - } - - async fn rollback_inner(&self) -> Result { - let _operation_guard = self - .operation_lock - .try_lock() - .map_err(|_| system_conflict("系统操作正在执行中"))?; - if let Some(reason) = self.config.update_support_error() { - return Err(system_conflict(reason)); - } - let operation_id = system_update_operation_id("rollback"); - let file_lock = SystemUpdateFileLock::acquire(&self.config.update_lock_file)?; - system_update_set_operation_running( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - &self.config.version, - )?; - self.events.emit( - SystemUpdateLogLevel::Info, - Some(&operation_id), - Some("rollback"), - "正在恢复上一版本", - ); - let result = system_update_rollback_release(&self.config); - match &result { - Ok(()) => self.events.emit_terminal( - SystemUpdateLogLevel::Success, - Some(&operation_id), - Some("done"), - "上一版本已恢复，等待服务重启生效", - ), - Err(error) => self.events.emit_terminal( - SystemUpdateLogLevel::Error, - Some(&operation_id), - Some("failed"), - error.to_string(), - ), - } - system_update_finish_operation( - &self.config.update_state_file, - &operation_id, - SystemOperationKind::Rollback, - None, - result.as_ref().err().map(ToString::to_string), - ); - drop(file_lock); - result?; - Ok(operation_id) - } -} - -fn deployment_mode_label(mode: &str) -> &'static str { - match mode { - "docker" => "Docker", - "binary" => "二进制", - _ => "源码运行", - } -} - -#[async_trait] -impl SystemAdminService for ProcessSystemAdminService { - async fn version(&self) -> Result { - self.version_value().await - } - - async fn update_detail(&self, refresh: bool) -> Result { - system_update_json_value( - system_update_check_latest_release(&self.release_cache, &self.config, refresh).await, - ) - } - - fn update_events(&self) -> SystemUpdateEventStream { - let receiver = self.events.subscribe(); - Box::pin(futures::stream::unfold( - (receiver, false), - |(mut receiver, close_after_send)| async move { - if close_after_send { - return None; - } - loop { - match receiver.recv().await { - Ok((id, data)) => { - let terminal = data - .get("terminal") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); - return Some((SystemUpdateEvent { id, data }, (receiver, terminal))); - } - Err(broadcast::error::RecvError::Lagged(_)) => continue, - Err(broadcast::error::RecvError::Closed) => return None, - } - } - }, - )) - } - - async fn perform_update( - &self, - target_version: Option, - ) -> Result { - system_update_json_value(self.perform_update_inner(target_version).await?) - } - - async fn update_status(&self) -> Result { - system_update_json_value(system_update_read_state(&self.config.update_state_file)?) - } - - async fn rollback(&self) -> Result { - let operation_id = self.rollback_inner().await?; - Ok(serde_json::json!({ - "message": "回滚完成，请重启服务。", - "needRestart": true, - "operationId": operation_id, - })) - } - - async fn restart(&self) -> Result { - if !self.config.self_restart_enabled { - return Err(system_conflict( - "自重启未启用，请设置 CPR_ENABLE_SELF_RESTART=true", - )); - } - let message = if self.config.deployment_mode == "docker" { - "已安排进程内重启" - } else { - system_update_spawn_replacement_process(&self.config)?; - "已安排自重启" - }; - let operation_id = system_update_operation_id("restart"); - let shutdown = self.shutdown.clone(); - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(500)).await; - shutdown.cancel(); - }); - Ok(serde_json::json!({ - "message": message, - "operationId": operation_id, - })) - } -} - -fn system_conflict(message: impl Into) -> SystemAdminError { - SystemAdminError::new(SystemAdminErrorKind::Conflict, message) -} - -fn system_invalid(message: impl Into) -> SystemAdminError { - SystemAdminError::new(SystemAdminErrorKind::Invalid, message) -} - -fn system_bad_gateway(message: impl Into) -> SystemAdminError { - SystemAdminError::new(SystemAdminErrorKind::BadGateway, message) -} - -fn system_internal(message: impl Into) -> SystemAdminError { - SystemAdminError::new(SystemAdminErrorKind::Internal, message) -} - -fn system_update_internal_with( - context: &'static str, -) -> impl FnOnce(E) -> SystemAdminError { - move |error| system_internal(format!("{context}: {error}")) -} - -fn system_update_bad_gateway_with( - context: &'static str, -) -> impl FnOnce(E) -> SystemAdminError { - move |error| system_bad_gateway(format!("{context}: {error}")) -} - -fn system_update_json_value(value: impl Serialize) -> Result { - serde_json::to_value(value).map_err(system_update_internal_with( - "Failed to encode system update response", - )) -} - -async fn system_update_check_latest_release( - cache: &SystemReleaseCache, - config: &SystemUpdateConfig, - force: bool, -) -> SystemUpdateInfo { - let Some(repository) = config.update_repository.as_deref() else { - return system_update_unavailable_info( - config, - "检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned(), - None, - ); - }; - if let Err(error) = system_update_validate_repository(repository) { - return system_update_unavailable_info(config, error.to_string(), None); - } - if let Err(reason) = system_update_validate_github_api_base(&config.github_api_base) { - return system_update_unavailable_info(config, reason, None); - } - let cache_key = config.release_cache_key(); - if !force && let Some(info) = system_update_cached_release_info(cache, &cache_key).await { - return info; - } - - match system_update_fetch_latest_release(&config.github_api_base, repository).await { - Ok(release) => { - let info = system_update_info_from_release(config, release); - system_update_cache_release_info(cache, cache_key, &info).await; - info - } - Err(error) => system_update_cached_release_info(cache, &cache_key) - .await - .unwrap_or_else(|| { - let mut info = system_update_base_info(config); - info.update_supported = false; - info.unsupported_reason = config.update_support_error(); - info.warning = Some(error.to_string()); - info - }), - } -} - -fn system_update_base_info(config: &SystemUpdateConfig) -> SystemUpdateInfo { - SystemUpdateInfo { - current_version: config.version.clone(), - latest_version: config.version.clone(), - has_update: false, - deployment_mode: config.deployment_mode.clone(), - deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_owned(), - build_type: config.build_type.clone(), - build_type_label: system_update_build_type_label(&config.build_type).to_owned(), - release_url: None, - notes: None, - cached: false, - update_supported: config.update_support_error().is_none(), - unsupported_reason: config.update_support_error(), - warning: None, - } -} - -fn system_update_unavailable_info( - config: &SystemUpdateConfig, - reason: String, - warning: Option, -) -> SystemUpdateInfo { - let mut info = system_update_base_info(config); - info.update_supported = false; - info.unsupported_reason = Some(reason); - info.warning = warning; - info -} - -async fn system_update_cached_release_info( - cache: &SystemReleaseCache, - cache_key: &str, -) -> Option { - let entry = cache.entry.lock().await; - let cached = entry.as_ref()?; - if cached.key != cache_key || cached.cached_at.elapsed() > SYSTEM_UPDATE_CACHE_TTL { - return None; - } - let mut info = cached.info.clone(); - drop(entry); - info.cached = true; - Some(info) -} - -async fn system_update_cache_release_info( - cache: &SystemReleaseCache, - cache_key: String, - info: &SystemUpdateInfo, -) { - let mut entry = cache.entry.lock().await; - *entry = Some(CachedSystemUpdateInfo { - key: cache_key, - info: info.clone(), - cached_at: std::time::Instant::now(), - }); -} - -async fn system_update_fetch_latest_release( - api_base: &str, - repository: &str, -) -> Result { - system_update_validate_github_api_base(api_base).map_err(system_conflict)?; - system_update_validate_repository(repository)?; - let url = format!( - "{}/{repository}/releases/latest", - api_base.trim_end_matches('/') - ); - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(system_update_internal_with( - "Failed to create release HTTP client", - ))?; - let response = client - .get(url) - .header(reqwest::header::ACCEPT, "application/vnd.github+json") - .header(reqwest::header::USER_AGENT, SYSTEM_UPDATE_APP_BINARY_NAME) - .send() - .await - .map_err(system_update_bad_gateway_with( - "GitHub release check failed", - ))?; - let status = response.status(); - if !status.is_success() { - return Err(system_bad_gateway(format!( - "GitHub release check failed with {status}" - ))); - } - response - .json::() - .await - .map_err(system_update_bad_gateway_with( - "Invalid GitHub release response", - )) -} - -fn system_update_info_from_release( - config: &SystemUpdateConfig, - release: SystemGitHubRelease, -) -> SystemUpdateInfo { - let latest_version = system_update_normalize_version_tag(&release.tag_name); - let has_update = system_update_release_allowed(config, &release) - && system_update_version_is_newer(&config.version, &latest_version).unwrap_or(false); - let unsupported_reason = config.update_support_error(); - SystemUpdateInfo { - current_version: config.version.clone(), - latest_version, - has_update, - deployment_mode: config.deployment_mode.clone(), - deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_owned(), - build_type: config.build_type.clone(), - build_type_label: system_update_build_type_label(&config.build_type).to_owned(), - release_url: release.html_url, - notes: release.body.or(release.name), - cached: false, - update_supported: unsupported_reason.is_none(), - unsupported_reason, - warning: None, - } -} - -fn system_update_release_allowed( - config: &SystemUpdateConfig, - release: &SystemGitHubRelease, -) -> bool { - config.update_channel != "stable" || !release.prerelease -} - -fn system_update_version_is_newer(current: &str, latest: &str) -> Option { - let current = semver::Version::parse(&system_update_normalize_version_tag(current)).ok()?; - let latest = semver::Version::parse(&system_update_normalize_version_tag(latest)).ok()?; - Some(latest > current) -} - -fn system_update_confirmed_target( - target_version: Option, -) -> Result { - let Some(target_version) = target_version else { - return Err(system_conflict("更新前需要确认目标版本")); - }; - let target_version = system_update_normalize_version_tag(&target_version); - if target_version.is_empty() { - return Err(system_invalid("目标版本不能为空")); - } - semver::Version::parse(&target_version).map_err(|_| system_invalid("目标版本格式无效"))?; - Ok(target_version) -} - -fn system_update_normalize_version_tag(version: &str) -> String { - version.trim().trim_start_matches('v').to_owned() -} - -fn system_update_select_release_archive<'a>( - release: &'a SystemGitHubRelease, - version: &str, -) -> Result<&'a SystemGitHubAsset, SystemAdminError> { - let os_aliases = system_update_platform_os_aliases(); - let arch_aliases = system_update_platform_arch_aliases(); - let normalized = system_update_normalize_version_tag(version); - release - .assets - .iter() - .find(|asset| { - let name = asset.name.as_str(); - name.contains(SYSTEM_UPDATE_APP_BINARY_NAME) - && name.contains(&normalized) - && system_update_asset_matches_platform(name, os_aliases, arch_aliases) - && !name.ends_with(".txt") - }) - .or_else(|| { - release.assets.iter().find(|asset| { - let name = asset.name.as_str(); - system_update_asset_matches_platform(name, os_aliases, arch_aliases) - && !name.ends_with(".txt") - }) - }) - .ok_or_else(|| { - system_conflict(format!( - "No compatible release archive found for {}/{}", - env::consts::OS, - env::consts::ARCH - )) - }) -} - -fn system_update_platform_os_aliases() -> &'static [&'static str] { - match env::consts::OS { - "macos" => &["macos", "darwin"], - "windows" => &["windows", "win32"], - "linux" => &["linux"], - _ => &[env::consts::OS], - } -} - -fn system_update_platform_arch_aliases() -> &'static [&'static str] { - match env::consts::ARCH { - "x86_64" => &["x86_64", "amd64"], - "aarch64" => &["aarch64", "arm64"], - _ => &[env::consts::ARCH], - } -} - -fn system_update_asset_matches_platform( - name: &str, - os_aliases: &[&str], - arch_aliases: &[&str], -) -> bool { - os_aliases.iter().any(|alias| name.contains(alias)) - && arch_aliases.iter().any(|alias| name.contains(alias)) -} - -fn system_update_validate_repository(repository: &str) -> Result<(), SystemAdminError> { - let mut segments = repository.split('/'); - let owner = segments.next().unwrap_or_default(); - let name = segments.next().unwrap_or_default(); - if owner.is_empty() - || name.is_empty() - || segments.next().is_some() - || !owner.chars().all(system_update_repository_character) - || !name.chars().all(system_update_repository_character) - { - return Err(system_conflict( - "CPR_UPDATE_REPOSITORY must use the owner/repository form", - )); - } - Ok(()) -} - -fn system_update_repository_character(character: char) -> bool { - character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.') -} - -fn system_update_validate_github_api_base(raw_url: &str) -> Result<(), String> { - let url = reqwest::Url::parse(raw_url) - .map_err(|error| format!("Invalid GitHub API base: {error}"))?; - if !url.username().is_empty() - || url.password().is_some() - || url.query().is_some() - || url.fragment().is_some() - { - return Err("GitHub API base must not include credentials, query, or fragment".to_owned()); - } - if url.path().trim_end_matches('/') != "/repos" { - return Err("GitHub API base path must be /repos".to_owned()); - } - if system_update_url_host_is_loopback(&url) { - return if url.scheme() == "http" || url.scheme() == "https" { - Ok(()) - } else { - Err("Loopback GitHub API base must use HTTP or HTTPS".to_owned()) - }; - } - if url.scheme() != "https" || url.host_str() != Some("api.github.com") { - return Err("GitHub API base must be https://api.github.com/repos".to_owned()); - } - Ok(()) -} - -fn system_update_validate_download_url( - raw_url: &str, - github_api_base: &str, -) -> Result<(), SystemAdminError> { - let url = reqwest::Url::parse(raw_url) - .map_err(|error| system_invalid(format!("Invalid download URL: {error}")))?; - if system_update_local_download_allowed(&url, github_api_base) { - return Ok(()); - } - if url.scheme() != "https" { - return Err(system_invalid("Only HTTPS release downloads are allowed")); - } - let host = url - .host_str() - .ok_or_else(|| system_invalid("Download URL is missing host"))?; - if system_update_github_download_host_allowed(host) { - Ok(()) - } else { - Err(system_invalid(format!( - "Download host is not allowed: {host}" - ))) - } -} - -fn system_update_github_download_host_allowed(host: &str) -> bool { - host == "github.com" - || host.ends_with(".github.com") - || host == "objects.githubusercontent.com" - || host.ends_with(".objects.githubusercontent.com") -} - -fn system_update_local_download_allowed(url: &reqwest::Url, github_api_base: &str) -> bool { - if !matches!(url.scheme(), "http" | "https") || !system_update_url_host_is_loopback(url) { - return false; - } - let Ok(api_base) = reqwest::Url::parse(github_api_base) else { - return false; - }; - system_update_url_host_is_loopback(&api_base) - && url.scheme() == api_base.scheme() - && url.host_str() == api_base.host_str() - && url.port_or_known_default() == api_base.port_or_known_default() -} - -fn system_update_url_host_is_loopback(url: &reqwest::Url) -> bool { - url.host_str().is_some_and(|host| { - host == "localhost" - || host == "127.0.0.1" - || host == "::1" - || host - .parse::() - .is_ok_and(|address| address.is_loopback()) - }) -} - -fn system_update_download_client( - github_api_base: &str, - timeout: Duration, -) -> Result { - let github_api_base = github_api_base.to_owned(); - reqwest::Client::builder() - .timeout(timeout) - .redirect(reqwest::redirect::Policy::custom(move |attempt| { - if system_update_download_url_allowed(attempt.url(), &github_api_base) { - attempt.follow() - } else { - attempt.error(std::io::Error::other( - "release redirect target is not trusted", - )) - } - })) - .build() - .map_err(system_update_internal_with( - "Failed to create download HTTP client", - )) -} - -fn system_update_download_url_allowed(url: &reqwest::Url, github_api_base: &str) -> bool { - system_update_local_download_allowed(url, github_api_base) - || (url.scheme() == "https" - && url - .host_str() - .is_some_and(system_update_github_download_host_allowed)) -} - -#[derive(Clone, Copy)] -struct SystemDownloadProgress<'a> { - operation_id: &'a str, - events: &'a SystemUpdateEventSender, -} - -async fn system_update_download_file( - url: &str, - destination: &Path, - expected_size: u64, - github_api_base: &str, - progress: Option>, -) -> Result<(), SystemAdminError> { - let client = system_update_download_client(github_api_base, Duration::from_secs(120))?; - let mut response = client - .get(url) - .header(reqwest::header::ACCEPT_ENCODING, "identity") - .send() - .await - .map_err(system_update_bad_gateway_with("Download failed"))?; - if !response.status().is_success() { - return Err(system_bad_gateway(format!( - "Download failed with {}", - response.status() - ))); - } - if response - .content_length() - .is_some_and(|size| size != expected_size) - { - return Err(system_bad_gateway( - "Downloaded size does not match release metadata", - )); - } - - let mut file = fs::File::create(destination).map_err(system_update_internal_with( - "Failed to create download file", - ))?; - let mut downloaded = 0_u64; - let mut next_progress = 10_u64; - while let Some(chunk) = response - .chunk() - .await - .map_err(system_update_bad_gateway_with("Download stream failed"))? - { - downloaded = downloaded.saturating_add(chunk.len() as u64); - if downloaded > expected_size || downloaded > SYSTEM_UPDATE_MAX_DOWNLOAD_SIZE { - return Err(system_invalid("Download exceeds declared release size")); - } - file.write_all(&chunk) - .map_err(system_update_internal_with("Failed to write download"))?; - if let Some(progress) = progress { - let percent = system_update_download_percent(downloaded, expected_size); - if percent >= next_progress || downloaded == expected_size { - progress.events.emit( - SystemUpdateLogLevel::Info, - Some(progress.operation_id), - Some("download"), - format!( - "已下载 {} / {} ({percent}%)", - system_update_format_bytes(downloaded), - system_update_format_bytes(expected_size) - ), - ); - next_progress = percent.saturating_add(10); - } - } - } - file.sync_all() - .map_err(system_update_internal_with("Failed to sync download"))?; - if downloaded != expected_size { - return Err(system_bad_gateway( - "Downloaded size does not match release metadata", - )); - } - Ok(()) -} - -async fn system_update_verify_checksum( - file_path: &Path, - file_name: &str, - checksum_url: &str, - github_api_base: &str, -) -> Result<(), SystemAdminError> { - let client = system_update_download_client(github_api_base, Duration::from_secs(30))?; - let mut response = client - .get(checksum_url) - .header(reqwest::header::ACCEPT_ENCODING, "identity") - .send() - .await - .map_err(system_update_bad_gateway_with("Checksum download failed"))?; - if !response.status().is_success() { - return Err(system_bad_gateway(format!( - "Checksum download failed with {}", - response.status() - ))); - } - if response - .content_length() - .is_some_and(|size| size > SYSTEM_UPDATE_MAX_CHECKSUM_SIZE) - { - return Err(system_bad_gateway("Checksum document is too large")); - } - let mut body = Vec::new(); - while let Some(chunk) = response - .chunk() - .await - .map_err(system_update_bad_gateway_with("Checksum read failed"))? - { - if body.len().saturating_add(chunk.len()) > SYSTEM_UPDATE_MAX_CHECKSUM_SIZE as usize { - return Err(system_bad_gateway("Checksum document is too large")); - } - body.extend_from_slice(&chunk); - } - let body = std::str::from_utf8(&body) - .map_err(system_update_bad_gateway_with("Invalid checksum document"))?; - let expected = body.lines().find_map(|line| { - let mut parts = line.split_whitespace(); - let hash = parts.next()?; - let name = parts.next()?; - (Path::new(name).file_name()?.to_string_lossy() == file_name).then(|| hash.to_owned()) - }); - let expected = expected.ok_or_else(|| system_bad_gateway("Checksum not found"))?; - if expected.len() != 64 || !expected.bytes().all(|byte| byte.is_ascii_hexdigit()) { - return Err(system_bad_gateway("Invalid SHA-256 checksum")); - } - let actual = system_update_sha256_file(file_path)?; - if expected.eq_ignore_ascii_case(&actual) { - Ok(()) - } else { - Err(system_bad_gateway("Checksum mismatch")) - } -} - -fn system_update_sha256_file(path: &Path) -> Result { - use sha2::Digest as _; - - let mut file = fs::File::open(path) - .map_err(system_update_internal_with("Failed to open checksum file"))?; - let mut hasher = sha2::Sha256::new(); - let mut buffer = [0_u8; 8192]; - loop { - let read = std::io::Read::read(&mut file, &mut buffer) - .map_err(system_update_internal_with("Failed to read checksum file"))?; - if read == 0 { - break; - } - hasher.update(&buffer[..read]); - } - Ok(hex::encode(hasher.finalize())) -} - -fn system_update_download_percent(downloaded: u64, total_size: u64) -> u64 { - if total_size == 0 { - return 0; - } - downloaded - .saturating_mul(100) - .saturating_div(total_size) - .min(100) -} - -fn system_update_format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = KIB * 1024; - if bytes >= MIB { - return format!("{:.1} MiB", bytes as f64 / MIB as f64); - } - if bytes >= KIB { - return format!("{:.1} KiB", bytes as f64 / KIB as f64); - } - format!("{bytes} B") -} - -fn system_update_build_type_label(value: &str) -> &str { - match value { - "release" => "正式构建", - "source" => "源码构建", - "dev" => "开发构建", - _ => value, - } -} - -#[derive(Debug)] -struct SystemExtractedRelease { - binary_path: PathBuf, - web_dist_dir: Option, -} - -fn system_update_extract_release_archive( - archive_path: &Path, - temp_dir: &Path, -) -> Result { - let file = fs::File::open(archive_path) - .map_err(system_update_internal_with("Failed to open archive"))?; - let decoder = flate2::read::GzDecoder::new(file); - let mut archive = tar::Archive::new(decoder); - let binary_path = temp_dir.join(SYSTEM_UPDATE_APP_BINARY_NAME); - let web_dist_dir = temp_dir.join("web-dist"); - let mut found_binary = false; - let mut found_web = false; - let mut extracted_size = 0_u64; - let mut file_count = 0_usize; - - for entry in archive - .entries() - .map_err(system_update_internal_with("Failed to read archive"))? - { - let mut entry = entry.map_err(system_update_internal_with("Invalid archive entry"))?; - let path = entry - .path() - .map_err(system_update_internal_with("Invalid archive path"))? - .to_path_buf(); - if system_update_unsafe_archive_path(&path) { - return Err(system_invalid("Unsafe archive path")); - } - if !entry.header().entry_type().is_file() { - continue; - } - file_count = file_count.saturating_add(1); - extracted_size = extracted_size.saturating_add(entry.header().size().unwrap_or(u64::MAX)); - if file_count > SYSTEM_UPDATE_MAX_ARCHIVE_FILES - || extracted_size > SYSTEM_UPDATE_MAX_EXTRACTED_SIZE - { - return Err(system_invalid( - "Release archive expands beyond safety limits", - )); - } - - if path - .file_name() - .is_some_and(|name| name == SYSTEM_UPDATE_APP_BINARY_NAME) - { - if found_binary { - return Err(system_invalid( - "Release archive contains duplicate binaries", - )); - } - entry - .unpack(&binary_path) - .map_err(system_update_internal_with("Failed to extract binary"))?; - found_binary = true; - continue; - } - - if let Some(relative) = system_update_web_dist_relative_path(&path) { - if relative.as_os_str().is_empty() { - continue; - } - let target = web_dist_dir.join(relative); - if let Some(parent) = target.parent() { - fs::create_dir_all(parent).map_err(system_update_internal_with( - "Failed to create web asset dir", - ))?; - } - entry - .unpack(&target) - .map_err(system_update_internal_with("Failed to extract web asset"))?; - found_web = true; - } - } - - if !found_binary { - return Err(system_invalid( - "Release archive does not contain codex-proxy-rs", - )); - } - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt as _; - - fs::set_permissions(&binary_path, fs::Permissions::from_mode(0o755)) - .map_err(system_update_internal_with("Failed to chmod binary"))?; - } - Ok(SystemExtractedRelease { - binary_path, - web_dist_dir: found_web.then_some(web_dist_dir), - }) -} - -fn system_update_replace_release_files( - executable_path: &Path, - web_dist_dir: &Path, - extracted: SystemExtractedRelease, -) -> Result<(), SystemAdminError> { - let web_backup = system_update_backup_path_for(web_dist_dir); - let web_replaced = if let Some(new_web) = extracted.web_dist_dir { - system_update_replace_dir(web_dist_dir, &web_backup, &new_web)?; - true - } else { - false - }; - - let binary_backup = system_update_backup_path_for(executable_path); - if binary_backup.exists() - && let Err(error) = fs::remove_file(&binary_backup) - { - let mut rollback_errors = Vec::new(); - if web_replaced { - system_update_collect_rollback_error( - &mut rollback_errors, - "restore web assets", - system_update_restore_dir(web_dist_dir, &web_backup), - ); - } - return Err(system_update_error_with_rollback( - "Failed to remove old binary backup", - error, - rollback_errors, - )); - } - if let Err(error) = system_update_move_file(executable_path, &binary_backup) { - let mut rollback_errors = Vec::new(); - if web_replaced { - system_update_collect_rollback_error( - &mut rollback_errors, - "restore web assets", - system_update_restore_dir(web_dist_dir, &web_backup), - ); - } - return Err(system_update_error_with_rollback( - "Binary backup failed", - error, - rollback_errors, - )); - } - if let Err(error) = system_update_move_file(&extracted.binary_path, executable_path) { - let mut rollback_errors = Vec::new(); - if executable_path.exists() { - system_update_collect_rollback_error( - &mut rollback_errors, - "remove partial replacement binary", - fs::remove_file(executable_path), - ); - } - system_update_collect_rollback_error( - &mut rollback_errors, - "restore previous binary", - system_update_move_file(&binary_backup, executable_path), - ); - if web_replaced { - system_update_collect_rollback_error( - &mut rollback_errors, - "restore web assets", - system_update_restore_dir(web_dist_dir, &web_backup), - ); - } - return Err(system_update_error_with_rollback( - "Binary replace failed", - error, - rollback_errors, - )); - } - Ok(()) -} - -fn system_update_replace_dir( - current: &Path, - backup: &Path, - replacement: &Path, -) -> Result<(), SystemAdminError> { - if backup.exists() { - fs::remove_dir_all(backup).map_err(system_update_internal_with( - "Failed to remove old web backup", - ))?; - } - if current.exists() { - system_update_move_dir(current, backup) - .map_err(system_update_internal_with("Failed to backup web assets"))?; - } - if let Err(error) = system_update_move_dir(replacement, current) { - let mut rollback_errors = Vec::new(); - if backup.exists() { - system_update_collect_rollback_error( - &mut rollback_errors, - "restore previous web assets", - system_update_restore_dir(current, backup), - ); - } - return Err(system_update_error_with_rollback( - "Failed to replace web assets", - error, - rollback_errors, - )); - } - Ok(()) -} - -fn system_update_rollback_release(config: &SystemUpdateConfig) -> Result<(), SystemAdminError> { - let executable_path = config.executable_path()?; - let binary_backup = system_update_backup_path_for(&executable_path); - if !binary_backup.exists() { - return Err(system_conflict("No binary backup found for rollback")); - } - system_update_swap_file(&executable_path, &binary_backup) - .map_err(system_update_internal_with("Binary rollback failed"))?; - - let web_backup = system_update_backup_path_for(&config.web_dist_dir); - if web_backup.exists() - && let Err(error) = system_update_swap_dir(&config.web_dist_dir, &web_backup) - { - let binary_restore = system_update_swap_file(&executable_path, &binary_backup); - let mut rollback_errors = Vec::new(); - system_update_collect_rollback_error( - &mut rollback_errors, - "restore binary after web rollback failure", - binary_restore, - ); - return Err(system_update_error_with_rollback( - "Web rollback failed", - error, - rollback_errors, - )); - } - Ok(()) -} - -fn system_update_swap_file(current: &Path, backup: &Path) -> io::Result<()> { - if !current.exists() { - return Err(io::Error::new( - io::ErrorKind::NotFound, - format!("current binary is missing: {}", current.display()), - )); - } - let swap = system_update_swap_path_for(current); - if swap.exists() { - fs::remove_file(&swap)?; - } - system_update_move_file(current, &swap)?; - if let Err(error) = system_update_move_file(backup, current) { - let _ = system_update_move_file(&swap, current); - return Err(error); - } - if let Err(error) = system_update_move_file(&swap, backup) { - let _ = system_update_move_file(current, &swap); - let _ = system_update_move_file(backup, current); - let _ = system_update_move_file(&swap, backup); - return Err(error); - } - Ok(()) -} - -fn system_update_swap_dir(current: &Path, backup: &Path) -> io::Result<()> { - if !current.exists() { - return Err(io::Error::new( - io::ErrorKind::NotFound, - format!("current web directory is missing: {}", current.display()), - )); - } - let swap = system_update_swap_path_for(current); - if swap.exists() { - fs::remove_dir_all(&swap)?; - } - system_update_move_dir(current, &swap)?; - if let Err(error) = system_update_move_dir(backup, current) { - let _ = system_update_move_dir(&swap, current); - return Err(error); - } - if let Err(error) = system_update_move_dir(&swap, backup) { - let _ = system_update_move_dir(current, &swap); - let _ = system_update_move_dir(backup, current); - let _ = system_update_move_dir(&swap, backup); - return Err(error); - } - Ok(()) -} - -fn system_update_collect_rollback_error( - errors: &mut Vec, - action: &'static str, - result: io::Result<()>, -) { - if let Err(error) = result { - errors.push(format!("{action}: {error}")); - } -} - -fn system_update_error_with_rollback( - context: &'static str, - error: impl fmt::Display, - rollback_errors: Vec, -) -> SystemAdminError { - if rollback_errors.is_empty() { - return system_internal(format!("{context}: {error}")); - } - system_internal(format!( - "{context}: {error}; rollback failed: {}", - rollback_errors.join("; ") - )) -} - -fn system_update_restore_dir(current: &Path, backup: &Path) -> io::Result<()> { - if current.exists() { - fs::remove_dir_all(current)?; - } - if backup.exists() { - system_update_move_dir(backup, current)?; - } - Ok(()) -} - -fn system_update_move_file(from: &Path, to: &Path) -> io::Result<()> { - match fs::rename(from, to) { - Ok(()) => Ok(()), - Err(error) if error.kind() == io::ErrorKind::CrossesDevices => { - if let Err(copy_error) = fs::copy(from, to) { - let _ = fs::remove_file(to); - return Err(copy_error); - } - fs::remove_file(from) - } - Err(error) => Err(error), - } -} - -fn system_update_move_dir(from: &Path, to: &Path) -> io::Result<()> { - match fs::rename(from, to) { - Ok(()) => Ok(()), - Err(error) if error.kind() == io::ErrorKind::CrossesDevices => { - if let Err(copy_error) = system_update_copy_dir_all(from, to) { - let _ = fs::remove_dir_all(to); - return Err(copy_error); - } - fs::remove_dir_all(from) - } - Err(error) => Err(error), - } -} - -fn system_update_copy_dir_all(from: &Path, to: &Path) -> io::Result<()> { - fs::create_dir_all(to)?; - for entry in fs::read_dir(from)? { - let entry = entry?; - let file_type = entry.file_type()?; - let target = to.join(entry.file_name()); - if file_type.is_dir() { - system_update_copy_dir_all(&entry.path(), &target)?; - } else if file_type.is_file() { - fs::copy(entry.path(), &target)?; - fs::set_permissions(&target, entry.metadata()?.permissions())?; - } else { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - format!("unsupported file type in {}", entry.path().display()), - )); - } - } - Ok(()) -} - -fn system_update_backup_path_for(path: &Path) -> PathBuf { - let mut backup = path.as_os_str().to_os_string(); - backup.push(".backup"); - PathBuf::from(backup) -} - -fn system_update_swap_path_for(path: &Path) -> PathBuf { - let mut swap = path.as_os_str().to_os_string(); - swap.push(".rollback-swap"); - PathBuf::from(swap) -} - -fn system_update_web_dist_relative_path(path: &Path) -> Option { - let components = path - .components() - .filter_map(|component| match component { - std::path::Component::Normal(value) => Some(value.to_os_string()), - _ => None, - }) - .collect::>(); - for index in 0..components.len() { - if components[index] == "web" - && components - .get(index + 1) - .is_some_and(|value| value == "dist") - { - return Some(components[index + 2..].iter().collect()); - } - if components[index] == "dist" { - return Some(components[index + 1..].iter().collect()); - } - } - None -} - -fn system_update_unsafe_archive_path(path: &Path) -> bool { - path.is_absolute() - || path.components().any(|component| { - matches!( - component, - std::path::Component::ParentDir - | std::path::Component::Prefix(_) - | std::path::Component::RootDir - ) - }) -} - -struct SystemUpdateTempDir { - path: PathBuf, -} - -impl SystemUpdateTempDir { - const fn new(path: PathBuf) -> Self { - Self { path } - } -} - -impl Drop for SystemUpdateTempDir { - fn drop(&mut self) { - if let Err(error) = fs::remove_dir_all(&self.path) - && error.kind() != io::ErrorKind::NotFound - { - tracing::warn!( - path = %self.path.display(), - error = %error, - "清理系统更新临时目录失败" - ); - } - } -} - -fn system_update_unique_temp_dir(parent: &Path) -> Result { - for attempt in 0..100_u8 { - let path = parent.join(format!( - ".codex-proxy-rs-update-{}-{attempt}", - Utc::now().timestamp_nanos_opt().unwrap_or_default() - )); - match fs::create_dir(&path) { - Ok(()) => return Ok(path), - Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue, - Err(error) => { - return Err(system_internal(format!( - "Failed to create update temp dir: {error}" - ))); - } - } - } - Err(system_internal("Failed to create unique update temp dir")) -} - -#[derive(Debug)] -struct SystemUpdateFileLock { - path: PathBuf, -} - -impl SystemUpdateFileLock { - fn acquire(path: &Path) -> Result { - if let Some(parent) = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - { - fs::create_dir_all(parent).map_err(system_update_internal_with( - "Failed to create update lock directory", - ))?; - } - match Self::try_create(path) { - Ok(()) => Ok(Self { - path: path.to_path_buf(), - }), - Err(error) if error.kind() == io::ErrorKind::AlreadyExists => { - if !system_update_stale_lock(path)? { - return Err(system_conflict("System update already running")); - } - fs::remove_file(path).map_err(system_update_internal_with( - "Failed to remove stale update lock", - ))?; - match Self::try_create(path) { - Ok(()) => Ok(Self { - path: path.to_path_buf(), - }), - Err(error) if error.kind() == io::ErrorKind::AlreadyExists => { - Err(system_conflict("System update already running")) - } - Err(error) => Err(system_internal(format!( - "Failed to create update lock: {error}" - ))), - } - } - Err(error) => Err(system_internal(format!( - "Failed to create update lock: {error}" - ))), - } - } - - fn try_create(path: &Path) -> io::Result<()> { - let mut file = fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(path)?; - writeln!( - file, - "pid={}\ncreated_at={}", - std::process::id(), - Utc::now().to_rfc3339() - )?; - file.sync_all() - } -} - -impl Drop for SystemUpdateFileLock { - fn drop(&mut self) { - if let Err(error) = fs::remove_file(&self.path) - && error.kind() != io::ErrorKind::NotFound - { - tracing::warn!( - path = %self.path.display(), - error = %error, - "清理系统更新文件锁失败" - ); - } - } -} - -fn system_update_stale_lock(path: &Path) -> Result { - let metadata = - fs::metadata(path).map_err(system_update_internal_with("Failed to read update lock"))?; - let modified = metadata.modified().map_err(system_update_internal_with( - "Failed to read update lock timestamp", - ))?; - modified - .elapsed() - .map(|age| age > Duration::from_secs(30 * 60)) - .map_err(system_update_internal_with( - "Failed to calculate update lock age", - )) -} - -fn system_update_read_state(path: &Path) -> Result { - if !path.exists() { - return Ok(SystemUpdateStatusData::default()); - } - let data = fs::read_to_string(path) - .map_err(system_update_internal_with("Failed to read update state"))?; - serde_json::from_str(&data).map_err(system_update_internal_with("Invalid update state")) -} - -fn system_update_write_state( - path: &Path, - state: &SystemUpdateStatusData, -) -> Result<(), SystemAdminError> { - if let Some(parent) = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - { - fs::create_dir_all(parent).map_err(system_update_internal_with( - "Failed to create update state directory", - ))?; - } - let data = serde_json::to_vec_pretty(state) - .map_err(system_update_internal_with("Failed to encode update state"))?; - let temporary_path = system_update_state_temporary_path(path); - let result = (|| { - let mut file = fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&temporary_path)?; - file.write_all(&data)?; - file.sync_all()?; - fs::rename(&temporary_path, path) - })(); - if result.is_err() { - let _ = fs::remove_file(&temporary_path); - } - result.map_err(system_update_internal_with("Failed to write update state")) -} - -fn system_update_state_temporary_path(path: &Path) -> PathBuf { - let mut temporary = path.as_os_str().to_os_string(); - temporary.push(format!( - ".tmp-{}-{}", - std::process::id(), - Utc::now().timestamp_nanos_opt().unwrap_or_default() - )); - PathBuf::from(temporary) -} - -fn system_update_set_operation_running( - path: &Path, - operation_id: &str, - kind: SystemOperationKind, - version: Option<&str>, - current_version: &str, -) -> Result<(), SystemAdminError> { - let mut state = system_update_read_state(path)?; - if state.current_version.is_none() { - state.current_version = Some(current_version.to_owned()); - } - state.operation = SystemOperationState { - operation_id: Some(operation_id.to_owned()), - kind: Some(kind), - status: SystemOperationStatus::Running, - target_version: version.map(ToOwned::to_owned), - message: Some("operation running".to_owned()), - error: None, - started_at: Some(Utc::now().to_rfc3339()), - finished_at: None, - }; - system_update_write_state(path, &state) -} - -fn system_update_finish_operation( - path: &Path, - operation_id: &str, - kind: SystemOperationKind, - version: Option, - error: Option, -) { - let mut state = match system_update_read_state(path) { - Ok(state) => state, - Err(error) => { - tracing::warn!(error = %error, "读取系统更新状态失败"); - return; - } - }; - if state.operation.operation_id.as_deref() != Some(operation_id) { - return; - } - if let Some(error) = error { - state.operation.status = SystemOperationStatus::Failed; - state.operation.message = Some("operation failed".to_owned()); - state.operation.error = Some(error); - } else { - state.operation.status = SystemOperationStatus::Succeeded; - state.operation.message = Some("operation succeeded".to_owned()); - state.operation.error = None; - match kind { - SystemOperationKind::Update => { - state.previous_version = state.current_version.take(); - state.current_version = version.clone(); - } - SystemOperationKind::Rollback => { - let current = state.current_version.take(); - state.current_version = state.previous_version.take(); - state.previous_version = current; - } - } - state.operation.target_version = version; - } - state.operation.finished_at = Some(Utc::now().to_rfc3339()); - if let Err(error) = system_update_write_state(path, &state) { - tracing::warn!(error = %error, "写入系统更新状态失败"); - } -} - -fn system_update_operation_id(kind: &str) -> String { - format!("sysop-{kind}-{}", Utc::now().timestamp_millis()) -} - -fn system_update_spawn_replacement_process( - config: &SystemUpdateConfig, -) -> Result<(), SystemAdminError> { - let executable_path = config.executable_path()?; - let metadata = fs::metadata(&executable_path).map_err(system_update_internal_with( - "Failed to schedule replacement process", - ))?; - if !metadata.is_file() { - return Err(system_internal( - "Failed to schedule replacement process: executable path is not a file", - )); - } - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt as _; - - if metadata.permissions().mode() & 0o111 == 0 { - return Err(system_internal( - "Failed to schedule replacement process: executable is not executable", - )); - } - } - - let delay_ms = system_update_environment_value(SYSTEM_UPDATE_RESTART_DELAY_ENV) - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or_else(|| { - SYSTEM_UPDATE_REPLACEMENT_START_DELAY_MS - .parse::() - .unwrap_or(1_200) - }); - #[cfg(unix)] - let mut command = { - let delay = format!("{}.{:03}", delay_ms / 1_000, delay_ms % 1_000); - let mut command = std::process::Command::new("/bin/sh"); - command - .arg("-c") - .arg("sleep \"$1\"; shift; exec \"$@\"") - .arg("codex-proxy-rs-restart") - .arg(delay) - .arg(&executable_path) - .args(env::args_os().skip(1)); - command - }; - #[cfg(not(unix))] - let mut command = { - let mut command = std::process::Command::new(&executable_path); - command - .args(env::args_os().skip(1)) - .env(SYSTEM_UPDATE_RESTART_DELAY_ENV, delay_ms.to_string()); - command - }; - command.stdin(std::process::Stdio::null()); - command - .spawn() - .map(|_| ()) - .map_err(system_update_internal_with( - "Failed to schedule replacement process", - )) -} - -fn system_update_default_temp_dir(state_file: &Path) -> PathBuf { - state_file - .parent() - .map(|parent| parent.join("update-tmp")) - .unwrap_or_else(|| env::temp_dir().join("codex-proxy-rs-update")) -} - -fn system_update_environment_value(key: &str) -> Option { - env::var(key) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -/// OpenAI 与 Admin router 共用的应用状态。 -pub struct ApplicationState { - admin: Arc, - openai: O, - health: Arc, -} - -impl ApplicationState { - #[must_use] - pub const fn new(admin: Arc, openai: O, health: Arc) -> Self { - Self { - admin, - openai, - health, - } - } -} - -impl Clone for ApplicationState -where - O: Clone, -{ - fn clone(&self) -> Self { - Self { - admin: Arc::clone(&self.admin), - openai: self.openai.clone(), - health: Arc::clone(&self.health), - } - } -} - -impl AdminSessionState for ApplicationState { - fn admin_session_resolver(&self) -> &dyn AdminSessionResolver { - self.admin.sessions.as_ref() - } -} - -impl AdminAuthState for ApplicationState { - fn admin_auth_service(&self) -> &dyn AdminAuthService { - self.admin.sessions.as_ref() - } -} - -impl AccountAdminState for ApplicationState { - fn account_admin_service(&self) -> &dyn AccountAdminService { - self.admin.accounts.as_ref() - } -} - -impl CatalogAdminState for ApplicationState { - fn catalog_admin_service(&self) -> &dyn CatalogAdminService { - self.admin.catalog.as_ref() - } -} - -impl ClientKeyAdminState for ApplicationState { - fn client_key_admin_service(&self) -> &dyn ClientKeyAdminService { - self.admin.client_keys.as_ref() - } -} - -impl CodexAdminState for ApplicationState { - fn codex_admin_service(&self) -> &dyn CodexAdminService { - self.admin.codex.as_ref() - } -} - -impl ObservabilityAdminState for ApplicationState { - fn observability_admin_service(&self) -> &dyn ObservabilityAdminService { - self.admin.observability.as_ref() - } -} - -impl AdminSettingsState for ApplicationState { - fn admin_settings_service(&self) -> &dyn AdminSettingsService { - self.admin.settings.as_ref() - } -} - -impl SystemAdminState for ApplicationState { - fn system_admin_service(&self) -> &dyn SystemAdminService { - self.admin.system.as_ref() - } -} - -impl XaiAdminState for ApplicationState { - fn xai_admin_service(&self) -> &dyn XaiAdminService { - self.admin.xai.as_ref() - } -} - -impl OpenAiApiState for ApplicationState -where - O: OpenAiClientService, -{ - type Service = O; - - fn openai_client_api(&self) -> Self::Service { - self.openai.clone() - } -} - -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::num::NonZeroU32; -use std::sync::RwLock; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; - -use gateway_api::openai::{ConnectionTask, auth::ClientApiKeyAuthError}; -use gateway_core::{ - engine::{ - credential::{AccountSelectionPolicy, RotationStrategy}, - provider::ProviderRegistry, - }, - policy::ClientPolicy, - routing::{ - ConfigRevision, InstanceHealth, ProviderInstance, ProviderInstanceId, ProviderModel, - PublicModelId, RuntimeSnapshot, UpstreamModelId, - }, -}; -use gateway_store::postgres::{RuntimeSnapshotData, RuntimeSnapshotRepository}; -use tokio_util::task::TaskTracker; - -/// 一个配置 revision 的完整数据面快照。 -#[derive(Debug)] -pub struct CompiledRuntimeSnapshot { - routing: RuntimeSnapshot, -} - -impl CompiledRuntimeSnapshot { - #[must_use] - pub const fn new(routing: RuntimeSnapshot) -> Self { - Self { routing } - } - - #[must_use] - pub const fn routing(&self) -> &RuntimeSnapshot { - &self.routing - } -} - -/// PostgreSQL 一致性事实与 Provider 实时目录的唯一快照编译器。 -pub struct RuntimeSnapshotCompiler { - repository: Arc, - providers: ProviderRegistry, -} - -impl RuntimeSnapshotCompiler { - #[must_use] - pub const fn new( - repository: Arc, - providers: ProviderRegistry, - ) -> Self { - Self { - repository, - providers, - } - } - - /// 读取一个 PostgreSQL revision，并为启用的 instance 查询实时模型目录。 - pub async fn compile(&self) -> Result { - let data = self - .repository - .load_runtime_snapshot() - .await - .map_err(|_| RuntimeSnapshotCompileError::StoreUnavailable)?; - if data.config_revision != data.observed_current_revision { - return Err(RuntimeSnapshotCompileError::RevisionChanged); - } - compile_runtime_snapshot(data, &self.providers).await - } -} - -/// 快照未发布时可安全记录的稳定错误。 -#[derive(Debug, thiserror::Error)] -pub enum RuntimeSnapshotCompileError { - #[error("runtime snapshot store is unavailable")] - StoreUnavailable, - #[error("runtime configuration changed while the snapshot was loading")] - RevisionChanged, - #[error("runtime snapshot contains invalid frozen data")] - InvalidData, -} - -async fn compile_runtime_snapshot( - data: RuntimeSnapshotData, - providers: &ProviderRegistry, -) -> Result { - let mut instances_by_id = BTreeMap::new(); - for record in &data.provider_instances { - let instance = ProviderInstance::new( - ProviderInstanceId::new(record.id.clone()) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?, - gateway_core::routing::ProviderKind::new(record.provider_kind.clone()) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?, - record.base_url.clone(), - record.enabled, - InstanceHealth::Healthy, - ); - instances_by_id.insert(record.id.clone(), instance.clone()); - } - - // Provider catalog 只用于公开模型列表和能力提示；未知模型仍允许透传， - // 最终是否支持由对应上游返回。 - let mut provider_models = Vec::new(); - for record in &data.provider_instances { - let instance = instances_by_id - .get(record.id.as_str()) - .ok_or(RuntimeSnapshotCompileError::InvalidData)?; - let models = match providers.query_model_capabilities(instance).await { - Ok(models) => models, - Err(error) => { - tracing::warn!( - provider_instance_id = %record.id, - provider_kind = %record.provider_kind, - error = ?error, - "Provider model catalog unavailable; requests remain pass-through" - ); - continue; - } - }; - for model in models { - provider_models.push(ProviderModel::new( - instance.id().clone(), - model.upstream_model().clone(), - model.capabilities().clone(), - )); - } - } - - let provider_model_mappings = data - .settings - .provider_model_mappings - .into_iter() - .map(|(provider, mapping)| { - Ok(( - gateway_core::routing::ProviderKind::new(provider) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?, - mapping, - )) - }) - .collect::, RuntimeSnapshotCompileError>>()?; - - let rotation_strategy = match data.settings.rotation_strategy.as_str() { - "smart" => RotationStrategy::Smart, - "quota_reset_priority" => RotationStrategy::QuotaResetPriority, - "round_robin" => RotationStrategy::RoundRobin, - "sticky" => RotationStrategy::Sticky, - _ => return Err(RuntimeSnapshotCompileError::InvalidData), - }; - let selection_policy = AccountSelectionPolicy::new( - rotation_strategy, - NonZeroU32::new(data.settings.max_concurrent_per_account) - .ok_or(RuntimeSnapshotCompileError::InvalidData)?, - Duration::from_millis(data.settings.request_interval_ms), - ); - let client_policies = data - .client_api_keys - .into_iter() - .map(|key| { - let provider_kind = gateway_core::routing::ProviderKind::new(key.provider_kind) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?; - Ok(ClientPolicy::new( - key.id, - key.plaintext_key, - provider_kind, - true, - key.limits, - )) - }) - .collect::, RuntimeSnapshotCompileError>>()?; - let instances = instances_by_id.into_values().collect(); - let routing = RuntimeSnapshot::new( - ConfigRevision::new(data.config_revision.get()) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?, - selection_policy, - instances, - provider_models, - client_policies, - ) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)? - .with_provider_model_mappings(provider_model_mappings); - Ok(CompiledRuntimeSnapshot::new(routing)) -} - -/// RuntimeSnapshot 发布和请求级冻结句柄。 -#[derive(Clone, Default)] -pub struct RuntimeSnapshotHandle { - current: Arc>>>, -} - -impl RuntimeSnapshotHandle { - #[must_use] - pub fn new(initial: CompiledRuntimeSnapshot) -> Self { - Self { - current: Arc::new(RwLock::new(Some(Arc::new(initial)))), - } - } - - pub fn publish(&self, snapshot: CompiledRuntimeSnapshot) { - *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); - } - - pub fn suspend(&self) { - *write_unpoisoned(&self.current) = None; - } - - #[must_use] - pub fn revision(&self) -> Option { - read_unpoisoned(&self.current) - .as_ref() - .map(|snapshot| snapshot.routing().revision().get()) - } - - fn acquire(&self) -> Result, ClientApiKeyAuthError> { - read_unpoisoned(&self.current) - .clone() - .ok_or(ClientApiKeyAuthError::RuntimeUnavailable) - } -} - -/// 配置提交后的本进程快照发布与跨进程失效通知。 -#[derive(Clone)] -pub struct RuntimeSnapshotPublisher { - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - runtime_changes: Arc, -} - -impl RuntimeSnapshotPublisher { - #[must_use] - pub const fn new( - compiler: Arc, - snapshots: RuntimeSnapshotHandle, - runtime_changes: Arc, - ) -> Self { - Self { - compiler, - snapshots, - runtime_changes, - } - } - - /// 重新编译并原子替换本进程快照。 - pub async fn refresh(&self) -> Result { - let snapshot = self.compiler.compile().await?; - let revision = Revision::new(snapshot.routing().revision().get()) - .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?; - self.snapshots.publish(snapshot); - Ok(revision) - } - - fn suspend(&self) { - self.snapshots.suspend(); - } - - fn published_revision(&self) -> Option { - self.snapshots.revision() - } - - /// 数据库提交不能被外部目录或 Redis 的暂时故障伪装成回滚。 - pub async fn publish_committed(&self, committed_revision: Revision) { - match self.refresh().await { - Ok(published_revision) if published_revision == committed_revision => {} - Ok(published_revision) => { - tracing::warn!( - committed_revision = committed_revision.get(), - published_revision = published_revision.get(), - "本进程跳过了已经被更新版本取代的配置快照" - ); - } - Err(error) => { - self.suspend(); - tracing::error!( - committed_revision = committed_revision.get(), - error = %error, - "配置已经提交，但本进程快照暂未刷新" - ); - } - } - let change = gateway_store::redis::RuntimeChange::SnapshotPublished { - config_revision: committed_revision, - }; - if self - .runtime_changes - .publish_runtime_change(&change) - .await - .is_err() - { - tracing::warn!( - config_revision = committed_revision.get(), - "配置已经提交，但 Redis 失效通知发布失败" - ); - } - } -} - -/// 已认证 Client 与认证时冻结的完整配置 revision。 -#[derive(Clone)] -pub struct AuthenticatedClient { - snapshot: Arc, - policy: ClientPolicy, -} - -impl AuthenticatedClient { - #[must_use] - pub const fn snapshot(&self) -> &Arc { - &self.snapshot - } - - #[must_use] - pub const fn policy(&self) -> &ClientPolicy { - &self.policy - } -} - -impl fmt::Debug for AuthenticatedClient { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("AuthenticatedClient") - .field("key_id", &self.policy.key_id()) - .field("revision", &self.snapshot.routing.revision()) - .finish_non_exhaustive() - } -} - -fn authenticate_client_key( - snapshots: &RuntimeSnapshotHandle, - plaintext: &str, -) -> Result { - if !valid_client_api_key_shape(plaintext) { - return Err(ClientApiKeyAuthError::InvalidKey); - } - let snapshot = snapshots.acquire()?; - let mut matched = None; - for policy in snapshot.routing.client_policies() { - let expected = policy.plaintext_key().expose_for_auth(); - let equal = plaintext.len() == expected.len() - && bool::from(plaintext.as_bytes().ct_eq(expected.as_bytes())); - if equal && policy.authorize().is_ok() { - matched = Some(policy.clone()); - } - } - matched - .map(|policy| AuthenticatedClient { snapshot, policy }) - .ok_or(ClientApiKeyAuthError::InvalidKey) -} - -fn valid_client_api_key_shape(value: &str) -> bool { - let Some(encoded) = value.strip_prefix("sk_") else { - return false; - }; - URL_SAFE_NO_PAD - .decode(encoded) - .is_ok_and(|decoded| decoded.len() == 32) -} - -/// HTTP/WS 长连接的统一接入、编号和 drain 生命周期。 -#[derive(Clone)] -pub struct ConnectionLifecycle { - shutting_down: Arc, - sequence: Arc, - tasks: TaskTracker, -} - -impl Default for ConnectionLifecycle { - fn default() -> Self { - Self::new() - } -} - -impl ConnectionLifecycle { - #[must_use] - pub fn new() -> Self { - Self { - shutting_down: Arc::new(AtomicBool::new(false)), - sequence: Arc::new(AtomicU64::new(0)), - tasks: TaskTracker::new(), - } - } - - #[must_use] - pub fn is_shutting_down(&self) -> bool { - self.shutting_down.load(Ordering::Acquire) - } - - pub fn spawn(&self, task: ConnectionTask) { - if self.is_shutting_down() { - return; - } - self.tasks.spawn(task); - } - - #[must_use] - pub fn next_id(&self, prefix: &str) -> String { - let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); - format!("{prefix}_{}_{sequence}", Uuid::now_v7().simple()) - } - - pub async fn shutdown(&self) { - self.shutting_down.store(true, Ordering::Release); - self.tasks.close(); - self.tasks.wait().await; - } -} - -fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { - lock.read() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { - lock.write() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - -use axum::{Router, extract::State, http::StatusCode, routing::get}; -use tower_http::{ - request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer}, - services::{ServeDir, ServeFile}, - trace::TraceLayer, -}; - -use crate::workers::WorkerHealthRegistry; - -/// `/healthz` 同时验证快照、worker 新鲜度和 PostgreSQL/Redis 真实连通性。 -pub struct HealthStatus { - snapshots: RuntimeSnapshotHandle, - workers: WorkerHealthRegistry, - postgres: sqlx::PgPool, - redis: redis::aio::ConnectionManager, -} - -impl HealthStatus { - #[must_use] - pub fn new( - snapshots: RuntimeSnapshotHandle, - workers: WorkerHealthRegistry, - postgres: sqlx::PgPool, - redis: redis::aio::ConnectionManager, - ) -> Self { - Self { - snapshots, - workers, - postgres, - redis, - } - } - - #[must_use] - pub async fn healthy(&self) -> bool { - if self.snapshots.acquire().is_err() || !self.workers.all_healthy() { - return false; - } - let postgres = async { - let mut connection = self.postgres.acquire().await?; - sqlx::Connection::ping(&mut *connection).await - }; - let mut redis = self.redis.clone(); - let infrastructure = tokio::time::timeout(Duration::from_secs(2), async move { - let ping = redis::cmd("PING"); - let redis_ping = ping.query_async::(&mut redis); - tokio::join!(postgres, redis_ping) - }) - .await; - matches!( - infrastructure, - Ok((Ok(()), Ok(response))) if response == "PONG" - ) - } -} - -/// 组合固定 OpenAI/Admin 路由、健康检查与旧前端静态资源。 -pub fn application_router( - state: ApplicationState, - asset_directory: impl Into, -) -> Router -where - O: OpenAiClientService, -{ - let asset_directory = asset_directory.into(); - let index = asset_directory.join("index.html"); - let request_id_header = axum::http::HeaderName::from_static("x-request-id"); - Router::new() - .route("/healthz", get(healthz::)) - .merge(gateway_api::openai::router::>()) - .merge(gateway_api::admin::router::>()) - .fallback_service(ServeDir::new(asset_directory).fallback(ServeFile::new(index))) - .layer(PropagateRequestIdLayer::new(request_id_header.clone())) - .layer(SetRequestIdLayer::new(request_id_header, MakeRequestUuid)) - .layer(TraceLayer::new_for_http()) - .with_state(state) -} - -async fn healthz(State(state): State>) -> StatusCode { - if state.health.healthy().await { - StatusCode::NO_CONTENT - } else { - StatusCode::SERVICE_UNAVAILABLE - } -} - -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use gateway_api::openai::{ - DeliveryEvent, ResponseExecutionSession as ApiResponseExecutionSession, ResponsesTransport, - StartedResponse, - error::gateway_error_from_engine, - responses::{ContinuationIntent, DecodedResponsesRequest}, -}; -use gateway_core::{ - engine::{ - AttemptCoordinator, CancellationToken as CoreCancellationToken, EngineError, GatewayEngine, - ModelRequestId, NewModelRequest, ProviderAttemptOutcome, - continuation::NativeContinuationPin, - coordinator::ResponseExecutionSession as CoreResponseExecutionSession, - }, - error::{GatewayError, GatewayErrorKind, ProviderErrorKind}, - policy::ClientApiKeyId, - routing::RoutingContext, -}; -use gateway_store::{ - postgres::PgExecutionStore, - redis::{ - ClientAdmissionDecision, ClientAdmissionLimits, ClientAdmissionRejection, - ClientAdmissionRepository, ClientAdmissionRequest, ProviderCircuitDecision, - ProviderCircuitRepository, - }, -}; - -const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60); - -/// 已认证 previous-response 到 Core pin 的唯一应用端口。 -#[async_trait] -pub trait NativeContinuationResolver: Send + Sync { - async fn resolve_native_continuation( - &self, - client_api_key_id: &ClientApiKeyId, - client_response_id: &str, - ) -> Result; -} - -/// PostgreSQL 历史事实到 Codex native previous-response pin 的组合 adapter。 -pub struct PgNativeContinuationResolver { - history: Arc, -} - -impl PgNativeContinuationResolver { - #[must_use] - pub const fn new( - history: Arc, - ) -> Self { - Self { history } - } -} - -#[async_trait] -impl NativeContinuationResolver for PgNativeContinuationResolver { - async fn resolve_native_continuation( - &self, - client_api_key_id: &ClientApiKeyId, - client_response_id: &str, - ) -> Result { - let history = self - .history - .find_model_request_by_client_response_id( - client_response_id, - client_api_key_id.as_str(), - ) - .await - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "previous response history is temporarily unavailable", - ) - })? - .ok_or_else(|| { - GatewayError::new( - GatewayErrorKind::InvalidRequest, - "previous response was not found", - ) - })?; - if history.provider_kind.as_deref() != Some("openai") { - return Err(GatewayError::new( - GatewayErrorKind::Unsupported, - "the previous response does not support native continuation", - )); - } - self.history - .resolve_native_continuation_pin( - client_response_id, - client_api_key_id.as_str(), - gateway_core::engine::continuation::NativeContinuationReuse::Reusable, - ) - .await - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "previous response history is temporarily unavailable", - ) - })? - .ok_or_else(|| { - GatewayError::new( - GatewayErrorKind::InvalidRequest, - "previous response is not eligible for native continuation", - ) - }) - } -} - -/// OpenAI API 与 Core streaming session 之间的薄组合 adapter。 -#[derive(Clone)] -pub struct GatewayOpenAiService { - snapshots: RuntimeSnapshotHandle, - coordinator: Arc>, - admissions: Arc, - circuits: Arc, - history: Arc, - connections: ConnectionLifecycle, -} - -impl GatewayOpenAiService { - #[must_use] - pub fn new( - snapshots: RuntimeSnapshotHandle, - execution_store: Arc, - providers: ProviderRegistry, - admissions: Arc, - circuits: Arc, - history: Arc, - connections: ConnectionLifecycle, - ) -> Self { - let engine = GatewayEngine::new(execution_store, providers); - Self { - snapshots, - coordinator: Arc::new(AttemptCoordinator::new(engine)), - admissions, - circuits, - history, - connections, - } - } - - /// 管理端账号探测仍走唯一 Router -> Engine -> Provider -> persistence 链， - /// 仅增加 required-account fence，并禁用账号重试与 Provider instance fallback。 - async fn test_account( - &self, - account_id: gateway_core::engine::credential::ProviderAccountId, - provider_instance_id: ProviderInstanceId, - upstream_model: UpstreamModelId, - ) -> Result, GatewayError> { - let snapshot = self.snapshots.acquire().map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "runtime snapshot is unavailable", - ) - })?; - let message = gateway_core::operation::Message::new( - gateway_core::operation::MessageRole::User, - vec![gateway_core::operation::ContentPart::Text( - "Reply with exactly OK.".to_owned(), - )], - ) - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "connection test request is invalid", - ) - })?; - let operation = gateway_core::operation::Operation::Generate( - gateway_core::operation::GenerateRequest::new(vec![message]).map_err(|_| { - GatewayError::new( - GatewayErrorKind::Internal, - "connection test request is invalid", - ) - })?, - ); - let model_request_id = ModelRequestId::new(self.connections.next_id("req")) - .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "invalid request ID"))?; - let started_at = SystemTime::now(); - let deadline_at = started_at - .checked_add(MODEL_REQUEST_DEADLINE) - .ok_or_else(|| GatewayError::new(GatewayErrorKind::Internal, "invalid system clock"))?; - let provider_kind = snapshot - .routing() - .provider_for_instance(&provider_instance_id) - .cloned() - .ok_or_else(|| { - GatewayError::new( - GatewayErrorKind::Unsupported, - "Provider instance is unavailable", - ) - })?; - let routing_context = RoutingContext { - provider_kind: Some(provider_kind), - allowed_instances: Some(BTreeSet::from([provider_instance_id])), - ..RoutingContext::default() - }; - let public_model = - PublicModelId::new(upstream_model.as_str().to_owned()).map_err(|_| { - GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid") - })?; - let plan = snapshot - .routing() - .plan(&public_model, &operation, &routing_context) - .map_err(map_routing_error)?; - let client_ref = ClientApiKeyId::new("admin_connection_test") - .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "invalid admin actor"))?; - let request = NewModelRequest { - id: model_request_id.clone(), - client_api_key_id: None, - client_api_key_ref: client_ref, - config_revision: plan.config_revision(), - protocol: "admin_connection_test".to_owned(), - operation: operation.kind(), - endpoint: "/api/admin/accounts/test".to_owned(), - client_transport: "internal".to_owned(), - requested_model: public_model, - input_token_estimate: operation.capability_requirements().minimum_context_tokens(), - client_ip: None, - user_agent: None, - reasoning_effort: None, - reasoning_preset: None, - request_kind: Some("account_connection_test".to_owned()), - subagent_kind: None, - compact: false, - started_at, - deadline_at, - }; - let cancellation = CoreCancellationToken::new(); - let mut session = self - .coordinator - .start( - request, - operation, - plan, - Some(account_id), - None, - cancellation, - ) - .await - .map_err(|error| gateway_error_from_engine(&error))?; - let collected = session.collect_uncommitted().await; - publish_provider_attempt_outcomes( - self.circuits.as_ref(), - session.provider_attempt_outcomes(), - ) - .await; - let events = collected.map_err(|error| gateway_error_from_engine(&error))?; - session - .commit_downstream(Some(200)) - .await - .map_err(|error| gateway_error_from_engine(&error))?; - Ok(events - .into_iter() - .filter_map(|event| match event { - gateway_core::event::GatewayEvent::TextDelta(delta) => Some(delta.text), - _ => None, - }) - .collect()) - } - - async fn route_context( - &self, - snapshot: &RuntimeSnapshot, - provider_kind: &gateway_core::routing::ProviderKind, - model: &PublicModelId, - operation: &gateway_core::operation::Operation, - ) -> Result { - let preliminary = snapshot - .plan( - model, - operation, - &RoutingContext { - provider_kind: Some(provider_kind.clone()), - allowed_instances: Some(snapshot.instance_ids_for_provider(provider_kind)), - ..RoutingContext::default() - }, - ) - .map_err(map_routing_error)?; - let mut blocked_instances = std::collections::BTreeSet::new(); - let mut checked = std::collections::BTreeSet::::new(); - for candidate in preliminary.candidates() { - let instance = candidate.instance(); - if !checked.insert(instance.clone()) { - continue; - } - match self - .circuits - .provider_circuit_decision(instance.as_str()) - .await - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "provider health state is temporarily unavailable", - ) - })? { - ProviderCircuitDecision::Allow => {} - ProviderCircuitDecision::BlockedUntil(_) => { - blocked_instances.insert(instance.clone()); - } - } - } - Ok(RoutingContext { - provider_kind: Some(provider_kind.clone()), - allowed_instances: Some(snapshot.instance_ids_for_provider(provider_kind)), - blocked_instances, - ..RoutingContext::default() - }) - } -} - -/// 只有 Provider instance 可归因的瞬态故障才影响 circuit。 -#[must_use] -pub const fn provider_failure_affects_circuit(error_kind: ProviderErrorKind) -> bool { - matches!( - error_kind, - ProviderErrorKind::Timeout - | ProviderErrorKind::Transport - | ProviderErrorKind::Protocol - | ProviderErrorKind::Unavailable - ) -} - -async fn publish_provider_attempt_outcomes( - circuits: &dyn ProviderCircuitRepository, - outcomes: &[ProviderAttemptOutcome], -) { - for outcome in outcomes { - let instance_id = outcome.provider_instance_id().as_str(); - let result = match outcome.error_kind() { - None => circuits.observe_provider_success(instance_id).await, - Some(error_kind) if provider_failure_affects_circuit(error_kind) => circuits - .observe_provider_failure(instance_id) - .await - .map(|_| ()), - Some(_) => continue, - }; - if result.is_err() { - tracing::warn!( - provider_instance_id = instance_id, - "Provider circuit 结果写入失败" - ); - } - } -} - -struct ClientAdmissionLease { - repository: Arc, - client_api_key_ref: String, - model_request_id: String, -} - -impl ClientAdmissionLease { - async fn release(self) { - if self - .repository - .release_client_request(&self.client_api_key_ref, &self.model_request_id) - .await - .is_err() - { - tracing::warn!("客户端准入租约释放失败，将由 Redis TTL 收敛"); - } - } -} - -/// API handler 持有的 Core session 与 Redis admission lease。 -pub struct GatewayResponseExecution { - core: Option>, - admission: Option, - circuits: Arc, - observed_provider_outcomes: usize, -} - -impl GatewayResponseExecution { - #[must_use] - fn new( - core: CoreResponseExecutionSession, - admission: ClientAdmissionLease, - circuits: Arc, - ) -> Self { - Self { - core: Some(core), - admission: Some(admission), - circuits, - observed_provider_outcomes: 0, - } - } - - fn core_mut( - &mut self, - ) -> Result<&mut CoreResponseExecutionSession, EngineError> { - self.core.as_mut().ok_or(EngineError::InvalidDeliveryState) - } - - async fn settle_if_finalized(&mut self) { - if !self - .core - .as_ref() - .is_some_and(CoreResponseExecutionSession::is_finalized) - { - return; - } - if let Some(admission) = self.admission.take() { - admission.release().await; - } - } - - async fn observe_provider_outcomes(&mut self) { - let Some(core) = self.core.as_ref() else { - return; - }; - let outcomes = core.provider_attempt_outcomes(); - let new_outcomes = outcomes - .get(self.observed_provider_outcomes..) - .unwrap_or_default() - .to_vec(); - self.observed_provider_outcomes = outcomes.len(); - publish_provider_attempt_outcomes(self.circuits.as_ref(), &new_outcomes).await; - } - - fn detach_inner(&mut self) { - let Some(mut core) = self.core.take() else { - return; - }; - core.cancel(); - let admission = self.admission.take(); - let circuits = Arc::clone(&self.circuits); - let observed_provider_outcomes = self.observed_provider_outcomes; - let Ok(runtime) = tokio::runtime::Handle::try_current() else { - return; - }; - std::mem::drop(runtime.spawn(async move { - if !core.is_finalized() { - let _ = core.cancel_and_finalize().await; - } - let pending_outcomes = core - .provider_attempt_outcomes() - .get(observed_provider_outcomes..) - .unwrap_or_default(); - publish_provider_attempt_outcomes(circuits.as_ref(), pending_outcomes).await; - if let Some(admission) = admission { - admission.release().await; - } - })); - } -} - -impl Drop for GatewayResponseExecution { - fn drop(&mut self) { - self.detach_inner(); - } -} - -#[async_trait] -impl ApiResponseExecutionSession for GatewayResponseExecution { - async fn next_delivery_event(&mut self) -> Result, EngineError> { - let event = self.core_mut()?.next_event().await.map(|event| { - event.map(|event| { - let requirement = event.commit_requirement(); - DeliveryEvent::new(event.into_event(), requirement) - }) - }); - self.observe_provider_outcomes().await; - self.settle_if_finalized().await; - event - } - - async fn collect_uncommitted( - &mut self, - ) -> Result, EngineError> { - let events = self.core_mut()?.collect_uncommitted().await; - self.observe_provider_outcomes().await; - self.settle_if_finalized().await; - events - } - - async fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> Result<(), EngineError> { - let result = self.core_mut()?.commit_downstream(client_status_code).await; - self.observe_provider_outcomes().await; - self.settle_if_finalized().await; - result - } - - async fn record_client_status(&mut self, client_status_code: u16) -> Result<(), EngineError> { - let result = self - .core_mut()? - .record_client_status(client_status_code) - .await; - self.observe_provider_outcomes().await; - self.settle_if_finalized().await; - result - } - - fn is_finalized(&self) -> bool { - self.core - .as_ref() - .is_none_or(CoreResponseExecutionSession::is_finalized) - } - - fn cancel(&self) { - if let Some(core) = &self.core { - core.cancel(); - } - } - - fn detach_finalize(mut self) { - self.detach_inner(); - } -} - -#[async_trait] -impl OpenAiClientService for GatewayOpenAiService { - type Client = AuthenticatedClient; - type Session = GatewayResponseExecution; - - fn authenticate(&self, plaintext: &str) -> Result { - authenticate_client_key(&self.snapshots, plaintext) - } - - fn public_models(&self, client: &Self::Client) -> Vec { - client - .snapshot() - .routing() - .public_models_for_provider(client.policy().provider_kind()) - .into_iter() - .map(|model| model.as_str().to_owned()) - .collect() - } - - fn contains_public_model(&self, client: &Self::Client, model: &PublicModelId) -> bool { - client - .snapshot() - .routing() - .contains_public_model_for_provider(model, client.policy().provider_kind()) - } - - async fn start_response( - &self, - client: Self::Client, - request: DecodedResponsesRequest, - transport: ResponsesTransport, - ) -> Result, GatewayError> { - client.policy().authorize().map_err(|_| { - GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled") - })?; - let (operation, metadata) = request.into_parts(); - let (reasoning_effort, semantics) = match &operation { - gateway_core::operation::Operation::Generate(request) => ( - request - .reasoning() - .and_then(|reasoning| reasoning.effort) - .map(reasoning_effort_name) - .map(str::to_owned), - provider_openai::codex_request_semantics(request), - ), - _ => ( - None, - provider_openai::transport::protocol::responses::CodexRequestSemantics::default(), - ), - }; - let public_model = - PublicModelId::new(metadata.public_model().to_owned()).map_err(|_| { - GatewayError::new( - GatewayErrorKind::ModelNotFound, - "requested model was not found", - ) - })?; - let model_request_id = - ModelRequestId::new(self.connections.next_id("req")).map_err(|_| { - GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID") - })?; - let started_at = SystemTime::now(); - let deadline_at = started_at - .checked_add(MODEL_REQUEST_DEADLINE) - .ok_or_else(|| { - GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") - })?; - let routing_context = self - .route_context( - client.snapshot().routing(), - client.policy().provider_kind(), - &public_model, - &operation, - ) - .await?; - let plan = client - .snapshot() - .routing() - .plan(&public_model, &operation, &routing_context) - .map_err(map_routing_error)?; - let continuation = match metadata.continuation() { - ContinuationIntent::None => None, - ContinuationIntent::PreviousResponseId(response_id) => Some( - self.history - .resolve_native_continuation(client.policy().key_id(), response_id) - .await?, - ), - }; - let limits = client.policy().limits(); - let admission_request = ClientAdmissionRequest { - model_request_id: model_request_id.as_str().to_owned(), - client_api_key_ref: client.policy().key_id().as_str().to_owned(), - input_token_estimate: operation.capability_requirements().minimum_context_tokens(), - lease_ttl: MODEL_REQUEST_DEADLINE, - limits: ClientAdmissionLimits { - max_concurrency: limits.max_concurrency, - requests_per_minute: limits.requests_per_minute, - tokens_per_minute: limits.tokens_per_minute, - }, - }; - match self - .admissions - .admit_client_request(&admission_request) - .await - .map_err(|_| { - GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "request admission is temporarily unavailable", - ) - })? { - ClientAdmissionDecision::Granted => {} - ClientAdmissionDecision::Rejected( - ClientAdmissionRejection::RateLimited - | ClientAdmissionRejection::ConcurrencyLimited, - ) => { - return Err(GatewayError::new( - GatewayErrorKind::RateLimited, - "request exceeds client API key limits", - )); - } - } - let admission = ClientAdmissionLease { - repository: Arc::clone(&self.admissions), - client_api_key_ref: client.policy().key_id().as_str().to_owned(), - model_request_id: model_request_id.as_str().to_owned(), - }; - let client_transport = match (transport, metadata.stream()) { - (ResponsesTransport::WebSocket, _) => "websocket", - (ResponsesTransport::Http, true) => "http_sse", - (ResponsesTransport::Http, false) => "http_json", - }; - let new_request = NewModelRequest { - id: model_request_id.clone(), - client_api_key_id: Some(client.policy().key_id().clone()), - client_api_key_ref: client.policy().key_id().clone(), - config_revision: plan.config_revision(), - protocol: "openai_responses".to_owned(), - operation: operation.kind(), - endpoint: "/v1/responses".to_owned(), - client_transport: client_transport.to_owned(), - requested_model: public_model, - input_token_estimate: admission_request.input_token_estimate, - client_ip: metadata.client_ip(), - user_agent: metadata.user_agent().map(str::to_owned), - reasoning_effort, - reasoning_preset: semantics.reasoning_preset.map(str::to_owned), - request_kind: semantics.request_kind, - subagent_kind: semantics.subagent_kind, - compact: semantics.compact, - started_at, - deadline_at, - }; - let cancellation = CoreCancellationToken::new(); - let core = match self - .coordinator - .start( - new_request, - operation, - plan, - None, - continuation, - cancellation, - ) - .await - { - Ok(core) => core, - Err(error) => { - admission.release().await; - return Err(gateway_error_from_engine(&error)); - } - }; - let created_at = started_at - .duration_since(UNIX_EPOCH) - .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid"))? - .as_secs(); - Ok(StartedResponse::new( - model_request_id.as_str().to_owned(), - GatewayResponseExecution::new(core, admission, Arc::clone(&self.circuits)), - created_at, - metadata.stream(), - )) - } - - fn is_shutting_down(&self) -> bool { - self.connections.is_shutting_down() - } - - fn spawn_connection(&self, task: ConnectionTask) { - self.connections.spawn(task); - } - - fn next_connection_id(&self) -> String { - self.connections.next_id("ws") - } - - fn next_request_id(&self) -> String { - self.connections.next_id("req") - } -} - -const fn reasoning_effort_name(effort: gateway_core::operation::ReasoningEffort) -> &'static str { - match effort { - gateway_core::operation::ReasoningEffort::Minimal => "minimal", - gateway_core::operation::ReasoningEffort::Low => "low", - gateway_core::operation::ReasoningEffort::Medium => "medium", - gateway_core::operation::ReasoningEffort::High => "high", - gateway_core::operation::ReasoningEffort::ExtraHigh => "xhigh", - } -} - -fn map_routing_error(error: gateway_core::error::RoutingError) -> GatewayError { - match error { - gateway_core::error::RoutingError::NoCapableProvider { .. } => GatewayError::new( - GatewayErrorKind::NoAvailableProvider, - "no Provider instance can execute this request", - ), - _ => GatewayError::new( - GatewayErrorKind::Internal, - "runtime routing configuration is invalid", - ), - } -} - -/// 监听端口前完成的客户端准入恢复结果。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ClientAdmissionStartupRecoveryReport { - pub expired_model_requests: u64, - pub restored_clients: u64, - pub restored_recent_requests: u64, - pub restored_running_requests: u64, -} - -/// 先收敛 PostgreSQL 过期请求，再把一分钟窗口和仍在运行的租约恢复到 Redis。 -/// -/// 任一步失败都直接返回错误，调用方不得启动监听端口。 -pub async fn restore_client_admission_startup( - model_requests: &dyn gateway_store::postgres::ModelRequestRepository, - recovery_repository: &dyn gateway_store::postgres::ClientAdmissionRecoveryRepository, - admissions: &dyn gateway_store::redis::ClientAdmissionRepository, - now: DateTime, -) -> gateway_store::StoreResult { - let expired = model_requests.recover_expired_model_requests(now).await?; - let recoveries = recovery_repository - .load_client_admission_recovery(now - ChronoDuration::seconds(61)) - .await?; - let restored_clients = u64::try_from(recoveries.len()).unwrap_or(u64::MAX); - let mut restored_recent_requests = 0_u64; - let mut restored_running_requests = 0_u64; - for recovery in recoveries { - let restore = gateway_store::redis::ClientAdmissionRestore { - client_api_key_ref: recovery.client_api_key_ref, - recent_requests: recovery - .recent_requests - .into_iter() - .map( - |request| gateway_store::redis::ClientAdmissionRecentRequest { - model_request_id: request.model_request_id, - started_at: request.started_at, - input_token_estimate: request.input_token_estimate, - }, - ) - .collect(), - running_requests: recovery - .running_requests - .into_iter() - .map( - |request| gateway_store::redis::ClientAdmissionRunningRequest { - model_request_id: request.model_request_id, - expires_at: request.deadline_at, - }, - ) - .collect(), - }; - let result = admissions.restore_client_admission(&restore).await?; - restored_recent_requests = - restored_recent_requests.saturating_add(result.restored_recent_requests); - restored_running_requests = - restored_running_requests.saturating_add(result.restored_running_requests); - } - Ok(ClientAdmissionStartupRecoveryReport { - expired_model_requests: expired.requests, - restored_clients, - restored_recent_requests, - restored_running_requests, - }) -} - -/// 装配并运行唯一生产网关进程。 -pub async fn run(config: BootstrapConfig) -> Result<(), Box> { - const REDIS_NAMESPACE: &str = "codex-proxy-rs"; - - let BootstrapConfigParts { - app, - database_url, - redis_url, - admin_default_password, - } = config.into_parts(); - let _log_guard = initialize_logging(&app.logging)?; - let pool = gateway_store::postgres::connect_and_migrate(database_url.expose_secret()).await?; - let redis_client = redis::Client::open(redis_url.expose_secret().clone())?; - let redis_connection = redis_client.get_connection_manager().await?; - - let control_plane = Arc::new(gateway_store::postgres::PgControlPlaneRepository::new( - pool.clone(), - )); - let instances = Arc::new(gateway_store::postgres::PgConfigCatalogRepository::new( - pool.clone(), - )); - let runtime_settings = Arc::new(gateway_store::postgres::PgRuntimeSettingsRepository::new( - pool.clone(), - )); - let client_keys = Arc::new(gateway_store::postgres::PgClientApiKeyRepository::new( - pool.clone(), - )); - let security = - Arc::new(gateway_store::postgres::PgAdminSecurityAuditRepository::new(pool.clone())); - let provider_accounts = Arc::new(gateway_store::postgres::PgProviderAccountRepository::new( - pool.clone(), - )); - let snapshot_repository = Arc::new(gateway_store::postgres::PgRuntimeSnapshotRepository::new( - pool.clone(), - )); - let execution_store = Arc::new(gateway_store::postgres::PgExecutionStore::new(pool.clone())); - let history = Arc::new(gateway_store::postgres::PgHistoryRepository::new( - pool.clone(), - )); - let observability = Arc::new(gateway_store::postgres::PgObservabilityRepository::new( - pool.clone(), - )); - let retention = Arc::new(gateway_store::postgres::PgRetentionRepository::new( - pool.clone(), - )); - - let admin_auth_state = Arc::new(gateway_store::redis::RedisAdminAuthStateRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?); - let credential_leases = gateway_store::redis::RedisCredentialLeaseRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?; - let credential_cooldowns = Arc::new( - gateway_store::redis::RedisCredentialCooldownRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?, - ); - let core_account_store = Arc::new( - gateway_store::redis::CooldownCachingProviderAccountStore::new( - provider_accounts.clone(), - credential_cooldowns, - ), - ); - let hydrated_cooldowns = core_account_store.hydrate(SystemTime::now()).await; - tracing::info!(hydrated_cooldowns, "账号 cooldown Redis 热缓存重建完成"); - let provider_leases = Arc::new(ProviderLeaseAdapter::new(credential_leases.clone())); - let credential_state = Arc::new(gateway_store::redis::RedisCredentialStateRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?); - let xai_cache = Arc::new(XaiCatalogCacheAdapter::new(credential_state)); - let admissions = Arc::new(gateway_store::redis::RedisClientAdmissionRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - )?); - let admission_recovery = - gateway_store::postgres::PgClientAdmissionRecoveryRepository::new(pool.clone()); - let admission_recovery_report = restore_client_admission_startup( - execution_store.as_ref(), - &admission_recovery, - admissions.as_ref(), - Utc::now(), +/// 按冻结顺序初始化全部 Bundle，并把进程阻塞权交给 Host。 +pub async fn run() -> Result<(), BootstrapError> { + let config = gateway_host::load_config::()?; + let GatewayConfig { + schema_version: _, + host, + store, + admin, + api, + openai, + xai, + } = config; + + let host = gateway_host::initialize(host).await?; + let mut store = gateway_store::initialize(store).await?; + let provider_ports = store.provider_ports(); + let mut openai = provider_openai::initialize(openai, provider_ports.clone()).await?; + let mut xai = provider_xai::initialize(xai, provider_ports).await?; + let providers = ProviderRegistry::new([openai.core_provider(), xai.core_provider()])?; + let mut core = gateway_core::initialize(store.core_ports(), providers).await?; + let admin = gateway_admin::initialize( + admin, + store.admin_ports(), + vec![openai.admin_provider(), xai.admin_provider()], + core.snapshot_control(), + core.account_probe(), + host.system_operations(), ) .await?; - tracing::info!( - expired_model_requests = admission_recovery_report.expired_model_requests, - restored_clients = admission_recovery_report.restored_clients, - restored_recent_requests = admission_recovery_report.restored_recent_requests, - restored_running_requests = admission_recovery_report.restored_running_requests, - "客户端准入热状态恢复完成" - ); - let circuits = Arc::new(gateway_store::redis::RedisProviderCircuitRepository::new( - redis_connection.clone(), - REDIS_NAMESPACE, - gateway_store::redis::ProviderCircuitPolicy::default(), - )?); - let runtime_changes = Arc::new(gateway_store::redis::RedisRuntimeChangeRepository::new( - redis_client, - REDIS_NAMESPACE, - )?); - let settings = runtime_settings.load_runtime_settings().await?; - let refresh_margin = Duration::from_secs(settings.refresh_margin_seconds); - let refresh_concurrency = settings.refresh_concurrency; - let wire_profile = provider_openai::transport::profile::CodexWireProfileState::new( - provider_openai::transport::profile::CodexWireProfile { - originator: app.wire_profile.originator.clone(), - codex_version: app.wire_profile.codex_version.clone(), - desktop_version: app.wire_profile.desktop_version.clone(), - desktop_build: app.wire_profile.desktop_build.clone(), - os_type: app.wire_profile.os_type.clone(), - os_version: app.wire_profile.os_version.clone(), - arch: app.wire_profile.arch.clone(), - terminal: app.wire_profile.terminal.clone(), - verified_at: app.wire_profile.verified_at, - }, - ); - let core_account_store: Arc = core_account_store; - let codex_repository = provider_openai::credential::CodexCredentialRepository::new(Arc::clone( - &core_account_store, - )); - let codex_catalog = Arc::new( - provider_openai::credential::CodexCredentialCatalogService::new( - codex_repository.clone(), - wire_profile.clone(), - )?, - ); - let codex_selector = Arc::new(provider_openai::credential::CodexCredentialSelector::new( - codex_repository.clone(), - provider_leases.clone(), - provider_openai::credential::CodexCookiePolicy::official()?, - )); - let codex_provider = Arc::new(provider_openai::CodexProvider::new( - codex_selector, - Arc::clone(&codex_catalog), - wire_profile.clone(), - )?); - let codex_token_client = - Arc::new(provider_openai::credential::token_client::official_openai_token_client()?); - let codex_verifier = Arc::new(provider_openai::credential::CodexJwtIdentityVerifier::new( - Box::new(provider_openai::credential::ReqwestOpenAiJwksSource::new()?), - )); - let codex_owner = Arc::new( - provider_openai::credential::CodexCredentialAdminService::new( - codex_repository.clone(), - codex_token_client.clone(), - codex_verifier.clone(), - provider_leases.clone(), - refresh_margin, - )?, - ); - let codex_runtime_refresh = Arc::new( - provider_openai::credential::CodexCredentialRefreshService::new( - codex_repository.clone(), - codex_token_client.clone(), - provider_leases.clone(), - refresh_margin, - )?, - ); - let codex_quota = Arc::new( - provider_openai::credential::CodexCredentialQuotaService::new( - codex_repository.clone(), - wire_profile, - )?, - ); - let codex_oauth: Arc = - Arc::new(provider_openai::credential::CodexOAuthAdminService::new( - Arc::new(InMemoryCodexOAuthPendingStore::new()), - codex_token_client, - codex_verifier.clone(), - Arc::clone(&core_account_store), - provider_openai::credential::CodexCredentialAdmin, - )); + let mut probes = store.health_probes(); + probes.extend(core.health_probes()); + let api = gateway_api::initialize( + api, + core.execution_service(), + admin.services(), + probes, + host.worker_health(), + host.connection_lifecycle(), + )?; - let xai_oauth = Arc::new(provider_xai::GrokOAuthClient::new( - provider_xai::GrokOAuthConfig::official(provider_xai::transport::GROK_CLIENT_VERSION)?, - Arc::new(provider_xai::ReqwestOAuthTransport::new()?), - Arc::new(provider_xai::ReqwestOidcTokenVerifier::new()?), - )); - let xai_repository = - provider_xai::GrokCredentialRepository::new(Arc::clone(&core_account_store)); - let xai_catalog_transport = Arc::new(provider_xai::ReqwestGrokModelCatalogTransport::new()?); - let xai_catalog = Arc::new(provider_xai::GrokCredentialCatalogService::new( - xai_repository.clone(), - xai_catalog_transport.clone(), - xai_cache.clone(), - )); - let xai_quota = Arc::new(provider_xai::GrokCredentialQuotaService::new( - xai_repository.clone(), - xai_catalog_transport, - )); - let xai_refresh = Arc::new(provider_xai::GrokCredentialRefreshService::new( - xai_repository.clone(), - Arc::new(provider_xai::GrokOAuthRefreshClient::new(Arc::clone( - &xai_oauth, - ))), - Arc::clone(&xai_catalog), - provider_leases.clone(), - refresh_margin, - )?); - let xai_selector = Arc::new(provider_xai::GrokAccountSessionSelector::new( - xai_repository.clone(), - xai_cache, - provider_leases, - )); - let xai_provider = Arc::new(provider_xai::GrokBuildProvider::new( - xai_selector, - Arc::new(provider_xai::ReqwestGrokInferenceTransport::new()?), - Arc::clone(&xai_catalog), - )); - - let mut registry = ProviderRegistry::builder(); - registry.register(codex_provider)?; - registry.register(xai_provider)?; - let providers = registry.build(); - let compiler = Arc::new(RuntimeSnapshotCompiler::new( - snapshot_repository, - providers.clone(), - )); - let snapshots = RuntimeSnapshotHandle::new(compiler.compile().await?); - let publisher = RuntimeSnapshotPublisher::new( - Arc::clone(&compiler), - snapshots.clone(), - runtime_changes.clone(), - ); - let connections = ConnectionLifecycle::new(); - let openai = Arc::new(GatewayOpenAiService::new( - snapshots.clone(), - execution_store.clone(), - providers, - admissions, - circuits, - Arc::new(PgNativeContinuationResolver::new(history)), - connections.clone(), - )); - - let admin_backend = Arc::new(StoreAdminAuthBackend::new( - security.clone(), - runtime_settings.clone(), - admin_auth_state, - )); - let sessions = Arc::new(DefaultAdminAuthService::new( - app.admin.default_username.clone(), - app.admin.session_ttl_minutes, - admin_backend, - )); - sessions - .ensure_default_admin(admin_default_password.expose_secret()) - .await?; - let shutdown = TokioCancellationToken::new(); - let accounts_admin = Arc::new(AccountAdminAdapter::new(AccountAdminPorts { - accounts: provider_accounts.clone(), - admin_accounts: provider_accounts.clone(), - core_store: core_account_store, - control_plane: control_plane.clone(), - instances: instances.clone(), - observability: observability.clone(), - security, - codex_owner: Arc::clone(&codex_owner), - codex_quota: Arc::clone(&codex_quota), - codex_catalog: Arc::clone(&codex_catalog), - xai_repository, - xai_refresh: Arc::clone(&xai_refresh), - xai_quota: Arc::clone(&xai_quota), - xai_catalog: Arc::clone(&xai_catalog), - connection_test: Arc::clone(&openai), - publisher: publisher.clone(), - })); - let admin = Arc::new(AdminServices::new(AdminServicePorts { - sessions, - accounts: accounts_admin, - catalog: Arc::new(CatalogAdminAdapter::new( - control_plane.clone(), - instances.clone(), - publisher.clone(), - )), - client_keys: Arc::new(ClientKeyAdminAdapter::new( - control_plane.clone(), - client_keys, - publisher.clone(), - )), - codex: Arc::new(CodexAdminAdapter::new(CodexAdminPorts { - accounts: provider_accounts.clone(), - admin_accounts: provider_accounts.clone(), - core_store: provider_accounts.clone(), - control_plane: control_plane.clone(), - verifier: codex_verifier, - oauth: codex_oauth, - owner: codex_owner, - publisher: publisher.clone(), - })), - observability: Arc::new(ObservabilityAdminAdapter::new( - observability, - provider_accounts.clone(), - runtime_settings.clone(), - app.wire_profile.dashboard_view(), - )), - settings: Arc::new(RuntimeSettingsAdminAdapter::new( - control_plane.clone(), - publisher.clone(), - )), - system: Arc::new(ProcessSystemAdminService::new(shutdown.clone())), - xai: Arc::new(XaiAdminAdapter::new( - provider_accounts.clone(), - provider_accounts, - control_plane, - Arc::new(gateway_store::postgres::PgRuntimeSettingsRepository::new( - pool.clone(), - )), - xai_oauth, - publisher.clone(), - )), - })); - - let supervisor = start_workers(WorkerOwners { - leases: credential_leases, - codex_refresh: codex_runtime_refresh, - codex_quota, - codex_catalog, - xai_refresh, - xai_quota, - xai_catalog, - snapshot_publisher: publisher.clone(), - instances: Arc::new(gateway_store::postgres::PgConfigCatalogRepository::new( - pool.clone(), - )), - accounts: Arc::new(gateway_store::postgres::PgProviderAccountRepository::new( - pool.clone(), - )), - ops_events: Arc::new(gateway_store::postgres::PgOpsEventRepository::new( - pool.clone(), - )), - refresh_concurrency, - execution: execution_store, - retention, - })?; - let health = Arc::new(HealthStatus::new( - snapshots.clone(), - supervisor.health().clone(), - pool.clone(), - redis_connection, - )); - let runtime_settings_reader: Arc = runtime_settings; - let runtime_change_task = spawn_runtime_change_consumer( - runtime_changes, - publisher, - runtime_settings_reader, - shutdown.child_token(), - ); - spawn_shutdown_signal(shutdown.clone()); - - let state = ApplicationState::new(admin, openai.as_ref().clone(), health); - let assets = env::var_os("CPR_WEB_DIST_DIR") - .map(PathBuf::from) - .unwrap_or_else(default_asset_directory); - let router = application_router(state, assets); - let listener = - tokio::net::TcpListener::bind((app.server.host.as_str(), app.server.port)).await?; - tracing::info!(host = %app.server.host, port = app.server.port, "网关开始监听"); - axum::serve( - listener, - router.into_make_service_with_connect_info::(), - ) - .with_graceful_shutdown(shutdown.clone().cancelled_owned()) - .await?; - - shutdown.cancel(); - connections.shutdown().await; - supervisor.shutdown(Duration::from_secs(30)).await; - if let Err(error) = runtime_change_task.await { - tracing::warn!(error = %error, "runtime change consumer did not stop cleanly"); - } - pool.close().await; + let mut plan = store.take_worker_contributions(); + plan.extend(core.take_worker_contributions()); + plan.extend(openai.take_worker_contributions()); + plan.extend(xai.take_worker_contributions()); + host.start_workers(plan, store.worker_leader_lease())?; + host.serve(api.router()).await?; Ok(()) } -fn default_asset_directory() -> PathBuf { - env::current_dir() - .ok() - .and_then(|directory| { - directory - .ancestors() - .map(|ancestor| ancestor.join("frontend/dist")) - .find(|candidate| candidate.is_dir()) - }) - .unwrap_or_else(|| PathBuf::from("web/dist")) -} - -fn spawn_shutdown_signal(shutdown: TokioCancellationToken) { - drop(tokio::spawn(async move { - #[cfg(unix)] - { - let terminate = - tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()); - match terminate { - Ok(mut terminate) => { - tokio::select! { - _ = tokio::signal::ctrl_c() => {} - _ = terminate.recv() => {} - _ = shutdown.cancelled() => return, - } - } - Err(_) => { - tokio::select! { - _ = tokio::signal::ctrl_c() => {} - _ = shutdown.cancelled() => return, - } - } - } - } - #[cfg(not(unix))] - tokio::select! { - _ = tokio::signal::ctrl_c() => {} - _ = shutdown.cancelled() => return, - } - shutdown.cancel(); - })); -} - -fn spawn_runtime_change_consumer( - repository: Arc, - publisher: RuntimeSnapshotPublisher, - runtime_settings: Arc, - shutdown: TokioCancellationToken, -) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - let notification_consumer = - consume_runtime_change_notifications(repository, publisher.clone(), shutdown.clone()); - let revision_reconciler = reconcile_runtime_revision(runtime_settings, publisher, shutdown); - tokio::join!(notification_consumer, revision_reconciler); - }) -} - -async fn consume_runtime_change_notifications( - repository: Arc, - publisher: RuntimeSnapshotPublisher, - shutdown: TokioCancellationToken, -) { - let mut retry_delay = Duration::from_secs(1); - loop { - if shutdown.is_cancelled() { - return; - } - let subscription = repository.subscribe_runtime_changes().await; - let mut subscription = match subscription { - Ok(subscription) => { - retry_delay = Duration::from_secs(1); - subscription - } - Err(_) => { - tracing::warn!("Redis runtime change 订阅失败，稍后重连"); - tokio::select! { - _ = shutdown.cancelled() => return, - _ = tokio::time::sleep(retry_delay) => {} - } - retry_delay = (retry_delay * 2).min(Duration::from_secs(30)); - continue; - } - }; - loop { - let next = tokio::select! { - _ = shutdown.cancelled() => return, - next = futures::StreamExt::next(&mut subscription) => next, - }; - match next { - Some(Ok(gateway_store::redis::RuntimeChange::SnapshotPublished { - config_revision, - })) => match publisher.refresh().await { - Ok(revision) if revision == config_revision => {} - Ok(revision) => tracing::debug!( - notified_revision = config_revision.get(), - loaded_revision = revision.get(), - "runtime change 通知已被更新 revision 覆盖" - ), - Err(error) => { - publisher.suspend(); - tracing::warn!(error = %error, "跨进程 RuntimeSnapshot 刷新失败，数据面已暂停") - } - }, - Some(Ok(gateway_store::redis::RuntimeChange::ProviderAccountChanged { - credential_revision, - .. - })) => tracing::debug!( - credential_revision = credential_revision.get(), - "已消费账号 revision 通知；Provider 下次读取使用新 revision" - ), - Some(Err(_)) | None => break, - } - } - } -} - -async fn reconcile_runtime_revision( - runtime_settings: Arc, - publisher: RuntimeSnapshotPublisher, - shutdown: TokioCancellationToken, -) { - let mut interval = tokio::time::interval(Duration::from_secs(5)); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - interval.tick().await; - loop { - tokio::select! { - _ = shutdown.cancelled() => return, - _ = interval.tick() => {} - } - let settings = match runtime_settings.load_runtime_settings().await { - Ok(settings) => settings, - Err(_) => { - publisher.suspend(); - tracing::error!("Runtime revision 对账无法读取 PostgreSQL，数据面已暂停"); - continue; - } - }; - if !runtime_revision_needs_refresh( - publisher.published_revision(), - settings.config_revision.get(), - ) { - continue; - } - match publisher.refresh().await { - Ok(revision) => tracing::info!( - config_revision = revision.get(), - "RuntimeSnapshot 已按 PostgreSQL revision 对账" - ), - Err(error) => { - publisher.suspend(); - tracing::error!(error = %error, "Runtime revision 对账刷新失败，数据面已暂停"); - } - } - } -} - -#[must_use] -pub fn runtime_revision_needs_refresh( - published_revision: Option, - persisted_revision: u64, -) -> bool { - published_revision != Some(persisted_revision) -} - -struct RedisWorkerLeaderLeasePort { - repository: gateway_store::redis::RedisCredentialLeaseRepository, - owner_id: String, -} - -impl RedisWorkerLeaderLeasePort { - fn new(repository: gateway_store::redis::RedisCredentialLeaseRepository) -> Self { - Self { - repository, - owner_id: format!("worker-process-{}", Uuid::now_v7().simple()), - } - } -} - -struct RedisWorkerLeaderLeaseGuard { - repository: gateway_store::redis::RedisCredentialLeaseRepository, - request: gateway_store::redis::CredentialLeaseRequest, - grant: Option, - token: crate::workers::WorkerFencingToken, -} - -#[async_trait] -impl crate::workers::WorkerLeaderLeaseGuard for RedisWorkerLeaderLeaseGuard { - fn fencing_token(&self) -> crate::workers::WorkerFencingToken { - self.token - } - - async fn renew(&mut self) -> Result<(), crate::workers::WorkerLeaseError> { - let current = self - .grant - .as_ref() - .ok_or_else(|| crate::workers::WorkerLeaseError::safe("worker lease is released"))?; - let renewed = gateway_store::redis::CredentialLeaseRepository::renew_credential_lease( - &self.repository, - &self.request, - current, - ) - .await - .map_err(|_| crate::workers::WorkerLeaseError::safe("worker lease renewal failed"))? - .ok_or_else(|| crate::workers::WorkerLeaseError::safe("worker lease was lost"))?; - self.grant = Some(renewed); - Ok(()) - } -} - -impl Drop for RedisWorkerLeaderLeaseGuard { - fn drop(&mut self) { - let Some(grant) = self.grant.take() else { - return; - }; - let repository = self.repository.clone(); - let request = self.request.clone(); - if let Ok(runtime) = tokio::runtime::Handle::try_current() { - drop(runtime.spawn(async move { - let _ = gateway_store::redis::CredentialLeaseRepository::release_credential_lease( - &repository, - &request, - &grant, - ) - .await; - })); - } - } -} - -#[async_trait] -impl crate::workers::WorkerLeaderLeasePort for RedisWorkerLeaderLeasePort { - async fn try_acquire( - &self, - request: crate::workers::WorkerLeaseRequest, - ) -> Result { - let lease_request = gateway_store::redis::CredentialLeaseRequest { - scope: gateway_store::redis::CredentialLeaseScope::ProviderTask, - resource_id: request.worker().to_string(), - owner_id: self.owner_id.clone(), - ttl: request.ttl(), - }; - let grant = gateway_store::redis::CredentialLeaseRepository::acquire_credential_lease( - &self.repository, - &lease_request, - ) - .await - .map_err(|_| crate::workers::WorkerLeaseError::safe("worker lease acquisition failed"))?; - let Some(grant) = grant else { - return Ok(crate::workers::WorkerLeaseAcquisition::Busy { retry_after: None }); - }; - let token = std::num::NonZeroU64::new(grant.fencing_token.get()) - .map(crate::workers::WorkerFencingToken::new) - .ok_or_else(|| crate::workers::WorkerLeaseError::safe("worker fence is invalid"))?; - Ok(crate::workers::WorkerLeaseAcquisition::Acquired(Box::new( - RedisWorkerLeaderLeaseGuard { - repository: self.repository.clone(), - request: lease_request, - grant: Some(grant), - token, - }, - ))) - } -} - -struct CodexOAuthRefreshTask { - service: Arc, - limit: u32, - permits: Arc, - ops_events: Arc, -} - -#[async_trait] -impl crate::workers::WorkerTask for CodexOAuthRefreshTask { - async fn run_cycle( - &self, - context: crate::workers::WorkerCycleContext, - ) -> Result<(), crate::workers::WorkerTaskError> { - if context.is_cancelled() { - return Ok(()); - } - let _permit = self - .permits - .acquire() - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("OAuth refresh stopped"))?; - let outcomes = self - .service - .refresh_due(self.limit) - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("Codex OAuth refresh failed"))?; - for outcome in &outcomes { - if let Some((account_id, failure_kind)) = codex_refresh_failure(outcome) { - append_worker_account_ops_event( - self.ops_events.as_ref(), - "openai", - None, - account_id, - "oauth_refresh", - failure_kind, - ) - .await?; - } - } - Ok(()) - } -} - -struct XaiOAuthRefreshTask { - service: Arc, - limit: u32, - permits: Arc, - ops_events: Arc, -} - -#[async_trait] -impl crate::workers::WorkerTask for XaiOAuthRefreshTask { - async fn run_cycle( - &self, - context: crate::workers::WorkerCycleContext, - ) -> Result<(), crate::workers::WorkerTaskError> { - if context.is_cancelled() { - return Ok(()); - } - let _permit = self - .permits - .acquire() - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("OAuth refresh stopped"))?; - let outcomes = self - .service - .refresh_due(self.limit) - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("xAI OAuth refresh failed"))?; - for outcome in &outcomes { - if let Some((account_id, failure_kind)) = xai_refresh_failure(outcome) { - append_worker_account_ops_event( - self.ops_events.as_ref(), - "xai", - None, - account_id, - "oauth_refresh", - failure_kind, - ) - .await?; - } - } - Ok(()) - } -} - -fn codex_refresh_failure( - outcome: &provider_openai::credential::CodexCredentialRefreshOutcome, -) -> Option<(&str, &'static str)> { - use provider_openai::credential::CodexCredentialRefreshOutcome as Outcome; - match outcome { - Outcome::Invalidated { account_id } => Some((account_id, "refresh_invalidated")), - Outcome::Banned { account_id } => Some((account_id, "refresh_banned")), - Outcome::Transient { account_id } => Some((account_id, "refresh_transient")), - Outcome::Ambiguous { account_id } => Some((account_id, "refresh_ambiguous")), - Outcome::Failed { account_id } => Some((account_id, "refresh_failed")), - Outcome::Refreshed { .. } | Outcome::LeaseUnavailable { .. } | Outcome::Stale { .. } => { - None - } - } -} - -fn xai_refresh_failure( - outcome: &provider_xai::GrokCredentialRefreshOutcome, -) -> Option<(&str, &'static str)> { - use provider_xai::GrokCredentialRefreshOutcome as Outcome; - match outcome { - Outcome::Invalidated { account_id } => Some((account_id.as_str(), "refresh_invalidated")), - Outcome::Ambiguous { account_id } => Some((account_id.as_str(), "refresh_ambiguous")), - Outcome::Transient { account_id } => Some((account_id.as_str(), "refresh_transient")), - Outcome::Rejected { account_id } => Some((account_id.as_str(), "refresh_rejected")), - Outcome::Failed { account_id } => Some((account_id.as_str(), "refresh_failed")), - Outcome::Refreshed { .. } | Outcome::LeaseUnavailable { .. } | Outcome::Stale { .. } => { - None - } - } -} - -struct ProviderQuotaCatalogTask { - instances: Arc, - accounts: Arc, - ops_events: Arc, - codex_quota: Arc, - codex_catalog: Arc, - xai_quota: Arc, - xai_catalog: Arc, - snapshot_publisher: RuntimeSnapshotPublisher, -} - -#[async_trait] -impl crate::workers::WorkerTask for ProviderQuotaCatalogTask { - async fn run_cycle( - &self, - context: crate::workers::WorkerCycleContext, - ) -> Result<(), crate::workers::WorkerTaskError> { - let instances = self - .instances - .list_provider_instances(false) - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("Provider catalog read failed"))?; - for instance in instances { - if context.is_cancelled() { - return Ok(()); - } - let accounts = self - .accounts - .list_provider_accounts(Some(&instance.id), false) - .await - .map_err(|_| { - crate::workers::WorkerTaskError::safe("Provider account catalog read failed") - })?; - if accounts.is_empty() { - continue; - } - match instance.provider_kind.as_str() { - "openai" => { - let core_instance = core_provider_instance(&instance).map_err(|_| { - crate::workers::WorkerTaskError::safe("Codex instance is invalid") - })?; - for account in accounts { - if context.is_cancelled() { - return Ok(()); - } - let account_id = - match gateway_core::engine::credential::ProviderAccountId::new( - account.id.clone(), - ) { - Ok(account_id) => account_id, - Err(_) => { - self.record_account_failure( - &instance, - &account.id, - "account_validation", - "invalid_account_id", - ) - .await?; - continue; - } - }; - if self - .codex_quota - .refresh_account(&core_instance, &account_id) - .await - .is_err() - { - self.record_account_failure( - &instance, - account_id.as_str(), - "quota_refresh", - "provider_account_quota_refresh_failed", - ) - .await?; - } - if self - .codex_catalog - .synchronize_account(&core_instance, &account_id) - .await - .is_err() - { - self.record_account_failure( - &instance, - account_id.as_str(), - "catalog_refresh", - "provider_account_catalog_refresh_failed", - ) - .await?; - } - } - } - "xai" => { - for account in accounts { - if context.is_cancelled() { - return Ok(()); - } - let account_id = - match gateway_core::engine::credential::ProviderAccountId::new( - account.id.clone(), - ) { - Ok(account_id) => account_id, - Err(_) => { - self.record_account_failure( - &instance, - &account.id, - "account_validation", - "invalid_account_id", - ) - .await?; - continue; - } - }; - if self.xai_quota.refresh_account(&account_id).await.is_err() { - self.record_account_failure( - &instance, - account_id.as_str(), - "quota_refresh", - "provider_account_quota_refresh_failed", - ) - .await?; - } - if self - .xai_catalog - .refresh_account_catalog( - &account_id, - provider_xai::transport::GROK_CLIENT_VERSION, - ) - .await - .is_err() - { - self.record_account_failure( - &instance, - account_id.as_str(), - "catalog_refresh", - "provider_account_catalog_refresh_failed", - ) - .await?; - } - } - } - _ => {} - } - } - self.snapshot_publisher.refresh().await.map_err(|error| { - tracing::warn!(error = %error, "Provider catalog cycle could not refresh RuntimeSnapshot"); - crate::workers::WorkerTaskError::safe("Runtime snapshot catalog refresh failed") - })?; - Ok(()) - } -} - -impl ProviderQuotaCatalogTask { - async fn record_account_failure( - &self, - instance: &gateway_store::postgres::ProviderInstanceRecord, - account_id: &str, - operation: &'static str, - failure_kind: &'static str, - ) -> Result<(), crate::workers::WorkerTaskError> { - append_worker_account_ops_event( - self.ops_events.as_ref(), - &instance.provider_kind, - Some(&instance.id), - account_id, - operation, - failure_kind, - ) - .await - } -} - -async fn append_worker_account_ops_event( - ops_events: &dyn gateway_store::postgres::OpsEventRepository, - provider_kind: &str, - provider_instance_id: Option<&str>, - account_id: &str, - operation: &'static str, - failure_kind: &'static str, -) -> Result<(), crate::workers::WorkerTaskError> { - tracing::warn!( - provider_kind, - operation, - "Provider 单账号后台同步失败，继续处理其他账号" - ); - ops_events - .append_ops_event(gateway_store::postgres::OpsEvent { - id: format!("ops_{}", Uuid::now_v7().simple()), - model_request_id: None, - attempt_index: None, - level: gateway_store::postgres::OpsEventLevel::Warning, - component: "worker".to_owned(), - operation: operation.to_owned(), - provider_instance_id: provider_instance_id.map(str::to_owned), - provider_kind: Some(provider_kind.to_owned()), - provider_account_id: Some(account_id.to_owned()), - provider_account_ref: Some(account_id.to_owned()), - upstream_model_id: None, - failure_kind: failure_kind.to_owned(), - status_code: None, - provider_error_code: None, - retry_after_ms: None, - upstream_request_id: None, - latency_ms: None, - message: "Provider account background synchronization failed".to_owned(), - occurrence_count: 1, - created_at: Utc::now(), - }) - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("worker ops event write failed")) -} - -struct StaleModelRequestRecoveryTask { - repository: Arc, -} - -#[async_trait] -impl crate::workers::WorkerTask for StaleModelRequestRecoveryTask { - async fn run_cycle( - &self, - _context: crate::workers::WorkerCycleContext, - ) -> Result<(), crate::workers::WorkerTaskError> { - gateway_store::postgres::ModelRequestRepository::recover_expired_model_requests( - self.repository.as_ref(), - Utc::now(), - ) - .await - .map(|_| ()) - .map_err(|_| crate::workers::WorkerTaskError::safe("stale request recovery failed")) - } -} - -struct RetentionTask { - repository: Arc, -} - -#[async_trait] -impl crate::workers::WorkerTask for RetentionTask { - async fn run_cycle( - &self, - _context: crate::workers::WorkerCycleContext, - ) -> Result<(), crate::workers::WorkerTaskError> { - let settings = gateway_store::postgres::RetentionRepository::load_retention_settings( - self.repository.as_ref(), - ) - .await - .map_err(|_| crate::workers::WorkerTaskError::safe("retention settings read failed"))?; - gateway_store::postgres::RetentionRepository::apply_retention( - self.repository.as_ref(), - Utc::now(), - settings, - ) - .await - .map(|_| ()) - .map_err(|_| crate::workers::WorkerTaskError::safe("retention cleanup failed")) - } -} - -struct WorkerOwners { - leases: gateway_store::redis::RedisCredentialLeaseRepository, - codex_refresh: Arc, - codex_quota: Arc, - codex_catalog: Arc, - xai_refresh: Arc, - xai_quota: Arc, - xai_catalog: Arc, - snapshot_publisher: RuntimeSnapshotPublisher, - instances: Arc, - accounts: Arc, - ops_events: Arc, - refresh_concurrency: u32, - execution: Arc, - retention: Arc, -} - -fn start_workers( - owners: WorkerOwners, -) -> Result { - use crate::workers::{WorkerDisabledReason, WorkerKind, WorkerRegistry, WorkerSchedule}; - let schedule = |interval| { - WorkerSchedule::new( - interval, - Duration::from_secs(1), - Duration::from_secs(60), - Duration::from_secs(15 * 60), - ) - }; - let mut registry = WorkerRegistry::new(); - let refresh_permits = Arc::new(tokio::sync::Semaphore::new( - usize::try_from(owners.refresh_concurrency).unwrap_or(usize::MAX), - )); - registry.register( - WorkerKind::OAuthRefresh, - "openai", - Arc::new(CodexOAuthRefreshTask { - service: owners.codex_refresh, - limit: owners.refresh_concurrency, - permits: Arc::clone(&refresh_permits), - ops_events: Arc::clone(&owners.ops_events), - }), - schedule(Duration::from_secs(30))?, - )?; - registry.register( - WorkerKind::OAuthRefresh, - "xai", - Arc::new(XaiOAuthRefreshTask { - service: owners.xai_refresh, - limit: owners.refresh_concurrency, - permits: refresh_permits, - ops_events: Arc::clone(&owners.ops_events), - }), - schedule(Duration::from_secs(30))?, - )?; - registry.register( - WorkerKind::QuotaCatalogHealth, - "provider-registry", - Arc::new(ProviderQuotaCatalogTask { - instances: owners.instances, - accounts: owners.accounts, - ops_events: owners.ops_events, - codex_quota: owners.codex_quota, - codex_catalog: owners.codex_catalog, - xai_quota: owners.xai_quota, - xai_catalog: owners.xai_catalog, - snapshot_publisher: owners.snapshot_publisher, - }), - schedule(Duration::from_secs(5 * 60))?, - )?; - registry.disable(WorkerDisabledReason::NoPersistentNativeClaimState)?; - registry.register( - WorkerKind::StaleModelRequestRecovery, - "postgres", - Arc::new(StaleModelRequestRecoveryTask { - repository: owners.execution, - }), - schedule(Duration::from_secs(30))?, - )?; - registry.register( - WorkerKind::Retention, - "postgres", - Arc::new(RetentionTask { - repository: owners.retention, - }), - schedule(Duration::from_secs(60 * 60))?, - )?; - registry.disable(WorkerDisabledReason::NoBufferedOpsEvents)?; - registry.start(Arc::new(RedisWorkerLeaderLeasePort::new(owners.leases))) -} - -use std::fs::{self, File, OpenOptions}; -use std::io::{self, Write}; - -use tracing_appender::non_blocking::{NonBlockingBuilder, WorkerGuard}; -use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt as _}; - -const LOG_FILE_PREFIX: &str = "codex-proxy-rs"; - -/// non-blocking 日志 writer 的进程级守卫。 -pub struct LogGuard { - _writers: Vec, -} - +/// 组合根只保留包级错误分类，不展开内部实现或敏感配置。 #[derive(Debug, thiserror::Error)] -pub enum LogError { - #[error("log IO failed")] - Io(#[from] io::Error), - #[error("logging filter is invalid")] - InvalidFilter, - #[error("global tracing subscriber is already initialized")] - AlreadyInitialized, - #[error("logging size limit is too large")] - SizeOverflow, -} - -/// 按自然日、单文件大小、保留天数和文件总数初始化结构化日志。 -pub fn initialize_logging(config: &LoggingConfig) -> Result { - let directive = env::var("RUST_LOG").unwrap_or_else(|_| config.level.clone()); - let filter = EnvFilter::try_new(directive).map_err(|_| LogError::InvalidFilter)?; - let mut guards = Vec::new(); - - let stdout_writer = config.stdout.then(|| { - let (writer, guard) = NonBlockingBuilder::default() - .thread_name("gateway-log-stdout") - .finish(io::stdout()); - guards.push(guard); - writer - }); - let file_writer = if config.file.enabled { - let maximum_bytes = config - .file - .max_file_size_mb - .checked_mul(1024 * 1024) - .ok_or(LogError::SizeOverflow)?; - let writer = RotatingLogWriter::open( - config.file.directory.clone(), - maximum_bytes, - config.file.retention_days, - config.file.max_files, - )?; - let (writer, guard) = NonBlockingBuilder::default() - .thread_name("gateway-log-file") - .finish(writer); - guards.push(guard); - Some(writer) - } else { - None - }; - - let stdout_layer = stdout_writer.map(|writer| { - tracing_subscriber::fmt::layer() - .compact() - .with_writer(writer) - .with_target(true) - .with_ansi(false) - }); - let file_layer = file_writer.map(|writer| { - tracing_subscriber::fmt::layer() - .json() - .with_writer(writer) - .with_target(true) - .with_file(true) - .with_line_number(true) - .with_thread_ids(true) - .with_thread_names(true) - .with_current_span(true) - .with_span_list(true) - }); - tracing_subscriber::registry() - .with(filter) - .with(stdout_layer) - .with(file_layer) - .try_init() - .map_err(|_| LogError::AlreadyInitialized)?; - Ok(LogGuard { _writers: guards }) -} - -struct RotatingLogWriter { - directory: PathBuf, - maximum_bytes: u64, - retention_days: usize, - maximum_files: usize, - date: chrono::NaiveDate, - segment: usize, - bytes_written: u64, - file: File, -} - -impl RotatingLogWriter { - fn open( - directory: PathBuf, - maximum_bytes: u64, - retention_days: usize, - maximum_files: usize, - ) -> io::Result { - fs::create_dir_all(&directory)?; - let date = Utc::now().date_naive(); - cleanup_log_files(&directory, date, retention_days, maximum_files)?; - let (segment, bytes_written) = writable_log_segment(&directory, date, maximum_bytes)?; - let file = open_log_segment(&directory, date, segment)?; - Ok(Self { - directory, - maximum_bytes, - retention_days, - maximum_files, - date, - segment, - bytes_written, - file, - }) - } - - fn rotate_if_required(&mut self, incoming_bytes: usize) -> io::Result<()> { - let date = Utc::now().date_naive(); - let incoming_bytes = u64::try_from(incoming_bytes).unwrap_or(u64::MAX); - let day_changed = date != self.date; - let size_exceeded = !day_changed - && self.bytes_written > 0 - && self.bytes_written.saturating_add(incoming_bytes) > self.maximum_bytes; - if !day_changed && !size_exceeded { - return Ok(()); - } - self.file.flush()?; - if day_changed { - self.date = date; - self.segment = 0; - } else { - self.segment = self.segment.saturating_add(1); - } - self.file = open_log_segment(&self.directory, self.date, self.segment)?; - self.bytes_written = self.file.metadata()?.len(); - cleanup_log_files( - &self.directory, - self.date, - self.retention_days, - self.maximum_files, - ) - } -} - -impl Write for RotatingLogWriter { - fn write(&mut self, buffer: &[u8]) -> io::Result { - self.rotate_if_required(buffer.len())?; - let written = self.file.write(buffer)?; - self.bytes_written = self - .bytes_written - .saturating_add(u64::try_from(written).unwrap_or(u64::MAX)); - Ok(written) - } - - fn flush(&mut self) -> io::Result<()> { - self.file.flush() - } -} - -#[derive(Debug)] -struct ManagedLogFile { - date: chrono::NaiveDate, - segment: usize, - path: PathBuf, -} - -fn writable_log_segment( - directory: &Path, - date: chrono::NaiveDate, - maximum_bytes: u64, -) -> io::Result<(usize, u64)> { - let latest = managed_log_files(directory)? - .into_iter() - .filter(|entry| entry.date == date) - .max_by_key(|entry| entry.segment); - let Some(latest) = latest else { - return Ok((0, 0)); - }; - let length = latest.path.metadata()?.len(); - if length >= maximum_bytes { - Ok((latest.segment.saturating_add(1), 0)) - } else { - Ok((latest.segment, length)) - } -} - -fn open_log_segment(directory: &Path, date: chrono::NaiveDate, segment: usize) -> io::Result { - OpenOptions::new() - .create(true) - .append(true) - .open(directory.join(log_file_name(date, segment))) -} - -fn log_file_name(date: chrono::NaiveDate, segment: usize) -> String { - if segment == 0 { - format!("{LOG_FILE_PREFIX}.{date}.log") - } else { - format!("{LOG_FILE_PREFIX}.{date}.{segment}.log") - } -} - -fn managed_log_files(directory: &Path) -> io::Result> { - Ok(fs::read_dir(directory)? - .filter_map(Result::ok) - .filter_map(|entry| { - parse_log_file_name(entry.file_name().to_string_lossy().as_ref()).map( - |(date, segment)| ManagedLogFile { - date, - segment, - path: entry.path(), - }, - ) - }) - .collect()) -} - -fn parse_log_file_name(name: &str) -> Option<(chrono::NaiveDate, usize)> { - let body = name - .strip_prefix(&format!("{LOG_FILE_PREFIX}."))? - .strip_suffix(".log")?; - let (date, segment) = body.split_once('.').map_or((body, 0), |(date, segment)| { - (date, segment.parse().ok().unwrap_or(usize::MAX)) - }); - if segment == usize::MAX { - return None; - } - chrono::NaiveDate::parse_from_str(date, "%Y-%m-%d") - .ok() - .map(|date| (date, segment)) -} - -fn cleanup_log_files( - directory: &Path, - today: chrono::NaiveDate, - retention_days: usize, - maximum_files: usize, -) -> io::Result<()> { - let retention_offset = i64::try_from(retention_days.saturating_sub(1)).unwrap_or(i64::MAX); - let cutoff = today - .checked_sub_signed(chrono::Duration::days(retention_offset)) - .unwrap_or(chrono::NaiveDate::MIN); - let mut retained = Vec::new(); - for entry in managed_log_files(directory)? { - if entry.date < cutoff { - fs::remove_file(entry.path)?; - } else { - retained.push(entry); - } - } - retained.sort_by(|left, right| { - right - .date - .cmp(&left.date) - .then_with(|| right.segment.cmp(&left.segment)) - }); - for entry in retained.into_iter().skip(maximum_files) { - fs::remove_file(entry.path)?; - } - Ok(()) +pub enum BootstrapError { + #[error(transparent)] + Config(#[from] gateway_host::ConfigError), + #[error(transparent)] + Host(#[from] gateway_host::HostError), + #[error(transparent)] + Store(#[from] gateway_store::StoreError), + #[error(transparent)] + OpenAi(#[from] provider_openai::OpenAiInitializeError), + #[error(transparent)] + Xai(#[from] provider_xai::XaiInitializeError), + #[error(transparent)] + Registry(#[from] RegistryError), + #[error(transparent)] + Core(#[from] gateway_core::CoreError), + #[error(transparent)] + Admin(#[from] gateway_admin::model::AdminError), + #[error(transparent)] + Api(#[from] gateway_api::ApiError), } diff --git a/backend/apps/gateway/src/lib.rs b/backend/apps/gateway/src/lib.rs new file mode 100644 index 000000000..615bb4e77 --- /dev/null +++ b/backend/apps/gateway/src/lib.rs @@ -0,0 +1,3 @@ +//! 网关进程组装根。 + +pub mod bootstrap; diff --git a/backend/apps/gateway/src/main.rs b/backend/apps/gateway/src/main.rs index 548fe8230..a1a79f653 100644 --- a/backend/apps/gateway/src/main.rs +++ b/backend/apps/gateway/src/main.rs @@ -1,14 +1,12 @@ use std::{env, error::Error, io}; -use codex_proxy_rs::{BootstrapConfig, run}; - fn main() -> Result<(), Box> { let mut arguments = env::args_os().skip(1); match arguments.next().as_deref().and_then(|value| value.to_str()) { None | Some("serve") => { reject_extra_arguments(arguments)?; - let config = BootstrapConfig::load()?; - runtime()?.block_on(run(config)) + runtime()?.block_on(codex_proxy_rs::bootstrap::run())?; + Ok(()) } Some(command) => { Err(invalid_cli(&format!("unknown command {command:?}; expected serve")).into()) diff --git a/backend/apps/gateway/src/workers.rs b/backend/apps/gateway/src/workers.rs deleted file mode 100644 index 61b694e0e..000000000 --- a/backend/apps/gateway/src/workers.rs +++ /dev/null @@ -1,894 +0,0 @@ -//! 后台任务注册、leader lease、监督、退避、健康状态与优雅关闭。 - -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; -use std::num::NonZeroU64; -use std::panic::AssertUnwindSafe; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, SystemTime}; - -use async_trait::async_trait; -use futures::FutureExt as _; -use tokio::task::JoinHandle; -use tokio_util::sync::CancellationToken; - -/// 冻结架构允许出现的六类后台任务组。 -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum WorkerKind { - OAuthRefresh, - QuotaCatalogHealth, - NativeClaimRecovery, - StaleModelRequestRecovery, - Retention, - OpsFlush, -} - -impl WorkerKind { - pub const ALL: [Self; 6] = [ - Self::OAuthRefresh, - Self::QuotaCatalogHealth, - Self::NativeClaimRecovery, - Self::StaleModelRequestRecovery, - Self::Retention, - Self::OpsFlush, - ]; - - #[must_use] - pub const fn as_str(self) -> &'static str { - match self { - Self::OAuthRefresh => "oauth_refresh", - Self::QuotaCatalogHealth => "quota_catalog_health", - Self::NativeClaimRecovery => "native_claim_recovery", - Self::StaleModelRequestRecovery => "stale_model_request_recovery", - Self::Retention => "retention", - Self::OpsFlush => "ops_flush", - } - } -} - -impl fmt::Display for WorkerKind { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(self.as_str()) - } -} - -/// 同类任务中的真实 owner 标识；一个任务组可以组合多个 owner。 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct WorkerId { - kind: WorkerKind, - owner: String, -} - -impl WorkerId { - pub fn new(kind: WorkerKind, owner: impl Into) -> Result { - let owner = owner.into(); - if !valid_owner(&owner) { - return Err(WorkerRegistryError::InvalidOwner(owner)); - } - Ok(Self { kind, owner }) - } - - #[must_use] - pub const fn kind(&self) -> WorkerKind { - self.kind - } - - #[must_use] - pub fn owner(&self) -> &str { - &self.owner - } -} - -impl fmt::Display for WorkerId { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(formatter, "{}:{}", self.kind, self.owner) - } -} - -/// Final DB 明确不存在对应状态时允许使用的两种终态。 -/// -/// 该枚举刻意不提供通用字符串分支，避免把缺少 owner 的真实任务伪装成禁用。 -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum WorkerDisabledReason { - NoPersistentNativeClaimState, - NoBufferedOpsEvents, -} - -impl WorkerDisabledReason { - #[must_use] - pub const fn kind(self) -> WorkerKind { - match self { - Self::NoPersistentNativeClaimState => WorkerKind::NativeClaimRecovery, - Self::NoBufferedOpsEvents => WorkerKind::OpsFlush, - } - } - - #[must_use] - pub const fn as_str(self) -> &'static str { - match self { - Self::NoPersistentNativeClaimState => { - "native continuation has no persistent claim state" - } - Self::NoBufferedOpsEvents => "ops events have no flush buffer", - } - } -} - -impl fmt::Display for WorkerDisabledReason { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(self.as_str()) - } -} - -/// 正整数 fencing token。 -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct WorkerFencingToken(NonZeroU64); - -impl WorkerFencingToken { - #[must_use] - pub const fn new(value: NonZeroU64) -> Self { - Self(value) - } - - #[must_use] - pub const fn get(self) -> u64 { - self.0.get() - } -} - -/// 一次 leader lease 申请；resource 由任务 kind 与 owner 共同确定。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct WorkerLeaseRequest { - worker: WorkerId, - ttl: Duration, -} - -impl WorkerLeaseRequest { - #[must_use] - pub const fn worker(&self) -> &WorkerId { - &self.worker - } - - #[must_use] - pub const fn ttl(&self) -> Duration { - self.ttl - } -} - -/// Leader lease guard 暴露 fencing token 与续租操作，并在 Drop 时释放。 -#[async_trait] -pub trait WorkerLeaderLeaseGuard: Send + Sync + 'static { - fn fencing_token(&self) -> WorkerFencingToken; - - async fn renew(&mut self) -> Result<(), WorkerLeaseError>; -} - -/// Leader lease 的稳定获取结果。 -pub enum WorkerLeaseAcquisition { - Acquired(Box), - Busy { retry_after: Option }, -} - -impl fmt::Debug for WorkerLeaseAcquisition { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Acquired(_) => formatter.write_str("Acquired([LEASE_GUARD])"), - Self::Busy { retry_after } => formatter - .debug_struct("Busy") - .field("retry_after", retry_after) - .finish(), - } - } -} - -/// 多实例 worker leader lease 的 App 抽象端口。 -#[async_trait] -pub trait WorkerLeaderLeasePort: Send + Sync + 'static { - async fn try_acquire( - &self, - request: WorkerLeaseRequest, - ) -> Result; -} - -/// 不暴露 Redis 原文或资源标识的 lease 错误。 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("{message}")] -pub struct WorkerLeaseError { - message: String, -} - -impl WorkerLeaseError { - #[must_use] - pub fn safe(message: impl Into) -> Self { - Self { - message: message.into(), - } - } - - #[must_use] - pub fn as_safe_str(&self) -> &str { - &self.message - } -} - -/// 任务算法收到的单周期上下文。 -#[derive(Clone)] -pub struct WorkerCycleContext { - worker: WorkerId, - fencing_token: WorkerFencingToken, - cancellation: CancellationToken, -} - -impl WorkerCycleContext { - #[must_use] - pub const fn worker(&self) -> &WorkerId { - &self.worker - } - - #[must_use] - pub const fn fencing_token(&self) -> WorkerFencingToken { - self.fencing_token - } - - #[must_use] - pub fn is_cancelled(&self) -> bool { - self.cancellation.is_cancelled() - } - - pub async fn cancelled(&self) { - self.cancellation.cancelled().await; - } -} - -/// 任务算法的外层端口。实现必须位于对应 Store 或 Provider owner。 -#[async_trait] -pub trait WorkerTask: Send + Sync + 'static { - async fn run_cycle(&self, context: WorkerCycleContext) -> Result<(), WorkerTaskError>; -} - -/// 不暴露 Provider 原文或存储细节的后台任务错误。 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("{message}")] -pub struct WorkerTaskError { - message: String, -} - -impl WorkerTaskError { - #[must_use] - pub fn safe(message: impl Into) -> Self { - Self { - message: message.into(), - } - } - - #[must_use] - pub fn as_safe_str(&self) -> &str { - &self.message - } -} - -/// 单个 owner 任务的固定监督策略。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct WorkerSchedule { - interval: Duration, - initial_backoff: Duration, - maximum_backoff: Duration, - leader_lease_ttl: Duration, - leader_lease_renewal_interval: Duration, -} - -impl WorkerSchedule { - /// 所有时长都必须大于零，且最大退避不能小于初始退避。 - pub fn new( - interval: Duration, - initial_backoff: Duration, - maximum_backoff: Duration, - leader_lease_ttl: Duration, - ) -> Result { - let Some(leader_lease_renewal_interval) = leader_lease_ttl.checked_div(3) else { - return Err(WorkerRegistryError::InvalidSchedule); - }; - if interval.is_zero() - || initial_backoff.is_zero() - || maximum_backoff < initial_backoff - || leader_lease_ttl.is_zero() - || leader_lease_renewal_interval.is_zero() - { - return Err(WorkerRegistryError::InvalidSchedule); - } - Ok(Self { - interval, - initial_backoff, - maximum_backoff, - leader_lease_ttl, - leader_lease_renewal_interval, - }) - } - - fn success_freshness(self) -> Duration { - self.interval - .saturating_mul(2) - .saturating_add(self.maximum_backoff) - } -} - -struct WorkerRegistration { - id: WorkerId, - task: Arc, - schedule: WorkerSchedule, -} - -/// 六类终态 worker 的唯一注册表。 -#[derive(Default)] -pub struct WorkerRegistry { - registrations: BTreeMap, - disabled: BTreeMap, -} - -impl WorkerRegistry { - #[must_use] - pub const fn new() -> Self { - Self { - registrations: BTreeMap::new(), - disabled: BTreeMap::new(), - } - } - - /// 注册一个真实 Store/Provider-owned 周期任务。 - pub fn register( - &mut self, - kind: WorkerKind, - owner: impl Into, - task: Arc, - schedule: WorkerSchedule, - ) -> Result<(), WorkerRegistryError> { - if self.disabled.contains_key(&kind) { - return Err(WorkerRegistryError::KindDisabled(kind)); - } - let id = WorkerId::new(kind, owner)?; - if self.registrations.contains_key(&id) { - return Err(WorkerRegistryError::Duplicate(id)); - } - self.registrations - .insert(id.clone(), WorkerRegistration { id, task, schedule }); - Ok(()) - } - - /// 声明 Final DB 不存在相应可恢复或可 flush 状态。 - pub fn disable(&mut self, reason: WorkerDisabledReason) -> Result<(), WorkerRegistryError> { - let kind = reason.kind(); - if self - .registrations - .keys() - .any(|registered| registered.kind() == kind) - { - return Err(WorkerRegistryError::KindHasTasks(kind)); - } - if self.disabled.insert(kind, reason).is_some() { - return Err(WorkerRegistryError::DuplicateDisabled(kind)); - } - Ok(()) - } - - /// 校验六类 worker 都有真实 owner 或合法终态后启动监督器。 - pub fn start( - self, - leader_leases: Arc, - ) -> Result { - let active_kinds = self - .registrations - .keys() - .map(WorkerId::kind) - .collect::>(); - let missing = WorkerKind::ALL - .into_iter() - .filter(|kind| !active_kinds.contains(kind) && !self.disabled.contains_key(kind)) - .collect::>(); - if !missing.is_empty() { - return Err(WorkerRegistryError::Missing(missing)); - } - - let cancellation = CancellationToken::new(); - let health = WorkerHealthRegistry::new(&self.registrations, &self.disabled); - let active_ids = self.registrations.keys().cloned().collect::>(); - let handles = self - .registrations - .into_values() - .map(|registration| { - let id = registration.id.clone(); - let handle = spawn_worker( - registration, - Arc::clone(&leader_leases), - cancellation.child_token(), - health.clone(), - ); - (id, handle) - }) - .collect(); - Ok(WorkerSupervisor { - cancellation, - health, - active_ids, - handles, - }) - } -} - -/// 注册或启动监督器失败。 -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -pub enum WorkerRegistryError { - #[error("worker schedule is invalid")] - InvalidSchedule, - #[error("worker owner is invalid: {0}")] - InvalidOwner(String), - #[error("worker `{0}` was registered more than once")] - Duplicate(WorkerId), - #[error("worker kind `{0}` is already disabled")] - KindDisabled(WorkerKind), - #[error("worker kind `{0}` already has real tasks")] - KindHasTasks(WorkerKind), - #[error("worker kind `{0}` was disabled more than once")] - DuplicateDisabled(WorkerKind), - #[error("required workers are missing: {0:?}")] - Missing(Vec), -} - -/// 一个 owner 任务的运行状态。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum WorkerRuntimeState { - Starting, - AcquiringLease, - Standby, - Running, - Idle, - BackingOff, - Stopped, - Disabled(WorkerDisabledReason), -} - -/// 健康快照中的稳定 key。 -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -pub enum WorkerHealthKey { - Task(WorkerId), - Disabled(WorkerKind), -} - -/// 一个任务或终态 worker kind 的监督状态。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct WorkerHealth { - pub state: WorkerRuntimeState, - pub consecutive_failures: u32, - pub completed_cycles: u64, - pub last_fencing_token: Option, - pub last_success_at: Option, - pub last_failure_at: Option, - pub last_error: Option, - success_freshness: Duration, -} - -impl WorkerHealth { - fn active(schedule: WorkerSchedule) -> Self { - Self { - state: WorkerRuntimeState::Starting, - consecutive_failures: 0, - completed_cycles: 0, - last_fencing_token: None, - last_success_at: None, - last_failure_at: None, - last_error: None, - success_freshness: schedule.success_freshness(), - } - } - - const fn disabled(reason: WorkerDisabledReason) -> Self { - Self { - state: WorkerRuntimeState::Disabled(reason), - consecutive_failures: 0, - completed_cycles: 0, - last_fencing_token: None, - last_success_at: None, - last_failure_at: None, - last_error: None, - success_freshness: Duration::ZERO, - } - } - - fn is_healthy(&self, now: SystemTime) -> bool { - match self.state { - WorkerRuntimeState::Disabled(_) | WorkerRuntimeState::Standby => true, - WorkerRuntimeState::AcquiringLease - | WorkerRuntimeState::Running - | WorkerRuntimeState::Idle => { - self.consecutive_failures == 0 - && self.last_success_at.is_some_and(|last_success| { - now.duration_since(last_success) - .is_ok_and(|age| age <= self.success_freshness) - }) - } - WorkerRuntimeState::Starting - | WorkerRuntimeState::BackingOff - | WorkerRuntimeState::Stopped => false, - } - } -} - -/// 可供健康检查读取的 worker 状态快照。 -#[derive(Clone)] -pub struct WorkerHealthRegistry { - states: Arc>>, -} - -impl WorkerHealthRegistry { - fn new( - registrations: &BTreeMap, - disabled: &BTreeMap, - ) -> Self { - let mut states = registrations - .values() - .map(|registration| { - ( - WorkerHealthKey::Task(registration.id.clone()), - WorkerHealth::active(registration.schedule), - ) - }) - .collect::>(); - states.extend(disabled.iter().map(|(kind, reason)| { - ( - WorkerHealthKey::Disabled(*kind), - WorkerHealth::disabled(*reason), - ) - })); - Self { - states: Arc::new(Mutex::new(states)), - } - } - - #[must_use] - pub fn snapshot(&self) -> BTreeMap { - lock_unpoisoned(&self.states).clone() - } - - #[must_use] - pub fn all_healthy(&self) -> bool { - let now = SystemTime::now(); - lock_unpoisoned(&self.states) - .values() - .all(|health| health.is_healthy(now)) - } - - fn update_task(&self, id: &WorkerId, update: impl FnOnce(&mut WorkerHealth)) { - if let Some(state) = - lock_unpoisoned(&self.states).get_mut(&WorkerHealthKey::Task(id.clone())) - { - update(state); - } - } - - fn acquiring_lease(&self, id: &WorkerId) { - self.update_task(id, |state| { - state.state = WorkerRuntimeState::AcquiringLease; - }); - } - - fn standby(&self, id: &WorkerId) { - self.update_task(id, |state| { - state.state = WorkerRuntimeState::Standby; - state.consecutive_failures = 0; - state.last_error = None; - }); - } - - fn running(&self, id: &WorkerId, fencing_token: WorkerFencingToken) { - self.update_task(id, |state| { - state.state = WorkerRuntimeState::Running; - state.last_fencing_token = Some(fencing_token); - }); - } - - fn succeeded(&self, id: &WorkerId) { - self.update_task(id, |state| { - state.state = WorkerRuntimeState::Idle; - state.consecutive_failures = 0; - state.completed_cycles = state.completed_cycles.saturating_add(1); - state.last_success_at = Some(SystemTime::now()); - state.last_error = None; - }); - } - - fn failed(&self, id: &WorkerId, error: String) -> u32 { - let mut failures = 0; - self.update_task(id, |state| { - state.state = WorkerRuntimeState::BackingOff; - state.consecutive_failures = state.consecutive_failures.saturating_add(1); - state.last_failure_at = Some(SystemTime::now()); - state.last_error = Some(error); - failures = state.consecutive_failures; - }); - failures - } - - fn stopped(&self, id: &WorkerId) { - self.update_task(id, |state| { - state.state = WorkerRuntimeState::Stopped; - }); - } -} - -/// 所有后台任务的统一 shutdown/join 句柄。 -pub struct WorkerSupervisor { - cancellation: CancellationToken, - health: WorkerHealthRegistry, - active_ids: Vec, - handles: Vec<(WorkerId, JoinHandle<()>)>, -} - -impl WorkerSupervisor { - #[must_use] - pub const fn health(&self) -> &WorkerHealthRegistry { - &self.health - } - - /// 先协作取消；超时后强制 abort，并等待 guard 完成 Drop。 - pub async fn shutdown(mut self, timeout: Duration) { - self.cancellation.cancel(); - let mut handles = std::mem::take(&mut self.handles); - let completed = tokio::time::timeout(timeout, async { - for (_, handle) in &mut handles { - let _ = handle.await; - } - }) - .await - .is_ok(); - if !completed { - for (_, handle) in &handles { - handle.abort(); - } - for (_, handle) in handles { - let _ = handle.await; - } - } - for id in &self.active_ids { - self.health.stopped(id); - } - } -} - -impl Drop for WorkerSupervisor { - fn drop(&mut self) { - self.cancellation.cancel(); - for (_, handle) in &self.handles { - handle.abort(); - } - for id in &self.active_ids { - self.health.stopped(id); - } - } -} - -fn spawn_worker( - registration: WorkerRegistration, - leader_leases: Arc, - cancellation: CancellationToken, - health: WorkerHealthRegistry, -) -> JoinHandle<()> { - tokio::spawn(async move { - supervise_worker(registration, leader_leases, cancellation, &health).await; - }) -} - -async fn supervise_worker( - registration: WorkerRegistration, - leader_leases: Arc, - cancellation: CancellationToken, - health: &WorkerHealthRegistry, -) { - let WorkerRegistration { id, task, schedule } = registration; - let mut backoff = schedule.initial_backoff; - - loop { - health.acquiring_lease(&id); - let request = WorkerLeaseRequest { - worker: id.clone(), - ttl: schedule.leader_lease_ttl, - }; - let acquisition = AssertUnwindSafe(leader_leases.try_acquire(request)).catch_unwind(); - tokio::pin!(acquisition); - let acquisition = tokio::select! { - () = cancellation.cancelled() => break, - acquisition = &mut acquisition => acquisition, - }; - - let delay = match acquisition { - Ok(Ok(WorkerLeaseAcquisition::Acquired(guard))) => { - match run_leased_cycle( - &id, - &task, - guard, - schedule.leader_lease_renewal_interval, - &cancellation, - health, - ) - .await - { - LeasedCycleOutcome::Completed(TaskCycleOutcome::Succeeded) => { - health.succeeded(&id); - backoff = schedule.initial_backoff; - schedule.interval - } - LeasedCycleOutcome::Completed(TaskCycleOutcome::Failed(error)) => { - let failures = health.failed(&id, error.as_safe_str().to_owned()); - tracing::warn!( - worker = %id, - failures, - "后台任务周期执行失败" - ); - take_backoff(&mut backoff, schedule.maximum_backoff) - } - LeasedCycleOutcome::Completed(TaskCycleOutcome::Panicked) => { - let failures = health.failed(&id, "worker panicked".to_owned()); - tracing::error!(worker = %id, failures, "后台任务发生 panic"); - take_backoff(&mut backoff, schedule.maximum_backoff) - } - LeasedCycleOutcome::LeaseLost(error) => { - let failures = health.failed(&id, error); - tracing::warn!(worker = %id, failures, "后台任务 leader lease 已失效"); - take_backoff(&mut backoff, schedule.maximum_backoff) - } - LeasedCycleOutcome::ShuttingDown => break, - } - } - Ok(Ok(WorkerLeaseAcquisition::Busy { retry_after })) => { - health.standby(&id); - retry_after - .filter(|delay| !delay.is_zero()) - .unwrap_or(schedule.interval) - .min(schedule.maximum_backoff) - } - Ok(Err(error)) => { - let failures = health.failed(&id, error.as_safe_str().to_owned()); - tracing::warn!(worker = %id, failures, "后台任务 leader lease 获取失败"); - take_backoff(&mut backoff, schedule.maximum_backoff) - } - Err(_) => { - let failures = health.failed(&id, "leader lease port panicked".to_owned()); - tracing::error!(worker = %id, failures, "后台任务 leader lease 端口发生 panic"); - take_backoff(&mut backoff, schedule.maximum_backoff) - } - }; - - tokio::select! { - () = cancellation.cancelled() => break, - () = tokio::time::sleep(delay) => {} - } - } - health.stopped(&id); -} - -enum TaskCycleOutcome { - Succeeded, - Failed(WorkerTaskError), - Panicked, -} - -enum LeasedCycleOutcome { - Completed(TaskCycleOutcome), - LeaseLost(String), - ShuttingDown, -} - -async fn run_leased_cycle( - id: &WorkerId, - task: &Arc, - mut guard: Box, - renewal_interval: Duration, - cancellation: &CancellationToken, - health: &WorkerHealthRegistry, -) -> LeasedCycleOutcome { - let fencing_token = match std::panic::catch_unwind(AssertUnwindSafe(|| guard.fencing_token())) { - Ok(fencing_token) => fencing_token, - Err(_) => return LeasedCycleOutcome::LeaseLost("leader lease guard panicked".to_owned()), - }; - health.running(id, fencing_token); - let cycle_cancellation = cancellation.child_token(); - let context = WorkerCycleContext { - worker: id.clone(), - fencing_token, - cancellation: cycle_cancellation.clone(), - }; - let cycle = run_task_cycle(task, context); - tokio::pin!(cycle); - let renewal_sleep = tokio::time::sleep(renewal_interval); - tokio::pin!(renewal_sleep); - - loop { - tokio::select! { - biased; - outcome = &mut cycle => return LeasedCycleOutcome::Completed(outcome), - () = cancellation.cancelled() => { - cycle_cancellation.cancel(); - let _ = cycle.await; - return LeasedCycleOutcome::ShuttingDown; - } - () = &mut renewal_sleep => { - let renewal = AssertUnwindSafe(guard.renew()).catch_unwind(); - tokio::pin!(renewal); - let renewal_timeout = tokio::time::sleep(renewal_interval); - tokio::pin!(renewal_timeout); - let renewal = tokio::select! { - biased; - outcome = &mut cycle => { - return LeasedCycleOutcome::Completed(outcome); - } - () = cancellation.cancelled() => { - cycle_cancellation.cancel(); - let _ = cycle.await; - return LeasedCycleOutcome::ShuttingDown; - } - renewal = &mut renewal => Some(renewal), - () = &mut renewal_timeout => None, - }; - match renewal { - Some(Ok(Ok(()))) => { - renewal_sleep - .as_mut() - .reset(tokio::time::Instant::now() + renewal_interval); - } - Some(Ok(Err(error))) => { - cycle_cancellation.cancel(); - return LeasedCycleOutcome::LeaseLost(error.as_safe_str().to_owned()); - } - Some(Err(_)) => { - cycle_cancellation.cancel(); - return LeasedCycleOutcome::LeaseLost( - "leader lease renewal panicked".to_owned(), - ); - } - None => { - cycle_cancellation.cancel(); - return LeasedCycleOutcome::LeaseLost( - "leader lease renewal timed out".to_owned(), - ); - } - } - } - } - } -} - -async fn run_task_cycle( - task: &Arc, - context: WorkerCycleContext, -) -> TaskCycleOutcome { - match AssertUnwindSafe(task.run_cycle(context)) - .catch_unwind() - .await - { - Ok(Ok(())) => TaskCycleOutcome::Succeeded, - Ok(Err(error)) => TaskCycleOutcome::Failed(error), - Err(_) => TaskCycleOutcome::Panicked, - } -} - -fn take_backoff(backoff: &mut Duration, maximum: Duration) -> Duration { - let delay = *backoff; - *backoff = backoff.checked_mul(2).unwrap_or(maximum).min(maximum); - delay -} - -fn valid_owner(owner: &str) -> bool { - !owner.is_empty() - && owner.len() <= 64 - && owner.bytes().all(|byte| { - byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"-_.".contains(&byte) - }) - && owner - .as_bytes() - .first() - .is_some_and(u8::is_ascii_alphanumeric) -} - -fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { - mutex - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} diff --git a/backend/apps/gateway/tests/bootstrap.rs b/backend/apps/gateway/tests/bootstrap.rs index 7a8649202..285e4b4c2 100644 --- a/backend/apps/gateway/tests/bootstrap.rs +++ b/backend/apps/gateway/tests/bootstrap.rs @@ -1,352 +1,73 @@ -use std::{ - fs, - io::Write, - path::{Path, PathBuf}, - sync::{Arc, Mutex}, - time::{Duration as StdDuration, SystemTime}, -}; +use std::{fs, process::Command}; -use async_trait::async_trait; -use axum::{ - body::{Body, to_bytes}, - http::{Request, StatusCode, header::CONTENT_TYPE}, -}; -use chrono::{DateTime, Duration, TimeZone, Utc}; -use codex_proxy_rs::{ - BootstrapConfig, ConfigError, InMemoryCodexOAuthPendingStore, ProcessSystemAdminService, - SystemUpdateConfig, TopologyOverrides, external_observability_range, health_timeline_view_at, - provider_failure_affects_circuit, restore_client_admission_startup, - runtime_revision_needs_refresh, -}; -use flate2::{Compression, write::GzEncoder}; -use gateway_api::admin::{ - AdminServiceError, AdminSessionResolver, AdminSessionState, - system::{self, SystemAdminErrorKind, SystemAdminService, SystemAdminState}, -}; -use gateway_core::error::ProviderErrorKind; -use gateway_store::postgres::{ObservationGranularity, RequestMetricPoint, RequestMetrics}; -use provider_openai::credential::{ - CodexOAuthPendingStore, CodexPendingAuthorization, StoredCodexPendingAuthorization, -}; -use secrecy::SecretString; -use serde_json::json; -use sha2::{Digest as _, Sha256}; -use tar::{Builder, EntryType, Header}; -use tokio_util::sync::CancellationToken; -use tower::ServiceExt as _; -use wiremock::{ - Mock, MockServer, ResponseTemplate, - matchers::{method, path}, -}; +use codex_proxy_rs::bootstrap::GatewayConfig; +use gateway_host::LoadableConfig; const CONFIG_EXAMPLE: &str = include_str!("../../../../deploy/config.example.yaml"); const POSTGRES_PASSWORD: &str = "111111111111111111111111111111111111111111111111"; const REDIS_PASSWORD: &str = "222222222222222222222222222222222222222222222222"; const ADMIN_PASSWORD: &str = "test-admin-password"; - -#[test] -fn health_timeline_should_keep_exactly_china_day_quarter_hour_slots() { - let now = utc_time(2026, 7, 19, 2, 22); - let points = vec![ - health_metric_point( - utc_time(2026, 7, 18, 15, 45), - RequestMetrics { - success_count: 100, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 19, 2, 15), - RequestMetrics { - success_count: 2, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 19, 2, 30), - RequestMetrics { - success_count: 100, - ..RequestMetrics::default() - }, - ), - ]; - - let timeline = health_timeline_view_at(&points, now); - - assert_eq!(timeline.description, "有效请求可用性"); - assert_eq!(timeline.points.len(), 96); - assert_eq!( - timeline.points.first().map(|point| point.time.as_str()), - Some("00:00") - ); - assert_eq!( - timeline.points.last().map(|point| point.time.as_str()), - Some("23:45") - ); - assert_eq!(timeline.success_requests, 2); - assert_eq!(timeline.points[41].status, "low_sample"); - assert_eq!(timeline.points[42].status, "future"); - assert_eq!(timeline.points[42].success_requests, 0); -} - -#[test] -fn health_timeline_should_match_legacy_status_precedence_and_thresholds() { - let now = utc_time(2026, 7, 19, 4, 0); - let points = vec![ - health_metric_point( - utc_time(2026, 7, 18, 16, 0), - RequestMetrics { - failure_count: 1, - cancelled_count: 1, - incomplete_count: 1, - caller_error_count: 1, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 18, 16, 15), - RequestMetrics { - failure_count: 3, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 18, 16, 30), - RequestMetrics { - success_count: 1, - failure_count: 1, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 18, 16, 45), - RequestMetrics { - success_count: 98, - failure_count: 2, - ..RequestMetrics::default() - }, - ), - health_metric_point( - utc_time(2026, 7, 18, 17, 0), - RequestMetrics { - success_count: 99, - failure_count: 1, - ..RequestMetrics::default() - }, - ), - ]; - - let timeline = health_timeline_view_at(&points, now); - let statuses = timeline - .points - .iter() - .take(5) - .map(|point| point.status.as_str()) - .collect::>(); - - assert_eq!( - statuses, - vec!["no_data", "unavailable", "low_sample", "unstable", "stable"] - ); - assert_eq!(timeline.points[0].reliability_display, "-"); - assert_eq!(timeline.points[4].reliability_display, "99.0%"); -} - -fn utc_time(year: i32, month: u32, day: u32, hour: u32, minute: u32) -> DateTime { - Utc.with_ymd_and_hms(year, month, day, hour, minute, 0) - .single() - .expect("valid UTC test time") -} - -fn health_metric_point(bucket_start: DateTime, metrics: RequestMetrics) -> RequestMetricPoint { - RequestMetricPoint { - bucket_start, - granularity: ObservationGranularity::FifteenMinutes, - metrics, - cost_coverage: gateway_store::postgres::CostCoverage::default(), - costs: Vec::new(), - } -} - -#[tokio::test] -async fn codex_pending_flow_wrong_owner_does_not_consume_it() { - let store = InMemoryCodexOAuthPendingStore::new(); - let pending = CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization { - flow_id: "flow_owner_bound".to_owned(), - owner_ref: "admin:owner".to_owned(), - started_request_ref: "request_owner_bound".to_owned(), - provider_instance_id: "provider_openai".to_owned(), - name: "Owner-bound OAuth".to_owned(), - expires_at: Utc::now() + Duration::minutes(10), - state: SecretString::from("state-owner-bound"), - nonce: SecretString::from("nonce-owner-bound"), - code_verifier: SecretString::from("verifier-owner-bound"), - reauthorization_account_id: None, - reauthorization_credential_revision: None, - }) - .expect("valid pending Codex OAuth flow"); - store.create(&pending).await.expect("store pending flow"); - - assert!( - store - .take("admin:other", "flow_owner_bound") - .await - .expect("wrong owner lookup") - .is_none() - ); - assert_eq!( - store - .take("admin:owner", "flow_owner_bound") - .await - .expect("owner lookup") - .map(|flow| flow.owner_ref().to_owned()), - Some("admin:owner".to_owned()) - ); - assert!( - store - .take("admin:owner", "flow_owner_bound") - .await - .expect("single-use lookup") - .is_none() - ); -} - -#[test] -fn external_observability_range_accepts_exactly_366_days() { - let end = Utc::now(); - let range = external_observability_range(end - Duration::days(366), end) - .expect("366-day external range should be accepted"); - assert_eq!(range.end, end); -} - -#[test] -fn external_observability_range_rejects_over_366_days_and_reversed_range() { - let end = Utc::now(); - assert!(external_observability_range(end - Duration::days(367), end).is_err()); - assert!(external_observability_range(end, end).is_err()); - assert!(external_observability_range(end + Duration::seconds(1), end).is_err()); -} - -#[test] -fn provider_circuit_should_only_count_instance_attributable_failures() { - for error_kind in [ - ProviderErrorKind::Timeout, - ProviderErrorKind::Transport, - ProviderErrorKind::Protocol, - ProviderErrorKind::Unavailable, - ] { - assert!(provider_failure_affects_circuit(error_kind)); - } - for error_kind in [ - ProviderErrorKind::InvalidRequest, - ProviderErrorKind::Unsupported, - ProviderErrorKind::Unauthorized, - ProviderErrorKind::PermissionDenied, - ProviderErrorKind::RateLimited, - ProviderErrorKind::QuotaExhausted, - ProviderErrorKind::Cancelled, - ProviderErrorKind::ProcessTerminated, - ] { - assert!(!provider_failure_affects_circuit(error_kind)); - } -} - -#[test] -fn runtime_revision_reconciliation_should_refresh_missing_or_stale_snapshot() { - assert!(!runtime_revision_needs_refresh(Some(7), 7)); - assert!(runtime_revision_needs_refresh(Some(6), 7)); - assert!(runtime_revision_needs_refresh(Some(8), 7)); - assert!(runtime_revision_needs_refresh(None, 7)); -} +const TOPOLOGY_CHILD_ENV: &str = "CPR_TEST_TOPOLOGY_CHILD"; #[test] fn config_loader_should_load_complete_terminal_example() { - let directory = write_config(&complete_config()); - let loaded = BootstrapConfig::load_from_path(directory.path().join("config.yaml")) - .expect("terminal configuration should load"); - - assert_eq!( - ( - loaded.app().server.host.as_str(), - loaded.app().server.port, - loaded.app().wire_profile.codex_version.as_str(), - loaded.app().wire_profile.desktop_version.as_str(), - loaded.app().wire_profile.desktop_build.as_str(), - loaded.app().admin.default_username.as_str(), - ), - ( - "127.0.0.1", - 8080, - "0.144.2", - "26.707.72221", - "5307", - "admin@cpr.local", - ) - ); + parse_config(&valid_config()).expect("terminal config example"); } #[test] fn config_loader_should_resolve_paths_relative_to_config_file() { - let directory = write_config(&complete_config()); - let loaded = BootstrapConfig::load_from_path(directory.path().join("config.yaml")) - .expect("terminal configuration should load"); - - assert_eq!( - loaded.app().logging.file.directory.as_path(), - directory.path().join("../.runtime/logs").as_path() - ); + let (config, _directory) = parse_config(&valid_config()).expect("resolved config"); + let debug = format!("{config:?}"); + assert!(debug.contains(".runtime/logs")); + assert!(debug.contains("frontend/dist")); } #[test] -fn config_loader_should_inject_service_passwords_into_connection_urls() { - let directory = write_config(&complete_config()); - let loaded = BootstrapConfig::load_from_path(directory.path().join("config.yaml")) - .expect("terminal configuration should load"); - - assert_eq!( - (loaded.database_url(), loaded.redis_url()), - ( - "postgres://codex_proxy:111111111111111111111111111111111111111111111111@127.0.0.1:5432/codex_proxy", - "redis://:222222222222222222222222222222222222222222222222@127.0.0.1:6379/", - ) - ); +fn config_loader_should_inject_connection_passwords_into_urls() { + parse_config(&valid_config()).expect("Store validates password injection into both URLs"); } #[test] fn config_loader_should_apply_only_explicit_topology_overrides() { - let directory = write_config(&complete_config()); - let overrides = TopologyOverrides { - server_host: Some("0.0.0.0".to_string()), - server_port: Some(18080), - database_url: Some("postgres://codex_proxy@postgres:5432/codex_proxy".to_string()), - redis_url: Some("redis://redis:6379/".to_string()), - }; - - let loaded = BootstrapConfig::load_from_path_with_overrides( - directory.path().join("config.yaml"), - overrides, - ) - .expect("explicit topology overrides should load"); - - assert_eq!(loaded.app().server.host, "0.0.0.0"); - assert_eq!(loaded.app().server.port, 18080); - assert_eq!( - loaded.database_url(), - "postgres://codex_proxy:111111111111111111111111111111111111111111111111@postgres:5432/codex_proxy" - ); - assert_eq!( - loaded.redis_url(), - "redis://:222222222222222222222222222222222222222222222222@redis:6379/" - ); + let invalid = valid_config() + .replace("host: '127.0.0.1'", "host: ''") + .replace("port: 8080", "port: 0") + .replace( + "url: 'postgres://codex_proxy@127.0.0.1:5432/codex_proxy'", + "url: 'invalid-postgres-url'", + ) + .replace("url: 'redis://127.0.0.1:6379/'", "url: 'invalid-redis-url'") + .replace(POSTGRES_PASSWORD, "invalid-postgres-password") + .replace(REDIS_PASSWORD, "invalid-redis-password"); + if std::env::var_os(TOPOLOGY_CHILD_ENV).is_some() { + parse_config(&invalid).expect("explicit package-owned environment overrides"); + return; + } + assert!(parse_config(&invalid).is_err()); + let status = Command::new(std::env::current_exe().expect("current test executable")) + .args([ + "--exact", + "bootstrap::config_loader_should_apply_only_explicit_topology_overrides", + ]) + .env(TOPOLOGY_CHILD_ENV, "1") + .env("CPR_SERVER_HOST", "127.0.0.1") + .env("CPR_SERVER_PORT", "8080") + .env( + "CPR_DATABASE_URL", + "postgres://codex_proxy@127.0.0.1:5432/codex_proxy", + ) + .env("CPR_REDIS_URL", "redis://127.0.0.1:6379/") + .env("CPR_DATABASE_PASSWORD", POSTGRES_PASSWORD) + .env("CPR_REDIS_PASSWORD", REDIS_PASSWORD) + .status() + .expect("run isolated environment override test"); + assert!(status.success()); } #[test] fn bootstrap_config_debug_should_redact_all_passwords() { - let directory = write_config(&complete_config()); - let loaded = BootstrapConfig::load_from_path(directory.path().join("config.yaml")) - .expect("terminal configuration should load"); - let debug = format!("{loaded:?}"); - + let (config, _directory) = parse_config(&valid_config()).expect("config"); + let debug = format!("{config:?}"); assert!(!debug.contains(POSTGRES_PASSWORD)); assert!(!debug.contains(REDIS_PASSWORD)); assert!(!debug.contains(ADMIN_PASSWORD)); @@ -355,1332 +76,118 @@ fn bootstrap_config_debug_should_redact_all_passwords() { #[test] fn config_loader_should_reject_unknown_fields() { - let yaml = complete_config().replace(" port: 8080", " port: 8080\n unexpected: true"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidDocument { .. }) + assert_rejected(format!( + "{}\nunknown_terminal_field: true\n", + valid_config() )); } #[test] fn config_loader_should_reject_removed_tls_section() { - let yaml = complete_config().replace( - " wire_profile:", - " tls:\n force_http11: false\n\n wire_profile:", - ); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidDocument { .. }) - )); + assert_rejected(valid_config().replace("openai:\n", "openai:\n tls: {}\n")); } #[test] fn config_loader_should_reject_missing_explicit_fields() { - let yaml = complete_config().replacen(" logging:\n", "", 1); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidDocument { .. }) - )); + assert_rejected(valid_config().replace(" request_id_header: 'x-request-id'\n", "")); } #[test] fn config_loader_should_reject_unsupported_schema_version() { - let yaml = complete_config().replace("schema_version: 1", "schema_version: 0"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::UnsupportedSchemaVersion) - )); + assert_rejected(valid_config().replace("schema_version: 1", "schema_version: 2")); } #[test] fn config_loader_should_reject_embedded_database_password() { - let yaml = complete_config().replace( - "postgres://codex_proxy@127.0.0.1", - "postgres://codex_proxy:duplicated@127.0.0.1", - ); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::PasswordInUrl("database.url")) + assert_rejected(valid_config().replace( + "postgres://codex_proxy@127.0.0.1:5432/codex_proxy", + "postgres://codex_proxy:embedded@127.0.0.1:5432/codex_proxy", )); } #[test] fn config_loader_should_reject_non_hex_postgres_password() { - let yaml = complete_config().replace(POSTGRES_PASSWORD, "not-a-hex-password"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidServicePassword("database.password")) - )); + assert_rejected(valid_config().replace(POSTGRES_PASSWORD, &"g".repeat(48))); } #[test] fn config_loader_should_reject_wrong_length_redis_password() { - let yaml = complete_config().replace(REDIS_PASSWORD, "2222222222222222"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidServicePassword("redis.password")) - )); + assert_rejected(valid_config().replace(REDIS_PASSWORD, "1234")); } #[test] fn config_loader_should_reject_weak_admin_password() { - let yaml = complete_config().replace(ADMIN_PASSWORD, "password"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::WeakAdminPassword) - )); + assert_rejected(valid_config().replace(ADMIN_PASSWORD, "password")); } #[test] fn config_loader_should_reject_admin_password_with_compose_interpolation() { - let yaml = complete_config().replace(ADMIN_PASSWORD, "strong$password-value"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::WeakAdminPassword) - )); + assert_rejected(valid_config().replace(ADMIN_PASSWORD, "unsafe$password")); } #[test] fn config_loader_should_reject_removed_fingerprint_section() { - let yaml = complete_config().replace(" wire_profile:", " fingerprint:"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidDocument { .. }) + assert_rejected(valid_config().replace( + "openai:\n", + "openai:\n fingerprint:\n browser: removed\n", )); } #[test] fn config_loader_should_reject_invalid_codex_core_version() { - let yaml = complete_config().replace("codex_version: '0.144.2'", "codex_version: 'Desktop'"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidField("wire_profile.codex_version")) - )); + assert_rejected(valid_config().replace("codex_version: '0.144.2'", "codex_version: 'latest'")); } #[test] fn config_loader_should_reject_disabled_all_log_outputs() { - let yaml = complete_config() - .replace(" stdout: true", " stdout: false") - .replace(" enabled: true", " enabled: false"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidField("logging")) - )); + assert_rejected( + valid_config() + .replace("stdout: true", "stdout: false") + .replace("enabled: true", "enabled: false"), + ); } #[test] fn config_loader_should_reject_zero_server_port() { - let yaml = complete_config().replace(" port: 8080", " port: 0"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidField("server.port")) - )); + assert_rejected(valid_config().replace("port: 8080", "port: 0")); } #[test] fn config_loader_should_reject_invalid_desktop_profile_fields() { - let yaml = complete_config().replace("desktop_build: '5307'", "desktop_build: 'build'"); - assert!(matches!( - load_yaml(&yaml), - Err(ConfigError::InvalidField("wire_profile.desktop_build")) - )); + assert_rejected(valid_config().replace("desktop_build: '5307'", "desktop_build: 'build'")); } -fn complete_config() -> String { +fn assert_rejected(config: String) { + assert!(parse_config(&config).is_err()); +} + +fn valid_config() -> String { CONFIG_EXAMPLE - .replace( - "password: &postgres_password ''", - &format!("password: &postgres_password '{POSTGRES_PASSWORD}'"), - ) - .replace( - "password: &redis_password ''", - &format!("password: &redis_password '{REDIS_PASSWORD}'"), + .replacen( + "password: ''", + &format!("password: '{POSTGRES_PASSWORD}'"), + 1, ) + .replacen("password: ''", &format!("password: '{REDIS_PASSWORD}'"), 1) .replace( "default_password: ''", &format!("default_password: '{ADMIN_PASSWORD}'"), ) } -fn load_yaml(yaml: &str) -> Result { - let directory = write_config(yaml); - BootstrapConfig::load_from_path(directory.path().join("config.yaml")) -} - -fn write_config(yaml: &str) -> tempfile::TempDir { - let directory = tempfile::tempdir().expect("temporary configuration directory"); - fs::write(directory.path().join("config.yaml"), yaml).expect("write test configuration"); - directory -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum StartupFailure { - Recover, - Load, - Restore, -} - -#[derive(Default)] -struct StartupRecoveryState { - calls: Mutex>, - restored: Mutex>, - window_started_at: Mutex>>, -} - -struct StartupModelRequests { - state: Arc, - failure: Option, -} - -#[async_trait] -impl gateway_store::postgres::ModelRequestRepository for StartupModelRequests { - async fn insert_model_request( - &self, - _request: gateway_store::postgres::NewModelRequest, - ) -> gateway_store::StoreResult<()> { - unreachable!("startup recovery never inserts requests") - } - - async fn begin_model_request_attempt( - &self, - _attempt: gateway_store::postgres::ModelRequestAttemptStart, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never starts attempts") - } - - async fn mark_upstream_send_state( - &self, - _model_request_id: &str, - _state: gateway_store::postgres::UpstreamSendState, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never changes send state") - } - - async fn mark_downstream_committed( - &self, - _model_request_id: &str, - _committed_at: chrono::DateTime, - _client_status_code: Option, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never commits downstream delivery") - } - - async fn record_client_status_code( - &self, - _model_request_id: &str, - _client_status_code: u16, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never records client status") - } - - async fn finalize_model_request( - &self, - _finalization: gateway_store::postgres::ModelRequestFinalization, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never finalizes live requests") - } - - async fn recover_expired_model_requests( - &self, - _now: chrono::DateTime, - ) -> gateway_store::StoreResult { - self.state.calls.lock().expect("calls lock").push("recover"); - if self.failure == Some(StartupFailure::Recover) { - return Err(startup_store_error(gateway_store::StoreBackend::PostgreSql)); - } - Ok(gateway_store::postgres::ModelRequestRecoveryReport { requests: 2 }) - } -} - -struct StartupRecoveryRepository { - state: Arc, - failure: Option, - now: chrono::DateTime, -} - -#[async_trait] -impl gateway_store::postgres::ClientAdmissionRecoveryRepository for StartupRecoveryRepository { - async fn load_client_admission_recovery( - &self, - window_started_at: chrono::DateTime, - ) -> gateway_store::StoreResult> { - self.state.calls.lock().expect("calls lock").push("load"); - *self.state.window_started_at.lock().expect("window lock") = Some(window_started_at); - if self.failure == Some(StartupFailure::Load) { - return Err(startup_store_error(gateway_store::StoreBackend::PostgreSql)); - } - Ok(vec![gateway_store::postgres::ClientAdmissionRecovery { - client_api_key_ref: "key_client_1".to_owned(), - recent_requests: vec![gateway_store::postgres::ClientAdmissionRecentRequest { - model_request_id: "req_recent_1".to_owned(), - started_at: self.now - Duration::seconds(30), - input_token_estimate: 12, - }], - running_requests: vec![gateway_store::postgres::ClientAdmissionRunningRequest { - model_request_id: "req_running_1".to_owned(), - deadline_at: self.now + Duration::seconds(60), - }], - }]) - } -} - -struct StartupAdmissions { - state: Arc, - failure: Option, -} - -#[async_trait] -impl gateway_store::redis::ClientAdmissionRepository for StartupAdmissions { - async fn admit_client_request( - &self, - _request: &gateway_store::redis::ClientAdmissionRequest, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never admits a new request") - } - - async fn release_client_request( - &self, - _client_api_key_ref: &str, - _model_request_id: &str, - ) -> gateway_store::StoreResult { - unreachable!("startup recovery never releases a request") - } - - async fn restore_client_admission( - &self, - recovery: &gateway_store::redis::ClientAdmissionRestore, - ) -> gateway_store::StoreResult { - self.state.calls.lock().expect("calls lock").push("restore"); - self.state - .restored - .lock() - .expect("restored lock") - .push(recovery.clone()); - if self.failure == Some(StartupFailure::Restore) { - return Err(startup_store_error(gateway_store::StoreBackend::Redis)); - } - Ok(gateway_store::redis::ClientAdmissionRestoreResult { - restored_recent_requests: 1, - restored_running_requests: 1, - }) - } - - async fn clear_client_admission( - &self, - _client_api_key_ref: &str, - ) -> gateway_store::StoreResult<()> { - unreachable!("startup recovery never clears client state") - } -} - -#[tokio::test] -async fn client_admission_startup_recovery_should_preserve_order_and_exact_facts() { - let now = Utc::now(); - let state = Arc::new(StartupRecoveryState::default()); - let model_requests = StartupModelRequests { - state: Arc::clone(&state), - failure: None, - }; - let recovery = StartupRecoveryRepository { - state: Arc::clone(&state), - failure: None, - now, - }; - let admissions = StartupAdmissions { - state: Arc::clone(&state), - failure: None, - }; - - let report = restore_client_admission_startup(&model_requests, &recovery, &admissions, now) - .await - .expect("startup recovery"); - - assert_eq!( - *state.calls.lock().expect("calls lock"), - ["recover", "load", "restore"] - ); - assert_eq!( - *state.window_started_at.lock().expect("window lock"), - Some(now - Duration::seconds(61)) - ); - let restored = state.restored.lock().expect("restored lock"); - assert_eq!(restored[0].recent_requests[0].input_token_estimate, 12); - assert_eq!( - restored[0].running_requests[0].expires_at, - now + Duration::seconds(60) - ); - assert_eq!(report.expired_model_requests, 2); - assert_eq!(report.restored_clients, 1); - assert_eq!(report.restored_recent_requests, 1); - assert_eq!(report.restored_running_requests, 1); -} - -#[tokio::test] -async fn client_admission_startup_recovery_should_fail_closed_at_each_boundary() { - let now = Utc::now(); - for (failure, expected_calls) in [ - (StartupFailure::Recover, vec!["recover"]), - (StartupFailure::Load, vec!["recover", "load"]), - (StartupFailure::Restore, vec!["recover", "load", "restore"]), - ] { - let state = Arc::new(StartupRecoveryState::default()); - let result = restore_client_admission_startup( - &StartupModelRequests { - state: Arc::clone(&state), - failure: Some(failure), - }, - &StartupRecoveryRepository { - state: Arc::clone(&state), - failure: Some(failure), - now, - }, - &StartupAdmissions { - state: Arc::clone(&state), - failure: Some(failure), - }, - now, - ) - .await; - - assert!(result.is_err()); - assert_eq!(*state.calls.lock().expect("calls lock"), expected_calls); - } -} - -fn startup_store_error(backend: gateway_store::StoreBackend) -> gateway_store::StoreError { - gateway_store::StoreError::Unavailable { - backend, - message: "startup recovery unavailable".to_owned(), - } -} - -const SYSTEM_VERSION_OLD: &str = "0.1.0"; -const SYSTEM_VERSION_NEW: &str = "0.2.0"; - -#[tokio::test] -async fn version_should_return_backend_build_metadata() { - let repository = "zyycn/codex-proxy-rs-version-route"; - let github = MockServer::start().await; - mount_system_latest_release(&github, repository, SYSTEM_VERSION_NEW).await; - let mut config = system_update_test_config(repository, &github.uri()); - config.git_sha = "version-test-sha".to_owned(); - config.build_time = "2026-07-01T00:00:00Z".to_owned(); - let (service, _) = system_update_test_service(config); - - let value = service.version().await.expect("version response"); - - assert!( - value["version"] == SYSTEM_VERSION_OLD - && value["gitSha"] == "version-test-sha" - && value["buildTime"] == "2026-07-01T00:00:00Z" - && value["deploymentMode"] == "docker" - && value["deploymentModeLabel"] == "Docker" - && value["updateChannel"] == "stable" - && value["latestVersion"] == SYSTEM_VERSION_NEW - && value["hasUpdate"] == true - && value["updateCached"] == false - && value["updateWarning"].is_null() - ); -} - -#[tokio::test] -async fn update_detail_should_use_cached_release_when_not_refreshed() { - let repository = "zyycn/codex-proxy-rs-cache-route"; - let github = MockServer::start().await; - mount_system_latest_release(&github, repository, "0.2.0").await; - let (service, _) = - system_update_test_service(system_update_test_config(repository, &github.uri())); - let initial = service.update_detail(true).await.expect("initial release"); - - github.reset().await; - let cached = service.update_detail(false).await.expect("cached release"); - - assert!( - initial["hasUpdate"] == true - && cached["cached"] == true - && cached["latestVersion"] == "0.2.0" - ); -} - -#[tokio::test] -async fn update_detail_should_fallback_to_cache_when_refresh_fetch_fails() { - let repository = "zyycn/codex-proxy-rs-cache-forced-route"; - let github = MockServer::start().await; - mount_system_latest_release(&github, repository, "0.3.0").await; - let (service, _) = - system_update_test_service(system_update_test_config(repository, &github.uri())); - let initial = service.update_detail(true).await.expect("initial release"); - - github.reset().await; - Mock::given(method("GET")) - .and(path(format!("/repos/{repository}/releases/latest"))) - .respond_with(ResponseTemplate::new(500)) - .mount(&github) - .await; - let fallback = service.update_detail(true).await.expect("cached fallback"); - - assert!( - initial["hasUpdate"] == true - && fallback["cached"] == true - && fallback["latestVersion"] == "0.3.0" - && fallback["warning"].is_null() - ); -} - -#[tokio::test] -async fn update_detail_should_report_source_build_as_unsupported() { - let repository = "zyycn/codex-proxy-rs-source-route"; - let github = MockServer::start().await; - mount_system_latest_release(&github, repository, SYSTEM_VERSION_NEW).await; - let mut config = system_update_test_config(repository, &github.uri()); - config.build_type = "source".to_owned(); - let (service, _) = system_update_test_service(config); - - let info = service.update_detail(true).await.expect("update detail"); - - assert!( - info["currentVersion"] == SYSTEM_VERSION_OLD - && info["latestVersion"] == SYSTEM_VERSION_NEW - && info["hasUpdate"] == true - && info["buildType"] == "source" - && info["buildTypeLabel"] == "源码构建" - && info["updateSupported"] == false - && info["unsupportedReason"] == "一键更新需要正式构建包" - ); -} - -#[tokio::test] -async fn update_detail_should_reject_untrusted_github_api_base() { - let config = system_update_test_config( - "zyycn/codex-proxy-rs-untrusted-api-base-route", - "https://mirror.example", - ); - let (service, _) = system_update_test_service(config); - - let info = service.update_detail(true).await.expect("update detail"); - - assert!( - info["hasUpdate"] == false - && info["updateSupported"] == false - && info["unsupportedReason"] == "GitHub API base must be https://api.github.com/repos" - ); -} - -#[tokio::test] -async fn restart_should_request_process_restart_inside_docker() { - let mut config = - system_update_test_config("zyycn/codex-proxy-rs-restart-route", "http://127.0.0.1:9"); - config.self_restart_enabled = true; - let (service, shutdown) = system_update_test_service(config); - - let value = service.restart().await.expect("schedule docker restart"); - let cancelled = tokio::time::timeout(StdDuration::from_secs(2), shutdown.cancelled()).await; - - assert!( - value["operationId"] - .as_str() - .is_some_and(|id| id.starts_with("sysop-restart-")) - && value["message"] == "已安排进程内重启" - && cancelled.is_ok() - ); -} - -#[tokio::test] -async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { - let replacement = Path::new("/bin/true"); - if !replacement.exists() { - return; - } - let mut config = system_update_test_config( - "zyycn/codex-proxy-rs-binary-restart-route", - "http://127.0.0.1:9", - ); - config.deployment_mode = "binary".to_owned(); - config.self_restart_enabled = true; - config.executable_path = Some(replacement.to_path_buf()); - let (service, shutdown) = system_update_test_service(config); - - let value = service.restart().await.expect("schedule binary restart"); - let cancelled = tokio::time::timeout(StdDuration::from_secs(2), shutdown.cancelled()).await; - - assert!(value["message"] == "已安排自重启" && cancelled.is_ok()); -} - -#[tokio::test] -async fn restart_should_not_shutdown_when_replacement_spawn_fails() { - let root = tempfile::tempdir().expect("restart root"); - let mut config = system_update_test_config( - "zyycn/codex-proxy-rs-restart-spawn-failure-route", - "http://127.0.0.1:9", - ); - config.deployment_mode = "binary".to_owned(); - config.self_restart_enabled = true; - config.executable_path = Some(root.path().join("missing-binary")); - let (service, shutdown) = system_update_test_service(config); - - let error = service.restart().await.expect_err("replacement must fail"); - let cancelled = tokio::time::timeout(StdDuration::from_millis(100), shutdown.cancelled()).await; - - assert!( - error.kind() == SystemAdminErrorKind::Internal - && error - .to_string() - .starts_with("Failed to schedule replacement process:") - && cancelled.is_err() - ); -} - -#[tokio::test] -async fn update_should_replace_local_release_files_with_latest_asset() { - let repository = "zyycn/codex-proxy-rs-update-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - - let value = service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect("perform update"); - let status = service.update_status().await.expect("update status"); - - assert!( - value["operationId"] - .as_str() - .is_some_and(|id| id.starts_with("sysop-update-")) - && value["targetVersion"] == "0.4.0" - && value["needRestart"] == true - && status["operation"]["status"] == "succeeded" - && status["currentVersion"] == "0.4.0" - && fs::read_to_string(&layout.executable).expect("new binary") == "new-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("new web") - == "new-web" - && fs::read_to_string(layout.root.path().join("codex-proxy-rs.backup")) - .expect("binary backup") - == "old-binary" - && fs::read_to_string(layout.root.path().join("web/dist.backup/index.html")) - .expect("web backup") - == "old-web" - ); -} - -#[tokio::test] -async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { - let repository = "zyycn/codex-proxy-rs-update-target-changed-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_latest_release(&github, repository, "0.4.0").await; - let (service, _) = system_update_test_service(layout.config.clone()); - let initial = service.update_detail(true).await.expect("initial release"); - github.reset().await; - mount_system_latest_release(&github, repository, "0.5.0").await; - - let error = service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect_err("changed target must fail"); - - assert!( - initial["latestVersion"] == "0.4.0" - && error.kind() == SystemAdminErrorKind::Conflict - && error.to_string() == "远端最新版本已变更为 v0.5.0，请重新检查并确认" - && fs::read_to_string(&layout.executable).expect("old binary") == "old-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("old web") - == "old-web" - ); -} - -#[tokio::test] -async fn update_should_reject_untrusted_github_api_base() { - let layout = SystemUpdateTestLayout::new( - "zyycn/codex-proxy-rs-update-untrusted-api-base-route", - "https://mirror.example", - ); - let (service, _) = system_update_test_service(layout.config.clone()); - - let error = service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect_err("untrusted API base must fail"); - - assert!( - error.kind() == SystemAdminErrorKind::Conflict - && error.to_string() == "GitHub API base must be https://api.github.com/repos" - && fs::read_to_string(&layout.executable).expect("old binary") == "old-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("old web") - == "old-web" - ); -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_is_missing() { - let repository = "zyycn/codex-proxy-rs-missing-checksum-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_archive( - &github, - repository, - "0.4.0", - system_release_archive("new-binary", "new-web"), - None, - ) - .await; - let (service, _) = system_update_test_service(layout.config.clone()); - - assert_system_update_failure_preserves_files( - service.as_ref(), - SystemAdminErrorKind::BadGateway, - "Release checksums.txt is required", - &layout, - ) - .await; -} - -#[tokio::test] -async fn update_should_fail_when_release_checksum_mismatches() { - let repository = "zyycn/codex-proxy-rs-bad-checksum-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - let archive = system_release_archive("new-binary", "new-web"); - let checksum = format!("{} {}\n", "0".repeat(64), system_archive_name("0.4.0")); - mount_system_release_archive(&github, repository, "0.4.0", archive, Some(checksum)).await; - let (service, _) = system_update_test_service(layout.config.clone()); - - assert_system_update_failure_preserves_files( - service.as_ref(), - SystemAdminErrorKind::BadGateway, - "Checksum mismatch", - &layout, - ) - .await; -} - -#[tokio::test] -async fn update_should_reject_release_archive_from_untrusted_host() { - let repository = "zyycn/codex-proxy-rs-untrusted-download-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_with_urls( - &github, - repository, - "0.4.0", - "https://evil.example/download/codex-proxy-rs.tar.gz", - "https://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/checksums.txt", - ) - .await; - let (service, _) = system_update_test_service(layout.config.clone()); - - assert_system_update_failure_preserves_files( - service.as_ref(), - SystemAdminErrorKind::Invalid, - "Download host is not allowed: evil.example", - &layout, - ) - .await; -} - -#[tokio::test] -async fn update_should_reject_insecure_release_archive() { - let repository = "zyycn/codex-proxy-rs-insecure-download-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_with_urls( - &github, - repository, - "0.4.0", - "http://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/archive.tar.gz", - "https://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/checksums.txt", - ) - .await; - let (service, _) = system_update_test_service(layout.config.clone()); - - assert_system_update_failure_preserves_files( - service.as_ref(), - SystemAdminErrorKind::Invalid, - "Only HTTPS release downloads are allowed", - &layout, - ) - .await; -} - -#[tokio::test] -async fn update_should_reject_release_archive_with_unsafe_path() { - let repository = "zyycn/codex-proxy-rs-unsafe-archive-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - let archive = - system_release_archive_with_extra("new-binary", "new-web", "../escape", b"unsafe"); - let checksum = system_archive_checksum("0.4.0", &archive); - mount_system_release_archive(&github, repository, "0.4.0", archive, Some(checksum)).await; - let (service, _) = system_update_test_service(layout.config.clone()); - - assert_system_update_failure_preserves_files( - service.as_ref(), - SystemAdminErrorKind::Invalid, - "Unsafe archive path", - &layout, - ) - .await; -} - -#[tokio::test] -async fn update_should_restore_web_assets_when_binary_backup_fails() { - let repository = "zyycn/codex-proxy-rs-binary-backup-failure-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - fs::create_dir(layout.root.path().join("codex-proxy-rs.backup")) - .expect("blocking backup directory"); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - - let error = service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect_err("binary backup must fail"); - let status = service.update_status().await.expect("failed status"); - - assert!( - error.kind() == SystemAdminErrorKind::Internal - && error - .to_string() - .starts_with("Failed to remove old binary backup:") - && status["operation"]["error"] - .as_str() - .is_some_and(|message| message.starts_with("Failed to remove old binary backup:")) - && fs::read_to_string(&layout.executable).expect("old binary") == "old-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("old web") - == "old-web" - && !layout.root.path().join("web/dist.backup").exists() - ); -} - -#[tokio::test] -async fn update_should_remove_stale_file_lock_and_continue() { - let repository = "zyycn/codex-proxy-rs-stale-lock-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - fs::create_dir_all( - layout - .config - .update_lock_file - .parent() - .expect("lock parent"), - ) - .expect("create lock parent"); - fs::write(&layout.config.update_lock_file, "pid=1\ncreated_at=stale\n") - .expect("write stale lock"); - fs::File::options() - .write(true) - .open(&layout.config.update_lock_file) - .expect("open stale lock") - .set_modified(SystemTime::now() - StdDuration::from_secs(31 * 60)) - .expect("age stale lock"); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - - service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect("update past stale lock"); - - assert!( - fs::read_to_string(&layout.executable).expect("new binary") == "new-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("new web") - == "new-web" - && !layout.config.update_lock_file.exists() - ); -} - -#[tokio::test] -async fn update_should_replace_web_assets_across_filesystems() { - let Ok(web_root) = tempfile::tempdir_in("/dev/shm") else { - return; - }; - let repository = "zyycn/codex-proxy-rs-cross-device-update-route"; - let github = MockServer::start().await; - let mut layout = SystemUpdateTestLayout::new(repository, &github.uri()); - let web_dist = web_root.path().join("web/dist"); - fs::create_dir_all(&web_dist).expect("cross-device web directory"); - fs::write(web_dist.join("index.html"), "old-web").expect("cross-device old web"); - layout.config.web_dist_dir = web_dist.clone(); - layout.web_dist = web_dist.clone(); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - - service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect("cross-device update"); - - assert!( - fs::read_to_string(&layout.executable).expect("new binary") == "new-binary" - && fs::read_to_string(web_dist.join("index.html")).expect("new web") == "new-web" - && fs::read_to_string(web_root.path().join("web/dist.backup/index.html")) - .expect("old web backup") - == "old-web" - ); -} - -#[tokio::test] -async fn update_status_should_read_local_update_state() { - let layout = - SystemUpdateTestLayout::new("zyycn/codex-proxy-rs-status-route", "http://127.0.0.1:9"); - fs::create_dir_all( - layout - .config - .update_state_file - .parent() - .expect("state parent"), - ) - .expect("create state parent"); - fs::write( - &layout.config.update_state_file, - json!({ - "previousVersion": SYSTEM_VERSION_OLD, - "currentVersion": SYSTEM_VERSION_NEW, - "operation": { - "operationId": "system-update-status", - "kind": "update", - "status": "running", - "targetVersion": SYSTEM_VERSION_NEW, - "message": "operation running", - "error": null, - "startedAt": "2026-07-01T00:00:00Z", - "finishedAt": null - } - }) - .to_string(), - ) - .expect("write update state"); - let (service, _) = system_update_test_service(layout.config.clone()); - - let status = service.update_status().await.expect("update status"); - - assert!( - status["operation"]["operationId"] == "system-update-status" - && status["operation"]["status"] == "running" - && status["currentVersion"] == SYSTEM_VERSION_NEW - ); -} - -#[tokio::test] -async fn update_events_should_open_authenticated_sse_stream() { - let layout = - SystemUpdateTestLayout::new("zyycn/codex-proxy-rs-events-route", "http://127.0.0.1:9"); - let (service, _) = system_update_test_service(layout.config.clone()); - let app = system_update_test_router(service); - - let unauthorized = app - .clone() - .oneshot( - Request::builder() - .uri("/api/admin/system/update-events") - .header("x-request-id", "req-system-events-denied") - .body(Body::empty()) - .expect("unauthorized request"), - ) - .await - .expect("unauthorized response"); - let response = app - .oneshot(system_update_authenticated_request( - "GET", - "/api/admin/system/update-events", - Body::empty(), - )) - .await - .expect("SSE response"); - let content_type = response - .headers() - .get(CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .unwrap_or_default(); - - assert!( - unauthorized.status() == StatusCode::UNAUTHORIZED - && response.status() == StatusCode::OK - && content_type.starts_with("text/event-stream") - ); -} - -#[tokio::test] -async fn update_events_should_close_after_terminal_update_log() { - let repository = "zyycn/codex-proxy-rs-events-terminal-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - let app = system_update_test_router(Arc::clone(&service)); - let response = app - .oneshot(system_update_authenticated_request( - "GET", - "/api/admin/system/update-events", - Body::empty(), - )) - .await - .expect("SSE response"); - let update_service = Arc::clone(&service); - let update = tokio::spawn(async move { - update_service - .perform_update(Some("0.4.0".to_owned())) - .await - }); - - let body = tokio::time::timeout( - StdDuration::from_secs(2), - to_bytes(response.into_body(), 128 * 1024), - ) - .await - .expect("terminal event should close SSE") - .expect("read SSE body"); - update - .await - .expect("join update") - .expect("perform terminal update"); - let text = String::from_utf8(body.to_vec()).expect("UTF-8 SSE"); - - assert!( - text.contains("\"step\":\"done\"") - && text.contains("\"terminal\":true") - && text.contains("更新文件已替换，等待服务重启生效") - ); -} - -#[tokio::test] -async fn rollback_should_restore_binary_web_and_version_state() { - let repository = "zyycn/codex-proxy-rs-rollback-route"; - let github = MockServer::start().await; - let layout = SystemUpdateTestLayout::new(repository, &github.uri()); - mount_system_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; - let (service, _) = system_update_test_service(layout.config.clone()); - service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect("perform update before rollback"); - - let value = service.rollback().await.expect("rollback release"); - let status = service.update_status().await.expect("rollback status"); - - assert!( - value["needRestart"] == true - && value["operationId"] - .as_str() - .is_some_and(|id| id.starts_with("sysop-rollback-")) - && status["operation"]["status"] == "succeeded" - && status["currentVersion"] == SYSTEM_VERSION_OLD - && status["previousVersion"] == "0.4.0" - && fs::read_to_string(&layout.executable).expect("rolled-back binary") == "old-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("rolled-back web") - == "old-web" - && fs::read_to_string(layout.root.path().join("codex-proxy-rs.backup")) - .expect("new binary backup") - == "new-binary" - ); -} - -struct SystemUpdateTestLayout { - root: tempfile::TempDir, - executable: PathBuf, - web_dist: PathBuf, - config: SystemUpdateConfig, -} - -impl SystemUpdateTestLayout { - fn new(repository: &str, github_api_base: &str) -> Self { - let root = tempfile::tempdir().expect("system update root"); - let executable = root.path().join("codex-proxy-rs"); - let web_dist = root.path().join("web/dist"); - fs::create_dir_all(&web_dist).expect("create old web directory"); - fs::write(&executable, "old-binary").expect("write old binary"); - fs::write(web_dist.join("index.html"), "old-web").expect("write old web"); - let mut config = system_update_test_config(repository, github_api_base); - config.executable_path = Some(executable.clone()); - config.web_dist_dir = web_dist.clone(); - config.update_state_file = root.path().join(".runtime/update-state.json"); - config.update_lock_file = root.path().join(".runtime/update.lock"); - config.update_temp_dir = root.path().join(".runtime/update-tmp"); - Self { - root, - executable, - web_dist, - config, - } - } -} - -fn system_update_test_config(repository: &str, github_api_base: &str) -> SystemUpdateConfig { - SystemUpdateConfig { - version: SYSTEM_VERSION_OLD.to_owned(), - git_sha: "test-git-sha".to_owned(), - build_time: "2026-07-01T00:00:00Z".to_owned(), - deployment_mode: "docker".to_owned(), - build_type: "release".to_owned(), - update_channel: "stable".to_owned(), - update_repository: Some(repository.to_owned()), - github_api_base: format!("{github_api_base}/repos"), - executable_path: Some(PathBuf::from("/app/bin/codex-proxy-rs")), - web_dist_dir: PathBuf::from("/app/web/dist"), - update_state_file: PathBuf::from("/app/.runtime/data/update-state.json"), - update_lock_file: PathBuf::from("/app/.runtime/data/update.lock"), - update_temp_dir: PathBuf::from("/app/.runtime/data/update-tmp"), - self_restart_enabled: false, - } -} - -fn system_update_test_service( - config: SystemUpdateConfig, -) -> (Arc, CancellationToken) { - let shutdown = CancellationToken::new(); - ( - Arc::new(ProcessSystemAdminService::with_config( - shutdown.clone(), - config, - )), - shutdown, - ) -} - -async fn assert_system_update_failure_preserves_files( - service: &ProcessSystemAdminService, - expected_kind: SystemAdminErrorKind, - expected_message: &str, - layout: &SystemUpdateTestLayout, -) { - let error = service - .perform_update(Some("0.4.0".to_owned())) - .await - .expect_err("update must fail"); - let status = service.update_status().await.expect("failed update status"); - assert!( - error.kind() == expected_kind - && error.to_string() == expected_message - && status["operation"]["status"] == "failed" - && status["operation"]["error"] == expected_message - && fs::read_to_string(&layout.executable).expect("old binary") == "old-binary" - && fs::read_to_string(layout.web_dist.join("index.html")).expect("old web") - == "old-web" - ); -} - -async fn mount_system_latest_release(server: &MockServer, repository: &str, version: &str) { - Mock::given(method("GET")) - .and(path(format!("/repos/{repository}/releases/latest"))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "tag_name": format!("v{version}"), - "name": format!("v{version}"), - "body": "## Changes\n- Update package", - "html_url": format!("https://github.com/{repository}/releases/tag/v{version}"), - "prerelease": false, - "assets": [] - }))) - .mount(server) - .await; -} - -async fn mount_system_release_with_archive( - server: &MockServer, - repository: &str, - version: &str, - binary: &str, - index: &str, -) { - let archive = system_release_archive(binary, index); - let checksum = system_archive_checksum(version, &archive); - mount_system_release_archive(server, repository, version, archive, Some(checksum)).await; -} - -async fn mount_system_release_archive( - server: &MockServer, - repository: &str, - version: &str, - archive: Vec, - checksum: Option, -) { - let archive_name = system_archive_name(version); - let mut assets = vec![json!({ - "name": archive_name, - "browser_download_url": format!("{}/download/{archive_name}", server.uri()), - "size": archive.len() - })]; - if let Some(checksum) = checksum.as_ref() { - assets.push(json!({ - "name": "checksums.txt", - "browser_download_url": format!("{}/download/checksums.txt", server.uri()), - "size": checksum.len() - })); - } - Mock::given(method("GET")) - .and(path(format!("/repos/{repository}/releases/latest"))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "tag_name": format!("v{version}"), - "name": format!("v{version}"), - "body": "## Changes\n- Update package", - "html_url": format!("https://github.com/{repository}/releases/tag/v{version}"), - "prerelease": false, - "assets": assets - }))) - .mount(server) - .await; - Mock::given(method("GET")) - .and(path(format!("/download/{archive_name}"))) - .respond_with(ResponseTemplate::new(200).set_body_bytes(archive)) - .mount(server) - .await; - if let Some(checksum) = checksum { - Mock::given(method("GET")) - .and(path("/download/checksums.txt")) - .respond_with(ResponseTemplate::new(200).set_body_string(checksum)) - .mount(server) - .await; - } -} - -async fn mount_system_release_with_urls( - server: &MockServer, - repository: &str, - version: &str, - archive_url: &str, - checksum_url: &str, -) { - Mock::given(method("GET")) - .and(path(format!("/repos/{repository}/releases/latest"))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "tag_name": format!("v{version}"), - "name": format!("v{version}"), - "body": "## Changes\n- Update package", - "html_url": format!("https://github.com/{repository}/releases/tag/v{version}"), - "prerelease": false, - "assets": [ - { - "name": system_archive_name(version), - "browser_download_url": archive_url, - "size": 128 - }, - { - "name": "checksums.txt", - "browser_download_url": checksum_url, - "size": 128 - } - ] - }))) - .mount(server) - .await; -} - -fn system_archive_name(version: &str) -> String { - format!( - "codex-proxy-rs_{version}_{}_{}.tar.gz", - std::env::consts::OS, - system_release_arch() - ) -} - -fn system_release_arch() -> &'static str { - match std::env::consts::ARCH { - "x86_64" => "amd64", - "aarch64" => "arm64", - arch => arch, - } -} - -fn system_archive_checksum(version: &str, archive: &[u8]) -> String { - format!( - "{} {}\n", - hex::encode(Sha256::digest(archive)), - system_archive_name(version) - ) -} - -fn system_release_archive(binary: &str, index: &str) -> Vec { - system_release_archive_with_extra(binary, index, "", &[]) -} - -fn system_release_archive_with_extra( - binary: &str, - index: &str, - extra_path: &str, - extra_bytes: &[u8], -) -> Vec { - let mut data = Vec::new(); - { - let encoder = GzEncoder::new(&mut data, Compression::default()); - let mut builder = Builder::new(encoder); - append_system_tar_file(&mut builder, "codex-proxy-rs", binary.as_bytes()); - append_system_tar_file(&mut builder, "web/dist/index.html", index.as_bytes()); - if !extra_path.is_empty() { - append_system_raw_tar_file(&mut builder, extra_path, extra_bytes); - } - builder.finish().expect("finish release archive"); - } - data -} - -fn append_system_tar_file(builder: &mut Builder, path: &str, bytes: &[u8]) { - let mut header = Header::new_gnu(); - header.set_size(bytes.len() as u64); - header.set_mode(0o644); - header.set_cksum(); - builder - .append_data(&mut header, path, bytes) - .expect("append release file"); -} - -fn append_system_raw_tar_file(builder: &mut Builder, path: &str, bytes: &[u8]) { - let mut header = Header::new_old(); - header.set_size(bytes.len() as u64); - header.set_mode(0o644); - header.set_entry_type(EntryType::Regular); - let name = path.as_bytes(); - header.as_old_mut().name[..name.len()].copy_from_slice(name); - header.set_cksum(); - builder - .append(&header, bytes) - .expect("append raw release file"); -} - -#[derive(Clone)] -struct SystemUpdateTestState { - service: Arc, - auth: Arc, -} - -struct SystemUpdateTestAuth; - -#[async_trait] -impl AdminSessionResolver for SystemUpdateTestAuth { - async fn resolve_admin_user_id( - &self, - session_id: Option<&str>, - ) -> Result, AdminServiceError> { - Ok((session_id == Some("session_1")).then(|| "admin_1".to_owned())) - } -} - -impl AdminSessionState for SystemUpdateTestState { - fn admin_session_resolver(&self) -> &dyn AdminSessionResolver { - self.auth.as_ref() - } -} - -impl SystemAdminState for SystemUpdateTestState { - fn system_admin_service(&self) -> &dyn SystemAdminService { - self.service.as_ref() - } -} - -fn system_update_test_router(service: Arc) -> axum::Router { - system::router::().with_state(SystemUpdateTestState { - service, - auth: Arc::new(SystemUpdateTestAuth), - }) -} - -fn system_update_authenticated_request(method: &str, uri: &str, body: Body) -> Request { - Request::builder() - .method(method) - .uri(uri) - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req-system-update") - .body(body) - .expect("authenticated system request") +fn parse_config(config: &str) -> Result<(GatewayConfig, tempfile::TempDir), String> { + let directory = tempfile::tempdir().map_err(|error| error.to_string())?; + let deploy = directory.path().join("deploy"); + fs::create_dir(&deploy).map_err(|error| error.to_string())?; + let path = deploy.join("config.yaml"); + fs::write(&path, config).map_err(|error| error.to_string())?; + let mut config = config::Config::builder() + .add_source(config::File::from(path).required(true)) + .build() + .and_then(config::Config::try_deserialize::) + .map_err(|error| error.to_string())?; + config + .resolve_and_validate(&deploy) + .map_err(|error| error.to_string())?; + Ok((config, directory)) } diff --git a/backend/apps/gateway/tests/main.rs b/backend/apps/gateway/tests/main.rs index 1fbc67478..9602f1c15 100644 --- a/backend/apps/gateway/tests/main.rs +++ b/backend/apps/gateway/tests/main.rs @@ -1,45 +1,46 @@ mod bootstrap; -mod workers; -use std::collections::BTreeSet; -use std::fs; -use std::path::{Path, PathBuf}; +use std::{ + collections::BTreeSet, + fs, + path::{Path, PathBuf}, +}; use syn::{Attribute, Item, Meta}; #[test] -fn app_tree_is_exactly_the_frozen_six_rust_files() { +fn app_tree_matches_frozen_terminal_manifest() { let root = Path::new(env!("CARGO_MANIFEST_DIR")); assert_eq!( rust_files(&root.join("src")), BTreeSet::from([ PathBuf::from("bootstrap.rs"), + PathBuf::from("lib.rs"), PathBuf::from("main.rs"), - PathBuf::from("workers.rs"), ]), ); assert_eq!( rust_files(&root.join("tests")), - BTreeSet::from([ - PathBuf::from("bootstrap.rs"), - PathBuf::from("main.rs"), - PathBuf::from("workers.rs"), - ]), + BTreeSet::from([PathBuf::from("bootstrap.rs"), PathBuf::from("main.rs")]), ); } #[test] -fn cargo_library_root_is_bootstrap() { +fn cargo_library_root_is_conventional_lib() { let manifest = fs::read_to_string(Path::new(env!("CARGO_MANIFEST_DIR")).join("Cargo.toml")) .expect("read app Cargo.toml"); - let normalized = manifest.split_whitespace().collect::(); - assert!(normalized.contains("[lib]path=\"src/bootstrap.rs\"")); + assert!(!manifest.contains("[lib]")); + assert!( + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("src/lib.rs") + .is_file() + ); } #[test] fn production_files_have_no_hidden_modules_or_test_hooks() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); - for relative in ["bootstrap.rs", "main.rs", "workers.rs"] { + for relative in ["bootstrap.rs", "lib.rs", "main.rs"] { let path = root.join(relative); let source = fs::read_to_string(&path).expect("read production source"); assert!( @@ -55,11 +56,6 @@ fn production_files_have_no_hidden_modules_or_test_hooks() { "{} has an inline module", path.display() ); - assert!( - !module.attrs.iter().any(is_path_or_test_cfg), - "{} has a hidden module hook", - path.display(), - ); } for attribute in item_attrs(item) { assert!( @@ -73,24 +69,25 @@ fn production_files_have_no_hidden_modules_or_test_hooks() { } #[test] -fn bootstrap_and_workers_do_not_own_storage_or_provider_wire_algorithms() { +fn bootstrap_owns_only_bundle_wiring() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); let bootstrap = fs::read_to_string(root.join("bootstrap.rs")).expect("read bootstrap"); - let workers = fs::read_to_string(root.join("workers.rs")).expect("read workers"); for forbidden in [ - "sqlx::query", - "sqlx::query_as", - "redis::Script", - "redis.call(", - "access_token\"", - "refresh_token\"", + "sqlx::", + "redis::", + "Repository", + "impl Provider", + "impl ExecutionService", + "tokio::spawn", + "access_token", + "refresh_token", ] { assert!( !bootstrap.contains(forbidden), "bootstrap owns `{forbidden}`" ); - assert!(!workers.contains(forbidden), "workers owns `{forbidden}`"); } + assert!(bootstrap.lines().count() <= 300); } fn rust_files(root: &Path) -> BTreeSet { diff --git a/backend/apps/gateway/tests/workers.rs b/backend/apps/gateway/tests/workers.rs deleted file mode 100644 index a46ba464c..000000000 --- a/backend/apps/gateway/tests/workers.rs +++ /dev/null @@ -1,853 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet, VecDeque}; -use std::future::pending; -use std::num::NonZeroU64; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; -use std::time::Duration; - -use async_trait::async_trait; -use codex_proxy_rs::workers::{ - WorkerCycleContext, WorkerDisabledReason, WorkerFencingToken, WorkerHealthKey, WorkerId, - WorkerKind, WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, - WorkerLeaseError, WorkerLeaseRequest, WorkerRegistry, WorkerRegistryError, WorkerRuntimeState, - WorkerSchedule, WorkerTask, WorkerTaskError, -}; -use tokio::sync::Notify; - -const ACTIVE_KINDS: [WorkerKind; 4] = [ - WorkerKind::OAuthRefresh, - WorkerKind::QuotaCatalogHealth, - WorkerKind::StaleModelRequestRecovery, - WorkerKind::Retention, -]; - -struct CountingTask { - runs: AtomicUsize, - fail_until: usize, - fences: Mutex>, -} - -impl CountingTask { - fn new(fail_until: usize) -> Self { - Self { - runs: AtomicUsize::new(0), - fail_until, - fences: Mutex::new(Vec::new()), - } - } -} - -#[async_trait] -impl WorkerTask for CountingTask { - async fn run_cycle(&self, context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - lock_unpoisoned(&self.fences).push(context.fencing_token().get()); - let run = self.runs.fetch_add(1, Ordering::SeqCst) + 1; - if run <= self.fail_until { - Err(WorkerTaskError::safe("expected test failure")) - } else { - Ok(()) - } - } -} - -struct PanicOnceTask { - runs: AtomicUsize, -} - -#[async_trait] -impl WorkerTask for PanicOnceTask { - async fn run_cycle(&self, _context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - if self.runs.fetch_add(1, Ordering::SeqCst) == 0 { - panic!("expected isolated worker panic"); - } - Ok(()) - } -} - -struct CancellationTask { - entered: AtomicUsize, - observed: AtomicUsize, - notification: Notify, -} - -#[async_trait] -impl WorkerTask for CancellationTask { - async fn run_cycle(&self, context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - self.entered.fetch_add(1, Ordering::SeqCst); - self.notification.notify_one(); - context.cancelled().await; - self.observed.fetch_add(1, Ordering::SeqCst); - Ok(()) - } -} - -struct HungTask { - entered: AtomicBool, - notification: Notify, -} - -struct SucceedsThenHangsTask { - runs: AtomicUsize, - second_cycle: Notify, -} - -#[async_trait] -impl WorkerTask for SucceedsThenHangsTask { - async fn run_cycle(&self, _context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - if self.runs.fetch_add(1, Ordering::SeqCst) == 0 { - return Ok(()); - } - self.second_cycle.notify_one(); - pending::<()>().await; - Ok(()) - } -} - -#[async_trait] -impl WorkerTask for HungTask { - async fn run_cycle(&self, _context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - self.entered.store(true, Ordering::SeqCst); - self.notification.notify_one(); - pending::<()>().await; - Ok(()) - } -} - -struct LeaseAwareTask { - entered: Notify, - release: Notify, - completed: AtomicBool, - saw_live_guard: AtomicBool, - expected_worker: WorkerId, - active_guards: Arc>>, - fence: AtomicU64, -} - -#[async_trait] -impl WorkerTask for LeaseAwareTask { - async fn run_cycle(&self, context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - self.saw_live_guard.store( - lock_unpoisoned(&self.active_guards).contains(context.worker()), - Ordering::SeqCst, - ); - if context.worker() != &self.expected_worker { - return Err(WorkerTaskError::safe("unexpected worker identity")); - } - self.fence - .store(context.fencing_token().get(), Ordering::SeqCst); - self.entered.notify_one(); - self.release.notified().await; - self.completed.store(true, Ordering::SeqCst); - Ok(()) - } -} - -struct ConcurrentCycleTask { - starts: AtomicUsize, - active: AtomicUsize, - maximum_active: AtomicUsize, -} - -struct ActiveCycle<'a> { - active: &'a AtomicUsize, -} - -impl Drop for ActiveCycle<'_> { - fn drop(&mut self) { - self.active.fetch_sub(1, Ordering::SeqCst); - } -} - -#[async_trait] -impl WorkerTask for ConcurrentCycleTask { - async fn run_cycle(&self, context: WorkerCycleContext) -> Result<(), WorkerTaskError> { - self.starts.fetch_add(1, Ordering::SeqCst); - let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; - self.maximum_active.fetch_max(active, Ordering::SeqCst); - let _active_cycle = ActiveCycle { - active: &self.active, - }; - context.cancelled().await; - Ok(()) - } -} - -#[derive(Debug, Clone)] -enum LeaseStep { - Acquired, - Busy(Option), - Error, - Panic, -} - -#[derive(Debug, Clone, Copy)] -enum RenewStep { - Success, - Error, -} - -#[derive(Default)] -struct FakeLeasePort { - scripts: Mutex>>, - renew_scripts: Arc>>>, - requests: Mutex>, - renewals: Arc>>, - active_guards: Arc>>, - next_fence: AtomicU64, -} - -impl FakeLeasePort { - fn script(&self, worker: WorkerId, steps: impl IntoIterator) { - lock_unpoisoned(&self.scripts).insert(worker, steps.into_iter().collect()); - } - - fn renew_script(&self, worker: WorkerId, steps: impl IntoIterator) { - lock_unpoisoned(&self.renew_scripts).insert(worker, steps.into_iter().collect()); - } - - fn requests(&self) -> Vec { - lock_unpoisoned(&self.requests).clone() - } - - fn renewal_count(&self, worker: &WorkerId) -> usize { - lock_unpoisoned(&self.renewals) - .iter() - .filter(|renewed| *renewed == worker) - .count() - } -} - -struct FakeLeaseGuard { - worker: WorkerId, - fencing_token: WorkerFencingToken, - active_guards: Arc>>, - renew_scripts: Arc>>>, - renewals: Arc>>, -} - -#[async_trait] -impl WorkerLeaderLeaseGuard for FakeLeaseGuard { - fn fencing_token(&self) -> WorkerFencingToken { - self.fencing_token - } - - async fn renew(&mut self) -> Result<(), WorkerLeaseError> { - lock_unpoisoned(&self.renewals).push(self.worker.clone()); - match lock_unpoisoned(&self.renew_scripts) - .get_mut(&self.worker) - .and_then(VecDeque::pop_front) - .unwrap_or(RenewStep::Success) - { - RenewStep::Success => Ok(()), - RenewStep::Error => Err(WorkerLeaseError::safe("expected renewal failure")), - } - } -} - -impl Drop for FakeLeaseGuard { - fn drop(&mut self) { - lock_unpoisoned(&self.active_guards).remove(&self.worker); - } -} - -#[async_trait] -impl WorkerLeaderLeasePort for FakeLeasePort { - async fn try_acquire( - &self, - request: WorkerLeaseRequest, - ) -> Result { - lock_unpoisoned(&self.requests).push(request.clone()); - let step = lock_unpoisoned(&self.scripts) - .get_mut(request.worker()) - .and_then(VecDeque::pop_front) - .unwrap_or(LeaseStep::Acquired); - match step { - LeaseStep::Acquired => { - let fence = self.next_fence.fetch_add(1, Ordering::SeqCst) + 1; - lock_unpoisoned(&self.active_guards).insert(request.worker().clone()); - Ok(WorkerLeaseAcquisition::Acquired(Box::new(FakeLeaseGuard { - worker: request.worker().clone(), - fencing_token: WorkerFencingToken::new( - NonZeroU64::new(fence).expect("positive test fencing token"), - ), - active_guards: Arc::clone(&self.active_guards), - renew_scripts: Arc::clone(&self.renew_scripts), - renewals: Arc::clone(&self.renewals), - }))) - } - LeaseStep::Busy(retry_after) => Ok(WorkerLeaseAcquisition::Busy { retry_after }), - LeaseStep::Error => Err(WorkerLeaseError::safe("expected lease failure")), - LeaseStep::Panic => panic!("expected isolated lease port panic"), - } - } -} - -#[test] -fn schedule_rejects_every_invalid_duration_boundary() { - let valid = Duration::from_millis(1); - for values in [ - (Duration::ZERO, valid, valid, valid), - (valid, Duration::ZERO, valid, valid), - (valid, valid, Duration::ZERO, valid), - (valid, valid, valid, Duration::ZERO), - (valid, valid, valid, Duration::from_nanos(2)), - ( - valid, - Duration::from_millis(2), - Duration::from_millis(1), - valid, - ), - ] { - assert_eq!( - WorkerSchedule::new(values.0, values.1, values.2, values.3), - Err(WorkerRegistryError::InvalidSchedule), - ); - } -} - -#[test] -fn worker_owner_rejects_empty_unsafe_or_oversized_names() { - for owner in [ - String::new(), - "Codex".to_owned(), - "provider/codex".to_owned(), - "x".repeat(65), - ] { - assert!(matches!( - WorkerId::new(WorkerKind::OAuthRefresh, owner), - Err(WorkerRegistryError::InvalidOwner(_)) - )); - } -} - -#[test] -fn registry_composes_multiple_owners_but_rejects_duplicate_identity() { - let mut registry = WorkerRegistry::new(); - let task = Arc::new(CountingTask::new(0)); - registry - .register( - WorkerKind::OAuthRefresh, - "openai", - task.clone(), - test_schedule(), - ) - .expect("first owner registration"); - registry - .register( - WorkerKind::OAuthRefresh, - "xai", - task.clone(), - test_schedule(), - ) - .expect("second owner registration"); - assert!(matches!( - registry.register( - WorkerKind::OAuthRefresh, - "openai", - task, - test_schedule(), - ), - Err(WorkerRegistryError::Duplicate(id)) if id.owner() == "openai" - )); -} - -#[test] -fn only_database_fact_reasons_can_disable_the_two_nonexistent_workers() { - let task = Arc::new(CountingTask::new(0)); - let mut disabled = WorkerRegistry::new(); - disabled - .disable(WorkerDisabledReason::NoPersistentNativeClaimState) - .expect("native claim terminal fact"); - assert_eq!( - disabled.register( - WorkerKind::NativeClaimRecovery, - "store", - task.clone(), - test_schedule(), - ), - Err(WorkerRegistryError::KindDisabled( - WorkerKind::NativeClaimRecovery - )), - ); - - let mut active = WorkerRegistry::new(); - active - .register(WorkerKind::OpsFlush, "store", task, test_schedule()) - .expect("real owner registration"); - assert_eq!( - active.disable(WorkerDisabledReason::NoBufferedOpsEvents), - Err(WorkerRegistryError::KindHasTasks(WorkerKind::OpsFlush)), - ); -} - -#[test] -fn registry_refuses_start_when_any_real_owner_kind_is_missing() { - let task = Arc::new(CountingTask::new(0)); - let mut registry = WorkerRegistry::new(); - for kind in ACTIVE_KINDS - .into_iter() - .filter(|kind| *kind != WorkerKind::Retention) - { - registry - .register(kind, kind.as_str(), task.clone(), test_schedule()) - .expect("unique active worker"); - } - disable_nonexistent_workers(&mut registry); - let lease_port: Arc = Arc::new(FakeLeasePort::default()); - assert!(matches!( - registry.start(lease_port), - Err(WorkerRegistryError::Missing(missing)) if missing == vec![WorkerKind::Retention] - )); -} - -#[tokio::test] -async fn disabled_database_facts_never_execute_or_acquire_a_lease() { - let task = Arc::new(CountingTask::new(0)); - let lease_port = Arc::new(FakeLeasePort::default()); - let supervisor = complete_registry(task) - .start(lease_port.clone()) - .expect("complete registry starts"); - - wait_until(|| lease_port.requests().len() >= ACTIVE_KINDS.len()).await; - assert!(lease_port.requests().iter().all(|request| { - !matches!( - request.worker().kind(), - WorkerKind::NativeClaimRecovery | WorkerKind::OpsFlush - ) - })); - let snapshot = supervisor.health().snapshot(); - assert_eq!( - snapshot - .get(&WorkerHealthKey::Disabled(WorkerKind::NativeClaimRecovery)) - .map(|health| health.state), - Some(WorkerRuntimeState::Disabled( - WorkerDisabledReason::NoPersistentNativeClaimState - )), - ); - assert_eq!( - snapshot - .get(&WorkerHealthKey::Disabled(WorkerKind::OpsFlush)) - .map(|health| health.state), - Some(WorkerRuntimeState::Disabled( - WorkerDisabledReason::NoBufferedOpsEvents - )), - ); - assert!(supervisor.health().all_healthy()); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test] -async fn one_kind_runs_all_registered_owner_tasks() { - let common = Arc::new(CountingTask::new(0)); - let second_oauth_owner = Arc::new(CountingTask::new(0)); - let mut registry = complete_registry(common); - registry - .register( - WorkerKind::OAuthRefresh, - "xai", - second_oauth_owner.clone(), - test_schedule(), - ) - .expect("second OAuth owner"); - let supervisor = registry - .start(Arc::new(FakeLeasePort::default())) - .expect("complete registry starts"); - - wait_until(|| second_oauth_owner.runs.load(Ordering::SeqCst) >= 1).await; - assert_eq!( - supervisor - .health() - .snapshot() - .keys() - .filter(|key| matches!(key, WorkerHealthKey::Task(id) if id.kind() == WorkerKind::OAuthRefresh)) - .count(), - 2, - ); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test] -async fn leader_guard_lives_through_cycle_and_fence_reaches_owner() { - let expected_worker = - WorkerId::new(WorkerKind::OAuthRefresh, "openai").expect("valid expected worker identity"); - let lease_port = Arc::new(FakeLeasePort::default()); - let lease_task = Arc::new(LeaseAwareTask { - entered: Notify::new(), - release: Notify::new(), - completed: AtomicBool::new(false), - saw_live_guard: AtomicBool::new(false), - expected_worker: expected_worker.clone(), - active_guards: Arc::clone(&lease_port.active_guards), - fence: AtomicU64::new(0), - }); - let other = Arc::new(CountingTask::new(0)); - let mut registry = WorkerRegistry::new(); - registry - .register( - WorkerKind::OAuthRefresh, - "openai", - lease_task.clone(), - long_interval_schedule(), - ) - .expect("lease-aware worker"); - for kind in ACTIVE_KINDS - .into_iter() - .filter(|kind| *kind != WorkerKind::OAuthRefresh) - { - registry - .register(kind, kind.as_str(), other.clone(), long_interval_schedule()) - .expect("other active worker"); - } - disable_nonexistent_workers(&mut registry); - let supervisor = registry - .start(lease_port.clone()) - .expect("complete registry starts"); - - lease_task.entered.notified().await; - assert!(lease_task.saw_live_guard.load(Ordering::SeqCst)); - assert!( - lock_unpoisoned(&lease_port.active_guards).contains(&expected_worker), - "guard must remain alive while owner task is running" - ); - lease_task.release.notify_one(); - wait_until(|| lease_task.completed.load(Ordering::SeqCst)).await; - wait_until(|| !lock_unpoisoned(&lease_port.active_guards).contains(&expected_worker)).await; - assert!(lease_task.fence.load(Ordering::SeqCst) > 0); - let request = lease_port - .requests() - .into_iter() - .find(|request| request.worker() == &expected_worker) - .expect("leader lease request recorded"); - assert_eq!(request.ttl(), Duration::from_secs(30)); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test(start_paused = true)] -async fn long_cycle_renews_leader_lease_before_ttl() { - let expected_worker = target_worker(); - let lease_port = Arc::new(FakeLeasePort::default()); - let lease_task = Arc::new(LeaseAwareTask { - entered: Notify::new(), - release: Notify::new(), - completed: AtomicBool::new(false), - saw_live_guard: AtomicBool::new(false), - expected_worker: expected_worker.clone(), - active_guards: Arc::clone(&lease_port.active_guards), - fence: AtomicU64::new(0), - }); - let supervisor = registry_with_target(lease_task.clone(), renewal_schedule()) - .start(lease_port.clone()) - .expect("complete registry starts"); - - lease_task.entered.notified().await; - tokio::time::advance(Duration::from_millis(10)).await; - yield_until(|| lease_port.renewal_count(&expected_worker) >= 1).await; - assert!(lock_unpoisoned(&lease_port.active_guards).contains(&expected_worker)); - lease_task.release.notify_one(); - yield_until(|| lease_task.completed.load(Ordering::SeqCst)).await; - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test(start_paused = true)] -async fn renewal_failure_cancels_old_cycle_before_another_leader_cycle_starts() { - let task = Arc::new(ConcurrentCycleTask { - starts: AtomicUsize::new(0), - active: AtomicUsize::new(0), - maximum_active: AtomicUsize::new(0), - }); - let lease_port = Arc::new(FakeLeasePort::default()); - lease_port.renew_script(target_worker(), [RenewStep::Error]); - let supervisor = registry_with_target(task.clone(), renewal_schedule()) - .start(lease_port) - .expect("complete registry starts"); - - yield_until(|| task.starts.load(Ordering::SeqCst) == 1).await; - tokio::time::advance(Duration::from_millis(10)).await; - yield_until(|| task.active.load(Ordering::SeqCst) == 0).await; - assert_eq!( - task_health(&supervisor.health().snapshot(), target_worker()).state, - WorkerRuntimeState::BackingOff, - ); - tokio::time::advance(Duration::from_millis(5)).await; - yield_until(|| task.starts.load(Ordering::SeqCst) == 2).await; - assert_eq!(task.maximum_active.load(Ordering::SeqCst), 1); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test(start_paused = true)] -async fn supervisor_uses_exponential_backoff_and_recovers_health() { - let recovering = Arc::new(CountingTask::new(2)); - let supervisor = registry_with_target(recovering.clone(), backoff_schedule()) - .start(Arc::new(FakeLeasePort::default())) - .expect("complete registry starts"); - - yield_until(|| recovering.runs.load(Ordering::SeqCst) == 1).await; - tokio::time::advance(Duration::from_millis(9)).await; - yield_many().await; - assert_eq!(recovering.runs.load(Ordering::SeqCst), 1); - tokio::time::advance(Duration::from_millis(1)).await; - yield_until(|| recovering.runs.load(Ordering::SeqCst) == 2).await; - tokio::time::advance(Duration::from_millis(19)).await; - yield_many().await; - assert_eq!(recovering.runs.load(Ordering::SeqCst), 2); - tokio::time::advance(Duration::from_millis(1)).await; - yield_until(|| recovering.runs.load(Ordering::SeqCst) == 3).await; - - let health = task_health(&supervisor.health().snapshot(), target_worker()); - assert_eq!(health.consecutive_failures, 0); - assert_eq!(health.completed_cycles, 1); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test(start_paused = true)] -async fn task_and_lease_port_panics_are_isolated_and_retried() { - let panic_task = Arc::new(PanicOnceTask { - runs: AtomicUsize::new(0), - }); - let lease_port = Arc::new(FakeLeasePort::default()); - let quota_worker = WorkerId::new(WorkerKind::QuotaCatalogHealth, "quota_catalog_health") - .expect("valid quota worker"); - lease_port.script( - quota_worker, - [LeaseStep::Error, LeaseStep::Panic, LeaseStep::Acquired], - ); - let supervisor = registry_with_target(panic_task.clone(), backoff_schedule()) - .start(lease_port) - .expect("complete registry starts"); - - yield_until(|| panic_task.runs.load(Ordering::SeqCst) == 1).await; - tokio::time::advance(Duration::from_millis(10)).await; - yield_until(|| panic_task.runs.load(Ordering::SeqCst) == 2).await; - tokio::time::advance(Duration::from_millis(20)).await; - yield_many().await; - assert!(panic_task.runs.load(Ordering::SeqCst) >= 2); - let health = task_health(&supervisor.health().snapshot(), target_worker()); - assert_eq!(health.consecutive_failures, 0); - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test(start_paused = true)] -async fn lease_busy_is_healthy_standby_and_retries_without_running_owner() { - let task = Arc::new(CountingTask::new(0)); - let lease_port = Arc::new(FakeLeasePort::default()); - lease_port.script( - target_worker(), - [ - LeaseStep::Busy(Some(Duration::from_millis(7))), - LeaseStep::Acquired, - ], - ); - let supervisor = registry_with_target(task.clone(), backoff_schedule()) - .start(lease_port) - .expect("complete registry starts"); - - yield_until(|| { - task_health(&supervisor.health().snapshot(), target_worker()).state - == WorkerRuntimeState::Standby - }) - .await; - assert_eq!(task.runs.load(Ordering::SeqCst), 0); - assert!(supervisor.health().all_healthy()); - tokio::time::advance(Duration::from_millis(7)).await; - yield_until(|| task.runs.load(Ordering::SeqCst) == 1).await; - supervisor.shutdown(Duration::from_secs(1)).await; -} - -#[tokio::test] -async fn cooperative_shutdown_propagates_cancel_before_dropping_all_guards() { - let task = Arc::new(CancellationTask { - entered: AtomicUsize::new(0), - observed: AtomicUsize::new(0), - notification: Notify::new(), - }); - let lease_port = Arc::new(FakeLeasePort::default()); - let supervisor = complete_registry(task.clone()) - .start(lease_port.clone()) - .expect("complete registry starts"); - - wait_until(|| task.entered.load(Ordering::SeqCst) == ACTIVE_KINDS.len()).await; - let health = supervisor.health().clone(); - supervisor.shutdown(Duration::from_secs(1)).await; - assert_eq!(task.observed.load(Ordering::SeqCst), ACTIVE_KINDS.len()); - assert!(lock_unpoisoned(&lease_port.active_guards).is_empty()); - assert!(health.snapshot().iter().all(|(key, value)| { - matches!(key, WorkerHealthKey::Disabled(_)) || value.state == WorkerRuntimeState::Stopped - })); -} - -#[tokio::test] -async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() { - let hung = Arc::new(HungTask { - entered: AtomicBool::new(false), - notification: Notify::new(), - }); - let lease_port = Arc::new(FakeLeasePort::default()); - let supervisor = registry_with_target(hung.clone(), long_interval_schedule()) - .start(lease_port.clone()) - .expect("complete registry starts"); - - hung.notification.notified().await; - assert!(hung.entered.load(Ordering::SeqCst)); - assert!( - !supervisor.health().all_healthy(), - "a running worker without one successful cycle is not healthy" - ); - let health = supervisor.health().clone(); - supervisor.shutdown(Duration::from_millis(10)).await; - assert!(lock_unpoisoned(&lease_port.active_guards).is_empty()); - assert_eq!( - task_health(&health.snapshot(), target_worker()).state, - WorkerRuntimeState::Stopped, - ); -} - -#[tokio::test] -async fn running_worker_becomes_unhealthy_when_last_success_is_stale() { - let task = Arc::new(SucceedsThenHangsTask { - runs: AtomicUsize::new(0), - second_cycle: Notify::new(), - }); - let schedule = WorkerSchedule::new( - Duration::from_millis(10), - Duration::from_millis(5), - Duration::from_millis(10), - Duration::from_secs(30), - ) - .expect("freshness schedule"); - let supervisor = registry_with_target(task.clone(), schedule) - .start(Arc::new(FakeLeasePort::default())) - .expect("complete registry starts"); - - task.second_cycle.notified().await; - assert_eq!(task.runs.load(Ordering::SeqCst), 2); - tokio::time::sleep(Duration::from_millis(40)).await; - assert!( - !supervisor.health().all_healthy(), - "running cannot mask an expired last-success timestamp" - ); - supervisor.shutdown(Duration::from_millis(10)).await; -} - -fn complete_registry(task: Arc) -> WorkerRegistry { - let mut registry = WorkerRegistry::new(); - for kind in ACTIVE_KINDS { - registry - .register(kind, kind.as_str(), task.clone(), test_schedule()) - .expect("unique real worker registration"); - } - disable_nonexistent_workers(&mut registry); - registry -} - -fn registry_with_target(target: Arc, schedule: WorkerSchedule) -> WorkerRegistry { - let other = Arc::new(CountingTask::new(0)); - let mut registry = WorkerRegistry::new(); - registry - .register(WorkerKind::OAuthRefresh, "openai", target, schedule) - .expect("target worker registration"); - for kind in ACTIVE_KINDS - .into_iter() - .filter(|kind| *kind != WorkerKind::OAuthRefresh) - { - registry - .register(kind, kind.as_str(), other.clone(), schedule) - .expect("other worker registration"); - } - disable_nonexistent_workers(&mut registry); - registry -} - -fn disable_nonexistent_workers(registry: &mut WorkerRegistry) { - registry - .disable(WorkerDisabledReason::NoPersistentNativeClaimState) - .expect("native claim has no persistent state"); - registry - .disable(WorkerDisabledReason::NoBufferedOpsEvents) - .expect("ops events have no buffer"); -} - -fn target_worker() -> WorkerId { - WorkerId::new(WorkerKind::OAuthRefresh, "openai").expect("valid target worker") -} - -fn task_health( - snapshot: &BTreeMap, - id: WorkerId, -) -> codex_proxy_rs::workers::WorkerHealth { - snapshot - .get(&WorkerHealthKey::Task(id)) - .expect("worker health exists") - .clone() -} - -fn test_schedule() -> WorkerSchedule { - WorkerSchedule::new( - Duration::from_millis(20), - Duration::from_millis(5), - Duration::from_millis(20), - Duration::from_secs(30), - ) - .expect("valid test schedule") -} - -fn long_interval_schedule() -> WorkerSchedule { - WorkerSchedule::new( - Duration::from_secs(60), - Duration::from_millis(5), - Duration::from_millis(20), - Duration::from_secs(30), - ) - .expect("valid long-interval schedule") -} - -fn backoff_schedule() -> WorkerSchedule { - WorkerSchedule::new( - Duration::from_secs(60), - Duration::from_millis(10), - Duration::from_millis(20), - Duration::from_secs(30), - ) - .expect("valid backoff schedule") -} - -fn renewal_schedule() -> WorkerSchedule { - WorkerSchedule::new( - Duration::from_secs(60), - Duration::from_millis(5), - Duration::from_millis(20), - Duration::from_millis(30), - ) - .expect("valid renewal schedule") -} - -async fn wait_until(mut condition: impl FnMut() -> bool) { - tokio::time::timeout(Duration::from_secs(1), async { - while !condition() { - tokio::task::yield_now().await; - } - }) - .await - .expect("worker test timeout"); -} - -async fn yield_until(mut condition: impl FnMut() -> bool) { - for _ in 0..100 { - if condition() { - return; - } - tokio::task::yield_now().await; - } - assert!(condition(), "condition did not become true after yielding"); -} - -async fn yield_many() { - for _ in 0..20 { - tokio::task::yield_now().await; - } -} - -fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { - mutex - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} diff --git a/backend/crates/gateway-admin/Cargo.toml b/backend/crates/gateway-admin/Cargo.toml new file mode 100644 index 000000000..4355f1623 --- /dev/null +++ b/backend/crates/gateway-admin/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "gateway-admin" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true +autotests = false + +[[test]] +name = "main" +path = "tests/main.rs" + +[dependencies] +argon2 = "0.5.3" +async-trait = "0.1.89" +base64 = "0.22.1" +chrono = "0.4.45" +futures = "0.3.32" +gateway-core = { path = "../gateway-core" } +hex = "0.4.3" +rand_core = { version = "0.6.4", features = ["getrandom"] } +serde = { version = "1.0.228", features = ["derive"] } +secrecy = { version = "0.10.3", features = ["serde"] } +subtle = "2.6.1" +thiserror = "2.0.18" +uuid = { version = "1.23.3", features = ["v7"] } + +[dev-dependencies] +tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] } + +[lints] +workspace = true diff --git a/backend/crates/gateway-admin/src/lib.rs b/backend/crates/gateway-admin/src/lib.rs new file mode 100644 index 000000000..16a28302d --- /dev/null +++ b/backend/crates/gateway-admin/src/lib.rs @@ -0,0 +1,301 @@ +//! 管理控制面的语义模型、用例与外部能力端口。 +//! +//! 本 crate 不包含 HTTP wire、数据库实现或具体 Provider 实现。 + +use std::{fmt, path::Path, sync::Arc}; + +use gateway_core::{ + engine::probe::AccountProbe, + routing::{ProviderKind, snapshot::SnapshotControl}, +}; +use secrecy::{ExposeSecret as _, SecretString}; +use serde::Deserialize; + +pub mod model; +pub mod ports; +mod use_case; + +pub use use_case::{ + accounts::AccountsService, auth::AuthService, catalog::CatalogService, + client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService, + settings::SettingsService, system::SystemService, xai::XaiService, +}; + +use model::{AdminError, AdminErrorKind}; +use ports::{ + provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind, ProviderAdminRegistry}, + store::AdminStorePorts, + system::SystemOperations, +}; +use use_case::{ + accounts::DefaultAccountsService, auth::DefaultAuthService, catalog::DefaultCatalogService, + client_keys::DefaultClientKeyService, observability::DefaultObservabilityService, + openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService, + xai::DefaultXaiService, +}; + +const OPENAI_PROVIDER_KIND: &str = "openai"; +const XAI_PROVIDER_KIND: &str = "xai"; +const MINIMUM_INITIAL_PASSWORD_BYTES: usize = 12; +const WEAK_INITIAL_PASSWORDS: &[&str] = &[ + "", + "admin", + "123456", + "password", + "changeme", + "change-me", + "replace-me", + "codex-proxy-rs", +]; + +/// 只用于首次幂等创建默认管理员的启动密码。 +#[derive(Clone, Deserialize)] +#[serde(transparent)] +pub struct InitialAdminPassword(SecretString); + +impl InitialAdminPassword { + #[must_use] + pub fn new(value: impl Into) -> Self { + Self(SecretString::from(value.into())) + } + + fn expose(&self) -> &str { + self.0.expose_secret() + } +} + +impl PartialEq for InitialAdminPassword { + fn eq(&self, other: &Self) -> bool { + self.expose() == other.expose() + } +} + +impl Eq for InitialAdminPassword {} + +impl fmt::Debug for InitialAdminPassword { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("InitialAdminPassword([REDACTED])") + } +} + +/// 管理控制面的启动配置。 +#[derive(Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct AdminConfig { + pub session_ttl_minutes: u64, + pub default_username: String, + pub default_password: InitialAdminPassword, +} + +impl AdminConfig { + /// 校验 Admin-owned 字段；当前配置不含相对路径。 + /// + /// # Errors + /// + /// 用户名、会话有效期或初始密码不满足安全约束时返回错误。 + pub fn resolve_and_validate(&mut self, _source_dir: &Path) -> Result<(), AdminConfigError> { + if self.default_username.trim().is_empty() + || self.default_username.chars().any(char::is_control) + { + return Err(AdminConfigError::InvalidField("admin.default_username")); + } + if self.session_ttl_minutes == 0 || i64::try_from(self.session_ttl_minutes).is_err() { + return Err(AdminConfigError::InvalidField("admin.session_ttl_minutes")); + } + let password = self.default_password.expose().trim(); + if password.len() < MINIMUM_INITIAL_PASSWORD_BYTES + || password.contains('$') + || WEAK_INITIAL_PASSWORDS.contains(&password.to_ascii_lowercase().as_str()) + { + return Err(AdminConfigError::WeakInitialPassword); + } + Ok(()) + } +} + +impl fmt::Debug for AdminConfig { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AdminConfig") + .field("session_ttl_minutes", &self.session_ttl_minutes) + .field("default_username", &self.default_username) + .field("default_password", &"[REDACTED]") + .finish() + } +} + +/// Admin-owned 启动配置错误；不回显任何配置值。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum AdminConfigError { + #[error("configuration field `{0}` is invalid")] + InvalidField(&'static str), + #[error("admin.default_password does not meet the initial password policy")] + WeakInitialPassword, +} + +/// API 持有的管理资源能力集合。 +/// +/// 字段全部私有；调用方经 accessor 直接调用能力，不需要命名内部 `use_case` 模块。 +#[derive(Clone)] +pub struct AdminServices { + auth: Arc, + accounts: Arc, + catalog: Arc, + client_keys: Arc, + observability: Arc, + settings: Arc, + system: Arc, + openai: Arc, + xai: Arc, +} + +impl AdminServices { + #[must_use] + pub fn auth(&self) -> &dyn AuthService { + self.auth.as_ref() + } + + #[must_use] + pub fn accounts(&self) -> &dyn AccountsService { + self.accounts.as_ref() + } + + #[must_use] + pub fn catalog(&self) -> &dyn CatalogService { + self.catalog.as_ref() + } + + #[must_use] + pub fn client_keys(&self) -> &dyn ClientKeyService { + self.client_keys.as_ref() + } + + #[must_use] + pub fn observability(&self) -> &dyn ObservabilityService { + self.observability.as_ref() + } + + #[must_use] + pub fn settings(&self) -> &dyn SettingsService { + self.settings.as_ref() + } + + #[must_use] + pub fn system(&self) -> &dyn SystemService { + self.system.as_ref() + } + + #[must_use] + pub fn openai(&self) -> &dyn OpenAiService { + self.openai.as_ref() + } + + #[must_use] + pub fn xai(&self) -> &dyn XaiService { + self.xai.as_ref() + } +} + +/// Admin 初始化完成后的封闭能力包。 +pub struct AdminBundle { + services: AdminServices, +} + +impl AdminBundle { + #[must_use] + pub fn services(&self) -> AdminServices { + self.services.clone() + } +} + +/// 校验配置、建立动态 Provider 注册表并完成默认管理员幂等初始化。 +/// +/// # Errors +/// +/// 配置非法、Provider 注册冲突/缺失或默认管理员初始化失败时返回错误。 +pub async fn initialize( + mut config: AdminConfig, + store: AdminStorePorts, + providers: Vec>, + snapshot: Arc, + probe: Arc, + system: Arc, +) -> Result { + config + .resolve_and_validate(Path::new(".")) + .map_err(|error| AdminError::invalid(error.to_string()))?; + let registry = ProviderAdminRegistry::new(providers).map_err(map_provider_registry_error)?; + let openai = registry + .require(&provider_kind(OPENAI_PROVIDER_KIND)?) + .map_err(map_provider_registry_error)?; + let xai = registry + .require(&provider_kind(XAI_PROVIDER_KIND)?) + .map_err(map_provider_registry_error)?; + + let auth = Arc::new(DefaultAuthService::new( + config.default_username, + config.session_ttl_minutes, + store.auth(), + )); + auth.ensure_default_admin(config.default_password.expose()) + .await?; + + let accounts = Arc::new(DefaultAccountsService::new( + store.accounts(), + store.catalog(), + store.settings(), + registry.clone(), + snapshot.clone(), + probe.clone(), + )); + let services = AdminServices { + auth, + accounts, + catalog: Arc::new(DefaultCatalogService::new( + store.catalog(), + snapshot.clone(), + )), + client_keys: Arc::new(DefaultClientKeyService::new( + store.client_keys(), + snapshot.clone(), + )), + observability: Arc::new(DefaultObservabilityService::new( + store.observability(), + store.settings(), + registry, + )), + settings: Arc::new(DefaultSettingsService::new( + store.settings(), + snapshot.clone(), + )), + system: Arc::new(DefaultSystemService::new(system)), + openai: Arc::new(DefaultOpenAiService::new( + openai, + store.accounts(), + snapshot.clone(), + )), + xai: Arc::new(DefaultXaiService::new( + xai, + store.accounts(), + snapshot.clone(), + )), + }; + Ok(AdminBundle { services }) +} + +fn provider_kind(value: &'static str) -> Result { + ProviderKind::new(value).map_err(|_| AdminError::internal("Built-in Provider kind is invalid")) +} + +fn map_provider_registry_error(error: ProviderAdminError) -> AdminError { + let kind = match error.kind() { + ProviderAdminErrorKind::Invalid | ProviderAdminErrorKind::Unsupported => { + AdminErrorKind::Invalid + } + ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, + ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, + ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, + }; + AdminError::new(kind, "Provider registry initialization failed") +} diff --git a/backend/crates/gateway-admin/src/model/accounts.rs b/backend/crates/gateway-admin/src/model/accounts.rs new file mode 100644 index 000000000..96fd5fd3e --- /dev/null +++ b/backend/crates/gateway-admin/src/model/accounts.rs @@ -0,0 +1,197 @@ +//! 多 Provider 账号目录与连接测试的公共事实。 + +use std::pin::Pin; + +use chrono::{DateTime, Utc}; +use futures::Stream; + +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; + +use super::{PageSize, Revision}; + +/// Provider 账号当前可用性。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AccountAvailability { + Unknown, + Ready, + Cooldown, + QuotaExhausted, + Expired, + Banned, + Invalid, +} + +/// 管理页使用的归一化账号状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AccountStatus { + Active, + Expired, + QuotaExhausted, + Disabled, + Banned, + Attention, +} + +/// 账号列表排序字段。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AccountSortField { + Email, + Status, + PlanType, + Usage, + LastUsedAt, + ExpiresAt, +} + +/// 账号列表排序方向。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SortDirection { + Asc, + Desc, +} + +/// 一组完整的账号排序规则。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountSort { + pub field: AccountSortField, + pub direction: SortDirection, +} + +/// 账号列表的存储查询条件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountListQuery { + pub page: u32, + pub page_size: PageSize, + pub provider_kind: Option, + pub search: Option, + pub status: Option, + pub sort: Option, +} + +/// 账号公共存储投影；Provider 专属字段不进入此结构。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountRecord { + pub id: String, + pub provider_instance_id: ProviderInstanceId, + pub provider_kind: ProviderKind, + pub name: String, + pub email: Option, + pub upstream_user_id: String, + pub upstream_account_id: Option, + pub plan_type: Option, + pub credential_revision: Revision, + pub has_refresh_token: bool, + pub access_token_expires_at: DateTime, + pub next_refresh_at: Option>, + pub enabled: bool, + pub availability: AccountAvailability, + pub availability_reason: Option, + pub cooldown_until: Option>, + pub availability_observed_at: DateTime, + pub quota_observed_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// 单一货币的账号成本聚合。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountCost { + pub currency: String, + pub amount: super::observability::DecimalAmount, +} + +/// 账号在一个模型上的历史用量。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountModelUsage { + pub model: String, + pub request_count: u64, + pub success_count: u64, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_coverage: super::observability::CostCoverage, + pub costs: Vec, + pub last_used_at: DateTime, +} + +/// 账号历史用量聚合。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountUsage { + pub account_id: String, + pub request_count: u64, + pub success_count: u64, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_coverage: super::observability::CostCoverage, + pub costs: Vec, + pub last_used_at: Option>, + pub models: Vec, +} + +/// 账号列表页所需的完整存储事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountPage { + pub config_revision: Revision, + pub items: Vec, + pub total: u64, + pub summary: AccountSummary, +} + +/// 统一账号目录的全局状态计数，不受当前筛选和分页影响。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountSummary { + pub total: u64, + pub active: u64, + pub quota_exhausted: u64, + pub attention: u64, +} + +/// 账号启停写入命令。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SetAccountEnabled { + pub expected_config_revision: Revision, + pub account_id: String, + pub enabled: bool, +} + +/// 账号删除命令。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeleteAccount { + pub expected_config_revision: Revision, + pub account_id: String, +} + +/// 账号连接测试的语义事件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AccountConnectionTestEvent { + Started { + model: String, + }, + Request { + model: String, + input_text: String, + stream: bool, + store: bool, + }, + Content { + text: String, + }, + Completed { + account_status: AccountStatus, + }, + Failed { + message: String, + account_status: AccountStatus, + }, +} + +/// 每次连接测试独占的有限事件流。 +pub type AccountConnectionTestEventStream = + Pin + Send + 'static>>; diff --git a/backend/crates/gateway-admin/src/model/auth.rs b/backend/crates/gateway-admin/src/model/auth.rs new file mode 100644 index 000000000..6e989ae5a --- /dev/null +++ b/backend/crates/gateway-admin/src/model/auth.rs @@ -0,0 +1,104 @@ +//! 管理员认证、会话和安全审计事实。 + +use chrono::{DateTime, Utc}; + +use super::{MutationActor, MutationContext, Revision}; + +/// 已认证的管理主体。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AdminPrincipal { + Session { admin_user_id: String }, + ApiKey, +} + +/// 传给管理用例的安全请求上下文。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminRequestContext { + pub principal: AdminPrincipal, + pub request_id: String, +} + +impl AdminRequestContext { + #[must_use] + pub fn mutation_context(&self) -> MutationContext { + let actor = match &self.principal { + AdminPrincipal::Session { admin_user_id } => MutationActor::AdminSession { + admin_user_id: admin_user_id.clone(), + }, + AdminPrincipal::ApiKey => MutationActor::AdminApiKey, + }; + MutationContext { + actor, + request_id: self.request_id.clone(), + } + } +} + +/// 管理员登录命令。 +#[derive(Clone, PartialEq, Eq)] +pub struct LoginCommand { + pub username: Option, + pub password: String, + pub source: String, +} + +impl std::fmt::Debug for LoginCommand { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LoginCommand") + .field("username", &self.username) + .field("password", &"[REDACTED]") + .field("source", &self.source) + .finish() + } +} + +/// 登录成功后返回的会话事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LoginResult { + pub session_id: String, + pub expires_at: DateTime, +} + +/// 登录状态机可被 API 精确映射的失败类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum LoginError { + #[error("invalid administrator credentials")] + InvalidCredentials, + #[error("administrator login is throttled")] + Throttled, + #[error("administrator authentication is unavailable")] + Unavailable, +} + +/// Redis 中可恢复的管理员会话。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminSession { + pub admin_user_id: String, + pub expires_at: DateTime, +} + +/// 安全审计事件类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuditActorKind { + AdminSession, + AdminApiKey, + System, + Anonymous, +} + +/// 管理写操作留下的最小审计事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAuditEvent { + pub id: String, + pub actor_kind: AuditActorKind, + pub actor_admin_user_id: Option, + pub actor_ref: String, + pub request_id: Option, + pub action: String, + pub entity_kind: String, + pub entity_ref: String, + pub config_revision: Option, + pub changed_fields: Vec, + pub occurred_at: DateTime, +} diff --git a/backend/crates/gateway-admin/src/model/catalog.rs b/backend/crates/gateway-admin/src/model/catalog.rs new file mode 100644 index 000000000..c7c13885e --- /dev/null +++ b/backend/crates/gateway-admin/src/model/catalog.rs @@ -0,0 +1,89 @@ +//! Provider instance 目录的 Command 与 Result。 + +use chrono::{DateTime, Utc}; + +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; + +use super::{PageSize, Revision}; + +/// 一个可管理的 Provider instance。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstance { + pub id: ProviderInstanceId, + pub provider_kind: ProviderKind, + pub name: String, + pub base_url: String, + pub enabled: bool, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Provider instance 列表结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstanceCatalog { + pub config_revision: Revision, + pub items: Vec, +} + +/// Provider instance 详情及其所属控制面 revision。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstanceDetail { + pub config_revision: Revision, + pub item: ProviderInstance, +} + +/// Provider instance 的稳定游标查询。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CatalogListQuery { + pub cursor: Option, + pub page_size: PageSize, +} + +/// Provider instance 分页结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstancePage { + pub config_revision: Revision, + pub items: Vec, + pub next_cursor: Option, +} + +/// 创建 Provider instance。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CreateProviderInstance { + pub expected_config_revision: Revision, + pub id: ProviderInstanceId, + pub provider_kind: ProviderKind, + pub name: String, + pub base_url: String, +} + +/// 修改 Provider instance 的公共字段。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UpdateProviderInstance { + pub expected_config_revision: Revision, + pub id: ProviderInstanceId, + pub name: String, + pub base_url: String, +} + +/// 修改 Provider instance 的启用状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SetProviderInstanceEnabled { + pub expected_config_revision: Revision, + pub id: ProviderInstanceId, + pub enabled: bool, +} + +/// 删除 Provider instance。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeleteProviderInstance { + pub expected_config_revision: Revision, + pub id: ProviderInstanceId, +} + +/// Provider instance 写入后的状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstanceMutation { + pub config_revision: Revision, + pub instance: Option, +} diff --git a/backend/crates/gateway-admin/src/model/client_keys.rs b/backend/crates/gateway-admin/src/model/client_keys.rs new file mode 100644 index 000000000..f431147f8 --- /dev/null +++ b/backend/crates/gateway-admin/src/model/client_keys.rs @@ -0,0 +1,223 @@ +//! Client API Key 的 Command、Result 与安全秘密类型。 + +use std::{fmt, num::NonZeroU16}; + +use chrono::{DateTime, Utc}; + +use gateway_core::{ + policy::{ClientApiKeyId, RateLimits}, + routing::ProviderKind, +}; + +use super::{AdminModelError, Revision}; + +/// Client Key 列表保持旧 HTTP 合同允许的完整非零 `u16` 页大小。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ClientKeyPageSize(NonZeroU16); + +impl ClientKeyPageSize { + /// 创建 1 至 65535 的 Client Key 页大小。 + /// + /// # Errors + /// + /// `value` 为零时返回 [`AdminModelError::InvalidClientKeyPageSize`]。 + pub fn new(value: u16) -> Result { + NonZeroU16::new(value) + .map(Self) + .ok_or(AdminModelError::InvalidClientKeyPageSize) + } + + #[must_use] + pub const fn get(self) -> u16 { + self.0.get() + } +} + +/// Client Key 列表排序字段。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientKeySortField { + Name, + Enabled, + CreatedAt, + LastUsedAt, +} + +/// Client Key 列表排序方向。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SortDirection { + Asc, + Desc, +} + +/// Client Key 列表排序规则。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClientKeySort { + pub field: ClientKeySortField, + pub direction: SortDirection, +} + +/// 与排序字段绑定的游标值。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ClientKeyCursorValue { + Name(String), + Enabled(bool), + CreatedAt(DateTime), + LastUsedAt(Option>), +} + +/// Client Key 的稳定键集游标。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientKeyCursor { + pub sort: ClientKeySort, + pub value: ClientKeyCursorValue, + pub id: ClientApiKeyId, +} + +/// Client Key 列表查询。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientKeyListQuery { + pub cursor: Option, + pub page_size: ClientKeyPageSize, + pub search: Option, + pub sort: ClientKeySort, +} + +/// 不含完整明文 Key 的管理投影。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientKeyRecord { + pub id: ClientApiKeyId, + pub name: String, + pub label: Option, + pub provider_kind: ProviderKind, + pub prefix: String, + pub enabled: bool, + pub limits: RateLimits, + pub last_used_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Client Key 列表页。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientKeyPage { + pub config_revision: Revision, + pub items: Vec, + pub total: u64, + pub next_cursor: Option, +} + +/// 仅在创建或显式 reveal 时跨越管理边界的明文 Key。 +#[derive(Clone, PartialEq, Eq)] +pub struct ClientKeySecret { + pub record: ClientKeyRecord, + plaintext: String, +} + +impl ClientKeySecret { + #[must_use] + pub fn new(record: ClientKeyRecord, plaintext: impl Into) -> Self { + Self { + record, + plaintext: plaintext.into(), + } + } + + #[must_use] + pub fn expose_for_response(&self) -> &str { + &self.plaintext + } +} + +impl fmt::Debug for ClientKeySecret { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ClientKeySecret") + .field("record", &self.record) + .field("plaintext", &"[REDACTED]") + .finish() + } +} + +/// API 提交的 Client Key 创建命令。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CreateClientKey { + pub expected_config_revision: Revision, + pub name: String, + pub label: Option, + pub provider_kind: ProviderKind, + pub limits: RateLimits, +} + +/// 管理用例生成 ID 与明文后的持久化命令。 +#[derive(Clone, PartialEq, Eq)] +pub struct NewClientKey { + pub expected_config_revision: Revision, + pub id: ClientApiKeyId, + pub name: String, + pub label: Option, + pub provider_kind: ProviderKind, + pub limits: RateLimits, + pub plaintext: String, +} + +impl fmt::Debug for NewClientKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("NewClientKey") + .field("id", &self.id) + .field("name", &self.name) + .field("provider_kind", &self.provider_kind) + .field("plaintext", &"[REDACTED]") + .finish_non_exhaustive() + } +} + +/// 修改 Client Key 的公开策略字段。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UpdateClientKey { + pub expected_config_revision: Revision, + pub id: ClientApiKeyId, + pub name: String, + pub label: Option, + pub provider_kind: ProviderKind, + pub limits: RateLimits, +} + +/// 修改 Client Key 启用状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SetClientKeyEnabled { + pub expected_config_revision: Revision, + pub id: ClientApiKeyId, + pub enabled: bool, +} + +/// 删除 Client Key。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeleteClientKey { + pub expected_config_revision: Revision, + pub id: ClientApiKeyId, +} + +/// Client Key 创建结果；完整明文仅存在于该一次性结果中。 +pub struct CreatedClientKey { + pub config_revision: Revision, + pub secret: ClientKeySecret, +} + +impl fmt::Debug for CreatedClientKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CreatedClientKey") + .field("config_revision", &self.config_revision) + .field("secret", &self.secret) + .finish() + } +} + +/// Client Key 普通写操作结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientKeyMutation { + pub config_revision: Revision, + pub record: Option, + pub id: ClientApiKeyId, +} diff --git a/backend/crates/gateway-admin/src/model/mod.rs b/backend/crates/gateway-admin/src/model/mod.rs new file mode 100644 index 000000000..bfa3d08ca --- /dev/null +++ b/backend/crates/gateway-admin/src/model/mod.rs @@ -0,0 +1,156 @@ +//! 管理控制面使用的 Command、Result 与稳定值对象。 + +use std::num::{NonZeroU16, NonZeroU64}; + +pub mod accounts; +pub mod auth; +pub mod catalog; +pub mod client_keys; +pub mod observability; +pub mod provider_credentials; +pub mod settings; +pub mod system; + +/// 管理用例对外返回的稳定错误分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminErrorKind { + Invalid, + Unauthorized, + NotFound, + Conflict, + RateLimited, + BadGateway, + Unavailable, + Internal, +} + +/// 不携带基础设施细节的管理用例错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{message}")] +pub struct AdminError { + kind: AdminErrorKind, + message: String, +} + +impl AdminError { + #[must_use] + pub fn new(kind: AdminErrorKind, message: impl Into) -> Self { + Self { + kind, + message: message.into(), + } + } + + #[must_use] + pub const fn kind(&self) -> AdminErrorKind { + self.kind + } + + #[must_use] + pub fn invalid(message: impl Into) -> Self { + Self::new(AdminErrorKind::Invalid, message) + } + + #[must_use] + pub fn not_found(message: impl Into) -> Self { + Self::new(AdminErrorKind::NotFound, message) + } + + #[must_use] + pub fn conflict(message: impl Into) -> Self { + Self::new(AdminErrorKind::Conflict, message) + } + + #[must_use] + pub fn unavailable(message: impl Into) -> Self { + Self::new(AdminErrorKind::Unavailable, message) + } + + #[must_use] + pub fn internal(message: impl Into) -> Self { + Self::new(AdminErrorKind::Internal, message) + } +} + +/// PostgreSQL 中所有正整数 revision 的管理层表示。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Revision(NonZeroU64); + +impl Revision { + /// 创建正整数 revision。 + /// + /// # Errors + /// + /// `value` 为零时返回 [`AdminModelError::ZeroRevision`]。 + pub fn new(value: u64) -> Result { + NonZeroU64::new(value) + .map(Self) + .ok_or(AdminModelError::ZeroRevision) + } + + #[must_use] + pub const fn get(self) -> u64 { + self.0.get() + } +} + +/// 管理列表统一使用的受限页大小。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PageSize(NonZeroU16); + +impl PageSize { + pub const MAX: u16 = 200; + + /// 创建 1 至 200 的页大小。 + /// + /// # Errors + /// + /// `value` 不在有效范围时返回 [`AdminModelError::InvalidPageSize`]。 + pub fn new(value: u16) -> Result { + if value > Self::MAX { + return Err(AdminModelError::InvalidPageSize(value)); + } + NonZeroU16::new(value) + .map(Self) + .ok_or(AdminModelError::InvalidPageSize(value)) + } + + #[must_use] + pub const fn get(self) -> u16 { + self.0.get() + } +} + +/// 可审计管理写操作的发起者。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MutationActor { + AdminSession { admin_user_id: String }, + AdminApiKey, + System, +} + +/// 管理写操作必须携带的审计上下文。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MutationContext { + pub actor: MutationActor, + pub request_id: String, +} + +/// 领域值对象构造失败。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum AdminModelError { + #[error("revision must be greater than zero")] + ZeroRevision, + #[error("page size {0} is outside 1..=200")] + InvalidPageSize(u16), + #[error("client key page size must be inside 1..=65535")] + InvalidClientKeyPageSize, + #[error("request outcome must be 1..=256 bytes without control characters")] + InvalidRequestOutcome, + #[error("time range must be positive and no longer than 366 days")] + InvalidTimeRange, + #[error("decimal amount is not a non-negative numeric(20,10) value")] + InvalidDecimalAmount, + #[error("latency percentile must be finite and non-negative")] + InvalidLatencyPercentile, +} diff --git a/backend/crates/gateway-admin/src/model/observability.rs b/backend/crates/gateway-admin/src/model/observability.rs new file mode 100644 index 000000000..c815cae29 --- /dev/null +++ b/backend/crates/gateway-admin/src/model/observability.rs @@ -0,0 +1,955 @@ +//! 用量、成本、健康与错误诊断的 UTC 语义事实。 + +use std::{num::NonZeroU32, str::FromStr}; + +use chrono::{DateTime, TimeDelta, Utc}; + +use super::{AdminModelError, PageSize}; + +/// 外部观测查询的 UTC 时间范围。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TimeRange { + pub start: DateTime, + pub end: DateTime, +} + +impl TimeRange { + /// 创建最长 366 天的正时间范围。 + /// + /// # Errors + /// + /// 范围为空、反向或超过 366 天时返回错误。 + pub fn new(start: DateTime, end: DateTime) -> Result { + let duration = end.signed_duration_since(start); + if duration <= TimeDelta::zero() || duration > TimeDelta::days(366) { + return Err(AdminModelError::InvalidTimeRange); + } + Ok(Self { start, end }) + } +} + +/// 从一开始的观测页码。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PageNumber(NonZeroU32); + +impl PageNumber { + #[must_use] + pub const fn new(value: NonZeroU32) -> Self { + Self(value) + } + + #[must_use] + pub const fn get(self) -> u32 { + self.0.get() + } +} + +/// 观测记录的稳定键集游标。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ObservabilityCursor { + pub observed_at: DateTime, + pub stable_id: String, +} + +/// 请求结果状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RequestOutcome { + Running, + Succeeded, + Failed, + Cancelled, + Incomplete, + Other(OtherRequestOutcome), +} + +/// 未知但有界的请求结果值；私有字段禁止绕过 [`RequestOutcome::new`] 构造。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OtherRequestOutcome(String); + +impl RequestOutcome { + pub const MAX_BYTES: usize = 256; + + /// 从持久化或 wire 字符串创建可扩展结果语义。 + /// + /// # Errors + /// + /// 空值、超过 256 字节或含控制字符时返回 + /// [`AdminModelError::InvalidRequestOutcome`]。 + pub fn new(value: impl Into) -> Result { + let value = value.into(); + match value.as_str() { + "running" => Ok(Self::Running), + "succeeded" => Ok(Self::Succeeded), + "failed" => Ok(Self::Failed), + "cancelled" => Ok(Self::Cancelled), + "incomplete" => Ok(Self::Incomplete), + _ if value.trim().is_empty() + || value.len() > Self::MAX_BYTES + || value.chars().any(char::is_control) => + { + Err(AdminModelError::InvalidRequestOutcome) + } + _ => Ok(Self::Other(OtherRequestOutcome(value))), + } + } + + #[must_use] + pub fn as_str(&self) -> &str { + match self { + Self::Running => "running", + Self::Succeeded => "succeeded", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::Incomplete => "incomplete", + Self::Other(value) => value.as_str(), + } + } +} + +impl OtherRequestOutcome { + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// 用量记录过滤条件。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct UsageFilter { + pub client_api_key_ref: Option, + pub request_id: Option, + pub provider_account_ref: Option, + pub operation: Option, + pub provider_kind: Option, + pub model: Option, + pub outcome: Option, + pub status_code: Option, + pub transport: Option, + pub attempt_index: Option, + pub response_id: Option, + pub upstream_request_id: Option, + pub search: Option, +} + +/// 用量记录分页查询。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageQuery { + pub range: TimeRange, + pub filter: UsageFilter, + pub cursor: Option, + pub page: PageNumber, + pub page_size: PageSize, +} + +/// 运维错误过滤条件。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct OpsErrorFilter { + pub client_api_key_ref: Option, + pub request_id: Option, + pub provider_kind: Option, + pub provider_account_ref: Option, + pub operation: Option, + pub model: Option, + pub transport: Option, + pub attempt_index: Option, + pub response_id: Option, + pub upstream_request_id: Option, + pub failure_kind: Option, + pub status_code: Option, + pub search: Option, +} + +/// 运维错误分页查询。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpsErrorQuery { + pub range: TimeRange, + pub filter: OpsErrorFilter, + pub cursor: Option, + pub page: PageNumber, + pub page_size: PageSize, +} + +/// 用量诊断维度。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DiagnosticDimension { + Provider, + Model, + Account, + ApiKey, + Transport, + Failure, + Status, +} + +/// `numeric(20,10)` 的非负规范金额。 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct DecimalAmount(String); + +impl DecimalAmount { + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } + + /// 精确相加两个 `numeric(20,10)` 金额。 + #[must_use] + pub fn checked_add(&self, other: &Self) -> Option { + let sum = scaled_amount(&self.0)?.checked_add(scaled_amount(&other.0)?)?; + if sum >= 10_u128.pow(20) { + return None; + } + let whole = sum / 10_u128.pow(10); + let fraction = sum % 10_u128.pow(10); + let value = if fraction == 0 { + whole.to_string() + } else { + format!("{whole}.{fraction:010}") + .trim_end_matches('0') + .to_owned() + }; + Some(Self(value)) + } +} + +fn scaled_amount(value: &str) -> Option { + let (whole, fraction) = value.split_once('.').unwrap_or((value, "")); + let whole = whole.parse::().ok()?; + let fraction = format!("{fraction:0<10}").parse::().ok()?; + whole.checked_mul(10_u128.pow(10))?.checked_add(fraction) +} + +impl FromStr for DecimalAmount { + type Err = AdminModelError; + + fn from_str(input: &str) -> Result { + let input = input.trim(); + let mut parts = input.split('.'); + let whole = parts.next().unwrap_or_default(); + let fraction = parts.next(); + let valid = !whole.is_empty() + && whole.len() <= 10 + && whole.bytes().all(|byte| byte.is_ascii_digit()) + && parts.next().is_none() + && fraction.is_none_or(|value| { + !value.is_empty() + && value.len() <= 10 + && value.bytes().all(|byte| byte.is_ascii_digit()) + }); + if !valid { + return Err(AdminModelError::InvalidDecimalAmount); + } + let whole = whole.trim_start_matches('0'); + let whole = if whole.is_empty() { "0" } else { whole }; + let fraction = fraction.unwrap_or_default().trim_end_matches('0'); + let canonical = if fraction.is_empty() { + whole.to_owned() + } else { + format!("{whole}.{fraction}") + }; + Ok(Self(canonical)) + } +} + +impl std::fmt::Display for DecimalAmount { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} + +/// 单一币种的成本合计。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CurrencyCost { + pub currency: String, + pub amount: DecimalAmount, +} + +/// PostgreSQL 连续百分位返回的非负有限毫秒值。 +/// +/// 以 IEEE-754 bits 保存，既不丢失插值小数，也可安全实现 `Eq`。 +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub struct PercentileMilliseconds(u64); + +impl PercentileMilliseconds { + /// 校验并保存一个毫秒百分位。 + /// + /// # Errors + /// + /// 非有限值或负值返回 [`AdminModelError::InvalidLatencyPercentile`]。 + pub fn new(value: f64) -> Result { + if !value.is_finite() || value < 0.0 { + return Err(AdminModelError::InvalidLatencyPercentile); + } + Ok(Self(value.to_bits())) + } + + #[must_use] + pub const fn as_f64(self) -> f64 { + f64::from_bits(self.0) + } +} + +impl std::fmt::Debug for PercentileMilliseconds { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.as_f64().fmt(formatter) + } +} + +/// 延迟分布的三个稳定百分位。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct LatencyPercentiles { + pub p50_ms: Option, + pub p95_ms: Option, + pub p99_ms: Option, +} + +/// Provider 价格规则计算所需的持久请求事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderBillingInput { + pub upstream_model_id: String, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub total: CurrencyCost, +} + +/// Provider 已确认的逐项费用与单价。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CalculatedBillingBreakdown { + pub input_amount: CurrencyCost, + pub output_amount: CurrencyCost, + pub cache_read_amount: CurrencyCost, + pub cache_write_amount: CurrencyCost, + pub standard_amount: CurrencyCost, + pub total_amount: CurrencyCost, + pub input_price_per_million: CurrencyCost, + pub output_price_per_million: CurrencyCost, + pub cache_read_price_per_million: CurrencyCost, + pub cache_write_price_per_million: CurrencyCost, + pub service_tier: Option, + pub multiplier_percent: u32, +} + +/// 单次请求的费用语义。 +/// +/// Provider 上报费用或无法恢复逐项价格时保留总额；Provider 验证成功后升级为完整分解。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UsageBilling { + Total { source: String, total: CurrencyCost }, + Calculated(Box), +} + +/// 计费数据覆盖情况。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CostCoverage { + pub provider_reported_count: u64, + pub calculated_count: u64, + pub partial_count: u64, + pub unavailable_count: u64, + pub not_billable_count: u64, +} + +impl CostCoverage { + /// Provider 直接上报或 Provider 规则完整计算的已知成本数。 + #[must_use] + pub fn known_count(&self) -> u64 { + self.provider_reported_count + .saturating_add(self.calculated_count) + } +} + +/// 请求级聚合指标。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RequestMetrics { + pub request_count: u64, + pub success_count: u64, + pub failure_count: u64, + pub cancelled_count: u64, + pub incomplete_count: u64, + pub caller_error_count: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub cache_write_tokens: u64, + pub reasoning_tokens: u64, + pub total_tokens: u64, + pub first_token_latency_sum_ms: u64, + pub first_token_latency_count: u64, + pub latency_sum_ms: u64, + pub latency_count: u64, + pub min_latency_ms: Option, + pub max_latency_ms: Option, + pub latency_percentiles: LatencyPercentiles, + pub first_token_latency_percentiles: LatencyPercentiles, + pub cache_eligible_request_count: u64, + pub cache_hit_request_count: u64, +} + +/// 上游 attempt 聚合指标。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AttemptMetrics { + pub attempt_count: u64, + pub success_count: u64, + pub failure_count: u64, + pub cancelled_count: u64, + pub incomplete_count: u64, + pub rate_limited_count: u64, + pub auth_failure_count: u64, + pub provider_5xx_count: u64, + pub cost_coverage: CostCoverage, + pub costs: Vec, +} + +/// 时间序列粒度。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Granularity { + FifteenMinutes, + Hour, + Day, +} + +/// 一段时间桶内的请求指标。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RequestMetricPoint { + pub bucket_start: DateTime, + pub granularity: Granularity, + pub metrics: RequestMetrics, + pub cost_coverage: CostCoverage, + pub costs: Vec, +} + +/// 账号池容量统计。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AccountPoolMetrics { + pub total: u64, + pub enabled: u64, + pub unavailable: u64, + pub active: u64, + pub expired: u64, + pub quota_exhausted: u64, + pub refreshing: Option, + pub disabled: u64, + pub banned: u64, +} + +/// Dashboard 中一个账号的模型级用量事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardAccountModelUsage { + pub model: String, + pub request_count: u64, + pub success_count: u64, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_coverage: CostCoverage, + pub costs: Vec, + pub last_used_at: DateTime, +} + +/// Dashboard 中一个账号的完整用量与公共状态事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardAccountUsage { + pub account_id: String, + pub provider_instance_id: String, + pub provider_kind: String, + pub name: String, + pub email: Option, + pub plan_type: Option, + pub enabled: bool, + pub availability: String, + pub request_count: u64, + pub success_count: u64, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_coverage: CostCoverage, + pub costs: Vec, + pub last_used_at: Option>, + pub models: Vec, +} + +/// 仪表盘所需的公共观测事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardObservation { + pub range: TimeRange, + pub requests: RequestMetrics, + pub attempts: AttemptMetrics, + pub provider_accounts: AccountPoolMetrics, + pub trend: Vec, + pub account_usage: Vec, + pub recent_requests: Vec, +} + +/// 一次完整模型请求的公共观测记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageRecord { + pub id: String, + pub client_api_key_ref: String, + pub config_revision: u64, + pub protocol: String, + pub operation: String, + pub endpoint: String, + pub client_transport: String, + pub requested_model_id: String, + pub input_token_estimate: u64, + pub provider_instance_id: Option, + pub provider_kind: Option, + pub provider_account_ref: Option, + pub provider_account_name: Option, + pub provider_account_email: Option, + pub upstream_model_id: Option, + pub upstream_transport: Option, + pub http_version: Option, + pub attempt_count: u32, + pub upstream_send_state: String, + pub downstream_committed_at: Option>, + pub outcome: RequestOutcome, + pub client_status_code: Option, + pub upstream_status_code: Option, + pub client_response_id: Option, + pub upstream_request_id: Option, + pub upstream_response_id: Option, + pub error_kind: Option, + pub provider_error_code: Option, + pub error_message: Option, + pub retry_after_ms: Option, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_source: String, + pub cost_amount: Option, + pub cost_currency: Option, + pub billing: Option, + pub transport_decision_wait_ms: Option, + pub connect_ms: Option, + pub headers_ms: Option, + pub first_event_ms: Option, + pub first_reasoning_ms: Option, + pub first_text_ms: Option, + pub first_token_ms: Option, + pub provider_processing_ms: Option, + pub latency_ms: Option, + pub client_ip: Option, + pub user_agent: Option, + pub reasoning_effort: Option, + pub reasoning_preset: Option, + pub request_kind: Option, + pub subagent_kind: Option, + pub compact: bool, + pub started_at: DateTime, + pub deadline_at: DateTime, + pub completed_at: Option>, +} + +/// 用量分页结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsagePage { + pub items: Vec, + pub total: u64, + pub next_cursor: Option, +} + +/// 请求中的一次上游尝试或运维事件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageAttempt { + pub source: String, + pub id: String, + pub attempt_index: u32, + pub component: String, + pub operation: String, + pub provider_instance_id: Option, + pub provider_kind: Option, + pub provider_account_ref: Option, + pub upstream_model_id: Option, + pub upstream_transport: Option, + pub upstream_send_state: Option, + pub outcome: RequestOutcome, + pub downstream_committed: bool, + pub status_code: Option, + pub provider_error_code: Option, + pub failure_kind: Option, + pub retry_after_ms: Option, + pub upstream_request_id: Option, + pub latency_ms: Option, + pub message: Option, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_tokens: Option, + pub cache_write_tokens: Option, + pub reasoning_tokens: Option, + pub total_tokens: Option, + pub cost_source: Option, + pub cost_amount: Option, + pub cost_currency: Option, + pub occurred_at: DateTime, +} + +/// 一条请求及其全部尝试。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageDetail { + pub request: UsageRecord, + pub attempts: Vec, +} + +/// 一个 Provider 的聚合用量。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderObservation { + pub provider_kind: String, + pub request_count: u64, + pub attempt_count: u64, + pub failure_count: u64, + pub total_tokens: u64, +} + +/// 用量总览。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageOverview { + pub range: TimeRange, + pub requests: RequestMetrics, + pub attempts: AttemptMetrics, + pub providers: Vec, +} + +/// 用量摘要的用例结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageSummary { + pub overview: UsageOverview, + pub average_latency_ms: Option, +} + +/// 单个诊断维度值的聚合结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DiagnosticObservation { + pub name: String, + pub request_count: u64, + pub success_count: u64, + pub failure_count: u64, + pub attempt_count: u64, + pub total_tokens: u64, + pub average_latency_ms: Option, + pub cost_coverage: CostCoverage, + pub costs: Vec, +} + +/// 统一运维错误记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpsError { + pub source: String, + pub event_id: String, + pub request_id: Option, + pub attempt_index: Option, + pub client_api_key_ref: Option, + pub component: String, + pub operation: String, + pub provider_instance_id: Option, + pub provider_kind: Option, + pub provider_account_ref: Option, + pub upstream_model_id: Option, + pub upstream_transport: Option, + pub failure_kind: String, + pub status_code: Option, + pub provider_error_code: Option, + pub client_response_id: Option, + pub upstream_request_id: Option, + pub latency_ms: Option, + pub message: String, + pub occurrence_count: u32, + pub occurred_at: DateTime, + pub stable_sort_id: String, +} + +/// 运维错误分页结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpsErrorPage { + pub items: Vec, + pub total: u64, + pub next_cursor: Option, +} + +/// 仪表盘趋势指标。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TrendKind { + Usage, + Latency, + Errors, +} + +/// 一个趋势桶的全部已计算语义，API 只负责时区标签与展示格式。 +#[derive(Debug, Clone, PartialEq)] +pub struct TrendPoint { + pub bucket_start: DateTime, + pub granularity: Granularity, + pub metrics: RequestMetrics, + pub service_failure_count: u64, + pub average_latency_ms: Option, + pub average_first_token_latency_ms: Option, + pub cached_token_rate: f64, + pub cache_hit_request_rate: Option, + pub success_rate: Option, + pub cost_coverage: CostCoverage, + pub costs: Vec, +} + +/// 趋势的无格式化汇总。 +#[derive(Debug, Clone, PartialEq)] +pub struct TrendSummary { + pub request_count: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub service_failure_count: u64, + pub average_latency_ms: Option, + pub max_latency_ms: Option, + pub min_latency_ms: Option, + pub success_rate: Option, + pub cache_hit_request_rate: Option, + pub costs: Vec, + pub cost_coverage: CostCoverage, +} + +/// 趋势结果；API 再决定图表标签与数字格式。 +#[derive(Debug, Clone, PartialEq)] +pub struct Trend { + pub kind: TrendKind, + pub points: Vec, + pub summary: TrendSummary, +} + +/// 健康时间桶状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HealthStatus { + Future, + NoData, + Unavailable, + LowSample, + Unstable, + Stable, +} + +/// 一个 15 分钟健康桶的语义结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct HealthTimelinePoint { + pub bucket_start: DateTime, + pub status: HealthStatus, + pub reliability_percent: Option, + pub success_requests: u64, + pub failed_requests: u64, + pub cancelled_requests: u64, + pub incomplete_requests: u64, + pub caller_error_requests: u64, +} + +/// 中国自然日的固定 96 槽健康时间线。 +#[derive(Debug, Clone, PartialEq)] +pub struct HealthTimeline { + pub reliability_percent: Option, + pub status: HealthStatus, + pub success_requests: u64, + pub failed_requests: u64, + pub cancelled_requests: u64, + pub incomplete_requests: u64, + pub caller_error_requests: u64, + pub points: Vec, +} + +/// Dashboard 请求画像的目标平台。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardWireTarget { + pub os_type: String, + pub os_version: String, + pub arch: String, + pub terminal: String, +} + +/// Desktop 发布检查状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DesktopReleaseStatus { + Unchecked, + Current, + UpdateAvailable, + Failed, +} + +/// 与请求画像分离的 Desktop 发布检查事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardDesktopRelease { + pub status: DesktopReleaseStatus, + pub checked_at: Option>, + pub latest_version: Option, + pub latest_build: Option, + pub published_at: Option>, + pub minimum_system_version: Option, + pub hardware_requirements: Option, + pub download_url: Option, + pub download_size: Option, + pub signature_present: Option, + pub error: Option, +} + +/// Provider 拥有并实际用于上游请求的 wire 画像快照。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardWireProfile { + pub originator: String, + pub codex_version: String, + pub desktop_version: String, + pub desktop_build: String, + pub target: DashboardWireTarget, + pub user_agent: String, + pub verified_at: DateTime, + pub release: DashboardDesktopRelease, +} + +/// 当前账号池的并发调度容量。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DashboardCapacity { + pub max_concurrent_per_account: u64, + pub total_slots: u64, + pub used_slots: Option, + pub available_slots: Option, +} + +/// Dashboard 某个自然日区间内的卡片计数。 +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct DashboardPeriodMetrics { + pub request_count: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub cached_token_rate: f64, + pub observed_cached_token_rate: Option, +} + +/// 仪表盘聚合结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct DashboardResult { + pub observation: DashboardObservation, + pub today: DashboardPeriodMetrics, + pub yesterday: DashboardPeriodMetrics, + pub total_billing_usd: Option, + pub total_cached_token_rate: f64, + pub average_first_token_latency_ms: Option, + pub trend: Trend, + pub health_timeline: HealthTimeline, + pub wire_profile: DashboardWireProfile, + pub capacity: DashboardCapacity, + pub rotation_strategy: super::settings::RotationStrategy, +} + +/// 洞察健康时间点的语义结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsightsHealthPoint { + pub bucket_start: DateTime, + pub success_requests: u64, + pub failed_requests: u64, + pub cancelled_requests: u64, + pub incomplete_requests: u64, + pub caller_error_requests: u64, + pub error_rate: f64, +} + +/// 洞察健康汇总；失败数已排除调用方错误。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsightsHealth { + pub total_requests: u64, + pub success_requests: u64, + pub failed_requests: u64, + pub cancelled_requests: u64, + pub incomplete_requests: u64, + pub caller_error_requests: u64, + pub success_rate: f64, + pub points: Vec, +} + +/// 洞察性能时间点。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageInsightsPerformancePoint { + pub bucket_start: DateTime, + pub latency_percentiles: LatencyPercentiles, + pub first_token_latency_percentiles: LatencyPercentiles, +} + +/// 洞察性能汇总与可观测覆盖率。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsightsPerformance { + pub latency_percentiles: LatencyPercentiles, + pub first_token_latency_percentiles: LatencyPercentiles, + pub latency_coverage: f64, + pub first_token_coverage: f64, + pub points: Vec, +} + +/// 洞察成本时间点。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsightsCostPoint { + pub bucket_start: DateTime, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub estimated_cost: Option, + pub standard_cost: Option, + pub cached_token_rate: f64, + pub cache_hit_request_rate: Option, +} + +/// 洞察成本汇总。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsightsCost { + pub estimated_cost: Option, + pub standard_cost: Option, + pub cost_per_request: Option, + pub tokens_per_request: f64, + pub cached_token_rate: f64, + pub cache_hit_request_rate: Option, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub points: Vec, + pub costs: Vec, + pub coverage: CostCoverage, +} + +/// 用量洞察的完整用例结果；API 只负责标签、时区与字符串格式化。 +#[derive(Debug, Clone, PartialEq)] +pub struct UsageInsights { + pub granularity: Granularity, + pub health: UsageInsightsHealth, + pub performance: UsageInsightsPerformance, + pub cost: UsageInsightsCost, + pub attempts: AttemptMetrics, + pub providers: Vec, +} + +/// 已完成分母计算的诊断项。 +#[derive(Debug, Clone, PartialEq)] +pub struct DiagnosticsItem { + pub name: String, + pub request_count: u64, + pub success_count: u64, + pub error_count: u64, + pub error_rate: f64, + pub request_share: f64, + pub average_latency_ms: Option, + pub estimated_cost: Option, + pub attempt_count: u64, + pub total_tokens: u64, +} + +/// 诊断结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct DiagnosticsResult { + pub dimension: DiagnosticDimension, + pub items: Vec, +} diff --git a/backend/crates/gateway-admin/src/model/provider_credentials.rs b/backend/crates/gateway-admin/src/model/provider_credentials.rs new file mode 100644 index 000000000..3f8833ad1 --- /dev/null +++ b/backend/crates/gateway-admin/src/model/provider_credentials.rs @@ -0,0 +1,643 @@ +//! Provider 管理能力交换的中立 Command 与 Result。 + +use std::fmt; + +use chrono::{DateTime, Utc}; +use gateway_core::{ + engine::credential::{OpaqueProviderData, ProviderAccountId}, + routing::{ProviderInstanceId, ProviderKind, UpstreamModelId}, +}; + +use super::{ + MutationActor, MutationContext, PageSize, Revision, + accounts::{AccountAvailability, AccountRecord, AccountStatus, AccountSummary}, +}; + +/// Provider-owned JSON；公共层只搬运且 Debug 不输出值。 +#[derive(Clone, PartialEq)] +pub struct ProviderDocument(OpaqueProviderData); + +impl ProviderDocument { + #[must_use] + pub const fn new(data: OpaqueProviderData) -> Self { + Self(data) + } + + /// 仅具体 Provider 可以解释内部字段。 + #[must_use] + pub const fn expose_to_provider(&self) -> &OpaqueProviderData { + &self.0 + } + + #[must_use] + pub fn into_provider_data(self) -> OpaqueProviderData { + self.0 + } +} + +impl fmt::Debug for ProviderDocument { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderDocument([PROVIDER_OWNED])") + } +} + +/// Credential 列表稳定游标。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialCursor { + pub created_at: DateTime, + pub account_id: ProviderAccountId, +} + +/// Provider credential 列表查询。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialListQuery { + pub provider_instance_id: Option, + pub availability: Option, + pub enabled: Option, + pub window: CredentialListWindow, +} + +/// Credential 可用性筛选；支持一个 wire 值对应多个持久状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CredentialAvailabilityFilter { + Exact(AccountAvailability), + AnyOf(Vec), +} + +impl CredentialAvailabilityFilter { + #[must_use] + pub fn matches(&self, availability: AccountAvailability) -> bool { + match self { + Self::Exact(expected) => *expected == availability, + Self::AnyOf(expected) => expected.contains(&availability), + } + } +} + +/// Credential 目录的集合窗口；完整列表与游标分页互斥。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CredentialListWindow { + All, + Page { + cursor: Option, + page_size: PageSize, + }, +} + +/// Provider credential 列表。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialPage { + pub config_revision: Revision, + pub items: Vec, + pub next_cursor: Option, +} + +/// Provider credential 详情。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialDetails { + pub config_revision: Revision, + pub credential: AccountRecord, +} + +/// Provider 正式文档批量导入命令。 +pub struct ImportCredentials { + pub context: MutationContext, + pub expected_config_revision: Revision, + pub provider_instance_id: ProviderInstanceId, + pub document: ProviderDocument, +} + +impl fmt::Debug for ImportCredentials { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ImportCredentials") + .field("context", &self.context) + .field("expected_config_revision", &self.expected_config_revision) + .field("provider_instance_id", &self.provider_instance_id) + .field("document", &self.document) + .finish() + } +} + +/// 批量导入提交结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialImportResult { + pub config_revision: Revision, + pub credential_ids: Vec, +} + +/// Provider 解析导入文档时只接收实例与不透明文档，不接触 revision 或审计上下文。 +pub struct PrepareCredentialImport { + pub provider_instance_id: ProviderInstanceId, + pub document: ProviderDocument, +} + +impl fmt::Debug for PrepareCredentialImport { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PrepareCredentialImport") + .field("provider_instance_id", &self.provider_instance_id) + .field("document", &self.document) + .finish() + } +} + +/// Provider 已验证、可由 Store 原子创建的一份 credential。 +#[derive(Debug, Clone, PartialEq)] +pub struct PreparedCredentialCreate { + pub account_id: ProviderAccountId, + pub provider_instance_id: ProviderInstanceId, + pub provider_kind: ProviderKind, + pub name: String, + pub email: Option, + pub upstream_user_id: String, + pub upstream_account_id: Option, + pub plan_type: Option, + pub provider_material: ProviderDocument, + pub has_refresh_token: bool, + pub access_token_expires_at: DateTime, + pub next_refresh_at: Option>, + pub enabled: bool, + pub availability: AccountAvailability, + pub availability_reason: Option, + pub cooldown_until: Option>, + pub availability_observed_at: DateTime, +} + +/// Provider 对一份导入文档的完整验证结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct PreparedCredentialImport { + pub provider_kind: ProviderKind, + pub provider_instance_id: ProviderInstanceId, + pub credentials: Vec, +} + +/// Admin 交给 Store 的导入事务命令。 +#[derive(Debug, Clone, PartialEq)] +pub struct CredentialImportCommit { + pub expected_config_revision: Revision, + pub prepared: PreparedCredentialImport, +} + +/// 可选的重新授权目标。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReauthorizationTarget { + pub account_id: ProviderAccountId, + pub credential_revision: Revision, +} + +/// OAuth pending owner 的中立身份；不编码具体 Provider 的 Redis key 或 JSON。 +#[derive(Clone, PartialEq, Eq)] +pub enum AuthorizationOwner { + AdminSession { admin_user_id: String }, + AdminApiKey, + System, +} + +impl fmt::Debug for AuthorizationOwner { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("AuthorizationOwner([REDACTED])") + } +} + +/// Provider 写入 pending payload 与 Redis owner binding 所需的全部中立字段。 +#[derive(Clone, PartialEq, Eq)] +pub struct AuthorizationOwnerBinding { + owner: AuthorizationOwner, + started_request_id: String, +} + +impl AuthorizationOwnerBinding { + #[must_use] + pub fn from_context(context: &MutationContext) -> Self { + let owner = match &context.actor { + MutationActor::AdminSession { admin_user_id } => AuthorizationOwner::AdminSession { + admin_user_id: admin_user_id.clone(), + }, + MutationActor::AdminApiKey => AuthorizationOwner::AdminApiKey, + MutationActor::System => AuthorizationOwner::System, + }; + Self { + owner, + started_request_id: context.request_id.clone(), + } + } + + #[must_use] + pub const fn owner(&self) -> &AuthorizationOwner { + &self.owner + } + + #[must_use] + pub fn started_request_id(&self) -> &str { + &self.started_request_id + } + + #[must_use] + pub fn matches_context(&self, context: &MutationContext) -> bool { + Self::from_context(context).owner == self.owner + } +} + +impl fmt::Debug for AuthorizationOwnerBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("AuthorizationOwnerBinding([REDACTED])") + } +} + +/// OAuth 完成时应创建新账号还是 CAS 更新既有 credential。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AuthorizationMutationTarget { + Create { + provider_instance_id: ProviderInstanceId, + name: String, + }, + Reauthorize { + provider_instance_id: ProviderInstanceId, + account_id: ProviderAccountId, + expected_credential_revision: Revision, + }, +} + +/// 必须完整进入 Provider opaque pending payload、并在 complete 后原样恢复的事务信封。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PendingAuthorizationMutation { + expected_config_revision: Revision, + provider_kind: ProviderKind, + target: AuthorizationMutationTarget, + owner_binding: AuthorizationOwnerBinding, +} + +impl PendingAuthorizationMutation { + #[must_use] + pub const fn new( + expected_config_revision: Revision, + provider_kind: ProviderKind, + target: AuthorizationMutationTarget, + owner_binding: AuthorizationOwnerBinding, + ) -> Self { + Self { + expected_config_revision, + provider_kind, + target, + owner_binding, + } + } + + #[must_use] + pub const fn expected_config_revision(&self) -> Revision { + self.expected_config_revision + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn target(&self) -> &AuthorizationMutationTarget { + &self.target + } + + #[must_use] + pub const fn owner_binding(&self) -> &AuthorizationOwnerBinding { + &self.owner_binding + } +} + +/// 启动 Provider OAuth Authorization Code 流程。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StartAuthorization { + pub context: MutationContext, + pub expected_config_revision: Revision, + pub provider_instance_id: ProviderInstanceId, + pub name: String, + pub reauthorization: Option, +} + +/// OAuth 流程启动结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AuthorizationStarted { + pub flow_id: String, + pub authorization_url: String, + pub expires_at: DateTime, +} + +/// refresh/rotate/reauthorize 的 lease 或 completion 生命周期。 +/// +/// 该 guard 不可 Clone。Admin 只能在 Store CAS 与审计事务成功后调用 `finish`；失败路径直接 +/// drop，使 Provider 可以释放 lease 或执行补偿。 +pub trait CredentialCommitGuard: Send + 'static { + fn finish(self: Box); +} + +/// OAuth complete 后由 Provider 返回的准备结果；Store 仍是唯一提交者。 +pub enum PreparedAuthorizationCredential { + Create(PreparedCredentialCreate), + Reauthorize(PreparedCredentialRotation), +} + +/// Provider 从 opaque pending payload 恢复的信封与已验证 credential 必须一起返回。 +pub struct PreparedAuthorizationCommit { + pub pending: PendingAuthorizationMutation, + pub credential: PreparedAuthorizationCredential, +} + +impl fmt::Debug for PreparedAuthorizationCommit { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PreparedAuthorizationCommit") + .field("pending", &self.pending) + .field("credential", &"[PREPARED]") + .finish() + } +} + +/// Store 可持久化的 OAuth credential facts，不携带 Provider guard。 +#[derive(Debug, Clone, PartialEq)] +pub enum AuthorizationCredentialCommit { + Create(PreparedCredentialCreate), + Reauthorize(PreparedCredentialRotationFacts), +} + +/// Admin 交给 Store 的 OAuth 原子事务命令。 +#[derive(Debug, Clone, PartialEq)] +pub struct AuthorizationCommit { + pub pending: PendingAuthorizationMutation, + pub credential: AuthorizationCredentialCommit, +} + +impl PreparedAuthorizationCommit { + #[must_use] + pub fn into_commit(self) -> (AuthorizationCommit, Option>) { + let (credential, guard) = match self.credential { + PreparedAuthorizationCredential::Create(credential) => { + (AuthorizationCredentialCommit::Create(credential), None) + } + PreparedAuthorizationCredential::Reauthorize(prepared) => { + let (facts, guard) = prepared.into_parts(); + ( + AuthorizationCredentialCommit::Reauthorize(facts), + Some(guard), + ) + } + }; + ( + AuthorizationCommit { + pending: self.pending, + credential, + }, + guard, + ) + } +} + +/// 完成 Provider OAuth 流程。 +#[derive(Clone, PartialEq, Eq)] +pub struct CompleteAuthorization { + pub context: MutationContext, + pub flow_id: String, + pub callback_url: String, +} + +impl fmt::Debug for CompleteAuthorization { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CompleteAuthorization") + .field("context", &self.context) + .field("flow_id", &"[REDACTED]") + .field("callback_url", &"[REDACTED]") + .finish() + } +} + +/// Credential 生命周期写操作。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialMutation { + pub context: MutationContext, + pub expected_config_revision: Revision, + pub account_id: ProviderAccountId, +} + +/// Credential 写入提交结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CredentialMutationResult { + pub config_revision: Revision, + pub account_id: ProviderAccountId, + pub credential_revision: Option, +} + +/// Provider-owned token 轮换命令。 +pub struct RotateCredential { + pub mutation: CredentialMutation, + pub expected_credential_revision: Revision, + pub provider_material: ProviderDocument, +} + +/// Provider 校验手工轮换材料时所需的非事务输入。 +pub struct PrepareCredentialRotation { + pub account: AccountRecord, + pub expected_credential_revision: Revision, + pub provider_material: ProviderDocument, +} + +impl fmt::Debug for PrepareCredentialRotation { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PrepareCredentialRotation") + .field("account", &self.account) + .field( + "expected_credential_revision", + &self.expected_credential_revision, + ) + .field("provider_material", &self.provider_material) + .finish() + } +} + +/// Provider 已验证、可由 Store 以 credential revision CAS 原子提交的轮换 facts。 +#[derive(Debug, Clone, PartialEq)] +pub struct PreparedCredentialRotationFacts { + pub account_id: ProviderAccountId, + pub provider_instance_id: ProviderInstanceId, + pub provider_kind: ProviderKind, + pub expected_credential_revision: Revision, + pub name: String, + pub email: Option, + pub plan_type: Option, + pub provider_material: ProviderDocument, + pub has_refresh_token: bool, + pub access_token_expires_at: DateTime, + pub next_refresh_at: Option>, +} + +/// Provider 返回的轮换准备结果；guard 必须覆盖后续 Store CAS 与审计事务。 +pub struct PreparedCredentialRotation { + facts: PreparedCredentialRotationFacts, + guard: Box, +} + +impl PreparedCredentialRotation { + #[must_use] + pub fn new( + facts: PreparedCredentialRotationFacts, + guard: Box, + ) -> Self { + Self { facts, guard } + } + + #[must_use] + pub const fn facts(&self) -> &PreparedCredentialRotationFacts { + &self.facts + } + + #[must_use] + pub fn into_parts( + self, + ) -> ( + PreparedCredentialRotationFacts, + Box, + ) { + (self.facts, self.guard) + } +} + +impl fmt::Debug for PreparedCredentialRotation { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PreparedCredentialRotation") + .field("facts", &self.facts) + .field("guard", &"[COMPLETION-GUARD]") + .finish() + } +} + +/// Admin 交给 Store 的轮换或 refresh 事务命令。 +#[derive(Debug, Clone, PartialEq)] +pub struct CredentialRotationCommit { + pub expected_config_revision: Revision, + pub prepared: PreparedCredentialRotationFacts, +} + +impl fmt::Debug for RotateCredential { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RotateCredential") + .field("mutation", &self.mutation) + .field( + "expected_credential_revision", + &self.expected_credential_revision, + ) + .field("provider_material", &self.provider_material) + .finish() + } +} + +/// 一个 Provider quota 窗口的公共投影。 +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderQuotaWindow { + pub key: String, + pub group: String, + pub source: Option, + pub window_seconds: Option, + pub used_percent: Option, + pub reset_at: Option>, + pub provider_data: Option, +} + +/// Provider 已解析的 quota 结果及其不透明差异字段。 +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderQuota { + pub observed_at: Option>, + pub refresh_token_expires_at: Option>, + pub windows: Vec, + pub provider_data: Option, +} + +/// Provider 实时模型目录的一项。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderModel { + pub id: UpstreamModelId, + pub name: String, +} + +/// Provider 实时模型目录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderModels { + pub models: Vec, + pub observed_at: Option>, +} + +/// Provider 执行 refresh 时所需的当前公共账号事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PrepareCredentialRefresh { + pub account: AccountRecord, +} + +/// Provider 敏感导出结果。 +pub struct ProviderExport { + pub provider_kind: ProviderKind, + pub account_ids: Vec, + pub document: ProviderDocument, +} + +/// Store 为 Provider 导出序列化准备的最小输入；material 对公共层保持不透明。 +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderExportCredentialInput { + pub account: AccountRecord, + pub provider_material: ProviderDocument, +} + +impl fmt::Debug for ProviderExport { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderExport") + .field("provider_kind", &self.provider_kind) + .field("account_ids", &self.account_ids) + .field("document", &self.document) + .finish() + } +} + +/// 统一账号目录的一行完整结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct AccountDirectoryItem { + pub account: AccountRecord, + pub provider_instance_name: String, + pub status: AccountStatus, + pub usage: Option, + pub rolling_usage: Option, + pub quota: ProviderQuota, +} + +/// 统一账号目录页。 +#[derive(Debug, Clone, PartialEq)] +pub struct AccountDirectoryPage { + pub config_revision: Revision, + pub items: Vec, + pub total: u64, + pub summary: AccountSummary, +} + +/// 凭据刷新提交后的完整账号结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct AccountRefreshResult { + pub config_revision: Revision, + pub account: AccountDirectoryItem, +} + +/// 多 Provider 导出文档集合。 +pub struct AccountExportBundle { + pub exported_at: DateTime, + pub documents: Vec, +} + +impl fmt::Debug for AccountExportBundle { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AccountExportBundle") + .field("exported_at", &self.exported_at) + .field("document_count", &self.documents.len()) + .finish() + } +} diff --git a/backend/crates/gateway-admin/src/model/settings.rs b/backend/crates/gateway-admin/src/model/settings.rs new file mode 100644 index 000000000..7df2a34cf --- /dev/null +++ b/backend/crates/gateway-admin/src/model/settings.rs @@ -0,0 +1,103 @@ +//! Runtime settings 与明文管理员 API Key 的语义模型。 + +use std::{collections::BTreeMap, fmt}; + +use chrono::{DateTime, Utc}; + +use gateway_core::routing::{ProviderKind, PublicModelId, UpstreamModelId}; + +use super::Revision; + +/// Provider → 客户端模型 → 上游模型的精确映射。 +pub type ProviderModelMappings = BTreeMap>; + +/// 账号调度策略。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RotationStrategy { + Smart, + QuotaResetPriority, + RoundRobin, + Sticky, +} + +/// 完整运行设置事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RuntimeSettings { + pub config_revision: Revision, + pub provider_model_mappings: ProviderModelMappings, + pub refresh_margin_seconds: u64, + pub refresh_concurrency: u32, + pub max_concurrent_per_account: u32, + pub request_interval_ms: u64, + pub rotation_strategy: RotationStrategy, + pub usage_retention_days: u32, + pub ops_event_retention_days: u32, + pub audit_retention_days: u32, + pub updated_at: DateTime, +} + +/// 原子替换运行设置的命令。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReplaceRuntimeSettings { + pub expected_config_revision: Revision, + pub provider_model_mappings: ProviderModelMappings, + pub refresh_margin_seconds: u64, + pub refresh_concurrency: u32, + pub max_concurrent_per_account: u32, + pub request_interval_ms: u64, + pub rotation_strategy: RotationStrategy, + pub usage_retention_days: u32, + pub ops_event_retention_days: u32, + pub audit_retention_days: u32, +} + +/// 明文管理员 API Key；按产品约束明文落库，但禁止 Debug 泄漏。 +#[derive(Clone, PartialEq, Eq)] +pub struct AdminApiKey(String); + +impl AdminApiKey { + #[must_use] + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + #[must_use] + pub fn expose_for_auth(&self) -> &str { + &self.0 + } + + /// 仅供显式 regenerate 响应读取一次。 + #[must_use] + pub fn expose_for_response(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for AdminApiKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("AdminApiKey([REDACTED])") + } +} + +/// 管理员 API Key 更新结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdminApiKeyMutation { + pub config_revision: Revision, + pub exists: bool, +} + +/// 新管理员 API Key 的一次性返回结果。 +pub struct RegeneratedAdminApiKey { + pub mutation: AdminApiKeyMutation, + pub key: AdminApiKey, +} + +impl fmt::Debug for RegeneratedAdminApiKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RegeneratedAdminApiKey") + .field("mutation", &self.mutation) + .field("key", &"[REDACTED]") + .finish() + } +} diff --git a/backend/crates/gateway-admin/src/model/system.rs b/backend/crates/gateway-admin/src/model/system.rs new file mode 100644 index 000000000..1db90356a --- /dev/null +++ b/backend/crates/gateway-admin/src/model/system.rs @@ -0,0 +1,129 @@ +//! 版本、自更新、回滚与进程重启的 UTC 语义模型。 + +use chrono::{DateTime, Utc}; + +/// 当前运行版本及更新检查摘要。 +/// +/// `deployment_mode_label` 之类的本地化字符串由 API 根据原始枚举值生成。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SystemVersion { + pub version: String, + pub git_sha: String, + pub build_time: String, + pub deployment_mode: String, + pub update_channel: String, + pub latest_version: String, + pub has_update: bool, + pub update_cached: bool, + pub update_warning: Option, +} + +/// 发布源提供的完整更新详情。 +/// +/// 部署模式与构建类型的展示标签归 API；其余字段均为 Host 已确认的原始事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SystemUpdateDetail { + pub current_version: String, + pub latest_version: String, + pub has_update: bool, + pub deployment_mode: String, + pub build_type: String, + pub release_url: Option, + pub notes: Option, + pub cached: bool, + pub update_supported: bool, + pub unsupported_reason: Option, + pub warning: Option, +} + +/// 系统操作类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SystemOperationKind { + Update, + Rollback, + Restart, +} + +/// 系统操作状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SystemOperationStatus { + Idle, + Running, + Succeeded, + Failed, +} + +/// 持久化系统操作的完整状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SystemOperationState { + pub operation_id: Option, + pub kind: Option, + pub status: SystemOperationStatus, + pub target_version: Option, + pub message: Option, + pub error: Option, + pub started_at: Option>, + pub finished_at: Option>, +} + +/// 最近一次系统操作及版本交换状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SystemUpdateStatus { + pub previous_version: Option, + pub current_version: Option, + pub operation: SystemOperationState, +} + +/// 自更新日志级别。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SystemUpdateEventLevel { + Info, + Warning, + Success, + Error, +} + +/// 一条已脱敏的系统更新事件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SystemUpdateEvent { + pub id: String, + pub operation_id: Option, + pub level: SystemUpdateEventLevel, + pub step: Option, + pub message: String, + pub terminal: bool, + pub progress_percent: Option, + pub occurred_at: DateTime, +} + +/// 已受理的异步系统操作；不同 wire shape 由枚举保持为合法状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SystemOperationAccepted { + Update { + operation_id: String, + deployment_mode: String, + message: String, + need_restart: bool, + target_version: String, + }, + Rollback { + operation_id: String, + message: String, + need_restart: bool, + }, + Restart { + operation_id: String, + message: String, + }, +} + +impl SystemOperationAccepted { + #[must_use] + pub const fn kind(&self) -> SystemOperationKind { + match self { + Self::Update { .. } => SystemOperationKind::Update, + Self::Rollback { .. } => SystemOperationKind::Rollback, + Self::Restart { .. } => SystemOperationKind::Restart, + } + } +} diff --git a/backend/crates/gateway-admin/src/ports/mod.rs b/backend/crates/gateway-admin/src/ports/mod.rs new file mode 100644 index 000000000..8de9dd195 --- /dev/null +++ b/backend/crates/gateway-admin/src/ports/mod.rs @@ -0,0 +1,5 @@ +//! 管理用例依赖的外部能力端口。 + +pub mod provider; +pub mod store; +pub mod system; diff --git a/backend/crates/gateway-admin/src/ports/provider.rs b/backend/crates/gateway-admin/src/ports/provider.rs new file mode 100644 index 000000000..7977494d1 --- /dev/null +++ b/backend/crates/gateway-admin/src/ports/provider.rs @@ -0,0 +1,178 @@ +//! Provider 管理能力与动态注册表。 + +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use gateway_core::{engine::credential::ProviderAccountId, routing::ProviderKind}; + +use crate::model::observability::{ + CalculatedBillingBreakdown, DashboardWireProfile, ProviderBillingInput, +}; +use crate::model::provider_credentials::{ + AuthorizationStarted, CompleteAuthorization, PendingAuthorizationMutation, + PrepareCredentialImport, PrepareCredentialRefresh, PrepareCredentialRotation, + PreparedAuthorizationCommit, PreparedCredentialImport, PreparedCredentialRotation, + ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota, +}; + +/// Provider 管理失败的稳定分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderAdminErrorKind { + Invalid, + Unsupported, + NotFound, + Conflict, + Unavailable, + Internal, +} + +/// 不泄漏 OAuth material 与 Provider 响应体的管理错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("provider admin operation failed: {kind:?}")] +pub struct ProviderAdminError { + kind: ProviderAdminErrorKind, +} + +impl ProviderAdminError { + #[must_use] + pub const fn new(kind: ProviderAdminErrorKind) -> Self { + Self { kind } + } + + #[must_use] + pub const fn kind(&self) -> ProviderAdminErrorKind { + self.kind + } +} + +/// 一个具体 Provider 对管理控制面提供的解析、验证、上游交互与运行时资源回收能力。 +/// +/// 数据变更由 Provider 返回 prepared facts；config revision、审计与 PostgreSQL 事务 +/// 全部由 [`crate::ports::store::AccountStore`] 提交。运行时资源通知只在事务成功后发生。 +#[async_trait] +pub trait ProviderAdmin: Send + Sync { + fn provider_kind(&self) -> &ProviderKind; + + /// 账号已经由控制面提交为不可调度状态，释放 Provider 持有的账号级运行时资源。 + /// + /// 无账号级运行时资源的 Provider 不需要执行额外操作。该通知发生在 Store 事务 + /// 成功之后，不参与事务成败，也不得恢复或改写已经提交的账号状态。 + async fn account_unavailable(&self, account_id: &ProviderAccountId); + + /// 返回该 Provider 实际持有的 Dashboard wire 画像；没有画像能力时返回 `None`。 + fn dashboard_wire_profile(&self) -> Option; + + /// 使用 Provider-owned 价格规则恢复持久请求的逐项费用。 + fn calculated_billing( + &self, + input: &ProviderBillingInput, + ) -> Result, ProviderAdminError>; + + async fn prepare_import( + &self, + command: PrepareCredentialImport, + ) -> Result; + + async fn start_authorization( + &self, + pending: PendingAuthorizationMutation, + ) -> Result; + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result; + + async fn prepare_rotation( + &self, + command: PrepareCredentialRotation, + ) -> Result; + + async fn prepare_refresh( + &self, + command: PrepareCredentialRefresh, + ) -> Result; + + async fn quota( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result; + + async fn models( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result; + + async fn export_credentials( + &self, + credentials: Vec, + ) -> Result; +} + +/// 按 ProviderKind 动态发现管理能力；不含具体 Provider 分支。 +#[derive(Clone)] +pub struct ProviderAdminRegistry { + providers: Arc>>, +} + +impl ProviderAdminRegistry { + /// 创建无重复 ProviderKind 的注册表。 + /// + /// # Errors + /// + /// 重复注册同一 ProviderKind 时返回 Conflict。 + pub fn new( + providers: impl IntoIterator>, + ) -> Result { + let mut registered = BTreeMap::new(); + for provider in providers { + let kind = provider.provider_kind().clone(); + if registered.insert(kind, provider).is_some() { + return Err(ProviderAdminError::new(ProviderAdminErrorKind::Conflict)); + } + } + Ok(Self { + providers: Arc::new(registered), + }) + } + + pub fn require( + &self, + provider_kind: &ProviderKind, + ) -> Result, ProviderAdminError> { + self.providers + .get(provider_kind) + .cloned() + .ok_or_else(|| ProviderAdminError::new(ProviderAdminErrorKind::Unsupported)) + } + + /// 返回唯一注册的 Dashboard wire 画像。 + /// + /// # Errors + /// + /// 没有 Provider 提供画像时返回 Unsupported；多个 Provider 同时声明画像时返回 Conflict。 + pub fn dashboard_wire_profile(&self) -> Result { + let mut profiles = self + .providers + .values() + .filter_map(|provider| provider.dashboard_wire_profile()); + let profile = profiles + .next() + .ok_or_else(|| ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))?; + if profiles.next().is_some() { + return Err(ProviderAdminError::new(ProviderAdminErrorKind::Conflict)); + } + Ok(profile) + } + + /// 动态分派 Provider-owned 费用规则，不含任何具体 Provider 分支。 + pub fn calculated_billing( + &self, + provider_kind: &ProviderKind, + input: &ProviderBillingInput, + ) -> Result, ProviderAdminError> { + self.require(provider_kind)?.calculated_billing(input) + } +} diff --git a/backend/crates/gateway-admin/src/ports/store.rs b/backend/crates/gateway-admin/src/ports/store.rs new file mode 100644 index 000000000..e5dba7b53 --- /dev/null +++ b/backend/crates/gateway-admin/src/ports/store.rs @@ -0,0 +1,392 @@ +//! 管理控制面所需的持久化能力。 +//! +//! 端口按业务资源拆分，方法使用领域模型，不暴露连接池、事务或 Redis client。 + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::model::{ + MutationContext, Revision, + accounts::{ + AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount, + SetAccountEnabled, + }, + auth::{AdminAuditEvent, AdminSession}, + catalog::{ + CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog, + ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled, + UpdateProviderInstance, + }, + client_keys::{ + ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey, + NewClientKey, SetClientKeyEnabled, UpdateClientKey, + }, + observability::{ + DashboardObservation, DiagnosticDimension, DiagnosticObservation, OpsErrorPage, + OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter, UsageOverview, + UsagePage, UsageQuery, + }, + provider_credentials::{ + AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult, + CredentialListQuery, CredentialMutationResult, CredentialPage, CredentialRotationCommit, + ProviderExportCredentialInput, + }, + settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings}, +}; + +/// 管理端可判定的持久化失败类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminStoreErrorKind { + Invalid, + NotFound, + StaleRevision, + Conflict, + Unavailable, +} + +/// 隐藏数据库实现细节的持久化错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{resource} store operation failed: {message}")] +pub struct AdminStoreError { + kind: AdminStoreErrorKind, + resource: &'static str, + message: String, +} + +impl AdminStoreError { + #[must_use] + pub fn new( + kind: AdminStoreErrorKind, + resource: &'static str, + message: impl Into, + ) -> Self { + Self { + kind, + resource, + message: message.into(), + } + } + + #[must_use] + pub const fn kind(&self) -> AdminStoreErrorKind { + self.kind + } + + #[must_use] + pub const fn resource(&self) -> &'static str { + self.resource + } +} + +pub type AdminStoreResult = Result; + +/// 账号目录与公共账号写操作。 +#[async_trait] +pub trait AccountStore: Send + Sync { + async fn list_accounts(&self, query: AccountListQuery) -> AdminStoreResult; + + async fn load_account(&self, account_id: &str) -> AdminStoreResult>; + + async fn load_account_usage( + &self, + range: TimeRange, + account_ids: &[String], + ) -> AdminStoreResult>; + + async fn list_credentials( + &self, + provider_kind: &gateway_core::routing::ProviderKind, + query: CredentialListQuery, + ) -> AdminStoreResult; + + async fn credential_details( + &self, + provider_kind: &gateway_core::routing::ProviderKind, + account_id: &gateway_core::engine::credential::ProviderAccountId, + ) -> AdminStoreResult>; + + async fn load_credentials_for_export( + &self, + provider_kind: &gateway_core::routing::ProviderKind, + account_ids: &[gateway_core::engine::credential::ProviderAccountId], + ) -> AdminStoreResult>; + + async fn commit_credential_import( + &self, + command: CredentialImportCommit, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn commit_authorization( + &self, + command: AuthorizationCommit, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn commit_credential_rotation( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn commit_credential_refresh( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn set_account_enabled( + &self, + command: SetAccountEnabled, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn delete_account( + &self, + command: DeleteAccount, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn record_credential_export( + &self, + account_ids: &[gateway_core::engine::credential::ProviderAccountId], + context: &MutationContext, + ) -> AdminStoreResult<()>; +} + +/// 管理员密码、会话、登录限流和安全审计。 +#[async_trait] +pub trait AuthStore: Send + Sync { + async fn load_password_hash(&self, admin_user_id: &str) -> AdminStoreResult>; + + async fn create_password_hash_if_absent( + &self, + admin_user_id: &str, + password_hash: &str, + ) -> AdminStoreResult; + + async fn load_admin_api_key(&self) -> AdminStoreResult>; + + async fn load_session(&self, session_id: &str) -> AdminStoreResult>; + + async fn store_session(&self, session_id: &str, session: &AdminSession) + -> AdminStoreResult<()>; + + async fn delete_session(&self, session_id: &str) -> AdminStoreResult>; + + async fn login_source_is_throttled( + &self, + source: &str, + failure_limit: u32, + window_seconds: u64, + ) -> AdminStoreResult; + + async fn record_login_failure( + &self, + source: &str, + failure_limit: u32, + window_seconds: u64, + ) -> AdminStoreResult; + + async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()>; + + async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()>; +} + +/// Provider instance 目录与 revision CAS 写入。 +#[async_trait] +pub trait CatalogStore: Send + Sync { + async fn list_provider_instances( + &self, + include_disabled: bool, + ) -> AdminStoreResult; + + async fn load_provider_instance( + &self, + id: &gateway_core::routing::ProviderInstanceId, + ) -> AdminStoreResult>; + + async fn create_provider_instance( + &self, + command: CreateProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn update_provider_instance( + &self, + command: UpdateProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn set_provider_instance_enabled( + &self, + command: SetProviderInstanceEnabled, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn delete_provider_instance( + &self, + command: DeleteProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult; +} + +/// Client API Key 管理及其 revision CAS 写入。 +#[async_trait] +pub trait ClientKeyStore: Send + Sync { + async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult; + + async fn reveal_client_key( + &self, + id: &gateway_core::policy::ClientApiKeyId, + ) -> AdminStoreResult>; + + async fn create_client_key( + &self, + command: NewClientKey, + context: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)>; + + async fn update_client_key( + &self, + command: UpdateClientKey, + context: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)>; + + async fn set_client_key_enabled( + &self, + command: SetClientKeyEnabled, + context: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)>; + + async fn delete_client_key( + &self, + command: DeleteClientKey, + context: &MutationContext, + ) -> AdminStoreResult; +} + +/// 用量、趋势、诊断与运维错误的只读能力。 +#[async_trait] +pub trait ObservabilityStore: Send + Sync { + async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult; + + async fn dashboard_trend(&self, range: TimeRange) -> AdminStoreResult>; + + async fn usage_trend( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> AdminStoreResult>; + + async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult; + + async fn usage_record_detail(&self, request_id: &str) -> AdminStoreResult; + + async fn usage_summary( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> AdminStoreResult; + + async fn usage_diagnostics( + &self, + range: TimeRange, + filter: UsageFilter, + dimension: DiagnosticDimension, + ) -> AdminStoreResult>; + + async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult; +} + +/// Runtime settings 与管理员 API Key 的 revision CAS 写入。 +#[async_trait] +pub trait SettingsStore: Send + Sync { + async fn load_runtime_settings(&self) -> AdminStoreResult; + + async fn admin_api_key_exists(&self) -> AdminStoreResult; + + async fn replace_runtime_settings( + &self, + command: ReplaceRuntimeSettings, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn replace_admin_api_key( + &self, + expected_config_revision: Revision, + key: AdminApiKey, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn delete_admin_api_key( + &self, + expected_config_revision: Revision, + context: &MutationContext, + ) -> AdminStoreResult; +} + +/// 管理用例所需能力的封闭集合。 +/// +/// 字段保持私有，每个 getter 只交出一种明确能力。该类型不提供通用拆包入口。 +#[derive(Clone)] +pub struct AdminStorePorts { + accounts: Arc, + auth: Arc, + catalog: Arc, + client_keys: Arc, + observability: Arc, + settings: Arc, +} + +impl AdminStorePorts { + #[must_use] + pub fn new( + accounts: Arc, + auth: Arc, + catalog: Arc, + client_keys: Arc, + observability: Arc, + settings: Arc, + ) -> Self { + Self { + accounts, + auth, + catalog, + client_keys, + observability, + settings, + } + } + + #[must_use] + pub fn accounts(&self) -> Arc { + self.accounts.clone() + } + + #[must_use] + pub fn auth(&self) -> Arc { + self.auth.clone() + } + + #[must_use] + pub fn catalog(&self) -> Arc { + self.catalog.clone() + } + + #[must_use] + pub fn client_keys(&self) -> Arc { + self.client_keys.clone() + } + + #[must_use] + pub fn observability(&self) -> Arc { + self.observability.clone() + } + + #[must_use] + pub fn settings(&self) -> Arc { + self.settings.clone() + } +} diff --git a/backend/crates/gateway-admin/src/ports/system.rs b/backend/crates/gateway-admin/src/ports/system.rs new file mode 100644 index 000000000..7196a89ce --- /dev/null +++ b/backend/crates/gateway-admin/src/ports/system.rs @@ -0,0 +1,76 @@ +//! Host 为系统管理用例提供的进程与操作系统能力。 + +use std::pin::Pin; + +use async_trait::async_trait; +use futures::Stream; + +use crate::model::system::{ + SystemOperationAccepted, SystemUpdateDetail, SystemUpdateEvent, SystemUpdateStatus, + SystemVersion, +}; + +/// Host 系统操作失败类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SystemOperationErrorKind { + Invalid, + Conflict, + Upstream, + Internal, +} + +/// 不泄漏路径、命令行或发布凭据的系统操作错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("system operation failed: {message}")] +pub struct SystemOperationError { + kind: SystemOperationErrorKind, + message: String, +} + +impl SystemOperationError { + #[must_use] + pub fn new(kind: SystemOperationErrorKind, message: impl Into) -> Self { + Self { + kind, + message: message.into(), + } + } + + #[must_use] + pub const fn kind(&self) -> SystemOperationErrorKind { + self.kind + } + + /// 返回 Host 已完成脱敏的客户端安全消息。 + #[must_use] + pub fn message(&self) -> &str { + &self.message + } +} + +/// 每个订阅者独占的系统事件流。 +pub type SystemUpdateEventStream = Pin + Send + 'static>>; + +/// 版本、自更新、回滚和重启能力；实现唯一归 gateway-host。 +#[async_trait] +pub trait SystemOperations: Send + Sync { + async fn version(&self) -> Result; + + async fn update_detail( + &self, + refresh: bool, + ) -> Result; + + fn update_events(&self) -> SystemUpdateEventStream; + + async fn perform_update( + &self, + target_version: Option, + ) -> Result; + + async fn update_status(&self) -> Result; + + async fn rollback(&self) -> Result; + + async fn restart(&self) -> Result; +} diff --git a/backend/crates/gateway-admin/src/use_case/accounts.rs b/backend/crates/gateway-admin/src/use_case/accounts.rs new file mode 100644 index 000000000..79d0867eb --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/accounts.rs @@ -0,0 +1,471 @@ +//! 统一账号目录与跨 Provider 动态分派。 + +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use chrono::{Duration, Utc}; +use futures::StreamExt as _; +use gateway_core::{ + engine::{ + credential::ProviderAccountId, + probe::{AccountProbe, AccountProbeRequest}, + }, + error::GatewayErrorKind, + routing::{ProviderKind, UpstreamModelId, snapshot::SnapshotControl}, +}; + +use crate::{ + model::{ + AdminError, MutationContext, Revision, + accounts::{ + AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery, + AccountRecord, AccountStatus, + }, + observability::TimeRange, + provider_credentials::{ + AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult, + PrepareCredentialRefresh, ProviderModels, + }, + }, + ports::{ + provider::ProviderAdminRegistry, + store::{AccountStore, CatalogStore, SettingsStore}, + }, +}; + +use super::{ + commit_credential_refresh, map_provider_error, map_store_error, publish_committed, + validate_prepared_rotation, +}; + +const CONNECTION_TEST_INPUT: &str = "Reply with exactly OK."; + +/// 统一账号页消费的服务。 +#[async_trait] +pub trait AccountsService: Send + Sync { + async fn list(&self, query: AccountListQuery) -> Result; + + async fn export( + &self, + context: &MutationContext, + account_ids: Vec, + ) -> Result; + + async fn refresh( + &self, + context: &MutationContext, + expected_config_revision: Revision, + account_id: ProviderAccountId, + ) -> Result; + + async fn quota( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result; + + async fn models( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result; + + async fn test_connection( + &self, + account_id: ProviderAccountId, + upstream_model: UpstreamModelId, + ) -> Result; +} + +pub(crate) struct DefaultAccountsService { + accounts: Arc, + catalog: Arc, + settings: Arc, + providers: ProviderAdminRegistry, + snapshot: Arc, + probe: Arc, +} + +impl DefaultAccountsService { + #[must_use] + pub(crate) const fn new( + accounts: Arc, + catalog: Arc, + settings: Arc, + providers: ProviderAdminRegistry, + snapshot: Arc, + probe: Arc, + ) -> Self { + Self { + accounts, + catalog, + settings, + providers, + snapshot, + probe, + } + } + + async fn load_account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.accounts + .load_account(account_id.as_str()) + .await + .map_err(|error| map_store_error(error, "provider account"))? + .ok_or_else(|| AdminError::not_found("Provider account was not found")) + } + + async fn provider_for_account( + &self, + account_id: &ProviderAccountId, + ) -> Result< + ( + AccountRecord, + Arc, + ), + AdminError, + > { + let account = self.load_account(account_id).await?; + let provider = self + .providers + .require(&account.provider_kind) + .map_err(|error| map_provider_error(error, "provider account"))?; + Ok((account, provider)) + } + + async fn load_directory_item( + &self, + account_id: &ProviderAccountId, + refresh_quota: bool, + ) -> Result { + let (account, provider) = self.provider_for_account(account_id).await?; + let settings = self + .settings + .load_runtime_settings() + .await + .map_err(|error| map_store_error(error, "runtime settings"))?; + let instance = self + .catalog + .load_provider_instance(&account.provider_instance_id) + .await + .map_err(|error| map_store_error(error, "Provider instance"))? + .ok_or_else(|| AdminError::not_found("Provider instance was not found"))?; + if instance.item.provider_kind != account.provider_kind { + return Err(AdminError::invalid( + "Provider account and instance kinds do not match", + )); + } + let now = Utc::now(); + let retained_range = retained_usage_range(now, settings.usage_retention_days); + let rolling_range = TimeRange { + start: now - Duration::hours(24), + end: now, + }; + let ids = vec![account.id.clone()]; + let (usage, rolling_usage, quota) = futures::try_join!( + async { + self.accounts + .load_account_usage(retained_range, &ids) + .await + .map_err(|error| map_store_error(error, "account usage")) + }, + async { + self.accounts + .load_account_usage(rolling_range, &ids) + .await + .map_err(|error| map_store_error(error, "rolling account usage")) + }, + async { + provider + .quota(account_id, refresh_quota) + .await + .map_err(|error| map_provider_error(error, "provider quota")) + }, + )?; + Ok(AccountDirectoryItem { + provider_instance_name: instance.item.name, + status: account_status(&account, now), + usage: usage.into_iter().next(), + rolling_usage: rolling_usage.into_iter().next(), + account, + quota, + }) + } +} + +#[async_trait] +impl AccountsService for DefaultAccountsService { + async fn list(&self, query: AccountListQuery) -> Result { + let (page, settings, instances) = futures::try_join!( + self.accounts.list_accounts(query), + self.settings.load_runtime_settings(), + self.catalog.list_provider_instances(true), + ) + .map_err(|error| map_store_error(error, "account directory"))?; + let now = Utc::now(); + let retained_range = retained_usage_range(now, settings.usage_retention_days); + let rolling_range = TimeRange { + start: now - Duration::hours(24), + end: now, + }; + let ids = page + .items + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let (usage, rolling_usage) = futures::try_join!( + self.accounts.load_account_usage(retained_range, &ids), + self.accounts.load_account_usage(rolling_range, &ids), + ) + .map_err(|error| map_store_error(error, "account usage"))?; + let usage = usage + .into_iter() + .map(|usage| (usage.account_id.clone(), usage)) + .collect::>(); + let rolling_usage = rolling_usage + .into_iter() + .map(|usage| (usage.account_id.clone(), usage)) + .collect::>(); + let instances = instances + .items + .into_iter() + .map(|instance| (instance.id.clone(), instance)) + .collect::>(); + let quotas = futures::future::join_all(page.items.iter().map(|account| async { + let account_id = ProviderAccountId::new(account.id.clone()) + .map_err(|_| AdminError::invalid("Invalid provider account ID"))?; + let provider = self + .providers + .require(&account.provider_kind) + .map_err(|error| map_provider_error(error, "provider quota"))?; + provider + .quota(&account_id, false) + .await + .map_err(|error| map_provider_error(error, "provider quota")) + })) + .await + .into_iter() + .collect::, AdminError>>()?; + let items = page + .items + .into_iter() + .zip(quotas) + .map(|(account, quota)| { + let instance = instances + .get(&account.provider_instance_id) + .ok_or_else(|| AdminError::not_found("Provider instance was not found"))?; + if instance.provider_kind != account.provider_kind { + return Err(AdminError::invalid( + "Provider account and instance kinds do not match", + )); + } + Ok(AccountDirectoryItem { + provider_instance_name: instance.name.clone(), + status: account_status(&account, now), + usage: usage.get(&account.id).cloned(), + rolling_usage: rolling_usage.get(&account.id).cloned(), + account, + quota, + }) + }) + .collect::, AdminError>>()?; + Ok(AccountDirectoryPage { + config_revision: page.config_revision, + items, + total: page.total, + summary: page.summary, + }) + } + + async fn export( + &self, + context: &MutationContext, + account_ids: Vec, + ) -> Result { + if account_ids.is_empty() || account_ids.len() > 200 { + return Err(AdminError::invalid( + "Account export requires between 1 and 200 accounts", + )); + } + let exported_ids = account_ids.clone(); + let mut grouped = BTreeMap::>::new(); + for account_id in account_ids { + let account = self.load_account(&account_id).await?; + grouped + .entry(account.provider_kind) + .or_default() + .push(account_id); + } + if grouped.values().any(|ids| { + let unique = ids.iter().collect::>(); + unique.len() != ids.len() + }) { + return Err(AdminError::invalid("Account export contains duplicate IDs")); + } + let mut documents = Vec::with_capacity(grouped.len()); + for (provider_kind, ids) in grouped { + let provider = self + .providers + .require(&provider_kind) + .map_err(|error| map_provider_error(error, "provider account export"))?; + let credentials = self + .accounts + .load_credentials_for_export(&provider_kind, &ids) + .await + .map_err(|error| map_store_error(error, "provider account export"))?; + documents.push( + provider + .export_credentials(credentials) + .await + .map_err(|error| map_provider_error(error, "provider account export"))?, + ); + } + self.accounts + .record_credential_export(&exported_ids, context) + .await + .map_err(|error| map_store_error(error, "provider account export audit"))?; + Ok(AccountExportBundle { + exported_at: Utc::now(), + documents, + }) + } + + async fn refresh( + &self, + context: &MutationContext, + expected_config_revision: Revision, + account_id: ProviderAccountId, + ) -> Result { + let (account, provider) = self.provider_for_account(&account_id).await?; + let prepared = provider + .prepare_refresh(PrepareCredentialRefresh { + account: account.clone(), + }) + .await + .map_err(|error| map_provider_error(error, "provider credential refresh"))?; + validate_prepared_rotation(&account, &prepared, "provider credential refresh")?; + let result = commit_credential_refresh( + self.accounts.as_ref(), + expected_config_revision, + prepared, + context, + "provider credential refresh", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + let account = self.load_directory_item(&result.account_id, false).await?; + Ok(AccountRefreshResult { + config_revision: result.config_revision, + account, + }) + } + + async fn quota( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + self.load_directory_item(account_id, refresh).await + } + + async fn models( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + let (_, provider) = self.provider_for_account(account_id).await?; + provider + .models(account_id, refresh) + .await + .map_err(|error| map_provider_error(error, "provider model catalog")) + } + + async fn test_connection( + &self, + account_id: ProviderAccountId, + upstream_model: UpstreamModelId, + ) -> Result { + let account = self.load_account(&account_id).await?; + let initial_status = account_status(&account, Utc::now()); + let model = upstream_model.as_str().to_owned(); + let initial = vec![ + AccountConnectionTestEvent::Started { + model: model.clone(), + }, + AccountConnectionTestEvent::Request { + model, + input_text: CONNECTION_TEST_INPUT.to_owned(), + stream: true, + store: false, + }, + ]; + let probe = Arc::clone(&self.probe); + let accounts = Arc::clone(&self.accounts); + let stored_account_id = account.id.clone(); + let terminal = futures::stream::once(async move { + let result = probe + .probe(AccountProbeRequest { + account_id, + provider_instance_id: account.provider_instance_id, + upstream_model, + }) + .await; + let status = accounts + .load_account(&stored_account_id) + .await + .ok() + .flatten() + .map_or(initial_status, |account| { + account_status(&account, Utc::now()) + }); + match result { + Ok(result) => result + .text + .into_iter() + .map(|text| AccountConnectionTestEvent::Content { text }) + .chain(std::iter::once(AccountConnectionTestEvent::Completed { + account_status: status, + })) + .collect(), + Err(error) => vec![AccountConnectionTestEvent::Failed { + message: match error.kind() { + GatewayErrorKind::InvalidRequest + | GatewayErrorKind::Unsupported + | GatewayErrorKind::ModelNotFound => error.safe_message().to_owned(), + _ => "Provider connection test failed".to_owned(), + }, + account_status: status, + }], + } + }) + .flat_map(futures::stream::iter); + Ok(Box::pin(futures::stream::iter(initial).chain(terminal))) + } +} + +fn account_status(account: &AccountRecord, now: chrono::DateTime) -> AccountStatus { + if !account.enabled { + AccountStatus::Disabled + } else { + match account.availability { + crate::model::accounts::AccountAvailability::Banned => AccountStatus::Banned, + crate::model::accounts::AccountAvailability::QuotaExhausted => { + AccountStatus::QuotaExhausted + } + crate::model::accounts::AccountAvailability::Expired => AccountStatus::Expired, + _ if account.access_token_expires_at <= now => AccountStatus::Expired, + crate::model::accounts::AccountAvailability::Ready => AccountStatus::Active, + _ => AccountStatus::Attention, + } + } +} + +fn retained_usage_range(now: chrono::DateTime, retention_days: u32) -> TimeRange { + // 账号留存投影是内部查询，不受外部观测接口 366 天上限约束。 + TimeRange { + start: now - Duration::days(i64::from(retention_days)), + end: now, + } +} diff --git a/backend/crates/gateway-admin/src/use_case/auth.rs b/backend/crates/gateway-admin/src/use_case/auth.rs new file mode 100644 index 000000000..0b64cdcea --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/auth.rs @@ -0,0 +1,253 @@ +//! 管理员认证状态机。 + +use std::sync::Arc; + +use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier, password_hash::SaltString}; +use async_trait::async_trait; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{Duration, Utc}; +use rand_core::{OsRng, RngCore as _}; +use subtle::ConstantTimeEq as _; +use uuid::Uuid; + +use crate::{ + model::{ + AdminError, + auth::{ + AdminAuditEvent, AdminSession, AuditActorKind, LoginCommand, LoginError, LoginResult, + }, + }, + ports::store::AuthStore, +}; + +use super::map_store_error; + +const LOGIN_FAILURE_LIMIT: u32 = 5; +const LOGIN_FAILURE_WINDOW_SECONDS: u64 = 15 * 60; + +/// API 鉴权与管理员登录消费的统一服务。 +#[async_trait] +pub trait AuthService: Send + Sync { + async fn ensure_default_admin(&self, password: &str) -> Result; + async fn resolve_admin_user_id( + &self, + session_id: Option<&str>, + ) -> Result, AdminError>; + async fn verify_admin_api_key(&self, key: &str) -> Result; + async fn login(&self, command: LoginCommand) -> Result; + async fn validate_session(&self, session_id: Option<&str>) -> Result; + async fn logout(&self, session_id: &str) -> Result<(), AdminError>; +} + +/// 单管理员认证策略的最终实现。 +pub(crate) struct DefaultAuthService { + default_admin_user_id: String, + session_ttl: Duration, + store: Arc, +} + +impl DefaultAuthService { + #[must_use] + pub(crate) fn new( + default_admin_user_id: impl Into, + session_ttl_minutes: u64, + store: Arc, + ) -> Self { + let minutes = i64::try_from(session_ttl_minutes) + .unwrap_or(i64::MAX) + .max(1); + Self { + default_admin_user_id: default_admin_user_id.into(), + session_ttl: Duration::minutes(minutes), + store, + } + } + + async fn record_failed_login(&self, source: &str) -> Result { + let throttled = self + .store + .record_login_failure(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS) + .await + .map_err(|_| LoginError::Unavailable)?; + Ok(if throttled { + LoginError::Throttled + } else { + LoginError::InvalidCredentials + }) + } + + fn auth_audit(&self, action: &str, occurred_at: chrono::DateTime) -> AdminAuditEvent { + AdminAuditEvent { + id: format!("audit_{}", Uuid::now_v7().simple()), + actor_kind: AuditActorKind::AdminSession, + actor_admin_user_id: Some(self.default_admin_user_id.clone()), + actor_ref: format!("admin:{}", self.default_admin_user_id), + request_id: None, + action: action.to_owned(), + entity_kind: "admin_session".to_owned(), + entity_ref: self.default_admin_user_id.clone(), + config_revision: None, + changed_fields: Vec::new(), + occurred_at, + } + } +} + +#[async_trait] +impl AuthService for DefaultAuthService { + async fn ensure_default_admin(&self, password: &str) -> Result { + let hash = hash_admin_password(password)?; + self.store + .create_password_hash_if_absent(&self.default_admin_user_id, &hash) + .await + .map_err(|error| map_store_error(error, "administrator")) + } + + async fn resolve_admin_user_id( + &self, + session_id: Option<&str>, + ) -> Result, AdminError> { + let Some(session_id) = session_id else { + return Ok(None); + }; + self.store + .load_session(session_id) + .await + .map(|session| { + session + .filter(|session| session.expires_at > Utc::now()) + .map(|session| session.admin_user_id) + }) + .map_err(|error| map_store_error(error, "administrator session")) + } + + async fn verify_admin_api_key(&self, key: &str) -> Result { + if !valid_admin_api_key_shape(key) { + return Ok(false); + } + let stored = self + .store + .load_admin_api_key() + .await + .map_err(|error| map_store_error(error, "administrator API key"))?; + Ok(stored.as_ref().is_some_and(|stored| { + let stored = stored.expose_for_auth(); + key.len() == stored.len() && bool::from(key.as_bytes().ct_eq(stored.as_bytes())) + })) + } + + async fn login(&self, command: LoginCommand) -> Result { + let source = normalized_login_source(&command.source); + if self + .store + .login_source_is_throttled(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS) + .await + .map_err(|_| LoginError::Unavailable)? + { + return Err(LoginError::Throttled); + } + if command + .username + .as_deref() + .unwrap_or(&self.default_admin_user_id) + != self.default_admin_user_id + { + return Err(self.record_failed_login(source).await?); + } + let hash = self + .store + .load_password_hash(&self.default_admin_user_id) + .await + .map_err(|_| LoginError::Unavailable)? + .ok_or(LoginError::InvalidCredentials)?; + if !verify_admin_password(&command.password, &hash).map_err(|_| LoginError::Unavailable)? { + return Err(self.record_failed_login(source).await?); + } + + self.store + .clear_login_failures(source) + .await + .map_err(|_| LoginError::Unavailable)?; + let session_id = random_session_token(); + let expires_at = Utc::now() + self.session_ttl; + self.store + .store_session( + &session_id, + &AdminSession { + admin_user_id: self.default_admin_user_id.clone(), + expires_at, + }, + ) + .await + .map_err(|_| LoginError::Unavailable)?; + if self + .store + .append_audit_event(self.auth_audit("admin.login", Utc::now())) + .await + .is_err() + { + let _ = self.store.delete_session(&session_id).await; + return Err(LoginError::Unavailable); + } + Ok(LoginResult { + session_id, + expires_at, + }) + } + + async fn validate_session(&self, session_id: Option<&str>) -> Result { + Ok(self.resolve_admin_user_id(session_id).await?.is_some()) + } + + async fn logout(&self, session_id: &str) -> Result<(), AdminError> { + let session = self + .store + .delete_session(session_id) + .await + .map_err(|error| map_store_error(error, "administrator session"))?; + if let Some(session) = session { + let mut event = self.auth_audit("admin.logout", Utc::now()); + event.actor_admin_user_id = Some(session.admin_user_id.clone()); + event.actor_ref = format!("admin:{}", session.admin_user_id); + event.entity_ref = session.admin_user_id; + self.store + .append_audit_event(event) + .await + .map_err(|error| map_store_error(error, "administrator audit"))?; + } + Ok(()) + } +} + +fn hash_admin_password(password: &str) -> Result { + let salt = SaltString::generate(&mut OsRng); + Argon2::default() + .hash_password(password.as_bytes(), &salt) + .map(|hash| hash.to_string()) + .map_err(|_| AdminError::internal("Failed to hash administrator password")) +} + +fn verify_admin_password(password: &str, encoded: &str) -> Result { + let hash = PasswordHash::new(encoded) + .map_err(|_| AdminError::internal("Stored administrator password hash is invalid"))?; + Ok(Argon2::default() + .verify_password(password.as_bytes(), &hash) + .is_ok()) +} + +fn random_session_token() -> String { + let mut bytes = [0_u8; 32]; + OsRng.fill_bytes(&mut bytes); + format!("session_{}", URL_SAFE_NO_PAD.encode(bytes)) +} + +fn valid_admin_api_key_shape(value: &str) -> bool { + value.len() == 70 + && value.starts_with("admin-") + && value[6..].bytes().all(|byte| byte.is_ascii_hexdigit()) +} + +fn normalized_login_source(source: &str) -> &str { + let source = source.trim(); + if source.is_empty() { "unknown" } else { source } +} diff --git a/backend/crates/gateway-admin/src/use_case/catalog.rs b/backend/crates/gateway-admin/src/use_case/catalog.rs new file mode 100644 index 000000000..a2f3a291f --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/catalog.rs @@ -0,0 +1,159 @@ +//! Provider instance 目录用例。 + +use std::sync::Arc; + +use async_trait::async_trait; +use gateway_core::routing::{ProviderInstanceId, snapshot::SnapshotControl}; + +use crate::{ + model::{ + AdminError, MutationContext, Revision, + catalog::{ + CatalogListQuery, CreateProviderInstance, DeleteProviderInstance, + ProviderInstanceDetail, ProviderInstanceMutation, ProviderInstancePage, + SetProviderInstanceEnabled, UpdateProviderInstance, + }, + }, + ports::store::CatalogStore, +}; + +use super::{map_store_error, publish_committed}; + +/// API 消费的 Provider instance 管理服务。 +#[async_trait] +pub trait CatalogService: Send + Sync { + async fn list(&self, query: CatalogListQuery) -> Result; + async fn get(&self, id: &ProviderInstanceId) -> Result; + async fn create( + &self, + context: &MutationContext, + command: CreateProviderInstance, + ) -> Result; + async fn update( + &self, + context: &MutationContext, + command: UpdateProviderInstance, + ) -> Result; + async fn set_enabled( + &self, + context: &MutationContext, + command: SetProviderInstanceEnabled, + ) -> Result; + async fn delete( + &self, + context: &MutationContext, + command: DeleteProviderInstance, + ) -> Result; +} + +pub(crate) struct DefaultCatalogService { + store: Arc, + snapshot: Arc, +} + +impl DefaultCatalogService { + #[must_use] + pub(crate) fn new(store: Arc, snapshot: Arc) -> Self { + Self { store, snapshot } + } +} + +#[async_trait] +impl CatalogService for DefaultCatalogService { + async fn list(&self, query: CatalogListQuery) -> Result { + let catalog = self + .store + .list_provider_instances(true) + .await + .map_err(|error| map_store_error(error, "provider instance catalog"))?; + let start = match query.cursor.as_ref() { + Some(cursor) => catalog + .items + .iter() + .position(|item| &item.id == cursor) + .map(|index| index + 1) + .ok_or_else(|| AdminError::invalid("Invalid provider instance cursor"))?, + None => 0, + }; + let limit = usize::from(query.page_size.get()); + let has_more = catalog.items.len().saturating_sub(start) > limit; + let items = catalog + .items + .into_iter() + .skip(start) + .take(limit) + .collect::>(); + let next_cursor = has_more + .then(|| items.last().map(|item| item.id.clone())) + .flatten(); + Ok(ProviderInstancePage { + config_revision: catalog.config_revision, + items, + next_cursor, + }) + } + + async fn get(&self, id: &ProviderInstanceId) -> Result { + self.store + .load_provider_instance(id) + .await + .map_err(|error| map_store_error(error, "provider instance"))? + .ok_or_else(|| AdminError::not_found("Provider instance was not found")) + } + + async fn create( + &self, + context: &MutationContext, + command: CreateProviderInstance, + ) -> Result { + let result = self + .store + .create_provider_instance(command, context) + .await + .map_err(|error| map_store_error(error, "provider instance"))?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn update( + &self, + context: &MutationContext, + command: UpdateProviderInstance, + ) -> Result { + let result = self + .store + .update_provider_instance(command, context) + .await + .map_err(|error| map_store_error(error, "provider instance"))?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn set_enabled( + &self, + context: &MutationContext, + command: SetProviderInstanceEnabled, + ) -> Result { + let result = self + .store + .set_provider_instance_enabled(command, context) + .await + .map_err(|error| map_store_error(error, "provider instance"))?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn delete( + &self, + context: &MutationContext, + command: DeleteProviderInstance, + ) -> Result { + let revision = self + .store + .delete_provider_instance(command, context) + .await + .map_err(|error| map_store_error(error, "provider instance"))?; + publish_committed(self.snapshot.as_ref(), revision).await?; + Ok(revision) + } +} diff --git a/backend/crates/gateway-admin/src/use_case/client_keys.rs b/backend/crates/gateway-admin/src/use_case/client_keys.rs new file mode 100644 index 000000000..3243843f2 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/client_keys.rs @@ -0,0 +1,204 @@ +//! Client API Key 管理用例。 + +use std::sync::Arc; + +use async_trait::async_trait; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use gateway_core::policy::ClientApiKeyId; +use gateway_core::routing::snapshot::SnapshotControl; +use rand_core::{OsRng, RngCore as _}; +use uuid::Uuid; + +use crate::{ + model::{ + AdminError, MutationContext, + client_keys::{ + ClientKeyCursorValue, ClientKeyListQuery, ClientKeyMutation, ClientKeyPage, + ClientKeySecret, ClientKeySortField, CreateClientKey, CreatedClientKey, + DeleteClientKey, NewClientKey, SetClientKeyEnabled, UpdateClientKey, + }, + }, + ports::store::ClientKeyStore, +}; + +use super::{map_store_error, publish_committed}; + +/// API 消费的 Client Key 管理服务。 +#[async_trait] +pub trait ClientKeyService: Send + Sync { + async fn list(&self, query: ClientKeyListQuery) -> Result; + async fn reveal(&self, id: &ClientApiKeyId) -> Result; + async fn create( + &self, + context: &MutationContext, + command: CreateClientKey, + ) -> Result; + async fn update( + &self, + context: &MutationContext, + command: UpdateClientKey, + ) -> Result; + async fn set_enabled( + &self, + context: &MutationContext, + command: SetClientKeyEnabled, + ) -> Result; + async fn delete( + &self, + context: &MutationContext, + command: DeleteClientKey, + ) -> Result; +} + +pub(crate) struct DefaultClientKeyService { + store: Arc, + snapshot: Arc, +} + +impl DefaultClientKeyService { + #[must_use] + pub(crate) fn new(store: Arc, snapshot: Arc) -> Self { + Self { store, snapshot } + } +} + +#[async_trait] +impl ClientKeyService for DefaultClientKeyService { + async fn list(&self, query: ClientKeyListQuery) -> Result { + validate_cursor(&query)?; + self.store + .list_client_keys(query) + .await + .map_err(|error| map_store_error(error, "client API key")) + } + + async fn reveal(&self, id: &ClientApiKeyId) -> Result { + self.store + .reveal_client_key(id) + .await + .map_err(|error| map_store_error(error, "client API key"))? + .ok_or_else(|| AdminError::not_found("Client API key was not found")) + } + + async fn create( + &self, + context: &MutationContext, + command: CreateClientKey, + ) -> Result { + let id = ClientApiKeyId::new(format!("key_{}", Uuid::now_v7().simple())) + .map_err(|_| AdminError::internal("Failed to create Client API key ID"))?; + let mut bytes = [0_u8; 32]; + OsRng.fill_bytes(&mut bytes); + let plaintext = format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes)); + let (config_revision, record) = self + .store + .create_client_key( + NewClientKey { + expected_config_revision: command.expected_config_revision, + id, + name: command.name, + label: command.label, + provider_kind: command.provider_kind, + limits: command.limits, + plaintext: plaintext.clone(), + }, + context, + ) + .await + .map_err(|error| map_store_error(error, "client API key"))?; + publish_committed(self.snapshot.as_ref(), config_revision).await?; + Ok(CreatedClientKey { + config_revision, + secret: ClientKeySecret::new(record, plaintext), + }) + } + + async fn update( + &self, + context: &MutationContext, + command: UpdateClientKey, + ) -> Result { + let id = command.id.clone(); + let (config_revision, record) = self + .store + .update_client_key(command, context) + .await + .map_err(|error| map_store_error(error, "client API key"))?; + publish_committed(self.snapshot.as_ref(), config_revision).await?; + Ok(ClientKeyMutation { + config_revision, + record: Some(record), + id, + }) + } + + async fn set_enabled( + &self, + context: &MutationContext, + command: SetClientKeyEnabled, + ) -> Result { + let id = command.id.clone(); + let (config_revision, record) = + self.store + .set_client_key_enabled(command, context) + .await + .map_err(|error| map_store_error(error, "client API key"))?; + publish_committed(self.snapshot.as_ref(), config_revision).await?; + Ok(ClientKeyMutation { + config_revision, + record: Some(record), + id, + }) + } + + async fn delete( + &self, + context: &MutationContext, + command: DeleteClientKey, + ) -> Result { + let id = command.id.clone(); + let config_revision = self + .store + .delete_client_key(command, context) + .await + .map_err(|error| map_store_error(error, "client API key"))?; + publish_committed(self.snapshot.as_ref(), config_revision).await?; + Ok(ClientKeyMutation { + config_revision, + record: None, + id, + }) + } +} + +fn validate_cursor(query: &ClientKeyListQuery) -> Result<(), AdminError> { + let Some(cursor) = &query.cursor else { + return Ok(()); + }; + if cursor.sort != query.sort { + return Err(AdminError::invalid( + "Client API key cursor sort does not match the query", + )); + } + let matches = matches!( + (cursor.sort.field, &cursor.value), + (ClientKeySortField::Name, ClientKeyCursorValue::Name(value)) if !value.trim().is_empty() + ) || matches!( + (cursor.sort.field, &cursor.value), + ( + ClientKeySortField::Enabled, + ClientKeyCursorValue::Enabled(_) + ) | ( + ClientKeySortField::CreatedAt, + ClientKeyCursorValue::CreatedAt(_) + ) | ( + ClientKeySortField::LastUsedAt, + ClientKeyCursorValue::LastUsedAt(_) + ) + ); + if matches { + Ok(()) + } else { + Err(AdminError::invalid("Invalid Client API key cursor")) + } +} diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs new file mode 100644 index 000000000..5190e0c8b --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -0,0 +1,355 @@ +//! 管理控制面的用例实现。 + +pub mod accounts; +pub mod auth; +pub mod catalog; +pub mod client_keys; +pub mod observability; +pub mod openai; +pub mod settings; +pub mod system; +pub mod xai; + +use crate::{ + model::{ + AdminError, AdminErrorKind, MutationContext, + accounts::{DeleteAccount, SetAccountEnabled}, + provider_credentials::{ + AuthorizationMutationTarget, CredentialDetails, CredentialMutation, + CredentialMutationResult, CredentialRotationCommit, PendingAuthorizationMutation, + PreparedAuthorizationCommit, PreparedAuthorizationCredential, PreparedCredentialImport, + PreparedCredentialRotation, StartAuthorization, + }, + }, + ports::{ + provider::ProviderAdmin, + store::{AccountStore, AdminStoreError, AdminStoreErrorKind}, + }, +}; +use gateway_core::{ + engine::credential::ProviderAccountId, + routing::{ConfigRevision, ProviderKind, snapshot::SnapshotControl}, +}; + +fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { + let kind = match error.kind() { + AdminStoreErrorKind::Invalid => AdminErrorKind::Invalid, + AdminStoreErrorKind::NotFound => AdminErrorKind::NotFound, + AdminStoreErrorKind::StaleRevision | AdminStoreErrorKind::Conflict => { + AdminErrorKind::Conflict + } + AdminStoreErrorKind::Unavailable => AdminErrorKind::Unavailable, + }; + AdminError::new(kind, format!("{resource} operation failed")) +} + +fn map_provider_error( + error: crate::ports::provider::ProviderAdminError, + resource: &'static str, +) -> AdminError { + use crate::ports::provider::ProviderAdminErrorKind; + + let kind = match error.kind() { + ProviderAdminErrorKind::Invalid => AdminErrorKind::Invalid, + ProviderAdminErrorKind::Unsupported => AdminErrorKind::Invalid, + ProviderAdminErrorKind::NotFound => AdminErrorKind::NotFound, + ProviderAdminErrorKind::Conflict => AdminErrorKind::Conflict, + ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable, + ProviderAdminErrorKind::Internal => AdminErrorKind::Internal, + }; + AdminError::new(kind, format!("{resource} operation failed")) +} + +async fn publish_committed( + snapshot: &dyn SnapshotControl, + revision: crate::model::Revision, +) -> Result<(), AdminError> { + let revision = ConfigRevision::new(revision.get()) + .map_err(|_| AdminError::internal("Committed configuration revision is invalid"))?; + snapshot.publish_committed(revision).await; + Ok(()) +} + +async fn required_credential( + accounts: &dyn AccountStore, + provider_kind: &ProviderKind, + account_id: &ProviderAccountId, + resource: &'static str, +) -> Result { + accounts + .credential_details(provider_kind, account_id) + .await + .map_err(|error| map_store_error(error, resource))? + .ok_or_else(|| AdminError::not_found(format!("{resource} was not found"))) +} + +async fn pending_authorization( + accounts: &dyn AccountStore, + provider_kind: &ProviderKind, + command: &StartAuthorization, + resource: &'static str, +) -> Result { + let target = match &command.reauthorization { + Some(target) => { + let details = + required_credential(accounts, provider_kind, &target.account_id, resource).await?; + if details.credential.provider_instance_id != command.provider_instance_id + || details.credential.credential_revision != target.credential_revision + { + return Err(AdminError::conflict(format!( + "{resource} reauthorization target is stale" + ))); + } + AuthorizationMutationTarget::Reauthorize { + provider_instance_id: command.provider_instance_id.clone(), + account_id: target.account_id.clone(), + expected_credential_revision: target.credential_revision, + } + } + None => AuthorizationMutationTarget::Create { + provider_instance_id: command.provider_instance_id.clone(), + name: command.name.clone(), + }, + }; + Ok(PendingAuthorizationMutation::new( + command.expected_config_revision, + provider_kind.clone(), + target, + crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( + &command.context, + ), + )) +} + +fn validate_prepared_import( + provider_kind: &ProviderKind, + provider_instance_id: &gateway_core::routing::ProviderInstanceId, + prepared: &PreparedCredentialImport, + resource: &'static str, +) -> Result<(), AdminError> { + if prepared.provider_kind != *provider_kind + || prepared.provider_instance_id != *provider_instance_id + || prepared.credentials.iter().any(|credential| { + credential.provider_kind != *provider_kind + || credential.provider_instance_id != *provider_instance_id + }) + { + return Err(AdminError::conflict(format!( + "{resource} prepared facts do not match the requested Provider scope" + ))); + } + Ok(()) +} + +fn validate_prepared_rotation( + account: &crate::model::accounts::AccountRecord, + prepared: &PreparedCredentialRotation, + resource: &'static str, +) -> Result<(), AdminError> { + let facts = prepared.facts(); + if facts.account_id.as_str() != account.id.as_str() + || facts.provider_kind != account.provider_kind + || facts.provider_instance_id != account.provider_instance_id + || facts.expected_credential_revision != account.credential_revision + { + return Err(AdminError::conflict(format!( + "{resource} prepared facts do not match the current credential" + ))); + } + Ok(()) +} + +fn validate_authorization_commit( + provider_kind: &ProviderKind, + context: &MutationContext, + prepared: &PreparedAuthorizationCommit, + resource: &'static str, +) -> Result<(), AdminError> { + if prepared.pending.provider_kind() != provider_kind + || !prepared.pending.owner_binding().matches_context(context) + { + return Err(AdminError::conflict(format!( + "{resource} pending authorization binding is invalid" + ))); + } + let matches_target = match (prepared.pending.target(), &prepared.credential) { + ( + AuthorizationMutationTarget::Create { + provider_instance_id, + .. + }, + PreparedAuthorizationCredential::Create(credential), + ) => { + credential.provider_kind == *provider_kind + && credential.provider_instance_id == *provider_instance_id + } + ( + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + }, + PreparedAuthorizationCredential::Reauthorize(credential), + ) => { + let facts = credential.facts(); + facts.provider_kind == *provider_kind + && facts.provider_instance_id == *provider_instance_id + && facts.account_id == *account_id + && facts.expected_credential_revision == *expected_credential_revision + } + _ => false, + }; + if !matches_target { + return Err(AdminError::conflict(format!( + "{resource} prepared credential does not match its pending target" + ))); + } + Ok(()) +} + +async fn commit_authorization( + accounts: &dyn AccountStore, + prepared: PreparedAuthorizationCommit, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (command, guard) = prepared.into_commit(); + match accounts.commit_authorization(command, context).await { + Ok(result) => { + if let Some(guard) = guard { + guard.finish(); + } + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +async fn commit_credential_rotation( + accounts: &dyn AccountStore, + expected_config_revision: crate::model::Revision, + prepared: PreparedCredentialRotation, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_rotation( + CredentialRotationCommit { + expected_config_revision, + prepared: facts, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +async fn commit_credential_refresh( + accounts: &dyn AccountStore, + expected_config_revision: crate::model::Revision, + prepared: PreparedCredentialRotation, + context: &MutationContext, + resource: &'static str, +) -> Result { + let (facts, guard) = prepared.into_parts(); + match accounts + .commit_credential_refresh( + CredentialRotationCommit { + expected_config_revision, + prepared: facts, + }, + context, + ) + .await + { + Ok(result) => { + guard.finish(); + Ok(result) + } + Err(error) => { + drop(guard); + Err(map_store_error(error, resource)) + } + } +} + +async fn set_credential_enabled( + accounts: &dyn AccountStore, + provider: &dyn ProviderAdmin, + command: CredentialMutation, + enabled: bool, + resource: &'static str, +) -> Result { + required_credential( + accounts, + provider.provider_kind(), + &command.account_id, + resource, + ) + .await?; + let account_id = command.account_id; + let revision = accounts + .set_account_enabled( + SetAccountEnabled { + expected_config_revision: command.expected_config_revision, + account_id: account_id.as_str().to_owned(), + enabled, + }, + &command.context, + ) + .await + .map_err(|error| map_store_error(error, resource))?; + if !enabled { + provider.account_unavailable(&account_id).await; + } + Ok(CredentialMutationResult { + config_revision: revision, + account_id, + credential_revision: None, + }) +} + +async fn delete_credential( + accounts: &dyn AccountStore, + provider: &dyn ProviderAdmin, + command: CredentialMutation, + resource: &'static str, +) -> Result { + required_credential( + accounts, + provider.provider_kind(), + &command.account_id, + resource, + ) + .await?; + let account_id = command.account_id; + let revision = accounts + .delete_account( + DeleteAccount { + expected_config_revision: command.expected_config_revision, + account_id: account_id.as_str().to_owned(), + }, + &command.context, + ) + .await + .map_err(|error| map_store_error(error, resource))?; + provider.account_unavailable(&account_id).await; + Ok(CredentialMutationResult { + config_revision: revision, + account_id, + credential_revision: None, + }) +} diff --git a/backend/crates/gateway-admin/src/use_case/observability.rs b/backend/crates/gateway-admin/src/use_case/observability.rs new file mode 100644 index 000000000..fa21641a8 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/observability.rs @@ -0,0 +1,727 @@ +//! 观测查询、趋势、费用覆盖与健康阈值规则。 + +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use chrono::{DateTime, Duration, Timelike as _, Utc}; +use gateway_core::routing::ProviderKind; + +use crate::{ + model::{ + AdminError, + observability::{ + CostCoverage, CurrencyCost, DashboardCapacity, DashboardPeriodMetrics, DashboardResult, + DecimalAmount, DiagnosticDimension, DiagnosticsItem, DiagnosticsResult, HealthStatus, + HealthTimeline, HealthTimelinePoint, OpsErrorPage, OpsErrorQuery, ProviderBillingInput, + RequestMetricPoint, RequestMetrics, TimeRange, Trend, TrendKind, TrendPoint, + TrendSummary, UsageBilling, UsageDetail, UsageFilter, UsageInsights, UsageInsightsCost, + UsageInsightsCostPoint, UsageInsightsHealth, UsageInsightsHealthPoint, + UsageInsightsPerformance, UsageInsightsPerformancePoint, UsageOverview, UsagePage, + UsageQuery, UsageSummary, + }, + }, + ports::{ + provider::ProviderAdminRegistry, + store::{ObservabilityStore, SettingsStore}, + }, +}; + +use super::{map_provider_error, map_store_error}; + +const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15; +const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4; +const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10; +const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3; +const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 99.0; +const CHINA_OFFSET_SECONDS: i64 = 8 * 60 * 60; + +/// API 消费的观测控制面服务。 +#[async_trait] +pub trait ObservabilityService: Send + Sync { + async fn dashboard_summary( + &self, + range: TimeRange, + kind: TrendKind, + ) -> Result; + async fn dashboard_trend(&self, range: TimeRange, kind: TrendKind) + -> Result; + async fn usage_records(&self, query: UsageQuery) -> Result; + async fn usage_record_detail(&self, request_id: &str) -> Result; + async fn usage_summary( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> Result; + async fn usage_insights( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> Result; + async fn diagnostics( + &self, + range: TimeRange, + filter: UsageFilter, + dimension: DiagnosticDimension, + ) -> Result; + async fn ops_errors(&self, query: OpsErrorQuery) -> Result; +} + +pub(crate) struct DefaultObservabilityService { + store: Arc, + settings: Arc, + providers: ProviderAdminRegistry, +} + +impl DefaultObservabilityService { + #[must_use] + pub(crate) fn new( + store: Arc, + settings: Arc, + providers: ProviderAdminRegistry, + ) -> Self { + Self { + store, + settings, + providers, + } + } +} + +#[async_trait] +impl ObservabilityService for DefaultObservabilityService { + async fn dashboard_summary( + &self, + range: TimeRange, + kind: TrendKind, + ) -> Result { + let (mut observation, settings) = futures::try_join!( + self.store.dashboard_summary(range), + self.settings.load_runtime_settings(), + ) + .map_err(|error| map_store_error(error, "dashboard"))?; + self.enrich_billing(&mut observation.recent_requests)?; + let today_start = china_day_start(observation.range.end); + let yesterday_start = today_start - Duration::days(1); + let today = + dashboard_period_metrics(&observation.trend, today_start, observation.range.end); + let yesterday = dashboard_period_metrics(&observation.trend, yesterday_start, today_start); + let total_billing_usd = observation + .attempts + .costs + .iter() + .find(|cost| cost.currency == "USD") + .map(|cost| cost.amount.clone()); + let total_cached_token_rate = rate_or_zero( + observation.requests.cached_tokens, + observation.requests.input_tokens, + ); + let average_first_token_latency_ms = average( + observation.requests.first_token_latency_sum_ms, + observation.requests.first_token_latency_count, + ); + let trend = trend(kind, observation.trend.clone())?; + let health_timeline = health_timeline_at(&observation.trend, Utc::now()); + let wire_profile = self + .providers + .dashboard_wire_profile() + .map_err(|error| map_provider_error(error, "dashboard wire profile"))?; + let max_concurrent_per_account = u64::from(settings.max_concurrent_per_account); + Ok(DashboardResult { + capacity: DashboardCapacity { + max_concurrent_per_account, + total_slots: observation + .provider_accounts + .active + .saturating_mul(max_concurrent_per_account), + used_slots: None, + available_slots: None, + }, + rotation_strategy: settings.rotation_strategy, + observation, + today, + yesterday, + total_billing_usd, + total_cached_token_rate, + average_first_token_latency_ms, + trend, + health_timeline, + wire_profile, + }) + } + + async fn dashboard_trend( + &self, + range: TimeRange, + kind: TrendKind, + ) -> Result { + let points = self + .store + .dashboard_trend(range) + .await + .map_err(|error| map_store_error(error, "dashboard trend"))?; + trend(kind, points) + } + + async fn usage_records(&self, query: UsageQuery) -> Result { + let mut page = self + .store + .list_usage_records(query) + .await + .map_err(|error| map_store_error(error, "usage records"))?; + self.enrich_billing(&mut page.items)?; + Ok(page) + } + + async fn usage_record_detail(&self, request_id: &str) -> Result { + if request_id.trim().is_empty() { + return Err(AdminError::invalid("Usage record ID is required")); + } + let mut detail = self + .store + .usage_record_detail(request_id) + .await + .map_err(|error| map_store_error(error, "usage record"))?; + self.enrich_billing(std::slice::from_mut(&mut detail.request))?; + Ok(detail) + } + + async fn usage_summary( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> Result { + let overview = self + .store + .usage_summary(range, filter) + .await + .map_err(|error| map_store_error(error, "usage summary"))?; + let average_latency_ms = average( + overview.requests.latency_sum_ms, + overview.requests.latency_count, + ); + Ok(UsageSummary { + overview, + average_latency_ms, + }) + } + + async fn usage_insights( + &self, + range: TimeRange, + filter: UsageFilter, + ) -> Result { + let (overview, trend) = futures::try_join!( + self.store.usage_summary(range, filter.clone()), + self.store.usage_trend(range, filter), + ) + .map_err(|error| map_store_error(error, "usage insights"))?; + build_usage_insights(overview, trend) + } + + async fn diagnostics( + &self, + range: TimeRange, + filter: UsageFilter, + dimension: DiagnosticDimension, + ) -> Result { + let items = self + .store + .usage_diagnostics(range, filter, dimension) + .await + .map_err(|error| map_store_error(error, "usage diagnostics"))?; + let total_requests = items.iter().fold(0_u64, |total, item| { + total.saturating_add(item.request_count) + }); + Ok(DiagnosticsResult { + dimension, + items: items + .into_iter() + .map(|item| DiagnosticsItem { + name: item.name, + request_count: item.request_count, + success_count: item.request_count.saturating_sub(item.failure_count), + error_count: item.failure_count, + error_rate: rate_or_zero(item.failure_count, item.request_count), + request_share: rate_or_zero(item.request_count, total_requests), + average_latency_ms: item.average_latency_ms, + estimated_cost: None, + attempt_count: item.attempt_count, + total_tokens: item.total_tokens, + }) + .collect(), + }) + } + + async fn ops_errors(&self, query: OpsErrorQuery) -> Result { + self.store + .list_ops_errors(query) + .await + .map_err(|error| map_store_error(error, "operations errors")) + } +} + +fn build_usage_insights( + overview: UsageOverview, + trend: Vec, +) -> Result { + let granularity = trend.first().map_or( + crate::model::observability::Granularity::FifteenMinutes, + |point| point.granularity, + ); + if trend.iter().any(|point| point.granularity != granularity) { + return Err(AdminError::internal( + "Usage insight trend contains mixed granularities", + )); + } + + let requests = &overview.requests; + let failed_requests = service_failure_count(requests); + let health_requests = requests.success_count.saturating_add(failed_requests); + let health = UsageInsightsHealth { + total_requests: health_requests, + success_requests: requests.success_count, + failed_requests, + cancelled_requests: requests.cancelled_count, + incomplete_requests: requests.incomplete_count, + caller_error_requests: requests.caller_error_count, + success_rate: rate_or_zero(requests.success_count, health_requests), + points: trend + .iter() + .map(|point| { + let failed_requests = service_failure_count(&point.metrics); + let health_requests = point.metrics.success_count.saturating_add(failed_requests); + UsageInsightsHealthPoint { + bucket_start: point.bucket_start, + success_requests: point.metrics.success_count, + failed_requests, + cancelled_requests: point.metrics.cancelled_count, + incomplete_requests: point.metrics.incomplete_count, + caller_error_requests: point.metrics.caller_error_count, + error_rate: rate_or_zero(failed_requests, health_requests), + } + }) + .collect(), + }; + let performance = UsageInsightsPerformance { + latency_percentiles: requests.latency_percentiles, + first_token_latency_percentiles: requests.first_token_latency_percentiles, + latency_coverage: rate_or_zero(requests.latency_count, requests.request_count), + first_token_coverage: rate_or_zero( + requests.first_token_latency_count, + requests.request_count, + ), + points: trend + .iter() + .map(|point| UsageInsightsPerformancePoint { + bucket_start: point.bucket_start, + latency_percentiles: point.metrics.latency_percentiles, + first_token_latency_percentiles: point.metrics.first_token_latency_percentiles, + }) + .collect(), + }; + let estimated_cost = overview + .attempts + .costs + .iter() + .find(|cost| cost.currency.eq_ignore_ascii_case("USD")) + .map(|cost| cost.amount.clone()); + let cost = UsageInsightsCost { + estimated_cost, + standard_cost: None, + cost_per_request: None, + tokens_per_request: rate_or_zero(requests.total_tokens, requests.request_count), + cached_token_rate: rate_or_zero(requests.cached_tokens, requests.input_tokens), + cache_hit_request_rate: ratio( + requests.cache_hit_request_count, + requests.cache_eligible_request_count, + ), + input_tokens: requests.input_tokens, + output_tokens: requests.output_tokens, + cached_tokens: requests.cached_tokens, + total_tokens: requests.total_tokens, + points: trend + .iter() + .map(|point| UsageInsightsCostPoint { + bucket_start: point.bucket_start, + input_tokens: point.metrics.input_tokens, + output_tokens: point.metrics.output_tokens, + cached_tokens: point.metrics.cached_tokens, + total_tokens: point.metrics.total_tokens, + estimated_cost: None, + standard_cost: None, + cached_token_rate: rate_or_zero( + point.metrics.cached_tokens, + point.metrics.input_tokens, + ), + cache_hit_request_rate: ratio( + point.metrics.cache_hit_request_count, + point.metrics.cache_eligible_request_count, + ), + }) + .collect(), + costs: overview.attempts.costs.clone(), + coverage: overview.attempts.cost_coverage.clone(), + }; + Ok(UsageInsights { + granularity, + health, + performance, + cost, + attempts: overview.attempts, + providers: overview.providers, + }) +} + +impl DefaultObservabilityService { + fn enrich_billing( + &self, + records: &mut [crate::model::observability::UsageRecord], + ) -> Result<(), AdminError> { + for record in records { + let Some(UsageBilling::Total { source, total }) = record.billing.as_ref() else { + continue; + }; + if source != "calculated" { + continue; + } + let (Some(provider), Some(upstream_model_id)) = ( + record.provider_kind.as_deref(), + record.upstream_model_id.as_ref(), + ) else { + continue; + }; + let provider_kind = ProviderKind::new(provider.to_owned()) + .map_err(|_| AdminError::internal("Stored Provider kind is invalid"))?; + let input = ProviderBillingInput { + upstream_model_id: upstream_model_id.clone(), + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cached_tokens: record.cached_tokens, + cache_write_tokens: record.cache_write_tokens, + total: total.clone(), + }; + if let Some(breakdown) = self + .providers + .calculated_billing(&provider_kind, &input) + .map_err(|error| map_provider_error(error, "usage billing"))? + { + record.billing = Some(UsageBilling::Calculated(Box::new(breakdown))); + } + } + Ok(()) + } +} + +/// 按指定时刻计算中国自然日的 96 个 15 分钟健康桶。 +#[must_use] +fn health_timeline_at(records: &[RequestMetricPoint], now: DateTime) -> HealthTimeline { + let current_slot = quarter_hour_start(now); + let start = china_day_start(now); + let mut buckets = (0..HEALTH_TIMELINE_SLOTS) + .map(|index| { + ( + start + Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index), + HealthWindow::default(), + ) + }) + .collect::>(); + for record in records { + if record.bucket_start < start || record.bucket_start > now { + continue; + } + let record_slot = quarter_hour_start(record.bucket_start); + if let Some((_, bucket)) = buckets + .iter_mut() + .find(|(bucket_start, _)| *bucket_start == record_slot) + { + bucket.add_metrics(&record.metrics); + } + } + + let totals = buckets + .iter() + .filter(|(bucket_start, _)| *bucket_start <= current_slot) + .fold(HealthWindow::default(), |mut totals, (_, bucket)| { + totals.add_window(*bucket); + totals + }); + HealthTimeline { + reliability_percent: health_reliability(totals), + status: health_status(totals, false), + success_requests: totals.success_requests, + failed_requests: totals.failed_requests, + cancelled_requests: totals.cancelled_requests, + incomplete_requests: totals.incomplete_requests, + caller_error_requests: totals.caller_error_requests, + points: buckets + .into_iter() + .map(|(bucket_start, bucket)| HealthTimelinePoint { + bucket_start, + status: health_status(bucket, bucket_start > current_slot), + reliability_percent: health_reliability(bucket), + success_requests: bucket.success_requests, + failed_requests: bucket.failed_requests, + cancelled_requests: bucket.cancelled_requests, + incomplete_requests: bucket.incomplete_requests, + caller_error_requests: bucket.caller_error_requests, + }) + .collect(), + } +} + +fn trend(kind: TrendKind, points: Vec) -> Result { + let mut metrics = RequestMetrics::default(); + let mut coverage = CostCoverage::default(); + let mut costs = BTreeMap::::new(); + for point in &points { + add_request_metrics(&mut metrics, &point.metrics); + coverage.provider_reported_count = coverage + .provider_reported_count + .saturating_add(point.cost_coverage.provider_reported_count); + coverage.calculated_count = coverage + .calculated_count + .saturating_add(point.cost_coverage.calculated_count); + coverage.partial_count = coverage + .partial_count + .saturating_add(point.cost_coverage.partial_count); + coverage.unavailable_count = coverage + .unavailable_count + .saturating_add(point.cost_coverage.unavailable_count); + coverage.not_billable_count = coverage + .not_billable_count + .saturating_add(point.cost_coverage.not_billable_count); + for cost in &point.costs { + if let Some(amount) = costs.get_mut(&cost.currency) { + *amount = amount.checked_add(&cost.amount).ok_or_else(|| { + AdminError::internal("Cost aggregation exceeded numeric bounds") + })?; + } else { + costs.insert(cost.currency.clone(), cost.amount.clone()); + } + } + } + let average_latency_ms = average(metrics.latency_sum_ms, metrics.latency_count); + let service_failure_count = service_failure_count(&metrics); + let success_rate = ratio( + metrics.success_count, + metrics.success_count.saturating_add(service_failure_count), + ); + let cache_hit_request_rate = ratio( + metrics.cache_hit_request_count, + metrics.cache_eligible_request_count, + ); + Ok(Trend { + kind, + points: points.into_iter().map(trend_point).collect(), + summary: TrendSummary { + request_count: metrics.request_count, + input_tokens: metrics.input_tokens, + output_tokens: metrics.output_tokens, + cached_tokens: metrics.cached_tokens, + total_tokens: metrics.total_tokens, + service_failure_count, + average_latency_ms, + max_latency_ms: metrics.max_latency_ms, + min_latency_ms: metrics.min_latency_ms, + success_rate, + cache_hit_request_rate, + costs: costs + .into_iter() + .map(|(currency, amount)| CurrencyCost { currency, amount }) + .collect(), + cost_coverage: coverage, + }, + }) +} + +fn trend_point(point: RequestMetricPoint) -> TrendPoint { + let service_failure_count = service_failure_count(&point.metrics); + TrendPoint { + bucket_start: point.bucket_start, + granularity: point.granularity, + service_failure_count, + average_latency_ms: average(point.metrics.latency_sum_ms, point.metrics.latency_count), + average_first_token_latency_ms: average( + point.metrics.first_token_latency_sum_ms, + point.metrics.first_token_latency_count, + ), + cached_token_rate: rate_or_zero(point.metrics.cached_tokens, point.metrics.input_tokens), + cache_hit_request_rate: ratio( + point.metrics.cache_hit_request_count, + point.metrics.cache_eligible_request_count, + ), + success_rate: ratio(point.metrics.success_count, point.metrics.request_count), + metrics: point.metrics, + cost_coverage: point.cost_coverage, + costs: point.costs, + } +} + +fn dashboard_period_metrics( + points: &[RequestMetricPoint], + start: DateTime, + end: DateTime, +) -> DashboardPeriodMetrics { + let mut result = DashboardPeriodMetrics::default(); + for metrics in points + .iter() + .filter(|point| point.bucket_start >= start && point.bucket_start < end) + .map(|point| &point.metrics) + { + result.request_count = result.request_count.saturating_add(metrics.request_count); + result.input_tokens = result.input_tokens.saturating_add(metrics.input_tokens); + result.output_tokens = result.output_tokens.saturating_add(metrics.output_tokens); + result.cached_tokens = result.cached_tokens.saturating_add(metrics.cached_tokens); + result.total_tokens = result.total_tokens.saturating_add(metrics.total_tokens); + } + result.cached_token_rate = rate_or_zero(result.cached_tokens, result.input_tokens); + result.observed_cached_token_rate = ratio(result.cached_tokens, result.input_tokens); + result +} + +fn add_request_metrics(total: &mut RequestMetrics, value: &RequestMetrics) { + total.request_count = total.request_count.saturating_add(value.request_count); + total.success_count = total.success_count.saturating_add(value.success_count); + total.failure_count = total.failure_count.saturating_add(value.failure_count); + total.cancelled_count = total.cancelled_count.saturating_add(value.cancelled_count); + total.incomplete_count = total + .incomplete_count + .saturating_add(value.incomplete_count); + total.caller_error_count = total + .caller_error_count + .saturating_add(value.caller_error_count); + total.input_tokens = total.input_tokens.saturating_add(value.input_tokens); + total.output_tokens = total.output_tokens.saturating_add(value.output_tokens); + total.cached_tokens = total.cached_tokens.saturating_add(value.cached_tokens); + total.cache_write_tokens = total + .cache_write_tokens + .saturating_add(value.cache_write_tokens); + total.reasoning_tokens = total + .reasoning_tokens + .saturating_add(value.reasoning_tokens); + total.total_tokens = total.total_tokens.saturating_add(value.total_tokens); + total.first_token_latency_sum_ms = total + .first_token_latency_sum_ms + .saturating_add(value.first_token_latency_sum_ms); + total.first_token_latency_count = total + .first_token_latency_count + .saturating_add(value.first_token_latency_count); + total.latency_sum_ms = total.latency_sum_ms.saturating_add(value.latency_sum_ms); + total.latency_count = total.latency_count.saturating_add(value.latency_count); + total.min_latency_ms = match (total.min_latency_ms, value.min_latency_ms) { + (Some(left), Some(right)) => Some(left.min(right)), + (left, right) => left.or(right), + }; + total.max_latency_ms = match (total.max_latency_ms, value.max_latency_ms) { + (Some(left), Some(right)) => Some(left.max(right)), + (left, right) => left.or(right), + }; + total.cache_eligible_request_count = total + .cache_eligible_request_count + .saturating_add(value.cache_eligible_request_count); + total.cache_hit_request_count = total + .cache_hit_request_count + .saturating_add(value.cache_hit_request_count); +} + +fn average(sum: u64, count: u64) -> Option { + (count > 0).then(|| sum / count) +} + +fn ratio(numerator: u64, denominator: u64) -> Option { + (denominator > 0).then(|| numerator as f64 / denominator as f64) +} + +fn rate_or_zero(numerator: u64, denominator: u64) -> f64 { + ratio(numerator, denominator).unwrap_or(0.0) +} + +#[derive(Debug, Clone, Copy, Default)] +struct HealthWindow { + success_requests: u64, + failed_requests: u64, + cancelled_requests: u64, + incomplete_requests: u64, + caller_error_requests: u64, +} + +impl HealthWindow { + fn add_metrics(&mut self, metrics: &RequestMetrics) { + self.success_requests = self.success_requests.saturating_add(metrics.success_count); + self.failed_requests = self + .failed_requests + .saturating_add(service_failure_count(metrics)); + self.cancelled_requests = self + .cancelled_requests + .saturating_add(metrics.cancelled_count); + self.incomplete_requests = self + .incomplete_requests + .saturating_add(metrics.incomplete_count); + self.caller_error_requests = self + .caller_error_requests + .saturating_add(metrics.caller_error_count); + } + + fn add_window(&mut self, other: Self) { + self.success_requests = self.success_requests.saturating_add(other.success_requests); + self.failed_requests = self.failed_requests.saturating_add(other.failed_requests); + self.cancelled_requests = self + .cancelled_requests + .saturating_add(other.cancelled_requests); + self.incomplete_requests = self + .incomplete_requests + .saturating_add(other.incomplete_requests); + self.caller_error_requests = self + .caller_error_requests + .saturating_add(other.caller_error_requests); + } +} + +fn health_status(bucket: HealthWindow, is_future: bool) -> HealthStatus { + let eligible_requests = bucket + .success_requests + .saturating_add(bucket.failed_requests); + if is_future { + HealthStatus::Future + } else if eligible_requests == 0 { + HealthStatus::NoData + } else if bucket.success_requests == 0 + && bucket.failed_requests >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD + { + HealthStatus::Unavailable + } else if eligible_requests < HEALTH_TIMELINE_MIN_SAMPLE_SIZE { + HealthStatus::LowSample + } else if health_reliability(bucket) + .is_some_and(|reliability| reliability < HEALTH_TIMELINE_STABLE_RELIABILITY) + { + HealthStatus::Unstable + } else { + HealthStatus::Stable + } +} + +fn health_reliability(bucket: HealthWindow) -> Option { + let eligible_requests = bucket + .success_requests + .saturating_add(bucket.failed_requests); + (eligible_requests > 0) + .then(|| bucket.success_requests as f64 / eligible_requests as f64 * 100.0) +} + +fn service_failure_count(metrics: &RequestMetrics) -> u64 { + metrics + .failure_count + .saturating_sub(metrics.caller_error_count) +} + +fn china_day_start(value: DateTime) -> DateTime { + let elapsed = (value.timestamp() + CHINA_OFFSET_SECONDS).rem_euclid(24 * 60 * 60); + value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond())) +} + +fn quarter_hour_start(value: DateTime) -> DateTime { + let elapsed = value + .timestamp() + .rem_euclid(HEALTH_TIMELINE_SLOT_MINUTES * 60); + value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond())) +} diff --git a/backend/crates/gateway-admin/src/use_case/openai.rs b/backend/crates/gateway-admin/src/use_case/openai.rs new file mode 100644 index 000000000..6449e0270 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/openai.rs @@ -0,0 +1,277 @@ +//! OpenAI 管理资源的中立 ProviderAdmin 委托。 + +use std::sync::Arc; + +use async_trait::async_trait; +use gateway_core::{engine::credential::ProviderAccountId, routing::snapshot::SnapshotControl}; + +use crate::{ + model::{ + AdminError, + provider_credentials::{ + AuthorizationStarted, CompleteAuthorization, CredentialDetails, CredentialImportCommit, + CredentialImportResult, CredentialListQuery, CredentialMutation, + CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport, + PrepareCredentialRotation, RotateCredential, StartAuthorization, + }, + }, + ports::{provider::ProviderAdmin, store::AccountStore}, +}; + +use super::{ + commit_authorization, commit_credential_rotation, delete_credential, map_provider_error, + map_store_error, pending_authorization, publish_committed, required_credential, + set_credential_enabled, validate_authorization_commit, validate_prepared_import, + validate_prepared_rotation, +}; + +/// OpenAI 固定管理路由消费的服务。 +#[async_trait] +pub trait OpenAiService: Send + Sync { + async fn list(&self, query: CredentialListQuery) -> Result; + async fn details( + &self, + account_id: &ProviderAccountId, + ) -> Result; + async fn import_document( + &self, + command: ImportCredentials, + ) -> Result; + async fn start_authorization( + &self, + command: StartAuthorization, + ) -> Result; + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result; + async fn rotate( + &self, + command: RotateCredential, + ) -> Result; + async fn enable( + &self, + command: CredentialMutation, + ) -> Result; + async fn disable( + &self, + command: CredentialMutation, + ) -> Result; + async fn delete( + &self, + command: CredentialMutation, + ) -> Result; +} + +pub(crate) struct DefaultOpenAiService { + provider: Arc, + accounts: Arc, + snapshot: Arc, +} + +impl DefaultOpenAiService { + #[must_use] + pub(crate) fn new( + provider: Arc, + accounts: Arc, + snapshot: Arc, + ) -> Self { + Self { + provider, + accounts, + snapshot, + } + } +} + +#[async_trait] +impl OpenAiService for DefaultOpenAiService { + async fn list(&self, query: CredentialListQuery) -> Result { + self.accounts + .list_credentials(self.provider.provider_kind(), query) + .await + .map_err(|error| map_store_error(error, "OpenAI credential")) + } + + async fn details( + &self, + account_id: &ProviderAccountId, + ) -> Result { + required_credential( + self.accounts.as_ref(), + self.provider.provider_kind(), + account_id, + "OpenAI credential", + ) + .await + } + + async fn import_document( + &self, + command: ImportCredentials, + ) -> Result { + let context = command.context; + let expected_config_revision = command.expected_config_revision; + let provider_instance_id = command.provider_instance_id; + let prepared = self + .provider + .prepare_import(PrepareCredentialImport { + provider_instance_id: provider_instance_id.clone(), + document: command.document, + }) + .await + .map_err(|error| map_provider_error(error, "OpenAI credential import"))?; + validate_prepared_import( + self.provider.provider_kind(), + &provider_instance_id, + &prepared, + "OpenAI credential import", + )?; + let result = self + .accounts + .commit_credential_import( + CredentialImportCommit { + expected_config_revision, + prepared, + }, + &context, + ) + .await + .map_err(|error| map_store_error(error, "OpenAI credential import"))?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn start_authorization( + &self, + command: StartAuthorization, + ) -> Result { + let pending = pending_authorization( + self.accounts.as_ref(), + self.provider.provider_kind(), + &command, + "OpenAI credential", + ) + .await?; + self.provider + .start_authorization(pending) + .await + .map_err(|error| map_provider_error(error, "OpenAI authorization")) + } + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result { + let context = command.context.clone(); + let prepared = self + .provider + .complete_authorization(command) + .await + .map_err(|error| map_provider_error(error, "OpenAI authorization"))?; + validate_authorization_commit( + self.provider.provider_kind(), + &context, + &prepared, + "OpenAI authorization", + )?; + let result = commit_authorization( + self.accounts.as_ref(), + prepared, + &context, + "OpenAI authorization", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn rotate( + &self, + command: RotateCredential, + ) -> Result { + let context = command.mutation.context; + let expected_config_revision = command.mutation.expected_config_revision; + let account_id = command.mutation.account_id; + let details = required_credential( + self.accounts.as_ref(), + self.provider.provider_kind(), + &account_id, + "OpenAI credential rotation", + ) + .await?; + if details.credential.credential_revision != command.expected_credential_revision { + return Err(AdminError::conflict( + "OpenAI credential rotation revision is stale", + )); + } + let account = details.credential; + let prepared = self + .provider + .prepare_rotation(PrepareCredentialRotation { + account: account.clone(), + expected_credential_revision: command.expected_credential_revision, + provider_material: command.provider_material, + }) + .await + .map_err(|error| map_provider_error(error, "OpenAI credential rotation"))?; + validate_prepared_rotation(&account, &prepared, "OpenAI credential rotation")?; + let result = commit_credential_rotation( + self.accounts.as_ref(), + expected_config_revision, + prepared, + &context, + "OpenAI credential rotation", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn enable( + &self, + command: CredentialMutation, + ) -> Result { + let result = set_credential_enabled( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + true, + "OpenAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn disable( + &self, + command: CredentialMutation, + ) -> Result { + let result = set_credential_enabled( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + false, + "OpenAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn delete( + &self, + command: CredentialMutation, + ) -> Result { + let result = delete_credential( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + "OpenAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } +} diff --git a/backend/crates/gateway-admin/src/use_case/settings.rs b/backend/crates/gateway-admin/src/use_case/settings.rs new file mode 100644 index 000000000..e41961f1d --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/settings.rs @@ -0,0 +1,130 @@ +//! Runtime settings 与管理员 API Key 用例。 + +use std::sync::Arc; + +use async_trait::async_trait; +use gateway_core::routing::snapshot::SnapshotControl; +use rand_core::{OsRng, RngCore as _}; + +use crate::{ + model::{ + AdminError, MutationContext, + settings::{ + AdminApiKey, AdminApiKeyMutation, RegeneratedAdminApiKey, ReplaceRuntimeSettings, + RuntimeSettings, + }, + }, + ports::store::SettingsStore, +}; + +use super::{map_store_error, publish_committed}; + +/// API 消费的 Runtime settings 管理服务。 +#[async_trait] +pub trait SettingsService: Send + Sync { + async fn load(&self) -> Result; + async fn replace( + &self, + context: &MutationContext, + command: ReplaceRuntimeSettings, + ) -> Result; + async fn admin_api_key_exists(&self) -> Result; + async fn regenerate_admin_api_key( + &self, + context: &MutationContext, + ) -> Result; + async fn delete_admin_api_key( + &self, + context: &MutationContext, + ) -> Result; +} + +pub(crate) struct DefaultSettingsService { + store: Arc, + snapshot: Arc, +} + +impl DefaultSettingsService { + #[must_use] + pub(crate) fn new(store: Arc, snapshot: Arc) -> Self { + Self { store, snapshot } + } +} + +#[async_trait] +impl SettingsService for DefaultSettingsService { + async fn load(&self) -> Result { + self.store + .load_runtime_settings() + .await + .map_err(|error| map_store_error(error, "runtime settings")) + } + + async fn replace( + &self, + context: &MutationContext, + command: ReplaceRuntimeSettings, + ) -> Result { + validate_settings(&command)?; + let settings = self + .store + .replace_runtime_settings(command, context) + .await + .map_err(|error| map_store_error(error, "runtime settings"))?; + publish_committed(self.snapshot.as_ref(), settings.config_revision).await?; + Ok(settings) + } + + async fn admin_api_key_exists(&self) -> Result { + self.store + .admin_api_key_exists() + .await + .map_err(|error| map_store_error(error, "administrator API key")) + } + + async fn regenerate_admin_api_key( + &self, + context: &MutationContext, + ) -> Result { + let settings = self.load().await?; + let mut bytes = [0_u8; 32]; + OsRng.fill_bytes(&mut bytes); + let key = AdminApiKey::new(format!("admin-{}", hex::encode(bytes))); + let mutation = self + .store + .replace_admin_api_key(settings.config_revision, key.clone(), context) + .await + .map_err(|error| map_store_error(error, "administrator API key"))?; + publish_committed(self.snapshot.as_ref(), mutation.config_revision).await?; + Ok(RegeneratedAdminApiKey { mutation, key }) + } + + async fn delete_admin_api_key( + &self, + context: &MutationContext, + ) -> Result { + let settings = self.load().await?; + let mutation = self + .store + .delete_admin_api_key(settings.config_revision, context) + .await + .map_err(|error| map_store_error(error, "administrator API key"))?; + publish_committed(self.snapshot.as_ref(), mutation.config_revision).await?; + Ok(mutation) + } +} + +fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError> { + let valid = command.refresh_margin_seconds > 0 + && command.refresh_concurrency > 0 + && command.max_concurrent_per_account > 0 + && command.usage_retention_days >= 31 + && command.ops_event_retention_days > 0 + && command.audit_retention_days > 0 + && i64::try_from(command.request_interval_ms).is_ok(); + if valid { + Ok(()) + } else { + Err(AdminError::invalid("Runtime settings violate constraints")) + } +} diff --git a/backend/crates/gateway-admin/src/use_case/system.rs b/backend/crates/gateway-admin/src/use_case/system.rs new file mode 100644 index 000000000..e7fdc8342 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/system.rs @@ -0,0 +1,98 @@ +//! 系统管理用例。 + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::{ + model::{ + AdminError, AdminErrorKind, + system::{SystemOperationAccepted, SystemUpdateDetail, SystemUpdateStatus, SystemVersion}, + }, + ports::system::{ + SystemOperationError, SystemOperationErrorKind, SystemOperations, SystemUpdateEventStream, + }, +}; + +/// API 消费的系统管理服务。 +#[async_trait] +pub trait SystemService: Send + Sync { + async fn version(&self) -> Result; + async fn update_detail(&self, refresh: bool) -> Result; + fn update_events(&self) -> SystemUpdateEventStream; + async fn perform_update( + &self, + target_version: Option, + ) -> Result; + async fn update_status(&self) -> Result; + async fn rollback(&self) -> Result; + async fn restart(&self) -> Result; +} + +/// 保持 Host 能力窄边界的默认系统用例。 +pub(crate) struct DefaultSystemService { + operations: Arc, +} + +impl DefaultSystemService { + #[must_use] + pub(crate) const fn new(operations: Arc) -> Self { + Self { operations } + } +} + +#[async_trait] +impl SystemService for DefaultSystemService { + async fn version(&self) -> Result { + self.operations.version().await.map_err(map_system_error) + } + + async fn update_detail(&self, refresh: bool) -> Result { + self.operations + .update_detail(refresh) + .await + .map_err(map_system_error) + } + + fn update_events(&self) -> SystemUpdateEventStream { + self.operations.update_events() + } + + async fn perform_update( + &self, + target_version: Option, + ) -> Result { + let target_version = target_version + .map(|version| version.trim().to_owned()) + .filter(|version| !version.is_empty()); + self.operations + .perform_update(target_version) + .await + .map_err(map_system_error) + } + + async fn update_status(&self) -> Result { + self.operations + .update_status() + .await + .map_err(map_system_error) + } + + async fn rollback(&self) -> Result { + self.operations.rollback().await.map_err(map_system_error) + } + + async fn restart(&self) -> Result { + self.operations.restart().await.map_err(map_system_error) + } +} + +fn map_system_error(error: SystemOperationError) -> AdminError { + let kind = match error.kind() { + SystemOperationErrorKind::Invalid => AdminErrorKind::Invalid, + SystemOperationErrorKind::Conflict => AdminErrorKind::Conflict, + SystemOperationErrorKind::Upstream => AdminErrorKind::BadGateway, + SystemOperationErrorKind::Internal => AdminErrorKind::Internal, + }; + AdminError::new(kind, error.message()) +} diff --git a/backend/crates/gateway-admin/src/use_case/xai.rs b/backend/crates/gateway-admin/src/use_case/xai.rs new file mode 100644 index 000000000..6a721fe87 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/xai.rs @@ -0,0 +1,213 @@ +//! xAI 管理资源的中立 ProviderAdmin 委托。 + +use std::sync::Arc; + +use async_trait::async_trait; +use gateway_core::routing::snapshot::SnapshotControl; + +use crate::{ + model::{ + AdminError, + provider_credentials::{ + AuthorizationStarted, CompleteAuthorization, CredentialImportCommit, + CredentialImportResult, CredentialListQuery, CredentialMutation, + CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport, + StartAuthorization, + }, + }, + ports::{provider::ProviderAdmin, store::AccountStore}, +}; + +use super::{ + commit_authorization, delete_credential, map_provider_error, map_store_error, + pending_authorization, publish_committed, set_credential_enabled, + validate_authorization_commit, validate_prepared_import, +}; + +/// xAI 固定管理路由消费的服务。 +#[async_trait] +pub trait XaiService: Send + Sync { + async fn list(&self, query: CredentialListQuery) -> Result; + async fn import_document( + &self, + command: ImportCredentials, + ) -> Result; + async fn start_authorization( + &self, + command: StartAuthorization, + ) -> Result; + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result; + async fn enable( + &self, + command: CredentialMutation, + ) -> Result; + async fn disable( + &self, + command: CredentialMutation, + ) -> Result; + async fn delete( + &self, + command: CredentialMutation, + ) -> Result; +} + +pub(crate) struct DefaultXaiService { + provider: Arc, + accounts: Arc, + snapshot: Arc, +} + +impl DefaultXaiService { + #[must_use] + pub(crate) fn new( + provider: Arc, + accounts: Arc, + snapshot: Arc, + ) -> Self { + Self { + provider, + accounts, + snapshot, + } + } +} + +#[async_trait] +impl XaiService for DefaultXaiService { + async fn list(&self, query: CredentialListQuery) -> Result { + self.accounts + .list_credentials(self.provider.provider_kind(), query) + .await + .map_err(|error| map_store_error(error, "xAI credential")) + } + + async fn import_document( + &self, + command: ImportCredentials, + ) -> Result { + let context = command.context; + let expected_config_revision = command.expected_config_revision; + let provider_instance_id = command.provider_instance_id; + let prepared = self + .provider + .prepare_import(PrepareCredentialImport { + provider_instance_id: provider_instance_id.clone(), + document: command.document, + }) + .await + .map_err(|error| map_provider_error(error, "xAI credential import"))?; + validate_prepared_import( + self.provider.provider_kind(), + &provider_instance_id, + &prepared, + "xAI credential import", + )?; + let result = self + .accounts + .commit_credential_import( + CredentialImportCommit { + expected_config_revision, + prepared, + }, + &context, + ) + .await + .map_err(|error| map_store_error(error, "xAI credential import"))?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn start_authorization( + &self, + command: StartAuthorization, + ) -> Result { + let pending = pending_authorization( + self.accounts.as_ref(), + self.provider.provider_kind(), + &command, + "xAI credential", + ) + .await?; + self.provider + .start_authorization(pending) + .await + .map_err(|error| map_provider_error(error, "xAI authorization")) + } + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result { + let context = command.context.clone(); + let prepared = self + .provider + .complete_authorization(command) + .await + .map_err(|error| map_provider_error(error, "xAI authorization"))?; + validate_authorization_commit( + self.provider.provider_kind(), + &context, + &prepared, + "xAI authorization", + )?; + let result = commit_authorization( + self.accounts.as_ref(), + prepared, + &context, + "xAI authorization", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn enable( + &self, + command: CredentialMutation, + ) -> Result { + let result = set_credential_enabled( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + true, + "xAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn disable( + &self, + command: CredentialMutation, + ) -> Result { + let result = set_credential_enabled( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + false, + "xAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } + + async fn delete( + &self, + command: CredentialMutation, + ) -> Result { + let result = delete_credential( + self.accounts.as_ref(), + self.provider.as_ref(), + command, + "xAI credential", + ) + .await?; + publish_committed(self.snapshot.as_ref(), result.config_revision).await?; + Ok(result) + } +} diff --git a/backend/crates/gateway-admin/tests/main.rs b/backend/crates/gateway-admin/tests/main.rs new file mode 100644 index 000000000..558a620a6 --- /dev/null +++ b/backend/crates/gateway-admin/tests/main.rs @@ -0,0 +1,117 @@ +use std::str::FromStr as _; + +use chrono::{TimeDelta, Utc}; +use gateway_core::engine::credential::OpaqueProviderData; + +use gateway_admin::{ + model::{ + PageSize, Revision, + auth::LoginCommand, + client_keys::ClientKeyPageSize, + observability::{DecimalAmount, RequestOutcome, TimeRange}, + provider_credentials::{CredentialCommitGuard, ProviderDocument}, + settings::AdminApiKey, + }, + ports::store::{ + AccountStore, AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore, + }, +}; + +mod use_case; + +#[test] +fn revision_should_reject_zero() { + assert!(Revision::new(0).is_err()); +} + +#[test] +fn page_size_should_accept_frozen_upper_bound() { + assert_eq!(PageSize::new(200).map(PageSize::get), Ok(200)); +} + +#[test] +fn page_size_should_reject_value_above_frozen_upper_bound() { + assert!(PageSize::new(201).is_err()); +} + +#[test] +fn client_key_page_size_should_keep_the_full_nonzero_u16_contract() { + assert!(ClientKeyPageSize::new(0).is_err()); + assert_eq!( + ClientKeyPageSize::new(u16::MAX).map(ClientKeyPageSize::get), + Ok(u16::MAX) + ); +} + +#[test] +fn request_outcome_filter_should_preserve_known_and_bounded_other_values() { + let known = RequestOutcome::new("succeeded").expect("known outcome"); + let other = RequestOutcome::new("provider_future_state").expect("other outcome"); + + assert_eq!(known, RequestOutcome::Succeeded); + assert!(matches!(&other, RequestOutcome::Other(_))); + assert_eq!(other.as_str(), "provider_future_state"); + assert!(RequestOutcome::new("").is_err()); + assert!(RequestOutcome::new("a".repeat(RequestOutcome::MAX_BYTES + 1)).is_err()); + assert!(RequestOutcome::new("future\nstate").is_err()); +} + +#[test] +fn time_range_should_accept_exactly_366_days() { + let end = Utc::now(); + assert!(TimeRange::new(end - TimeDelta::days(366), end).is_ok()); +} + +#[test] +fn time_range_should_reject_more_than_366_days() { + let end = Utc::now(); + assert!(TimeRange::new(end - TimeDelta::days(366) - TimeDelta::seconds(1), end).is_err()); +} + +#[test] +fn decimal_amount_should_canonicalize_redundant_zeroes() { + assert_eq!( + DecimalAmount::from_str("00012.34000").map(|amount| amount.to_string()), + Ok("12.34".to_owned()) + ); +} + +#[test] +fn admin_api_key_debug_should_redact_plaintext() { + let key = AdminApiKey::new("secret-admin-key"); + assert!(!format!("{key:?}").contains("secret-admin-key")); +} + +#[test] +fn login_command_debug_should_redact_password() { + let command = LoginCommand { + username: None, + password: "secret-password".to_owned(), + source: "127.0.0.1".to_owned(), + }; + assert!(!format!("{command:?}").contains("secret-password")); +} + +#[test] +fn all_store_capability_traits_should_be_object_safe() { + fn assert_object_safe() {} + + assert_object_safe::(); + assert_object_safe::(); + assert_object_safe::(); + assert_object_safe::(); + assert_object_safe::(); + assert_object_safe::(); + + fn assert_send_object_safe() {} + assert_send_object_safe::(); +} + +#[test] +fn provider_document_debug_should_not_expose_opaque_material() { + let document = ProviderDocument::new(OpaqueProviderData::new(Default::default())); + assert_eq!( + format!("{document:?}"), + "ProviderDocument([PROVIDER_OWNED])" + ); +} diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs new file mode 100644 index 000000000..ba69a3d14 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -0,0 +1,964 @@ +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use chrono::{TimeDelta, Utc}; +use futures::future::BoxFuture; +use gateway_core::{ + engine::{ + credential::{OpaqueProviderData, ProviderAccountId}, + probe::{AccountProbe, AccountProbeRequest, AccountProbeResult}, + }, + error::GatewayError, + routing::{ProviderInstanceId, ProviderKind}, +}; + +use gateway_admin::{ + AdminServices, + model::{ + MutationContext, Revision, + accounts::{ + AccountAvailability, AccountListQuery, AccountPage, AccountRecord, AccountSummary, + AccountUsage, DeleteAccount, SetAccountEnabled, + }, + catalog::{ + CreateProviderInstance, DeleteProviderInstance, ProviderInstance, + ProviderInstanceCatalog, ProviderInstanceDetail, ProviderInstanceMutation, + SetProviderInstanceEnabled, UpdateProviderInstance, + }, + observability::TimeRange, + provider_credentials::{ + AuthorizationCommit, AuthorizationCredentialCommit, AuthorizationMutationTarget, + AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard, CredentialDetails, + CredentialImportCommit, CredentialImportResult, CredentialListQuery, + CredentialMutationResult, CredentialPage, CredentialRotationCommit, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, PreparedAuthorizationCommit, + PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport, + PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument, + ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota, + }, + settings::{ + AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy, + RuntimeSettings, + }, + }, + ports::{ + provider::{ + ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind, ProviderAdminRegistry, + }, + store::{ + AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult, CatalogStore, + SettingsStore, + }, + }, +}; + +pub(super) type EventLog = Arc>>; + +pub(super) struct FakeProviderAdmin { + kind: ProviderKind, + events: EventLog, + failure: Mutex>, + pending: Mutex>, + export_inputs: Mutex>, +} + +impl FakeProviderAdmin { + pub(super) fn new(kind: &str, events: EventLog) -> Arc { + Arc::new(Self { + kind: ProviderKind::new(kind).expect("provider kind"), + events, + failure: Mutex::new(None), + pending: Mutex::new(None), + export_inputs: Mutex::new(Vec::new()), + }) + } + + pub(super) fn fail_next(&self, kind: ProviderAdminErrorKind) { + *self.failure.lock().expect("provider failure") = Some(kind); + } + + pub(super) fn pending(&self) -> Option { + self.pending.lock().expect("pending authorization").clone() + } + + fn export_inputs(&self) -> Vec { + self.export_inputs.lock().expect("export inputs").clone() + } + + fn record(&self, event: &'static str) { + self.events.lock().expect("provider events").push(event); + } + + fn require_available(&self) -> Result<(), ProviderAdminError> { + match self.failure.lock().expect("provider failure").take() { + Some(kind) => Err(ProviderAdminError::new(kind)), + None => Ok(()), + } + } + + fn prepared_rotation(&self, account: &AccountRecord) -> PreparedCredentialRotation { + PreparedCredentialRotation::new( + PreparedCredentialRotationFacts { + account_id: ProviderAccountId::new(account.id.clone()).expect("account ID"), + provider_instance_id: account.provider_instance_id.clone(), + provider_kind: account.provider_kind.clone(), + expected_credential_revision: account.credential_revision, + name: account.name.clone(), + email: account.email.clone(), + plan_type: account.plan_type.clone(), + provider_material: document(), + has_refresh_token: account.has_refresh_token, + access_token_expires_at: account.access_token_expires_at + TimeDelta::hours(1), + next_refresh_at: account.next_refresh_at, + }, + Box::new(RecordingGuard::new(self.events.clone())), + ) + } +} + +#[async_trait] +impl ProviderAdmin for FakeProviderAdmin { + fn provider_kind(&self) -> &ProviderKind { + &self.kind + } + + async fn account_unavailable(&self, _: &ProviderAccountId) { + self.record("provider.account_unavailable"); + } + + fn dashboard_wire_profile( + &self, + ) -> Option { + None + } + + fn calculated_billing( + &self, + _: &gateway_admin::model::observability::ProviderBillingInput, + ) -> Result< + Option, + ProviderAdminError, + > { + Ok(None) + } + + async fn prepare_import( + &self, + command: PrepareCredentialImport, + ) -> Result { + self.record("provider.prepare_import"); + self.require_available()?; + Ok(PreparedCredentialImport { + provider_kind: self.kind.clone(), + provider_instance_id: command.provider_instance_id.clone(), + credentials: vec![prepared_create( + self.kind.clone(), + command.provider_instance_id, + "prepared-import", + )], + }) + } + + async fn start_authorization( + &self, + pending: PendingAuthorizationMutation, + ) -> Result { + self.record("provider.start_authorization"); + self.require_available()?; + *self.pending.lock().expect("pending authorization") = Some(pending); + Ok(AuthorizationStarted { + flow_id: "flow-test".to_owned(), + authorization_url: "https://example.invalid/oauth".to_owned(), + expires_at: Utc::now() + TimeDelta::minutes(10), + }) + } + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result { + self.record("provider.complete_authorization"); + self.require_available()?; + let pending = self + .pending + .lock() + .expect("pending authorization") + .take() + .ok_or_else(unsupported)?; + if !pending.owner_binding().matches_context(&command.context) { + return Err(ProviderAdminError::new(ProviderAdminErrorKind::NotFound)); + } + let credential = match pending.target() { + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + } => PreparedAuthorizationCredential::Create(prepared_create( + self.kind.clone(), + provider_instance_id.clone(), + name, + )), + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => { + let mut account = account_record(self.kind.as_str()); + account.id = account_id.as_str().to_owned(); + account.provider_instance_id = provider_instance_id.clone(); + account.credential_revision = *expected_credential_revision; + PreparedAuthorizationCredential::Reauthorize(self.prepared_rotation(&account)) + } + }; + Ok(PreparedAuthorizationCommit { + pending, + credential, + }) + } + + async fn prepare_rotation( + &self, + command: PrepareCredentialRotation, + ) -> Result { + self.record("provider.prepare_rotation"); + self.require_available()?; + Ok(self.prepared_rotation(&command.account)) + } + + async fn prepare_refresh( + &self, + command: PrepareCredentialRefresh, + ) -> Result { + self.record("provider.prepare_refresh"); + self.require_available()?; + Ok(self.prepared_rotation(&command.account)) + } + + async fn quota( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Ok(ProviderQuota { + observed_at: None, + refresh_token_expires_at: None, + windows: Vec::new(), + provider_data: None, + }) + } + + async fn models( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Ok(ProviderModels { + models: Vec::new(), + observed_at: None, + }) + } + + async fn export_credentials( + &self, + credentials: Vec, + ) -> Result { + self.record("provider.export"); + self.require_available()?; + *self.export_inputs.lock().expect("export inputs") = credentials.clone(); + Ok(ProviderExport { + provider_kind: self.kind.clone(), + account_ids: credentials + .into_iter() + .map(|credential| { + ProviderAccountId::new(credential.account.id).expect("stored account ID") + }) + .collect(), + document: document(), + }) + } +} + +pub(super) struct FakeAccountStore { + events: EventLog, + account: AccountRecord, + fail_commit: Mutex, + audit_requests: Mutex>, + authorization_revision: Mutex>, +} + +impl FakeAccountStore { + pub(super) fn new(kind: &str, events: EventLog) -> Arc { + Arc::new(Self { + events, + account: account_record(kind), + fail_commit: Mutex::new(false), + audit_requests: Mutex::new(Vec::new()), + authorization_revision: Mutex::new(None), + }) + } + + pub(super) fn fail_next_commit(&self) { + *self.fail_commit.lock().expect("store failure") = true; + } + + pub(super) fn audit_requests(&self) -> Vec { + self.audit_requests.lock().expect("audit requests").clone() + } + + pub(super) fn authorization_revision(&self) -> Option { + *self + .authorization_revision + .lock() + .expect("authorization revision") + } + + fn record(&self, event: &'static str) { + self.events.lock().expect("store events").push(event); + } + + fn require_commit(&self) -> AdminStoreResult<()> { + let mut failure = self.fail_commit.lock().expect("store failure"); + if std::mem::take(&mut *failure) { + Err(store_unavailable()) + } else { + Ok(()) + } + } + + fn record_context(&self, context: &MutationContext) { + self.audit_requests + .lock() + .expect("audit requests") + .push(context.request_id.clone()); + } +} + +#[async_trait] +impl AccountStore for FakeAccountStore { + async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult { + Ok(AccountPage { + config_revision: revision(1), + items: vec![self.account.clone()], + total: 1, + summary: AccountSummary { + total: 1, + active: 1, + quota_exhausted: 0, + attention: 0, + }, + }) + } + + async fn load_account(&self, account_id: &str) -> AdminStoreResult> { + self.record("store.load_account"); + Ok((account_id == self.account.id).then(|| self.account.clone())) + } + + async fn load_account_usage( + &self, + _: TimeRange, + _: &[String], + ) -> AdminStoreResult> { + Ok(Vec::new()) + } + + async fn list_credentials( + &self, + provider_kind: &ProviderKind, + _: CredentialListQuery, + ) -> AdminStoreResult { + self.record("store.list_credentials"); + Ok(CredentialPage { + config_revision: revision(1), + items: (provider_kind == &self.account.provider_kind) + .then(|| self.account.clone()) + .into_iter() + .collect(), + next_cursor: None, + }) + } + + async fn credential_details( + &self, + provider_kind: &ProviderKind, + account_id: &ProviderAccountId, + ) -> AdminStoreResult> { + self.record("store.credential_details"); + Ok( + (provider_kind == &self.account.provider_kind + && account_id.as_str() == self.account.id) + .then(|| CredentialDetails { + config_revision: revision(1), + credential: self.account.clone(), + }), + ) + } + + async fn load_credentials_for_export( + &self, + provider_kind: &ProviderKind, + account_ids: &[ProviderAccountId], + ) -> AdminStoreResult> { + self.record("store.load_credentials_for_export"); + if provider_kind != &self.account.provider_kind + || account_ids + .iter() + .any(|account_id| account_id.as_str() != self.account.id) + { + return Err(AdminStoreError::new( + AdminStoreErrorKind::NotFound, + "test account", + "credential not found", + )); + } + Ok(vec![ProviderExportCredentialInput { + account: self.account.clone(), + provider_material: document(), + }]) + } + + async fn commit_credential_import( + &self, + command: CredentialImportCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.commit_import"); + self.record_context(context); + self.require_commit()?; + Ok(CredentialImportResult { + config_revision: revision(2), + credential_ids: command + .prepared + .credentials + .into_iter() + .map(|credential| credential.account_id) + .collect(), + }) + } + + async fn commit_authorization( + &self, + command: AuthorizationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.commit_authorization"); + self.record_context(context); + *self + .authorization_revision + .lock() + .expect("authorization revision") = Some(command.pending.expected_config_revision()); + self.require_commit()?; + let (account_id, credential_revision) = match command.credential { + AuthorizationCredentialCommit::Create(credential) => (credential.account_id, None), + AuthorizationCredentialCommit::Reauthorize(credential) => ( + credential.account_id, + Some(revision(credential.expected_credential_revision.get() + 1)), + ), + }; + Ok(CredentialMutationResult { + config_revision: revision(2), + account_id, + credential_revision, + }) + } + + async fn commit_credential_rotation( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.commit_rotation"); + self.record_context(context); + self.require_commit()?; + Ok(rotation_result(command)) + } + + async fn commit_credential_refresh( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.commit_refresh"); + self.record_context(context); + self.require_commit()?; + Ok(rotation_result(command)) + } + + async fn set_account_enabled( + &self, + _: SetAccountEnabled, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.set_enabled"); + self.record_context(context); + self.require_commit()?; + Ok(revision(2)) + } + + async fn delete_account( + &self, + _: DeleteAccount, + context: &MutationContext, + ) -> AdminStoreResult { + self.record("store.delete"); + self.record_context(context); + self.require_commit()?; + Ok(revision(2)) + } + + async fn record_credential_export( + &self, + _: &[ProviderAccountId], + context: &MutationContext, + ) -> AdminStoreResult<()> { + self.record("store.audit_export"); + self.record_context(context); + Ok(()) + } +} + +struct StaticSettingsStore; + +#[async_trait] +impl SettingsStore for StaticSettingsStore { + async fn load_runtime_settings(&self) -> AdminStoreResult { + Ok(RuntimeSettings { + config_revision: revision(1), + provider_model_mappings: Default::default(), + refresh_margin_seconds: 300, + refresh_concurrency: 2, + max_concurrent_per_account: 1, + request_interval_ms: 0, + rotation_strategy: RotationStrategy::Smart, + usage_retention_days: 30, + ops_event_retention_days: 30, + audit_retention_days: 30, + updated_at: Utc::now(), + }) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn replace_runtime_settings( + &self, + _: ReplaceRuntimeSettings, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn replace_admin_api_key( + &self, + _: Revision, + _: AdminApiKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn delete_admin_api_key( + &self, + _: Revision, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } +} + +struct StaticCatalogStore { + instance: ProviderInstance, +} + +#[async_trait] +impl CatalogStore for StaticCatalogStore { + async fn list_provider_instances(&self, _: bool) -> AdminStoreResult { + Ok(ProviderInstanceCatalog { + config_revision: revision(1), + items: vec![self.instance.clone()], + }) + } + + async fn load_provider_instance( + &self, + id: &ProviderInstanceId, + ) -> AdminStoreResult> { + Ok((id == &self.instance.id).then(|| ProviderInstanceDetail { + config_revision: revision(1), + item: self.instance.clone(), + })) + } + + async fn create_provider_instance( + &self, + _: CreateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn update_provider_instance( + &self, + _: UpdateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn set_provider_instance_enabled( + &self, + _: SetProviderInstanceEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } + + async fn delete_provider_instance( + &self, + _: DeleteProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(store_unavailable()) + } +} + +struct RecordingGuard { + events: EventLog, + finished: bool, +} + +impl RecordingGuard { + fn new(events: EventLog) -> Self { + Self { + events, + finished: false, + } + } +} + +impl CredentialCommitGuard for RecordingGuard { + fn finish(mut self: Box) { + self.events + .lock() + .expect("guard events") + .push("guard.finish"); + self.finished = true; + } +} + +impl Drop for RecordingGuard { + fn drop(&mut self) { + if !self.finished { + self.events.lock().expect("guard events").push("guard.drop"); + } + } +} + +#[test] +fn provider_registry_should_resolve_custom_kind_without_central_match() { + let provider = FakeProviderAdmin::new("custom-provider", events()); + let registry = ProviderAdminRegistry::new([provider as Arc]) + .expect("provider registry"); + let resolved = registry + .require(&ProviderKind::new("custom-provider").expect("provider kind")) + .expect("registered provider"); + assert_eq!(resolved.provider_kind().as_str(), "custom-provider"); +} + +#[test] +fn provider_registry_should_reject_duplicate_kind() { + let first = FakeProviderAdmin::new("duplicate", events()); + let second = FakeProviderAdmin::new("duplicate", events()); + let result = ProviderAdminRegistry::new([ + first as Arc, + second as Arc, + ]); + assert!(matches!( + result, + Err(error) if error.kind() == ProviderAdminErrorKind::Conflict + )); +} + +#[tokio::test] +async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() { + let provider = FakeProviderAdmin::new("openai", events()); + let store = FakeAccountStore::new("openai", events()); + let expected = store.account.clone(); + accounts_service(provider.clone(), store) + .await + .accounts() + .export( + &context("export-complete-input"), + vec![ProviderAccountId::new(expected.id.clone()).expect("account ID")], + ) + .await + .expect("export credentials"); + + let inputs = provider.export_inputs(); + let input = inputs.first().expect("provider export input"); + assert_eq!( + ( + input.account.created_at, + input.account.updated_at, + &input.provider_material, + ), + (expected.created_at, expected.updated_at, &document()), + ); +} + +#[tokio::test] +async fn accounts_refresh_should_keep_guard_through_store_commit() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = accounts_service(provider, store.clone()).await; + + let result = services + .accounts() + .refresh( + &context("refresh-request"), + revision(1), + ProviderAccountId::new("acct_test").expect("account ID"), + ) + .await + .expect("refresh credential"); + assert_eq!(result.config_revision, revision(2)); + assert_eq!(result.account.provider_instance_name, "Test Provider"); + assert_eq!( + result.account.status, + gateway_admin::model::accounts::AccountStatus::Active + ); + + assert_eq!( + recorded(&events), + [ + "store.load_account", + "provider.prepare_refresh", + "store.commit_refresh", + "guard.finish", + "store.load_account", + ] + ); + assert_eq!(store.audit_requests(), ["refresh-request"]); +} + +#[tokio::test] +async fn accounts_list_should_return_complete_directory_semantics() { + let provider = FakeProviderAdmin::new("openai", events()); + let store = FakeAccountStore::new("openai", events()); + let page = accounts_service(provider, store) + .await + .accounts() + .list(AccountListQuery { + page: 1, + page_size: gateway_admin::model::PageSize::new(20).expect("page size"), + provider_kind: None, + search: None, + status: None, + sort: None, + }) + .await + .expect("complete account directory"); + + assert_eq!(page.summary.total, 1); + assert_eq!(page.summary.active, 1); + let account = page.items.first().expect("account item"); + assert_eq!(account.provider_instance_name, "Test Provider"); + assert_eq!( + account.status, + gateway_admin::model::accounts::AccountStatus::Active + ); +} + +#[tokio::test] +async fn accounts_refresh_provider_failure_should_not_call_store_commit() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + provider.fail_next(ProviderAdminErrorKind::Unavailable); + let store = FakeAccountStore::new("openai", events.clone()); + let services = accounts_service(provider, store).await; + + services + .accounts() + .refresh( + &context("refresh-provider-error"), + revision(1), + ProviderAccountId::new("acct_test").expect("account ID"), + ) + .await + .expect_err("Provider preparation must fail"); + + assert_eq!( + recorded(&events), + ["store.load_account", "provider.prepare_refresh"] + ); +} + +#[tokio::test] +async fn accounts_refresh_store_failure_should_drop_guard_after_cas_attempt() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + store.fail_next_commit(); + let services = accounts_service(provider, store).await; + + services + .accounts() + .refresh( + &context("refresh-store-error"), + revision(1), + ProviderAccountId::new("acct_test").expect("account ID"), + ) + .await + .expect_err("Store CAS must fail"); + + assert_eq!( + recorded(&events), + [ + "store.load_account", + "provider.prepare_refresh", + "store.commit_refresh", + "guard.drop", + ] + ); +} + +pub(super) fn events() -> EventLog { + Arc::new(Mutex::new(Vec::new())) +} + +pub(super) fn recorded(events: &EventLog) -> Vec<&'static str> { + events.lock().expect("recorded events").clone() +} + +pub(super) fn context(request_id: &str) -> MutationContext { + MutationContext { + actor: gateway_admin::model::MutationActor::AdminSession { + admin_user_id: "admin-test".to_owned(), + }, + request_id: request_id.to_owned(), + } +} + +pub(super) fn document() -> ProviderDocument { + ProviderDocument::new(OpaqueProviderData::new(Default::default())) +} + +pub(super) fn account_record(kind: &str) -> AccountRecord { + let now = Utc::now(); + AccountRecord { + id: "acct_test".to_owned(), + provider_instance_id: ProviderInstanceId::new(format!("inst_{kind}")).expect("instance ID"), + provider_kind: ProviderKind::new(kind).expect("provider kind"), + name: "test account".to_owned(), + email: Some("test@example.invalid".to_owned()), + upstream_user_id: "upstream-user".to_owned(), + upstream_account_id: None, + plan_type: Some("test".to_owned()), + credential_revision: revision(1), + has_refresh_token: true, + access_token_expires_at: now + TimeDelta::hours(1), + next_refresh_at: Some(now + TimeDelta::minutes(30)), + enabled: true, + availability: AccountAvailability::Ready, + availability_reason: None, + cooldown_until: None, + availability_observed_at: now, + quota_observed_at: None, + created_at: now, + updated_at: now, + } +} + +pub(super) fn revision(value: u64) -> Revision { + Revision::new(value).expect("positive revision") +} + +fn prepared_create( + provider_kind: ProviderKind, + provider_instance_id: ProviderInstanceId, + name: &str, +) -> PreparedCredentialCreate { + let now = Utc::now(); + PreparedCredentialCreate { + account_id: ProviderAccountId::new("acct_prepared").expect("prepared account ID"), + provider_instance_id, + provider_kind, + name: name.to_owned(), + email: Some("prepared@example.invalid".to_owned()), + upstream_user_id: "prepared-user".to_owned(), + upstream_account_id: None, + plan_type: Some("test".to_owned()), + provider_material: document(), + has_refresh_token: true, + access_token_expires_at: now + TimeDelta::hours(1), + next_refresh_at: Some(now + TimeDelta::minutes(30)), + enabled: true, + availability: AccountAvailability::Ready, + availability_reason: None, + cooldown_until: None, + availability_observed_at: now, + } +} + +fn rotation_result(command: CredentialRotationCommit) -> CredentialMutationResult { + CredentialMutationResult { + config_revision: revision(command.expected_config_revision.get() + 1), + account_id: command.prepared.account_id, + credential_revision: Some(revision( + command.prepared.expected_credential_revision.get() + 1, + )), + } +} + +async fn accounts_service( + provider: Arc, + store: Arc, +) -> AdminServices { + let now = Utc::now(); + let instance = ProviderInstance { + id: store.account.provider_instance_id.clone(), + provider_kind: store.account.provider_kind.clone(), + name: "Test Provider".to_owned(), + base_url: "https://example.invalid".to_owned(), + enabled: true, + created_at: now, + updated_at: now, + }; + super::AdminHarness::new() + .accounts(store) + .catalog(Arc::new(StaticCatalogStore { instance })) + .settings(Arc::new(StaticSettingsStore)) + .provider(provider) + .probe(Arc::new(SuccessfulAccountProbe)) + .build() + .await +} + +struct SuccessfulAccountProbe; + +impl AccountProbe for SuccessfulAccountProbe { + fn probe( + &self, + _: AccountProbeRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Ok(AccountProbeResult { + text: vec!["OK".to_owned()], + }) + }) + } +} + +fn store_unavailable() -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "test account", + "unavailable", + ) +} + +fn unsupported() -> ProviderAdminError { + ProviderAdminError::new(ProviderAdminErrorKind::Unsupported) +} diff --git a/backend/crates/gateway-admin/tests/use_case/auth.rs b/backend/crates/gateway-admin/tests/use_case/auth.rs new file mode 100644 index 000000000..65d4c3908 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/auth.rs @@ -0,0 +1,150 @@ +use std::{collections::BTreeMap, sync::Mutex}; + +use async_trait::async_trait; + +use gateway_admin::{ + model::{ + auth::{AdminAuditEvent, AdminSession, LoginCommand}, + settings::AdminApiKey, + }, + ports::store::{AdminStoreResult, AuthStore}, +}; + +#[derive(Default)] +struct MemoryAuthStore { + password_hash: Mutex>, + sessions: Mutex>, + audits: Mutex>, + failures: Mutex, +} + +#[async_trait] +impl AuthStore for MemoryAuthStore { + async fn load_password_hash(&self, _: &str) -> AdminStoreResult> { + Ok(self.password_hash.lock().expect("password hash").clone()) + } + + async fn create_password_hash_if_absent( + &self, + _: &str, + password_hash: &str, + ) -> AdminStoreResult { + let mut stored = self.password_hash.lock().expect("password hash"); + if stored.is_some() { + return Ok(false); + } + *stored = Some(password_hash.to_owned()); + Ok(true) + } + + async fn load_admin_api_key(&self) -> AdminStoreResult> { + Ok(None) + } + + async fn load_session(&self, session_id: &str) -> AdminStoreResult> { + Ok(self + .sessions + .lock() + .expect("sessions") + .get(session_id) + .cloned()) + } + + async fn store_session( + &self, + session_id: &str, + session: &AdminSession, + ) -> AdminStoreResult<()> { + self.sessions + .lock() + .expect("sessions") + .insert(session_id.to_owned(), session.clone()); + Ok(()) + } + + async fn delete_session(&self, session_id: &str) -> AdminStoreResult> { + Ok(self.sessions.lock().expect("sessions").remove(session_id)) + } + + async fn login_source_is_throttled( + &self, + _: &str, + failure_limit: u32, + _: u64, + ) -> AdminStoreResult { + Ok(*self.failures.lock().expect("failures") >= failure_limit) + } + + async fn record_login_failure( + &self, + _: &str, + failure_limit: u32, + _: u64, + ) -> AdminStoreResult { + let mut failures = self.failures.lock().expect("failures"); + *failures += 1; + Ok(*failures >= failure_limit) + } + + async fn clear_login_failures(&self, _: &str) -> AdminStoreResult<()> { + *self.failures.lock().expect("failures") = 0; + Ok(()) + } + + async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()> { + self.audits.lock().expect("audits").push(event); + Ok(()) + } +} + +#[tokio::test] +async fn successful_login_should_create_expiring_session_and_audit() { + let store = std::sync::Arc::new(MemoryAuthStore::default()); + let services = super::AdminHarness::new().auth(store.clone()).build().await; + + let result = services + .auth() + .login(LoginCommand { + username: Some("admin".to_owned()), + password: "strong-test-password".to_owned(), + source: "127.0.0.1".to_owned(), + }) + .await + .expect("login"); + + assert!( + services + .auth() + .validate_session(Some(&result.session_id)) + .await + .expect("validate") + ); + assert_eq!(store.audits.lock().expect("audits").len(), 1); +} + +#[tokio::test] +async fn repeated_default_initialization_should_not_replace_password() { + let store = std::sync::Arc::new(MemoryAuthStore::default()); + super::AdminHarness::new() + .auth(store.clone()) + .default_password("first-strong-password") + .build() + .await; + let services = super::AdminHarness::new() + .auth(store) + .default_password("second-strong-password") + .build() + .await; + + assert!( + services + .auth() + .login(LoginCommand { + username: Some("admin".to_owned()), + password: "first-strong-password".to_owned(), + source: "127.0.0.1".to_owned(), + }) + .await + .is_ok() + ); +} diff --git a/backend/crates/gateway-admin/tests/use_case/catalog.rs b/backend/crates/gateway-admin/tests/use_case/catalog.rs new file mode 100644 index 000000000..afacc8084 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/catalog.rs @@ -0,0 +1,144 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::Utc; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; + +use gateway_admin::{ + model::{ + MutationContext, PageSize, Revision, + catalog::{ + CatalogListQuery, CreateProviderInstance, DeleteProviderInstance, ProviderInstance, + ProviderInstanceCatalog, ProviderInstanceDetail, ProviderInstanceMutation, + SetProviderInstanceEnabled, UpdateProviderInstance, + }, + }, + ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult, CatalogStore}, +}; + +struct CatalogFixtureStore { + items: Vec, +} + +#[async_trait] +impl CatalogStore for CatalogFixtureStore { + async fn list_provider_instances(&self, _: bool) -> AdminStoreResult { + Ok(ProviderInstanceCatalog { + config_revision: Revision::new(7).expect("revision"), + items: self.items.clone(), + }) + } + + async fn load_provider_instance( + &self, + id: &ProviderInstanceId, + ) -> AdminStoreResult> { + Ok(self + .items + .iter() + .find(|item| &item.id == id) + .cloned() + .map(|item| ProviderInstanceDetail { + config_revision: Revision::new(7).expect("revision"), + item, + })) + } + + async fn create_provider_instance( + &self, + _: CreateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } + + async fn update_provider_instance( + &self, + _: UpdateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } + + async fn set_provider_instance_enabled( + &self, + _: SetProviderInstanceEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } + + async fn delete_provider_instance( + &self, + _: DeleteProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } +} + +#[tokio::test] +async fn catalog_cursor_should_resume_after_exact_instance() { + let items = ["inst_a", "inst_b", "inst_c"] + .into_iter() + .map(instance) + .collect(); + let services = super::AdminHarness::new() + .catalog(Arc::new(CatalogFixtureStore { items })) + .build() + .await; + + let page = services + .catalog() + .list(CatalogListQuery { + cursor: Some(ProviderInstanceId::new("inst_a").expect("cursor")), + page_size: PageSize::new(1).expect("page size"), + }) + .await + .expect("catalog page"); + + assert_eq!(page.items[0].id.as_str(), "inst_b"); + assert_eq!( + page.next_cursor.as_ref().map(ProviderInstanceId::as_str), + Some("inst_b") + ); +} + +#[tokio::test] +async fn catalog_detail_should_keep_its_control_plane_revision() { + let services = super::AdminHarness::new() + .catalog(Arc::new(CatalogFixtureStore { + items: vec![instance("inst_detail")], + })) + .build() + .await; + + let detail = services + .catalog() + .get(&ProviderInstanceId::new("inst_detail").expect("instance ID")) + .await + .expect("catalog detail"); + + assert_eq!(detail.config_revision.get(), 7); + assert_eq!(detail.item.id.as_str(), "inst_detail"); +} + +fn instance(id: &str) -> ProviderInstance { + ProviderInstance { + id: ProviderInstanceId::new(id).expect("instance ID"), + provider_kind: ProviderKind::new("openai").expect("provider"), + name: id.to_owned(), + base_url: "https://example.com".to_owned(), + enabled: true, + created_at: Utc::now(), + updated_at: Utc::now(), + } +} + +fn unused() -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "catalog", + "unused in this test", + ) +} diff --git a/backend/crates/gateway-admin/tests/use_case/client_keys.rs b/backend/crates/gateway-admin/tests/use_case/client_keys.rs new file mode 100644 index 000000000..ae15cd563 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/client_keys.rs @@ -0,0 +1,127 @@ +use async_trait::async_trait; +use gateway_core::policy::ClientApiKeyId; + +use gateway_admin::{ + model::{ + AdminErrorKind, MutationContext, Revision, + client_keys::{ + ClientKeyCursor, ClientKeyCursorValue, ClientKeyListQuery, ClientKeyPage, + ClientKeyPageSize, ClientKeyRecord, ClientKeySecret, ClientKeySort, ClientKeySortField, + DeleteClientKey, NewClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey, + }, + }, + ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult, ClientKeyStore}, +}; + +struct UnusedClientKeyStore; + +#[async_trait] +impl ClientKeyStore for UnusedClientKeyStore { + async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult { + assert_eq!(query.page_size.get(), u16::MAX); + Ok(ClientKeyPage { + config_revision: Revision::new(1).expect("revision"), + items: Vec::new(), + total: 0, + next_cursor: None, + }) + } + + async fn reveal_client_key( + &self, + _: &ClientApiKeyId, + ) -> AdminStoreResult> { + Err(unused()) + } + + async fn create_client_key( + &self, + _: NewClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unused()) + } + + async fn update_client_key( + &self, + _: UpdateClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unused()) + } + + async fn set_client_key_enabled( + &self, + _: SetClientKeyEnabled, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unused()) + } + + async fn delete_client_key( + &self, + _: DeleteClientKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } +} + +#[tokio::test] +async fn client_key_cursor_should_reject_value_that_does_not_match_sort() { + let services = super::AdminHarness::new() + .client_keys(std::sync::Arc::new(UnusedClientKeyStore)) + .build() + .await; + let sort = ClientKeySort { + field: ClientKeySortField::Name, + direction: SortDirection::Asc, + }; + let error = services + .client_keys() + .list(ClientKeyListQuery { + cursor: Some(ClientKeyCursor { + sort, + value: ClientKeyCursorValue::Enabled(true), + id: ClientApiKeyId::new("key_cursor").expect("key ID"), + }), + page_size: ClientKeyPageSize::new(50).expect("page size"), + search: None, + sort, + }) + .await + .expect_err("mismatched cursor must fail"); + + assert_eq!(error.kind(), AdminErrorKind::Invalid); +} + +#[tokio::test] +async fn client_key_list_should_forward_the_full_nonzero_u16_page_size() { + let services = super::AdminHarness::new() + .client_keys(std::sync::Arc::new(UnusedClientKeyStore)) + .build() + .await; + let page = services + .client_keys() + .list(ClientKeyListQuery { + cursor: None, + page_size: ClientKeyPageSize::new(u16::MAX).expect("maximum page size"), + search: None, + sort: ClientKeySort { + field: ClientKeySortField::CreatedAt, + direction: SortDirection::Desc, + }, + }) + .await + .expect("maximum page size should reach store"); + + assert_eq!(page.total, 0); +} + +fn unused() -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "client key", + "unused in this test", + ) +} diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs new file mode 100644 index 000000000..89d199fe5 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -0,0 +1,719 @@ +mod accounts; +mod auth; +mod catalog; +mod client_keys; +mod observability; +mod openai; +mod settings; +mod system; +mod xai; + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use futures::future::BoxFuture; +use gateway_admin::{ + AdminConfig, AdminServices, InitialAdminPassword, + model::{ + MutationContext, Revision, + accounts::{ + AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount, + SetAccountEnabled, + }, + auth::{AdminAuditEvent, AdminSession}, + catalog::{ + CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog, + ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled, + UpdateProviderInstance, + }, + client_keys::{ + ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey, + NewClientKey, SetClientKeyEnabled, UpdateClientKey, + }, + observability::{ + DashboardDesktopRelease, DashboardObservation, DashboardWireProfile, + DashboardWireTarget, DesktopReleaseStatus, DiagnosticDimension, DiagnosticObservation, + OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter, + UsageOverview, UsagePage, UsageQuery, + }, + provider_credentials::{ + AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails, + CredentialImportCommit, CredentialImportResult, CredentialListQuery, + CredentialMutationResult, CredentialPage, CredentialRotationCommit, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport, + PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, + ProviderModels, ProviderQuota, + }, + settings::{AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RuntimeSettings}, + system::{SystemOperationAccepted, SystemUpdateDetail, SystemUpdateStatus, SystemVersion}, + }, + ports::{ + provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind}, + store::{ + AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, + AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore, + }, + system::{ + SystemOperationError, SystemOperationErrorKind, SystemOperations, + SystemUpdateEventStream, + }, + }, +}; +use gateway_core::{ + engine::{ + credential::ProviderAccountId, + probe::{AccountProbe, AccountProbeRequest, AccountProbeResult}, + }, + error::{GatewayError, GatewayErrorKind}, + policy::ClientApiKeyId, + routing::{ConfigRevision, ProviderInstanceId, ProviderKind, snapshot::SnapshotControl}, +}; + +pub(super) struct AdminHarness { + default_password: String, + accounts: Arc, + auth: Arc, + catalog: Arc, + client_keys: Arc, + observability: Arc, + settings: Arc, + providers: Vec>, + probe: Arc, + system: Arc, +} + +impl AdminHarness { + pub(super) fn new() -> Self { + let unavailable = Arc::new(UnavailableStore); + Self { + default_password: "strong-test-password".to_owned(), + accounts: unavailable.clone(), + auth: Arc::new(BootstrapAuthStore::default()), + catalog: unavailable.clone(), + client_keys: unavailable.clone(), + observability: unavailable.clone(), + settings: unavailable, + providers: vec![ + Arc::new(UnavailableProvider::new("openai")), + Arc::new(UnavailableProvider::new("xai")), + ], + probe: Arc::new(UnavailableProbe), + system: Arc::new(UnavailableSystem), + } + } + + pub(super) fn default_password(mut self, password: &str) -> Self { + self.default_password = password.to_owned(); + self + } + + pub(super) fn accounts(mut self, store: Arc) -> Self { + self.accounts = store; + self + } + + pub(super) fn auth(mut self, store: Arc) -> Self { + self.auth = store; + self + } + + pub(super) fn catalog(mut self, store: Arc) -> Self { + self.catalog = store; + self + } + + pub(super) fn client_keys(mut self, store: Arc) -> Self { + self.client_keys = store; + self + } + + pub(super) fn observability(mut self, store: Arc) -> Self { + self.observability = store; + self + } + + pub(super) fn settings(mut self, store: Arc) -> Self { + self.settings = store; + self + } + + pub(super) fn provider(mut self, provider: Arc) -> Self { + self.providers + .retain(|registered| registered.provider_kind() != provider.provider_kind()); + self.providers.push(provider); + self + } + + pub(super) fn probe(mut self, probe: Arc) -> Self { + self.probe = probe; + self + } + + pub(super) fn system(mut self, system: Arc) -> Self { + self.system = system; + self + } + + pub(super) async fn build(self) -> AdminServices { + gateway_admin::initialize( + AdminConfig { + session_ttl_minutes: 60, + default_username: "admin".to_owned(), + default_password: InitialAdminPassword::new(self.default_password), + }, + AdminStorePorts::new( + self.accounts, + self.auth, + self.catalog, + self.client_keys, + self.observability, + self.settings, + ), + self.providers, + Arc::new(NoopSnapshot), + self.probe, + self.system, + ) + .await + .expect("initialize admin test harness") + .services() + } +} + +#[derive(Default)] +struct BootstrapAuthStore { + password_hash: Mutex>, +} + +#[async_trait] +impl AuthStore for BootstrapAuthStore { + async fn load_password_hash(&self, _: &str) -> AdminStoreResult> { + Ok(self.password_hash.lock().expect("password hash").clone()) + } + + async fn create_password_hash_if_absent( + &self, + _: &str, + password_hash: &str, + ) -> AdminStoreResult { + let mut stored = self.password_hash.lock().expect("password hash"); + if stored.is_some() { + return Ok(false); + } + *stored = Some(password_hash.to_owned()); + Ok(true) + } + + async fn load_admin_api_key(&self) -> AdminStoreResult> { + Ok(None) + } + + async fn load_session(&self, _: &str) -> AdminStoreResult> { + Ok(None) + } + + async fn store_session(&self, _: &str, _: &AdminSession) -> AdminStoreResult<()> { + Err(unavailable("admin session")) + } + + async fn delete_session(&self, _: &str) -> AdminStoreResult> { + Err(unavailable("admin session")) + } + + async fn login_source_is_throttled(&self, _: &str, _: u32, _: u64) -> AdminStoreResult { + Ok(false) + } + + async fn record_login_failure(&self, _: &str, _: u32, _: u64) -> AdminStoreResult { + Err(unavailable("login failure")) + } + + async fn clear_login_failures(&self, _: &str) -> AdminStoreResult<()> { + Ok(()) + } + + async fn append_audit_event(&self, _: AdminAuditEvent) -> AdminStoreResult<()> { + Ok(()) + } +} + +struct UnavailableStore; + +#[async_trait] +impl AccountStore for UnavailableStore { + async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult { + Err(unavailable("accounts")) + } + + async fn load_account(&self, _: &str) -> AdminStoreResult> { + Err(unavailable("account")) + } + + async fn load_account_usage( + &self, + _: TimeRange, + _: &[String], + ) -> AdminStoreResult> { + Err(unavailable("account usage")) + } + + async fn list_credentials( + &self, + _: &ProviderKind, + _: CredentialListQuery, + ) -> AdminStoreResult { + Err(unavailable("credentials")) + } + + async fn credential_details( + &self, + _: &ProviderKind, + _: &ProviderAccountId, + ) -> AdminStoreResult> { + Err(unavailable("credential")) + } + + async fn load_credentials_for_export( + &self, + _: &ProviderKind, + _: &[ProviderAccountId], + ) -> AdminStoreResult> { + Err(unavailable("credential export")) + } + + async fn commit_credential_import( + &self, + _: CredentialImportCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential import")) + } + + async fn commit_authorization( + &self, + _: AuthorizationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("authorization")) + } + + async fn commit_credential_rotation( + &self, + _: CredentialRotationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential rotation")) + } + + async fn commit_credential_refresh( + &self, + _: CredentialRotationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential refresh")) + } + + async fn set_account_enabled( + &self, + _: SetAccountEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account enabled")) + } + + async fn delete_account( + &self, + _: DeleteAccount, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account delete")) + } + + async fn record_credential_export( + &self, + _: &[ProviderAccountId], + _: &MutationContext, + ) -> AdminStoreResult<()> { + Err(unavailable("credential export audit")) + } +} + +#[async_trait] +impl CatalogStore for UnavailableStore { + async fn list_provider_instances(&self, _: bool) -> AdminStoreResult { + Err(unavailable("provider catalog")) + } + + async fn load_provider_instance( + &self, + _: &ProviderInstanceId, + ) -> AdminStoreResult> { + Err(unavailable("provider instance")) + } + + async fn create_provider_instance( + &self, + _: CreateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider create")) + } + + async fn update_provider_instance( + &self, + _: UpdateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider update")) + } + + async fn set_provider_instance_enabled( + &self, + _: SetProviderInstanceEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider enabled")) + } + + async fn delete_provider_instance( + &self, + _: DeleteProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider delete")) + } +} + +#[async_trait] +impl ClientKeyStore for UnavailableStore { + async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult { + Err(unavailable("client key list")) + } + + async fn reveal_client_key( + &self, + _: &ClientApiKeyId, + ) -> AdminStoreResult> { + Err(unavailable("client key")) + } + + async fn create_client_key( + &self, + _: NewClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key create")) + } + + async fn update_client_key( + &self, + _: UpdateClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key update")) + } + + async fn set_client_key_enabled( + &self, + _: SetClientKeyEnabled, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key enabled")) + } + + async fn delete_client_key( + &self, + _: DeleteClientKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("client key delete")) + } +} + +#[async_trait] +impl ObservabilityStore for UnavailableStore { + async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult { + Err(unavailable("dashboard")) + } + + async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult> { + Err(unavailable("dashboard trend")) + } + + async fn usage_trend( + &self, + _: TimeRange, + _: UsageFilter, + ) -> AdminStoreResult> { + Err(unavailable("usage trend")) + } + + async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult { + Err(unavailable("usage records")) + } + + async fn usage_record_detail(&self, _: &str) -> AdminStoreResult { + Err(unavailable("usage detail")) + } + + async fn usage_summary(&self, _: TimeRange, _: UsageFilter) -> AdminStoreResult { + Err(unavailable("usage summary")) + } + + async fn usage_diagnostics( + &self, + _: TimeRange, + _: UsageFilter, + _: DiagnosticDimension, + ) -> AdminStoreResult> { + Err(unavailable("usage diagnostics")) + } + + async fn list_ops_errors(&self, _: OpsErrorQuery) -> AdminStoreResult { + Err(unavailable("ops errors")) + } +} + +#[async_trait] +impl SettingsStore for UnavailableStore { + async fn load_runtime_settings(&self) -> AdminStoreResult { + Err(unavailable("settings")) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + Err(unavailable("admin API key")) + } + + async fn replace_runtime_settings( + &self, + _: ReplaceRuntimeSettings, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("settings")) + } + + async fn replace_admin_api_key( + &self, + _: Revision, + _: AdminApiKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("admin API key")) + } + + async fn delete_admin_api_key( + &self, + _: Revision, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("admin API key")) + } +} + +struct UnavailableProvider { + kind: ProviderKind, + dashboard_profile: Option, +} + +impl UnavailableProvider { + fn new(kind: &str) -> Self { + Self { + kind: ProviderKind::new(kind).expect("provider kind"), + dashboard_profile: None, + } + } +} + +#[async_trait] +impl ProviderAdmin for UnavailableProvider { + fn provider_kind(&self) -> &ProviderKind { + &self.kind + } + + async fn account_unavailable(&self, _: &ProviderAccountId) {} + + fn dashboard_wire_profile(&self) -> Option { + self.dashboard_profile.clone() + } + + fn calculated_billing( + &self, + _: &gateway_admin::model::observability::ProviderBillingInput, + ) -> Result< + Option, + ProviderAdminError, + > { + Ok(None) + } + + async fn prepare_import( + &self, + _: PrepareCredentialImport, + ) -> Result { + Err(unsupported_provider()) + } + + async fn start_authorization( + &self, + _: PendingAuthorizationMutation, + ) -> Result { + Err(unsupported_provider()) + } + + async fn complete_authorization( + &self, + _: CompleteAuthorization, + ) -> Result { + Err(unsupported_provider()) + } + + async fn prepare_rotation( + &self, + _: PrepareCredentialRotation, + ) -> Result { + Err(unsupported_provider()) + } + + async fn prepare_refresh( + &self, + _: PrepareCredentialRefresh, + ) -> Result { + Err(unsupported_provider()) + } + + async fn quota( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Err(unsupported_provider()) + } + + async fn models( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Err(unsupported_provider()) + } + + async fn export_credentials( + &self, + _: Vec, + ) -> Result { + Err(unsupported_provider()) + } +} + +pub(super) fn dashboard_profile_provider() -> Arc { + Arc::new(UnavailableProvider { + kind: ProviderKind::new("openai").expect("provider kind"), + dashboard_profile: Some(DashboardWireProfile { + originator: "gateway-admin-test".to_owned(), + codex_version: "test".to_owned(), + desktop_version: "test".to_owned(), + desktop_build: "test".to_owned(), + target: DashboardWireTarget { + os_type: "linux".to_owned(), + os_version: "test".to_owned(), + arch: "x86_64".to_owned(), + terminal: "test".to_owned(), + }, + user_agent: "gateway-admin-test".to_owned(), + verified_at: chrono::Utc::now(), + release: DashboardDesktopRelease { + status: DesktopReleaseStatus::Unchecked, + checked_at: None, + latest_version: None, + latest_build: None, + published_at: None, + minimum_system_version: None, + hardware_requirements: None, + download_url: None, + download_size: None, + signature_present: None, + error: None, + }, + }), + }) +} + +struct NoopSnapshot; + +impl SnapshotControl for NoopSnapshot { + fn publish_committed(&self, _: ConfigRevision) -> BoxFuture<'_, ()> { + Box::pin(async {}) + } +} + +struct UnavailableProbe; + +impl AccountProbe for UnavailableProbe { + fn probe( + &self, + _: AccountProbeRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(GatewayError::new( + GatewayErrorKind::Internal, + "test account probe is unavailable", + )) + }) + } +} + +struct UnavailableSystem; + +#[async_trait] +impl SystemOperations for UnavailableSystem { + async fn version(&self) -> Result { + Err(unavailable_system()) + } + + async fn update_detail(&self, _: bool) -> Result { + Err(unavailable_system()) + } + + fn update_events(&self) -> SystemUpdateEventStream { + Box::pin(futures::stream::empty()) + } + + async fn perform_update( + &self, + _: Option, + ) -> Result { + Err(unavailable_system()) + } + + async fn update_status(&self) -> Result { + Err(unavailable_system()) + } + + async fn rollback(&self) -> Result { + Err(unavailable_system()) + } + + async fn restart(&self) -> Result { + Err(unavailable_system()) + } +} + +fn unavailable(resource: &'static str) -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + resource, + "unavailable in this test", + ) +} + +fn unsupported_provider() -> ProviderAdminError { + ProviderAdminError::new(ProviderAdminErrorKind::Unsupported) +} + +fn unavailable_system() -> SystemOperationError { + SystemOperationError::new( + SystemOperationErrorKind::Internal, + "unavailable in this test", + ) +} diff --git a/backend/crates/gateway-admin/tests/use_case/observability.rs b/backend/crates/gateway-admin/tests/use_case/observability.rs new file mode 100644 index 000000000..550dcedfc --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/observability.rs @@ -0,0 +1,430 @@ +use std::{ + str::FromStr as _, + sync::{Arc, Mutex}, +}; + +use async_trait::async_trait; +use chrono::{DateTime, Duration, Timelike as _, Utc}; + +use gateway_admin::{ + AdminServices, + model::{ + MutationContext, Revision, + observability::{ + AccountPoolMetrics, AttemptMetrics, CostCoverage, CurrencyCost, DashboardObservation, + DiagnosticDimension, DiagnosticObservation, Granularity, HealthStatus, OpsErrorPage, + OpsErrorQuery, RequestMetricPoint, RequestMetrics, TimeRange, TrendKind, UsageDetail, + UsageFilter, UsageOverview, UsagePage, UsageQuery, + }, + settings::{ + AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy, + RuntimeSettings, + }, + }, + ports::store::{AdminStoreResult, ObservabilityStore, SettingsStore}, +}; + +#[test] +fn external_observability_range_accepts_exactly_366_days() { + let end = Utc::now(); + let range = TimeRange::new(end - Duration::days(366), end) + .expect("366-day external range should be accepted"); + assert_eq!(range.end, end); +} + +#[test] +fn external_observability_range_rejects_over_366_days_and_reversed_range() { + let end = Utc::now(); + assert!(TimeRange::new(end - Duration::days(367), end).is_err()); + assert!(TimeRange::new(end, end).is_err()); + assert!(TimeRange::new(end + Duration::seconds(1), end).is_err()); +} + +#[tokio::test] +async fn health_timeline_should_keep_exactly_china_day_quarter_hour_slots() { + let now = Utc::now(); + let day_start = china_day_start(now); + let current_slot = quarter_hour_start(now); + let store = Arc::new(FixtureObservabilityStore::new(observation_range(now))); + store.replace_trend(vec![ + health_metric_point( + day_start - Duration::minutes(15), + RequestMetrics { + success_count: 100, + ..RequestMetrics::default() + }, + ), + health_metric_point( + current_slot, + RequestMetrics { + success_count: 2, + ..RequestMetrics::default() + }, + ), + health_metric_point( + current_slot + Duration::minutes(15), + RequestMetrics { + success_count: 100, + ..RequestMetrics::default() + }, + ), + ]); + let services = observability_services(store).await; + + let timeline = services + .observability() + .dashboard_summary(observation_range(now), TrendKind::Usage) + .await + .expect("dashboard summary") + .health_timeline; + + assert_eq!(timeline.points.len(), 96); + assert_eq!( + timeline.points.first().map(|point| point.bucket_start), + Some(day_start) + ); + assert_eq!( + timeline.points.last().map(|point| point.bucket_start), + Some(day_start + Duration::minutes(95 * 15)) + ); + assert_eq!(timeline.success_requests, 2); + assert_eq!( + timeline + .points + .iter() + .find(|point| point.bucket_start == current_slot) + .map(|point| point.status), + Some(HealthStatus::LowSample) + ); + assert!(timeline.points.iter().all(|point| { + point.bucket_start <= quarter_hour_start(Utc::now()) + || (point.status == HealthStatus::Future && point.success_requests == 0) + })); +} + +#[tokio::test] +async fn health_timeline_should_match_legacy_status_precedence_and_thresholds() { + let now = Utc::now(); + let current_slot = quarter_hour_start(now); + let store = Arc::new(FixtureObservabilityStore::new(observation_range(now))); + let services = observability_services(store.clone()).await; + let fixtures = [ + ( + RequestMetrics { + failure_count: 1, + cancelled_count: 1, + incomplete_count: 1, + caller_error_count: 1, + ..RequestMetrics::default() + }, + HealthStatus::NoData, + None, + ), + ( + RequestMetrics { + failure_count: 3, + ..RequestMetrics::default() + }, + HealthStatus::Unavailable, + Some(0.0), + ), + ( + RequestMetrics { + success_count: 1, + failure_count: 1, + ..RequestMetrics::default() + }, + HealthStatus::LowSample, + Some(50.0), + ), + ( + RequestMetrics { + success_count: 98, + failure_count: 2, + ..RequestMetrics::default() + }, + HealthStatus::Unstable, + Some(98.0), + ), + ( + RequestMetrics { + success_count: 99, + failure_count: 1, + ..RequestMetrics::default() + }, + HealthStatus::Stable, + Some(99.0), + ), + ]; + + for (metrics, expected_status, expected_reliability) in fixtures { + store.replace_trend(vec![health_metric_point(current_slot, metrics)]); + let timeline = services + .observability() + .dashboard_summary(observation_range(now), TrendKind::Usage) + .await + .expect("dashboard summary") + .health_timeline; + let point = timeline + .points + .iter() + .find(|point| point.bucket_start == current_slot) + .expect("current health slot"); + assert_eq!(point.status, expected_status); + assert_eq!(point.reliability_percent, expected_reliability); + } +} + +#[tokio::test] +async fn observability_services_should_calculate_usage_insights_and_diagnostic_shares() { + let now = Utc::now(); + let range = observation_range(now); + let store = Arc::new(FixtureObservabilityStore::new(range)); + let metrics = RequestMetrics { + request_count: 10, + success_count: 6, + failure_count: 4, + caller_error_count: 2, + input_tokens: 800, + output_tokens: 200, + total_tokens: 1_000, + latency_count: 5, + ..RequestMetrics::default() + }; + store.replace_overview(UsageOverview { + range, + requests: metrics.clone(), + attempts: AttemptMetrics { + attempt_count: 12, + cost_coverage: CostCoverage { + calculated_count: 10, + ..CostCoverage::default() + }, + costs: vec![CurrencyCost { + currency: "USD".to_owned(), + amount: gateway_admin::model::observability::DecimalAmount::from_str("1.25") + .expect("USD cost"), + }], + ..AttemptMetrics::default() + }, + providers: Vec::new(), + }); + store.replace_trend(vec![RequestMetricPoint { + bucket_start: quarter_hour_start(now), + granularity: Granularity::FifteenMinutes, + metrics, + cost_coverage: CostCoverage::default(), + costs: Vec::new(), + }]); + store.replace_diagnostics(vec![diagnostic("openai", 3), diagnostic("xai", 1)]); + let services = observability_services(store).await; + + let insights = services + .observability() + .usage_insights(range, UsageFilter::default()) + .await + .expect("usage insights"); + let diagnostics = services + .observability() + .diagnostics(range, UsageFilter::default(), DiagnosticDimension::Provider) + .await + .expect("usage diagnostics"); + + assert_eq!(insights.health.total_requests, 8); + assert_eq!(insights.health.failed_requests, 2); + assert_eq!(insights.granularity, Granularity::FifteenMinutes); + assert_eq!(insights.performance.latency_coverage, 0.5); + assert_eq!(insights.cost.tokens_per_request, 100.0); + assert_eq!( + insights + .cost + .estimated_cost + .as_ref() + .map(gateway_admin::model::observability::DecimalAmount::as_str), + Some("1.25") + ); + assert_eq!(diagnostics.items[0].request_share, 0.75); + assert_eq!(diagnostics.items[1].request_share, 0.25); +} + +struct FixtureObservabilityStore { + trend: Mutex>, + overview: Mutex, + diagnostics: Mutex>, +} + +impl FixtureObservabilityStore { + fn new(range: TimeRange) -> Self { + Self { + trend: Mutex::new(Vec::new()), + overview: Mutex::new(UsageOverview { + range, + requests: RequestMetrics::default(), + attempts: AttemptMetrics::default(), + providers: Vec::new(), + }), + diagnostics: Mutex::new(Vec::new()), + } + } + + fn replace_trend(&self, trend: Vec) { + *self.trend.lock().expect("trend") = trend; + } + + fn replace_overview(&self, overview: UsageOverview) { + *self.overview.lock().expect("overview") = overview; + } + + fn replace_diagnostics(&self, diagnostics: Vec) { + *self.diagnostics.lock().expect("diagnostics") = diagnostics; + } +} + +#[async_trait] +impl ObservabilityStore for FixtureObservabilityStore { + async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult { + Ok(DashboardObservation { + range, + requests: RequestMetrics::default(), + attempts: AttemptMetrics::default(), + provider_accounts: AccountPoolMetrics::default(), + trend: self.trend.lock().expect("trend").clone(), + account_usage: Vec::new(), + recent_requests: Vec::new(), + }) + } + + async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult> { + Ok(self.trend.lock().expect("trend").clone()) + } + + async fn usage_trend( + &self, + _: TimeRange, + _: UsageFilter, + ) -> AdminStoreResult> { + Ok(self.trend.lock().expect("trend").clone()) + } + + async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult { + Err(super::unavailable("usage records")) + } + + async fn usage_record_detail(&self, _: &str) -> AdminStoreResult { + Err(super::unavailable("usage detail")) + } + + async fn usage_summary(&self, _: TimeRange, _: UsageFilter) -> AdminStoreResult { + Ok(self.overview.lock().expect("overview").clone()) + } + + async fn usage_diagnostics( + &self, + _: TimeRange, + _: UsageFilter, + _: DiagnosticDimension, + ) -> AdminStoreResult> { + Ok(self.diagnostics.lock().expect("diagnostics").clone()) + } + + async fn list_ops_errors(&self, _: OpsErrorQuery) -> AdminStoreResult { + Err(super::unavailable("ops errors")) + } +} + +struct FixtureSettingsStore; + +#[async_trait] +impl SettingsStore for FixtureSettingsStore { + async fn load_runtime_settings(&self) -> AdminStoreResult { + Ok(RuntimeSettings { + config_revision: Revision::new(1).expect("revision"), + provider_model_mappings: Default::default(), + refresh_margin_seconds: 300, + refresh_concurrency: 2, + max_concurrent_per_account: 1, + request_interval_ms: 0, + rotation_strategy: RotationStrategy::Smart, + usage_retention_days: 31, + ops_event_retention_days: 30, + audit_retention_days: 30, + updated_at: Utc::now(), + }) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + Err(super::unavailable("admin API key")) + } + + async fn replace_runtime_settings( + &self, + _: ReplaceRuntimeSettings, + _: &MutationContext, + ) -> AdminStoreResult { + Err(super::unavailable("settings")) + } + + async fn replace_admin_api_key( + &self, + _: Revision, + _: AdminApiKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(super::unavailable("admin API key")) + } + + async fn delete_admin_api_key( + &self, + _: Revision, + _: &MutationContext, + ) -> AdminStoreResult { + Err(super::unavailable("admin API key")) + } +} + +async fn observability_services(store: Arc) -> AdminServices { + super::AdminHarness::new() + .observability(store) + .settings(Arc::new(FixtureSettingsStore)) + .provider(super::dashboard_profile_provider()) + .build() + .await +} + +fn observation_range(end: DateTime) -> TimeRange { + TimeRange::new(end - Duration::hours(24), end).expect("observation range") +} + +fn health_metric_point(bucket_start: DateTime, metrics: RequestMetrics) -> RequestMetricPoint { + RequestMetricPoint { + bucket_start, + granularity: Granularity::FifteenMinutes, + metrics, + cost_coverage: CostCoverage::default(), + costs: Vec::new(), + } +} + +fn diagnostic(name: &str, request_count: u64) -> DiagnosticObservation { + DiagnosticObservation { + name: name.to_owned(), + request_count, + success_count: request_count, + failure_count: 0, + attempt_count: request_count, + total_tokens: request_count.saturating_mul(100), + average_latency_ms: Some(100), + cost_coverage: CostCoverage::default(), + costs: Vec::new(), + } +} + +fn china_day_start(value: DateTime) -> DateTime { + let elapsed = (value.timestamp() + 8 * 60 * 60).rem_euclid(24 * 60 * 60); + value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond())) +} + +fn quarter_hour_start(value: DateTime) -> DateTime { + let elapsed = value.timestamp().rem_euclid(15 * 60); + value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond())) +} diff --git a/backend/crates/gateway-admin/tests/use_case/openai.rs b/backend/crates/gateway-admin/tests/use_case/openai.rs new file mode 100644 index 000000000..3587f8669 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/openai.rs @@ -0,0 +1,236 @@ +use std::sync::Arc; + +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::routing::ProviderInstanceId; + +use gateway_admin::{ + AdminServices, + model::provider_credentials::{ + AuthorizationMutationTarget, CompleteAuthorization, CredentialMutation, ImportCredentials, + ReauthorizationTarget, StartAuthorization, + }, + ports::provider::ProviderAdminErrorKind, +}; + +use super::accounts::{ + FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision, +}; + +#[tokio::test] +async fn openai_enable_should_use_scoped_store_without_provider_mutation() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .openai() + .enable(mutation("acct_test")) + .await + .expect("enable credential"); + + assert_eq!( + recorded(&events), + ["store.credential_details", "store.set_enabled"] + ); + assert_eq!(store.audit_requests(), ["request-openai"]); +} + +#[tokio::test] +async fn openai_disable_should_commit_then_release_provider_resources() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .openai() + .disable(mutation("acct_test")) + .await + .expect("disable credential"); + + assert_eq!( + recorded(&events), + [ + "store.credential_details", + "store.set_enabled", + "provider.account_unavailable", + ] + ); + assert_eq!(store.audit_requests(), ["request-openai"]); +} + +#[tokio::test] +async fn openai_delete_should_commit_then_release_provider_resources() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .openai() + .delete(mutation("acct_test")) + .await + .expect("delete credential"); + + assert_eq!( + recorded(&events), + [ + "store.credential_details", + "store.delete", + "provider.account_unavailable", + ] + ); + assert_eq!(store.audit_requests(), ["request-openai"]); +} + +#[tokio::test] +async fn openai_disable_store_failure_should_not_release_provider_resources() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + store.fail_next_commit(); + let services = service(provider, store).await; + + services + .openai() + .disable(mutation("acct_test")) + .await + .expect_err("failed Store transaction"); + + assert_eq!( + recorded(&events), + ["store.credential_details", "store.set_enabled"] + ); +} + +#[tokio::test] +async fn openai_import_should_prepare_before_atomic_store_commit() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .openai() + .import_document(ImportCredentials { + context: context("import-openai"), + expected_config_revision: revision(1), + provider_instance_id: instance(), + document: document(), + }) + .await + .expect("import credential"); + + assert_eq!( + recorded(&events), + ["provider.prepare_import", "store.commit_import"] + ); + assert_eq!(store.audit_requests(), ["import-openai"]); +} + +#[tokio::test] +async fn openai_import_provider_error_should_not_touch_store_transaction() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + provider.fail_next(ProviderAdminErrorKind::Invalid); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider, store).await; + + services + .openai() + .import_document(ImportCredentials { + context: context("import-openai-error"), + expected_config_revision: revision(1), + provider_instance_id: instance(), + document: document(), + }) + .await + .expect_err("invalid Provider document"); + + assert_eq!(recorded(&events), ["provider.prepare_import"]); +} + +#[tokio::test] +async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_through_commit() { + let events = events(); + let provider = FakeProviderAdmin::new("openai", events.clone()); + let store = FakeAccountStore::new("openai", events.clone()); + let services = service(provider.clone(), store.clone()).await; + services + .openai() + .start_authorization(StartAuthorization { + context: context("oauth-start-openai"), + expected_config_revision: revision(7), + provider_instance_id: instance(), + name: "reauthorize".to_owned(), + reauthorization: Some(ReauthorizationTarget { + account_id: ProviderAccountId::new("acct_test").expect("account ID"), + credential_revision: revision(1), + }), + }) + .await + .expect("start reauthorization"); + + let pending = provider.pending().expect("pending envelope"); + assert_eq!(pending.expected_config_revision(), revision(7)); + assert_eq!(pending.provider_kind().as_str(), "openai"); + assert_eq!( + pending.owner_binding().started_request_id(), + "oauth-start-openai" + ); + assert!(matches!( + pending.target(), + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } if provider_instance_id == &instance() + && account_id.as_str() == "acct_test" + && *expected_credential_revision == revision(1) + )); + assert!(!format!("{pending:?}").contains("admin-test")); + + services + .openai() + .complete_authorization(CompleteAuthorization { + context: context("oauth-complete-openai"), + flow_id: "flow-test".to_owned(), + callback_url: "http://localhost/callback?code=test&state=test".to_owned(), + }) + .await + .expect("complete reauthorization"); + + assert_eq!( + recorded(&events), + [ + "store.credential_details", + "provider.start_authorization", + "provider.complete_authorization", + "store.commit_authorization", + "guard.finish", + ] + ); + assert_eq!(store.authorization_revision(), Some(revision(7))); + assert_eq!(store.audit_requests(), ["oauth-complete-openai"]); +} + +async fn service(provider: Arc, store: Arc) -> AdminServices { + super::AdminHarness::new() + .provider(provider) + .accounts(store) + .build() + .await +} + +fn mutation(account_id: &str) -> CredentialMutation { + CredentialMutation { + context: context("request-openai"), + expected_config_revision: revision(1), + account_id: ProviderAccountId::new(account_id).expect("account ID"), + } +} + +fn instance() -> ProviderInstanceId { + ProviderInstanceId::new("inst_openai").expect("instance ID") +} diff --git a/backend/crates/gateway-admin/tests/use_case/settings.rs b/backend/crates/gateway-admin/tests/use_case/settings.rs new file mode 100644 index 000000000..f13ea4af3 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/settings.rs @@ -0,0 +1,90 @@ +use async_trait::async_trait; + +use gateway_admin::{ + model::{ + AdminErrorKind, MutationContext, Revision, + settings::{ + AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy, + RuntimeSettings, + }, + }, + ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult, SettingsStore}, +}; + +struct UnusedSettingsStore; + +#[async_trait] +impl SettingsStore for UnusedSettingsStore { + async fn load_runtime_settings(&self) -> AdminStoreResult { + Err(unused()) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + Err(unused()) + } + + async fn replace_runtime_settings( + &self, + _: ReplaceRuntimeSettings, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } + + async fn replace_admin_api_key( + &self, + _: Revision, + _: AdminApiKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } + + async fn delete_admin_api_key( + &self, + _: Revision, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unused()) + } +} + +#[tokio::test] +async fn settings_should_reject_zero_refresh_margin_before_store_call() { + let services = super::AdminHarness::new() + .settings(std::sync::Arc::new(UnusedSettingsStore)) + .build() + .await; + let error = services + .settings() + .replace( + &MutationContext { + actor: gateway_admin::model::MutationActor::System, + request_id: "request-settings".to_owned(), + }, + ReplaceRuntimeSettings { + expected_config_revision: Revision::new(1).expect("revision"), + provider_model_mappings: Default::default(), + refresh_margin_seconds: 0, + refresh_concurrency: 1, + max_concurrent_per_account: 1, + request_interval_ms: 0, + rotation_strategy: RotationStrategy::Smart, + usage_retention_days: 31, + ops_event_retention_days: 30, + audit_retention_days: 30, + }, + ) + .await + .expect_err("invalid settings"); + + assert_eq!(error.kind(), AdminErrorKind::Invalid); +} + +fn unused() -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "settings", + "unused in this test", + ) +} diff --git a/backend/crates/gateway-admin/tests/use_case/system.rs b/backend/crates/gateway-admin/tests/use_case/system.rs new file mode 100644 index 000000000..7e87dae0c --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/system.rs @@ -0,0 +1,115 @@ +use std::sync::Mutex; + +use async_trait::async_trait; + +use gateway_admin::{ + model::system::{ + SystemOperationAccepted, SystemOperationState, SystemOperationStatus, SystemUpdateDetail, + SystemUpdateStatus, SystemVersion, + }, + ports::system::{SystemOperationError, SystemOperations, SystemUpdateEventStream}, +}; + +#[derive(Default)] +struct RecordingSystemOperations { + target: Mutex>>, +} + +#[async_trait] +impl SystemOperations for RecordingSystemOperations { + async fn version(&self) -> Result { + Ok(SystemVersion { + version: "1.0.0".to_owned(), + git_sha: "unknown".to_owned(), + build_time: "unknown".to_owned(), + deployment_mode: "source".to_owned(), + update_channel: "stable".to_owned(), + latest_version: "1.0.0".to_owned(), + has_update: false, + update_cached: false, + update_warning: None, + }) + } + + async fn update_detail(&self, _: bool) -> Result { + Ok(SystemUpdateDetail { + current_version: "1.0.0".to_owned(), + latest_version: "1.0.0".to_owned(), + has_update: false, + deployment_mode: "source".to_owned(), + build_type: "source".to_owned(), + release_url: None, + notes: None, + cached: false, + update_supported: false, + unsupported_reason: Some("source build".to_owned()), + warning: None, + }) + } + + fn update_events(&self) -> SystemUpdateEventStream { + Box::pin(futures::stream::empty()) + } + + async fn perform_update( + &self, + target_version: Option, + ) -> Result { + *self.target.lock().expect("target") = Some(target_version.clone()); + Ok(SystemOperationAccepted::Update { + operation_id: "operation-update".to_owned(), + deployment_mode: "source".to_owned(), + message: "accepted".to_owned(), + need_restart: true, + target_version: target_version.unwrap_or_else(|| "latest".to_owned()), + }) + } + + async fn update_status(&self) -> Result { + Ok(SystemUpdateStatus { + previous_version: None, + current_version: Some("1.0.0".to_owned()), + operation: SystemOperationState { + operation_id: None, + kind: None, + status: SystemOperationStatus::Idle, + target_version: None, + message: None, + error: None, + started_at: None, + finished_at: None, + }, + }) + } + + async fn rollback(&self) -> Result { + Ok(SystemOperationAccepted::Rollback { + operation_id: "operation-rollback".to_owned(), + message: "accepted".to_owned(), + need_restart: true, + }) + } + + async fn restart(&self) -> Result { + Ok(SystemOperationAccepted::Restart { + operation_id: "operation-restart".to_owned(), + message: "accepted".to_owned(), + }) + } +} + +#[tokio::test] +async fn system_update_should_normalize_blank_target_to_latest() { + let operations = std::sync::Arc::new(RecordingSystemOperations::default()); + let services = super::AdminHarness::new() + .system(operations.clone()) + .build() + .await; + services + .system() + .perform_update(Some(" ".to_owned())) + .await + .expect("perform update"); + + assert_eq!(*operations.target.lock().expect("target"), Some(None)); +} diff --git a/backend/crates/gateway-admin/tests/use_case/xai.rs b/backend/crates/gateway-admin/tests/use_case/xai.rs new file mode 100644 index 000000000..170eb28b1 --- /dev/null +++ b/backend/crates/gateway-admin/tests/use_case/xai.rs @@ -0,0 +1,77 @@ +use std::sync::Arc; + +use gateway_core::{engine::credential::ProviderAccountId, routing::ProviderInstanceId}; + +use gateway_admin::{ + AdminServices, + model::provider_credentials::{CredentialMutation, ImportCredentials}, +}; + +use super::accounts::{ + FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision, +}; + +#[tokio::test] +async fn xai_disable_should_commit_then_notify_the_stateless_provider() { + let events = events(); + let provider = FakeProviderAdmin::new("xai", events.clone()); + let store = FakeAccountStore::new("xai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .xai() + .disable(mutation("acct_test")) + .await + .expect("disable credential"); + + assert_eq!( + recorded(&events), + [ + "store.credential_details", + "store.set_enabled", + "provider.account_unavailable", + ] + ); + assert_eq!(store.audit_requests(), ["request-xai"]); +} + +#[tokio::test] +async fn xai_import_should_prepare_before_atomic_store_commit() { + let events = events(); + let provider = FakeProviderAdmin::new("xai", events.clone()); + let store = FakeAccountStore::new("xai", events.clone()); + let services = service(provider, store.clone()).await; + + services + .xai() + .import_document(ImportCredentials { + context: context("import-xai"), + expected_config_revision: revision(3), + provider_instance_id: ProviderInstanceId::new("inst_xai").expect("instance ID"), + document: document(), + }) + .await + .expect("import xAI credential"); + + assert_eq!( + recorded(&events), + ["provider.prepare_import", "store.commit_import"] + ); + assert_eq!(store.audit_requests(), ["import-xai"]); +} + +async fn service(provider: Arc, store: Arc) -> AdminServices { + super::AdminHarness::new() + .provider(provider) + .accounts(store) + .build() + .await +} + +fn mutation(account_id: &str) -> CredentialMutation { + CredentialMutation { + context: context("request-xai"), + expected_config_revision: revision(1), + account_id: ProviderAccountId::new(account_id).expect("account ID"), + } +} diff --git a/backend/crates/gateway-api/Cargo.toml b/backend/crates/gateway-api/Cargo.toml index d8acdd491..c9fbdbe41 100644 --- a/backend/crates/gateway-api/Cargo.toml +++ b/backend/crates/gateway-api/Cargo.toml @@ -12,23 +12,24 @@ name = "main" path = "tests/main.rs" [dependencies] -argon2 = "0.5.3" -async-trait = "0.1.89" axum = { version = "0.8.9", default-features = false, features = ["json", "query", "ws"] } base64 = "0.22.1" chrono = { version = "0.4.45", features = ["serde"] } futures = "0.3.32" +gateway-admin = { path = "../gateway-admin" } gateway-core = { path = "../gateway-core" } gateway-protocol = { path = "../gateway-protocol" } -rand_core = { version = "0.6.4", features = ["getrandom"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150" thiserror = "2.0.18" -tokio = { version = "1.52.3", default-features = false, features = ["macros"] } +tokio = { version = "1.52.3", default-features = false, features = ["macros", "time"] } +tower-http = { version = "0.6.11", features = ["cors", "fs", "request-id", "timeout", "trace"] } tracing = "0.1.44" -subtle = "2.6.1" +url = "2.5.8" +uuid = { version = "1.23.3", features = ["v7"] } [dev-dependencies] +async-trait = "0.1.89" tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] } tower = { version = "0.5.3", features = ["util"] } diff --git a/backend/crates/gateway-api/src/admin/accounts.rs b/backend/crates/gateway-api/src/admin/accounts.rs index 7bdcc8c8c..3ee8de107 100644 --- a/backend/crates/gateway-api/src/admin/accounts.rs +++ b/backend/crates/gateway-api/src/admin/accounts.rs @@ -1,14 +1,42 @@ //! 管理端账号目录、查询校验与安全响应 wire。 -use std::{fmt, pin::Pin}; +use std::{collections::BTreeSet, convert::Infallible, fmt}; -use futures::Stream; -use serde::Deserialize; +use axum::{ + Json, Router, + extract::{Query, State}, + http::{HeaderValue, StatusCode, header::CACHE_CONTROL}, + response::{ + IntoResponse, Response, + sse::{Event, KeepAlive, Sse}, + }, + routing::{get, post}, +}; +use chrono::{DateTime, FixedOffset, Utc}; +use futures::{Stream, StreamExt as _}; +use gateway_admin::model::{ + AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision, + accounts::{ + AccountAvailability, AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, + AccountListQuery, AccountModelUsage, AccountSort, AccountSortField, + AccountStatus as DomainAccountStatus, AccountUsage, SortDirection, + }, + provider_credentials::{ + AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult, + ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow, + }, +}; +use gateway_core::{ + engine::credential::ProviderAccountId, + routing::{ProviderKind, UpstreamModelId}, +}; +use serde::{Deserialize, Serialize}; use serde_json::Value; -use gateway_core::routing::ProviderKind; - -use crate::admin::WireValidationError; +use super::{ + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta, + WireValidationError, +}; const DEFAULT_PAGE_SIZE: u32 = 50; const MAX_PAGE_SIZE: u32 = 200; @@ -28,8 +56,8 @@ pub struct ListQuery { } impl ListQuery { - /// 解析并校验全部 wire 字段。 - pub fn validate(self) -> Result { + /// 解析并校验全部 wire 字段，生成 Admin 查询命令。 + pub fn validate(self) -> Result { let page = self.page.unwrap_or(1); if page == 0 { return Err(WireValidationError::new("page")); @@ -38,8 +66,7 @@ impl ListQuery { if page_size == 0 || page_size > MAX_PAGE_SIZE { return Err(WireValidationError::new("pageSize")); } - let provider = ProviderFilter::parse(self.provider.as_deref().unwrap_or("all")) - .ok_or_else(|| WireValidationError::new("provider"))?; + let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?; let search = self .search .map(|value| value.trim().to_owned()) @@ -52,22 +79,25 @@ impl ListQuery { let status = match self.status.as_deref().map(str::trim) { None | Some("") => None, Some(value) => Some( - AccountStatus::parse(value).ok_or_else(|| WireValidationError::new("status"))?, + parse_account_status(value).ok_or_else(|| WireValidationError::new("status"))?, ), }; let sort = match (self.sort_by.as_deref(), self.sort_direction.as_deref()) { (None, None) => None, (Some(field), Some(direction)) => Some(AccountSort { - field: SortField::parse(field).ok_or_else(|| WireValidationError::new("sortBy"))?, - direction: SortDirection::parse(direction) + field: parse_sort_field(field).ok_or_else(|| WireValidationError::new("sortBy"))?, + direction: parse_sort_direction(direction) .ok_or_else(|| WireValidationError::new("sortDirection"))?, }), _ => return Err(WireValidationError::new("sort")), }; - Ok(ValidatedListQuery { + Ok(AccountListQuery { page, - page_size, - provider, + page_size: PageSize::new( + u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?, + ) + .map_err(|_| WireValidationError::new("pageSize"))?, + provider_kind, search, status, sort, @@ -75,110 +105,47 @@ impl ListQuery { } } -/// 已校验的账号查询。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ValidatedListQuery { - pub page: u32, - pub page_size: u32, - pub provider: ProviderFilter, - pub search: Option, - pub status: Option, - pub sort: Option, +fn parse_provider(value: &str) -> Result, WireValidationError> { + let value = value.trim(); + if value.is_empty() || value.eq_ignore_ascii_case("all") { + return Ok(None); + } + ProviderKind::new(value.to_owned()) + .map(Some) + .map_err(|_| WireValidationError::new("provider")) } -/// Provider 过滤值。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ProviderFilter { - All, - Provider(ProviderKind), -} - -impl ProviderFilter { - fn parse(value: &str) -> Option { - let value = value.trim(); - if value.is_empty() || value.eq_ignore_ascii_case("all") { - return Some(Self::All); - } - ProviderKind::new(value.to_owned()).ok().map(Self::Provider) +fn parse_account_status(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "active" => Some(DomainAccountStatus::Active), + "expired" => Some(DomainAccountStatus::Expired), + "quota_exhausted" => Some(DomainAccountStatus::QuotaExhausted), + "disabled" => Some(DomainAccountStatus::Disabled), + "banned" => Some(DomainAccountStatus::Banned), + _ => None, } } -/// 账号状态过滤值。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AccountStatus { - Active, - Expired, - QuotaExhausted, - Disabled, - Banned, -} - -impl AccountStatus { - fn parse(value: &str) -> Option { - match value.trim().to_ascii_lowercase().as_str() { - "active" => Some(Self::Active), - "expired" => Some(Self::Expired), - "quota_exhausted" => Some(Self::QuotaExhausted), - "disabled" => Some(Self::Disabled), - "banned" => Some(Self::Banned), - _ => None, - } +fn parse_sort_field(value: &str) -> Option { + match value.trim() { + "email" => Some(AccountSortField::Email), + "status" => Some(AccountSortField::Status), + "planType" => Some(AccountSortField::PlanType), + "usage" => Some(AccountSortField::Usage), + "lastUsedAt" => Some(AccountSortField::LastUsedAt), + "expiresAt" => Some(AccountSortField::ExpiresAt), + _ => None, } } -/// 账号排序字段。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SortField { - Email, - Status, - PlanType, - Usage, - LastUsedAt, - ExpiresAt, -} - -impl SortField { - fn parse(value: &str) -> Option { - match value.trim() { - "email" => Some(Self::Email), - "status" => Some(Self::Status), - "planType" => Some(Self::PlanType), - "usage" => Some(Self::Usage), - "lastUsedAt" => Some(Self::LastUsedAt), - "expiresAt" => Some(Self::ExpiresAt), - _ => None, - } +fn parse_sort_direction(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "asc" => Some(SortDirection::Asc), + "desc" => Some(SortDirection::Desc), + _ => None, } } -/// 账号排序方向。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SortDirection { - Asc, - Desc, -} - -impl SortDirection { - fn parse(value: &str) -> Option { - match value.trim().to_ascii_lowercase().as_str() { - "asc" => Some(Self::Asc), - "desc" => Some(Self::Desc), - _ => None, - } - } -} - -/// 已校验的排序组合。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AccountSort { - pub field: SortField, - pub direction: SortDirection, -} - -use serde::Serialize; - -use crate::admin::PageMeta; - /// 账号列表响应数据。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] @@ -333,6 +300,11 @@ impl AccountIdQuery { pub fn validate(&self) -> Result<(), WireValidationError> { require_account_id(&self.id) } + + fn into_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id")) + } } /// 敏感导出的固定 query；IDs 使用逗号分隔，禁止隐式导出全部账号。 @@ -344,7 +316,7 @@ pub struct AccountExportQuery { } impl AccountExportQuery { - pub fn into_ids(self) -> Result, WireValidationError> { + pub fn into_ids(self) -> Result, WireValidationError> { if self.confirm != "export_sensitive_accounts" { return Err(WireValidationError::new("confirm")); } @@ -359,11 +331,13 @@ impl AccountExportQuery { { return Err(WireValidationError::new("ids")); } - let unique = ids.iter().collect::>(); + let unique = ids.iter().collect::>(); if unique.len() != ids.len() { return Err(WireValidationError::new("ids")); } - Ok(ids) + ids.into_iter() + .map(|id| ProviderAccountId::new(id).map_err(|_| WireValidationError::new("ids"))) + .collect() } } @@ -378,6 +352,11 @@ impl AccountActionRequest { pub fn validate(&self) -> Result<(), WireValidationError> { require_account_id(&self.id) } + + fn into_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id")) + } } /// 手工 OAuth 刷新会变更持久 credential，因此必须携带配置 CAS revision。 @@ -396,6 +375,15 @@ impl AccountRefreshRequest { } Ok(()) } + + fn into_command(self) -> Result<(Revision, ProviderAccountId), WireValidationError> { + self.validate()?; + Ok(( + Revision::new(self.expected_config_revision) + .map_err(|_| WireValidationError::new("expectedConfigRevision"))?, + ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))?, + )) + } } /// 连接测试 query；测试仍经唯一 Core/Provider 模型请求路径执行。 @@ -414,6 +402,14 @@ impl AccountTestQuery { } Ok(()) } + + fn into_command(self) -> Result<(ProviderAccountId, UpstreamModelId), WireValidationError> { + self.validate()?; + Ok(( + ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))?, + UpstreamModelId::new(self.model_id).map_err(|_| WireValidationError::new("modelId"))?, + )) + } } #[derive(Debug, Clone, Serialize)] @@ -449,6 +445,23 @@ impl AccountExportData { pub const fn new(value: Value) -> Self { Self(value) } + + fn from_result(bundle: AccountExportBundle) -> Self { + let documents = bundle + .documents + .into_iter() + .map(|document| { + serde_json::json!({ + "provider": document.provider_kind.to_string(), + "document": provider_document_value(document.document), + }) + }) + .collect::>(); + Self::new(serde_json::json!({ + "exportedAt": bundle.exported_at.to_rfc3339(), + "documents": documents, + })) + } } impl fmt::Debug for AccountExportData { @@ -461,58 +474,52 @@ pub struct AccountConnectionTestEvent { pub data: Value, } -impl AccountConnectionTestEvent { - #[must_use] - pub fn started(model: impl Into) -> Self { - Self { - data: serde_json::json!({ +impl From for AccountConnectionTestEvent { + fn from(event: DomainConnectionTestEvent) -> Self { + let data = match event { + DomainConnectionTestEvent::Started { model } => serde_json::json!({ "type": "test_start", - "model": model.into(), + "model": model, "text": "正在连接上游 Responses" }), - } - } - - #[must_use] - pub fn request(payload: Value) -> Self { - Self { - data: serde_json::json!({ "type": "request", "payload": payload }), - } - } - - #[must_use] - pub fn content(text: impl Into) -> Self { - Self { - data: serde_json::json!({ "type": "content", "text": text.into() }), - } - } - - #[must_use] - pub fn completed(account_status: impl Into) -> Self { - Self { - data: serde_json::json!({ + DomainConnectionTestEvent::Request { + model, + input_text, + stream, + store, + } => serde_json::json!({ + "type": "request", + "payload": { + "model": model, + "input": [{ + "role": "user", + "content": [{ "type": "input_text", "text": input_text }] + }], + "stream": stream, + "store": store + } + }), + DomainConnectionTestEvent::Content { text } => { + serde_json::json!({ "type": "content", "text": text }) + } + DomainConnectionTestEvent::Completed { account_status } => serde_json::json!({ "type": "test_complete", "success": true, - "accountStatus": account_status.into() + "accountStatus": account_status_name(account_status) }), - } - } - - #[must_use] - pub fn failed(error: impl Into, account_status: impl Into) -> Self { - Self { - data: serde_json::json!({ + DomainConnectionTestEvent::Failed { + message, + account_status, + } => serde_json::json!({ "type": "error", - "error": error.into(), - "accountStatus": account_status.into() + "error": message, + "accountStatus": account_status_name(account_status) }), - } + }; + Self { data } } } -pub type AccountConnectionTestEventStream = - Pin + Send + 'static>>; - fn require_account_id(value: &str) -> Result<(), WireValidationError> { if value.trim().is_empty() || value.len() > 128 @@ -524,59 +531,10 @@ fn require_account_id(value: &str) -> Result<(), WireValidationError> { Ok(()) } -use async_trait::async_trait; -use axum::{ - Json, Router, - extract::{Query, State}, - http::{HeaderValue, StatusCode, header::CACHE_CONTROL}, - response::{ - IntoResponse, Response, - sse::{Event, KeepAlive, Sse}, - }, - routing::{get, post}, -}; -use futures::StreamExt as _; -use std::convert::Infallible; - -use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminServiceError, AdminServiceErrorKind, - AdminSessionState, -}; - -/// 统一账号目录应用端口。 -#[async_trait] -pub trait AccountAdminService: Send + Sync { - async fn list(&self, query: ValidatedListQuery) -> Result; - async fn export( - &self, - context: &super::AdminRequestContext, - ids: Vec, - ) -> Result; - async fn refresh( - &self, - context: &super::AdminRequestContext, - request: AccountRefreshRequest, - ) -> Result; - async fn quota(&self, id: String) -> Result; - async fn refresh_quota(&self, id: String) -> Result; - async fn models(&self, id: String) -> Result; - async fn refresh_models(&self, id: String) -> Result; - async fn test_connection( - &self, - id: String, - model_id: String, - ) -> Result; -} - -/// 账号目录 HTTP module 所需最小 state。 -pub trait AccountAdminState: AdminSessionState { - fn account_admin_service(&self) -> &dyn AccountAdminService; -} - /// 构造固定 GET 账号目录路由。 pub fn router() -> Router where - S: AccountAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/accounts", get(list_accounts::)) @@ -604,16 +562,18 @@ async fn list_accounts( Query(query): Query, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let query = query - .validate() - .map_err(|error| AdminError::bad_request(format!("Invalid {}", error.field())))?; - let data = state - .account_admin_service() - .list(query) + let command = query.validate().map_err(map_wire_error)?; + let page = command.page; + let page_size = command.page_size.get(); + let result = state + .admin_services() + .accounts() + .list(command) .await .map_err(map_service_error)?; + let data = account_page_data(result, page, page_size, Utc::now()); Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -623,14 +583,16 @@ async fn export_accounts( Query(query): Query, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { let ids = query.into_ids().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .export(auth.context(), ids) + let result = state + .admin_services() + .accounts() + .export(&auth.context().mutation_context(), ids) .await .map_err(map_service_error)?; + let data = AccountExportData::from_result(result); let mut response = AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)).into_response(); response .headers_mut() @@ -644,14 +606,20 @@ async fn refresh_account( Json(request): Json, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - request.validate().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .refresh(auth.context(), request) + let (expected_config_revision, account_id) = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .refresh( + &auth.context().mutation_context(), + expected_config_revision, + account_id, + ) .await .map_err(map_service_error)?; + let data = account_refresh_data(result, Utc::now()); Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -661,14 +629,18 @@ async fn account_quota( Query(query): Query, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .quota(query.id) + let account_id = query.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .quota(&account_id, false) .await .map_err(map_service_error)?; + let data = AccountQuotaData { + account: account_view(result, Utc::now()), + }; Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -678,14 +650,18 @@ async fn refresh_account_quota( Json(request): Json, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - request.validate().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .refresh_quota(request.id) + let account_id = request.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .quota(&account_id, true) .await .map_err(map_service_error)?; + let data = AccountQuotaData { + account: account_view(result, Utc::now()), + }; Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -695,14 +671,16 @@ async fn account_models( Query(query): Query, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .models(query.id) + let account_id = query.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .models(&account_id, false) .await .map_err(map_service_error)?; + let data = account_models_data(result); Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -712,14 +690,16 @@ async fn refresh_account_models( Json(request): Json, ) -> Result where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - request.validate().map_err(map_wire_error)?; - let data = state - .account_admin_service() - .refresh_models(request.id) + let account_id = request.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .models(&account_id, true) .await .map_err(map_service_error)?; + let data = account_models_data(result); Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -729,33 +709,508 @@ async fn test_account_connection( Query(query): Query, ) -> Result>>, AdminError> where - S: AccountAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; + let (account_id, upstream_model) = query.into_command().map_err(map_wire_error)?; let stream = state - .account_admin_service() - .test_connection(query.id, query.model_id) + .admin_services() + .accounts() + .test_connection(account_id, upstream_model) .await .map_err(map_service_error)? .map(|event| { + let event = AccountConnectionTestEvent::from(event); let data = serde_json::to_string(&event.data).unwrap_or_else(|_| "{}".to_owned()); Ok(Event::default().data(data)) }); Ok(Sse::new(stream).keep_alive(KeepAlive::default())) } +fn account_page_data( + result: AccountDirectoryPage, + page: u32, + page_size: u16, + now: DateTime, +) -> AccountPageData { + let total_pages = if result.total == 0 { + 0 + } else { + u32::try_from(result.total.div_ceil(u64::from(page_size))).unwrap_or(u32::MAX) + }; + AccountPageData { + config_revision: result.config_revision.get(), + items: result + .items + .into_iter() + .map(|item| account_view(item, now)) + .collect(), + page: PageMeta::new(page, u32::from(page_size), result.total, total_pages), + summary: AccountSummaryView { + total: result.summary.total, + active: result.summary.active, + quota_exhausted: result.summary.quota_exhausted, + attention: result.summary.attention, + }, + } +} + +fn account_refresh_data(result: AccountRefreshResult, now: DateTime) -> AccountRefreshData { + AccountRefreshData { + config_revision: result.config_revision.get(), + account: account_view(result.account, now), + } +} + +fn account_models_data(result: ProviderModels) -> AccountModelsData { + AccountModelsData { + models: result + .models + .into_iter() + .map(|model| { + let id = model.id.to_string(); + AccountModelView { + label: id.clone(), + id, + } + }) + .collect(), + } +} + +fn account_view(item: AccountDirectoryItem, now: DateTime) -> AccountView { + let AccountDirectoryItem { + account, + provider_instance_name, + status, + usage, + rolling_usage, + quota, + } = item; + let status = account_status_name(status).to_owned(); + let expires_at = china_rfc3339(&account.access_token_expires_at); + let added_at = china_rfc3339(&account.created_at); + let updated_at = china_rfc3339(&account.updated_at); + let (quota, refresh_token_expires_at) = + account_quota_view(&account.provider_kind, quota, rolling_usage.as_ref(), now); + AccountView { + id: account.id.clone(), + name: account.name, + provider: account.provider_kind.to_string(), + provider_instance_id: account.provider_instance_id.to_string(), + provider_instance_name, + resource_ref: account.id, + email: account.email, + account_id: account.upstream_account_id, + user_id: Some(account.upstream_user_id), + label: None, + plan_type: account.plan_type, + has_refresh_token: account.has_refresh_token, + status: status.clone(), + display_status: status, + token_refreshing: false, + availability: account_availability_name(account.availability).to_owned(), + enabled: account.enabled, + credential_revision: account.credential_revision.get(), + state_revision: None, + access_token_expires_at: Some(expires_at), + access_token_expires_at_display: Some(china_datetime(&account.access_token_expires_at)), + refresh_token_expires_at, + next_refresh_at: account.next_refresh_at.map(|value| china_rfc3339(&value)), + added_at, + added_at_display: china_datetime(&account.created_at), + updated_at, + updated_at_display: china_datetime(&account.updated_at), + quota, + usage: account_usage_view(usage, now), + } +} + +fn account_quota_view( + provider_kind: &ProviderKind, + quota: ProviderQuota, + rolling_usage: Option<&AccountUsage>, + now: DateTime, +) -> (AccountQuotaView, Option) { + let refresh_token_expires_at = quota + .refresh_token_expires_at + .map(|value| china_rfc3339(&value)); + let refreshed_at_display = quota + .observed_at + .map_or_else(|| "—".to_owned(), |value| relative_time(value, now)); + let windows = quota + .windows + .into_iter() + .map(|window| quota_window_view(provider_kind, window, rolling_usage)) + .collect(); + ( + AccountQuotaView { + refreshed_at_display, + windows, + }, + refresh_token_expires_at, + ) +} + +fn quota_window_view( + provider_kind: &ProviderKind, + window: ProviderQuotaWindow, + rolling_usage: Option<&AccountUsage>, +) -> AccountQuotaWindowView { + let ProviderQuotaWindow { + key, + group, + source, + window_seconds, + used_percent, + reset_at, + provider_data, + } = window; + let local_usage = if provider_kind.as_str() == "xai" { + provider_data.map(provider_document_value).or_else(|| { + (key == "free-rolling-24h") + .then(|| rolling_usage.map(free_quota_local_usage)) + .flatten() + }) + } else { + None + }; + AccountQuotaWindowView { + label_display: quota_window_label( + provider_kind, + &key, + &group, + source.as_deref(), + window_seconds, + ), + key, + group, + window_seconds, + used_percent, + used_percent_display: used_percent + .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), + local_usage, + reset_at_display: reset_at.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), + } +} + +fn quota_window_label( + provider_kind: &ProviderKind, + key: &str, + group: &str, + source: Option<&str>, + window_seconds: Option, +) -> String { + if provider_kind.as_str() == "xai" && key == "free-rolling-24h" { + return "24小时额度".to_owned(); + } + let base = match (provider_kind.as_str(), group) { + (_, "monthly") => "月限额".to_owned(), + ("openai", "shortTerm") if window_seconds.is_some_and(|seconds| seconds > 86_400) => { + "周限额".to_owned() + } + ("openai", "shortTerm") => "5小时限额".to_owned(), + ("xai", "shortTerm") => "周限额".to_owned(), + _ => custom_quota_window_label(window_seconds), + }; + if provider_kind.as_str() != "openai" { + return base; + } + let Some(source) = source else { + return base; + }; + if matches!(source, "core" | "spend_control" | "monthly_limit") { + return base; + } + let source = if is_codex_review_limit(source) { + "代码审查" + } else { + source + }; + format!("{source} · {base}") +} + +fn custom_quota_window_label(window_seconds: Option) -> String { + let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else { + return "额度".to_owned(); + }; + if seconds % 86_400 == 0 { + format!("{}天限额", seconds / 86_400) + } else if seconds % 3_600 == 0 { + format!("{}小时限额", seconds / 3_600) + } else { + format!("{}分钟限额", seconds.div_ceil(60)) + } +} + +fn is_codex_review_limit(value: &str) -> bool { + let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_"); + matches!( + normalized.as_str(), + "review" | "code_review" | "codex_review" | "codex_code_review" + ) || normalized.contains("code_review") + || normalized.contains("codex_review") +} + +fn free_quota_local_usage(usage: &AccountUsage) -> Value { + let total_tokens = usage.total_tokens.unwrap_or_default(); + serde_json::json!({ + "requestCount": usage.request_count, + "requestCountDisplay": format_number(usage.request_count), + "inputTokens": usage.input_tokens.unwrap_or_default(), + "inputTokensDisplay": display_optional_tokens(usage.input_tokens), + "outputTokens": usage.output_tokens.unwrap_or_default(), + "outputTokensDisplay": display_optional_tokens(usage.output_tokens), + "cachedTokens": usage.cached_tokens.unwrap_or_default(), + "cachedTokensDisplay": display_optional_tokens(usage.cached_tokens), + "totalTokens": total_tokens, + "totalTokensDisplay": format_tokens(total_tokens), + }) +} + +fn account_usage_view(usage: Option, now: DateTime) -> AccountUsageView { + let Some(usage) = usage else { + return empty_account_usage(); + }; + let known_count = usage + .cost_coverage + .provider_reported_count + .saturating_add(usage.cost_coverage.calculated_count); + let cost_estimate_status = if known_count == 0 { + "unknown" + } else if usage.cost_coverage.unavailable_count > 0 { + "partial" + } else { + "known" + }; + AccountUsageView { + request_count: Some(usage.request_count), + request_count_display: format_number(usage.request_count), + input_tokens: usage.input_tokens, + input_tokens_display: display_optional_tokens(usage.input_tokens), + output_tokens: usage.output_tokens, + output_tokens_display: display_optional_tokens(usage.output_tokens), + cached_tokens: usage.cached_tokens, + cached_tokens_display: display_optional_tokens(usage.cached_tokens), + total_tokens: usage.total_tokens, + total_tokens_display: display_optional_tokens(usage.total_tokens), + created_tokens: usage.cache_write_tokens, + created_tokens_display: display_optional_tokens(usage.cache_write_tokens), + read_tokens: usage.cached_tokens, + read_tokens_display: display_optional_tokens(usage.cached_tokens), + last_used_at: usage.last_used_at.map(|value| china_rfc3339(&value)), + last_used_at_display: usage + .last_used_at + .map_or_else(|| "—".to_owned(), |value| relative_time(value, now)), + cost_estimate_status: cost_estimate_status.to_owned(), + known_cost_count: Some(known_count), + partial_cost_count: Some(u64::from(cost_estimate_status == "partial")), + unknown_cost_count: Some(usage.cost_coverage.unavailable_count), + costs: usage.costs.iter().map(account_currency_cost_view).collect(), + models: usage + .models + .into_iter() + .map(|model| account_model_usage_view(model, now)) + .collect(), + } +} + +fn account_model_usage_view(usage: AccountModelUsage, now: DateTime) -> ModelUsageView { + let known_count = usage + .cost_coverage + .provider_reported_count + .saturating_add(usage.cost_coverage.calculated_count); + let cost_estimate_status = if known_count == 0 { + "unknown" + } else if usage.cost_coverage.unavailable_count > 0 { + "partial" + } else { + "known" + }; + let usd = usage + .costs + .iter() + .find(|cost| cost.currency.eq_ignore_ascii_case("USD")); + ModelUsageView { + model: usage.model, + request_count: usage.request_count, + request_count_display: format_number(usage.request_count), + success_rate: (usage.request_count > 0) + .then(|| usage.success_count as f64 * 100.0 / usage.request_count as f64), + success_rate_display: if usage.request_count == 0 { + "—".to_owned() + } else { + format!( + "{:.1}%", + usage.success_count as f64 * 100.0 / usage.request_count as f64 + ) + }, + input_tokens: usage.input_tokens, + input_tokens_display: display_optional_tokens(usage.input_tokens), + output_tokens: usage.output_tokens, + output_tokens_display: display_optional_tokens(usage.output_tokens), + cached_tokens: usage.cached_tokens, + cached_tokens_display: display_optional_tokens(usage.cached_tokens), + total_tokens: usage.total_tokens, + total_tokens_display: display_optional_tokens(usage.total_tokens), + billing_amount_usd: usd.map(|cost| cost.amount.as_str().to_owned()), + billing_amount_usd_display: usd.map_or_else( + || "—".to_owned(), + |cost| format!("${}", cost.amount.as_str()), + ), + cost_estimate_status: cost_estimate_status.to_owned(), + known_cost_count: known_count, + partial_cost_count: u64::from(cost_estimate_status == "partial"), + unknown_cost_count: usage.cost_coverage.unavailable_count, + costs: usage.costs.iter().map(account_currency_cost_view).collect(), + last_used_at: china_rfc3339(&usage.last_used_at), + last_used_at_display: relative_time(usage.last_used_at, now), + } +} + +fn account_currency_cost_view(cost: &AccountCost) -> CurrencyCostView { + CurrencyCostView { + currency: cost.currency.clone(), + estimated_amount: cost.amount.as_str().to_owned(), + estimated_amount_display: format!("{} {}", cost.currency, cost.amount.as_str()), + } +} + +fn display_optional_tokens(value: Option) -> String { + value.map_or_else(|| "—".to_owned(), format_tokens) +} + +fn empty_account_usage() -> AccountUsageView { + AccountUsageView { + request_count: None, + request_count_display: "—".to_owned(), + input_tokens: None, + input_tokens_display: "—".to_owned(), + output_tokens: None, + output_tokens_display: "—".to_owned(), + cached_tokens: None, + cached_tokens_display: "—".to_owned(), + total_tokens: None, + total_tokens_display: "—".to_owned(), + created_tokens: None, + created_tokens_display: "—".to_owned(), + read_tokens: None, + read_tokens_display: "—".to_owned(), + last_used_at: None, + last_used_at_display: "—".to_owned(), + cost_estimate_status: "unavailable".to_owned(), + known_cost_count: None, + partial_cost_count: None, + unknown_cost_count: None, + costs: Vec::new(), + models: Vec::new(), + } +} + +fn provider_document_value(document: ProviderDocument) -> Value { + Value::Object(document.into_provider_data().into_inner()) +} + +fn account_status_name(status: DomainAccountStatus) -> &'static str { + match status { + DomainAccountStatus::Active => "active", + DomainAccountStatus::Expired => "expired", + DomainAccountStatus::QuotaExhausted => "quota_exhausted", + DomainAccountStatus::Disabled => "disabled", + DomainAccountStatus::Banned => "banned", + DomainAccountStatus::Attention => "attention", + } +} + +fn account_availability_name(availability: AccountAvailability) -> &'static str { + match availability { + AccountAvailability::Unknown => "unknown", + AccountAvailability::Ready => "ready", + AccountAvailability::Cooldown => "cooldown", + AccountAvailability::QuotaExhausted => "quota_exhausted", + AccountAvailability::Expired => "expired", + AccountAvailability::Banned => "banned", + AccountAvailability::Invalid => "invalid", + } +} + +fn format_number(value: u64) -> String { + let text = value.to_string(); + let mut output = String::with_capacity(text.len() + text.len() / 3); + for (index, character) in text.chars().rev().enumerate() { + if index > 0 && index % 3 == 0 { + output.push(','); + } + output.push(character); + } + output.chars().rev().collect() +} + +fn format_tokens(value: u64) -> String { + if value < 1_000 { + return format_number(value); + } + for (suffix, threshold) in [ + ("P", 1_000_000_000_000_000_u64), + ("T", 1_000_000_000_000_u64), + ("B", 1_000_000_000_u64), + ("M", 1_000_000_u64), + ("K", 1_000_u64), + ] { + if value >= threshold { + let scaled = value as f64 / threshold as f64; + return format!("{scaled:.1}{suffix}").replace(".0", ""); + } + } + format_number(value) +} + +fn relative_time(value: DateTime, now: DateTime) -> String { + let elapsed = now.signed_duration_since(value); + if elapsed.num_seconds() < 0 { + return china_datetime(&value); + } + if elapsed.num_seconds() < 60 { + return "刚刚".to_owned(); + } + if elapsed.num_minutes() < 60 { + return format!("{} 分钟前", elapsed.num_minutes()); + } + if elapsed.num_hours() < 24 { + return format!("{} 小时前", elapsed.num_hours()); + } + format!("{} 天前", elapsed.num_days()) +} + +fn china_offset() -> FixedOffset { + FixedOffset::east_opt(8 * 60 * 60).expect("UTC+8 is a valid fixed offset") +} + +fn china_rfc3339(value: &DateTime) -> String { + value.with_timezone(&china_offset()).to_rfc3339() +} + +fn china_datetime(value: &DateTime) -> String { + value + .with_timezone(&china_offset()) + .format("%Y-%m-%d %H:%M:%S") + .to_string() +} + fn map_wire_error(error: WireValidationError) -> AdminError { AdminError::bad_request(format!("Invalid {}", error.field())) } fn map_service_error(error: AdminServiceError) -> AdminError { match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => { + AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()), + AdminErrorKind::Unauthorized => AdminError::admin_session_required(), + AdminErrorKind::NotFound => AdminError::not_found(error.to_string()), + AdminErrorKind::Conflict => AdminError::conflict(error.to_string()), + AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(), + AdminErrorKind::BadGateway => AdminError::bad_gateway(error.to_string()), + AdminErrorKind::Unavailable => { AdminError::service_unavailable("Account directory unavailable") } - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), + AdminErrorKind::Internal => AdminError::internal(error.to_string()), } } diff --git a/backend/crates/gateway-api/src/admin/auth.rs b/backend/crates/gateway-api/src/admin/auth.rs index 885bd98aa..54b5415d7 100644 --- a/backend/crates/gateway-api/src/admin/auth.rs +++ b/backend/crates/gateway-api/src/admin/auth.rs @@ -1,9 +1,7 @@ -//! 管理员登录、会话、principal 与认证服务端口。 +//! 管理员登录、会话与管理请求鉴权接线。 -use std::{fmt, sync::Arc}; +use std::fmt; -use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier, password_hash::SaltString}; -use async_trait::async_trait; use axum::{ Extension, Json, Router, extract::{FromRequestParts, State}, @@ -11,135 +9,21 @@ use axum::{ response::{IntoResponse, Response}, routing::{get, post}, }; -use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; -use chrono::{DateTime, Duration, Utc}; -use rand_core::{OsRng, RngCore}; +use gateway_admin::{ + AdminServices, + model::auth::{AdminPrincipal, AdminRequestContext, LoginCommand, LoginError}, +}; use serde::{Deserialize, Serialize}; -use subtle::ConstantTimeEq as _; use super::{AdminEnvelope, AdminError, AdminResponse}; -/// 管理端服务错误分类。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AdminServiceErrorKind { - Invalid, - NotFound, - Conflict, - Unavailable, - Internal, -} - -/// 管理端服务错误。 -#[derive(Debug, Clone, thiserror::Error)] -#[error("{message}")] -pub struct AdminServiceError { - kind: AdminServiceErrorKind, - message: String, -} - -impl AdminServiceError { - #[must_use] - pub fn new(kind: AdminServiceErrorKind, message: impl Into) -> Self { - Self { - kind, - message: message.into(), - } - } - - #[must_use] - pub fn invalid(message: impl Into) -> Self { - Self::new(AdminServiceErrorKind::Invalid, message) - } - - #[must_use] - pub fn not_found(message: impl Into) -> Self { - Self::new(AdminServiceErrorKind::NotFound, message) - } - - #[must_use] - pub fn conflict(message: impl Into) -> Self { - Self::new(AdminServiceErrorKind::Conflict, message) - } - - #[must_use] - pub fn unavailable(message: impl Into) -> Self { - Self::new(AdminServiceErrorKind::Unavailable, message) - } - - #[must_use] - pub fn internal(message: impl Into) -> Self { - Self::new(AdminServiceErrorKind::Internal, message) - } - - #[must_use] - pub fn kind(&self) -> AdminServiceErrorKind { - self.kind - } -} - -/// 已认证的管理端主体。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum AdminPrincipal { - Session { admin_user_id: String }, - ApiKey, -} - -/// 传给应用服务的安全请求上下文。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminRequestContext { - principal: AdminPrincipal, - request_id: String, -} - -impl AdminRequestContext { - #[must_use] - pub fn new(principal: AdminPrincipal, request_id: impl Into) -> Self { - Self { - principal, - request_id: request_id.into(), - } - } - - #[must_use] - pub fn principal(&self) -> &AdminPrincipal { - &self.principal - } - - #[must_use] - pub fn admin_user_id(&self) -> Option<&str> { - match &self.principal { - AdminPrincipal::Session { admin_user_id } => Some(admin_user_id), - AdminPrincipal::ApiKey => None, - } - } - - #[must_use] - pub fn request_id(&self) -> &str { - &self.request_id - } - - #[must_use] - pub fn is_api_key(&self) -> bool { - matches!(self.principal, AdminPrincipal::ApiKey) - } -} - const REQUEST_ID_HEADER: &str = "x-request-id"; +const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; +const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax"; -#[async_trait] -pub trait AdminSessionResolver: Send + Sync { - async fn resolve_admin_user_id( - &self, - session_id: Option<&str>, - ) -> Result, AdminServiceError>; - - async fn verify_admin_api_key(&self, _key: &str) -> Result { - Ok(false) - } -} - +/// 所有管理 HTTP 模块从 state 消费同一个认证用例端口。 pub trait AdminSessionState { - fn admin_session_resolver(&self) -> &dyn AdminSessionResolver; + fn admin_services(&self) -> &AdminServices; } /// 已通过管理员会话或部署级管理 API Key 鉴权的请求。 @@ -149,7 +33,7 @@ pub struct AdminAuth { impl AdminAuth { #[must_use] - pub fn context(&self) -> &AdminRequestContext { + pub const fn context(&self) -> &AdminRequestContext { &self.context } @@ -175,7 +59,10 @@ where .filter(|value| !value.is_empty()) .ok_or_else(|| AdminError::internal("Missing admin request context"))?; Ok(Self { - context: AdminRequestContext::new(principal, request_id), + context: AdminRequestContext { + principal, + request_id: request_id.to_owned(), + }, }) } } @@ -185,7 +72,8 @@ where S: AdminSessionState + Send + Sync, { match state - .admin_session_resolver() + .admin_services() + .auth() .resolve_admin_user_id(admin_session_cookie(headers).as_deref()) .await { @@ -201,7 +89,8 @@ where { if let Some(api_key) = admin_api_key_header(headers) { return match state - .admin_session_resolver() + .admin_services() + .auth() .verify_admin_api_key(&api_key) .await { @@ -221,10 +110,7 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option { (!value.is_empty()).then(|| value.to_owned()) } -const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; -const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax"; - -/// 管理员登录请求；密码不得进入 `Debug`。 +/// 管理员登录请求；密码不得进入 Debug。 #[derive(Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct AdminLoginRequest { @@ -233,7 +119,7 @@ pub struct AdminLoginRequest { } impl AdminLoginRequest { - /// 取出认证服务所需字段。 + /// 取出认证用例所需字段。 #[must_use] pub fn into_parts(self) -> (Option, String) { (self.username, self.password) @@ -258,7 +144,6 @@ pub struct AdminLoginData { } impl AdminLoginData { - /// 构造登录成功响应。 #[must_use] pub fn new(expires_at: String) -> Self { Self { expires_at } @@ -273,7 +158,6 @@ pub struct AdminSessionStatusData { } impl AdminSessionStatusData { - /// 构造会话状态响应。 #[must_use] pub const fn new(authenticated: bool) -> Self { Self { authenticated } @@ -287,7 +171,6 @@ pub struct AdminLogoutData { } impl AdminLogoutData { - /// 构造稳定登出响应。 #[must_use] pub const fn new() -> Self { Self { @@ -318,45 +201,10 @@ impl AdminLoginSource { } } -/// 登录成功后由应用层返回的 session 事实。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminLoginSession { - pub session_id: String, - pub expires_at: String, -} - -/// 登录特有的低基数失败分类。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AdminLoginError { - InvalidCredentials, - Throttled, - Unavailable, -} - -/// 管理员认证应用端口。 -#[async_trait] -pub trait AdminAuthService: Send + Sync { - async fn login( - &self, - source: &str, - username: Option, - password: String, - ) -> Result; - - async fn validate(&self, session_id: Option) -> Result; - - async fn logout(&self, session_id: String) -> Result<(), AdminServiceError>; -} - -/// 认证 HTTP module 的最小 state。 -pub trait AdminAuthState { - fn admin_auth_service(&self) -> &dyn AdminAuthService; -} - /// 构造固定 GET/POST 管理员认证路由。 pub fn router() -> Router where - S: AdminAuthState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/login", post(login::)) @@ -370,24 +218,25 @@ async fn login( Json(payload): Json, ) -> Result where - S: AdminAuthState + Send + Sync, + S: AdminSessionState + Send + Sync, { let source = source .as_ref() .map_or("unknown", |Extension(source)| source.as_str()); let (username, password) = payload.into_parts(); let session = state - .admin_auth_service() - .login(source, username, password) + .admin_services() + .auth() + .login(LoginCommand { + username, + password, + source: source.to_owned(), + }) .await - .map_err(|error| match error { - AdminLoginError::InvalidCredentials => AdminError::invalid_admin_credentials(), - AdminLoginError::Throttled => AdminError::too_many_login_attempts(), - AdminLoginError::Unavailable => AdminError::internal("Failed to create admin session"), - })?; + .map_err(map_login_error)?; let mut response = AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminLoginData::new(session.expires_at)), + AdminEnvelope::ok(AdminLoginData::new(session.expires_at.to_rfc3339())), ) .into_response(); let cookie = format!( @@ -407,11 +256,12 @@ async fn session_status( headers: HeaderMap, ) -> Result where - S: AdminAuthState + Send + Sync, + S: AdminSessionState + Send + Sync, { let authenticated = state - .admin_auth_service() - .validate(admin_session_cookie(&headers)) + .admin_services() + .auth() + .validate_session(admin_session_cookie(&headers).as_deref()) .await .map_err(|_| AdminError::internal("Failed to validate admin session"))?; Ok(AdminResponse::new( @@ -422,10 +272,10 @@ where async fn logout(State(state): State, headers: HeaderMap) -> Result where - S: AdminAuthState + Send + Sync, + S: AdminSessionState + Send + Sync, { if let Some(session_id) = admin_session_cookie(&headers) { - let _ = state.admin_auth_service().logout(session_id).await; + let _ = state.admin_services().auth().logout(&session_id).await; } let mut response = AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(AdminLogoutData::new())) @@ -447,271 +297,10 @@ fn admin_session_cookie(headers: &HeaderMap) -> Option { }) } -const LOGIN_FAILURE_LIMIT: u32 = 5; -const LOGIN_FAILURE_WINDOW_SECONDS: u64 = 15 * 60; - -/// 认证领域需要的外层持久端口；API crate 不知道 PostgreSQL 或 runtime settings。 -#[async_trait] -pub trait AdminAuthBackend: Send + Sync { - async fn password_hash(&self, admin_user_id: &str) - -> Result, AdminServiceError>; - async fn store_password_hash( - &self, - admin_user_id: &str, - password_hash: &str, - ) -> Result<(), AdminServiceError>; - async fn admin_api_key(&self) -> Result, AdminServiceError>; - async fn load_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError>; - async fn store_admin_session( - &self, - session_id: &str, - session: &AdminBackendSession, - ) -> Result<(), AdminServiceError>; - async fn delete_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError>; - async fn login_source_is_throttled( - &self, - source: &str, - failure_limit: u32, - window_seconds: u64, - ) -> Result; - async fn record_login_failure( - &self, - source: &str, - failure_limit: u32, - window_seconds: u64, - ) -> Result; - async fn clear_login_failures(&self, source: &str) -> Result<(), AdminServiceError>; - async fn append_auth_audit(&self, event: AdminAuthAuditEvent) -> Result<(), AdminServiceError>; -} - -/// 登录与登出的最小安全审计事实。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminAuthAuditEvent { - pub admin_user_id: String, - pub action: &'static str, - pub occurred_at: DateTime, -} - -/// Redis session owner 与认证领域交换的非秘密会话事实。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminBackendSession { - pub admin_user_id: String, - pub expires_at: DateTime, -} - -/// 冻结管理端使用的默认认证领域实现。 -pub struct DefaultAdminAuthService { - default_admin_user_id: String, - session_ttl: Duration, - backend: Arc, -} - -impl DefaultAdminAuthService { - #[must_use] - pub fn new( - default_admin_user_id: String, - session_ttl_minutes: u64, - backend: Arc, - ) -> Self { - let minutes = i64::try_from(session_ttl_minutes).unwrap_or(i64::MAX); - Self { - default_admin_user_id, - session_ttl: Duration::minutes(minutes.max(1)), - backend, - } - } - - /// 空库第一次启动时创建唯一默认管理员；已有用户绝不覆盖密码。 - pub async fn ensure_default_admin(&self, password: &str) -> Result { - if self - .backend - .password_hash(&self.default_admin_user_id) - .await? - .is_some() - { - return Ok(false); - } - let hash = hash_admin_password(password)?; - self.backend - .store_password_hash(&self.default_admin_user_id, &hash) - .await?; - Ok(true) +fn map_login_error(error: LoginError) -> AdminError { + match error { + LoginError::InvalidCredentials => AdminError::invalid_admin_credentials(), + LoginError::Throttled => AdminError::too_many_login_attempts(), + LoginError::Unavailable => AdminError::internal("Failed to create admin session"), } } - -#[async_trait] -impl AdminSessionResolver for DefaultAdminAuthService { - async fn resolve_admin_user_id( - &self, - session_id: Option<&str>, - ) -> Result, AdminServiceError> { - let Some(session_id) = session_id else { - return Ok(None); - }; - Ok(self - .backend - .load_admin_session(session_id) - .await? - .filter(|session| session.expires_at > Utc::now()) - .map(|session| session.admin_user_id)) - } - - async fn verify_admin_api_key(&self, key: &str) -> Result { - if !valid_admin_api_key_shape(key) { - return Ok(false); - } - let stored = self.backend.admin_api_key().await?; - Ok(stored.as_deref().is_some_and(|stored| { - key.len() == stored.len() && bool::from(key.as_bytes().ct_eq(stored.as_bytes())) - })) - } -} - -#[async_trait] -impl AdminAuthService for DefaultAdminAuthService { - async fn login( - &self, - source: &str, - username: Option, - password: String, - ) -> Result { - let source = normalized_login_source(source); - if self - .backend - .login_source_is_throttled(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS) - .await - .map_err(|_| AdminLoginError::Unavailable)? - { - return Err(AdminLoginError::Throttled); - } - if username.as_deref().unwrap_or(&self.default_admin_user_id) != self.default_admin_user_id - { - let throttled = self - .backend - .record_login_failure(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS) - .await - .map_err(|_| AdminLoginError::Unavailable)?; - return Err(if throttled { - AdminLoginError::Throttled - } else { - AdminLoginError::InvalidCredentials - }); - } - let hash = self - .backend - .password_hash(&self.default_admin_user_id) - .await - .map_err(|_| AdminLoginError::Unavailable)? - .ok_or(AdminLoginError::InvalidCredentials)?; - if !verify_admin_password(&password, &hash).map_err(|_| AdminLoginError::Unavailable)? { - let throttled = self - .backend - .record_login_failure(source, LOGIN_FAILURE_LIMIT, LOGIN_FAILURE_WINDOW_SECONDS) - .await - .map_err(|_| AdminLoginError::Unavailable)?; - return Err(if throttled { - AdminLoginError::Throttled - } else { - AdminLoginError::InvalidCredentials - }); - } - - self.backend - .clear_login_failures(source) - .await - .map_err(|_| AdminLoginError::Unavailable)?; - let session_id = random_session_token(); - let expires_at = Utc::now() + self.session_ttl; - self.backend - .store_admin_session( - &session_id, - &AdminBackendSession { - admin_user_id: self.default_admin_user_id.clone(), - expires_at, - }, - ) - .await - .map_err(|_| AdminLoginError::Unavailable)?; - if self - .backend - .append_auth_audit(AdminAuthAuditEvent { - admin_user_id: self.default_admin_user_id.clone(), - action: "admin.login", - occurred_at: Utc::now(), - }) - .await - .is_err() - { - let _ = self.backend.delete_admin_session(&session_id).await; - return Err(AdminLoginError::Unavailable); - } - Ok(AdminLoginSession { - session_id, - expires_at: expires_at.to_rfc3339(), - }) - } - - async fn validate(&self, session_id: Option) -> Result { - let Some(session_id) = session_id else { - return Ok(false); - }; - Ok(self - .backend - .load_admin_session(&session_id) - .await? - .is_some_and(|session| session.expires_at > Utc::now())) - } - - async fn logout(&self, session_id: String) -> Result<(), AdminServiceError> { - let session = self.backend.delete_admin_session(&session_id).await?; - if let Some(session) = session { - self.backend - .append_auth_audit(AdminAuthAuditEvent { - admin_user_id: session.admin_user_id, - action: "admin.logout", - occurred_at: Utc::now(), - }) - .await?; - } - Ok(()) - } -} - -fn hash_admin_password(password: &str) -> Result { - let salt = SaltString::generate(&mut OsRng); - Argon2::default() - .hash_password(password.as_bytes(), &salt) - .map(|hash| hash.to_string()) - .map_err(|_| AdminServiceError::internal("Failed to hash admin password")) -} - -fn verify_admin_password(password: &str, encoded: &str) -> Result { - let hash = PasswordHash::new(encoded) - .map_err(|_| AdminServiceError::internal("Stored admin password hash is invalid"))?; - Ok(Argon2::default() - .verify_password(password.as_bytes(), &hash) - .is_ok()) -} - -fn random_session_token() -> String { - let mut bytes = [0_u8; 32]; - OsRng.fill_bytes(&mut bytes); - format!("session_{}", URL_SAFE_NO_PAD.encode(bytes)) -} - -fn valid_admin_api_key_shape(value: &str) -> bool { - value.len() == 70 - && value.starts_with("admin-") - && value[6..].bytes().all(|byte| byte.is_ascii_hexdigit()) -} - -fn normalized_login_source(source: &str) -> &str { - let source = source.trim(); - if source.is_empty() { "unknown" } else { source } -} diff --git a/backend/crates/gateway-api/src/admin/catalog.rs b/backend/crates/gateway-api/src/admin/catalog.rs index 3fd1a250c..202c2aaee 100644 --- a/backend/crates/gateway-api/src/admin/catalog.rs +++ b/backend/crates/gateway-api/src/admin/catalog.rs @@ -1,6 +1,5 @@ //! Provider instance 管理边界。 -use async_trait::async_trait; use axum::{ Json, Router, extract::{Query, State}, @@ -9,12 +8,20 @@ use axum::{ routing::{get, post}, }; use chrono::{DateTime, Utc}; -use gateway_core::routing::ProviderInstanceId; +use gateway_admin::model::{ + PageSize, Revision, + catalog::{ + CatalogListQuery as CatalogQuery, CreateProviderInstance, DeleteProviderInstance, + ProviderInstance, ProviderInstanceDetail, ProviderInstanceMutation, ProviderInstancePage, + SetProviderInstanceEnabled, UpdateProviderInstance, + }, +}; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use serde::{Deserialize, Serialize}; use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminRequestContext, AdminResponse, AdminServiceError, - AdminServiceErrorKind, AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, + wire::map_admin_service_error, }; const DEFAULT_PAGE_SIZE: u16 = 50; @@ -39,6 +46,18 @@ impl CatalogListQuery { pub fn page_size(&self) -> u16 { self.limit.unwrap_or(DEFAULT_PAGE_SIZE) } + + fn into_command(self) -> Result { + self.validate()?; + let page_size = + PageSize::new(self.page_size()).map_err(|_| WireValidationError::new("limit"))?; + let cursor = self + .cursor + .map(ProviderInstanceId::new) + .transpose() + .map_err(|_| WireValidationError::new("catalogCursor"))?; + Ok(CatalogQuery { cursor, page_size }) + } } #[derive(Debug, Clone, Deserialize)] @@ -51,6 +70,12 @@ impl CatalogIdQuery { pub fn validate(&self) -> Result<(), WireValidationError> { require_id(&self.id, "id") } + + fn into_id(self) -> Result { + self.validate()?; + ProviderInstanceId::new(self.id) + .map_err(|_| WireValidationError::new("catalogDetailNotFound")) + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -78,6 +103,18 @@ impl CreateProviderInstanceRequest { require_text(&self.name, "name")?; require_text(&self.base_url, "baseUrl") } + + fn into_command(self) -> Result { + self.validate()?; + Ok(CreateProviderInstance { + expected_config_revision: revision(self.expected_config_revision)?, + id: ProviderInstanceId::new(self.id).map_err(|_| WireValidationError::new("id"))?, + provider_kind: ProviderKind::new(self.provider_kind) + .map_err(|_| WireValidationError::new("providerKind"))?, + name: self.name, + base_url: self.base_url, + }) + } } #[derive(Debug, Clone, Deserialize)] @@ -96,6 +133,16 @@ impl UpdateProviderInstanceRequest { require_text(&self.name, "name")?; require_text(&self.base_url, "baseUrl") } + + fn into_command(self) -> Result { + self.validate()?; + Ok(UpdateProviderInstance { + expected_config_revision: revision(self.expected_config_revision)?, + id: ProviderInstanceId::new(self.id).map_err(|_| WireValidationError::new("id"))?, + name: self.name, + base_url: self.base_url, + }) + } } #[derive(Debug, Clone, Deserialize)] @@ -110,6 +157,15 @@ impl CatalogRevisionRequest { require_id(&self.id, "id")?; require_revision(self.expected_config_revision) } + + fn into_parts(self) -> Result<(Revision, ProviderInstanceId), WireValidationError> { + self.validate()?; + Ok(( + revision(self.expected_config_revision)?, + ProviderInstanceId::new(self.id) + .map_err(|_| WireValidationError::new("catalogMutationNotFound"))?, + )) + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -124,6 +180,20 @@ pub struct ProviderInstanceView { pub updated_at: DateTime, } +impl From for ProviderInstanceView { + fn from(instance: ProviderInstance) -> Self { + Self { + id: instance.id.to_string(), + provider_kind: instance.provider_kind.to_string(), + name: instance.name, + base_url: instance.base_url, + enabled: instance.enabled, + created_at: instance.created_at, + updated_at: instance.updated_at, + } + } +} + #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct ProviderInstanceListData { @@ -139,50 +209,28 @@ pub struct ProviderInstanceDetailData { pub item: ProviderInstanceView, } -#[async_trait] -pub trait CatalogAdminService: Send + Sync { - async fn list_provider_instances( - &self, - query: CatalogListQuery, - ) -> Result; - async fn provider_instance( - &self, - id: String, - ) -> Result; - async fn create_provider_instance( - &self, - context: &AdminRequestContext, - request: CreateProviderInstanceRequest, - ) -> Result; - async fn update_provider_instance( - &self, - context: &AdminRequestContext, - request: UpdateProviderInstanceRequest, - ) -> Result; - async fn enable_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result; - async fn disable_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result; - async fn delete_provider_instance( - &self, - context: &AdminRequestContext, - request: CatalogRevisionRequest, - ) -> Result; +impl From for ProviderInstanceListData { + fn from(page: ProviderInstancePage) -> Self { + Self { + config_revision: page.config_revision.get(), + items: page.items.into_iter().map(Into::into).collect(), + next_cursor: page.next_cursor.map(|cursor| cursor.to_string()), + } + } } -pub trait CatalogAdminState: AdminSessionState { - fn catalog_admin_service(&self) -> &dyn CatalogAdminService; +impl From for ProviderInstanceDetailData { + fn from(detail: ProviderInstanceDetail) -> Self { + Self { + config_revision: detail.config_revision.get(), + item: detail.item.into(), + } + } } pub fn router() -> Router where - S: CatalogAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route( @@ -215,105 +263,181 @@ where ) } -macro_rules! get_handler { - ($name:ident, $query:ty, $method:ident) => { - async fn $name( - _auth: AdminAuth, - State(state): State, - Query(query): Query<$query>, - ) -> Result - where - S: CatalogAdminState + Send + Sync, - { - query.validate().map_err(map_wire_error)?; - let data = state - .catalog_admin_service() - .$method(query) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) - } - }; +async fn list_provider_instances( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let result = state + .admin_services() + .catalog() + .list(query.into_command().map_err(map_wire_error)?) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(ProviderInstanceListData::from(result)), + )) } -macro_rules! get_id_handler { - ($name:ident, $method:ident) => { - async fn $name( - _auth: AdminAuth, - State(state): State, - Query(query): Query, - ) -> Result - where - S: CatalogAdminState + Send + Sync, - { - query.validate().map_err(map_wire_error)?; - let data = state - .catalog_admin_service() - .$method(query.id) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) - } - }; +async fn provider_instance( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let id = query.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .catalog() + .get(&id) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(ProviderInstanceDetailData::from(result)), + )) } -macro_rules! post_handler { - ($name:ident, $request:ty, $method:ident, $status:expr) => { - async fn $name( - auth: AdminAuth, - State(state): State, - Json(request): Json<$request>, - ) -> Result - where - S: CatalogAdminState + Send + Sync, - { - request.validate().map_err(map_wire_error)?; - let data = state - .catalog_admin_service() - .$method(auth.context(), request) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new($status, AdminEnvelope::ok(data))) - } - }; +async fn create_provider_instance( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request.into_command().map_err(map_wire_error)?; + let id = command.id.to_string(); + let result = state + .admin_services() + .catalog() + .create(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + mutation_response(StatusCode::CREATED, result, id) } -get_handler!( - list_provider_instances, - CatalogListQuery, - list_provider_instances -); -get_id_handler!(provider_instance, provider_instance); -post_handler!( - create_provider_instance, - CreateProviderInstanceRequest, - create_provider_instance, - StatusCode::CREATED -); -post_handler!( - update_provider_instance, - UpdateProviderInstanceRequest, - update_provider_instance, - StatusCode::OK -); -post_handler!( - enable_provider_instance, - CatalogRevisionRequest, - enable_provider_instance, - StatusCode::OK -); -post_handler!( - disable_provider_instance, - CatalogRevisionRequest, - disable_provider_instance, - StatusCode::OK -); -post_handler!( - delete_provider_instance, - CatalogRevisionRequest, - delete_provider_instance, - StatusCode::OK -); +async fn update_provider_instance( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request.into_command().map_err(map_wire_error)?; + let id = command.id.to_string(); + let result = state + .admin_services() + .catalog() + .update(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + mutation_response(StatusCode::OK, result, id) +} + +async fn enable_provider_instance( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + set_provider_instance_enabled(auth, state, request, true).await +} + +async fn disable_provider_instance( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + set_provider_instance_enabled(auth, state, request, false).await +} + +async fn set_provider_instance_enabled( + auth: AdminAuth, + state: S, + request: CatalogRevisionRequest, + enabled: bool, +) -> Result>, AdminError> +where + S: AdminSessionState + Send + Sync, +{ + let (expected_config_revision, id) = request.into_parts().map_err(map_wire_error)?; + let response_id = id.to_string(); + let result = state + .admin_services() + .catalog() + .set_enabled( + &auth.context().mutation_context(), + SetProviderInstanceEnabled { + expected_config_revision, + id, + enabled, + }, + ) + .await + .map_err(map_service_error)?; + mutation_response(StatusCode::OK, result, response_id) +} + +async fn delete_provider_instance( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (expected_config_revision, id) = request.into_parts().map_err(map_wire_error)?; + let response_id = id.to_string(); + let config_revision = state + .admin_services() + .catalog() + .delete( + &auth.context().mutation_context(), + DeleteProviderInstance { + expected_config_revision, + id, + }, + ) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(CatalogMutationData { + config_revision: config_revision.get(), + id: response_id, + }), + )) +} + +fn mutation_response( + status: StatusCode, + result: ProviderInstanceMutation, + fallback_id: String, +) -> Result>, AdminError> { + let id = result + .instance + .map_or(fallback_id, |instance| instance.id.to_string()); + Ok(AdminResponse::new( + status, + AdminEnvelope::ok(CatalogMutationData { + config_revision: result.config_revision.get(), + id, + }), + )) +} fn require_revision(value: u64) -> Result<(), WireValidationError> { if value == 0 || i64::try_from(value).is_err() { @@ -322,6 +446,11 @@ fn require_revision(value: u64) -> Result<(), WireValidationError> { Ok(()) } +fn revision(value: u64) -> Result { + require_revision(value)?; + Revision::new(value).map_err(|_| WireValidationError::new("expectedConfigRevision")) +} + fn require_id(value: &str, field: &'static str) -> Result<(), WireValidationError> { require_text(value, field)?; if value.starts_with("__") @@ -376,17 +505,19 @@ fn validate_limit(limit: Option) -> Result<(), WireValidationError> { } fn map_wire_error(error: WireValidationError) -> AdminError { - AdminError::bad_request(format!("Invalid field: {}", error.field())) -} - -fn map_service_error(error: AdminServiceError) -> AdminError { - match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => { - AdminError::service_unavailable("Catalog unavailable") - } - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), + match error.field() { + "catalogCursor" => AdminError::bad_request("Invalid catalog cursor"), + "catalogDetailNotFound" => AdminError::not_found("Provider instance was not found"), + "catalogMutationNotFound" => AdminError::not_found("provider instance was not found"), + field => AdminError::bad_request(format!("Invalid field: {field}")), } } + +fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError { + if error.kind() == gateway_admin::model::AdminErrorKind::Invalid + && error.to_string() == "Invalid provider instance cursor" + { + return AdminError::bad_request("Invalid catalog cursor"); + } + map_admin_service_error(error, "Catalog unavailable") +} diff --git a/backend/crates/gateway-api/src/admin/client_keys.rs b/backend/crates/gateway-api/src/admin/client_keys.rs index e45bdae5f..6fde17fba 100644 --- a/backend/crates/gateway-api/src/admin/client_keys.rs +++ b/backend/crates/gateway-api/src/admin/client_keys.rs @@ -4,9 +4,21 @@ use std::fmt; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use chrono::{DateTime, Utc}; +use gateway_admin::model::{ + Revision, + client_keys::{ + ClientKeyCursor, ClientKeyCursorValue as DomainCursorValue, ClientKeyListQuery, + ClientKeyMutation, ClientKeyPage, ClientKeyPageSize, ClientKeyRecord, ClientKeySecret, + ClientKeySort as DomainSort, ClientKeySortField as DomainSortField, CreateClientKey, + CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey, + }, +}; +use gateway_core::{ + policy::{ClientApiKeyId, RateLimits}, + routing::ProviderKind, +}; use serde::{Deserialize, Serialize}; -use async_trait::async_trait; use axum::{ Json, Router, extract::{Query, State}, @@ -16,12 +28,13 @@ use axum::{ }; use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminRequestContext, AdminResponse, AdminServiceError, - AdminServiceErrorKind, AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, + wire::map_admin_service_error, }; const MAX_CURSOR_BYTES: usize = 512; const MAX_SEARCH_BYTES: usize = 256; +const DEFAULT_PAGE_SIZE: u16 = 50; /// Client Key 列表查询。 #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] @@ -35,8 +48,8 @@ pub struct ListClientKeysQuery { } impl ListClientKeysQuery { - /// 校验游标与页大小的 wire 边界并取出字段。 - pub fn into_parts(self) -> Result { + /// 校验 wire 边界并直接构造管理用例查询。 + pub fn into_command(self) -> Result { if self .cursor .as_deref() @@ -59,11 +72,20 @@ impl ListClientKeysQuery { self.sort_by.as_deref().unwrap_or("createdAt"), self.sort_direction.as_deref().unwrap_or("desc"), )?; - Ok(ListClientKeysFields { - cursor: self.cursor, - limit: self.limit, + let cursor = self + .cursor + .as_deref() + .map(decode_client_key_cursor) + .transpose()? + .map(domain_cursor) + .transpose()?; + let page_size = ClientKeyPageSize::new(self.limit.unwrap_or(DEFAULT_PAGE_SIZE)) + .map_err(|_| WireValidationError::new("limit"))?; + Ok(ClientKeyListQuery { + cursor, + page_size, search: search.filter(|search| !search.is_empty()), - sort, + sort: domain_sort(sort), }) } } @@ -112,15 +134,6 @@ impl ClientKeySort { } } -/// 已校验的 Client Key 列表查询字段。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ListClientKeysFields { - pub cursor: Option, - pub limit: Option, - pub search: Option, - pub sort: ClientKeySort, -} - /// 创建 Client Key 请求。 #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -135,8 +148,8 @@ pub struct CreateClientKeyRequest { } impl CreateClientKeyRequest { - /// 校验请求并转换为应用层可消费字段。 - pub fn into_fields(self) -> Result { + /// 校验 wire 边界并直接构造管理用例命令。 + pub fn into_command(self) -> Result { validate_revision(self.expected_config_revision)?; validate_required_text(&self.name, "name")?; validate_optional_text(self.label.as_deref(), "label")?; @@ -144,30 +157,21 @@ impl CreateClientKeyRequest { validate_limit(self.max_concurrency, "maxConcurrency")?; validate_limit(self.requests_per_minute, "requestsPerMinute")?; validate_limit(self.tokens_per_minute, "tokensPerMinute")?; - Ok(CreateClientKeyFields { - expected_config_revision: self.expected_config_revision, + Ok(CreateClientKey { + expected_config_revision: revision(self.expected_config_revision)?, name: self.name, label: self.label, - provider_kind: self.provider_kind, - max_concurrency: self.max_concurrency, - requests_per_minute: self.requests_per_minute, - tokens_per_minute: self.tokens_per_minute, + provider_kind: ProviderKind::new(self.provider_kind) + .map_err(|_| WireValidationError::new("providerKind"))?, + limits: RateLimits { + max_concurrency: self.max_concurrency, + requests_per_minute: self.requests_per_minute, + tokens_per_minute: self.tokens_per_minute, + }, }) } } -/// 已校验的 Client Key 创建字段。 -#[derive(Debug, Clone, PartialEq)] -pub struct CreateClientKeyFields { - pub expected_config_revision: u64, - pub name: String, - pub label: Option, - pub provider_kind: String, - pub max_concurrency: u64, - pub requests_per_minute: u64, - pub tokens_per_minute: u64, -} - /// 更新 Client Key 请求。 #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -183,8 +187,8 @@ pub struct UpdateClientKeyRequest { } impl UpdateClientKeyRequest { - /// 校验请求并转换为应用层可消费字段。 - pub fn into_fields(self) -> Result { + /// 校验 wire 边界并直接构造管理用例命令。 + pub fn into_command(self) -> Result { validate_required_text(&self.id, "id")?; validate_revision(self.expected_config_revision)?; validate_required_text(&self.name, "name")?; @@ -193,32 +197,22 @@ impl UpdateClientKeyRequest { validate_limit(self.max_concurrency, "maxConcurrency")?; validate_limit(self.requests_per_minute, "requestsPerMinute")?; validate_limit(self.tokens_per_minute, "tokensPerMinute")?; - Ok(UpdateClientKeyFields { - id: self.id, - expected_config_revision: self.expected_config_revision, + Ok(UpdateClientKey { + expected_config_revision: revision(self.expected_config_revision)?, + id: client_key_id(self.id, "clientKeyMutationNotFound")?, name: self.name, label: self.label, - provider_kind: self.provider_kind, - max_concurrency: self.max_concurrency, - requests_per_minute: self.requests_per_minute, - tokens_per_minute: self.tokens_per_minute, + provider_kind: ProviderKind::new(self.provider_kind) + .map_err(|_| WireValidationError::new("providerKind"))?, + limits: RateLimits { + max_concurrency: self.max_concurrency, + requests_per_minute: self.requests_per_minute, + tokens_per_minute: self.tokens_per_minute, + }, }) } } -/// 已校验的 Client Key 更新字段。 -#[derive(Debug, Clone, PartialEq)] -pub struct UpdateClientKeyFields { - pub id: String, - pub expected_config_revision: u64, - pub name: String, - pub label: Option, - pub provider_kind: String, - pub max_concurrency: u64, - pub requests_per_minute: u64, - pub tokens_per_minute: u64, -} - /// 只携带 ID 与配置 revision 的 Client Key mutation 请求。 #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -239,6 +233,10 @@ impl ClientKeyIdQuery { validate_required_text(&self.id, "id")?; Ok(self.id) } + + fn into_domain_id(self) -> Result { + client_key_id(self.into_id()?, "clientKeyRevealNotFound") + } } impl ClientKeyRevisionRequest { @@ -248,6 +246,14 @@ impl ClientKeyRevisionRequest { validate_revision(self.expected_config_revision)?; Ok((self.id, self.expected_config_revision)) } + + fn into_command_parts(self) -> Result<(ClientApiKeyId, Revision), WireValidationError> { + let (id, expected_config_revision) = self.into_parts()?; + Ok(( + client_key_id(id, "clientKeyMutationNotFound")?, + revision(expected_config_revision)?, + )) + } } /// 不含完整 Key 的管理端安全视图。 @@ -268,44 +274,25 @@ pub struct ClientKeyView { last_used_at: Option>, } -impl ClientKeyView { - /// 从应用层已经读取的安全字段构造视图。 - #[must_use] - pub fn new(fields: ClientKeyViewFields) -> Self { +impl From for ClientKeyView { + fn from(record: ClientKeyRecord) -> Self { Self { - id: fields.id, - name: fields.name, - label: fields.label, - provider_kind: fields.provider_kind, - prefix: fields.prefix, - enabled: fields.enabled, - max_concurrency: fields.max_concurrency, - requests_per_minute: fields.requests_per_minute, - tokens_per_minute: fields.tokens_per_minute, - created_at: fields.created_at, - updated_at: fields.updated_at, - last_used_at: fields.last_used_at, + id: record.id.to_string(), + name: record.name, + label: record.label, + provider_kind: record.provider_kind.to_string(), + prefix: record.prefix, + enabled: record.enabled, + max_concurrency: record.limits.max_concurrency, + requests_per_minute: record.limits.requests_per_minute, + tokens_per_minute: record.limits.tokens_per_minute, + created_at: record.created_at, + updated_at: record.updated_at, + last_used_at: record.last_used_at, } } } -/// Client Key 安全视图的应用层输入字段。 -#[derive(Debug, Clone, PartialEq)] -pub struct ClientKeyViewFields { - pub id: String, - pub name: String, - pub label: Option, - pub provider_kind: String, - pub prefix: String, - pub enabled: bool, - pub max_concurrency: u64, - pub requests_per_minute: u64, - pub tokens_per_minute: u64, - pub created_at: DateTime, - pub updated_at: DateTime, - pub last_used_at: Option>, -} - /// Client Key 列表响应数据。 #[derive(Debug, Clone, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] @@ -334,6 +321,26 @@ impl ClientKeyListData { } } +impl TryFrom for ClientKeyListData { + type Error = WireValidationError; + + fn try_from(page: ClientKeyPage) -> Result { + let next_cursor = page + .next_cursor + .map(wire_cursor) + .transpose()? + .as_ref() + .map(encode_client_key_cursor) + .transpose()?; + Ok(Self::new( + page.config_revision.get(), + page.items.into_iter().map(Into::into).collect(), + next_cursor, + page.total, + )) + } +} + /// Client Key 创建响应；完整值只允许出现在本次序列化结果中。 #[derive(Serialize)] #[serde(rename_all = "camelCase")] @@ -357,6 +364,17 @@ impl CreatedClientKeyData { } } +impl From for CreatedClientKeyData { + fn from(created: CreatedClientKey) -> Self { + Self::new( + created.config_revision.get(), + created.secret.record.id.to_string(), + created.secret.record.prefix.clone(), + created.secret.expose_for_response().to_owned(), + ) + } +} + impl fmt::Debug for CreatedClientKeyData { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -384,6 +402,15 @@ impl RevealedClientKeyData { } } +impl From for RevealedClientKeyData { + fn from(secret: ClientKeySecret) -> Self { + Self::new( + secret.record.id.to_string(), + secret.expose_for_response().to_owned(), + ) + } +} + impl fmt::Debug for RevealedClientKeyData { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -413,6 +440,12 @@ impl MutatedClientKeyData { } } +impl From for MutatedClientKeyData { + fn from(mutation: ClientKeyMutation) -> Self { + Self::new(mutation.config_revision.get(), mutation.id.to_string()) + } +} + /// 解码后的 Client Key 游标字段。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ClientKeyCursorData { @@ -497,6 +530,66 @@ fn validate_client_key_cursor(cursor: &ClientKeyCursorData) -> Result<(), WireVa } } +const fn domain_sort(sort: ClientKeySort) -> DomainSort { + DomainSort { + field: match sort.field { + ClientKeySortField::Name => DomainSortField::Name, + ClientKeySortField::Enabled => DomainSortField::Enabled, + ClientKeySortField::CreatedAt => DomainSortField::CreatedAt, + ClientKeySortField::LastUsedAt => DomainSortField::LastUsedAt, + }, + direction: match sort.direction { + ClientKeySortDirection::Asc => SortDirection::Asc, + ClientKeySortDirection::Desc => SortDirection::Desc, + }, + } +} + +const fn wire_sort(sort: DomainSort) -> ClientKeySort { + ClientKeySort { + field: match sort.field { + DomainSortField::Name => ClientKeySortField::Name, + DomainSortField::Enabled => ClientKeySortField::Enabled, + DomainSortField::CreatedAt => ClientKeySortField::CreatedAt, + DomainSortField::LastUsedAt => ClientKeySortField::LastUsedAt, + }, + direction: match sort.direction { + SortDirection::Asc => ClientKeySortDirection::Asc, + SortDirection::Desc => ClientKeySortDirection::Desc, + }, + } +} + +fn domain_cursor(cursor: ClientKeyCursorData) -> Result { + let value = match cursor.value { + ClientKeyCursorValue::Name(value) => DomainCursorValue::Name(value), + ClientKeyCursorValue::Enabled(value) => DomainCursorValue::Enabled(value), + ClientKeyCursorValue::CreatedAt(value) => DomainCursorValue::CreatedAt(value), + ClientKeyCursorValue::LastUsedAt(value) => DomainCursorValue::LastUsedAt(value), + }; + Ok(ClientKeyCursor { + sort: domain_sort(cursor.sort), + value, + id: client_key_id(cursor.id, "cursor")?, + }) +} + +fn wire_cursor(cursor: ClientKeyCursor) -> Result { + let value = match cursor.value { + DomainCursorValue::Name(value) => ClientKeyCursorValue::Name(value), + DomainCursorValue::Enabled(value) => ClientKeyCursorValue::Enabled(value), + DomainCursorValue::CreatedAt(value) => ClientKeyCursorValue::CreatedAt(value), + DomainCursorValue::LastUsedAt(value) => ClientKeyCursorValue::LastUsedAt(value), + }; + let cursor = ClientKeyCursorData { + sort: wire_sort(cursor.sort), + value, + id: cursor.id.to_string(), + }; + validate_client_key_cursor(&cursor)?; + Ok(cursor) +} + fn validate_revision(value: u64) -> Result<(), WireValidationError> { if value == 0 { return Err(WireValidationError::new("expectedConfigRevision")); @@ -504,6 +597,18 @@ fn validate_revision(value: u64) -> Result<(), WireValidationError> { Ok(()) } +fn revision(value: u64) -> Result { + validate_revision(value)?; + Revision::new(value).map_err(|_| WireValidationError::new("expectedConfigRevision")) +} + +fn client_key_id( + value: String, + field: &'static str, +) -> Result { + ClientApiKeyId::new(value).map_err(|_| WireValidationError::new(field)) +} + fn validate_limit(value: u64, field: &'static str) -> Result<(), WireValidationError> { if i64::try_from(value).is_err() { return Err(WireValidationError::new(field)); @@ -554,59 +659,10 @@ fn contains_client_key_material(value: &str) -> bool { }) } -/// Client API Key 管理应用端口。 -#[async_trait] -pub trait ClientKeyAdminService: Send + Sync { - async fn list( - &self, - query: ListClientKeysFields, - ) -> Result; - - async fn create( - &self, - context: &AdminRequestContext, - fields: CreateClientKeyFields, - ) -> Result; - - async fn reveal(&self, id: String) -> Result; - - async fn update( - &self, - context: &AdminRequestContext, - fields: UpdateClientKeyFields, - ) -> Result; - - async fn disable( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result; - - async fn enable( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result; - - async fn delete( - &self, - context: &AdminRequestContext, - id: String, - expected_config_revision: u64, - ) -> Result; -} - -/// Client API Key HTTP module 所需最小 state。 -pub trait ClientKeyAdminState: AdminSessionState { - fn client_key_admin_service(&self) -> &dyn ClientKeyAdminService; -} - /// 构造固定 GET/POST 且 ID 仅位于 query/body 的 Client API Key 路由。 pub fn router() -> Router where - S: ClientKeyAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route( @@ -638,14 +694,16 @@ async fn list_client_keys( Query(query): Query, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let fields = query.into_parts().map_err(map_wire_error)?; - let data = state - .client_key_admin_service() - .list(fields) + let result = state + .admin_services() + .client_keys() + .list(query.into_command().map_err(map_wire_error)?) .await .map_err(map_service_error)?; + let data = ClientKeyListData::try_from(result) + .map_err(|_| AdminError::internal("Failed to encode client key cursor"))?; Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) } @@ -655,17 +713,18 @@ async fn create_client_key( Json(payload): Json, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let fields = payload.into_fields().map_err(map_wire_error)?; - let data = state - .client_key_admin_service() - .create(auth.context(), fields) + let command = payload.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .client_keys() + .create(&auth.context().mutation_context(), command) .await .map_err(map_service_error)?; Ok(AdminResponse::new( StatusCode::CREATED, - AdminEnvelope::ok(data), + AdminEnvelope::ok(CreatedClientKeyData::from(result)), )) } @@ -675,15 +734,20 @@ async fn reveal_client_key( Query(query): Query, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let id = query.into_id().map_err(map_wire_error)?; - let data = state - .client_key_admin_service() - .reveal(id) + let id = query.into_domain_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .client_keys() + .reveal(&id) .await .map_err(map_service_error)?; - let mut response = AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)).into_response(); + let mut response = AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RevealedClientKeyData::from(result)), + ) + .into_response(); response .headers_mut() .insert(CACHE_CONTROL, HeaderValue::from_static("no-store")); @@ -696,13 +760,14 @@ async fn update_client_key( Json(payload): Json, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let fields = payload.into_fields().map_err(map_wire_error)?; + let command = payload.into_command().map_err(map_wire_error)?; mutation_response( state - .client_key_admin_service() - .update(auth.context(), fields) + .admin_services() + .client_keys() + .update(&auth.context().mutation_context(), command) .await, ) } @@ -713,13 +778,21 @@ async fn disable_client_key( Json(payload): Json, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let (id, revision) = payload.into_parts().map_err(map_wire_error)?; + let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?; mutation_response( state - .client_key_admin_service() - .disable(auth.context(), id, revision) + .admin_services() + .client_keys() + .set_enabled( + &auth.context().mutation_context(), + SetClientKeyEnabled { + expected_config_revision, + id, + enabled: false, + }, + ) .await, ) } @@ -730,13 +803,21 @@ async fn enable_client_key( Json(payload): Json, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let (id, revision) = payload.into_parts().map_err(map_wire_error)?; + let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?; mutation_response( state - .client_key_admin_service() - .enable(auth.context(), id, revision) + .admin_services() + .client_keys() + .set_enabled( + &auth.context().mutation_context(), + SetClientKeyEnabled { + expected_config_revision, + id, + enabled: true, + }, + ) .await, ) } @@ -747,27 +828,39 @@ async fn delete_client_key( Json(payload): Json, ) -> Result where - S: ClientKeyAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let (id, revision) = payload.into_parts().map_err(map_wire_error)?; + let (id, expected_config_revision) = payload.into_command_parts().map_err(map_wire_error)?; mutation_response( state - .client_key_admin_service() - .delete(auth.context(), id, revision) + .admin_services() + .client_keys() + .delete( + &auth.context().mutation_context(), + DeleteClientKey { + expected_config_revision, + id, + }, + ) .await, ) } fn mutation_response( - result: Result, + result: Result, ) -> Result>, AdminError> { let data = result.map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(MutatedClientKeyData::from(data)), + )) } fn map_wire_error(error: WireValidationError) -> AdminError { match error.field() { "cursor" => AdminError::bad_request("Invalid client key cursor"), + "clientKeyRevealNotFound" => AdminError::not_found("Client API key was not found"), + "clientKeyMutationNotFound" => AdminError::not_found("client API key was not found"), "expectedConfigRevision" => { AdminError::bad_request("expectedConfigRevision must be positive") } @@ -775,14 +868,6 @@ fn map_wire_error(error: WireValidationError) -> AdminError { } } -fn map_service_error(error: AdminServiceError) -> AdminError { - match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => { - AdminError::service_unavailable("Configuration repository unavailable") - } - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), - } +fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError { + map_admin_service_error(error, "Configuration repository unavailable") } diff --git a/backend/crates/gateway-api/src/admin/mod.rs b/backend/crates/gateway-api/src/admin/mod.rs index fa94cefde..011b11cd4 100644 --- a/backend/crates/gateway-api/src/admin/mod.rs +++ b/backend/crates/gateway-api/src/admin/mod.rs @@ -18,45 +18,16 @@ pub mod system; pub mod wire; pub mod xai; -pub use auth::{ - AdminAuth, AdminPrincipal, AdminRequestContext, AdminServiceError, AdminServiceErrorKind, - AdminSessionResolver, AdminSessionState, -}; -pub use settings::{ - AdminApiKeyStatus, DeletedAdminApiKey, ProviderModelMappings, RegeneratedAdminApiKey, - RuntimeSettingsView, UpdateRuntimeSettingsRequest, -}; +pub use auth::{AdminAuth, AdminSessionState}; pub use wire::{ ADMIN_OK_CODE, ADMIN_OK_MESSAGE, AdminEnvelope, AdminError, AdminErrorBody, AdminErrorCode, AdminPageData, AdminResponse, PageMeta, WireValidationError, }; -use accounts::AccountAdminState; -use auth::AdminAuthState; -use catalog::CatalogAdminState; -use client_keys::ClientKeyAdminState; -use observability::ObservabilityAdminState; -use openai::CodexAdminState; -use settings::AdminSettingsState; -use system::SystemAdminState; -use xai::XaiAdminState; - /// 构造完整且固定的 `/api/admin` 路由。 pub fn router() -> Router where - S: AccountAdminState - + AdminAuthState - + CatalogAdminState - + ClientKeyAdminState - + CodexAdminState - + ObservabilityAdminState - + AdminSettingsState - + SystemAdminState - + XaiAdminState - + Clone - + Send - + Sync - + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .merge(accounts::router::()) diff --git a/backend/crates/gateway-api/src/admin/observability.rs b/backend/crates/gateway-api/src/admin/observability.rs index c393b8519..1de00db6b 100644 --- a/backend/crates/gateway-api/src/admin/observability.rs +++ b/backend/crates/gateway-api/src/admin/observability.rs @@ -1,6 +1,7 @@ -//! Dashboard、用量与诊断查询的 wire、service port 和固定路由。 +//! Dashboard、用量与诊断查询的 wire 映射和固定路由。 + +use std::num::NonZeroU32; -use async_trait::async_trait; use axum::{ Router, extract::{Query, State}, @@ -8,12 +9,14 @@ use axum::{ response::IntoResponse, routing::get, }; -use chrono::{DateTime, Utc}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, Duration, Timelike as _, Utc}; +use gateway_admin::model::{PageSize as DomainPageSize, observability as domain}; use serde::{Deserialize, Serialize}; use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminServiceError, AdminServiceErrorKind, - AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, + wire::map_admin_service_error, }; /// 观测列表默认页大小。 @@ -281,6 +284,194 @@ fn trimmed(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } +fn domain_trend_kind(kind: TrendKind) -> domain::TrendKind { + match kind { + TrendKind::Usage => domain::TrendKind::Usage, + TrendKind::Latency => domain::TrendKind::Latency, + TrendKind::Errors => domain::TrendKind::Errors, + } +} + +fn domain_diagnostic_dimension(dimension: DiagnosticDimension) -> domain::DiagnosticDimension { + match dimension { + DiagnosticDimension::Model => domain::DiagnosticDimension::Model, + DiagnosticDimension::Account => domain::DiagnosticDimension::Account, + DiagnosticDimension::ApiKey => domain::DiagnosticDimension::ApiKey, + DiagnosticDimension::Provider => domain::DiagnosticDimension::Provider, + DiagnosticDimension::Transport => domain::DiagnosticDimension::Transport, + DiagnosticDimension::Failure => domain::DiagnosticDimension::Failure, + DiagnosticDimension::Status => domain::DiagnosticDimension::Status, + } +} + +fn request_outcome( + value: Option, +) -> Result, WireValidationError> { + non_empty(value) + .map(|value| { + domain::RequestOutcome::new(value).map_err(|_| WireValidationError::new("outcome")) + }) + .transpose() +} + +fn usage_range( + start: Option<&str>, + end: Option<&str>, +) -> Result { + let end = parse_datetime(end)?.unwrap_or_else(Utc::now); + let start = parse_datetime(start)?.unwrap_or(end - Duration::days(7)); + domain::TimeRange::new(start, end).map_err(|_| WireValidationError::new("timeRange")) +} + +fn dashboard_range( + start: Option<&str>, + end: Option<&str>, +) -> Result { + let end = parse_datetime(end)?.unwrap_or_else(Utc::now); + let start = parse_datetime(start)?.unwrap_or_else(|| china_day_start(end) - Duration::days(1)); + domain::TimeRange::new(start, end).map_err(|_| WireValidationError::new("timeRange")) +} + +fn dashboard_today_range( + start: Option<&str>, + end: Option<&str>, +) -> Result { + let end = parse_datetime(end)?.unwrap_or_else(Utc::now); + let start = parse_datetime(start)?.unwrap_or_else(|| china_day_start(end)); + domain::TimeRange::new(start, end).map_err(|_| WireValidationError::new("timeRange")) +} + +fn usage_filter(query: &UsageQuery) -> Result { + let outcome = query.outcome.clone().or_else(|| { + non_empty(query.kind.clone()).filter(|value| { + matches!( + value.as_str(), + "running" | "succeeded" | "failed" | "cancelled" | "incomplete" + ) + }) + }); + Ok(domain::UsageFilter { + client_api_key_ref: non_empty(query.client_api_key_id.clone()), + request_id: non_empty(query.request_id.clone()), + provider_account_ref: non_empty(query.account_id.clone()), + operation: non_empty(query.route.clone()), + provider_kind: non_empty(query.provider.clone()), + model: non_empty(query.model.clone()), + outcome: request_outcome(outcome)?, + status_code: parse_status(query.status_code)?, + transport: non_empty(query.transport.clone()), + attempt_index: parse_attempt_index(query.attempt_index)?, + response_id: non_empty(query.response_id.clone()), + upstream_request_id: non_empty(query.upstream_request_id.clone()), + search: non_empty(query.search.clone()), + }) +} + +fn usage_command( + query: &UsageQuery, +) -> Result<(domain::UsageQuery, u32, u16), WireValidationError> { + let (page, page_size) = query.validate_page()?; + query.validate_cursor()?; + let page_number = NonZeroU32::new(page) + .map(domain::PageNumber::new) + .ok_or_else(|| WireValidationError::new("page"))?; + let page_size_value = + DomainPageSize::new(page_size).map_err(|_| WireValidationError::new("pageSize"))?; + Ok(( + domain::UsageQuery { + range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?, + filter: usage_filter(query)?, + cursor: decode_observability_cursor(query.cursor.as_deref())?, + page: page_number, + page_size: page_size_value, + }, + page, + page_size, + )) +} + +fn ops_command(query: &OpsQuery) -> Result<(domain::OpsErrorQuery, u32, u16), WireValidationError> { + let (page, page_size) = query.validate_page()?; + query.validate_cursor()?; + let page_number = NonZeroU32::new(page) + .map(domain::PageNumber::new) + .ok_or_else(|| WireValidationError::new("page"))?; + let page_size_value = + DomainPageSize::new(page_size).map_err(|_| WireValidationError::new("pageSize"))?; + let status_code = parse_status( + query + .upstream_status_code + .or(query.client_status_code) + .or(query.status_code), + )?; + Ok(( + domain::OpsErrorQuery { + range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?, + filter: domain::OpsErrorFilter { + client_api_key_ref: non_empty(query.client_api_key_id.clone()), + request_id: non_empty(query.request_id.clone()), + provider_kind: non_empty(query.provider.clone()), + provider_account_ref: non_empty(query.account_id.clone()), + operation: non_empty(query.route.clone()).or_else(|| non_empty(query.kind.clone())), + model: non_empty(query.model.clone()), + transport: non_empty(query.transport.clone()), + attempt_index: parse_attempt_index(query.attempt_index)?, + response_id: non_empty(query.response_id.clone()), + upstream_request_id: non_empty(query.upstream_request_id.clone()), + failure_kind: non_empty(query.failure_class.clone()), + status_code, + search: non_empty(query.search.clone()), + }, + cursor: decode_observability_cursor(query.cursor.as_deref())?, + page: page_number, + page_size: page_size_value, + }, + page, + page_size, + )) +} + +fn decode_observability_cursor( + value: Option<&str>, +) -> Result, WireValidationError> { + value + .map(|encoded| { + if encoded.is_empty() || encoded.len() > MAX_CURSOR_BYTES { + return Err(WireValidationError::new("cursor")); + } + let bytes = URL_SAFE_NO_PAD + .decode(encoded) + .map_err(|_| WireValidationError::new("cursor"))?; + let wire: CursorWire = + serde_json::from_slice(&bytes).map_err(|_| WireValidationError::new("cursor"))?; + if wire.stable_id.trim().is_empty() { + return Err(WireValidationError::new("cursor")); + } + Ok(domain::ObservabilityCursor { + observed_at: wire.observed_at, + stable_id: wire.stable_id, + }) + }) + .transpose() +} + +fn encode_observability_cursor( + cursor: &domain::ObservabilityCursor, +) -> Result { + let bytes = serde_json::to_vec(&CursorWire { + observed_at: cursor.observed_at, + stable_id: cursor.stable_id.clone(), + }) + .map_err(|_| WireValidationError::new("cursor"))?; + Ok(URL_SAFE_NO_PAD.encode(bytes)) +} + +fn non_empty(value: Option) -> Option { + value + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + /// 观测列表游标的稳定 wire 形状。 #[derive(Clone, Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -719,7 +910,7 @@ pub struct DashboardWireTargetView { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct DashboardDesktopReleaseView { - pub status: String, + pub status: DashboardDesktopReleaseStatusView, #[serde(skip_serializing_if = "Option::is_none")] pub checked_at: Option>, #[serde(skip_serializing_if = "Option::is_none")] @@ -742,6 +933,27 @@ pub struct DashboardDesktopReleaseView { pub error: Option, } +/// Dashboard 发布检查的稳定 wire 状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum DashboardDesktopReleaseStatusView { + Unchecked, + Aligned, + ReviewRequired, + CheckFailed, +} + +impl From for DashboardDesktopReleaseStatusView { + fn from(status: domain::DesktopReleaseStatus) -> Self { + match status { + domain::DesktopReleaseStatus::Unchecked => Self::Unchecked, + domain::DesktopReleaseStatus::Current => Self::Aligned, + domain::DesktopReleaseStatus::UpdateAvailable => Self::ReviewRequired, + domain::DesktopReleaseStatus::Failed => Self::CheckFailed, + } + } +} + /// Dashboard 汇总响应数据。 #[derive(Serialize)] #[serde(rename_all = "camelCase")] @@ -955,46 +1167,9 @@ pub struct OpsErrorMetadataView { pub account_label: Option, } -#[async_trait] -pub trait ObservabilityAdminService: Send + Sync { - async fn dashboard_summary( - &self, - query: DashboardQuery, - ) -> Result; - async fn dashboard_trend(&self, query: DashboardQuery) -> Result; - async fn usage_records( - &self, - query: UsageQuery, - ) -> Result, AdminServiceError>; - async fn usage_record_detail( - &self, - query: DetailQuery, - ) -> Result; - async fn usage_records_summary( - &self, - query: UsageQuery, - ) -> Result; - async fn usage_insights_overview( - &self, - query: UsageQuery, - ) -> Result; - async fn usage_insights_diagnostics( - &self, - query: DiagnosticsQuery, - ) -> Result; - async fn ops_errors( - &self, - query: OpsQuery, - ) -> Result, AdminServiceError>; -} - -pub trait ObservabilityAdminState: AdminSessionState { - fn observability_admin_service(&self) -> &dyn ObservabilityAdminService; -} - pub fn router() -> Router where - S: ObservabilityAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/dashboard/summary", get(dashboard_summary::)) @@ -1019,73 +1194,1189 @@ where .route("/api/admin/ops/errors", get(ops_errors::)) } -macro_rules! query_handler { - ($name:ident, $query:ty, $data:ty, $method:ident) => { - async fn $name( - _auth: AdminAuth, - State(state): State, - Query(query): Query<$query>, - ) -> Result - where - S: ObservabilityAdminState + Send + Sync, - { - let data: $data = state - .observability_admin_service() - .$method(query) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) - } - }; +async fn dashboard_summary( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let kind = query.trend_kind().map_err(map_wire_error)?; + let range = dashboard_range(query.start_time.as_deref(), query.end_time.as_deref()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .dashboard_summary(range, domain_trend_kind(kind)) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(dashboard_view(result, kind)), + )) } -query_handler!( - dashboard_summary, - DashboardQuery, - DashboardDataView, - dashboard_summary -); -query_handler!(dashboard_trend, DashboardQuery, TrendData, dashboard_trend); -query_handler!( - usage_records, - UsageQuery, - PageData, - usage_records -); -query_handler!( - usage_record_detail, - DetailQuery, - UsageRecordDetailView, - usage_record_detail -); -query_handler!( - usage_records_summary, - UsageQuery, - UsageSummaryView, - usage_records_summary -); -query_handler!( - usage_insights_overview, - UsageQuery, - UsageInsightsOverviewView, - usage_insights_overview -); -query_handler!( - usage_insights_diagnostics, - DiagnosticsQuery, - DiagnosticsView, - usage_insights_diagnostics -); -query_handler!(ops_errors, OpsQuery, PageData, ops_errors); +async fn dashboard_trend( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let kind = query.trend_kind().map_err(map_wire_error)?; + let range = dashboard_today_range(query.start_time.as_deref(), query.end_time.as_deref()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .dashboard_trend(range, domain_trend_kind(kind)) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(trend_view(result, kind)), + )) +} -fn map_service_error(error: AdminServiceError) -> AdminError { - match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => { - AdminError::service_unavailable("Observability repository unavailable") +async fn usage_records( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (command, page, page_size) = usage_command(&query).map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .usage_records(command) + .await + .map_err(map_service_error)?; + let data = usage_page_view(result, page, page_size) + .map_err(|_| AdminError::internal("Failed to encode observability cursor"))?; + Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) +} + +async fn usage_record_detail( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + query.validate().map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .usage_record_detail(query.id.trim()) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(usage_detail_view(result)), + )) +} + +async fn usage_records_summary( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref()) + .map_err(map_wire_error)?; + let filter = usage_filter(&query).map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .usage_summary(range, filter) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(usage_summary_view(result)), + )) +} + +async fn usage_insights_overview( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref()) + .map_err(map_wire_error)?; + let filter = usage_filter(&query).map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .usage_insights(range, filter) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(usage_insights_view(result)), + )) +} + +async fn usage_insights_diagnostics( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let dimension = query.dimension().map_err(map_wire_error)?; + let range = usage_range(query.start_time.as_deref(), query.end_time.as_deref()) + .map_err(map_wire_error)?; + let filter = domain::UsageFilter { + provider_kind: non_empty(query.provider), + model: non_empty(query.model), + status_code: parse_status(query.status_code).map_err(map_wire_error)?, + search: non_empty(query.search), + ..domain::UsageFilter::default() + }; + let result = state + .admin_services() + .observability() + .diagnostics(range, filter, domain_diagnostic_dimension(dimension)) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(diagnostics_view(result, dimension)), + )) +} + +async fn ops_errors( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (command, page, page_size) = ops_command(&query).map_err(map_wire_error)?; + let result = state + .admin_services() + .observability() + .ops_errors(command) + .await + .map_err(map_service_error)?; + let data = ops_page_view(result, page, page_size) + .map_err(|_| AdminError::internal("Failed to encode observability cursor"))?; + Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) +} + +fn format_number(value: u64) -> String { + let text = value.to_string(); + let mut output = String::with_capacity(text.len() + text.len() / 3); + for (index, character) in text.chars().rev().enumerate() { + if index > 0 && index % 3 == 0 { + output.push(','); } - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), + output.push(character); + } + output.chars().rev().collect() +} + +fn format_compact_number(value: u64) -> String { + if value < 1_000 { + return format_number(value); + } + for (suffix, threshold) in [ + ("P", 1_000_000_000_000_000_u64), + ("T", 1_000_000_000_000_u64), + ("B", 1_000_000_000_u64), + ("M", 1_000_000_u64), + ("K", 1_000_u64), + ] { + if value >= threshold { + let scaled = value as f64 / threshold as f64; + return format!("{scaled:.1}{suffix}").replace(".0", ""); + } + } + format_number(value) +} + +fn display_duration(value: Option) -> String { + let Some(value) = value.and_then(|value| i64::try_from(value).ok()) else { + return "—".to_owned(); + }; + if value < 1_000 { + format!("{value} ms") + } else if value < 60_000 { + format!("{:.2} s", value as f64 / 1_000.0) + } else if value < 3_600_000 { + format!("{:.1} min", value as f64 / 60_000.0) + } else { + format!("{:.1} h", value as f64 / 3_600_000.0) } } + +fn display_rate(value: f64) -> String { + if value.is_finite() { + format!("{:.1}%", value * 100.0) + } else { + "—".to_owned() + } +} + +fn china_day_start(value: DateTime) -> DateTime { + const CHINA_OFFSET_SECONDS: i64 = 8 * 60 * 60; + let elapsed = (value.timestamp() + CHINA_OFFSET_SECONDS).rem_euclid(24 * 60 * 60); + value - Duration::seconds(elapsed) - Duration::nanoseconds(i64::from(value.nanosecond())) +} + +fn china_datetime(value: &DateTime) -> String { + (*value + Duration::hours(8)) + .format("%Y-%m-%d %H:%M:%S") + .to_string() +} + +fn china_label(value: DateTime, format: &str) -> String { + (value + Duration::hours(8)).format(format).to_string() +} + +fn outcome_name(outcome: &domain::RequestOutcome) -> &str { + outcome.as_str() +} + +fn request_metrics_view(metrics: &domain::RequestMetrics) -> RequestMetricsView { + RequestMetricsView { + request_count: metrics.request_count, + success_count: metrics.success_count, + failure_count: metrics.failure_count, + cancelled_count: metrics.cancelled_count, + incomplete_count: metrics.incomplete_count, + caller_error_count: metrics.caller_error_count, + input_tokens: metrics.input_tokens, + output_tokens: metrics.output_tokens, + cached_tokens: metrics.cached_tokens, + cache_write_tokens: metrics.cache_write_tokens, + reasoning_tokens: metrics.reasoning_tokens, + total_tokens: metrics.total_tokens, + } +} + +fn cost_coverage_view(coverage: &domain::CostCoverage) -> CostCoverageView { + CostCoverageView { + known: coverage.known_count(), + partial: coverage.partial_count, + unknown: coverage.unavailable_count, + not_billable: coverage.not_billable_count, + } +} + +fn cost_views(costs: &[domain::CurrencyCost]) -> Vec { + costs + .iter() + .map(|cost| CostView { + currency: cost.currency.clone(), + estimated_amount: cost.amount.as_str().to_owned(), + }) + .collect() +} + +fn attempt_metrics_view(metrics: &domain::AttemptMetrics) -> AttemptMetricsView { + AttemptMetricsView { + attempt_count: metrics.attempt_count, + success_count: metrics.success_count, + failure_count: metrics.failure_count, + cancelled_count: metrics.cancelled_count, + incomplete_count: metrics.incomplete_count, + rate_limited_count: metrics.rate_limited_count, + auth_failure_count: metrics.auth_failure_count, + provider5xx_count: metrics.provider_5xx_count, + cost_coverage: cost_coverage_view(&metrics.cost_coverage), + costs: cost_views(&metrics.costs), + } +} + +fn token_details(record: &domain::UsageRecord) -> TokenDetailsView { + TokenDetailsView { + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cached_tokens: record.cached_tokens, + cache_write_tokens: record.cache_write_tokens, + reasoning_tokens: record.reasoning_tokens, + total_tokens: record.total_tokens, + input_tokens_display: record + .input_tokens + .map_or_else(|| "-".to_owned(), format_number), + output_tokens_display: record + .output_tokens + .map_or_else(|| "-".to_owned(), format_number), + cached_tokens_display: record + .cached_tokens + .map_or_else(|| "-".to_owned(), format_compact_number), + cache_write_tokens_display: record + .cache_write_tokens + .map_or_else(|| "-".to_owned(), format_compact_number), + reasoning_tokens_display: record + .reasoning_tokens + .map_or_else(|| "-".to_owned(), format_number), + total_tokens_display: record + .total_tokens + .map_or_else(|| "-".to_owned(), format_number), + } +} + +fn format_decimal_currency(amount: &str, currency: &str) -> String { + if currency == "USD" { + let value = amount.parse::().map_or(0.0, |value| value); + let precision = if value != 0.0 && value.abs() < 0.01 { + 4 + } else { + 2 + }; + format!("${value:.precision$}") + } else { + format!("{currency} {amount}") + } +} + +fn format_money(cost: &domain::CurrencyCost) -> String { + format_decimal_currency(cost.amount.as_str(), &cost.currency) +} + +fn format_token_price(cost: &domain::CurrencyCost) -> String { + if cost.currency != "USD" { + return format!("{} {} / 1M Token", cost.currency, cost.amount.as_str()); + } + let value = cost + .amount + .as_str() + .parse::() + .map_or(0.0, |value| value); + format!("${value:.4} / 1M Token") +} + +fn format_service_tier(service_tier: Option<&str>) -> String { + match service_tier { + Some("priority" | "fast") => "Fast".to_owned(), + Some("flex") => "Flex".to_owned(), + Some("default") | None => "Default".to_owned(), + Some(other) => other.to_owned(), + } +} + +fn billing_view(billing: Option<&domain::UsageBilling>) -> Option { + match billing? { + domain::UsageBilling::Total { total, .. } => Some(BillingView { + input_amount_display: "—".to_owned(), + output_amount_display: "—".to_owned(), + cache_read_amount_display: "—".to_owned(), + cache_write_amount_display: "—".to_owned(), + standard_amount_display: "—".to_owned(), + total_amount_display: format_money(total), + input_price_display: "—".to_owned(), + output_price_display: "—".to_owned(), + cache_read_price_display: "—".to_owned(), + cache_write_price_display: "—".to_owned(), + service_tier_display: "—".to_owned(), + multiplier_display: "—".to_owned(), + }), + domain::UsageBilling::Calculated(value) => Some(BillingView { + input_amount_display: format_money(&value.input_amount), + output_amount_display: format_money(&value.output_amount), + cache_read_amount_display: format_money(&value.cache_read_amount), + cache_write_amount_display: format_money(&value.cache_write_amount), + standard_amount_display: format_money(&value.standard_amount), + total_amount_display: format_money(&value.total_amount), + input_price_display: format_token_price(&value.input_price_per_million), + output_price_display: format_token_price(&value.output_price_per_million), + cache_read_price_display: format_token_price(&value.cache_read_price_per_million), + cache_write_price_display: format_token_price(&value.cache_write_price_per_million), + service_tier_display: format_service_tier(value.service_tier.as_deref()), + multiplier_display: format!("{:.2}x", f64::from(value.multiplier_percent) / 100.0), + }), + } +} + +fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView { + let tokens = token_details(&record); + let billing = billing_view(record.billing.as_ref()); + let costs = record + .cost_amount + .as_ref() + .zip(record.cost_currency.as_ref()) + .map(|(amount, currency)| { + vec![CostView { + currency: currency.clone(), + estimated_amount: amount.as_str().to_owned(), + }] + }) + .unwrap_or_default(); + let status_code = record + .client_status_code + .or(record.upstream_status_code) + .map(i64::from); + let outcome = outcome_name(&record.outcome).to_owned(); + let message = record + .error_message + .clone() + .unwrap_or_else(|| outcome.clone()); + let first_token_display = display_duration(record.first_token_ms); + let latency_display = display_duration(record.latency_ms); + let cost_coverage = match record.cost_source.as_str() { + "provider_reported" | "calculated" => CostCoverageView { + known: 1, + partial: 0, + unknown: 0, + not_billable: 0, + }, + _ => CostCoverageView { + known: 0, + partial: 0, + unknown: 1, + not_billable: 0, + }, + }; + let model = record + .upstream_model_id + .clone() + .unwrap_or_else(|| record.requested_model_id.clone()); + let transport = record + .upstream_transport + .clone() + .or_else(|| Some(record.client_transport.clone())); + UsageRecordView { + id: record.id.clone(), + request_id: record.id, + client_api_key_id: Some(record.client_api_key_ref), + kind: record.operation, + provider: record.provider_kind, + account_id: record.provider_account_ref, + account_email: record.provider_account_email, + route: None, + model, + requested_model: Some(record.requested_model_id), + upstream_model: record.upstream_model_id, + service_tier: None, + status_code, + transport, + attempt_index: None, + attempt_count: u64::from(record.attempt_count), + response_id: record.client_response_id, + upstream_request_id: record.upstream_request_id, + latency_ms: record.latency_ms, + first_token_ms: record.first_token_ms, + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cached_tokens: record.cached_tokens, + cache_write_tokens: record.cache_write_tokens, + reasoning_tokens: record.reasoning_tokens, + message, + metadata: UsageRecordMetadataView { + protocol: record.protocol, + logical_outcome: outcome.clone(), + attempt_count: u64::from(record.attempt_count), + }, + created_at: record.started_at, + created_at_display: china_datetime(&record.started_at), + client_ip: record.client_ip, + user_agent: record.user_agent, + reasoning_effort: record.reasoning_effort, + reasoning_preset: record.reasoning_preset, + compact: Some(record.compact), + request_kind: record.request_kind, + subagent_kind: record.subagent_kind, + token_details: tokens, + billing, + costs, + cost_coverage, + first_token_latency_ms: record.first_token_ms, + first_token_latency_ms_display: first_token_display, + latency_ms_display: latency_display, + logical_outcome: outcome, + } +} + +fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptView { + let occurred_at = attempt.occurred_at; + UsageAttemptView { + id: attempt.id, + attempt_index: attempt.attempt_index, + trigger: attempt.source, + provider: attempt + .provider_kind + .unwrap_or_else(|| "unknown".to_owned()), + provider_instance_id: attempt + .provider_instance_id + .unwrap_or_else(|| "unknown".to_owned()), + model: attempt + .upstream_model_id + .unwrap_or_else(|| "unknown".to_owned()), + transport: attempt + .upstream_transport + .unwrap_or_else(|| "unknown".to_owned()), + send_state: attempt + .upstream_send_state + .unwrap_or_else(|| "unknown".to_owned()), + outcome: outcome_name(&attempt.outcome).to_owned(), + downstream_committed: attempt.downstream_committed, + status_code: attempt.status_code, + provider_error_code: attempt.provider_error_code, + failure_class: attempt.failure_kind, + cost_estimate_status: attempt + .cost_source + .clone() + .unwrap_or_else(|| "unavailable".to_owned()), + estimated_cost_amount: attempt.cost_amount.map(|amount| amount.as_str().to_owned()), + estimated_cost_currency: attempt.cost_currency, + input_tokens: attempt.input_tokens, + output_tokens: attempt.output_tokens, + cached_tokens: attempt.cached_tokens, + total_tokens: attempt.total_tokens, + first_token_ms: None, + latency_ms: attempt.latency_ms, + credential_name: attempt.provider_account_ref, + account_email: None, + started_at: occurred_at, + completed_at: Some(occurred_at), + } +} + +fn usage_detail_view(detail: domain::UsageDetail) -> UsageRecordDetailView { + UsageRecordDetailView { + request: usage_record_view(detail.request), + attempts: detail + .attempts + .into_iter() + .map(usage_attempt_view) + .collect(), + } +} + +fn page_meta(page: u32, page_size: u16, total: u64) -> PageMeta { + let page_size_u64 = u64::from(page_size); + let total_pages = total.saturating_add(page_size_u64 - 1) / page_size_u64; + PageMeta::new( + page, + page_size, + total, + u32::try_from(total_pages).unwrap_or(u32::MAX), + ) +} + +fn usage_page_view( + page: domain::UsagePage, + page_number: u32, + page_size: u16, +) -> Result, WireValidationError> { + Ok(PageData { + items: page.items.into_iter().map(usage_record_view).collect(), + page: page_meta(page_number, page_size, page.total), + next_cursor: page + .next_cursor + .as_ref() + .map(encode_observability_cursor) + .transpose()?, + }) +} + +fn ops_error_view(error: domain::OpsError) -> OpsErrorView { + let status = error.status_code.map(i64::from); + OpsErrorView { + id: error.event_id, + request_id: error.request_id, + client_api_key_id: error.client_api_key_ref, + kind: error.operation.clone(), + provider: error.provider_kind, + account_id: error.provider_account_ref, + route: error.operation, + model: error.upstream_model_id, + status_code: status, + client_status_code: None, + upstream_status_code: None, + transport: error.upstream_transport, + attempt_index: error.attempt_index, + failure_class: error.failure_kind, + response_id: error.client_response_id, + upstream_request_id: error.upstream_request_id, + latency_ms: error.latency_ms, + message: error.message, + metadata: OpsErrorMetadataView { + source: error.source, + component: error.component, + attempt_id: None, + provider_instance_id: error.provider_instance_id, + account_label: None, + }, + created_at: error.occurred_at, + created_at_display: china_datetime(&error.occurred_at), + } +} + +fn ops_page_view( + page: domain::OpsErrorPage, + page_number: u32, + page_size: u16, +) -> Result, WireValidationError> { + Ok(PageData { + items: page.items.into_iter().map(ops_error_view).collect(), + page: page_meta(page_number, page_size, page.total), + next_cursor: page + .next_cursor + .as_ref() + .map(encode_observability_cursor) + .transpose()?, + }) +} + +fn trend_point_view(point: domain::TrendPoint) -> TrendPointView { + let local_time = china_label(point.bucket_start, "%H:%M"); + let label = china_label(point.bucket_start, "%m-%d %H:%M"); + let success_rate_value = point.success_rate.map(|value| value * 100.0); + TrendPointView { + time: local_time, + bucket: point.bucket_start, + label, + requests: format_compact_number(point.metrics.request_count), + requests_value: point.metrics.request_count, + input_tokens: format_compact_number(point.metrics.input_tokens), + input_tokens_value: point.metrics.input_tokens, + output_tokens: format_compact_number(point.metrics.output_tokens), + output_tokens_value: point.metrics.output_tokens, + cached_tokens: format_compact_number(point.metrics.cached_tokens), + cached_tokens_value: point.metrics.cached_tokens, + cache_hit_rate_value: point.cached_token_rate, + tokens_value: point.metrics.total_tokens, + errors: format_compact_number(point.service_failure_count), + errors_value: point.service_failure_count, + latency: display_duration(point.average_latency_ms), + latency_value: point.average_latency_ms, + first_token_latency: display_duration(point.average_first_token_latency_ms), + first_token_latency_value: point.average_first_token_latency_ms, + max_latency: display_duration(point.metrics.max_latency_ms), + max_latency_value: point.metrics.max_latency_ms, + min_latency: display_duration(point.metrics.min_latency_ms), + min_latency_value: point.metrics.min_latency_ms, + success_rate: success_rate_value + .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), + success_rate_value, + } +} + +fn trend_summary_view(kind: TrendKind, summary: &domain::TrendSummary) -> Vec { + match kind { + TrendKind::Usage => vec![ + TrendSummaryView { + label: "输入".to_owned(), + value: format_compact_number(summary.input_tokens), + ratio: None, + }, + TrendSummaryView { + label: "输出".to_owned(), + value: format_compact_number(summary.output_tokens), + ratio: None, + }, + TrendSummaryView { + label: "缓存".to_owned(), + value: format_compact_number(summary.cached_tokens), + ratio: None, + }, + ], + TrendKind::Latency => vec![ + TrendSummaryView { + label: "平均".to_owned(), + value: display_duration(summary.average_latency_ms), + ratio: None, + }, + TrendSummaryView { + label: "最高".to_owned(), + value: display_duration(summary.max_latency_ms), + ratio: None, + }, + TrendSummaryView { + label: "最低".to_owned(), + value: display_duration(summary.min_latency_ms), + ratio: None, + }, + ], + TrendKind::Errors => vec![ + TrendSummaryView { + label: "错误数".to_owned(), + value: format_compact_number(summary.service_failure_count), + ratio: None, + }, + TrendSummaryView { + label: "成功率".to_owned(), + value: "—".to_owned(), + ratio: summary + .success_rate + .map(|value| format!("{:.1}%", value * 100.0)), + }, + TrendSummaryView { + label: "总请求".to_owned(), + value: format_compact_number(summary.request_count), + ratio: None, + }, + ], + } +} + +fn trend_view(trend: domain::Trend, kind: TrendKind) -> TrendData { + TrendData { + kind, + summary: trend_summary_view(kind, &trend.summary), + points: trend.points.into_iter().map(trend_point_view).collect(), + } +} + +fn health_status_name(status: domain::HealthStatus) -> &'static str { + match status { + domain::HealthStatus::Future => "future", + domain::HealthStatus::NoData => "no_data", + domain::HealthStatus::Unavailable => "unavailable", + domain::HealthStatus::LowSample => "low_sample", + domain::HealthStatus::Unstable => "unstable", + domain::HealthStatus::Stable => "stable", + } +} + +fn reliability_display(value: Option) -> String { + value.map_or_else(|| "-".to_owned(), |value| format!("{value:.1}%")) +} + +fn health_timeline_view(timeline: domain::HealthTimeline) -> HealthTimelineView { + HealthTimelineView { + title: "请求健康时间线".to_owned(), + description: "有效请求可用性".to_owned(), + reliability_display: reliability_display(timeline.reliability_percent), + status: health_status_name(timeline.status).to_owned(), + success_requests: timeline.success_requests, + failed_requests: timeline.failed_requests, + cancelled_requests: timeline.cancelled_requests, + incomplete_requests: timeline.incomplete_requests, + caller_error_requests: timeline.caller_error_requests, + points: timeline + .points + .into_iter() + .enumerate() + .map(|(index, point)| { + let elapsed_minutes = i64::try_from(index).unwrap_or(i64::MAX).saturating_mul(15); + HealthTimelinePointView { + time: format!("{:02}:{:02}", elapsed_minutes / 60, elapsed_minutes % 60), + status: health_status_name(point.status).to_owned(), + reliability_display: reliability_display(point.reliability_percent), + success_requests: point.success_requests, + failed_requests: point.failed_requests, + cancelled_requests: point.cancelled_requests, + incomplete_requests: point.incomplete_requests, + caller_error_requests: point.caller_error_requests, + } + }) + .collect(), + } +} + +fn wire_profile_view(profile: domain::DashboardWireProfile) -> DashboardWireProfileView { + DashboardWireProfileView { + originator: profile.originator, + codex_version: profile.codex_version, + desktop_version: profile.desktop_version, + desktop_build: profile.desktop_build, + target: DashboardWireTargetView { + os_type: profile.target.os_type, + os_version: profile.target.os_version, + arch: profile.target.arch, + terminal: profile.target.terminal, + }, + user_agent: profile.user_agent, + verified_at: profile.verified_at, + release: DashboardDesktopReleaseView { + status: profile.release.status.into(), + checked_at: profile.release.checked_at, + latest_version: profile.release.latest_version, + latest_build: profile.release.latest_build, + published_at: profile.release.published_at, + minimum_system_version: profile.release.minimum_system_version, + hardware_requirements: profile.release.hardware_requirements, + download_url: profile.release.download_url, + download_size: profile.release.download_size, + signature_present: profile.release.signature_present, + error: profile.release.error, + }, + } +} + +fn relative_time(value: Option>, now: DateTime) -> String { + let Some(value) = value else { + return "从未使用".to_owned(); + }; + let elapsed = now.signed_duration_since(value); + if elapsed.num_seconds() < 0 { + return china_datetime(&value); + } + if elapsed.num_seconds() < 60 { + return "刚刚".to_owned(); + } + if elapsed.num_minutes() < 60 { + return format!("{} 分钟前", elapsed.num_minutes()); + } + if elapsed.num_hours() < 24 { + return format!("{} 小时前", elapsed.num_hours()); + } + format!("{} 天前", elapsed.num_days()) +} + +fn rotation_strategy_name( + strategy: gateway_admin::model::settings::RotationStrategy, +) -> &'static str { + match strategy { + gateway_admin::model::settings::RotationStrategy::Smart => "smart", + gateway_admin::model::settings::RotationStrategy::QuotaResetPriority => { + "quota_reset_priority" + } + gateway_admin::model::settings::RotationStrategy::RoundRobin => "round_robin", + gateway_admin::model::settings::RotationStrategy::Sticky => "sticky", + } +} + +fn dashboard_view(result: domain::DashboardResult, kind: TrendKind) -> DashboardDataView { + let domain::DashboardResult { + observation, + today, + yesterday, + total_billing_usd, + total_cached_token_rate, + average_first_token_latency_ms, + trend, + health_timeline, + wire_profile, + capacity, + rotation_strategy, + } = result; + let domain::DashboardObservation { + range, + requests, + attempts, + provider_accounts, + trend: _, + account_usage, + recent_requests, + } = observation; + let mut account_usage_views = Vec::with_capacity(account_usage.len()); + let mut credential_usage_views = Vec::with_capacity(account_usage.len()); + for credential in account_usage { + account_usage_views.push(DashboardAccountUsageView { + id: credential.account_id.clone(), + email: credential + .email + .clone() + .unwrap_or_else(|| credential.name.clone()), + plan_type: credential.plan_type.clone(), + tokens: credential + .total_tokens + .map_or_else(|| "—".to_owned(), format_compact_number), + quota_used_percent: None, + last_used: relative_time(credential.last_used_at, range.end), + }); + credential_usage_views.push(DashboardCredentialUsageView { + id: credential.account_id, + display_name: credential.email.unwrap_or(credential.name), + plan_type: credential.plan_type, + tokens: credential + .total_tokens + .map_or_else(|| "-".to_owned(), format_compact_number), + tokens_value: credential.total_tokens, + last_used: credential + .last_used_at + .map_or_else(|| "-".to_owned(), |value| china_datetime(&value)), + provider: credential.provider_kind, + availability: credential.availability, + request_count: credential.request_count, + }); + } + DashboardDataView { + cards: DashboardCardsView { + credentials: DashboardCredentialsCardView { + total: format_compact_number(provider_accounts.total), + total_value: provider_accounts.total, + enabled: format_compact_number(provider_accounts.enabled), + enabled_value: provider_accounts.enabled, + unavailable: format_compact_number(provider_accounts.unavailable), + unavailable_value: provider_accounts.unavailable, + }, + traffic: DashboardTrafficCardView { + today_requests: format_compact_number(today.request_count), + today_requests_value: today.request_count, + yesterday_requests_value: yesterday.request_count, + total_requests: format_compact_number(requests.request_count), + }, + tokens: DashboardTokensCardView { + today_tokens: format_compact_number(today.total_tokens), + today_tokens_value: today.total_tokens, + yesterday_tokens_value: yesterday.total_tokens, + total_tokens: format_compact_number(requests.total_tokens), + total_billing_amount_usd: total_billing_usd + .as_ref() + .map_or_else(|| "—".to_owned(), |amount| format!("${}", amount.as_str())), + }, + cache: DashboardCacheCardView { + today_hit_rate: display_rate(today.cached_token_rate), + today_hit_rate_value: today.observed_cached_token_rate, + yesterday_hit_rate_value: yesterday.observed_cached_token_rate, + total_hit_rate: display_rate(total_cached_token_rate), + total_cached_tokens: format_compact_number(requests.cached_tokens), + average_first_token_latency_ms: display_duration(average_first_token_latency_ms), + }, + }, + trend: trend_view(trend, kind), + health_timeline: health_timeline_view(health_timeline), + wire_profile: wire_profile_view(wire_profile), + account_usage: account_usage_views, + credential_usage: credential_usage_views, + usage_records: recent_requests.into_iter().map(usage_record_view).collect(), + pool_summary: DashboardPoolSummaryView { + total: provider_accounts.total, + active: provider_accounts.active, + expired: provider_accounts.expired, + quota_exhausted: provider_accounts.quota_exhausted, + refreshing: provider_accounts.refreshing, + disabled: provider_accounts.disabled, + banned: provider_accounts.banned, + }, + capacity_info: DashboardCapacityInfoView { + max_concurrent_per_account: capacity.max_concurrent_per_account, + total_slots: capacity.total_slots, + used_slots: capacity.used_slots, + available_slots: capacity.available_slots, + }, + rotation_strategy: rotation_strategy_name(rotation_strategy).to_owned(), + logical_requests: request_metrics_view(&requests), + attempts: attempt_metrics_view(&attempts), + costs: cost_views(&attempts.costs), + } +} + +fn usage_summary_view(summary: domain::UsageSummary) -> UsageSummaryView { + let overview = summary.overview; + UsageSummaryView { + total_requests: format_compact_number(overview.requests.request_count), + input_tokens: format_compact_number(overview.requests.input_tokens), + output_tokens: format_compact_number(overview.requests.output_tokens), + cached_tokens: format_compact_number(overview.requests.cached_tokens), + cache_write_tokens: format_compact_number(overview.requests.cache_write_tokens), + total_tokens: format_compact_number(overview.requests.total_tokens), + average_latency_ms: display_duration(summary.average_latency_ms), + logical_requests: request_metrics_view(&overview.requests), + attempts: attempt_metrics_view(&overview.attempts), + } +} + +fn usage_insights_view(insights: domain::UsageInsights) -> UsageInsightsOverviewView { + let health_points = insights + .health + .points + .iter() + .map(|point| OverviewHealthPointView { + bucket: point.bucket_start, + label: china_label(point.bucket_start, "%m-%d %H:%M"), + success_requests: point.success_requests, + failed_requests: point.failed_requests, + cancelled_requests: point.cancelled_requests, + incomplete_requests: point.incomplete_requests, + caller_error_requests: point.caller_error_requests, + error_rate: point.error_rate, + }) + .collect(); + let performance_points = insights + .performance + .points + .iter() + .map(|point| OverviewPerformancePointView { + bucket: point.bucket_start, + label: china_label(point.bucket_start, "%m-%d %H:%M"), + latency_p50_ms: point.latency_percentiles.p50_ms.map(|value| value.as_f64()), + latency_p95_ms: point.latency_percentiles.p95_ms.map(|value| value.as_f64()), + latency_p99_ms: point.latency_percentiles.p99_ms.map(|value| value.as_f64()), + first_token_p50_ms: point + .first_token_latency_percentiles + .p50_ms + .map(|value| value.as_f64()), + first_token_p95_ms: point + .first_token_latency_percentiles + .p95_ms + .map(|value| value.as_f64()), + first_token_p99_ms: point + .first_token_latency_percentiles + .p99_ms + .map(|value| value.as_f64()), + }) + .collect(); + let cost_points = insights + .cost + .points + .iter() + .map(|point| OverviewCostPointView { + bucket: point.bucket_start, + label: china_label(point.bucket_start, "%m-%d %H:%M"), + input_tokens: point.input_tokens, + output_tokens: point.output_tokens, + cached_tokens: point.cached_tokens, + total_tokens: point.total_tokens, + estimated_cost: point.estimated_cost.as_ref().map(ToString::to_string), + standard_cost: point.standard_cost.as_ref().map(ToString::to_string), + cached_token_rate: point.cached_token_rate, + cache_hit_request_rate: point.cache_hit_request_rate, + }) + .collect(); + UsageInsightsOverviewView { + granularity: granularity_name(insights.granularity).to_owned(), + health: OverviewHealthView { + total_requests: insights.health.total_requests, + success_requests: insights.health.success_requests, + failed_requests: insights.health.failed_requests, + cancelled_requests: insights.health.cancelled_requests, + incomplete_requests: insights.health.incomplete_requests, + caller_error_requests: insights.health.caller_error_requests, + success_rate: insights.health.success_rate, + request_change_rate: None, + success_rate_change: None, + points: health_points, + }, + performance: OverviewPerformanceView { + latency_p50_ms: insights + .performance + .latency_percentiles + .p50_ms + .map(|value| value.as_f64()), + latency_p95_ms: insights + .performance + .latency_percentiles + .p95_ms + .map(|value| value.as_f64()), + latency_p99_ms: insights + .performance + .latency_percentiles + .p99_ms + .map(|value| value.as_f64()), + first_token_p50_ms: insights + .performance + .first_token_latency_percentiles + .p50_ms + .map(|value| value.as_f64()), + first_token_p95_ms: insights + .performance + .first_token_latency_percentiles + .p95_ms + .map(|value| value.as_f64()), + first_token_p99_ms: insights + .performance + .first_token_latency_percentiles + .p99_ms + .map(|value| value.as_f64()), + latency_coverage: insights.performance.latency_coverage, + first_token_coverage: insights.performance.first_token_coverage, + points: performance_points, + }, + cost: OverviewCostView { + estimated_cost: insights + .cost + .estimated_cost + .as_ref() + .map(ToString::to_string), + standard_cost: insights + .cost + .standard_cost + .as_ref() + .map(ToString::to_string), + cost_per_request: insights + .cost + .cost_per_request + .as_ref() + .map(ToString::to_string), + tokens_per_request: insights.cost.tokens_per_request, + cached_token_rate: insights.cost.cached_token_rate, + cache_hit_request_rate: insights.cost.cache_hit_request_rate, + input_tokens: insights.cost.input_tokens, + output_tokens: insights.cost.output_tokens, + cached_tokens: insights.cost.cached_tokens, + total_tokens: insights.cost.total_tokens, + points: cost_points, + costs: cost_views(&insights.cost.costs), + coverage: cost_coverage_view(&insights.cost.coverage), + }, + attempts: attempt_metrics_view(&insights.attempts), + providers: insights + .providers + .into_iter() + .map(|provider| ProviderOverviewView { + provider: provider.provider_kind, + request_count: provider.request_count, + attempt_count: provider.attempt_count, + failure_count: provider.failure_count, + total_tokens: provider.total_tokens, + }) + .collect(), + } +} + +fn granularity_name(granularity: domain::Granularity) -> &'static str { + match granularity { + domain::Granularity::FifteenMinutes => "15m", + domain::Granularity::Hour => "1h", + domain::Granularity::Day => "1d", + } +} + +fn diagnostics_view( + result: domain::DiagnosticsResult, + dimension: DiagnosticDimension, +) -> DiagnosticsView { + DiagnosticsView { + dimension: dimension.display_name().to_owned(), + items: result + .items + .into_iter() + .map(|item| DiagnosticItemView { + name: if item.name == "__none__" { + "未知".to_owned() + } else { + item.name + }, + request_count: item.request_count, + success_count: item.success_count, + error_count: item.error_count, + error_rate: item.error_rate, + request_share: item.request_share, + average_latency_ms: item.average_latency_ms, + estimated_cost: item.estimated_cost.as_ref().map(ToString::to_string), + attempt_count: item.attempt_count, + total_tokens: item.total_tokens, + }) + .collect(), + } +} + +fn map_wire_error(error: WireValidationError) -> AdminError { + let message = match error.field() { + "timeRange" => "Invalid time range", + "statusCode" => "Status code must be between 100 and 599", + "attemptIndex" => "Attempt index is out of range", + "kind" => "Invalid dashboard trend kind", + "dimension" => "Invalid diagnostics dimension", + "outcome" => "Invalid Observability query", + "id" => "Usage record ID is required", + "page" | "pageSize" => "Invalid Observability query", + "cursor" => "Invalid observability cursor", + _ => "Invalid observability query", + }; + AdminError::bad_request(message) +} + +fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError { + map_admin_service_error(error, "Observability repository unavailable") +} diff --git a/backend/crates/gateway-api/src/admin/openai.rs b/backend/crates/gateway-api/src/admin/openai.rs index 16e1b8ea2..2a11cd117 100644 --- a/backend/crates/gateway-api/src/admin/openai.rs +++ b/backend/crates/gateway-api/src/admin/openai.rs @@ -1,6 +1,5 @@ //! Codex OAuth credential 管理端 wire contract。 -use async_trait::async_trait; use axum::{ Json, Router, extract::{Query, State}, @@ -8,13 +7,28 @@ use axum::{ response::IntoResponse, routing::{get, post}, }; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use chrono::{DateTime, Utc}; +use gateway_admin::model::{ + AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision, + accounts::{AccountAvailability, AccountRecord}, + provider_credentials::{ + AuthorizationStarted, CompleteAuthorization, CredentialAvailabilityFilter, + CredentialCursor, CredentialDetails, CredentialImportResult, CredentialListQuery, + CredentialListWindow, CredentialMutation, CredentialMutationResult, CredentialPage, + ImportCredentials, ProviderDocument, ReauthorizationTarget, RotateCredential, + StartAuthorization, + }, +}; +use gateway_core::{ + engine::credential::{OpaqueProviderData, ProviderAccountId}, + routing::ProviderInstanceId, +}; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; use crate::admin::{ - AdminAuth, AdminEnvelope, AdminError, AdminRequestContext, AdminResponse, AdminServiceError, - AdminServiceErrorKind, AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, }; pub const DEFAULT_PAGE_SIZE: u16 = 50; @@ -67,6 +81,28 @@ impl ListCredentialsQuery { } Ok(()) } + + fn into_command(self) -> Result { + self.validate()?; + Ok(CredentialListQuery { + provider_instance_id: self + .provider_instance_id + .map(ProviderInstanceId::new) + .transpose() + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + availability: self + .availability + .as_deref() + .map(parse_availability) + .transpose()?, + enabled: self.enabled, + window: CredentialListWindow::Page { + cursor: self.cursor.as_deref().map(decode_cursor).transpose()?, + page_size: PageSize::new(self.limit.unwrap_or(DEFAULT_PAGE_SIZE)) + .map_err(|_| WireValidationError::new("limit"))?, + }, + }) + } } /// Codex credential 详情查询。 @@ -81,6 +117,12 @@ impl CredentialDetailsQuery { pub fn validate(&self) -> Result<(), WireValidationError> { require_id(&self.credential_id, "credentialId") } + + fn into_account_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.credential_id) + .map_err(|_| WireValidationError::new("credentialId")) + } } /// OpenAI Provider-owned 账号导入文档。 @@ -105,6 +147,23 @@ impl ImportCredentialsDocumentRequest { } Ok(()) } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(ImportCredentials { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + provider_instance_id: ProviderInstanceId::new(self.provider_instance_id) + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + document: provider_document(self.document, "document")?, + }) + } } /// 启动 Codex Authorization Code + PKCE；可创建新账号或绑定既有账号重新授权。 @@ -136,6 +195,33 @@ impl StartOAuthAuthorizationRequest { } Ok(()) } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + let reauthorization = match (self.credential_id, self.expected_credential_revision) { + (None, None) => None, + (Some(account_id), Some(credential_revision)) => Some(ReauthorizationTarget { + account_id: ProviderAccountId::new(account_id) + .map_err(|_| WireValidationError::new("credentialId"))?, + credential_revision: revision(credential_revision, "expectedCredentialRevision")?, + }), + _ => return Err(WireValidationError::new("reauthorization")), + }; + Ok(StartAuthorization { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + provider_instance_id: ProviderInstanceId::new(self.provider_instance_id) + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + name: self.name, + reauthorization, + }) + } } /// 完成 Codex Authorization Code flow；flow ID 与 callback 都放在 POST body。 @@ -148,9 +234,7 @@ pub struct CompleteOAuthAuthorizationRequest { impl CompleteOAuthAuthorizationRequest { pub fn validate(&self) -> Result<(), WireValidationError> { - use base64::Engine as _; - - if !base64::engine::general_purpose::URL_SAFE_NO_PAD + if !URL_SAFE_NO_PAD .decode(&self.flow_id) .is_ok_and(|decoded| decoded.len() == 32) { @@ -164,6 +248,18 @@ impl CompleteOAuthAuthorizationRequest { } Ok(()) } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(CompleteAuthorization { + context, + flow_id: self.flow_id, + callback_url: self.callback_url, + }) + } } /// Codex OAuth token 轮换请求。 @@ -188,6 +284,35 @@ impl RotateCredentialRequest { )?; validate_oauth_material(&self.access_token, self.refresh_token.as_deref()) } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + let mut material = Map::new(); + material.insert("access_token".to_owned(), Value::String(self.access_token)); + material.insert( + "refresh_token".to_owned(), + self.refresh_token.map_or(Value::Null, Value::String), + ); + Ok(RotateCredential { + mutation: CredentialMutation { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + account_id: ProviderAccountId::new(self.credential_id) + .map_err(|_| WireValidationError::new("credentialId"))?, + }, + expected_credential_revision: revision( + self.expected_credential_revision, + "expectedCredentialRevision", + )?, + provider_material: ProviderDocument::new(OpaqueProviderData::new(material)), + }) + } } /// Codex credential 生命周期 mutation 请求。 @@ -204,6 +329,22 @@ impl CredentialMutationRequest { require_id(&self.credential_id, "credentialId")?; require_positive(self.expected_config_revision, "expectedConfigRevision") } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(CredentialMutation { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + account_id: ProviderAccountId::new(self.credential_id) + .map_err(|_| WireValidationError::new("credentialId"))?, + }) + } } /// 列表游标的 JSON 形状；HTTP 组合层负责 base64url 编解码。 @@ -287,6 +428,72 @@ pub struct CodexCredentialRotationData { pub credential_revision: i64, } +impl From for CodexCredentialView { + fn from(account: AccountRecord) -> Self { + Self { + id: account.id, + provider_instance_id: account.provider_instance_id.to_string(), + name: account.name, + email: account.email, + upstream_user_id: account.upstream_user_id, + upstream_account_id: account.upstream_account_id, + enabled: account.enabled, + credential_revision: i64::try_from(account.credential_revision.get()) + .unwrap_or(i64::MAX), + has_refresh_token: account.has_refresh_token, + availability: codex_availability(account.availability).to_owned(), + availability_reason: account.availability_reason, + plan_type: account.plan_type, + access_token_expires_at: account.access_token_expires_at, + next_refresh_at: account.next_refresh_at, + cooldown_until: account.cooldown_until, + created_at: account.created_at, + updated_at: account.updated_at, + } + } +} + +impl From for CodexCredentialDetailsData { + fn from(details: CredentialDetails) -> Self { + Self { + config_revision: details.config_revision.get(), + credential: details.credential.into(), + } + } +} + +impl From for CodexCredentialsDocumentImportData { + fn from(result: CredentialImportResult) -> Self { + Self { + config_revision: result.config_revision.get(), + credential_ids: result + .credential_ids + .into_iter() + .map(|id| id.to_string()) + .collect(), + } + } +} + +impl From for CodexOAuthAuthorizationStartedData { + fn from(started: AuthorizationStarted) -> Self { + Self { + flow_id: started.flow_id, + authorization_url: started.authorization_url, + expires_at: started.expires_at, + } + } +} + +impl From for CodexCredentialMutationData { + fn from(result: CredentialMutationResult) -> Self { + Self { + config_revision: result.config_revision.get(), + credential_id: result.account_id.to_string(), + } + } +} + fn validate_oauth_material( access_token: &str, refresh_token: Option<&str>, @@ -346,63 +553,10 @@ fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidatio Ok(()) } -/// Codex OAuth credential 管理应用端口。 -#[async_trait] -pub trait CodexAdminService: Send + Sync { - async fn list( - &self, - query: ListCredentialsQuery, - ) -> Result; - async fn details( - &self, - credential_id: String, - ) -> Result; - async fn import_document( - &self, - context: &AdminRequestContext, - request: ImportCredentialsDocumentRequest, - ) -> Result; - async fn start_authorization( - &self, - context: &AdminRequestContext, - request: StartOAuthAuthorizationRequest, - ) -> Result; - async fn complete_authorization( - &self, - context: &AdminRequestContext, - request: CompleteOAuthAuthorizationRequest, - ) -> Result; - async fn rotate( - &self, - context: &AdminRequestContext, - request: RotateCredentialRequest, - ) -> Result; - async fn enable( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; - async fn disable( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; - async fn delete( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; -} - -/// Codex HTTP module 所需最小 state。 -pub trait CodexAdminState: AdminSessionState { - fn codex_admin_service(&self) -> &dyn CodexAdminService; -} - /// 构造固定 GET/POST Codex credential 管理路由。 pub fn router() -> Router where - S: CodexAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/openai/credentials", get(list_credentials::)) @@ -443,13 +597,16 @@ async fn list_credentials( Query(query): Query, ) -> Result where - S: CodexAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; - ok( - StatusCode::OK, - state.codex_admin_service().list(query).await, - ) + let command = query.into_command().map_err(map_wire_error)?; + let page = state + .admin_services() + .openai() + .list(command) + .await + .map_err(map_service_error)?; + ok(StatusCode::OK, credential_list_data(page)?) } async fn credential_details( @@ -458,88 +615,173 @@ async fn credential_details( Query(query): Query, ) -> Result where - S: CodexAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; + let account_id = query.into_account_id().map_err(map_wire_error)?; + let details = state + .admin_services() + .openai() + .details(&account_id) + .await + .map_err(map_service_error)?; + ok(StatusCode::OK, CodexCredentialDetailsData::from(details)) +} + +async fn import_credentials_document( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .openai() + .import_document(command) + .await + .map_err(map_service_error)?; ok( - StatusCode::OK, - state - .codex_admin_service() - .details(query.credential_id) - .await, + StatusCode::CREATED, + CodexCredentialsDocumentImportData::from(result), ) } -macro_rules! codex_post_handler { - ($name:ident, $request:ty, $method:ident, $status:expr) => { - async fn $name( - auth: AdminAuth, - State(state): State, - Json(request): Json<$request>, - ) -> Result - where - S: CodexAdminState + Send + Sync, - { - request.validate().map_err(map_wire_error)?; - ok( - $status, - state - .codex_admin_service() - .$method(auth.context(), request) - .await, - ) - } - }; +async fn start_oauth_authorization( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .openai() + .start_authorization(command) + .await + .map_err(map_service_error)?; + ok( + StatusCode::CREATED, + CodexOAuthAuthorizationStartedData::from(result), + ) } -codex_post_handler!( - import_credentials_document, - ImportCredentialsDocumentRequest, - import_document, - StatusCode::CREATED -); -codex_post_handler!( - start_oauth_authorization, - StartOAuthAuthorizationRequest, - start_authorization, - StatusCode::CREATED -); -codex_post_handler!( - complete_oauth_authorization, - CompleteOAuthAuthorizationRequest, - complete_authorization, - StatusCode::CREATED -); -codex_post_handler!( - rotate_credential, - RotateCredentialRequest, - rotate, - StatusCode::OK -); -codex_post_handler!( - enable_credential, - CredentialMutationRequest, - enable, - StatusCode::OK -); -codex_post_handler!( - disable_credential, - CredentialMutationRequest, - disable, - StatusCode::OK -); -codex_post_handler!( - delete_credential, - CredentialMutationRequest, - delete, - StatusCode::OK -); +async fn complete_oauth_authorization( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .openai() + .complete_authorization(command) + .await + .map_err(map_service_error)?; + ok( + StatusCode::CREATED, + CodexCredentialMutationData::from(result), + ) +} + +async fn rotate_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .openai() + .rotate(command) + .await + .map_err(map_service_error)?; + ok(StatusCode::OK, credential_rotation_data(result)?) +} + +async fn enable_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Enable).await +} + +async fn disable_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Disable).await +} + +async fn delete_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Delete).await +} + +#[derive(Clone, Copy)] +enum CredentialAction { + Enable, + Disable, + Delete, +} + +async fn mutate_credential( + auth: AdminAuth, + state: S, + request: CredentialMutationRequest, + action: CredentialAction, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let service = state.admin_services().openai(); + let result = match action { + CredentialAction::Enable => service.enable(command).await, + CredentialAction::Disable => service.disable(command).await, + CredentialAction::Delete => service.delete(command).await, + } + .map_err(map_service_error)?; + ok(StatusCode::OK, CodexCredentialMutationData::from(result)) +} fn ok( status: StatusCode, - result: Result, + data: T, ) -> Result>, AdminError> { - let data = result.map_err(map_service_error)?; Ok(AdminResponse::new(status, AdminEnvelope::ok(data))) } @@ -549,10 +791,105 @@ fn map_wire_error(error: WireValidationError) -> AdminError { fn map_service_error(error: AdminServiceError) -> AdminError { match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => AdminError::service_unavailable(error.to_string()), - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), + AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()), + AdminErrorKind::Unauthorized => AdminError::admin_session_required(), + AdminErrorKind::NotFound => AdminError::not_found(error.to_string()), + AdminErrorKind::Conflict => AdminError::conflict(error.to_string()), + AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(), + AdminErrorKind::BadGateway => AdminError::bad_gateway(error.to_string()), + AdminErrorKind::Unavailable => AdminError::service_unavailable(error.to_string()), + AdminErrorKind::Internal => AdminError::internal(error.to_string()), + } +} + +fn revision(value: u64, field: &'static str) -> Result { + Revision::new(value).map_err(|_| WireValidationError::new(field)) +} + +fn provider_document( + value: Value, + field: &'static str, +) -> Result { + match value { + Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))), + _ => Err(WireValidationError::new(field)), + } +} + +fn parse_availability(value: &str) -> Result { + match value { + "unknown" => Ok(CredentialAvailabilityFilter::Exact( + AccountAvailability::Unknown, + )), + "ready" => Ok(CredentialAvailabilityFilter::Exact( + AccountAvailability::Ready, + )), + "cooldown" => Ok(CredentialAvailabilityFilter::Exact( + AccountAvailability::Cooldown, + )), + "exhausted" => Ok(CredentialAvailabilityFilter::Exact( + AccountAvailability::QuotaExhausted, + )), + "invalid" => Ok(CredentialAvailabilityFilter::AnyOf(vec![ + AccountAvailability::Expired, + AccountAvailability::Banned, + AccountAvailability::Invalid, + ])), + _ => Err(WireValidationError::new("availability")), + } +} + +fn decode_cursor(value: &str) -> Result { + let encoded = URL_SAFE_NO_PAD + .decode(value) + .map_err(|_| WireValidationError::new("cursor"))?; + let wire: CredentialCursorWire = + serde_json::from_slice(&encoded).map_err(|_| WireValidationError::new("cursor"))?; + require_id(&wire.credential_id, "cursor")?; + Ok(CredentialCursor { + created_at: wire.created_at, + account_id: ProviderAccountId::new(wire.credential_id) + .map_err(|_| WireValidationError::new("cursor"))?, + }) +} + +fn encode_cursor(cursor: &CredentialCursor) -> Result { + let encoded = serde_json::to_vec(&CredentialCursorWire { + created_at: cursor.created_at, + credential_id: cursor.account_id.to_string(), + }) + .map_err(|_| AdminError::internal("Failed to encode credential cursor"))?; + Ok(URL_SAFE_NO_PAD.encode(encoded)) +} + +fn credential_list_data(page: CredentialPage) -> Result { + Ok(CodexCredentialListData { + config_revision: page.config_revision.get(), + items: page.items.into_iter().map(Into::into).collect(), + next_cursor: page.next_cursor.as_ref().map(encode_cursor).transpose()?, + }) +} + +fn credential_rotation_data( + result: CredentialMutationResult, +) -> Result { + let credential_revision = result + .credential_revision + .ok_or_else(|| AdminError::internal("Credential revision is missing"))?; + Ok(CodexCredentialRotationData { + credential_id: result.account_id.to_string(), + credential_revision: i64::try_from(credential_revision.get()).unwrap_or(i64::MAX), + }) +} + +const fn codex_availability(value: AccountAvailability) -> &'static str { + match value { + AccountAvailability::Unknown => "unknown", + AccountAvailability::Ready => "ready", + AccountAvailability::Cooldown => "cooldown", + AccountAvailability::QuotaExhausted => "exhausted", + AccountAvailability::Expired + | AccountAvailability::Banned + | AccountAvailability::Invalid => "invalid", } } diff --git a/backend/crates/gateway-api/src/admin/settings.rs b/backend/crates/gateway-api/src/admin/settings.rs index 12ae81b7d..81c894fe3 100644 --- a/backend/crates/gateway-api/src/admin/settings.rs +++ b/backend/crates/gateway-api/src/admin/settings.rs @@ -2,7 +2,6 @@ use std::{collections::BTreeMap, fmt}; -use async_trait::async_trait; use axum::{ Json, Router, extract::State, @@ -11,11 +10,19 @@ use axum::{ routing::{get, post}, }; use chrono::{DateTime, Utc}; +use gateway_admin::model::{ + Revision, + settings::{ + ProviderModelMappings as DomainProviderModelMappings, ReplaceRuntimeSettings, + RotationStrategy, RuntimeSettings, + }, +}; +use gateway_core::routing::{ProviderKind, PublicModelId, UpstreamModelId}; use serde::{Deserialize, Serialize}; use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminRequestContext, AdminResponse, AdminServiceError, - AdminServiceErrorKind, AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, + wire::map_admin_service_error, }; /// 按 Provider 分组的精确模型映射：Provider → 客户端模型 → 上游模型。 @@ -82,6 +89,47 @@ impl UpdateRuntimeSettingsRequest { } Ok(()) } + + fn into_command(self) -> Result { + self.validate()?; + Ok(ReplaceRuntimeSettings { + expected_config_revision: Revision::new(self.expected_config_revision) + .map_err(|_| WireValidationError::new("settingsRevisionOverflow"))?, + provider_model_mappings: domain_provider_model_mappings(self.provider_model_mappings)?, + refresh_margin_seconds: self.refresh_margin_seconds, + refresh_concurrency: u32::try_from(self.refresh_concurrency) + .map_err(|_| WireValidationError::new("settingsRefreshConcurrencyOverflow"))?, + max_concurrent_per_account: u32::try_from(self.max_concurrent_per_account) + .map_err(|_| WireValidationError::new("settingsMaxConcurrencyOverflow"))?, + request_interval_ms: self.request_interval_ms, + rotation_strategy: parse_rotation_strategy(&self.rotation_strategy) + .ok_or_else(|| WireValidationError::new("rotationStrategy"))?, + usage_retention_days: u32::try_from(self.usage_retention_days) + .map_err(|_| WireValidationError::new("settingsUsageRetentionOverflow"))?, + ops_event_retention_days: u32::try_from(self.ops_event_retention_days) + .map_err(|_| WireValidationError::new("settingsOpsRetentionOverflow"))?, + audit_retention_days: u32::try_from(self.audit_retention_days) + .map_err(|_| WireValidationError::new("settingsAuditRetentionOverflow"))?, + }) + } +} + +impl From for RuntimeSettingsView { + fn from(settings: RuntimeSettings) -> Self { + Self { + config_revision: settings.config_revision.get(), + provider_model_mappings: wire_provider_model_mappings(settings.provider_model_mappings), + refresh_margin_seconds: settings.refresh_margin_seconds, + refresh_concurrency: u64::from(settings.refresh_concurrency), + max_concurrent_per_account: u64::from(settings.max_concurrent_per_account), + request_interval_ms: settings.request_interval_ms, + rotation_strategy: rotation_strategy_name(settings.rotation_strategy).to_owned(), + usage_retention_days: u64::from(settings.usage_retention_days), + ops_event_retention_days: u64::from(settings.ops_event_retention_days), + audit_retention_days: u64::from(settings.audit_retention_days), + updated_at: settings.updated_at, + } + } } /// 管理 API Key 状态；状态读取不回显完整值。 @@ -120,39 +168,10 @@ impl Default for DeletedAdminApiKey { } } -/// 设置页聚合和明文 Admin API Key 管理应用端口。 -#[async_trait] -pub trait AdminSettingsService: Send + Sync { - async fn load(&self) -> Result; - - async fn replace( - &self, - context: &AdminRequestContext, - request: UpdateRuntimeSettingsRequest, - ) -> Result; - - async fn admin_api_key_exists(&self) -> Result; - - async fn regenerate_admin_api_key( - &self, - context: &AdminRequestContext, - ) -> Result; - - async fn delete_admin_api_key( - &self, - context: &AdminRequestContext, - ) -> Result<(), AdminServiceError>; -} - -/// 设置 HTTP module 所需最小 state。 -pub trait AdminSettingsState: AdminSessionState { - fn admin_settings_service(&self) -> &dyn AdminSettingsService; -} - /// 构造固定 GET/POST 设置路由。 pub fn router() -> Router where - S: AdminSettingsState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route( @@ -174,14 +193,18 @@ async fn settings( State(state): State, ) -> Result where - S: AdminSettingsState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let data = state - .admin_settings_service() + let result = state + .admin_services() + .settings() .load() .await .map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RuntimeSettingsView::from(result)), + )) } async fn update_settings( @@ -190,15 +213,19 @@ async fn update_settings( Json(request): Json, ) -> Result where - S: AdminSettingsState + Send + Sync, + S: AdminSessionState + Send + Sync, { - request.validate().map_err(map_wire_error)?; - let data = state - .admin_settings_service() - .replace(auth.context(), request) + let command = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .settings() + .replace(&auth.context().mutation_context(), command) .await .map_err(map_service_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RuntimeSettingsView::from(result)), + )) } async fn admin_api_key_status( @@ -206,10 +233,11 @@ async fn admin_api_key_status( State(state): State, ) -> Result where - S: AdminSettingsState + Send + Sync, + S: AdminSessionState + Send + Sync, { let exists = state - .admin_settings_service() + .admin_services() + .settings() .admin_api_key_exists() .await .map_err(map_service_error)?; @@ -224,16 +252,19 @@ async fn regenerate_admin_api_key( State(state): State, ) -> Result where - S: AdminSettingsState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let key = state - .admin_settings_service() - .regenerate_admin_api_key(auth.context()) + let result = state + .admin_services() + .settings() + .regenerate_admin_api_key(&auth.context().mutation_context()) .await .map_err(map_service_error)?; Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(RegeneratedAdminApiKey { key }), + AdminEnvelope::ok(RegeneratedAdminApiKey { + key: result.key.expose_for_response().to_owned(), + }), )) } @@ -242,11 +273,12 @@ async fn delete_admin_api_key( State(state): State, ) -> Result where - S: AdminSettingsState + Send + Sync, + S: AdminSessionState + Send + Sync, { state - .admin_settings_service() - .delete_admin_api_key(auth.context()) + .admin_services() + .settings() + .delete_admin_api_key(&auth.context().mutation_context()) .await .map_err(map_service_error)?; Ok(AdminResponse::new( @@ -281,6 +313,64 @@ fn validate_provider_model_mappings( Ok(()) } +fn domain_provider_model_mappings( + mappings: ProviderModelMappings, +) -> Result { + mappings + .into_iter() + .map(|(provider, entries)| { + let provider = ProviderKind::new(provider) + .map_err(|_| WireValidationError::new("providerModelMappings"))?; + let entries = entries + .into_iter() + .map(|(requested, upstream)| { + Ok(( + PublicModelId::new(requested) + .map_err(|_| WireValidationError::new("providerModelMappings"))?, + UpstreamModelId::new(upstream) + .map_err(|_| WireValidationError::new("providerModelMappings"))?, + )) + }) + .collect::>()?; + Ok((provider, entries)) + }) + .collect() +} + +fn wire_provider_model_mappings(mappings: DomainProviderModelMappings) -> ProviderModelMappings { + mappings + .into_iter() + .map(|(provider, entries)| { + ( + provider.to_string(), + entries + .into_iter() + .map(|(requested, upstream)| (requested.to_string(), upstream.to_string())) + .collect(), + ) + }) + .collect() +} + +fn parse_rotation_strategy(value: &str) -> Option { + match value { + "smart" => Some(RotationStrategy::Smart), + "quota_reset_priority" => Some(RotationStrategy::QuotaResetPriority), + "round_robin" => Some(RotationStrategy::RoundRobin), + "sticky" => Some(RotationStrategy::Sticky), + _ => None, + } +} + +const fn rotation_strategy_name(strategy: RotationStrategy) -> &'static str { + match strategy { + RotationStrategy::Smart => "smart", + RotationStrategy::QuotaResetPriority => "quota_reset_priority", + RotationStrategy::RoundRobin => "round_robin", + RotationStrategy::Sticky => "sticky", + } +} + fn valid_slug(value: &str, max_len: usize) -> bool { !value.is_empty() && value.len() <= max_len @@ -298,17 +388,18 @@ fn valid_model_name(value: &str, max_len: usize) -> bool { } fn map_wire_error(error: WireValidationError) -> AdminError { - AdminError::bad_request(format!("Invalid field: {}", error.field())) + let message = match error.field() { + "settingsRevisionOverflow" => "expectedConfigRevision must be positive".to_owned(), + "settingsRefreshConcurrencyOverflow" => "Invalid refreshConcurrency".to_owned(), + "settingsMaxConcurrencyOverflow" => "Invalid maxConcurrentPerAccount".to_owned(), + "settingsUsageRetentionOverflow" => "Invalid usageRetentionDays".to_owned(), + "settingsOpsRetentionOverflow" => "Invalid opsEventRetentionDays".to_owned(), + "settingsAuditRetentionOverflow" => "Invalid auditRetentionDays".to_owned(), + field => format!("Invalid field: {field}"), + }; + AdminError::bad_request(message) } -fn map_service_error(error: AdminServiceError) -> AdminError { - match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => { - AdminError::service_unavailable("Settings repository unavailable") - } - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), - } +fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError { + map_admin_service_error(error, "Settings repository unavailable") } diff --git a/backend/crates/gateway-api/src/admin/system.rs b/backend/crates/gateway-api/src/admin/system.rs index 50f2c5187..4eb973ad4 100644 --- a/backend/crates/gateway-api/src/admin/system.rs +++ b/backend/crates/gateway-api/src/admin/system.rs @@ -2,9 +2,8 @@ //! //! 这里不依赖更新服务或进程控制；应用层通过窄端口提供系统操作事实。 -use std::{convert::Infallible, pin::Pin}; +use std::convert::Infallible; -use async_trait::async_trait; use axum::{ Json, Router, extract::{Query, State}, @@ -16,10 +15,17 @@ use axum::{ routing::{get, post}, }; use futures::{Stream, StreamExt}; -use serde::Deserialize; -use serde_json::Value; +use gateway_admin::model::system::{ + SystemOperationAccepted, SystemOperationKind, SystemOperationState, SystemOperationStatus, + SystemUpdateDetail, SystemUpdateEvent, SystemUpdateEventLevel, SystemUpdateStatus, + SystemVersion, +}; +use serde::{Deserialize, Serialize}; -use super::{AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState}; +use super::{ + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, + wire::map_admin_service_error, +}; /// 更新详情查询参数。 #[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq)] @@ -51,71 +57,162 @@ impl UpdateRequest { } } -/// 一条已经脱敏的更新事件。 -pub struct SystemUpdateEvent { - pub id: String, - pub data: Value, +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct SystemVersionView { + version: String, + git_sha: String, + build_time: String, + deployment_mode: String, + deployment_mode_label: String, + update_channel: String, + latest_version: String, + has_update: bool, + update_cached: bool, + update_warning: Option, } -/// 每个 SSE 订阅独占的更新事件流。 -pub type SystemUpdateEventStream = Pin + Send + 'static>>; - -/// 系统操作失败分类。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SystemAdminErrorKind { - Invalid, - Conflict, - BadGateway, - Internal, -} - -/// 系统操作安全错误。 -#[derive(Debug, thiserror::Error)] -#[error("{message}")] -pub struct SystemAdminError { - kind: SystemAdminErrorKind, - message: String, -} - -impl SystemAdminError { - #[must_use] - pub fn new(kind: SystemAdminErrorKind, message: impl Into) -> Self { +impl From for SystemVersionView { + fn from(version: SystemVersion) -> Self { Self { - kind, - message: message.into(), + deployment_mode_label: deployment_mode_label(&version.deployment_mode).to_owned(), + version: version.version, + git_sha: version.git_sha, + build_time: version.build_time, + deployment_mode: version.deployment_mode, + update_channel: version.update_channel, + latest_version: version.latest_version, + has_update: version.has_update, + update_cached: version.update_cached, + update_warning: version.update_warning, } } +} - #[must_use] - pub const fn kind(&self) -> SystemAdminErrorKind { - self.kind +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct SystemUpdateDetailView { + current_version: String, + latest_version: String, + has_update: bool, + deployment_mode: String, + deployment_mode_label: String, + build_type: String, + build_type_label: String, + release_url: Option, + notes: Option, + cached: bool, + update_supported: bool, + unsupported_reason: Option, + warning: Option, +} + +impl From for SystemUpdateDetailView { + fn from(detail: SystemUpdateDetail) -> Self { + Self { + deployment_mode_label: deployment_mode_label(&detail.deployment_mode).to_owned(), + build_type_label: build_type_label(&detail.build_type), + current_version: detail.current_version, + latest_version: detail.latest_version, + has_update: detail.has_update, + deployment_mode: detail.deployment_mode, + build_type: detail.build_type, + release_url: detail.release_url, + notes: detail.notes, + cached: detail.cached, + update_supported: detail.update_supported, + unsupported_reason: detail.unsupported_reason, + warning: detail.warning, + } } } -/// 系统版本、更新与进程控制应用端口。 -#[async_trait] -pub trait SystemAdminService: Send + Sync { - async fn version(&self) -> Result; - async fn update_detail(&self, refresh: bool) -> Result; - fn update_events(&self) -> SystemUpdateEventStream; - async fn perform_update( - &self, - target_version: Option, - ) -> Result; - async fn update_status(&self) -> Result; - async fn rollback(&self) -> Result; - async fn restart(&self) -> Result; +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct SystemOperationStateView { + operation_id: Option, + kind: Option<&'static str>, + status: &'static str, + target_version: Option, + message: Option, + error: Option, + started_at: Option, + finished_at: Option, } -/// 系统管理 HTTP module 所需最小 state。 -pub trait SystemAdminState: AdminSessionState { - fn system_admin_service(&self) -> &dyn SystemAdminService; +impl From for SystemOperationStateView { + fn from(operation: SystemOperationState) -> Self { + Self { + operation_id: operation.operation_id, + kind: operation.kind.map(operation_kind_name), + status: operation_status_name(operation.status), + target_version: operation.target_version, + message: operation.message, + error: operation.error, + started_at: operation.started_at.map(|value| value.to_rfc3339()), + finished_at: operation.finished_at.map(|value| value.to_rfc3339()), + } + } +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct SystemUpdateStatusView { + previous_version: Option, + current_version: Option, + operation: SystemOperationStateView, +} + +impl From for SystemUpdateStatusView { + fn from(status: SystemUpdateStatus) -> Self { + Self { + previous_version: status.previous_version, + current_version: status.current_version, + operation: status.operation.into(), + } + } +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct UpdateAcceptedView { + operation_id: String, + deployment_mode: String, + message: String, + need_restart: bool, + target_version: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct RollbackAcceptedView { + message: String, + need_restart: bool, + operation_id: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct RestartAcceptedView { + message: String, + operation_id: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct SystemUpdateEventView { + operation_id: Option, + level: &'static str, + step: Option, + message: String, + terminal: bool, + at: String, } /// 构造固定 GET/POST 系统管理路由。 pub fn router() -> Router where - S: SystemAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/system/version", get(version::)) @@ -135,9 +232,18 @@ async fn version( State(state): State, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response(state.system_admin_service().version().await) + let version = state + .admin_services() + .system() + .version() + .await + .map_err(map_system_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(SystemVersionView::from(version)), + )) } async fn update_detail( @@ -146,14 +252,18 @@ async fn update_detail( Query(query): Query, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response( - state - .system_admin_service() - .update_detail(query.refresh()) - .await, - ) + let detail = state + .admin_services() + .system() + .update_detail(query.refresh()) + .await + .map_err(map_system_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(SystemUpdateDetailView::from(detail)), + )) } async fn update_event_stream( @@ -161,12 +271,17 @@ async fn update_event_stream( State(state): State, ) -> Result>>, AdminError> where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - let stream = state.system_admin_service().update_events().map(|message| { - let data = serde_json::to_string(&message.data).unwrap_or_else(|_| "{}".to_owned()); - Ok(Event::default().event("update").id(message.id).data(data)) - }); + let stream = state + .admin_services() + .system() + .update_events() + .map(|message| { + let id = message.id.clone(); + let data = SystemUpdateEventView::from(message).into_json().to_string(); + Ok(Event::default().event("update").id(id).data(data)) + }); Ok(Sse::new(stream).keep_alive(KeepAlive::default())) } @@ -176,14 +291,34 @@ async fn perform_update( payload: Option>, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response( - state - .system_admin_service() - .perform_update(payload.map(|Json(value)| value.into_target_version())) - .await, - ) + let result = state + .admin_services() + .system() + .perform_update(payload.map(|Json(value)| value.into_target_version())) + .await + .map_err(map_system_error)?; + let SystemOperationAccepted::Update { + operation_id, + deployment_mode, + message, + need_restart, + target_version, + } = result + else { + return Err(AdminError::internal("Invalid system update result")); + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(UpdateAcceptedView { + operation_id, + deployment_mode, + message, + need_restart, + target_version, + }), + )) } async fn update_status( @@ -191,9 +326,18 @@ async fn update_status( State(state): State, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response(state.system_admin_service().update_status().await) + let status = state + .admin_services() + .system() + .update_status() + .await + .map_err(map_system_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(SystemUpdateStatusView::from(status)), + )) } async fn rollback( @@ -201,9 +345,30 @@ async fn rollback( State(state): State, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response(state.system_admin_service().rollback().await) + let result = state + .admin_services() + .system() + .rollback() + .await + .map_err(map_system_error)?; + let SystemOperationAccepted::Rollback { + operation_id, + message, + need_restart, + } = result + else { + return Err(AdminError::internal("Invalid system rollback result")); + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RollbackAcceptedView { + message, + need_restart, + operation_id, + }), + )) } async fn restart( @@ -211,23 +376,100 @@ async fn restart( State(state): State, ) -> Result where - S: SystemAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - value_response(state.system_admin_service().restart().await) + let result = state + .admin_services() + .system() + .restart() + .await + .map_err(map_system_error)?; + let SystemOperationAccepted::Restart { + operation_id, + message, + } = result + else { + return Err(AdminError::internal("Invalid system restart result")); + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RestartAcceptedView { + message, + operation_id, + }), + )) } -fn value_response( - result: Result, -) -> Result>, AdminError> { - let data = result.map_err(map_system_error)?; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) -} - -fn map_system_error(error: SystemAdminError) -> AdminError { - match error.kind() { - SystemAdminErrorKind::Invalid => AdminError::bad_request(error.to_string()), - SystemAdminErrorKind::Conflict => AdminError::conflict(error.to_string()), - SystemAdminErrorKind::BadGateway => AdminError::bad_gateway(error.to_string()), - SystemAdminErrorKind::Internal => AdminError::internal(error.to_string()), +impl From for SystemUpdateEventView { + fn from(event: SystemUpdateEvent) -> Self { + Self { + operation_id: event.operation_id, + level: update_event_level_name(event.level), + step: event.step, + message: event.message, + terminal: event.terminal, + at: event.occurred_at.to_rfc3339(), + } } } + +impl SystemUpdateEventView { + fn into_json(self) -> serde_json::Value { + serde_json::json!({ + "operationId": self.operation_id, + "level": self.level, + "step": self.step, + "message": self.message, + "terminal": self.terminal, + "at": self.at, + }) + } +} + +fn deployment_mode_label(mode: &str) -> &'static str { + match mode { + "docker" => "Docker", + "binary" => "二进制", + _ => "源码运行", + } +} + +fn build_type_label(build_type: &str) -> String { + match build_type { + "release" => "正式构建", + "source" => "源码构建", + "dev" => "开发构建", + value => value, + } + .to_owned() +} + +const fn operation_kind_name(kind: SystemOperationKind) -> &'static str { + match kind { + SystemOperationKind::Update => "update", + SystemOperationKind::Rollback => "rollback", + SystemOperationKind::Restart => "restart", + } +} + +const fn operation_status_name(status: SystemOperationStatus) -> &'static str { + match status { + SystemOperationStatus::Idle => "idle", + SystemOperationStatus::Running => "running", + SystemOperationStatus::Succeeded => "succeeded", + SystemOperationStatus::Failed => "failed", + } +} + +const fn update_event_level_name(level: SystemUpdateEventLevel) -> &'static str { + match level { + SystemUpdateEventLevel::Info => "info", + SystemUpdateEventLevel::Warning => "warning", + SystemUpdateEventLevel::Success => "success", + SystemUpdateEventLevel::Error => "error", + } +} + +fn map_system_error(error: gateway_admin::model::AdminError) -> AdminError { + map_admin_service_error(error, "System operation unavailable") +} diff --git a/backend/crates/gateway-api/src/admin/wire.rs b/backend/crates/gateway-api/src/admin/wire.rs index f3c651a36..74e165bb4 100644 --- a/backend/crates/gateway-api/src/admin/wire.rs +++ b/backend/crates/gateway-api/src/admin/wire.rs @@ -246,6 +246,25 @@ impl AdminError { } } +/// 把管理用例的稳定错误分类映射到既有 HTTP 错误 contract。 +pub(crate) fn map_admin_service_error( + error: gateway_admin::model::AdminError, + unavailable_message: &'static str, +) -> AdminError { + use gateway_admin::model::AdminErrorKind; + + match error.kind() { + AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()), + AdminErrorKind::Unauthorized => AdminError::admin_session_required(), + AdminErrorKind::NotFound => AdminError::not_found(error.to_string()), + AdminErrorKind::Conflict => AdminError::conflict(error.to_string()), + AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(), + AdminErrorKind::BadGateway => AdminError::bad_gateway(error.to_string()), + AdminErrorKind::Unavailable => AdminError::service_unavailable(unavailable_message), + AdminErrorKind::Internal => AdminError::internal(error.to_string()), + } +} + impl IntoResponse for AdminError { fn into_response(self) -> Response { (self.status, Json(self.body)).into_response() diff --git a/backend/crates/gateway-api/src/admin/xai.rs b/backend/crates/gateway-api/src/admin/xai.rs index cbafdad2a..55a361672 100644 --- a/backend/crates/gateway-api/src/admin/xai.rs +++ b/backend/crates/gateway-api/src/admin/xai.rs @@ -1,6 +1,5 @@ //! xAI 管理端 wire contract。 -use async_trait::async_trait; use axum::{ Json, Router, extract::{Query, State}, @@ -9,12 +8,24 @@ use axum::{ routing::{get, post}, }; use chrono::{DateTime, Utc}; +use gateway_admin::model::{ + AdminError as AdminServiceError, AdminErrorKind, Revision, + accounts::{AccountAvailability, AccountRecord}, + provider_credentials::{ + AuthorizationStarted, CompleteAuthorization, CredentialImportResult, CredentialListQuery, + CredentialListWindow, CredentialMutation, CredentialMutationResult, CredentialPage, + ImportCredentials, ProviderDocument, StartAuthorization, + }, +}; +use gateway_core::{ + engine::credential::{OpaqueProviderData, ProviderAccountId}, + routing::ProviderInstanceId, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::admin::{ - AdminAuth, AdminEnvelope, AdminError, AdminRequestContext, AdminResponse, AdminServiceError, - AdminServiceErrorKind, AdminSessionState, WireValidationError, + AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError, }; const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1024 * 1024; @@ -41,6 +52,23 @@ impl XaiCredentialImportDocumentRequest { } Ok(()) } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(ImportCredentials { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + provider_instance_id: ProviderInstanceId::new(self.provider_instance_id) + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + document: provider_document(self.document, "document")?, + }) + } } /// 批量导入成功响应；不包含任何 token、subject、邮箱或 client ID。 @@ -80,6 +108,20 @@ impl ListCredentialsQuery { } Ok(()) } + + fn into_command(self) -> Result { + self.validate()?; + Ok(CredentialListQuery { + provider_instance_id: self + .provider_instance_id + .map(ProviderInstanceId::new) + .transpose() + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + availability: None, + enabled: None, + window: CredentialListWindow::All, + }) + } } #[derive(Debug, Clone, Deserialize)] @@ -110,6 +152,24 @@ impl StartAuthorizationRequest { pub fn validate(&self) -> Result<(), WireValidationError> { self.credential.validate() } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(StartAuthorization { + context, + expected_config_revision: revision( + self.credential.expected_config_revision, + "expectedConfigRevision", + )?, + provider_instance_id: ProviderInstanceId::new(self.credential.provider_instance_id) + .map_err(|_| WireValidationError::new("providerInstanceId"))?, + name: self.credential.name, + reauthorization: None, + }) + } } #[derive(Deserialize)] @@ -124,6 +184,18 @@ impl CompleteAuthorizationRequest { require_id(&self.flow_id, "flowId")?; require_text(&self.callback_url, MAX_CALLBACK_FIELD_BYTES, "callbackUrl") } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(CompleteAuthorization { + context, + flow_id: self.flow_id, + callback_url: self.callback_url, + }) + } } #[derive(Debug, Clone, Deserialize)] @@ -138,6 +210,22 @@ impl CredentialMutationRequest { require_id(&self.credential_id, "credentialId")?; require_positive(self.expected_config_revision, "expectedConfigRevision") } + + fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result { + self.validate()?; + Ok(CredentialMutation { + context, + expected_config_revision: revision( + self.expected_config_revision, + "expectedConfigRevision", + )?, + account_id: ProviderAccountId::new(self.credential_id) + .map_err(|_| WireValidationError::new("credentialId"))?, + }) + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -148,15 +236,6 @@ pub struct AuthorizationStartData { pub expires_at: DateTime, } -/// Authorization Code flow 启动的应用端结果。 -/// -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AuthorizationStartResult { - pub flow_id: String, - pub authorization_url: String, - pub expires_at: DateTime, -} - /// xAI OAuth credential 的安全管理视图。 #[derive(Debug, Clone, PartialEq, Eq, Serialize)] #[serde(rename_all = "camelCase")] @@ -196,54 +275,76 @@ pub struct XaiCredentialMutationData { pub credential_id: String, } -/// xAI OAuth credential 管理应用端口。 -#[async_trait] -pub trait XaiAdminService: Send + Sync { - async fn list( - &self, - query: ListCredentialsQuery, - ) -> Result; - async fn import_document( - &self, - context: &AdminRequestContext, - request: XaiCredentialImportDocumentRequest, - ) -> Result; - async fn start_authorization( - &self, - context: &AdminRequestContext, - request: StartAuthorizationRequest, - ) -> Result; - async fn complete_authorization( - &self, - context: &AdminRequestContext, - request: CompleteAuthorizationRequest, - ) -> Result; - async fn disable( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; - async fn enable( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; - async fn delete( - &self, - context: &AdminRequestContext, - request: CredentialMutationRequest, - ) -> Result; +impl From for XaiCredentialViewData { + fn from(account: AccountRecord) -> Self { + Self { + id: account.id, + provider_instance_id: account.provider_instance_id.to_string(), + name: account.name, + email: account.email, + upstream_user_id: account.upstream_user_id, + upstream_account_id: account.upstream_account_id, + plan_type: account.plan_type, + enabled: account.enabled, + credential_revision: i64::try_from(account.credential_revision.get()) + .unwrap_or(i64::MAX), + has_refresh_token: account.has_refresh_token, + availability: availability_name(account.availability).to_owned(), + availability_reason: account.availability_reason, + access_token_expires_at: account.access_token_expires_at, + next_refresh_at: account.next_refresh_at, + cooldown_until: account.cooldown_until, + created_at: account.created_at, + updated_at: account.updated_at, + } + } } -/// xAI HTTP module 所需最小 state。 -pub trait XaiAdminState: AdminSessionState { - fn xai_admin_service(&self) -> &dyn XaiAdminService; +impl From for XaiCredentialListData { + fn from(page: CredentialPage) -> Self { + Self { + config_revision: page.config_revision.get(), + items: page.items.into_iter().map(Into::into).collect(), + } + } +} + +impl From for XaiCredentialImportData { + fn from(result: CredentialImportResult) -> Self { + Self::new( + result.config_revision.get(), + result + .credential_ids + .into_iter() + .map(|id| id.to_string()) + .collect(), + ) + } +} + +impl From for AuthorizationStartData { + fn from(started: AuthorizationStarted) -> Self { + Self { + flow_id: started.flow_id, + authorization_url: started.authorization_url, + expires_at: started.expires_at, + } + } +} + +impl From for XaiCredentialMutationData { + fn from(result: CredentialMutationResult) -> Self { + Self { + config_revision: result.config_revision.get(), + credential_id: result.account_id.to_string(), + } + } } /// 构造固定 GET/POST xAI OAuth 管理路由。 pub fn router() -> Router where - S: XaiAdminState + Clone + Send + Sync + 'static, + S: AdminSessionState + Clone + Send + Sync + 'static, { Router::new() .route("/api/admin/xai/credentials", get(list_credentials::)) @@ -279,95 +380,145 @@ async fn list_credentials( Query(query): Query, ) -> Result where - S: XaiAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - query.validate().map_err(map_wire_error)?; - ok(StatusCode::OK, state.xai_admin_service().list(query).await) + let command = query.into_command().map_err(map_wire_error)?; + let page = state + .admin_services() + .xai() + .list(command) + .await + .map_err(map_service_error)?; + ok(StatusCode::OK, XaiCredentialListData::from(page)) } -macro_rules! xai_post_handler { - ($name:ident, $request:ty, $method:ident, $status:expr) => { - async fn $name( - auth: AdminAuth, - State(state): State, - Json(request): Json<$request>, - ) -> Result - where - S: XaiAdminState + Send + Sync, - { - request.validate().map_err(map_wire_error)?; - ok( - $status, - state - .xai_admin_service() - .$method(auth.context(), request) - .await, - ) - } - }; +async fn import_credentials_document( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .xai() + .import_document(command) + .await + .map_err(map_service_error)?; + ok(StatusCode::CREATED, XaiCredentialImportData::from(result)) } -xai_post_handler!( - import_credentials_document, - XaiCredentialImportDocumentRequest, - import_document, - StatusCode::CREATED -); async fn start_authorization_code_flow( auth: AdminAuth, State(state): State, Json(request): Json, ) -> Result where - S: XaiAdminState + Send + Sync, + S: AdminSessionState + Send + Sync, { - request.validate().map_err(map_wire_error)?; + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; let started = state - .xai_admin_service() - .start_authorization(auth.context(), request) + .admin_services() + .xai() + .start_authorization(command) .await .map_err(map_service_error)?; require_id(&started.flow_id, "flowId").map_err(map_wire_error)?; - let data = AuthorizationStartData { - authorization_url: started.authorization_url, - flow_id: started.flow_id, - expires_at: started.expires_at, - }; - Ok(AdminResponse::new( - StatusCode::CREATED, - AdminEnvelope::ok(data), - )) + ok(StatusCode::CREATED, AuthorizationStartData::from(started)) +} + +async fn complete_authorization_code_flow( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = state + .admin_services() + .xai() + .complete_authorization(command) + .await + .map_err(map_service_error)?; + ok(StatusCode::CREATED, XaiCredentialMutationData::from(result)) +} + +async fn disable_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Disable).await +} + +async fn enable_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Enable).await +} + +async fn delete_credential( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + mutate_credential(auth, state, request, CredentialAction::Delete).await +} + +#[derive(Clone, Copy)] +enum CredentialAction { + Disable, + Enable, + Delete, +} + +async fn mutate_credential( + auth: AdminAuth, + state: S, + request: CredentialMutationRequest, + action: CredentialAction, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let service = state.admin_services().xai(); + let result = match action { + CredentialAction::Disable => service.disable(command).await, + CredentialAction::Enable => service.enable(command).await, + CredentialAction::Delete => service.delete(command).await, + } + .map_err(map_service_error)?; + ok(StatusCode::OK, XaiCredentialMutationData::from(result)) } -xai_post_handler!( - complete_authorization_code_flow, - CompleteAuthorizationRequest, - complete_authorization, - StatusCode::CREATED -); -xai_post_handler!( - disable_credential, - CredentialMutationRequest, - disable, - StatusCode::OK -); -xai_post_handler!( - enable_credential, - CredentialMutationRequest, - enable, - StatusCode::OK -); -xai_post_handler!( - delete_credential, - CredentialMutationRequest, - delete, - StatusCode::OK -); fn ok( status: StatusCode, - result: Result, + data: T, ) -> Result>, AdminError> { - let data = result.map_err(map_service_error)?; Ok(AdminResponse::new(status, AdminEnvelope::ok(data))) } @@ -377,11 +528,41 @@ fn map_wire_error(error: WireValidationError) -> AdminError { fn map_service_error(error: AdminServiceError) -> AdminError { match error.kind() { - AdminServiceErrorKind::Invalid => AdminError::bad_request(error.to_string()), - AdminServiceErrorKind::NotFound => AdminError::not_found(error.to_string()), - AdminServiceErrorKind::Conflict => AdminError::conflict(error.to_string()), - AdminServiceErrorKind::Unavailable => AdminError::bad_gateway(error.to_string()), - AdminServiceErrorKind::Internal => AdminError::internal(error.to_string()), + AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()), + AdminErrorKind::Unauthorized => AdminError::admin_session_required(), + AdminErrorKind::NotFound => AdminError::not_found(error.to_string()), + AdminErrorKind::Conflict => AdminError::conflict(error.to_string()), + AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(), + AdminErrorKind::BadGateway | AdminErrorKind::Unavailable => { + AdminError::bad_gateway(error.to_string()) + } + AdminErrorKind::Internal => AdminError::internal(error.to_string()), + } +} + +fn revision(value: u64, field: &'static str) -> Result { + Revision::new(value).map_err(|_| WireValidationError::new(field)) +} + +fn provider_document( + value: Value, + field: &'static str, +) -> Result { + match value { + Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))), + _ => Err(WireValidationError::new(field)), + } +} + +const fn availability_name(value: AccountAvailability) -> &'static str { + match value { + AccountAvailability::Unknown => "unknown", + AccountAvailability::Ready => "ready", + AccountAvailability::Cooldown => "cooldown", + AccountAvailability::QuotaExhausted => "quota_exhausted", + AccountAvailability::Expired => "expired", + AccountAvailability::Banned => "banned", + AccountAvailability::Invalid => "invalid", } } diff --git a/backend/crates/gateway-api/src/health.rs b/backend/crates/gateway-api/src/health.rs new file mode 100644 index 000000000..ace024e0e --- /dev/null +++ b/backend/crates/gateway-api/src/health.rs @@ -0,0 +1,68 @@ +//! `/healthz` 对 Core、Store 与 Host worker 健康事实的聚合。 + +use std::sync::Arc; +use std::time::Duration; + +use axum::{extract::State, http::StatusCode}; +use futures::future::join_all; +use gateway_core::health::{ + HealthProbe, HealthState, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState, +}; + +use crate::ApiState; + +const HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(2); + +#[derive(Clone)] +pub(crate) struct HealthStatus { + probes: Arc<[Arc]>, + workers: Arc, +} + +impl HealthStatus { + #[must_use] + pub(crate) fn new( + probes: Vec>, + workers: Arc, + ) -> Self { + Self { + probes: probes.into(), + workers, + } + } + + pub(crate) async fn healthy(&self) -> bool { + if !self.workers.snapshot().iter().all(worker_is_healthy) { + return false; + } + let checks = self.probes.iter().map(|probe| probe.check()); + tokio::time::timeout(HEALTH_CHECK_TIMEOUT, join_all(checks)) + .await + .is_ok_and(|states| { + states + .into_iter() + .all(|state| matches!(state, HealthState::Healthy)) + }) + } +} + +pub(crate) async fn healthz(State(state): State) -> StatusCode { + if state.health().healthy().await { + StatusCode::NO_CONTENT + } else { + StatusCode::SERVICE_UNAVAILABLE + } +} + +fn worker_is_healthy(worker: &WorkerHealthSnapshot) -> bool { + match worker.state { + WorkerRuntimeState::Disabled | WorkerRuntimeState::Standby => true, + WorkerRuntimeState::Running => worker.consecutive_failures == 0, + WorkerRuntimeState::AcquiringLease | WorkerRuntimeState::Idle => { + worker.consecutive_failures == 0 && worker.last_success_at.is_some() + } + WorkerRuntimeState::Starting + | WorkerRuntimeState::BackingOff + | WorkerRuntimeState::Stopped => false, + } +} diff --git a/backend/crates/gateway-api/src/lib.rs b/backend/crates/gateway-api/src/lib.rs index 6478ebe5a..4a32026d6 100644 --- a/backend/crates/gateway-api/src/lib.rs +++ b/backend/crates/gateway-api/src/lib.rs @@ -1,9 +1,218 @@ -//! 客户端协议 adapter。 +//! 客户端与管理端 HTTP 协议 adapter。 //! -//! 本 crate 只负责把外部 wire contract 转换为 [`gateway_core`] 的稳定 -//! operation/event contract。它不执行路由、Provider 调用、重试或持久化。 +//! 本 crate 只负责请求解码、Core/Admin 调用和 HTTP/WS/SSE delivery。 -#![forbid(unsafe_code)] +use std::collections::BTreeSet; +use std::path::{Path, PathBuf}; +use std::str::FromStr as _; +use std::sync::Arc; +use std::time::Duration; + +use axum::Router; +use axum::http::{HeaderName, HeaderValue, Method, StatusCode}; +use axum::routing::get; +use gateway_admin::AdminServices; +use gateway_core::engine::execution::ExecutionService; +use gateway_core::health::{HealthProbe, WorkerHealthSource}; +use gateway_core::lifecycle::ConnectionLifecycle; +use serde::Deserialize; +use tower_http::cors::{Any, CorsLayer}; +use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer}; +use tower_http::services::{ServeDir, ServeFile}; +use tower_http::timeout::TimeoutLayer; +use tower_http::trace::TraceLayer; +use url::Url; + +const WEB_DIST_ENV: &str = "CPR_WEB_DIST_DIR"; + +use crate::health::HealthStatus; +use crate::openai::service::OpenAiService; pub mod admin; +mod health; pub mod openai; + +/// API-owned HTTP 与静态资源配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct ApiConfig { + pub asset_directory: PathBuf, + pub cors_allowed_origins: Vec, + pub request_timeout_seconds: Option, + pub request_id_header: String, +} + +impl ApiConfig { + /// 解析静态资源相对路径并校验全部 HTTP 配置。 + /// + /// # Errors + /// + /// 路径为空、origin/header 非法或 timeout 为零时返回脱敏错误。 + pub fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ApiConfigError> { + match std::env::var(WEB_DIST_ENV) { + Ok(value) if value.trim().is_empty() => { + return Err(ApiConfigError::InvalidAssetDirectory); + } + Ok(value) => self.asset_directory = PathBuf::from(value), + Err(std::env::VarError::NotPresent) => {} + Err(std::env::VarError::NotUnicode(_)) => { + return Err(ApiConfigError::InvalidAssetDirectory); + } + } + if self.asset_directory.as_os_str().is_empty() { + return Err(ApiConfigError::InvalidAssetDirectory); + } + if self.asset_directory.is_relative() { + self.asset_directory = source_dir.join(&self.asset_directory); + } + if self.request_timeout_seconds == Some(0) { + return Err(ApiConfigError::InvalidRequestTimeout); + } + let header = HeaderName::from_str(&self.request_id_header) + .map_err(|_| ApiConfigError::InvalidRequestIdHeader)?; + self.request_id_header = header.as_str().to_owned(); + + let mut origins = Vec::with_capacity(self.cors_allowed_origins.len()); + let mut unique = BTreeSet::new(); + for origin in &self.cors_allowed_origins { + let origin = validate_origin(origin)?; + if !unique.insert(origin.clone()) { + return Err(ApiConfigError::DuplicateCorsOrigin); + } + origins.push(origin); + } + self.cors_allowed_origins = origins; + Ok(()) + } +} + +/// API 配置非法的稳定分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum ApiConfigError { + #[error("API asset directory is invalid")] + InvalidAssetDirectory, + #[error("API CORS origin is invalid")] + InvalidCorsOrigin, + #[error("API CORS origin is duplicated")] + DuplicateCorsOrigin, + #[error("API request timeout is invalid")] + InvalidRequestTimeout, + #[error("API request ID header is invalid")] + InvalidRequestIdHeader, +} + +/// 完成组装的唯一 API router。 +pub struct ApiBundle { + router: Router, +} + +impl ApiBundle { + pub fn router(self) -> Router { + self.router + } +} + +/// 组装客户端、管理端、健康检查和静态资源路由。 +pub fn initialize( + mut config: ApiConfig, + execution: Arc, + admin: AdminServices, + probes: Vec>, + worker_health: Arc, + lifecycle: Arc, +) -> Result { + config + .resolve_and_validate(Path::new(".")) + .map_err(ApiError::Config)?; + let request_id_header = HeaderName::from_str(&config.request_id_header) + .map_err(|_| ApiError::Config(ApiConfigError::InvalidRequestIdHeader))?; + let state = ApiState { + admin, + openai: OpenAiService::new(execution, lifecycle), + health: HealthStatus::new(probes, worker_health), + }; + let index = config.asset_directory.join("index.html"); + let mut router = Router::new() + .route("/healthz", get(health::healthz)) + .merge(openai::router::router()) + .merge(admin::router::()) + .fallback_service(ServeDir::new(config.asset_directory).fallback(ServeFile::new(index))); + if !config.cors_allowed_origins.is_empty() { + let origins = config + .cors_allowed_origins + .iter() + .map(|origin| { + HeaderValue::from_str(origin) + .map_err(|_| ApiError::Config(ApiConfigError::InvalidCorsOrigin)) + }) + .collect::, _>>()?; + router = router.layer( + CorsLayer::new() + .allow_origin(origins) + .allow_methods([Method::GET, Method::POST]) + .allow_headers(Any) + .allow_credentials(true), + ); + } + if let Some(seconds) = config.request_timeout_seconds { + router = router.layer(TimeoutLayer::with_status_code( + StatusCode::REQUEST_TIMEOUT, + Duration::from_secs(seconds), + )); + } + let router = router + .layer(PropagateRequestIdLayer::new(request_id_header.clone())) + .layer(SetRequestIdLayer::new(request_id_header, MakeRequestUuid)) + .layer(TraceLayer::new_for_http()) + .with_state(state); + Ok(ApiBundle { router }) +} + +/// API 初始化失败的脱敏分类。 +#[derive(Debug, thiserror::Error)] +pub enum ApiError { + #[error(transparent)] + Config(ApiConfigError), +} + +#[derive(Clone)] +pub(crate) struct ApiState { + admin: AdminServices, + openai: OpenAiService, + health: HealthStatus, +} + +impl ApiState { + #[must_use] + pub(crate) const fn openai(&self) -> &OpenAiService { + &self.openai + } + + #[must_use] + pub(crate) const fn health(&self) -> &HealthStatus { + &self.health + } +} + +impl admin::AdminSessionState for ApiState { + fn admin_services(&self) -> &AdminServices { + &self.admin + } +} + +fn validate_origin(raw: &str) -> Result { + let url = Url::parse(raw).map_err(|_| ApiConfigError::InvalidCorsOrigin)?; + if !matches!(url.scheme(), "http" | "https") + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.path() != "/" + || url.query().is_some() + || url.fragment().is_some() + { + return Err(ApiConfigError::InvalidCorsOrigin); + } + let origin = url.origin().ascii_serialization(); + HeaderValue::from_str(&origin).map_err(|_| ApiConfigError::InvalidCorsOrigin)?; + Ok(origin) +} diff --git a/backend/crates/gateway-api/src/openai/auth.rs b/backend/crates/gateway-api/src/openai/auth.rs index 094018a33..af2b57c55 100644 --- a/backend/crates/gateway-api/src/openai/auth.rs +++ b/backend/crates/gateway-api/src/openai/auth.rs @@ -4,10 +4,11 @@ use axum::{ http::{HeaderMap, header::AUTHORIZATION}, response::{IntoResponse, Response}, }; +use gateway_core::engine::execution::AuthenticatedClient; use super::{ error::{missing_client_api_key_response, runtime_unavailable_response}, - service::OpenAiClientService, + service::OpenAiService, }; /// Client API key 鉴权失败原因。 @@ -63,13 +64,10 @@ pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAu Ok(token) } -pub(crate) fn authenticate_client( - service: &S, +pub(crate) fn authenticate_client( + service: &OpenAiService, headers: &HeaderMap, -) -> Result -where - S: OpenAiClientService, -{ +) -> Result { bearer_client_api_key(headers).and_then(|key| service.authenticate(key)) } diff --git a/backend/crates/gateway-api/src/openai/mod.rs b/backend/crates/gateway-api/src/openai/mod.rs index 4b0317c42..9459ca165 100644 --- a/backend/crates/gateway-api/src/openai/mod.rs +++ b/backend/crates/gateway-api/src/openai/mod.rs @@ -5,10 +5,4 @@ pub mod error; pub mod models; pub mod responses; pub mod router; -pub mod service; - -pub use router::router; -pub use service::{ - ConnectionTask, DeliveryEvent, OpenAiApiState, OpenAiClientService, ResponseExecutionSession, - ResponsesTransport, StartedResponse, -}; +pub(crate) mod service; diff --git a/backend/crates/gateway-api/src/openai/models.rs b/backend/crates/gateway-api/src/openai/models.rs index 33384aaed..4460e9890 100644 --- a/backend/crates/gateway-api/src/openai/models.rs +++ b/backend/crates/gateway-api/src/openai/models.rs @@ -9,21 +9,19 @@ use axum::{ use gateway_core::routing::PublicModelId; use serde_json::{Value, json}; +use crate::ApiState; + use super::{ auth::{authenticate_client, authentication_error_response}, error::model_not_found_response, - service::{OpenAiApiState, OpenAiClientService}, }; const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000; /// `GET /v1/models`。 -pub async fn models(State(state): State, headers: HeaderMap) -> Response -where - S: OpenAiApiState, -{ - let service = state.openai_client_api(); - let client = match authenticate_client(&service, &headers) { +pub(crate) async fn models(State(state): State, headers: HeaderMap) -> Response { + let service = state.openai(); + let client = match authenticate_client(service, &headers) { Ok(client) => client, Err(error) => return authentication_error_response(error), }; @@ -44,16 +42,13 @@ where } /// `GET /v1/models/{model_id}`。 -pub async fn model_detail( - State(state): State, +pub(crate) async fn model_detail( + State(state): State, headers: HeaderMap, Path(model_id): Path, -) -> Response -where - S: OpenAiApiState, -{ - let service = state.openai_client_api(); - let client = match authenticate_client(&service, &headers) { +) -> Response { + let service = state.openai(); + let client = match authenticate_client(service, &headers) { Ok(client) => client, Err(error) => return authentication_error_response(error), }; diff --git a/backend/crates/gateway-api/src/openai/responses/error.rs b/backend/crates/gateway-api/src/openai/responses/error.rs index ef7f066b0..6bdabaaa2 100644 --- a/backend/crates/gateway-api/src/openai/responses/error.rs +++ b/backend/crates/gateway-api/src/openai/responses/error.rs @@ -197,6 +197,15 @@ pub enum ResponseEncodeError { /// Adapter 自身生成的 event framing 无法还原为 WebSocket event JSON。 #[error("responses event encoder produced an invalid internal framing")] InvalidEventEncoding, + /// 同一响应混用了协议原生与 canonical 表达。 + #[error("responses stream changed its wire representation")] + MixedWireRepresentation, + /// 协议原生流缺少可返回的终态 response object。 + #[error("responses wire stream has no terminal response")] + MissingWireTerminal, + /// 协议原生事件携带了不一致的响应身份。 + #[error("responses wire stream changed its response identity")] + WireIdentityChanged, } impl ResponseEncodeError { diff --git a/backend/crates/gateway-api/src/openai/responses/http.rs b/backend/crates/gateway-api/src/openai/responses/http.rs index 566037a41..533b77d41 100644 --- a/backend/crates/gateway-api/src/openai/responses/http.rs +++ b/backend/crates/gateway-api/src/openai/responses/http.rs @@ -11,41 +11,68 @@ use axum::{ HeaderMap, HeaderValue, StatusCode, header::{CACHE_CONTROL, CONTENT_TYPE, USER_AGENT}, }, - response::Response, + response::{IntoResponse, Response}, }; use futures::stream; use gateway_core::engine::CommitRequirement; +use gateway_core::engine::execution::{ClientTransport, ExecutionSession, StartedExecution}; use gateway_core::error::{GatewayError, GatewayErrorKind}; -use gateway_core::event::GatewayEvent; +use gateway_core::event::{ProviderEvent, ProviderResponseHeader}; +use gateway_core::lifecycle::ConnectionGuard; use gateway_protocol::openai::sse::{DONE_SSE_FRAME, response_failed_sse_event}; +use crate::ApiState; use crate::openai::{ auth::{authenticate_client, authentication_error_response}, error::{ gateway_error_contract, gateway_error_from_engine, gateway_error_response, - protocol_error_response, + protocol_error_response, runtime_unavailable_response, }, - service::{OpenAiApiState, OpenAiClientService, ResponseExecutionSession, ResponsesTransport}, + service::created_at_unix_seconds, }; -use super::{ProtocolErrorBody, ResponseEncodeError, ResponsesCollector, decode_request}; +use super::{ + OpenAiResponsesEncoder, ProtocolErrorBody, ResponseEncodeError, + request::{decode_request_with_headers, decode_review_request_with_headers}, +}; /// `POST /v1/responses`。 -pub async fn responses( - State(state): State, +pub(crate) async fn responses( + State(state): State, connect_info: Option>>, headers: HeaderMap, body: Bytes, -) -> Response -where - S: OpenAiApiState, -{ - let service = state.openai_client_api(); - let client = match authenticate_client(&service, &headers) { +) -> Response { + handle_responses(state, connect_info, headers, body, false).await +} + +/// `POST /v1/responses/review`。 +pub(crate) async fn review_responses( + State(state): State, + connect_info: Option>>, + headers: HeaderMap, + body: Bytes, +) -> Response { + handle_responses(state, connect_info, headers, body, true).await +} + +async fn handle_responses( + state: ApiState, + connect_info: Option>>, + headers: HeaderMap, + body: Bytes, + review: bool, +) -> Response { + let service = state.openai(); + let client = match authenticate_client(service, &headers) { Ok(client) => client, Err(error) => return authentication_error_response(error), }; - let decoded = match decode_request(&body) { + let decoded = match if review { + decode_review_request_with_headers(&body, &headers) + } else { + decode_request_with_headers(&body, &headers) + } { Ok(decoded) => decoded, Err(error) => { return protocol_error_response(StatusCode::BAD_REQUEST, error.protocol_body()); @@ -56,17 +83,53 @@ where connect_info.map(|Extension(ConnectInfo(address))| address), ); let decoded = decoded.with_client_context(client_ip, user_agent); + let streaming = decoded.metadata().stream(); + let connection_guard = if streaming { + match service.try_register_connection() { + Ok(guard) => Some(guard), + Err(_) => return runtime_unavailable_response().into_response(), + } + } else { + None + }; let started = match service - .start_response(client, decoded, ResponsesTransport::Http) + .start_response( + client, + decoded, + if streaming { + ClientTransport::HttpSse + } else { + ClientTransport::HttpJson + }, + if review { + "/v1/responses/review" + } else { + "/v1/responses" + }, + ) .await { Ok(started) => started, Err(error) => return gateway_error_response(&error), }; - let (session, created_at, streaming) = started.into_parts(); - if streaming { - stream_execution_response(session, created_at).await + let StartedExecution { + created_at, + stream, + session, + .. + } = started; + let created_at = match created_at_unix_seconds(created_at) { + Ok(created_at) => created_at, + Err(error) => { + let mut execution = PendingExecution::new(session); + let response = gateway_error_response(&error); + return execution.record_response_status(response).await; + } + }; + if stream { + stream_execution_response(session, created_at, connection_guard).await } else { + drop(connection_guard); collect_execution_response(session, created_at).await } } @@ -121,10 +184,10 @@ const fn is_private_or_loopback(address: IpAddr) -> bool { } /// 编码完整 canonical event 集合，并在完整 JSON 成功后提交下游。 -pub async fn collect_execution_response(session: S, created_at: u64) -> Response -where - S: ResponseExecutionSession, -{ +pub async fn collect_execution_response( + session: Box, + created_at: u64, +) -> Response { let mut execution = PendingExecution::new(session); let Some(session) = execution.session_mut() else { return internal_gateway_response("gateway response session is unavailable"); @@ -147,7 +210,8 @@ where return response; } }; - let response = json_body_response(encoded); + let response_headers = session.response_headers().to_vec(); + let response = json_body_response(encoded, &response_headers); let Some(session) = execution.session_mut() else { return internal_gateway_response("gateway response session is unavailable"); @@ -185,39 +249,58 @@ impl BufferedResponseEncodeError { fn encode_collected_events( created_at: u64, - events: &[GatewayEvent], + events: &[ProviderEvent], ) -> Result, BufferedResponseEncodeError> { - let mut collector = ResponsesCollector::new(created_at); + let mut encoder = OpenAiResponsesEncoder::new(created_at); for event in events { - collector - .push(event) + encoder + .push_sse(event) .map_err(BufferedResponseEncodeError::Canonical)?; } - let response = collector + let response = encoder .finish() .map_err(BufferedResponseEncodeError::Canonical)?; serde_json::to_vec(&response).map_err(|_| BufferedResponseEncodeError::Json) } -fn json_body_response(encoded: Vec) -> Response { +fn json_body_response(encoded: Vec, response_headers: &[ProviderResponseHeader]) -> Response { let mut response = Response::new(Body::from(encoded)); *response.status_mut() = StatusCode::OK; response .headers_mut() .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + apply_safe_response_headers(response, response_headers) +} + +fn apply_safe_response_headers( + mut response: Response, + response_headers: &[ProviderResponseHeader], +) -> Response { + for header in response_headers { + let Some(name) = super::safe_response_header_name(header.name()) else { + continue; + }; + let Ok(value) = HeaderValue::from_str(header.value().as_str()) else { + continue; + }; + response + .headers_mut() + .insert(axum::http::header::HeaderName::from_static(name), value); + } response } /// 编码首个 SSE frame 后提交下游，再持续驱动同一执行会话。 -pub async fn stream_execution_response(session: S, created_at: u64) -> Response -where - S: ResponseExecutionSession, -{ +pub async fn stream_execution_response( + session: Box, + created_at: u64, + connection_guard: Option>, +) -> Response { let mut execution = PendingExecution::new(session); let Some(session) = execution.session_mut() else { return internal_gateway_response("gateway response session is unavailable"); }; - let first = match session.next_delivery_event().await { + let first = match session.next_event().await { Ok(Some(event)) => event, Ok(None) => { let response = gateway_error_response(&GatewayError::new( @@ -233,26 +316,32 @@ where return execution.record_response_status(response).await; } }; - let (first_event, first_requirement) = first.into_parts(); + let first_requirement = first.commit_requirement(); + let first_events = first.into_provider_events(); if first_requirement != CommitRequirement::CommitBeforeDelivery { let response = internal_gateway_response("gateway first event did not require commit"); let response = execution.record_response_status(response).await; execution.cancel_and_finalize().await; return response; } - let mut collector = ResponsesCollector::new(created_at); - let frames = match collector.push(&first_event) { - Ok(frames) => frames, - Err(error) => { - let response = - protocol_error_response(StatusCode::INTERNAL_SERVER_ERROR, error.protocol_body()); - let response = execution.record_response_status(response).await; - execution.cancel_and_finalize().await; - return response; + let mut encoder = OpenAiResponsesEncoder::new(created_at); + let mut frames = Vec::new(); + for event in &first_events { + match encoder.push_sse(event) { + Ok(encoded) => frames.extend(encoded), + Err(error) => { + let response = protocol_error_response( + StatusCode::INTERNAL_SERVER_ERROR, + error.protocol_body(), + ); + let response = execution.record_response_status(response).await; + execution.cancel_and_finalize().await; + return response; + } } - }; + } if frames.is_empty() { - let response = internal_gateway_response("gateway first event encoded no SSE frame"); + let response = internal_gateway_response("gateway commit batch encoded no output"); let response = execution.record_response_status(response).await; execution.cancel_and_finalize().await; return response; @@ -260,6 +349,7 @@ where let Some(session) = execution.session_mut() else { return internal_gateway_response("gateway response session is unavailable"); }; + let response_headers = session.response_headers().to_vec(); if let Err(error) = session .commit_downstream(Some(StatusCode::OK.as_u16())) .await @@ -271,7 +361,10 @@ where let Some(session) = execution.into_session() else { return internal_gateway_response("gateway response session is unavailable"); }; - let state = ResponsesStreamState::new(session, collector, frames); + let mut state = ResponsesStreamState::new(session, encoder, frames, connection_guard); + if state.encoder.is_completed() { + state.finish_completed(Vec::new()).await; + } let body = Body::from_stream(stream::unfold(state, |mut state| async move { loop { if let Some(chunk) = state.pending.pop_front() { @@ -283,10 +376,10 @@ where state.advance().await; } })); - event_stream_response(body) + event_stream_response(body, &response_headers) } -fn event_stream_response(body: Body) -> Response { +fn event_stream_response(body: Body, response_headers: &[ProviderResponseHeader]) -> Response { let mut response = Response::new(body); *response.status_mut() = StatusCode::OK; response @@ -295,39 +388,33 @@ fn event_stream_response(body: Body) -> Response { response .headers_mut() .insert(CACHE_CONTROL, HeaderValue::from_static("no-cache")); - response + apply_safe_response_headers(response, response_headers) } fn internal_gateway_response(message: &'static str) -> Response { gateway_error_response(&GatewayError::new(GatewayErrorKind::Internal, message)) } -pub(super) struct PendingExecution -where - S: ResponseExecutionSession, -{ - session: Option, +pub(super) struct PendingExecution { + session: Option>, } -impl PendingExecution -where - S: ResponseExecutionSession, -{ - pub(super) fn new(session: S) -> Self { +impl PendingExecution { + pub(super) fn new(session: Box) -> Self { Self { session: Some(session), } } - pub(super) fn session_mut(&mut self) -> Option<&mut S> { - self.session.as_mut() + pub(super) fn session_mut(&mut self) -> Option<&mut (dyn ExecutionSession + 'static)> { + self.session.as_deref_mut() } async fn cancel_and_finalize(&mut self) { if let Some(session) = self.session.as_mut() { session.cancel(); if !session.is_finalized() { - let _ = session.next_delivery_event().await; + let _ = session.next_event().await; } } } @@ -350,15 +437,12 @@ where self.session = None; } - fn into_session(mut self) -> Option { + fn into_session(mut self) -> Option> { self.session.take() } } -impl Drop for PendingExecution -where - S: ResponseExecutionSession, -{ +impl Drop for PendingExecution { fn drop(&mut self) { let Some(session) = self.session.take() else { return; @@ -367,38 +451,39 @@ where return; } session.cancel(); - session.detach_finalize(); + detach_finalize(session); } } -struct ResponsesStreamState -where - S: ResponseExecutionSession, -{ - session: Option, - collector: ResponsesCollector, +struct ResponsesStreamState { + session: Option>, + encoder: OpenAiResponsesEncoder, pending: VecDeque, output_finished: bool, execution_terminal: bool, + _connection_guard: Option>, } -impl ResponsesStreamState -where - S: ResponseExecutionSession, -{ - fn new(session: S, collector: ResponsesCollector, initial_frames: Vec) -> Self { +impl ResponsesStreamState { + fn new( + session: Box, + encoder: OpenAiResponsesEncoder, + initial_frames: Vec, + connection_guard: Option>, + ) -> Self { Self { session: Some(session), - collector, + encoder, pending: initial_frames.into_iter().map(Bytes::from).collect(), output_finished: false, execution_terminal: false, + _connection_guard: connection_guard, } } async fn advance(&mut self) { let next = match self.session.as_mut() { - Some(session) => session.next_delivery_event().await, + Some(session) => session.next_event().await, None => { self.finish_with_gateway_error(GatewayError::new( GatewayErrorKind::Internal, @@ -409,7 +494,8 @@ where }; match next { Ok(Some(event)) => { - let (event, requirement) = event.into_parts(); + let requirement = event.commit_requirement(); + let events = event.into_provider_events(); if requirement != CommitRequirement::AlreadyCommitted { self.cancel_execution().await; self.finish_with_gateway_error(GatewayError::new( @@ -417,14 +503,19 @@ where "gateway requested another downstream commit", )); } else { - self.push_event(event).await; + for event in events { + self.push_event(event).await; + if self.output_finished { + break; + } + } } } Ok(None) => { self.execution_terminal = self .session .as_ref() - .is_some_and(ResponseExecutionSession::is_finalized); + .is_some_and(|session| session.is_finalized()); self.finish_with_gateway_error(GatewayError::new( GatewayErrorKind::Internal, "gateway response ended without a terminal event", @@ -434,17 +525,16 @@ where self.execution_terminal = self .session .as_ref() - .is_some_and(ResponseExecutionSession::is_finalized); + .is_some_and(|session| session.is_finalized()); self.finish_with_gateway_error(gateway_error_from_engine(&error)); } } } - async fn push_event(&mut self, event: GatewayEvent) { - let completed = matches!(event, GatewayEvent::Completed(_)); - match self.collector.push(&event) { + async fn push_event(&mut self, event: ProviderEvent) { + match self.encoder.push_sse(&event) { Ok(frames) => { - if completed { + if self.encoder.is_completed() { self.finish_completed(frames).await; } else { self.pending.extend(frames.into_iter().map(Bytes::from)); @@ -459,7 +549,7 @@ where async fn finish_completed(&mut self, frames: Vec) { let next = match self.session.as_mut() { - Some(session) => session.next_delivery_event().await, + Some(session) => session.next_event().await, None => { self.finish_with_gateway_error(GatewayError::new( GatewayErrorKind::Internal, @@ -473,7 +563,7 @@ where if self .session .as_ref() - .is_some_and(ResponseExecutionSession::is_finalized) => + .is_some_and(|session| session.is_finalized()) => { self.pending.extend(frames.into_iter().map(Bytes::from)); self.execution_terminal = true; @@ -499,7 +589,7 @@ where self.execution_terminal = self .session .as_ref() - .is_some_and(ResponseExecutionSession::is_finalized); + .is_some_and(|session| session.is_finalized()); self.finish_with_gateway_error(gateway_error_from_engine(&error)); } } @@ -509,7 +599,7 @@ where if let Some(session) = self.session.as_mut() { session.cancel(); if !session.is_finalized() { - let _ = session.next_delivery_event().await; + let _ = session.next_event().await; } self.execution_terminal = session.is_finalized(); } @@ -542,10 +632,7 @@ where } } -impl Drop for ResponsesStreamState -where - S: ResponseExecutionSession, -{ +impl Drop for ResponsesStreamState { fn drop(&mut self) { if self.execution_terminal { return; @@ -554,7 +641,14 @@ where return; }; session.cancel(); - // Body drop 后由应用 runtime 继续状态机清理，HTTP body 生命周期不等待它。 - session.detach_finalize(); + // Body drop 后由 Core 继续状态机清理，HTTP body 生命周期不等待它。 + detach_finalize(session); + } +} + +fn detach_finalize(session: Box) { + let finalize = session.detach_finalize(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + drop(runtime.spawn(finalize)); } } diff --git a/backend/crates/gateway-api/src/openai/responses/mod.rs b/backend/crates/gateway-api/src/openai/responses/mod.rs index 59dee035b..367c70354 100644 --- a/backend/crates/gateway-api/src/openai/responses/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/mod.rs @@ -1,7 +1,4 @@ -//! OpenAI Responses 的请求 decoder 与 canonical event encoder。 -//! -//! 这里的“透明”指已声明语义逐项映射，而不是原始 JSON/SSE passthrough。 -//! 未实现语义会返回稳定的 typed error，不能静默删除。 +//! OpenAI Responses 的透明 wire adapter 与 Core canonical facts 投影。 mod error; mod http; @@ -10,13 +7,23 @@ mod response; mod websocket; pub use error::{ProtocolErrorBody, RequestDecodeError, ResponseEncodeError}; -pub use http::{collect_execution_response, responses, stream_execution_response}; +pub use http::{collect_execution_response, stream_execution_response}; +pub(crate) use http::{responses, review_responses}; pub use request::{ ContinuationIntent, DecodedResponsesRequest, PROVIDER_OPTIONS_VERSION, ResponsesRequestMetadata, decode_request, }; -pub use response::{CollectedResponses, ResponsesCollector}; -pub use websocket::{ - ResponseCreateFrameError, ResponsesWebSocketAdapter, decode_response_create, - responses_websocket, -}; +pub use response::{CollectedResponses, OpenAiResponsesEncoder, ResponsesCollector}; +pub(crate) use websocket::responses_websocket; +pub use websocket::{ResponseCreateFrameError, decode_response_create}; + +pub(super) fn safe_response_header_name(name: &str) -> Option<&'static str> { + match name { + "x-request-id" => Some("x-request-id"), + "openai-model" => Some("openai-model"), + "x-models-etag" => Some("x-models-etag"), + "x-reasoning-included" => Some("x-reasoning-included"), + "openai-processing-ms" => Some("openai-processing-ms"), + _ => None, + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/request.rs b/backend/crates/gateway-api/src/openai/responses/request.rs index 8adb568b2..a6f33df9e 100644 --- a/backend/crates/gateway-api/src/openai/responses/request.rs +++ b/backend/crates/gateway-api/src/openai/responses/request.rs @@ -1,11 +1,16 @@ -//! OpenAI Responses JSON 到 canonical operation 的严格解码。 +//! OpenAI Responses JSON 到 Core 路由事实与不透明 wire payload 的单一解码边界。 use std::{fmt, net::IpAddr}; +use axum::http::HeaderMap; use gateway_core::operation::{ - ContentPart, ContinuationMode, GenerateRequest, ImageSource, JsonSchemaFormat, Message, - MessageRole, Operation, OutputFormat, ProviderOptions, ReasoningEffort, ReasoningRequirement, - ReasoningSummary, ResponsePersistence, ToolDefinition, + ContentPart, ContinuationMode, Feature, GenerateRequest, ImageSource, JsonSchemaFormat, + Message, MessageRole, Operation, OutputFormat, ProtocolPayload, ProviderOptions, + ProviderSessionState, ReasoningEffort, ReasoningRequirement, ReasoningSummary, + ResponsePersistence, ToolDefinition, +}; +use gateway_protocol::openai::{ + X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER, }; use serde_json::{Map, Value}; @@ -14,43 +19,114 @@ use super::error::RequestDecodeError; /// Gateway 扩展 `provider_options` 的唯一受支持版本。 pub const PROVIDER_OPTIONS_VERSION: &str = "v1"; -const TOP_LEVEL_FIELDS: &[&str] = &[ - "input", - "instructions", - "max_output_tokens", - "model", - "previous_response_id", - "provider_options", - "reasoning", - "store", - "stream", - "text", - "tools", -]; +const OPENAI_PROTOCOL: &str = "openai"; +const OPENAI_TRANSPORT_OPTION: &str = "transport"; +const HTTP_SSE_TRANSPORT: &str = "http_sse"; +const WEBSOCKET_TRANSPORT: &str = "websocket"; +const REVIEW_SUBAGENT: &str = "review"; +const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent"; +const MAX_PROMPT_CACHE_KEY_BYTES: usize = 1_024; -const KNOWN_UNSUPPORTED_TOP_LEVEL_FIELDS: &[&str] = &[ - "background", - "conversation", - "context_management", - "include", - "max_tool_calls", - "metadata", - "moderation", - "parallel_tool_calls", - "prompt", - "prompt_cache_key", - "prompt_cache_options", - "prompt_cache_retention", - "safety_identifier", - "service_tier", - "stream_options", - "temperature", - "tool_choice", - "top_logprobs", - "top_p", - "truncation", - "user", -]; +#[derive(Clone, Default)] +pub(super) struct OpenAiRequestHeaders { + turn_state: Option, + turn_metadata: Option, + beta_features: Option, + version: Option, + include_timing_metrics: Option, + codex_window_id: Option, + parent_thread_id: Option, + + conversation_id: Option, + session_id: Option, + thread_id: Option, + client_request_id: Option, + turn_id: Option, + + responses_lite: Option, + memgen_request: Option, + subagent: Option, +} + +impl OpenAiRequestHeaders { + pub(super) fn from_headers(headers: &HeaderMap) -> Self { + Self { + turn_state: header_string(headers, "x-codex-turn-state"), + turn_metadata: header_string(headers, "x-codex-turn-metadata"), + beta_features: header_string(headers, "x-codex-beta-features"), + version: header_string(headers, "version"), + include_timing_metrics: header_string(headers, "x-responsesapi-include-timing-metrics"), + codex_window_id: header_string(headers, "x-codex-window-id"), + parent_thread_id: header_string(headers, "x-codex-parent-thread-id"), + conversation_id: header_string(headers, "conversation-id") + .or_else(|| header_string(headers, "conversation_id")), + session_id: header_string(headers, "session-id") + .or_else(|| header_string(headers, "session_id")), + thread_id: header_string(headers, "thread-id"), + client_request_id: header_string(headers, "x-client-request-id"), + turn_id: header_string(headers, "x-codex-turn-id"), + responses_lite: header_string(headers, X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER), + memgen_request: header_string(headers, X_OPENAI_MEMGEN_REQUEST_HEADER), + subagent: header_string(headers, OPENAI_SUBAGENT_KEY).filter(|value| { + matches!( + value.as_str(), + "review" | "compact" | "memory_consolidation" | "collab_spawn" + ) + }), + } + } + + fn apply_subagent(&self, body: &mut Map, forced_subagent: Option<&str>) { + if let Some(subagent) = forced_subagent.or(self.subagent.as_deref()) { + inject_subagent_metadata(body, subagent); + } + } + + fn provider_options(&self) -> Map { + let mut options = Map::new(); + insert_header_option(&mut options, "turn_state", self.turn_state.as_ref()); + insert_header_option(&mut options, "turn_metadata", self.turn_metadata.as_ref()); + insert_header_option(&mut options, "beta_features", self.beta_features.as_ref()); + insert_header_option(&mut options, "version", self.version.as_ref()); + insert_header_option( + &mut options, + "include_timing_metrics", + self.include_timing_metrics.as_ref(), + ); + insert_header_option( + &mut options, + "codex_window_id", + self.codex_window_id.as_ref(), + ); + insert_header_option( + &mut options, + "parent_thread_id", + self.parent_thread_id.as_ref(), + ); + insert_header_option( + &mut options, + "conversation_id", + self.conversation_id.as_ref(), + ); + insert_header_option(&mut options, "session_id", self.session_id.as_ref()); + insert_header_option(&mut options, "thread_id", self.thread_id.as_ref()); + insert_header_option( + &mut options, + "client_request_id", + self.client_request_id.as_ref(), + ); + insert_header_option(&mut options, "turn_id", self.turn_id.as_ref()); + insert_header_option(&mut options, "responses_lite", self.responses_lite.as_ref()); + insert_header_option(&mut options, "memgen_request", self.memgen_request.as_ref()); + options + } +} + +fn insert_header_option(options: &mut Map, field: &str, value: Option<&String>) { + if let Some(value) = value { + options.insert(field.to_owned(), Value::String(value.clone())); + } +} /// 客户端声明的 continuation 意图。 #[derive(Clone, PartialEq, Eq)] @@ -58,9 +134,6 @@ pub enum ContinuationIntent { /// 不使用先前响应。 None, /// 使用当前调用方可见的 OpenAI response ID。 - /// - /// History owner 必须在认证后把它解析为调用方隔离的 binding；不得直接 - /// 作为 Provider handle 使用。 PreviousResponseId(String), } @@ -110,9 +183,6 @@ impl ResponsesRequestMetadata { } /// 返回协议层公开模型候选。 - /// - /// Decoder 阶段它与 `requested_model` 相同；Router 冻结 Route Plan 后必须 - /// 以实际解析出的 public model 构造 logical request 事实。 #[must_use] pub fn public_model(&self) -> &str { &self.public_model @@ -172,6 +242,13 @@ pub struct DecodedResponsesRequest { } impl DecodedResponsesRequest { + /// 附着当前 WebSocket 连接保存的 Provider 私有上一轮状态。 + #[must_use] + pub fn with_provider_session_state(mut self, state: ProviderSessionState) -> Self { + self.operation = self.operation.with_provider_session_state(state); + self + } + /// 返回协议无关 operation。 #[must_use] pub const fn operation(&self) -> &Operation { @@ -213,95 +290,162 @@ impl fmt::Debug for DecodedResponsesRequest { } } -/// 严格解码 OpenAI Responses JSON object。 +/// 解码普通 `POST /v1/responses` 请求。 /// /// # Errors /// -/// JSON 非法、字段类型错误、出现未知字段或出现已知但未实现的语义时返回 -/// [`RequestDecodeError`]。错误及其 `Debug` 不保存原始 body 或 prompt。 +/// JSON 非法、顶层不是 object，或网关必须解释的路由字段无效时返回安全错误。 pub fn decode_request(body: &[u8]) -> Result { - // UTF-8 wire bytes 是 Provider 无关且不会低估 byte-level token 数量的保守上界； - // continuation owner 后续还要把恢复历史的 estimate 加入同一 canonical 事实。 + decode_request_inner(body, false, &OpenAiRequestHeaders::default()) +} + +/// 使用下游 OpenAI/Codex 请求头解码 `POST /v1/responses`。 +pub(super) fn decode_request_with_headers( + body: &[u8], + headers: &HeaderMap, +) -> Result { + decode_request_inner(body, false, &OpenAiRequestHeaders::from_headers(headers)) +} + +/// 解码 `POST /v1/responses/review` 请求并冻结 review subagent 语义。 +pub(super) fn decode_review_request_with_headers( + body: &[u8], + headers: &HeaderMap, +) -> Result { + decode_request_inner(body, true, &OpenAiRequestHeaders::from_headers(headers)) +} + +pub(super) fn decode_request_inner( + body: &[u8], + review: bool, + request_headers: &OpenAiRequestHeaders, +) -> Result { let input_token_estimate = u64::try_from(body.len()).map_err(|_| RequestDecodeError::InvalidValue { field: "input".to_owned(), })?; let value = serde_json::from_slice::(body).map_err(|_| RequestDecodeError::MalformedJson)?; - let object = value - .as_object() - .ok_or(RequestDecodeError::ExpectedObject)?; - reject_top_level_fields(object)?; + let Value::Object(mut object) = value else { + return Err(RequestDecodeError::ExpectedObject); + }; + let use_websocket = optional_bool(&object, "use_websocket", "use_websocket")?; + object.remove("use_websocket"); - let model = required_non_empty_string(object, "model", "model")?.to_owned(); + let model = required_non_empty_string(&object, "model", "model")?.to_owned(); if model.len() > 256 || model.chars().any(char::is_control) { return Err(RequestDecodeError::InvalidValue { field: "model".to_owned(), }); } - - let mut messages = Vec::new(); - if let Some(instructions) = optional_non_empty_string(object, "instructions", "instructions")? { - messages.push(canonical_message( - MessageRole::Developer, - vec![ContentPart::Text(instructions.to_owned())], - "instructions", - )?); - } - let input = object - .get("input") - .ok_or_else(|| RequestDecodeError::MissingField { + if !object.contains_key("input") { + return Err(RequestDecodeError::MissingField { field: "input".to_owned(), - })?; - messages.extend(parse_input(input)?); - - let mut request = GenerateRequest::new(messages) - .map_err(|_| RequestDecodeError::CanonicalContract { - field: "input".to_owned(), - })? - .with_estimated_context_tokens(input_token_estimate); - - if let Some(tools) = object.get("tools") { - request = request.with_tools(parse_tools(tools)?); - } - if let Some(text) = object.get("text") { - request = request.with_output_format(parse_text_format(text)?); - } - if let Some(reasoning) = object.get("reasoning") - && let Some(reasoning) = parse_reasoning(reasoning)? - { - request = request.with_reasoning(reasoning); - } - if let Some(tokens) = optional_u64(object, "max_output_tokens", "max_output_tokens")? { - if tokens == 0 { - return Err(RequestDecodeError::InvalidValue { - field: "max_output_tokens".to_owned(), - }); - } - request = request.with_max_output_tokens(tokens); + }); } + let stream = optional_bool(&object, "stream", "stream")?.unwrap_or(true); + let store = optional_bool(&object, "store", "store")?.unwrap_or(false); let continuation = - match optional_non_empty_string(object, "previous_response_id", "previous_response_id")? { + match optional_non_empty_string(&object, "previous_response_id", "previous_response_id")? { Some(response_id) => { validate_response_id(response_id)?; - request = request.with_continuation(ContinuationMode::Native); ContinuationIntent::PreviousResponseId(response_id.to_owned()) } None => ContinuationIntent::None, }; - if let Some(options) = object.get("provider_options") { - request = request.with_provider_options(parse_provider_options(options)?); + request_headers.apply_subagent(&mut object, review.then_some(REVIEW_SUBAGENT)); + let provider_options_document = object.remove("provider_options"); + let mut openai_fallback_options = request_headers.provider_options(); + if let Some(use_websocket) = use_websocket { + // 显式 provider_options 优先；body transport 只覆盖同名 header fallback。 + openai_fallback_options.insert( + OPENAI_TRANSPORT_OPTION.to_owned(), + Value::String( + if use_websocket { + WEBSOCKET_TRANSPORT + } else { + HTTP_SSE_TRANSPORT + } + .to_owned(), + ), + ); + } + let provider_options = + parse_provider_options(provider_options_document.as_ref(), openai_fallback_options)?; + + let messages = canonical_messages(&object); + let tools = canonical_function_tools(&object); + let tools_requested = object + .get("tools") + .and_then(Value::as_array) + .is_some_and(|tools| !tools.is_empty()); + let vision_requested = contains_type(object.get("input"), "input_image"); + let reasoning = canonical_reasoning(&object); + let reasoning_requested = object.get("reasoning").is_some_and(Value::is_object); + let output_format = canonical_output_format(&object); + let json_schema_requested = object + .get("text") + .and_then(|text| text.get("format")) + .and_then(|format| format.get("type")) + .and_then(Value::as_str) + == Some("json_schema"); + let max_output_tokens = object + .get("max_output_tokens") + .and_then(Value::as_u64) + .filter(|tokens| *tokens > 0); + let prompt_cache_key = object + .get("prompt_cache_key") + .and_then(Value::as_str) + .filter(|key| valid_prompt_cache_key(key)) + .map(ToOwned::to_owned); + let payload = ProtocolPayload::json_object(OPENAI_PROTOCOL, object).map_err(|_| { + RequestDecodeError::CanonicalContract { + field: "request".to_owned(), + } + })?; + let mut request = GenerateRequest::from_protocol_payload(messages, payload) + .with_estimated_context_tokens(input_token_estimate) + .with_provider_options(provider_options) + .with_response_persistence(if store { + ResponsePersistence::Store + } else { + ResponsePersistence::DoNotStore + }); + + if tools_requested { + request = request.require_feature(Feature::Tools); + } + if !tools.is_empty() { + request = request.with_tools(tools); + } + if vision_requested { + request = request.require_feature(Feature::Vision); + } + if let Some(reasoning) = reasoning { + request = request + .with_reasoning(reasoning) + .require_feature(Feature::Reasoning); + } else if reasoning_requested { + request = request.require_feature(Feature::Reasoning); + } + if json_schema_requested { + request = request.require_feature(Feature::JsonSchema); + } + if let Some(format) = output_format { + request = request.with_output_format(format); + } + if let Some(tokens) = max_output_tokens { + request = request.with_max_output_tokens(tokens); + } + if let Some(key) = prompt_cache_key { + request = request.with_prompt_cache_key(key); + } + if continuation.is_continuation() { + request = request.with_continuation(ContinuationMode::Native); } - let stream = optional_bool(object, "stream", "stream")?.unwrap_or(false); - let store = optional_bool(object, "store", "store")?.unwrap_or(true); - request = request.with_response_persistence(if store { - ResponsePersistence::Store - } else { - ResponsePersistence::DoNotStore - }); Ok(DecodedResponsesRequest { operation: Operation::Generate(request), metadata: ResponsesRequestMetadata { @@ -316,334 +460,247 @@ pub fn decode_request(body: &[u8]) -> Result) -> Result<(), RequestDecodeError> { - for field in object.keys() { - if TOP_LEVEL_FIELDS.contains(&field.as_str()) { - continue; - } - if KNOWN_UNSUPPORTED_TOP_LEVEL_FIELDS.contains(&field.as_str()) { - return Err(RequestDecodeError::UnsupportedField { - field: field.clone(), - }); - } - return Err(RequestDecodeError::UnknownField { - field: field.clone(), - }); +fn inject_subagent_metadata(body: &mut Map, subagent: &str) { + let metadata = body + .entry("client_metadata".to_owned()) + .or_insert_with(|| Value::Object(Map::new())); + if !metadata.is_object() { + *metadata = Value::Object(Map::new()); } - Ok(()) -} - -fn parse_input(value: &Value) -> Result, RequestDecodeError> { - match value { - Value::String(text) => { - require_non_empty(text, "input")?; - Ok(vec![canonical_message( - MessageRole::User, - vec![ContentPart::Text(text.clone())], - "input", - )?]) - } - Value::Array(items) if items.is_empty() => Err(RequestDecodeError::EmptyField { - field: "input".to_owned(), - }), - Value::Array(items) => items - .iter() - .enumerate() - .map(|(index, item)| parse_input_item(item, index)) - .collect(), - _ => Err(RequestDecodeError::InvalidType { - field: "input".to_owned(), - expected: "a string or array", - }), + if let Some(metadata) = metadata.as_object_mut() { + metadata.insert( + OPENAI_SUBAGENT_KEY.to_owned(), + Value::String(subagent.to_owned()), + ); } } -fn parse_input_item(value: &Value, index: usize) -> Result { - let path = format!("input[{index}]"); - let object = required_object(value, &path)?; +fn header_string(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn canonical_messages(body: &Map) -> Vec { + let mut messages = Vec::new(); + if let Some(instructions) = body + .get("instructions") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + && let Some(message) = message( + MessageRole::Developer, + vec![ContentPart::Text(instructions.to_owned())], + ) + { + messages.push(message); + } + match body.get("input") { + Some(Value::String(input)) if !input.is_empty() => { + if let Some(message) = + message(MessageRole::User, vec![ContentPart::Text(input.clone())]) + { + messages.push(message); + } + } + Some(Value::Array(input)) => { + messages.extend(input.iter().filter_map(canonical_input_item)); + } + _ => {} + } + messages +} + +fn canonical_input_item(value: &Value) -> Option { + let object = value.as_object()?; match object.get("type").and_then(Value::as_str) { - Some("function_call_output") => parse_function_call_output(object, &path), - Some("message") | None => parse_message(object, &path), - Some(_) => Err(RequestDecodeError::UnsupportedField { - field: format!("{path}.type"), - }), - } -} - -fn parse_message(object: &Map, path: &str) -> Result { - reject_nested_fields(object, &["content", "role", "type"], &[], path)?; - if let Some(kind) = object.get("type") { - require_exact_string(kind, "message", &format!("{path}.type"))?; - } - let role_path = format!("{path}.role"); - let role = match required_string_value(object, "role", &role_path)? { - "system" => MessageRole::System, - "developer" => MessageRole::Developer, - "user" => MessageRole::User, - "assistant" => MessageRole::Assistant, + Some("function_call_output") => message( + MessageRole::User, + vec![ContentPart::ToolResult { + call_id: object.get("call_id")?.as_str()?.to_owned(), + output: object.get("output")?.as_str()?.to_owned(), + }], + ), + Some("function_call") => message( + MessageRole::Assistant, + vec![ContentPart::ToolCall { + call_id: object.get("call_id")?.as_str()?.to_owned(), + name: object.get("name")?.as_str()?.to_owned(), + arguments: object.get("arguments")?.as_str()?.to_owned(), + }], + ), _ => { - return Err(RequestDecodeError::InvalidValue { field: role_path }); + let role = match object.get("role").and_then(Value::as_str)? { + "system" => MessageRole::System, + "developer" => MessageRole::Developer, + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + _ => return None, + }; + let content = canonical_content(object.get("content")?, role); + message(role, content) } - }; - let content_path = format!("{path}.content"); - let content = object - .get("content") - .ok_or_else(|| RequestDecodeError::MissingField { - field: content_path.clone(), - })?; - let parts = parse_message_content(content, role, &content_path)?; - canonical_message(role, parts, &content_path) + } } -fn parse_message_content( - value: &Value, - role: MessageRole, - path: &str, -) -> Result, RequestDecodeError> { +fn canonical_content(value: &Value, role: MessageRole) -> Vec { match value { - Value::String(text) => { - require_non_empty(text, path)?; - Ok(vec![ContentPart::Text(text.clone())]) - } - Value::Array(parts) if parts.is_empty() => Err(RequestDecodeError::EmptyField { - field: path.to_owned(), - }), + Value::String(text) if !text.is_empty() => vec![ContentPart::Text(text.clone())], Value::Array(parts) => parts .iter() - .enumerate() - .map(|(index, part)| parse_content_part(part, role, &format!("{path}[{index}]"))) + .filter_map(|part| canonical_content_part(part, role)) .collect(), - _ => Err(RequestDecodeError::InvalidType { - field: path.to_owned(), - expected: "a string or array", - }), + _ => Vec::new(), } } -fn parse_content_part( - value: &Value, - role: MessageRole, - path: &str, -) -> Result { - let object = required_object(value, path)?; - let type_path = format!("{path}.type"); - match required_string_value(object, "type", &type_path)? { - "input_text" => { - reject_nested_fields(object, &["text", "type"], &[], path)?; - let text = required_non_empty_string(object, "text", &format!("{path}.text"))?; - Ok(ContentPart::Text(text.to_owned())) - } - "output_text" if role == MessageRole::Assistant => { - reject_nested_fields( - object, - &["text", "type"], - &["annotations", "logprobs"], - path, - )?; - let text = required_non_empty_string(object, "text", &format!("{path}.text"))?; - Ok(ContentPart::Text(text.to_owned())) - } - "output_text" => Err(RequestDecodeError::InvalidValue { field: type_path }), - "input_image" => parse_input_image(object, path), - _ => Err(RequestDecodeError::UnsupportedField { field: type_path }), +fn canonical_content_part(value: &Value, role: MessageRole) -> Option { + let object = value.as_object()?; + match object.get("type").and_then(Value::as_str)? { + "input_text" | "output_text" => object + .get("text") + .and_then(Value::as_str) + .filter(|text| !text.is_empty()) + .map(|text| ContentPart::Text(text.to_owned())), + "input_image" if role == MessageRole::User => object + .get("image_url") + .and_then(Value::as_str) + .filter(|url| !url.is_empty()) + .map(|url| ContentPart::Image(ImageSource::Url(url.to_owned()))), + _ => None, } } -fn parse_input_image( - object: &Map, - path: &str, -) -> Result { - reject_nested_fields(object, &["image_url", "type"], &["detail", "file_id"], path)?; - let url_path = format!("{path}.image_url"); - let url = required_non_empty_string(object, "image_url", &url_path)?; - if url.len() > 8 * 1024 || !(url.starts_with("https://") || url.starts_with("http://")) { - return Err(RequestDecodeError::InvalidValue { field: url_path }); - } - Ok(ContentPart::Image(ImageSource::Url(url.to_owned()))) +fn message(role: MessageRole, content: Vec) -> Option { + (!content.is_empty()) + .then(|| Message::new(role, content).ok()) + .flatten() } -fn parse_function_call_output( - object: &Map, - path: &str, -) -> Result { - reject_nested_fields(object, &["call_id", "output", "type"], &[], path)?; - let call_id = required_non_empty_string(object, "call_id", &format!("{path}.call_id"))?; - let output = required_string_value(object, "output", &format!("{path}.output"))?; - canonical_message( - MessageRole::User, - vec![ContentPart::ToolResult { - call_id: call_id.to_owned(), - output: output.to_owned(), - }], - path, - ) -} - -fn parse_tools(value: &Value) -> Result, RequestDecodeError> { - let tools = value - .as_array() - .ok_or_else(|| RequestDecodeError::InvalidType { - field: "tools".to_owned(), - expected: "an array", - })?; - if tools.is_empty() { - return Ok(Vec::new()); - } - tools - .iter() - .enumerate() - .map(|(index, tool)| parse_tool(tool, index)) +fn canonical_function_tools(body: &Map) -> Vec { + body.get("tools") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|value| { + let tool = value.as_object()?; + if tool.get("type").and_then(Value::as_str) != Some("function") { + return None; + } + let name = tool.get("name")?.as_str()?; + let schema = tool.get("parameters")?.as_object()?.clone(); + ToolDefinition::new( + name, + tool.get("description") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + schema, + ) + .ok() + .map(|definition| { + definition.with_strict(tool.get("strict").and_then(Value::as_bool).unwrap_or(false)) + }) + }) .collect() } -fn parse_tool(value: &Value, index: usize) -> Result { - let path = format!("tools[{index}]"); - let object = required_object(value, &path)?; - reject_nested_fields( - object, - &["description", "name", "parameters", "strict", "type"], - &[], - &path, - )?; - require_exact_string( - object - .get("type") - .ok_or_else(|| RequestDecodeError::MissingField { - field: format!("{path}.type"), - })?, - "function", - &format!("{path}.type"), - )?; - let name = required_non_empty_string(object, "name", &format!("{path}.name"))?; - let description = - optional_non_empty_string(object, "description", &format!("{path}.description"))? - .map(ToOwned::to_owned); - let parameters_path = format!("{path}.parameters"); - let parameters = object - .get("parameters") - .ok_or_else(|| RequestDecodeError::MissingField { - field: parameters_path.clone(), - })? - .as_object() - .ok_or(RequestDecodeError::InvalidType { - field: parameters_path, - expected: "a JSON object", - })? - .clone(); - let strict = optional_bool(object, "strict", &format!("{path}.strict"))?.unwrap_or(false); - ToolDefinition::new(name, description, parameters) - .map(|tool| tool.with_strict(strict)) - .map_err(|_| RequestDecodeError::CanonicalContract { - field: format!("{path}.name"), - }) +fn canonical_reasoning(body: &Map) -> Option { + let reasoning = body.get("reasoning")?.as_object()?; + let effort = reasoning + .get("effort") + .and_then(Value::as_str) + .and_then(|value| match value { + "minimal" => Some(ReasoningEffort::Minimal), + "low" => Some(ReasoningEffort::Low), + "medium" => Some(ReasoningEffort::Medium), + "high" => Some(ReasoningEffort::High), + "xhigh" => Some(ReasoningEffort::ExtraHigh), + _ => None, + }); + let summary = reasoning + .get("summary") + .and_then(Value::as_str) + .and_then(|value| match value { + "auto" => Some(ReasoningSummary::Auto), + "concise" => Some(ReasoningSummary::Concise), + "detailed" => Some(ReasoningSummary::Detailed), + "none" => Some(ReasoningSummary::None), + _ => None, + }); + (effort.is_some() || summary.is_some()).then_some(ReasoningRequirement { effort, summary }) } -fn parse_text_format(value: &Value) -> Result { - let object = required_object(value, "text")?; - reject_nested_fields(object, &["format"], &["verbosity"], "text")?; - let Some(format) = object.get("format") else { - return Ok(OutputFormat::Text); - }; - let format = required_object(format, "text.format")?; - let format_type = required_string_value(format, "type", "text.format.type")?; - match format_type { - "text" => { - reject_nested_fields(format, &["type"], &[], "text.format")?; - Ok(OutputFormat::Text) - } - "json_object" => { - reject_nested_fields(format, &["type"], &[], "text.format")?; - Ok(OutputFormat::JsonObject) - } - "json_schema" => parse_json_schema_format(format), - _ => Err(RequestDecodeError::UnsupportedField { - field: "text.format.type".to_owned(), - }), +fn canonical_output_format(body: &Map) -> Option { + let format = body.get("text")?.get("format")?.as_object()?; + match format.get("type")?.as_str()? { + "text" => Some(OutputFormat::Text), + "json_object" => Some(OutputFormat::JsonObject), + "json_schema" => JsonSchemaFormat::new( + format.get("name")?.as_str()?, + format + .get("description") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + format.get("schema")?.as_object()?.clone(), + format + .get("strict") + .and_then(Value::as_bool) + .unwrap_or(false), + ) + .ok() + .map(OutputFormat::JsonSchema), + _ => None, } } -fn parse_json_schema_format( - object: &Map, -) -> Result { - reject_nested_fields( - object, - &["description", "name", "schema", "strict", "type"], - &[], - "text.format", - )?; - let name = required_non_empty_string(object, "name", "text.format.name")?; - let description = optional_non_empty_string(object, "description", "text.format.description")? - .map(ToOwned::to_owned); - let schema = object - .get("schema") - .ok_or_else(|| RequestDecodeError::MissingField { - field: "text.format.schema".to_owned(), - })? +fn contains_type(value: Option<&Value>, expected: &str) -> bool { + match value { + Some(Value::Array(values)) => values + .iter() + .any(|value| contains_type(Some(value), expected)), + Some(Value::Object(object)) => { + object.get("type").and_then(Value::as_str) == Some(expected) + || object + .values() + .any(|value| contains_type(Some(value), expected)) + } + _ => false, + } +} + +fn valid_prompt_cache_key(key: &str) -> bool { + !key.is_empty() + && key.len() <= MAX_PROMPT_CACHE_KEY_BYTES + && key.trim() == key + && !key.chars().any(char::is_control) +} + +fn parse_provider_options( + value: Option<&Value>, + mut openai_fallback_options: Map, +) -> Result { + let Some(value) = value else { + if openai_fallback_options.is_empty() { + return Ok(ProviderOptions::new()); + } + openai_fallback_options.insert("schema_version".to_owned(), Value::from(1)); + let mut options = ProviderOptions::new(); + options + .insert(OPENAI_PROTOCOL, openai_fallback_options) + .map_err(|_| RequestDecodeError::CanonicalContract { + field: "provider_options.providers.openai".to_owned(), + })?; + return Ok(options); + }; + let object = value .as_object() .ok_or_else(|| RequestDecodeError::InvalidType { - field: "text.format.schema".to_owned(), + field: "provider_options".to_owned(), expected: "a JSON object", - })? - .clone(); - let strict = optional_bool(object, "strict", "text.format.strict")?.unwrap_or(false); - JsonSchemaFormat::new(name, description, schema, strict) - .map(OutputFormat::JsonSchema) - .map_err(|_| RequestDecodeError::CanonicalContract { - field: "text.format.name".to_owned(), - }) -} - -fn parse_reasoning(value: &Value) -> Result, RequestDecodeError> { - let object = required_object(value, "reasoning")?; - reject_nested_fields(object, &["effort", "summary"], &[], "reasoning")?; - let effort = match object.get("effort") { - None | Some(Value::Null) => None, - Some(Value::String(value)) => Some(match value.as_str() { - "minimal" => ReasoningEffort::Minimal, - "low" => ReasoningEffort::Low, - "medium" => ReasoningEffort::Medium, - "high" => ReasoningEffort::High, - "xhigh" => ReasoningEffort::ExtraHigh, - _ => { - return Err(RequestDecodeError::InvalidValue { - field: "reasoning.effort".to_owned(), - }); - } - }), - Some(_) => { - return Err(RequestDecodeError::InvalidType { - field: "reasoning.effort".to_owned(), - expected: "a string", - }); - } - }; - let summary = match object.get("summary") { - None | Some(Value::Null) => None, - Some(Value::String(value)) => Some(match value.as_str() { - "auto" => ReasoningSummary::Auto, - "concise" => ReasoningSummary::Concise, - "detailed" => ReasoningSummary::Detailed, - "none" => ReasoningSummary::None, - _ => { - return Err(RequestDecodeError::InvalidValue { - field: "reasoning.summary".to_owned(), - }); - } - }), - Some(_) => { - return Err(RequestDecodeError::InvalidType { - field: "reasoning.summary".to_owned(), - expected: "a string", - }); - } - }; - Ok((effort.is_some() || summary.is_some()).then_some(ReasoningRequirement { effort, summary })) -} - -fn parse_provider_options(value: &Value) -> Result { - let object = required_object(value, "provider_options")?; - reject_nested_fields(object, &["providers", "version"], &[], "provider_options")?; - let version = required_string_value(object, "version", "provider_options.version")?; + })?; + let version = required_non_empty_string(object, "version", "provider_options.version")?; if version != PROVIDER_OPTIONS_VERSION { return Err(RequestDecodeError::UnsupportedProviderOptionsVersion); } @@ -659,27 +716,44 @@ fn parse_provider_options(value: &Value) -> Result, + fallback_options: &mut Map, +) { + for (field, value) in std::mem::take(fallback_options) { + provider_options.entry(field).or_insert(value); + } +} + fn validate_response_id(response_id: &str) -> Result<(), RequestDecodeError> { - if response_id.len() > 256 - || !response_id.starts_with("resp_") - || response_id.chars().any(char::is_control) - { + if response_id.len() > 256 || response_id.chars().any(char::is_control) { return Err(RequestDecodeError::InvalidValue { field: "previous_response_id".to_owned(), }); @@ -687,53 +761,12 @@ fn validate_response_id(response_id: &str) -> Result<(), RequestDecodeError> { Ok(()) } -fn canonical_message( - role: MessageRole, - parts: Vec, - field: &str, -) -> Result { - Message::new(role, parts).map_err(|_| RequestDecodeError::CanonicalContract { - field: field.to_owned(), - }) -} - -fn reject_nested_fields( - object: &Map, - supported: &[&str], - known_unsupported: &[&str], - path: &str, -) -> Result<(), RequestDecodeError> { - for field in object.keys() { - if supported.contains(&field.as_str()) { - continue; - } - let field_path = format!("{path}.{field}"); - if known_unsupported.contains(&field.as_str()) { - return Err(RequestDecodeError::UnsupportedField { field: field_path }); - } - return Err(RequestDecodeError::UnknownField { field: field_path }); - } - Ok(()) -} - -fn required_object<'a>( - value: &'a Value, - field: &str, -) -> Result<&'a Map, RequestDecodeError> { - value - .as_object() - .ok_or_else(|| RequestDecodeError::InvalidType { - field: field.to_owned(), - expected: "a JSON object", - }) -} - -fn required_string_value<'a>( +fn required_non_empty_string<'a>( object: &'a Map, key: &str, field: &str, ) -> Result<&'a str, RequestDecodeError> { - object + let value = object .get(key) .ok_or_else(|| RequestDecodeError::MissingField { field: field.to_owned(), @@ -742,16 +775,12 @@ fn required_string_value<'a>( .ok_or_else(|| RequestDecodeError::InvalidType { field: field.to_owned(), expected: "a string", - }) -} - -fn required_non_empty_string<'a>( - object: &'a Map, - key: &str, - field: &str, -) -> Result<&'a str, RequestDecodeError> { - let value = required_string_value(object, key, field)?; - require_non_empty(value, field)?; + })?; + if value.is_empty() { + return Err(RequestDecodeError::EmptyField { + field: field.to_owned(), + }); + } Ok(value) } @@ -762,10 +791,10 @@ fn optional_non_empty_string<'a>( ) -> Result, RequestDecodeError> { match object.get(key) { None | Some(Value::Null) => Ok(None), - Some(Value::String(value)) => { - require_non_empty(value, field)?; - Ok(Some(value)) - } + Some(Value::String(value)) if !value.is_empty() => Ok(Some(value)), + Some(Value::String(_)) => Err(RequestDecodeError::EmptyField { + field: field.to_owned(), + }), Some(_) => Err(RequestDecodeError::InvalidType { field: field.to_owned(), expected: "a string", @@ -787,54 +816,3 @@ fn optional_bool( }), } } - -fn optional_u64( - object: &Map, - key: &str, - field: &str, -) -> Result, RequestDecodeError> { - match object.get(key) { - None | Some(Value::Null) => Ok(None), - Some(Value::Number(value)) => { - value - .as_u64() - .map(Some) - .ok_or_else(|| RequestDecodeError::InvalidType { - field: field.to_owned(), - expected: "a non-negative integer", - }) - } - Some(_) => Err(RequestDecodeError::InvalidType { - field: field.to_owned(), - expected: "a non-negative integer", - }), - } -} - -fn require_non_empty(value: &str, field: &str) -> Result<(), RequestDecodeError> { - if value.is_empty() { - return Err(RequestDecodeError::EmptyField { - field: field.to_owned(), - }); - } - Ok(()) -} - -fn require_exact_string( - value: &Value, - expected: &'static str, - field: &str, -) -> Result<(), RequestDecodeError> { - let actual = value - .as_str() - .ok_or_else(|| RequestDecodeError::InvalidType { - field: field.to_owned(), - expected: "a string", - })?; - if actual != expected { - return Err(RequestDecodeError::UnsupportedField { - field: field.to_owned(), - }); - } - Ok(()) -} diff --git a/backend/crates/gateway-api/src/openai/responses/response.rs b/backend/crates/gateway-api/src/openai/responses/response.rs index 130d9a9da..b56f5991b 100644 --- a/backend/crates/gateway-api/src/openai/responses/response.rs +++ b/backend/crates/gateway-api/src/openai/responses/response.rs @@ -4,14 +4,254 @@ use std::collections::BTreeMap; use gateway_core::accounting::Usage; use gateway_core::event::{ - ContentItem, ContentKind, EventSequenceValidator, FinishReason, GatewayEvent, ReasoningDelta, - ResponseMeta, TextDelta, ToolCallDelta, + ContentItem, ContentKind, EventSequenceValidator, FinishReason, GatewayEvent, ProviderEvent, + ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta, }; use gateway_protocol::openai::sse::encode_sse_event; use serde_json::{Map, Value, json}; use super::error::ResponseEncodeError; +const OPENAI_PROTOCOL: &str = "openai"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ResponseRepresentation { + Undecided, + Canonical, + OpenAiWire, +} + +/// Provider event 封套到 OpenAI Responses wire 的唯一表达边界。 +/// +/// OpenAI Provider 带原生 wire 时只下发该 wire；其他 Provider 的 canonical +/// facts 仍由 [`ResponsesCollector`] 表达。同一响应不能在两种模式之间切换。 +#[derive(Debug)] +pub struct OpenAiResponsesEncoder { + canonical: ResponsesCollector, + representation: ResponseRepresentation, + upstream_response_id: Option, + gateway_response_id: Option, + wire_terminal: Option, + pending_wire: Vec<(Option, Value)>, + canonical_completed: bool, +} + +impl OpenAiResponsesEncoder { + /// 创建响应 encoder。 + #[must_use] + pub fn new(created_at: u64) -> Self { + Self { + canonical: ResponsesCollector::new(created_at), + representation: ResponseRepresentation::Undecided, + upstream_response_id: None, + gateway_response_id: None, + wire_terminal: None, + pending_wire: Vec::new(), + canonical_completed: false, + } + } + + /// 消费一个 Provider event，并返回 SSE frames。 + /// + /// # Errors + /// + /// canonical 顺序无效、响应身份变化或表达模式混用时返回错误。 + pub fn push_sse(&mut self, event: &ProviderEvent) -> Result, ResponseEncodeError> { + self.push(event, WireEncoding::Sse) + } + + /// 消费一个 Provider event，并返回 WebSocket JSON messages。 + /// + /// # Errors + /// + /// 继承 [`Self::push_sse`] 的错误。 + pub fn push_websocket( + &mut self, + event: &ProviderEvent, + ) -> Result, ResponseEncodeError> { + self.push(event, WireEncoding::WebSocket) + } + + /// 返回是否已经观察到 canonical 或 wire 终态。 + #[must_use] + pub fn is_completed(&self) -> bool { + (self.wire_terminal.is_some() && self.canonical_completed) + || self.canonical.final_response.is_some() + } + + /// 返回 Core 冻结并已写入客户端事件的网关响应 ID。 + #[must_use] + pub fn gateway_response_id(&self) -> Option<&str> { + self.gateway_response_id.as_deref() + } + + /// 校验完整响应并返回非流式 JSON。 + /// + /// # Errors + /// + /// 流不完整或原生 wire 缺少终态 response object 时返回错误。 + pub fn finish(self) -> Result { + match self.representation { + ResponseRepresentation::OpenAiWire if self.canonical_completed => self + .wire_terminal + .ok_or(ResponseEncodeError::MissingWireTerminal), + ResponseRepresentation::OpenAiWire => Err(ResponseEncodeError::MissingWireTerminal), + ResponseRepresentation::Undecided | ResponseRepresentation::Canonical => { + self.canonical.finish() + } + } + } + + fn push( + &mut self, + event: &ProviderEvent, + encoding: WireEncoding, + ) -> Result, ResponseEncodeError> { + self.observe_identity(event)?; + let openai_wire = event + .wire_event() + .filter(|wire| wire.protocol() == OPENAI_PROTOCOL); + if openai_wire.is_some() { + if self.representation == ResponseRepresentation::Canonical { + return Err(ResponseEncodeError::MixedWireRepresentation); + } + self.representation = ResponseRepresentation::OpenAiWire; + } else if event.has_canonical_facts() + && self.representation == ResponseRepresentation::Undecided + { + self.representation = ResponseRepresentation::Canonical; + } + + if self.representation == ResponseRepresentation::OpenAiWire { + let Some(wire) = openai_wire else { + return Ok(Vec::new()); + }; + self.pending_wire.push(( + wire.event_type().map(ToOwned::to_owned), + wire.data().clone(), + )); + if self.gateway_response_id.is_none() { + return Ok(Vec::new()); + } + let pending = std::mem::take(&mut self.pending_wire); + return pending + .into_iter() + .map(|(event_type, mut data)| { + self.rewrite_wire_identity(&mut data); + let effective_type = event_type + .as_deref() + .or_else(|| data.get("type").and_then(Value::as_str)); + if matches!( + effective_type, + Some("response.completed" | "response.incomplete") + ) { + self.wire_terminal = data.get("response").cloned(); + } + Ok(match encoding { + WireEncoding::Sse => encode_sse_event( + event_type.as_deref().unwrap_or_default(), + &data.to_string(), + ), + WireEncoding::WebSocket => data.to_string(), + }) + }) + .collect(); + } + + let mut canonical_frames = Vec::new(); + for fact in event.canonical_facts() { + canonical_frames.extend(self.canonical.push(fact)?); + } + + match encoding { + WireEncoding::Sse => Ok(canonical_frames), + WireEncoding::WebSocket => canonical_frames + .into_iter() + .map(|frame| { + frame + .lines() + .find_map(|line| line.strip_prefix("data: ")) + .map(ToOwned::to_owned) + .ok_or(ResponseEncodeError::InvalidEventEncoding) + }) + .collect(), + } + } + + fn observe_identity(&mut self, event: &ProviderEvent) -> Result<(), ResponseEncodeError> { + for fact in event.canonical_facts() { + let metadata = match fact { + GatewayEvent::Started(metadata) => metadata, + GatewayEvent::Completed(metadata) => { + self.canonical_completed = true; + metadata + } + _ => continue, + }; + let Some(upstream) = metadata.upstream_response_id() else { + continue; + }; + if self + .upstream_response_id + .as_deref() + .is_some_and(|current| current != upstream.as_str()) + || self + .gateway_response_id + .as_deref() + .is_some_and(|current| current != metadata.response_id()) + { + return Err(ResponseEncodeError::WireIdentityChanged); + } + self.upstream_response_id = Some(upstream.as_str().to_owned()); + self.gateway_response_id = Some(metadata.response_id().to_owned()); + } + Ok(()) + } + + fn rewrite_wire_identity(&self, data: &mut Value) { + let (Some(upstream), Some(gateway)) = ( + self.upstream_response_id.as_deref(), + self.gateway_response_id.as_deref(), + ) else { + return; + }; + rewrite_response_identity(data, upstream, gateway); + } +} + +#[derive(Debug, Clone, Copy)] +enum WireEncoding { + Sse, + WebSocket, +} + +fn rewrite_response_identity(value: &mut Value, upstream: &str, gateway: &str) { + match value { + Value::Array(values) => { + for value in values { + rewrite_response_identity(value, upstream, gateway); + } + } + Value::Object(object) => { + if let Some(Value::String(response_id)) = object.get_mut("response_id") + && response_id == upstream + { + *response_id = gateway.to_owned(); + } + if let Some(Value::Object(response)) = object.get_mut("response") + && let Some(Value::String(response_id)) = response.get_mut("id") + && response_id == upstream + { + *response_id = gateway.to_owned(); + } + for value in object.values_mut() { + rewrite_response_identity(value, upstream, gateway); + } + } + _ => {} + } +} + /// 同一 canonical event 流生成的两种 Responses 表达。 #[derive(Debug, Clone, PartialEq)] pub struct CollectedResponses { diff --git a/backend/crates/gateway-api/src/openai/responses/websocket.rs b/backend/crates/gateway-api/src/openai/responses/websocket.rs index ea3d7987d..492833e01 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket.rs @@ -1,6 +1,7 @@ //! OpenAI Responses WebSocket adapter。 use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; use axum::{ extract::{ @@ -12,42 +13,42 @@ use axum::{ response::{IntoResponse, Response}, }; use futures::StreamExt; -use gateway_core::engine::{CommitRequirement, EngineError}; +use gateway_core::engine::execution::{AuthenticatedClient, ClientTransport, StartedExecution}; +use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError}; use gateway_core::error::{GatewayError, GatewayErrorKind}; -use gateway_core::event::GatewayEvent; +use gateway_core::event::ProviderResponseHeader; +use gateway_core::lifecycle::{ConnectionGuard, ConnectionLifecycle}; +use gateway_core::operation::ProviderSessionState; use serde_json::{Map, Value, json}; use thiserror::Error; -use crate::openai::{ - auth::{authenticate_client, authentication_error_response}, - error::{gateway_error_contract, gateway_error_from_engine, runtime_unavailable_response}, - router::MAX_CLIENT_REQUEST_BODY_BYTES, - service::{ - DeliveryEvent, OpenAiApiState, OpenAiClientService, ResponseExecutionSession, - ResponsesTransport, StartedResponse, +use crate::{ + ApiState, + openai::{ + auth::{authenticate_client, authentication_error_response}, + error::{gateway_error_contract, gateway_error_from_engine, runtime_unavailable_response}, + router::MAX_CLIENT_REQUEST_BODY_BYTES, + service::{OpenAiService, created_at_unix_seconds}, }, }; use super::{ - DecodedResponsesRequest, ProtocolErrorBody, RequestDecodeError, ResponsesCollector, - decode_request, + DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody, RequestDecodeError, http::{PendingExecution, request_client_context}, + request::OpenAiRequestHeaders, }; const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection"; const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only"; /// 将已认证的 `GET /v1/responses` 升级为 Responses WebSocket。 -pub async fn responses_websocket( - State(state): State, +pub(crate) async fn responses_websocket( + State(state): State, connect_info: Option>>, headers: HeaderMap, websocket: WebSocketUpgrade, -) -> Response -where - S: OpenAiApiState, -{ - let service = state.openai_client_api(); +) -> Response { + let service = state.openai().clone(); let client = match authenticate_client(&service, &headers) { Ok(client) => client, Err(error) => return authentication_error_response(error), @@ -56,81 +57,87 @@ where &headers, connect_info.map(|Extension(ConnectInfo(address))| address), ); - ResponsesWebSocketAdapter::new(service) - .upgrade_with_client_context(websocket, client, client_ip, user_agent) + ResponsesWebSocketAdapter::new(service).upgrade_with_client_context( + websocket, + client, + client_ip, + user_agent, + OpenAiRequestHeaders::from_headers(&headers), + ) } /// 已鉴权 Responses WebSocket 升级边界。 #[derive(Clone)] -pub struct ResponsesWebSocketAdapter { - service: S, +pub(crate) struct ResponsesWebSocketAdapter { + service: OpenAiService, } -impl ResponsesWebSocketAdapter -where - S: OpenAiClientService, -{ +impl ResponsesWebSocketAdapter { /// 绑定应用提供的唯一 OpenAI 客户端服务端口。 #[must_use] - pub const fn new(service: S) -> Self { + pub const fn new(service: OpenAiService) -> Self { Self { service } } - /// 升级一个已认证的请求；明文 API key 不进入长连接 session。 - #[must_use] - pub fn upgrade(&self, websocket: WebSocketUpgrade, client: S::Client) -> Response { - self.upgrade_with_client_context(websocket, client, None, None) - } - fn upgrade_with_client_context( &self, websocket: WebSocketUpgrade, - client: S::Client, + client: AuthenticatedClient, client_ip: Option, user_agent: Option, + request_headers: OpenAiRequestHeaders, ) -> Response { - if self.service.is_shutting_down() { - return runtime_unavailable_response().into_response(); - } + let connection_guard = match self.service.try_register_connection() { + Ok(guard) => guard, + Err(_) => return runtime_unavailable_response().into_response(), + }; let session = ResponsesWebSocketSession { service: self.service.clone(), client, - connection_id: self.service.next_connection_id(), + connection_id: self.service.next_request_id().replacen("req_", "ws_", 1), client_ip, user_agent, + request_headers, + lifecycle: self.service.lifecycle(), + _connection_guard: connection_guard, }; - let lifecycle = self.service.clone(); websocket .max_message_size(MAX_CLIENT_REQUEST_BODY_BYTES) .max_frame_size(MAX_CLIENT_REQUEST_BODY_BYTES) .on_upgrade(move |socket| async move { - lifecycle.spawn_connection(Box::pin(serve_responses_websocket(socket, session))); + serve_responses_websocket(socket, session).await; }) } } -struct ResponsesWebSocketSession -where - S: OpenAiClientService, -{ - service: S, - client: S::Client, +struct ResponsesWebSocketSession { + service: OpenAiService, + client: AuthenticatedClient, connection_id: String, client_ip: Option, user_agent: Option, + request_headers: OpenAiRequestHeaders, + lifecycle: Arc, + _connection_guard: Box, } -async fn serve_responses_websocket(mut socket: WebSocket, session: ResponsesWebSocketSession) -where - S: OpenAiClientService, -{ +async fn serve_responses_websocket(mut socket: WebSocket, session: ResponsesWebSocketSession) { tracing::info!( websocket_connection_id = %session.connection_id, "Responses WebSocket connected" ); let mut request_count = 0_u64; + let mut replay = ConnectionReplaySnapshot::default(); - while let Some(message) = socket.next().await { + let cancellation = session.lifecycle.cancellation(); + loop { + let message = tokio::select! { + () = cancellation.cancelled() => break, + message = socket.next() => message, + }; + let Some(message) = message else { + break; + }; let payload = match message { Ok(Message::Text(payload)) => payload, Ok(Message::Ping(_) | Message::Pong(_)) => continue, @@ -150,37 +157,40 @@ where }; request_count = request_count.saturating_add(1); let correlation_id = session.service.next_request_id(); - let decoded = match decode_response_create(payload.as_str()) { - Ok(decoded) => { - decoded.with_client_context(session.client_ip, session.user_agent.clone()) - } - Err(error) => { - tracing::info!( - websocket_connection_id = %session.connection_id, - request_id = %correlation_id, - error = %error, - "Responses WebSocket request rejected" - ); - if send_protocol_error( - &mut socket, - StatusCode::BAD_REQUEST, - error.protocol_body(), - &correlation_id, - ) - .await - == ForwardOutcome::Disconnect - { - break; + let decoded = + match decode_response_create_with_context(payload.as_str(), &session.request_headers) { + Ok(decoded) => { + decoded.with_client_context(session.client_ip, session.user_agent.clone()) } - continue; - } - }; + Err(error) => { + tracing::info!( + websocket_connection_id = %session.connection_id, + request_id = %correlation_id, + error = %error, + "Responses WebSocket request rejected" + ); + if send_protocol_error( + &mut socket, + StatusCode::BAD_REQUEST, + error.protocol_body(), + &correlation_id, + ) + .await + == ForwardOutcome::Disconnect + { + break; + } + continue; + } + }; + let decoded = replay.prepare(decoded); let started = match session .service .start_response( session.client.clone(), decoded, - ResponsesTransport::WebSocket, + ClientTransport::WebSocket, + "/v1/responses", ) .await { @@ -195,7 +205,8 @@ where } }; - if forward_execution(&mut socket, started).await == ForwardOutcome::Disconnect { + if forward_execution(&mut socket, started, &mut replay).await == ForwardOutcome::Disconnect + { break; } } @@ -213,13 +224,44 @@ enum ForwardOutcome { Disconnect, } -async fn forward_execution(socket: &mut WebSocket, started: StartedResponse) -> ForwardOutcome -where - S: ResponseExecutionSession, -{ - let request_id = started.request_id().to_owned(); - let (session, created_at, streaming) = started.into_parts(); - let mut execution = PendingExecution::new(session); +#[derive(Default)] +struct ConnectionReplaySnapshot { + last_response_id: Option, + provider_state: Option, +} + +impl ConnectionReplaySnapshot { + fn prepare(&self, request: DecodedResponsesRequest) -> DecodedResponsesRequest { + match ( + request.metadata().continuation().previous_response_id(), + self.last_response_id.as_deref(), + self.provider_state.as_ref(), + ) { + (Some(previous), Some(last), Some(state)) if previous == last => { + request.with_provider_session_state(state.clone()) + } + _ => request, + } + } + + fn commit(&mut self, response_id: String, provider_state: Option) { + self.last_response_id = Some(response_id); + self.provider_state = provider_state; + } +} + +async fn forward_execution( + socket: &mut WebSocket, + started: StartedExecution, + replay: &mut ConnectionReplaySnapshot, +) -> ForwardOutcome { + let request_id = started.request_id.to_string(); + let streaming = started.stream; + let mut execution = PendingExecution::new(started.session); + let created_at = match created_at_unix_seconds(started.created_at) { + Ok(created_at) => created_at, + Err(error) => return send_gateway_error(socket, &error, &request_id).await, + }; if !streaming { let error = GatewayError::new( GatewayErrorKind::Internal, @@ -249,7 +291,8 @@ where } } }; - let (first, requirement) = first.into_parts(); + let requirement = first.commit_requirement(); + let mut first = first.into_provider_events(); if requirement != CommitRequirement::CommitBeforeDelivery { let error = GatewayError::new( GatewayErrorKind::Internal, @@ -257,34 +300,66 @@ where ); return send_gateway_error(socket, &error, &request_id).await; } - let mut collector = ResponsesCollector::new(created_at); - let first_messages = match collector.push_websocket_events(&first) { - Ok(messages) if !messages.is_empty() => messages, - Ok(_) => { - let error = GatewayError::new( - GatewayErrorKind::Internal, - "gateway first event encoded no WebSocket message", - ); - return send_gateway_error(socket, &error, &request_id).await; + let mut encoder = OpenAiResponsesEncoder::new(created_at); + let mut provider_state = None; + let mut first_messages = Vec::new(); + for event in &mut first { + if let Some(update) = event.take_session_update() { + provider_state = Some(update); } - Err(error) => { - return send_protocol_error( - socket, - StatusCode::INTERNAL_SERVER_ERROR, - error.protocol_body(), - &request_id, - ) - .await; + match encoder.push_websocket(event) { + Ok(messages) => first_messages.extend(messages), + Err(error) => { + return send_protocol_error( + socket, + StatusCode::INTERNAL_SERVER_ERROR, + error.protocol_body(), + &request_id, + ) + .await; + } } - }; + } + if first_messages.is_empty() { + let error = GatewayError::new( + GatewayErrorKind::Internal, + "gateway commit batch encoded no output", + ); + return send_gateway_error(socket, &error, &request_id).await; + } let Some(response_session) = execution.session_mut() else { return ForwardOutcome::Disconnect; }; + let response_headers = response_session.response_headers().to_vec(); if let Err(error) = response_session.commit_downstream(None).await { let error = gateway_error_from_engine(&error); return send_gateway_error(socket, &error, &request_id).await; } - if !send_text(socket, response_metadata_event(&request_id)).await + if encoder.is_completed() { + if let Err(outcome) = confirm_completed_execution(socket, &mut execution, &request_id).await + { + return outcome; + } + if let Err(error) = commit_connection_replay(replay, &encoder, provider_state.take()) { + return send_gateway_error(socket, &error, &request_id).await; + } + if !send_text( + socket, + response_metadata_event(&request_id, &response_headers), + ) + .await + || !send_messages(socket, first_messages).await + { + return ForwardOutcome::Disconnect; + } + execution.disarm(); + return ForwardOutcome::Continue; + } + if !send_text( + socket, + response_metadata_event(&request_id, &response_headers), + ) + .await || !send_messages(socket, first_messages).await { return ForwardOutcome::Disconnect; @@ -293,7 +368,8 @@ where loop { match next_active_input(socket, &mut execution).await { ActiveInput::Event(Ok(Some(delivery))) => { - let (event, requirement) = delivery.into_parts(); + let requirement = delivery.commit_requirement(); + let mut events = delivery.into_provider_events(); if requirement != CommitRequirement::AlreadyCommitted { let error = GatewayError::new( GatewayErrorKind::Internal, @@ -301,63 +377,35 @@ where ); return send_gateway_error(socket, &error, &request_id).await; } - let terminal = matches!(event, GatewayEvent::Completed(_)); - let messages = match collector.push_websocket_events(&event) { - Ok(messages) => messages, - Err(error) => { - return send_protocol_error( - socket, - StatusCode::INTERNAL_SERVER_ERROR, - error.protocol_body(), - &request_id, - ) - .await; + let mut messages = Vec::new(); + for event in &mut events { + if let Some(update) = event.take_session_update() { + provider_state = Some(update); } - }; - if terminal { - if collector.finish().is_err() { - let error = GatewayError::new( - GatewayErrorKind::Internal, - "gateway terminal event did not finish the response encoder", - ); - return send_gateway_error(socket, &error, &request_id).await; - } - loop { - match next_active_input(socket, &mut execution).await { - ActiveInput::Event(Ok(None)) - if execution - .session_mut() - .is_some_and(|session| session.is_finalized()) => - { - break; - } - ActiveInput::Event(Ok(None)) => { - let error = GatewayError::new( - GatewayErrorKind::Internal, - "gateway response was not finalized after its terminal event", - ); - return send_gateway_error(socket, &error, &request_id).await; - } - ActiveInput::Event(Ok(Some(_))) => { - let error = GatewayError::new( - GatewayErrorKind::Internal, - "gateway response continued after its terminal event", - ); - return send_gateway_error(socket, &error, &request_id).await; - } - ActiveInput::Event(Err(error)) => { - let error = gateway_error_from_engine(&error); - return send_gateway_error(socket, &error, &request_id).await; - } - ActiveInput::Control => {} - ActiveInput::Disconnect => return ForwardOutcome::Disconnect, - ActiveInput::ProtocolViolation => { - close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION) - .await; - return ForwardOutcome::Disconnect; - } + match encoder.push_websocket(event) { + Ok(encoded) => messages.extend(encoded), + Err(error) => { + return send_protocol_error( + socket, + StatusCode::INTERNAL_SERVER_ERROR, + error.protocol_body(), + &request_id, + ) + .await; } } + } + if encoder.is_completed() { + if let Err(outcome) = + confirm_completed_execution(socket, &mut execution, &request_id).await + { + return outcome; + } + if let Err(error) = + commit_connection_replay(replay, &encoder, provider_state.take()) + { + return send_gateway_error(socket, &error, &request_id).await; + } if !send_messages(socket, messages).await { return ForwardOutcome::Disconnect; } @@ -389,25 +437,79 @@ where } } +fn commit_connection_replay( + replay: &mut ConnectionReplaySnapshot, + encoder: &OpenAiResponsesEncoder, + provider_state: Option, +) -> Result<(), GatewayError> { + let response_id = encoder.gateway_response_id().ok_or_else(|| { + GatewayError::new( + GatewayErrorKind::Internal, + "gateway completed a response without an identity", + ) + })?; + replay.commit(response_id.to_owned(), provider_state); + Ok(()) +} + +async fn confirm_completed_execution( + socket: &mut WebSocket, + execution: &mut PendingExecution, + request_id: &str, +) -> Result<(), ForwardOutcome> { + loop { + match next_active_input(socket, execution).await { + ActiveInput::Event(Ok(None)) + if execution + .session_mut() + .is_some_and(|session| session.is_finalized()) => + { + return Ok(()); + } + ActiveInput::Event(Ok(None)) => { + let error = GatewayError::new( + GatewayErrorKind::Internal, + "gateway response was not finalized after its terminal event", + ); + return Err(send_gateway_error(socket, &error, request_id).await); + } + ActiveInput::Event(Ok(Some(_))) => { + let error = GatewayError::new( + GatewayErrorKind::Internal, + "gateway response continued after its terminal event", + ); + return Err(send_gateway_error(socket, &error, request_id).await); + } + ActiveInput::Event(Err(error)) => { + let error = gateway_error_from_engine(&error); + return Err(send_gateway_error(socket, &error, request_id).await); + } + ActiveInput::Control => {} + ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect), + ActiveInput::ProtocolViolation => { + close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await; + return Err(ForwardOutcome::Disconnect); + } + } + } +} + enum ActiveInput { - Event(Result, EngineError>), + Event(Result, EngineError>), Control, Disconnect, ProtocolViolation, } -async fn next_active_input( +async fn next_active_input( socket: &mut WebSocket, - execution: &mut PendingExecution, -) -> ActiveInput -where - S: ResponseExecutionSession, -{ + execution: &mut PendingExecution, +) -> ActiveInput { let Some(session) = execution.session_mut() else { return ActiveInput::Disconnect; }; tokio::select! { - event = session.next_delivery_event() => ActiveInput::Event(event), + event = session.next_event() => ActiveInput::Event(event), message = socket.next() => match message { Some(Ok(Message::Ping(_) | Message::Pong(_))) => ActiveInput::Control, Some(Ok(Message::Close(_))) | None | Some(Err(_)) => ActiveInput::Disconnect, @@ -426,6 +528,20 @@ where /// 无法映射到 canonical operation 时返回不包含正文内容的稳定错误。 pub fn decode_response_create( payload: &str, +) -> Result { + decode_response_create_inner(payload, &OpenAiRequestHeaders::default()) +} + +fn decode_response_create_with_context( + payload: &str, + request_headers: &OpenAiRequestHeaders, +) -> Result { + decode_response_create_inner(payload, request_headers) +} + +fn decode_response_create_inner( + payload: &str, + request_headers: &OpenAiRequestHeaders, ) -> Result { let Value::Object(mut body) = serde_json::from_str::(payload) .map_err(|_| ResponseCreateFrameError::InvalidJson)? @@ -453,7 +569,8 @@ pub fn decode_response_create( } let encoded = serde_json::to_vec(&Value::Object(body)) .map_err(|_| ResponseCreateFrameError::InvalidJson)?; - decode_request(&encoded).map_err(ResponseCreateFrameError::Request) + super::request::decode_request_inner(&encoded, false, request_headers) + .map_err(ResponseCreateFrameError::Request) } /// `response.create` 帧的稳定安全错误。 @@ -494,10 +611,28 @@ impl ResponseCreateFrameError { } } -fn response_metadata_event(request_id: &str) -> String { +fn response_metadata_event( + request_id: &str, + response_headers: &[ProviderResponseHeader], +) -> String { + let mut headers = response_headers + .iter() + .filter_map(|header| { + super::safe_response_header_name(header.name()).map(|name| { + ( + name.to_owned(), + Value::String(header.value().as_str().to_owned()), + ) + }) + }) + .collect::>(); + headers.insert( + "x-request-id".to_owned(), + Value::String(request_id.to_owned()), + ); json!({ "type": "response.metadata", - "headers": {"x-request-id": request_id}, + "headers": headers, }) .to_string() } diff --git a/backend/crates/gateway-api/src/openai/router.rs b/backend/crates/gateway-api/src/openai/router.rs index bae83acef..4840f6c97 100644 --- a/backend/crates/gateway-api/src/openai/router.rs +++ b/backend/crates/gateway-api/src/openai/router.rs @@ -1,28 +1,28 @@ //! OpenAI 客户端协议路由。 -use axum::{Router, extract::DefaultBodyLimit, routing::get}; +use axum::{ + Router, + extract::DefaultBodyLimit, + routing::{get, post}, +}; use super::{ models::{model_detail, models}, - responses::{responses, responses_websocket}, - service::OpenAiApiState, + responses::{responses, responses_websocket, review_responses}, }; +use crate::ApiState; + /// OpenAI 请求正文上限。 pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024; /// 构造 OpenAI 客户端协议路由。 -pub fn router() -> Router -where - S: OpenAiApiState, -{ +pub(crate) fn router() -> Router { Router::new() - .route( - "/v1/responses", - get(responses_websocket::).post(responses::), - ) - .route("/v1/models", get(models::)) + .route("/v1/responses", get(responses_websocket).post(responses)) + .route("/v1/responses/review", post(review_responses)) + .route("/v1/models", get(models)) // 官方 OpenAI 模型详情合同使用 path ID；它不属于 Admin API 约束。 - .route("/v1/models/{model_id}", get(model_detail::)) + .route("/v1/models/{model_id}", get(model_detail)) .layer(DefaultBodyLimit::max(MAX_CLIENT_REQUEST_BODY_BYTES)) } diff --git a/backend/crates/gateway-api/src/openai/service.rs b/backend/crates/gateway-api/src/openai/service.rs index 68d17bf70..dd3512abe 100644 --- a/backend/crates/gateway-api/src/openai/service.rs +++ b/backend/crates/gateway-api/src/openai/service.rs @@ -1,161 +1,136 @@ -//! OpenAI HTTP adapter 与应用组合根之间的窄端口。 +//! OpenAI wire adapter 到 Core 执行用例的唯一映射。 -use std::{future::Future, pin::Pin}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; -use async_trait::async_trait; -use gateway_core::{ - engine::{CommitRequirement, EngineError}, - error::GatewayError, - event::GatewayEvent, - routing::PublicModelId, +use gateway_core::engine::continuation::PreviousResponseId; +use gateway_core::engine::execution::{ + AuthenticatedClient, ClientAuthenticationError, ClientTransport, ExecutionRequestMetadata, + ExecutionService, StartExecution, StartedExecution, }; +use gateway_core::error::{GatewayError, GatewayErrorKind}; +use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle}; +use gateway_core::routing::PublicModelId; +use uuid::Uuid; -use super::{auth::ClientApiKeyAuthError, responses::DecodedResponsesRequest}; +use super::auth::ClientApiKeyAuthError; +use super::responses::{ContinuationIntent, DecodedResponsesRequest}; -/// 由应用生命周期控制器接管的长连接任务。 -pub type ConnectionTask = Pin + Send + 'static>>; - -/// Responses 请求使用的下游传输。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ResponsesTransport { - /// 普通 HTTP 或 HTTP SSE。 - Http, - /// OpenAI Responses WebSocket。 - WebSocket, +/// OpenAI HTTP/WS adapter 共享的 Core 与连接生命周期能力。 +#[derive(Clone)] +pub(crate) struct OpenAiService { + execution: Arc, + lifecycle: Arc, } -/// 一条带下游提交边界的 canonical delivery event。 -pub struct DeliveryEvent { - event: GatewayEvent, - commit_requirement: CommitRequirement, -} - -impl DeliveryEvent { - /// 绑定 canonical event 与 Engine 已冻结的提交要求。 +impl OpenAiService { #[must_use] - pub const fn new(event: GatewayEvent, commit_requirement: CommitRequirement) -> Self { - Self { - event, - commit_requirement, - } - } - - /// 拆分 canonical event 与提交要求。 - #[must_use] - pub fn into_parts(self) -> (GatewayEvent, CommitRequirement) { - (self.event, self.commit_requirement) - } -} - -/// OpenAI response delivery 只依赖的执行会话端口。 -#[async_trait] -pub trait ResponseExecutionSession: Send + 'static { - /// 读取下一条带提交语义的 canonical event。 - async fn next_delivery_event(&mut self) -> Result, EngineError>; - - /// 在提交下游前收集完整 canonical event 序列。 - async fn collect_uncommitted(&mut self) -> Result, EngineError>; - - /// 原子提交下游 delivery 边界。 - async fn commit_downstream( - &mut self, - client_status_code: Option, - ) -> Result<(), EngineError>; - - /// 首字节前失败已经终结后，补写协议 adapter 实际返回的 HTTP 状态。 - async fn record_client_status(&mut self, client_status_code: u16) -> Result<(), EngineError>; - - /// 返回执行事实是否已经终结。 - fn is_finalized(&self) -> bool; - - /// 触发请求取消；不得等待清理完成。 - fn cancel(&self); - - /// 将未完成会话交给应用 runtime 做脱离 HTTP 生命周期的最终清理。 - fn detach_finalize(self); -} - -/// 已完成认证、请求装配和 logical-request 持久化的执行结果。 -pub struct StartedResponse { - request_id: String, - session: S, - created_at_unix_seconds: u64, - streaming: bool, -} - -impl StartedResponse { - /// 构造协议 adapter 可消费的执行事实。 - #[must_use] - pub fn new( - request_id: String, - session: S, - created_at_unix_seconds: u64, - streaming: bool, + pub(crate) const fn new( + execution: Arc, + lifecycle: Arc, ) -> Self { Self { - request_id, - session, - created_at_unix_seconds, - streaming, + execution, + lifecycle, } } - /// 返回 logical request ID。 - #[must_use] - pub fn request_id(&self) -> &str { - &self.request_id - } - - /// 拆分 handler 编码所需的冻结事实。 - #[must_use] - pub fn into_parts(self) -> (S, u64, bool) { - (self.session, self.created_at_unix_seconds, self.streaming) - } -} - -/// OpenAI 客户端协议所需的应用服务端口。 -#[async_trait] -pub trait OpenAiClientService: Clone + Send + Sync + 'static { - /// 应用内部的已认证调用方快照。 - type Client: Clone + Send + Sync + 'static; - /// 应用数据面执行会话。 - type Session: ResponseExecutionSession; - - /// 用明文下游 API key 冻结一次认证快照。 - fn authenticate(&self, plaintext: &str) -> Result; - - /// 返回当前调用方可见的全部公开模型。 - fn public_models(&self, client: &Self::Client) -> Vec; - - /// 判断调用方是否可访问一个已验证语法的公开模型。 - fn contains_public_model(&self, client: &Self::Client, model: &PublicModelId) -> bool; - - /// 启动唯一数据面执行路径。 - async fn start_response( + pub(crate) fn authenticate( &self, - client: Self::Client, + plaintext: &str, + ) -> Result { + self.execution + .authenticate(plaintext) + .map_err(map_authentication_error) + } + + pub(crate) fn public_models(&self, client: &AuthenticatedClient) -> Vec { + self.execution + .public_models(client) + .into_iter() + .map(|model| model.as_str().to_owned()) + .collect() + } + + pub(crate) fn contains_public_model( + &self, + client: &AuthenticatedClient, + model: &PublicModelId, + ) -> bool { + self.execution.contains_public_model(client, model) + } + + pub(crate) async fn start_response( + &self, + client: AuthenticatedClient, request: DecodedResponsesRequest, - transport: ResponsesTransport, - ) -> Result, GatewayError>; + transport: ClientTransport, + endpoint: &'static str, + ) -> Result { + let (operation, metadata) = request.into_parts(); + let public_model = + PublicModelId::new(metadata.public_model().to_owned()).map_err(|_| { + GatewayError::new( + GatewayErrorKind::ModelNotFound, + "requested model was not found", + ) + })?; + let previous_response_id = match metadata.continuation() { + ContinuationIntent::None => None, + ContinuationIntent::PreviousResponseId(value) => { + Some(PreviousResponseId::new(value.clone()).map_err(|_| { + GatewayError::new( + GatewayErrorKind::InvalidRequest, + "previous_response_id is invalid", + ) + })?) + } + }; + self.execution + .start(StartExecution { + client, + public_model, + operation, + metadata: ExecutionRequestMetadata { + protocol: "openai".to_owned(), + endpoint: endpoint.to_owned(), + transport, + stream: metadata.stream(), + client_ip: metadata.client_ip(), + user_agent: metadata.user_agent().map(str::to_owned), + previous_response_id, + }, + }) + .await + } - /// 返回进程是否已停止接收新的长连接。 - fn is_shutting_down(&self) -> bool; + pub(crate) fn try_register_connection( + &self, + ) -> Result, ConnectionDraining> { + self.lifecycle.try_register() + } - /// 把长连接任务注册到应用生命周期控制器。 - fn spawn_connection(&self, task: ConnectionTask); + #[must_use] + pub(crate) fn lifecycle(&self) -> Arc { + Arc::clone(&self.lifecycle) + } - /// 分配仅用于诊断的连接 ID。 - fn next_connection_id(&self) -> String; - - /// 分配在数据面启动前使用的安全请求关联 ID。 - fn next_request_id(&self) -> String; + #[must_use] + pub(crate) fn next_request_id(&self) -> String { + format!("req_{}", Uuid::now_v7().simple()) + } } -/// Axum state 向协议 router 暴露的唯一组合入口。 -pub trait OpenAiApiState: Clone + Send + Sync + 'static { - /// 每个请求获取的轻量服务 adapter。 - type Service: OpenAiClientService; +pub(crate) fn created_at_unix_seconds(created_at: SystemTime) -> Result { + created_at + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid")) +} - /// 返回绑定同一应用服务集合的 OpenAI adapter。 - fn openai_client_api(&self) -> Self::Service; +const fn map_authentication_error(error: ClientAuthenticationError) -> ClientApiKeyAuthError { + match error { + ClientAuthenticationError::InvalidKey => ClientApiKeyAuthError::InvalidKey, + ClientAuthenticationError::SnapshotUnavailable => ClientApiKeyAuthError::RuntimeUnavailable, + } } diff --git a/backend/crates/gateway-api/tests/admin/accounts.rs b/backend/crates/gateway-api/tests/admin/accounts.rs index b4879d6ad..b3b722f53 100644 --- a/backend/crates/gateway-api/tests/admin/accounts.rs +++ b/backend/crates/gateway-api/tests/admin/accounts.rs @@ -1,7 +1,6 @@ mod query { - use gateway_api::admin::accounts::{ - AccountStatus, ListQuery, ProviderFilter, SortDirection, SortField, - }; + use gateway_admin::model::accounts::{AccountSortField, AccountStatus, SortDirection}; + use gateway_api::admin::accounts::ListQuery; use serde_json::json; #[test] @@ -18,14 +17,17 @@ mod query { .expect("deserialize account query"); let query = query.validate().expect("validate account query"); assert_eq!(query.page, 3); - assert_eq!(query.page_size, 20); + assert_eq!(query.page_size.get(), 20); assert!(matches!( - query.provider, - ProviderFilter::Provider(ref provider) if provider.as_str() == "xai" + query.provider_kind, + Some(ref provider) if provider.as_str() == "xai" )); assert_eq!(query.search.as_deref(), Some("operator")); assert_eq!(query.status, Some(AccountStatus::Active)); - assert_eq!(query.sort.expect("sort").field, SortField::LastUsedAt); + assert_eq!( + query.sort.expect("sort").field, + AccountSortField::LastUsedAt + ); assert_eq!( query.sort.expect("copy sort").direction, SortDirection::Desc @@ -92,6 +94,9 @@ mod response { } mod actions { + use gateway_admin::model::accounts::{ + AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus, + }; use gateway_api::admin::accounts::{ AccountActionRequest, AccountConnectionTestEvent, AccountExportData, AccountExportQuery, AccountIdQuery, AccountRefreshRequest, AccountTestQuery, @@ -106,7 +111,12 @@ mod actions { })) .expect("decode export query"); assert_eq!( - valid.into_ids().expect("valid export"), + valid + .into_ids() + .expect("valid export") + .into_iter() + .map(|id| id.to_string()) + .collect::>(), ["acct_1".to_owned(), "acct_2".to_owned()] ); @@ -156,19 +166,47 @@ mod actions { #[test] fn connection_test_events_should_preserve_the_existing_frontend_contract() { let events = [ - AccountConnectionTestEvent::started("grok-4.5"), - AccountConnectionTestEvent::request(json!({ "model": "grok-4.5" })), - AccountConnectionTestEvent::content("OK"), - AccountConnectionTestEvent::completed("active"), - AccountConnectionTestEvent::failed("upstream unavailable", "active"), + DomainConnectionTestEvent::Started { + model: "grok-4.5".to_owned(), + }, + DomainConnectionTestEvent::Request { + model: "grok-4.5".to_owned(), + input_text: "Reply with exactly OK.".to_owned(), + stream: true, + store: false, + }, + DomainConnectionTestEvent::Content { + text: "OK".to_owned(), + }, + DomainConnectionTestEvent::Completed { + account_status: AccountStatus::Active, + }, + DomainConnectionTestEvent::Failed { + message: "upstream unavailable".to_owned(), + account_status: AccountStatus::Active, + }, ] - .map(|event| event.data); + .map(|event| AccountConnectionTestEvent::from(event).data); assert_eq!( events, [ json!({ "type": "test_start", "model": "grok-4.5", "text": "正在连接上游 Responses" }), - json!({ "type": "request", "payload": { "model": "grok-4.5" } }), + json!({ + "type": "request", + "payload": { + "model": "grok-4.5", + "input": [{ + "role": "user", + "content": [{ + "type": "input_text", + "text": "Reply with exactly OK." + }] + }], + "stream": true, + "store": false + } + }), json!({ "type": "content", "text": "OK" }), json!({ "type": "test_complete", "success": true, "accountStatus": "active" }), json!({ "type": "error", "error": "upstream unavailable", "accountStatus": "active" }), diff --git a/backend/crates/gateway-api/tests/admin/auth.rs b/backend/crates/gateway-api/tests/admin/auth.rs index 6b8632b8f..64b01a9b3 100644 --- a/backend/crates/gateway-api/tests/admin/auth.rs +++ b/backend/crates/gateway-api/tests/admin/auth.rs @@ -1,16 +1,11 @@ -use std::collections::BTreeMap; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; -use gateway_api::admin::AdminServiceError; +use gateway_admin::model::auth::{LoginCommand, LoginError}; use gateway_api::admin::auth::{ - AdminAuthAuditEvent, AdminAuthBackend, AdminAuthService, AdminBackendSession, AdminLoginData, - AdminLoginError, AdminLoginRequest, AdminLogoutData, AdminSessionResolver, - AdminSessionStatusData, DefaultAdminAuthService, + AdminLoginData, AdminLoginRequest, AdminLogoutData, AdminSessionStatusData, }; use serde_json::json; +use super::AdminTestFixture; + #[test] fn login_request_should_deny_unknown_fields_and_redact_password_debug() { let password = "admin-password-must-not-leak"; @@ -50,200 +45,107 @@ fn auth_responses_should_keep_stable_wire_shapes() { ); } -#[derive(Default)] -struct MemoryAuthBackend { - password_hash: Mutex>, - admin_api_key: Mutex>, - sessions: Mutex>, - login_failures: Mutex>, - audits: Mutex>, - fail_audit: AtomicBool, -} - -#[async_trait] -impl AdminAuthBackend for MemoryAuthBackend { - async fn password_hash( - &self, - _admin_user_id: &str, - ) -> Result, AdminServiceError> { - Ok(self.password_hash.lock().expect("password lock").clone()) - } - - async fn store_password_hash( - &self, - _admin_user_id: &str, - password_hash: &str, - ) -> Result<(), AdminServiceError> { - *self.password_hash.lock().expect("password lock") = Some(password_hash.to_owned()); - Ok(()) - } - - async fn admin_api_key(&self) -> Result, AdminServiceError> { - Ok(self.admin_api_key.lock().expect("API key lock").clone()) - } - - async fn load_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError> { - Ok(self - .sessions - .lock() - .expect("session lock") - .get(session_id) - .cloned()) - } - - async fn store_admin_session( - &self, - session_id: &str, - session: &AdminBackendSession, - ) -> Result<(), AdminServiceError> { - self.sessions - .lock() - .expect("session lock") - .insert(session_id.to_owned(), session.clone()); - Ok(()) - } - - async fn delete_admin_session( - &self, - session_id: &str, - ) -> Result, AdminServiceError> { - Ok(self - .sessions - .lock() - .expect("session lock") - .remove(session_id)) - } - - async fn login_source_is_throttled( - &self, - source: &str, - failure_limit: u32, - _window_seconds: u64, - ) -> Result { - Ok(self - .login_failures - .lock() - .expect("failure lock") - .get(source) - .is_some_and(|count| *count >= failure_limit)) - } - - async fn record_login_failure( - &self, - source: &str, - failure_limit: u32, - _window_seconds: u64, - ) -> Result { - let mut failures = self.login_failures.lock().expect("failure lock"); - let count = failures.entry(source.to_owned()).or_default(); - *count = count.saturating_add(1); - Ok(*count >= failure_limit) - } - - async fn clear_login_failures(&self, source: &str) -> Result<(), AdminServiceError> { - self.login_failures - .lock() - .expect("failure lock") - .remove(source); - Ok(()) - } - - async fn append_auth_audit(&self, event: AdminAuthAuditEvent) -> Result<(), AdminServiceError> { - if self.fail_audit.load(Ordering::SeqCst) { - return Err(AdminServiceError::unavailable("audit unavailable")); - } - self.audits.lock().expect("audit lock").push(event); - Ok(()) - } -} - #[tokio::test] async fn default_auth_service_should_initialize_login_validate_and_logout() { - let backend = Arc::new(MemoryAuthBackend::default()); - let service = DefaultAdminAuthService::new("admin_1".to_owned(), 60, backend.clone()); - assert!( - service - .ensure_default_admin("strong-admin-password") - .await - .unwrap() - ); - assert!( - !service - .ensure_default_admin("different-password") - .await - .unwrap() - ); - + let fixture = AdminTestFixture::new().await; + let service = fixture.services.auth(); let session = service - .login( - "127.0.0.1", - Some("admin_1".to_owned()), - "strong-admin-password".to_owned(), - ) + .login(LoginCommand { + source: "127.0.0.1".to_owned(), + username: Some("admin_1".to_owned()), + password: "strong-admin-password".to_owned(), + }) .await .expect("login succeeds"); assert!( service - .validate(Some(session.session_id.clone())) + .validate_session(Some(&session.session_id)) .await - .unwrap() + .expect("validate session") ); assert_eq!( service .resolve_admin_user_id(Some(&session.session_id)) .await - .unwrap() + .expect("resolve session") .as_deref(), - Some("admin_1"), + Some("admin_1") ); - service.logout(session.session_id.clone()).await.unwrap(); - assert!(!service.validate(Some(session.session_id)).await.unwrap()); - assert_eq!(backend.audits.lock().expect("audit lock").len(), 2); + service + .logout(&session.session_id) + .await + .expect("logout session"); + assert!( + !service + .validate_session(Some(&session.session_id)) + .await + .expect("validate logged-out session") + ); + assert_eq!(fixture.auth.audit_count(), 2); } #[tokio::test] async fn default_auth_service_should_throttle_repeated_invalid_identity() { - let backend = Arc::new(MemoryAuthBackend::default()); - let service = DefaultAdminAuthService::new("admin_1".to_owned(), 60, backend.clone()); + let fixture = AdminTestFixture::new().await; for attempt in 1..=5 { - let error = service - .login( - "shared-source", - Some("wrong-user".to_owned()), - "wrong-password".to_owned(), - ) + let error = fixture + .services + .auth() + .login(LoginCommand { + source: "shared-source".to_owned(), + username: Some("wrong-user".to_owned()), + password: "wrong-password".to_owned(), + }) .await - .unwrap_err(); + .expect_err("invalid login"); if attempt < 5 { - assert_eq!(error, AdminLoginError::InvalidCredentials); + assert_eq!(error, LoginError::InvalidCredentials); } else { - assert_eq!(error, AdminLoginError::Throttled); + assert_eq!(error, LoginError::Throttled); } } assert_eq!( - service - .login("shared-source", None, "anything".to_owned()) + fixture + .services + .auth() + .login(LoginCommand { + source: "shared-source".to_owned(), + username: None, + password: "anything".to_owned(), + }) .await - .unwrap_err(), - AdminLoginError::Throttled, + .expect_err("source remains throttled"), + LoginError::Throttled ); - assert!(backend.sessions.lock().expect("session lock").is_empty()); + assert_eq!(fixture.auth.session_count(), 0); } #[tokio::test] async fn default_auth_service_should_verify_only_full_plaintext_admin_key() { - let backend = Arc::new(MemoryAuthBackend::default()); + let fixture = AdminTestFixture::new().await; let key = format!("admin-{}", "a".repeat(64)); - *backend.admin_api_key.lock().expect("API key lock") = Some(key.clone()); - let service = DefaultAdminAuthService::new("admin_1".to_owned(), 60, backend); + fixture.auth.set_api_key(&key); - assert!(service.verify_admin_api_key(&key).await.unwrap()); - assert!(!service.verify_admin_api_key("admin-short").await.unwrap()); assert!( - !service + fixture + .services + .auth() + .verify_admin_api_key(&key) + .await + .unwrap() + ); + assert!( + !fixture + .services + .auth() + .verify_admin_api_key("admin-short") + .await + .unwrap() + ); + assert!( + !fixture + .services + .auth() .verify_admin_api_key(&format!("admin-{}", "b".repeat(64))) .await .unwrap() @@ -252,19 +154,21 @@ async fn default_auth_service_should_verify_only_full_plaintext_admin_key() { #[tokio::test] async fn audit_failure_should_revoke_new_session_before_returning_it() { - let backend = Arc::new(MemoryAuthBackend::default()); - let service = DefaultAdminAuthService::new("admin_1".to_owned(), 60, backend.clone()); - service - .ensure_default_admin("strong-admin-password") - .await - .unwrap(); - backend.fail_audit.store(true, Ordering::SeqCst); + let fixture = AdminTestFixture::new().await; + fixture.auth.fail_audit(true); assert_eq!( - service - .login("source", None, "strong-admin-password".to_owned()) + fixture + .services + .auth() + .login(LoginCommand { + source: "source".to_owned(), + username: None, + password: "strong-admin-password".to_owned(), + }) .await - .unwrap_err(), - AdminLoginError::Unavailable, + .expect_err("audit failure rejects login"), + LoginError::Unavailable ); + assert_eq!(fixture.auth.session_count(), 0); } diff --git a/backend/crates/gateway-api/tests/admin/client_keys.rs b/backend/crates/gateway-api/tests/admin/client_keys.rs index 33398bda9..c55659a17 100644 --- a/backend/crates/gateway-api/tests/admin/client_keys.rs +++ b/backend/crates/gateway-api/tests/admin/client_keys.rs @@ -1,17 +1,15 @@ use chrono::{DateTime, TimeZone as _, Utc}; use gateway_api::admin::client_keys::{ - self, ClientKeyAdminService, ClientKeyAdminState, ClientKeyCursorData, ClientKeyCursorValue, - ClientKeyListData, ClientKeyRevisionRequest, ClientKeySort, ClientKeySortDirection, - ClientKeySortField, ClientKeyView, ClientKeyViewFields, CreateClientKeyFields, - CreateClientKeyRequest, CreatedClientKeyData, ListClientKeysFields, ListClientKeysQuery, - MutatedClientKeyData, RevealedClientKeyData, UpdateClientKeyFields, UpdateClientKeyRequest, - decode_client_key_cursor, encode_client_key_cursor, -}; -use gateway_api::admin::{ - AdminRequestContext, AdminServiceError, AdminSessionResolver, AdminSessionState, + self, ClientKeyCursorData, ClientKeyCursorValue, ClientKeyListData, ClientKeyRevisionRequest, + ClientKeySort, ClientKeySortDirection, ClientKeySortField, ClientKeyView, + CreateClientKeyRequest, CreatedClientKeyData, ListClientKeysQuery, MutatedClientKeyData, + RevealedClientKeyData, UpdateClientKeyRequest, decode_client_key_cursor, + encode_client_key_cursor, }; use serde_json::json; +use super::{AdminTestFixture, AdminTestState}; + #[test] fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() { let unknown = serde_json::from_value::(json!({ "other": true })); @@ -37,33 +35,33 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() { assert!(unknown.is_err()); assert_eq!( - zero.into_parts().expect_err("reject zero limit").field(), + zero.into_command().expect_err("reject zero limit").field(), "limit" ); assert_eq!( oversized - .into_parts() + .into_command() .expect_err("reject oversized cursor") .field(), "cursor" ); assert_eq!( oversized_search - .into_parts() + .into_command() .expect_err("reject oversized search") .field(), "search" ); assert_eq!( secret_search - .into_parts() + .into_command() .expect_err("reject full key in search") .field(), "search" ); assert_eq!( invalid_sort - .into_parts() + .into_command() .expect_err("reject invalid sort") .field(), "sortBy" @@ -73,23 +71,33 @@ fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() { "sortDirection": "asc" })) .expect("deserialize valid sort") - .into_parts() + .into_command() .expect("validate sort"); assert_eq!( - sorted.sort, - ClientKeySort { - field: ClientKeySortField::LastUsedAt, - direction: ClientKeySortDirection::Asc, - } + sorted.sort.field, + gateway_admin::model::client_keys::ClientKeySortField::LastUsedAt, + ); + assert_eq!( + sorted.sort.direction, + gateway_admin::model::client_keys::SortDirection::Asc, ); assert_eq!( serde_json::from_value::(json!({ "search": " " })) .expect("deserialize blank search") - .into_parts() + .into_command() .expect("normalize blank search") .search, None ); + assert_eq!( + serde_json::from_value::(json!({ "limit": u16::MAX })) + .expect("deserialize maximum u16 page size") + .into_command() + .expect("accept maximum u16 page size") + .page_size + .get(), + u16::MAX + ); } #[test] @@ -104,9 +112,9 @@ fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields( "tokensPerMinute": 100000 })) .expect("deserialize create") - .into_fields() + .into_command() .expect("validate create"); - assert_eq!(valid.expected_config_revision, 7); + assert_eq!(valid.expected_config_revision.get(), 7); for (payload, field) in [ ( @@ -147,7 +155,7 @@ fn client_key_mutations_should_validate_revision_text_limits_and_unknown_fields( .expect("deserialize invalid create shape"); assert_eq!( request - .into_fields() + .into_command() .expect_err("reject invalid create") .field(), field @@ -207,16 +215,18 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() { .with_ymd_and_hms(2026, 7, 18, 8, 0, 0) .single() .expect("valid time"); - let view = ClientKeyView::new(ClientKeyViewFields { - id: "key_visible".to_owned(), + let view = ClientKeyView::from(gateway_admin::model::client_keys::ClientKeyRecord { + id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"), name: "visible".to_owned(), label: None, - provider_kind: "openai".to_owned(), + provider_kind: gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"), prefix: "sk_visible12".to_owned(), enabled: true, - max_concurrency: 2, - requests_per_minute: 60, - tokens_per_minute: 100_000, + limits: gateway_core::policy::RateLimits { + max_concurrency: 2, + requests_per_minute: 60, + tokens_per_minute: 100_000, + }, created_at, updated_at: created_at, last_used_at: Some(created_at), @@ -265,93 +275,6 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() { ); } -struct RevealService; - -#[async_trait::async_trait] -impl AdminSessionResolver for RevealService { - async fn resolve_admin_user_id( - &self, - session_id: Option<&str>, - ) -> Result, AdminServiceError> { - Ok((session_id == Some("valid-session")).then(|| "admin_1".to_owned())) - } -} - -#[async_trait::async_trait] -impl ClientKeyAdminService for RevealService { - async fn list( - &self, - _query: ListClientKeysFields, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } - - async fn create( - &self, - _context: &AdminRequestContext, - _fields: CreateClientKeyFields, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } - - async fn reveal(&self, id: String) -> Result { - Ok(RevealedClientKeyData::new( - id, - format!("sk_{}", "a".repeat(43)), - )) - } - - async fn update( - &self, - _context: &AdminRequestContext, - _fields: UpdateClientKeyFields, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } - - async fn disable( - &self, - _context: &AdminRequestContext, - _id: String, - _expected_config_revision: u64, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } - - async fn enable( - &self, - _context: &AdminRequestContext, - _id: String, - _expected_config_revision: u64, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } - - async fn delete( - &self, - _context: &AdminRequestContext, - _id: String, - _expected_config_revision: u64, - ) -> Result { - Err(AdminServiceError::internal("unused test operation")) - } -} - -#[derive(Clone)] -struct RevealState(std::sync::Arc); - -impl AdminSessionState for RevealState { - fn admin_session_resolver(&self) -> &dyn AdminSessionResolver { - self.0.as_ref() - } -} - -impl ClientKeyAdminState for RevealState { - fn client_key_admin_service(&self) -> &dyn ClientKeyAdminService { - self.0.as_ref() - } -} - #[tokio::test] async fn reveal_route_should_use_query_id_and_no_store_response() { use axum::{ @@ -360,9 +283,10 @@ async fn reveal_route_should_use_query_id_and_no_store_response() { }; use tower::ServiceExt as _; - let state = RevealState(std::sync::Arc::new(RevealService)); - let response = client_keys::router::() - .with_state(state) + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = client_keys::router::() + .with_state(fixture.state()) .oneshot( Request::builder() .uri("/api/admin/client-keys/reveal?id=key_1") @@ -391,3 +315,29 @@ async fn reveal_route_should_use_query_id_and_no_store_response() { format!("sk_{}", "a".repeat(43)) ); } + +#[tokio::test] +async fn list_route_should_accept_the_full_nonzero_u16_page_size() { + use axum::{ + body::Body, + http::{Request, StatusCode, header}, + }; + use tower::ServiceExt as _; + + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = client_keys::router::() + .with_state(fixture.state()) + .oneshot( + Request::builder() + .uri("/api/admin/client-keys?limit=65535") + .header(header::COOKIE, "cpr_admin_session=valid-session") + .header("x-request-id", "req_client_keys_max_page") + .body(Body::empty()) + .expect("client key list request"), + ) + .await + .expect("client key list response"); + + assert_eq!(response.status(), StatusCode::OK); +} diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index 026532597..fb3332a65 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -1,3 +1,78 @@ +use std::{ + collections::BTreeMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, +}; + +use async_trait::async_trait; +use chrono::{Duration, Utc}; +use futures::future::BoxFuture; +use gateway_admin::{ + AdminConfig, AdminServices, InitialAdminPassword, + model::{ + MutationContext, Revision, + accounts::{ + AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount, + SetAccountEnabled, + }, + auth::{AdminAuditEvent, AdminSession}, + catalog::{ + CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog, + ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled, + UpdateProviderInstance, + }, + client_keys::{ + ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey, + NewClientKey, SetClientKeyEnabled, UpdateClientKey, + }, + observability::{ + DashboardObservation, DiagnosticDimension, DiagnosticObservation, OpsErrorPage, + OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter, UsageOverview, + UsagePage, UsageQuery, + }, + provider_credentials::{ + AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails, + CredentialImportCommit, CredentialImportResult, CredentialListQuery, + CredentialMutationResult, CredentialPage, CredentialRotationCommit, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport, + PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, + ProviderModels, ProviderQuota, + }, + settings::{ + AdminApiKey, AdminApiKeyMutation, ProviderModelMappings, ReplaceRuntimeSettings, + RotationStrategy, RuntimeSettings, + }, + system::{SystemOperationAccepted, SystemUpdateDetail, SystemUpdateStatus, SystemVersion}, + }, + ports::{ + provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind}, + store::{ + AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, + AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore, + }, + system::{ + SystemOperationError, SystemOperationErrorKind, SystemOperations, + SystemUpdateEventStream, + }, + }, +}; +use gateway_api::admin::AdminSessionState; +use gateway_core::{ + engine::{ + credential::ProviderAccountId, + probe::{AccountProbe, AccountProbeRequest, AccountProbeResult}, + }, + error::GatewayError, + policy::{ClientApiKeyId, RateLimits}, + routing::{ + ConfigRevision, ProviderInstanceId, ProviderKind, PublicModelId, UpstreamModelId, + snapshot::SnapshotControl, + }, +}; + mod accounts; mod auth; mod catalog; @@ -8,3 +83,740 @@ mod settings; mod system; mod wire; mod xai; + +pub(super) struct AdminTestFixture { + pub services: AdminServices, + pub auth: Arc, + pub settings: Arc, +} + +impl AdminTestFixture { + pub async fn new() -> Self { + let api_key = Arc::new(Mutex::new(None)); + let auth = Arc::new(MemoryAuthStore::new(api_key.clone())); + let settings = Arc::new(MemorySettingsStore::new(api_key)); + let client_keys = Arc::new(MemoryClientKeyStore); + let unused = Arc::new(UnusedStore); + let stores = AdminStorePorts::new( + unused.clone(), + auth.clone(), + unused.clone(), + client_keys.clone(), + unused, + settings.clone(), + ); + let providers: Vec> = vec![ + Arc::new(UnusedProvider::new("openai")), + Arc::new(UnusedProvider::new("xai")), + ]; + let bundle = gateway_admin::initialize( + AdminConfig { + session_ttl_minutes: 60, + default_username: "admin_1".to_owned(), + default_password: InitialAdminPassword::new("strong-admin-password"), + }, + stores, + providers, + Arc::new(NoopSnapshot), + Arc::new(NoopProbe), + Arc::new(UnusedSystem), + ) + .await + .expect("initialize test admin services"); + Self { + services: bundle.services(), + auth, + settings, + } + } + + pub fn state(&self) -> AdminTestState { + AdminTestState(self.services.clone()) + } +} + +#[derive(Clone)] +pub(super) struct AdminTestState(AdminServices); + +impl AdminSessionState for AdminTestState { + fn admin_services(&self) -> &AdminServices { + &self.0 + } +} + +pub(super) struct MemoryAuthStore { + password_hash: Mutex>, + sessions: Mutex>, + failures: Mutex>, + audits: Mutex>, + api_key: Arc>>, + fail_audit: AtomicBool, +} + +impl MemoryAuthStore { + fn new(api_key: Arc>>) -> Self { + Self { + password_hash: Mutex::new(None), + sessions: Mutex::new(BTreeMap::new()), + failures: Mutex::new(BTreeMap::new()), + audits: Mutex::new(Vec::new()), + api_key, + fail_audit: AtomicBool::new(false), + } + } + + pub fn insert_session(&self, session_id: &str) { + self.sessions.lock().expect("sessions").insert( + session_id.to_owned(), + AdminSession { + admin_user_id: "admin_1".to_owned(), + expires_at: Utc::now() + Duration::hours(1), + }, + ); + } + + pub fn set_api_key(&self, value: &str) { + *self.api_key.lock().expect("API key") = Some(AdminApiKey::new(value)); + } + + pub fn fail_audit(&self, fail: bool) { + self.fail_audit.store(fail, Ordering::SeqCst); + } + + pub fn session_count(&self) -> usize { + self.sessions.lock().expect("sessions").len() + } + + pub fn audit_count(&self) -> usize { + self.audits.lock().expect("audits").len() + } +} + +#[async_trait] +impl AuthStore for MemoryAuthStore { + async fn load_password_hash(&self, _: &str) -> AdminStoreResult> { + Ok(self.password_hash.lock().expect("password hash").clone()) + } + + async fn create_password_hash_if_absent( + &self, + _: &str, + password_hash: &str, + ) -> AdminStoreResult { + let mut stored = self.password_hash.lock().expect("password hash"); + if stored.is_some() { + return Ok(false); + } + *stored = Some(password_hash.to_owned()); + Ok(true) + } + + async fn load_admin_api_key(&self) -> AdminStoreResult> { + Ok(self.api_key.lock().expect("API key").clone()) + } + + async fn load_session(&self, session_id: &str) -> AdminStoreResult> { + Ok(self + .sessions + .lock() + .expect("sessions") + .get(session_id) + .cloned()) + } + + async fn store_session( + &self, + session_id: &str, + session: &AdminSession, + ) -> AdminStoreResult<()> { + self.sessions + .lock() + .expect("sessions") + .insert(session_id.to_owned(), session.clone()); + Ok(()) + } + + async fn delete_session(&self, session_id: &str) -> AdminStoreResult> { + Ok(self.sessions.lock().expect("sessions").remove(session_id)) + } + + async fn login_source_is_throttled( + &self, + source: &str, + failure_limit: u32, + _: u64, + ) -> AdminStoreResult { + Ok(self + .failures + .lock() + .expect("failures") + .get(source) + .is_some_and(|count| *count >= failure_limit)) + } + + async fn record_login_failure( + &self, + source: &str, + failure_limit: u32, + _: u64, + ) -> AdminStoreResult { + let mut failures = self.failures.lock().expect("failures"); + let count = failures.entry(source.to_owned()).or_default(); + *count = count.saturating_add(1); + Ok(*count >= failure_limit) + } + + async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()> { + self.failures.lock().expect("failures").remove(source); + Ok(()) + } + + async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()> { + if self.fail_audit.load(Ordering::SeqCst) { + return Err(unavailable("auth audit")); + } + self.audits.lock().expect("audits").push(event); + Ok(()) + } +} + +pub(super) struct MemorySettingsStore { + settings: Mutex, + api_key: Arc>>, +} + +impl MemorySettingsStore { + fn new(api_key: Arc>>) -> Self { + Self { + settings: Mutex::new(test_runtime_settings()), + api_key, + } + } + + pub fn set_api_key(&self, value: &str) { + *self.api_key.lock().expect("API key") = Some(AdminApiKey::new(value)); + } +} + +#[async_trait] +impl SettingsStore for MemorySettingsStore { + async fn load_runtime_settings(&self) -> AdminStoreResult { + Ok(self.settings.lock().expect("settings").clone()) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + Ok(self.api_key.lock().expect("API key").is_some()) + } + + async fn replace_runtime_settings( + &self, + command: ReplaceRuntimeSettings, + _: &MutationContext, + ) -> AdminStoreResult { + let updated = RuntimeSettings { + config_revision: next_revision(command.expected_config_revision), + provider_model_mappings: command.provider_model_mappings, + refresh_margin_seconds: command.refresh_margin_seconds, + refresh_concurrency: command.refresh_concurrency, + max_concurrent_per_account: command.max_concurrent_per_account, + request_interval_ms: command.request_interval_ms, + rotation_strategy: command.rotation_strategy, + usage_retention_days: command.usage_retention_days, + ops_event_retention_days: command.ops_event_retention_days, + audit_retention_days: command.audit_retention_days, + updated_at: Utc::now(), + }; + *self.settings.lock().expect("settings") = updated.clone(); + Ok(updated) + } + + async fn replace_admin_api_key( + &self, + expected_config_revision: Revision, + key: AdminApiKey, + _: &MutationContext, + ) -> AdminStoreResult { + *self.api_key.lock().expect("API key") = Some(key); + Ok(AdminApiKeyMutation { + config_revision: next_revision(expected_config_revision), + exists: true, + }) + } + + async fn delete_admin_api_key( + &self, + expected_config_revision: Revision, + _: &MutationContext, + ) -> AdminStoreResult { + *self.api_key.lock().expect("API key") = None; + Ok(AdminApiKeyMutation { + config_revision: next_revision(expected_config_revision), + exists: false, + }) + } +} + +pub(super) struct MemoryClientKeyStore; + +#[async_trait] +impl ClientKeyStore for MemoryClientKeyStore { + async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult { + Ok(ClientKeyPage { + config_revision: Revision::new(1).expect("revision"), + items: Vec::new(), + total: 0, + next_cursor: None, + }) + } + + async fn reveal_client_key( + &self, + id: &ClientApiKeyId, + ) -> AdminStoreResult> { + let now = Utc::now(); + Ok(Some(ClientKeySecret::new( + ClientKeyRecord { + id: id.clone(), + name: "revealed".to_owned(), + label: None, + provider_kind: ProviderKind::new("openai").expect("provider kind"), + prefix: "sk_aaaaaaaaa".to_owned(), + enabled: true, + limits: RateLimits::unlimited(), + last_used_at: None, + created_at: now, + updated_at: now, + }, + format!("sk_{}", "a".repeat(43)), + ))) + } + + async fn create_client_key( + &self, + _: NewClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key create")) + } + + async fn update_client_key( + &self, + _: UpdateClientKey, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key update")) + } + + async fn set_client_key_enabled( + &self, + _: SetClientKeyEnabled, + _: &MutationContext, + ) -> AdminStoreResult<(Revision, ClientKeyRecord)> { + Err(unavailable("client key enabled")) + } + + async fn delete_client_key( + &self, + _: DeleteClientKey, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("client key delete")) + } +} + +struct UnusedStore; + +#[async_trait] +impl AccountStore for UnusedStore { + async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult { + Err(unavailable("account list")) + } + + async fn load_account(&self, _: &str) -> AdminStoreResult> { + Err(unavailable("account")) + } + + async fn load_account_usage( + &self, + _: TimeRange, + _: &[String], + ) -> AdminStoreResult> { + Err(unavailable("account usage")) + } + + async fn list_credentials( + &self, + _: &ProviderKind, + _: CredentialListQuery, + ) -> AdminStoreResult { + Err(unavailable("credential list")) + } + + async fn credential_details( + &self, + _: &ProviderKind, + _: &ProviderAccountId, + ) -> AdminStoreResult> { + Err(unavailable("credential")) + } + + async fn load_credentials_for_export( + &self, + _: &ProviderKind, + _: &[ProviderAccountId], + ) -> AdminStoreResult> { + Err(unavailable("credential export")) + } + + async fn commit_credential_import( + &self, + _: CredentialImportCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential import")) + } + + async fn commit_authorization( + &self, + _: AuthorizationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("authorization")) + } + + async fn commit_credential_rotation( + &self, + _: CredentialRotationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential rotation")) + } + + async fn commit_credential_refresh( + &self, + _: CredentialRotationCommit, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("credential refresh")) + } + + async fn set_account_enabled( + &self, + _: SetAccountEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account enabled")) + } + + async fn delete_account( + &self, + _: DeleteAccount, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account delete")) + } + + async fn record_credential_export( + &self, + _: &[ProviderAccountId], + _: &MutationContext, + ) -> AdminStoreResult<()> { + Err(unavailable("credential export audit")) + } +} + +#[async_trait] +impl CatalogStore for UnusedStore { + async fn list_provider_instances(&self, _: bool) -> AdminStoreResult { + Err(unavailable("provider catalog")) + } + + async fn load_provider_instance( + &self, + _: &ProviderInstanceId, + ) -> AdminStoreResult> { + Err(unavailable("provider instance")) + } + + async fn create_provider_instance( + &self, + _: CreateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider create")) + } + + async fn update_provider_instance( + &self, + _: UpdateProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider update")) + } + + async fn set_provider_instance_enabled( + &self, + _: SetProviderInstanceEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider enabled")) + } + + async fn delete_provider_instance( + &self, + _: DeleteProviderInstance, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("provider delete")) + } +} + +#[async_trait] +impl ObservabilityStore for UnusedStore { + async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult { + Err(unavailable("dashboard")) + } + + async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult> { + Err(unavailable("dashboard trend")) + } + + async fn usage_trend( + &self, + _: TimeRange, + _: UsageFilter, + ) -> AdminStoreResult> { + Err(unavailable("usage trend")) + } + + async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult { + Ok(UsagePage { + items: Vec::new(), + total: 0, + next_cursor: None, + }) + } + + async fn usage_record_detail(&self, _: &str) -> AdminStoreResult { + Err(unavailable("usage detail")) + } + + async fn usage_summary(&self, _: TimeRange, _: UsageFilter) -> AdminStoreResult { + Err(unavailable("usage summary")) + } + + async fn usage_diagnostics( + &self, + _: TimeRange, + _: UsageFilter, + _: DiagnosticDimension, + ) -> AdminStoreResult> { + Err(unavailable("usage diagnostics")) + } + + async fn list_ops_errors(&self, _: OpsErrorQuery) -> AdminStoreResult { + Err(unavailable("ops errors")) + } +} + +struct UnusedProvider { + kind: ProviderKind, +} + +impl UnusedProvider { + fn new(kind: &str) -> Self { + Self { + kind: ProviderKind::new(kind).expect("provider kind"), + } + } +} + +#[async_trait] +impl ProviderAdmin for UnusedProvider { + fn provider_kind(&self) -> &ProviderKind { + &self.kind + } + + async fn account_unavailable(&self, _: &ProviderAccountId) {} + + fn dashboard_wire_profile( + &self, + ) -> Option { + None + } + + fn calculated_billing( + &self, + _: &gateway_admin::model::observability::ProviderBillingInput, + ) -> Result< + Option, + ProviderAdminError, + > { + Ok(None) + } + + async fn prepare_import( + &self, + _: PrepareCredentialImport, + ) -> Result { + Err(unsupported_provider()) + } + + async fn start_authorization( + &self, + _: PendingAuthorizationMutation, + ) -> Result { + Err(unsupported_provider()) + } + + async fn complete_authorization( + &self, + _: CompleteAuthorization, + ) -> Result { + Err(unsupported_provider()) + } + + async fn prepare_rotation( + &self, + _: PrepareCredentialRotation, + ) -> Result { + Err(unsupported_provider()) + } + + async fn prepare_refresh( + &self, + _: PrepareCredentialRefresh, + ) -> Result { + Err(unsupported_provider()) + } + + async fn quota( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Err(unsupported_provider()) + } + + async fn models( + &self, + _: &ProviderAccountId, + _: bool, + ) -> Result { + Err(unsupported_provider()) + } + + async fn export_credentials( + &self, + _: Vec, + ) -> Result { + Err(unsupported_provider()) + } +} + +struct NoopSnapshot; + +impl SnapshotControl for NoopSnapshot { + fn publish_committed(&self, _: ConfigRevision) -> BoxFuture<'_, ()> { + Box::pin(async {}) + } +} + +struct NoopProbe; + +impl AccountProbe for NoopProbe { + fn probe( + &self, + _: AccountProbeRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { panic!("unused account probe") }) + } +} + +struct UnusedSystem; + +#[async_trait] +impl SystemOperations for UnusedSystem { + async fn version(&self) -> Result { + Err(unavailable_system()) + } + + async fn update_detail(&self, _: bool) -> Result { + Err(unavailable_system()) + } + + fn update_events(&self) -> SystemUpdateEventStream { + Box::pin(futures::stream::empty()) + } + + async fn perform_update( + &self, + _: Option, + ) -> Result { + Err(unavailable_system()) + } + + async fn update_status(&self) -> Result { + Err(unavailable_system()) + } + + async fn rollback(&self) -> Result { + Err(unavailable_system()) + } + + async fn restart(&self) -> Result { + Err(unavailable_system()) + } +} + +fn test_runtime_settings() -> RuntimeSettings { + let mappings: ProviderModelMappings = BTreeMap::from([ + ( + ProviderKind::new("openai").expect("provider kind"), + BTreeMap::from([( + PublicModelId::new("coding-default").expect("public model"), + UpstreamModelId::new("gpt-5.4").expect("upstream model"), + )]), + ), + ( + ProviderKind::new("xai").expect("provider kind"), + BTreeMap::from([( + PublicModelId::new("grok-latest").expect("public model"), + UpstreamModelId::new("grok-4.5").expect("upstream model"), + )]), + ), + ]); + RuntimeSettings { + config_revision: Revision::new(7).expect("revision"), + provider_model_mappings: mappings, + refresh_margin_seconds: 3_600, + refresh_concurrency: 2, + max_concurrent_per_account: 3, + request_interval_ms: 50, + rotation_strategy: RotationStrategy::Smart, + usage_retention_days: 31, + ops_event_retention_days: 30, + audit_retention_days: 90, + updated_at: Utc::now(), + } +} + +fn next_revision(revision: Revision) -> Revision { + Revision::new(revision.get().saturating_add(1)).expect("next revision") +} + +fn unavailable(resource: &'static str) -> AdminStoreError { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + resource, + "unused test port", + ) +} + +fn unsupported_provider() -> ProviderAdminError { + ProviderAdminError::new(ProviderAdminErrorKind::Unsupported) +} + +fn unavailable_system() -> SystemOperationError { + SystemOperationError::new(SystemOperationErrorKind::Internal, "unused test operation") +} diff --git a/backend/crates/gateway-api/tests/admin/observability.rs b/backend/crates/gateway-api/tests/admin/observability.rs index b38dc4c41..36cb1f705 100644 --- a/backend/crates/gateway-api/tests/admin/observability.rs +++ b/backend/crates/gateway-api/tests/admin/observability.rs @@ -65,9 +65,10 @@ mod query { mod response { use chrono::{TimeZone, Utc}; + use gateway_admin::model::observability::DesktopReleaseStatus; use gateway_api::admin::observability::{ - BillingView, CostCoverageView, CursorWire, PageData, PageMeta, TokenDetailsView, TrendData, - TrendKind, TrendPointView, TrendSummaryView, + BillingView, CostCoverageView, CursorWire, DashboardDesktopReleaseStatusView, PageData, + PageMeta, TokenDetailsView, TrendData, TrendKind, TrendPointView, TrendSummaryView, }; use serde_json::json; @@ -180,4 +181,46 @@ mod response { assert_eq!(value["serviceTierDisplay"], "Fast"); assert_eq!(value["multiplierDisplay"], "1.00x"); } + + #[test] + fn desktop_release_status_should_preserve_the_existing_dashboard_wire_values() { + for (domain, expected) in [ + (DesktopReleaseStatus::Unchecked, "unchecked"), + (DesktopReleaseStatus::Current, "aligned"), + (DesktopReleaseStatus::UpdateAvailable, "review_required"), + (DesktopReleaseStatus::Failed, "check_failed"), + ] { + let status = DashboardDesktopReleaseStatusView::from(domain); + assert_eq!(serde_json::to_value(status).unwrap(), expected); + } + } +} + +#[tokio::test] +async fn usage_route_should_forward_a_bounded_unknown_outcome_filter() { + use axum::{ + body::Body, + http::{Request, StatusCode, header}, + }; + use gateway_api::admin::observability; + use tower::ServiceExt as _; + + use super::{AdminTestFixture, AdminTestState}; + + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = observability::router::() + .with_state(fixture.state()) + .oneshot( + Request::builder() + .uri("/api/admin/usage/records?outcome=provider_future_state") + .header(header::COOKIE, "cpr_admin_session=valid-session") + .header("x-request-id", "req_usage_other_outcome") + .body(Body::empty()) + .expect("usage request"), + ) + .await + .expect("usage response"); + + assert_eq!(response.status(), StatusCode::OK); } diff --git a/backend/crates/gateway-api/tests/admin/openai.rs b/backend/crates/gateway-api/tests/admin/openai.rs index 1e6b820e5..61484e1fc 100644 --- a/backend/crates/gateway-api/tests/admin/openai.rs +++ b/backend/crates/gateway-api/tests/admin/openai.rs @@ -1,8 +1,14 @@ -use gateway_api::admin::openai::{ - CompleteOAuthAuthorizationRequest, CredentialDetailsQuery, CredentialMutationRequest, - ImportCredentialsDocumentRequest, ListCredentialsQuery, RotateCredentialRequest, - StartOAuthAuthorizationRequest, +use chrono::{TimeZone, Utc}; +use gateway_admin::model::{ + Revision, + accounts::{AccountAvailability, AccountRecord}, }; +use gateway_api::admin::openai::{ + CodexCredentialView, CompleteOAuthAuthorizationRequest, CredentialDetailsQuery, + CredentialMutationRequest, ImportCredentialsDocumentRequest, ListCredentialsQuery, + RotateCredentialRequest, StartOAuthAuthorizationRequest, +}; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; fn compact_jwt() -> String { "header.payload.signature".to_owned() @@ -159,3 +165,38 @@ fn oauth_reauthorization_requires_account_and_revision_together() { "reauthorization" ); } + +#[test] +fn credential_result_should_map_to_the_existing_safe_wire_view() { + let timestamp = Utc + .with_ymd_and_hms(2026, 7, 18, 3, 0, 0) + .single() + .expect("valid fixture timestamp"); + let view = CodexCredentialView::from(AccountRecord { + id: "acct_openai_safe".to_owned(), + provider_instance_id: ProviderInstanceId::new("inst_openai").expect("provider instance ID"), + provider_kind: ProviderKind::new("openai").expect("provider kind"), + name: "Codex OAuth".to_owned(), + email: Some("verified@example.com".to_owned()), + upstream_user_id: "subject_openai".to_owned(), + upstream_account_id: None, + plan_type: Some("pro".to_owned()), + credential_revision: Revision::new(2).expect("credential revision"), + has_refresh_token: true, + access_token_expires_at: timestamp, + next_refresh_at: Some(timestamp), + enabled: true, + availability: AccountAvailability::Banned, + availability_reason: Some("account disabled".to_owned()), + cooldown_until: None, + availability_observed_at: timestamp, + quota_observed_at: None, + created_at: timestamp, + updated_at: timestamp, + }); + + let rendered = serde_json::to_value(view).expect("serialize safe credential view"); + assert_eq!(rendered["availability"], "invalid"); + assert_eq!(rendered["credentialRevision"], 2); + assert!(rendered.get("accessToken").is_none()); +} diff --git a/backend/crates/gateway-api/tests/admin/settings.rs b/backend/crates/gateway-api/tests/admin/settings.rs index 4635ca312..f7352d429 100644 --- a/backend/crates/gateway-api/tests/admin/settings.rs +++ b/backend/crates/gateway-api/tests/admin/settings.rs @@ -1,118 +1,16 @@ -use std::{ - collections::BTreeMap, - sync::{Arc, Mutex}, -}; - -use async_trait::async_trait; use axum::{ Router, body::{Body, to_bytes}, http::{Method, Request, StatusCode, header}, }; -use chrono::Utc; -use gateway_api::admin::{ - AdminRequestContext, AdminServiceError, AdminSessionResolver, AdminSessionState, - RuntimeSettingsView, UpdateRuntimeSettingsRequest, - settings::{self, AdminSettingsService, AdminSettingsState}, -}; +use gateway_api::admin::settings::{self, UpdateRuntimeSettingsRequest}; use serde_json::{Value, json}; use tower::ServiceExt; -struct FakeSettings { - settings: Mutex, - key_exists: Mutex, -} +use super::{AdminTestFixture, AdminTestState}; -impl Default for FakeSettings { - fn default() -> Self { - Self { - settings: Mutex::new(runtime_settings()), - key_exists: Mutex::new(false), - } - } -} - -#[async_trait] -impl AdminSessionResolver for FakeSettings { - async fn resolve_admin_user_id( - &self, - session_id: Option<&str>, - ) -> Result, AdminServiceError> { - Ok((session_id == Some("valid-session")).then(|| "admin_1".to_owned())) - } - - async fn verify_admin_api_key(&self, key: &str) -> Result { - Ok(key == "admin-valid-test-key") - } -} - -#[async_trait] -impl AdminSettingsService for FakeSettings { - async fn load(&self) -> Result { - Ok(self.settings.lock().expect("settings state").clone()) - } - - async fn replace( - &self, - context: &AdminRequestContext, - request: UpdateRuntimeSettingsRequest, - ) -> Result { - assert_eq!(context.admin_user_id(), Some("admin_1")); - let settings = RuntimeSettingsView { - config_revision: request.expected_config_revision + 1, - provider_model_mappings: request.provider_model_mappings, - refresh_margin_seconds: request.refresh_margin_seconds, - refresh_concurrency: request.refresh_concurrency, - max_concurrent_per_account: request.max_concurrent_per_account, - request_interval_ms: request.request_interval_ms, - rotation_strategy: request.rotation_strategy, - usage_retention_days: request.usage_retention_days, - ops_event_retention_days: request.ops_event_retention_days, - audit_retention_days: request.audit_retention_days, - updated_at: Utc::now(), - }; - *self.settings.lock().expect("settings state") = settings.clone(); - Ok(settings) - } - - async fn admin_api_key_exists(&self) -> Result { - Ok(*self.key_exists.lock().expect("key state")) - } - - async fn regenerate_admin_api_key( - &self, - _context: &AdminRequestContext, - ) -> Result { - *self.key_exists.lock().expect("key state") = true; - Ok("admin-generated-secret-once".to_owned()) - } - - async fn delete_admin_api_key( - &self, - _context: &AdminRequestContext, - ) -> Result<(), AdminServiceError> { - *self.key_exists.lock().expect("key state") = false; - Ok(()) - } -} - -#[derive(Clone)] -struct TestState(Arc); - -impl AdminSessionState for TestState { - fn admin_session_resolver(&self) -> &dyn AdminSessionResolver { - self.0.as_ref() - } -} - -impl AdminSettingsState for TestState { - fn admin_settings_service(&self) -> &dyn AdminSettingsService { - self.0.as_ref() - } -} - -fn app(settings: Arc) -> Router { - settings::router::().with_state(TestState(settings)) +fn app(state: AdminTestState) -> Router { + settings::router::().with_state(state) } fn request(method: Method, path: &str, body: Option) -> Request { @@ -137,31 +35,6 @@ async fn response_json(response: axum::response::Response) -> Value { serde_json::from_slice(&bytes).expect("parse response JSON") } -fn runtime_settings() -> RuntimeSettingsView { - RuntimeSettingsView { - config_revision: 7, - provider_model_mappings: BTreeMap::from([ - ( - "openai".to_owned(), - BTreeMap::from([("coding-default".to_owned(), "gpt-5.4".to_owned())]), - ), - ( - "xai".to_owned(), - BTreeMap::from([("grok-latest".to_owned(), "grok-4.5".to_owned())]), - ), - ]), - refresh_margin_seconds: 3600, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "smart".into(), - usage_retention_days: 31, - ops_event_retention_days: 30, - audit_retention_days: 90, - updated_at: Utc::now(), - } -} - fn update_body() -> Value { json!({ "expectedConfigRevision": 7, @@ -219,7 +92,9 @@ fn settings_request_should_reject_removed_bucket_retention() { #[tokio::test] async fn settings_get_should_preserve_provider_scoped_model_mappings() { - let response = app(Arc::new(FakeSettings::default())) + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = app(fixture.state()) .oneshot(request(Method::GET, "/api/admin/settings", None)) .await .expect("settings response"); @@ -237,7 +112,9 @@ async fn settings_get_should_preserve_provider_scoped_model_mappings() { #[tokio::test] async fn settings_post_should_replace_provider_scoped_model_mappings() { - let response = app(Arc::new(FakeSettings::default())) + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = app(fixture.state()) .oneshot(request( Method::POST, "/api/admin/settings", @@ -275,7 +152,8 @@ fn settings_request_should_reject_invalid_provider_mapping_slug() { #[tokio::test] async fn settings_should_require_admin_auth() { - let response = app(Arc::new(FakeSettings::default())) + let fixture = AdminTestFixture::new().await; + let response = app(fixture.state()) .oneshot( Request::builder() .uri("/api/admin/settings") @@ -291,7 +169,9 @@ async fn settings_should_require_admin_auth() { #[tokio::test] async fn admin_key_should_return_secret_only_on_regenerate() { - let response = app(Arc::new(FakeSettings::default())) + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let response = app(fixture.state()) .oneshot(request( Method::POST, "/api/admin/settings/admin-api-key/regenerate", @@ -301,14 +181,19 @@ async fn admin_key_should_return_secret_only_on_regenerate() { .expect("regenerate response"); let data = response_json(response).await["data"].clone(); - assert_eq!(data["key"], "admin-generated-secret-once"); + assert!( + data["key"] + .as_str() + .is_some_and(|key| key.starts_with("admin-") && key.len() == 70) + ); } #[tokio::test] async fn admin_key_delete_should_use_fixed_post_path() { - let settings = Arc::new(FakeSettings::default()); - *settings.key_exists.lock().expect("key state") = true; - let response = app(settings) + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + fixture.settings.set_api_key("admin-valid-test-key"); + let response = app(fixture.state()) .oneshot(request( Method::POST, "/api/admin/settings/admin-api-key", @@ -322,11 +207,14 @@ async fn admin_key_delete_should_use_fixed_post_path() { #[tokio::test] async fn settings_should_accept_admin_api_key_header() { - let response = app(Arc::new(FakeSettings::default())) + let fixture = AdminTestFixture::new().await; + let key = format!("admin-{}", "a".repeat(64)); + fixture.auth.set_api_key(&key); + let response = app(fixture.state()) .oneshot( Request::builder() .uri("/api/admin/settings") - .header("x-api-key", "admin-valid-test-key") + .header("x-api-key", key) .header("x-request-id", "req_api_key") .body(Body::empty()) .expect("api key request"), diff --git a/backend/crates/gateway-api/tests/admin/xai.rs b/backend/crates/gateway-api/tests/admin/xai.rs index a56bd2acb..4c7df596f 100644 --- a/backend/crates/gateway-api/tests/admin/xai.rs +++ b/backend/crates/gateway-api/tests/admin/xai.rs @@ -1,9 +1,14 @@ use chrono::{TimeZone, Utc}; +use gateway_admin::model::{ + Revision, + accounts::{AccountAvailability, AccountRecord}, +}; use gateway_api::admin::xai::{ AuthorizationStartData, CompleteAuthorizationRequest, NewCredentialRequest, StartAuthorizationRequest, XaiCredentialImportData, XaiCredentialImportDocumentRequest, XaiCredentialListData, XaiCredentialViewData, }; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use serde_json::{Value, json}; fn valid_request() -> Value { @@ -166,3 +171,38 @@ fn authorization_requests_require_server_owned_flow_and_full_callback_url() { .expect("deserialize callback completion"); complete.validate().expect("callback completion is valid"); } + +#[test] +fn credential_result_should_map_to_the_existing_safe_wire_view() { + let timestamp = Utc + .with_ymd_and_hms(2026, 7, 18, 3, 0, 0) + .single() + .expect("valid fixture timestamp"); + let view = XaiCredentialViewData::from(AccountRecord { + id: "acct_xai_safe".to_owned(), + provider_instance_id: ProviderInstanceId::new("inst_xai").expect("provider instance ID"), + provider_kind: ProviderKind::new("xai").expect("provider kind"), + name: "Grok OAuth".to_owned(), + email: Some("verified@example.com".to_owned()), + upstream_user_id: "subject_xai".to_owned(), + upstream_account_id: None, + plan_type: Some("pro".to_owned()), + credential_revision: Revision::new(2).expect("credential revision"), + has_refresh_token: true, + access_token_expires_at: timestamp, + next_refresh_at: Some(timestamp), + enabled: true, + availability: AccountAvailability::QuotaExhausted, + availability_reason: None, + cooldown_until: None, + availability_observed_at: timestamp, + quota_observed_at: None, + created_at: timestamp, + updated_at: timestamp, + }); + + let rendered = serde_json::to_value(view).expect("serialize safe credential view"); + assert_eq!(rendered["availability"], "quota_exhausted"); + assert_eq!(rendered["credentialRevision"], 2); + assert!(rendered.get("refreshToken").is_none()); +} diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index 0c7cc63e1..3c0b26f01 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -44,6 +44,7 @@ fn source_tree_should_match_frozen_machine_manifest() { "src/admin/system.rs", "src/admin/wire.rs", "src/admin/xai.rs", + "src/health.rs", "src/lib.rs", "src/openai/auth.rs", "src/openai/error.rs", @@ -81,6 +82,7 @@ fn test_tree_should_match_frozen_rust_mirror() { "tests/admin/wire.rs", "tests/admin/xai.rs", "tests/architecture.rs", + "tests/health.rs", "tests/main.rs", "tests/openai/auth.rs", "tests/openai/error.rs", diff --git a/backend/crates/gateway-api/tests/health.rs b/backend/crates/gateway-api/tests/health.rs new file mode 100644 index 000000000..f5cf2b0ab --- /dev/null +++ b/backend/crates/gateway-api/tests/health.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use futures::future::BoxFuture; +use gateway_core::{ + engine::execution::{ + AuthenticatedClient, ClientAuthenticationError, ExecutionService, StartExecution, + StartedExecution, + }, + error::GatewayError, + routing::PublicModelId, +}; +use tower::ServiceExt; + +#[tokio::test] +async fn healthz_should_return_no_content_when_all_inputs_are_healthy() { + let response = crate::openai::api_router(Arc::new(UnusedExecution)) + .await + .oneshot( + Request::builder() + .uri("/healthz") + .body(Body::empty()) + .expect("health request"), + ) + .await + .expect("health response"); + + assert_eq!(response.status(), StatusCode::NO_CONTENT); +} + +struct UnusedExecution; + +impl ExecutionService for UnusedExecution { + fn authenticate(&self, _: &str) -> Result { + unreachable!("health check does not authenticate") + } + + fn public_models(&self, _: &AuthenticatedClient) -> Vec { + unreachable!("health check does not list models") + } + + fn contains_public_model(&self, _: &AuthenticatedClient, _: &PublicModelId) -> bool { + unreachable!("health check does not inspect models") + } + + fn start(&self, _: StartExecution) -> BoxFuture<'_, Result> { + Box::pin(async { unreachable!("health check does not execute requests") }) + } +} diff --git a/backend/crates/gateway-api/tests/main.rs b/backend/crates/gateway-api/tests/main.rs index 2395aa526..7a1a93ab5 100644 --- a/backend/crates/gateway-api/tests/main.rs +++ b/backend/crates/gateway-api/tests/main.rs @@ -1,3 +1,4 @@ mod admin; mod architecture; +mod health; mod openai; diff --git a/backend/crates/gateway-api/tests/openai/mod.rs b/backend/crates/gateway-api/tests/openai/mod.rs index ddb78c14d..a263339c2 100644 --- a/backend/crates/gateway-api/tests/openai/mod.rs +++ b/backend/crates/gateway-api/tests/openai/mod.rs @@ -3,3 +3,253 @@ mod error; mod models; mod responses; mod router; + +use std::collections::BTreeSet; +use std::num::NonZeroU32; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; + +use async_trait::async_trait; +use futures::future::BoxFuture; +use gateway_core::engine::admission::{ + ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery, + ClientAdmissionRequest, ClientAdmissionRestoreResult, +}; +use gateway_core::engine::continuation::{ + NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId, +}; +use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy}; +use gateway_core::engine::execution::{ + AuthenticatedClient, DefaultExecutionService, ExecutionService, ProviderCircuitDecision, + ProviderCircuitError, ProviderCircuitPort, +}; +use gateway_core::engine::provider::ProviderRegistry; +use gateway_core::engine::{ + AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, + NewModelRequest, RecoveryReport, UpstreamSendState, +}; +use gateway_core::error::StoreError; +use gateway_core::health::{WorkerHealthSnapshot, WorkerHealthSource}; +use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle}; +use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::snapshot::RuntimeSnapshotHandle; +use gateway_core::routing::{ + ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId, + ProviderKind, ProviderModel, RuntimeSnapshot, UpstreamModelId, +}; + +pub(super) async fn api_router(execution: Arc) -> axum::Router { + let admin = crate::admin::AdminTestFixture::new().await; + gateway_api::initialize( + gateway_api::ApiConfig { + asset_directory: std::env::temp_dir(), + cors_allowed_origins: Vec::new(), + request_timeout_seconds: None, + request_id_header: "x-request-id".to_owned(), + }, + execution, + admin.services, + Vec::new(), + Arc::new(EmptyWorkerHealth), + Arc::new(TestLifecycle::default()), + ) + .expect("API bundle") + .router() +} + +pub(super) fn authenticated_client(plaintext: &str) -> AuthenticatedClient { + let source = DefaultExecutionService::new( + RuntimeSnapshotHandle::new(snapshot(plaintext)), + Arc::new(UnusedExecutionStore), + ProviderRegistry::default(), + Arc::new(UnusedAdmissions), + Arc::new(UnusedCircuits), + Arc::new(UnusedContinuation), + ); + source + .authenticate(plaintext) + .expect("authenticated client") +} + +fn snapshot(plaintext: &str) -> RuntimeSnapshot { + let provider = ProviderKind::new("openai").expect("provider"); + let instance_id = ProviderInstanceId::new("inst_openai_api_test").expect("instance"); + let capabilities = ModelCapabilities::new( + BTreeSet::from([gateway_core::operation::OperationKind::Generate]), + 128_000, + Some(16_000), + ); + RuntimeSnapshot::new( + ConfigRevision::new(1).expect("revision"), + AccountSelectionPolicy::new( + RotationStrategy::Smart, + NonZeroU32::new(2).expect("concurrency"), + Duration::from_millis(1), + ), + vec![ProviderInstance::new( + instance_id.clone(), + provider.clone(), + "https://api.example.invalid".to_owned(), + true, + InstanceHealth::Healthy, + )], + ["model-a", "model-b"] + .into_iter() + .map(|model| { + ProviderModel::new( + instance_id.clone(), + UpstreamModelId::new(model).expect("model"), + capabilities.clone(), + ) + }) + .collect(), + vec![ClientPolicy::new( + ClientApiKeyId::new("key_api_test").expect("key ID"), + PlaintextClientApiKey::new(plaintext).expect("plaintext key"), + provider, + true, + RateLimits::unlimited(), + )], + ) + .expect("runtime snapshot") +} + +#[derive(Default)] +struct TestLifecycle { + cancellation: gateway_core::engine::CancellationToken, +} + +struct TestConnectionGuard; + +impl ConnectionGuard for TestConnectionGuard {} + +impl ConnectionLifecycle for TestLifecycle { + fn try_register(&self) -> Result, ConnectionDraining> { + Ok(Box::new(TestConnectionGuard)) + } + + fn cancellation(&self) -> gateway_core::engine::CancellationToken { + self.cancellation.clone() + } + + fn is_draining(&self) -> bool { + self.cancellation.is_cancelled() + } +} + +struct EmptyWorkerHealth; + +impl WorkerHealthSource for EmptyWorkerHealth { + fn snapshot(&self) -> Vec { + Vec::new() + } +} + +struct UnusedExecutionStore; + +#[async_trait] +impl ExecutionStore for UnusedExecutionStore { + async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn mark_send_state( + &self, + _: &ModelRequestId, + _: UpstreamSendState, + ) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn mark_downstream_committed( + &self, + _: &ModelRequestId, + _: SystemTime, + _: Option, + ) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> { + unreachable!("authentication fixture does not execute") + } + + async fn recover_expired(&self, _: SystemTime) -> Result { + unreachable!("authentication fixture does not execute") + } +} + +struct UnusedAdmissions; + +impl ClientAdmissionPort for UnusedAdmissions { + fn admit( + &self, + _: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } + + fn release<'a>( + &'a self, + _: &'a ClientApiKeyId, + _: &'a ModelRequestId, + ) -> BoxFuture<'a, Result> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } + + fn restore( + &self, + _: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } +} + +struct UnusedCircuits; + +impl ProviderCircuitPort for UnusedCircuits { + fn decision<'a>( + &'a self, + _: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } + + fn observe_failure<'a>( + &'a self, + _: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } + + fn observe_success<'a>( + &'a self, + _: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } +} + +struct UnusedContinuation; + +impl NativeContinuationPort for UnusedContinuation { + fn resolve<'a>( + &'a self, + _: &'a ClientApiKeyId, + _: &'a PreviousResponseId, + ) -> BoxFuture<'a, Result, NativeContinuationStoreError>> { + Box::pin(async { unreachable!("authentication fixture does not execute") }) + } +} diff --git a/backend/crates/gateway-api/tests/openai/models.rs b/backend/crates/gateway-api/tests/openai/models.rs index 77052aa62..b2f64fea1 100644 --- a/backend/crates/gateway-api/tests/openai/models.rs +++ b/backend/crates/gateway-api/tests/openai/models.rs @@ -1,116 +1,62 @@ -use std::collections::BTreeSet; +use std::sync::Arc; -use async_trait::async_trait; use axum::{ body::{Body, to_bytes}, http::{Request, StatusCode, header::AUTHORIZATION}, }; -use gateway_api::openai::{ - ConnectionTask, DeliveryEvent, OpenAiApiState, OpenAiClientService, ResponseExecutionSession, - ResponsesTransport, StartedResponse, auth::ClientApiKeyAuthError, - responses::DecodedResponsesRequest, -}; -use gateway_core::{ - engine::EngineError, - error::{GatewayError, GatewayErrorKind}, - event::GatewayEvent, - routing::PublicModelId, +use futures::future::BoxFuture; +use gateway_core::engine::execution::{ + AuthenticatedClient, ClientAuthenticationError, ExecutionService, StartExecution, + StartedExecution, }; +use gateway_core::error::{GatewayError, GatewayErrorKind}; +use gateway_core::routing::PublicModelId; use tower::ServiceExt; -#[derive(Clone)] -struct ModelsState; +use super::{api_router, authenticated_client}; -#[derive(Clone)] -struct ModelsService; - -struct UnusedSession; - -#[async_trait] -impl ResponseExecutionSession for UnusedSession { - async fn next_delivery_event(&mut self) -> Result, EngineError> { - unreachable!("models endpoints do not start an execution") - } - - async fn collect_uncommitted(&mut self) -> Result, EngineError> { - unreachable!("models endpoints do not start an execution") - } - - async fn commit_downstream( - &mut self, - _client_status_code: Option, - ) -> Result<(), EngineError> { - unreachable!("models endpoints do not start an execution") - } - - async fn record_client_status(&mut self, _client_status_code: u16) -> Result<(), EngineError> { - unreachable!("models endpoints do not start an execution") - } - - fn is_finalized(&self) -> bool { - true - } - - fn cancel(&self) {} - - fn detach_finalize(self) {} +struct ModelsExecution { + client: AuthenticatedClient, } -#[async_trait] -impl OpenAiClientService for ModelsService { - type Client = (); - type Session = UnusedSession; +impl ModelsExecution { + fn new() -> Arc { + Arc::new(Self { + client: authenticated_client("sk_models_test"), + }) + } +} - fn authenticate(&self, plaintext: &str) -> Result { +impl ExecutionService for ModelsExecution { + fn authenticate( + &self, + plaintext: &str, + ) -> Result { if plaintext == "sk_models_test" { - Ok(()) + Ok(self.client.clone()) } else { - Err(ClientApiKeyAuthError::InvalidKey) + Err(ClientAuthenticationError::InvalidKey) } } - fn public_models(&self, _client: &Self::Client) -> Vec { - vec!["model-a".to_owned(), "model-b".to_owned()] + fn public_models(&self, _: &AuthenticatedClient) -> Vec { + ["model-a", "model-b"] + .into_iter() + .map(|model| PublicModelId::new(model).expect("model")) + .collect() } - fn contains_public_model(&self, _client: &Self::Client, model: &PublicModelId) -> bool { - BTreeSet::from(["model-a", "model-b"]).contains(model.as_str()) + fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool { + matches!(model.as_str(), "model-a" | "model-b") } - async fn start_response( - &self, - _client: Self::Client, - _request: DecodedResponsesRequest, - _transport: ResponsesTransport, - ) -> Result, GatewayError> { - Err(GatewayError::new( - GatewayErrorKind::Internal, - "models test must not start a response", - )) - } - - fn is_shutting_down(&self) -> bool { - false - } - - fn spawn_connection(&self, task: ConnectionTask) { - drop(task); - } - - fn next_connection_id(&self) -> String { - "ws_models_test".to_owned() - } - - fn next_request_id(&self) -> String { - "req_models_test".to_owned() - } -} - -impl OpenAiApiState for ModelsState { - type Service = ModelsService; - - fn openai_client_api(&self) -> Self::Service { - ModelsService + fn start(&self, _: StartExecution) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(GatewayError::new( + GatewayErrorKind::Internal, + "models test must not start a response", + )) + }) } } @@ -123,8 +69,8 @@ fn authorized_request(path: &str) -> Request { #[tokio::test] async fn models_should_encode_the_service_visible_catalog() { - let response = gateway_api::openai::router::() - .with_state(ModelsState) + let response = api_router(ModelsExecution::new()) + .await .oneshot(authorized_request("/v1/models")) .await .expect("list models response"); @@ -147,8 +93,8 @@ async fn models_should_encode_the_service_visible_catalog() { #[tokio::test] async fn model_detail_should_keep_the_official_path_id_contract() { - let response = gateway_api::openai::router::() - .with_state(ModelsState) + let response = api_router(ModelsExecution::new()) + .await .oneshot(authorized_request("/v1/models/model-a")) .await .expect("model detail response"); @@ -158,8 +104,8 @@ async fn model_detail_should_keep_the_official_path_id_contract() { #[tokio::test] async fn model_detail_should_hide_unknown_models() { - let response = gateway_api::openai::router::() - .with_state(ModelsState) + let response = api_router(ModelsExecution::new()) + .await .oneshot(authorized_request("/v1/models/model-private")) .await .expect("unknown model response"); @@ -169,8 +115,8 @@ async fn model_detail_should_hide_unknown_models() { #[tokio::test] async fn models_should_require_a_valid_bearer_client_key() { - let response = gateway_api::openai::router::() - .with_state(ModelsState) + let response = api_router(ModelsExecution::new()) + .await .oneshot( Request::get("/v1/models") .body(Body::empty()) diff --git a/backend/crates/gateway-api/tests/openai/responses/http.rs b/backend/crates/gateway-api/tests/openai/responses/http.rs index 1755228e6..e255be238 100644 --- a/backend/crates/gateway-api/tests/openai/responses/http.rs +++ b/backend/crates/gateway-api/tests/openai/responses/http.rs @@ -1,29 +1,36 @@ use std::collections::VecDeque; -use std::future::pending; use std::net::{IpAddr, SocketAddr}; use std::sync::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, }; -use async_trait::async_trait; use axum::{ - body::{Bytes, to_bytes}, - extract::{Extension, State, connect_info::ConnectInfo}, - http::{HeaderMap, StatusCode, header::AUTHORIZATION}, + body::{Body, to_bytes}, + extract::connect_info::ConnectInfo, + http::{HeaderMap, Request, StatusCode, header::AUTHORIZATION}, }; -use gateway_core::engine::{CommitRequirement, EngineError, UpstreamSendState}; -use gateway_core::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind}; -use gateway_core::event::{ContentItem, ContentKind, GatewayEvent, ResponseMeta, TextDelta}; +use futures::future::{BoxFuture, pending}; +use gateway_core::engine::execution::{ + AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession, + StartExecution, StartedExecution, +}; +use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, UpstreamSendState}; +use gateway_core::error::{ + GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind, SafeUpstreamValue, +}; +use gateway_core::event::{ + ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent, + ProviderResponseHeader, ResponseMeta, TextDelta, +}; +use gateway_core::operation::Operation; use gateway_core::routing::PublicModelId; +use serde_json::{Value, json}; -use gateway_api::openai::responses::{ - DecodedResponsesRequest, collect_execution_response, responses, stream_execution_response, -}; -use gateway_api::openai::{ - ConnectionTask, DeliveryEvent, OpenAiApiState, OpenAiClientService, ResponseExecutionSession, - ResponsesTransport, StartedResponse, auth::ClientApiKeyAuthError, -}; +use gateway_api::openai::responses::{collect_execution_response, stream_execution_response}; +use tower::ServiceExt; + +use crate::openai::{api_router, authenticated_client}; #[derive(Default)] struct Trace { @@ -65,7 +72,7 @@ impl Trace { } enum NextStep { - Event(DeliveryEvent), + Event(CoordinatedEvent), Error(EngineError), FinalizeCancelled, FinalizeSuccess, @@ -75,79 +82,81 @@ enum NextStep { struct FakeSession { trace: Arc, next: VecDeque, - collected: Option>, + collected: Option>, collect_error: Option, collect_pending: bool, finalize_on_commit: bool, fail_commit: bool, finalized: bool, + response_headers: Vec, } #[derive(Clone)] -struct ContextCaptureState { +struct ContextCaptureExecution { observed: Arc>>, + client: AuthenticatedClient, } -type CapturedClientContext = (Option, Option); +#[derive(Debug, Clone, PartialEq)] +struct CapturedClientContext { + client_ip: Option, + user_agent: Option, + endpoint: String, + client_metadata: Option, + openai_options: Option, +} -#[async_trait] -impl OpenAiClientService for ContextCaptureState { - type Client = (); - type Session = FakeSession; - - fn authenticate(&self, plaintext: &str) -> Result { - (plaintext == "sk_context_test") - .then_some(()) - .ok_or(ClientApiKeyAuthError::InvalidKey) +impl ExecutionService for ContextCaptureExecution { + fn authenticate( + &self, + plaintext: &str, + ) -> Result { + if plaintext == "sk_context_test" { + Ok(self.client.clone()) + } else { + Err(ClientAuthenticationError::InvalidKey) + } } - fn public_models(&self, _client: &Self::Client) -> Vec { + fn public_models(&self, _: &AuthenticatedClient) -> Vec { Vec::new() } - fn contains_public_model(&self, _client: &Self::Client, _model: &PublicModelId) -> bool { + fn contains_public_model(&self, _: &AuthenticatedClient, _: &PublicModelId) -> bool { true } - async fn start_response( + fn start( &self, - _client: Self::Client, - request: DecodedResponsesRequest, - _transport: ResponsesTransport, - ) -> Result, GatewayError> { - let metadata = request.metadata(); - *self.observed.lock().expect("context capture lock") = Some(( - metadata.client_ip(), - metadata.user_agent().map(str::to_owned), - )); - Err(GatewayError::new( - GatewayErrorKind::Internal, - "context capture completed", - )) - } - - fn is_shutting_down(&self) -> bool { - false - } - - fn spawn_connection(&self, task: ConnectionTask) { - drop(task); - } - - fn next_connection_id(&self) -> String { - "ws_context_test".to_owned() - } - - fn next_request_id(&self) -> String { - "req_context_test".to_owned() - } -} - -impl OpenAiApiState for ContextCaptureState { - type Service = Self; - - fn openai_client_api(&self) -> Self::Service { - self.clone() + request: StartExecution, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let (client_metadata, openai_options) = match &request.operation { + Operation::Generate(generate) => ( + generate + .protocol_payload() + .and_then(|payload| payload.body().get("client_metadata")) + .cloned(), + generate + .provider_options() + .get("openai") + .cloned() + .map(Value::Object), + ), + _ => (None, None), + }; + *self.observed.lock().expect("context capture lock") = Some(CapturedClientContext { + client_ip: request.metadata.client_ip, + user_agent: request.metadata.user_agent, + endpoint: request.metadata.endpoint, + client_metadata, + openai_options, + }); + Err(GatewayError::new( + GatewayErrorKind::Internal, + "context capture completed", + )) + }) } } @@ -162,6 +171,7 @@ impl FakeSession { finalize_on_commit: false, fail_commit: false, finalized: false, + response_headers: Vec::new(), } } @@ -169,12 +179,13 @@ impl FakeSession { Self { trace, next: VecDeque::from([NextStep::FinalizeCancelled]), - collected: Some(events), + collected: Some(events.into_iter().map(ProviderEvent::canonical).collect()), collect_error: None, collect_pending: false, finalize_on_commit: true, fail_commit: false, finalized: false, + response_headers: Vec::new(), } } @@ -188,6 +199,7 @@ impl FakeSession { finalize_on_commit: true, fail_commit: false, finalized: false, + response_headers: Vec::new(), } } @@ -201,68 +213,87 @@ impl FakeSession { self.collect_error = Some(error); self } + + fn with_response_headers(mut self, response_headers: Vec) -> Self { + self.response_headers = response_headers; + self + } } -#[async_trait] -impl ResponseExecutionSession for FakeSession { - async fn next_delivery_event(&mut self) -> Result, EngineError> { - match self.next.pop_front().unwrap_or(NextStep::End) { - NextStep::Event(event) => { - self.trace.push("next_event"); - Ok(Some(event)) +impl ExecutionSession for FakeSession { + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + match self.next.pop_front().unwrap_or(NextStep::End) { + NextStep::Event(event) => { + self.trace.push("next_event"); + Ok(Some(event)) + } + NextStep::Error(error) => { + self.trace.push("next_error"); + self.finalized = true; + Err(error) + } + NextStep::FinalizeCancelled => { + self.trace.push("cancel_finalize"); + self.finalized = true; + Err(EngineError::Cancelled) + } + NextStep::FinalizeSuccess => { + self.trace.push("next_end"); + self.finalized = true; + Ok(None) + } + NextStep::End => { + self.trace.push("next_end"); + Ok(None) + } } - NextStep::Error(error) => { - self.trace.push("next_error"); - self.finalized = true; - Err(error) - } - NextStep::FinalizeCancelled => { - self.trace.push("cancel_finalize"); - self.finalized = true; - Err(EngineError::Cancelled) - } - NextStep::FinalizeSuccess => { - self.trace.push("next_end"); - self.finalized = true; - Ok(None) - } - NextStep::End => { - self.trace.push("next_end"); - Ok(None) - } - } + }) } - async fn collect_uncommitted(&mut self) -> Result, EngineError> { - self.trace.push("collect"); - if self.collect_pending { - pending::<()>().await; - } - if let Some(error) = self.collect_error.take() { - self.finalized = true; - return Err(error); - } - Ok(self.collected.take().unwrap_or_default()) + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + self.trace.push("collect"); + if self.collect_pending { + pending::<()>().await; + } + if let Some(error) = self.collect_error.take() { + self.finalized = true; + return Err(error); + } + Ok(self.collected.take().unwrap_or_default()) + }) } - async fn commit_downstream( + fn response_headers(&self) -> &[ProviderResponseHeader] { + &self.response_headers + } + + fn commit_downstream( &mut self, client_status_code: Option, - ) -> Result<(), EngineError> { - self.trace.push("commit"); - if self.fail_commit { - return Err(EngineError::ProviderMetadataMismatch); - } - if let Some(status) = client_status_code { - self.trace.record_client_status(status); - } - self.finalized = self.finalize_on_commit; - Ok(()) + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + self.trace.push("commit"); + if self.fail_commit { + return Err(EngineError::ProviderMetadataMismatch); + } + if let Some(status) = client_status_code { + self.trace.record_client_status(status); + } + self.finalized = self.finalize_on_commit; + Ok(()) + }) } - async fn record_client_status(&mut self, client_status_code: u16) -> Result<(), EngineError> { - self.trace.record_client_status(client_status_code); - Ok(()) + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + self.trace.record_client_status(client_status_code); + Ok(()) + }) } fn is_finalized(&self) -> bool { @@ -273,13 +304,12 @@ impl ResponseExecutionSession for FakeSession { self.trace.cancel(); } - fn detach_finalize(mut self) { - if self.finalized { - return; - } - std::mem::drop(tokio::spawn(async move { - let _ = self.next_delivery_event().await; - })); + fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { + Box::pin(async move { + if !self.finalized { + let _ = self.next_event().await; + } + }) } } @@ -291,8 +321,8 @@ fn completed() -> GatewayEvent { GatewayEvent::Completed(ResponseMeta::new("resp_test", "public-model")) } -fn delivery(event: GatewayEvent, commit_requirement: CommitRequirement) -> DeliveryEvent { - DeliveryEvent::new(event, commit_requirement) +fn delivery(event: GatewayEvent, commit_requirement: CommitRequirement) -> CoordinatedEvent { + CoordinatedEvent::single(ProviderEvent::canonical(event), commit_requirement) } async fn captured_client_context( @@ -304,16 +334,20 @@ async fn captured_client_context( "Bearer sk_context_test".parse().expect("authorization"), ); let observed = Arc::new(Mutex::new(None)); - let state = ContextCaptureState { + let execution = Arc::new(ContextCaptureExecution { observed: Arc::clone(&observed), - }; - let response = responses::( - State(state), - Some(Extension(ConnectInfo(peer_address))), - headers, - Bytes::from_static(br#"{"model":"smart-code","input":"hello"}"#), - ) - .await; + client: authenticated_client("sk_context_test"), + }); + let mut request = Request::post("/v1/responses") + .body(Body::from(r#"{"model":"smart-code","input":"hello"}"#)) + .expect("context request"); + *request.headers_mut() = headers; + request.extensions_mut().insert(ConnectInfo(peer_address)); + let response = api_router(execution) + .await + .oneshot(request) + .await + .expect("context response"); assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); observed .lock() @@ -336,10 +370,13 @@ async fn request_context_should_resolve_forwarded_precedence_and_peer_fallback() headers.insert("user-agent", " Codex-CLI/1.0 ".parse().expect("user agent")); assert_eq!( captured_client_context(headers, peer).await, - ( - Some("198.51.100.1".parse().expect("expected IP")), - Some("Codex-CLI/1.0".to_owned()), - ) + CapturedClientContext { + client_ip: Some("198.51.100.1".parse().expect("expected IP")), + user_agent: Some("Codex-CLI/1.0".to_owned()), + endpoint: "/v1/responses".to_owned(), + client_metadata: None, + openai_options: None, + } ); let mut headers = HeaderMap::new(); @@ -348,16 +385,103 @@ async fn request_context_should_resolve_forwarded_precedence_and_peer_fallback() "10.0.0.2, 203.0.113.3".parse().expect("forwarded IPs"), ); assert_eq!( - captured_client_context(headers, peer).await.0, + captured_client_context(headers, peer).await.client_ip, Some("203.0.113.3".parse().expect("expected IP")) ); assert_eq!( - captured_client_context(HeaderMap::new(), peer).await.0, + captured_client_context(HeaderMap::new(), peer) + .await + .client_ip, Some("192.0.2.10".parse().expect("expected peer IP")) ); } +#[tokio::test] +async fn http_request_should_forward_safe_codex_headers_into_provider_options() { + let peer = "192.0.2.10:443".parse().expect("peer address"); + let mut headers = HeaderMap::new(); + headers.insert("x-codex-turn-state", "turn-state".parse().expect("header")); + headers.insert("conversation-id", "conversation-1".parse().expect("header")); + headers.insert("session-id", "session-1".parse().expect("header")); + headers.insert("x-openai-subagent", "compact".parse().expect("header")); + headers.insert( + "x-openai-internal-codex-responses-lite", + "true".parse().expect("header"), + ); + + let captured = captured_client_context(headers, peer).await; + let options = captured + .openai_options + .as_ref() + .and_then(Value::as_object) + .expect("OpenAI options"); + assert_eq!(options.get("turn_state"), Some(&json!("turn-state"))); + assert_eq!( + options.get("conversation_id"), + Some(&json!("conversation-1")) + ); + assert_eq!(options.get("session_id"), Some(&json!("session-1"))); + assert_eq!(options.get("responses_lite"), Some(&json!("true"))); + assert!(!options.contains_key("authorization")); + assert_eq!( + captured + .client_metadata + .as_ref() + .and_then(|metadata| metadata.get("x-openai-subagent")), + Some(&json!("compact")) + ); +} + +#[tokio::test] +async fn review_route_should_inject_review_subagent_and_record_its_endpoint() { + let observed = Arc::new(Mutex::new(None)); + let execution = Arc::new(ContextCaptureExecution { + observed: Arc::clone(&observed), + client: authenticated_client("sk_context_test"), + }); + let request = Request::post("/v1/responses/review") + .header(AUTHORIZATION, "Bearer sk_context_test") + .header("x-openai-subagent", "compact") + .body(Body::from( + json!({ + "model": "smart-code", + "input": "review this", + "client_metadata": {"existing": "preserved"} + }) + .to_string(), + )) + .expect("review request"); + + let response = api_router(execution) + .await + .oneshot(request) + .await + .expect("review response"); + let captured = observed + .lock() + .expect("review capture lock") + .clone() + .expect("captured review request"); + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(captured.endpoint, "/v1/responses/review"); + assert_eq!( + captured + .client_metadata + .as_ref() + .and_then(|metadata| metadata.get("x-openai-subagent")), + Some(&json!("review")) + ); + assert_eq!( + captured + .client_metadata + .as_ref() + .and_then(|metadata| metadata.get("existing")), + Some(&json!("preserved")) + ); +} + #[tokio::test] async fn streaming_encodes_first_frame_before_commit_and_http_delivery() { let trace = Arc::new(Trace::default()); @@ -369,7 +493,7 @@ async fn streaming_encodes_first_frame_before_commit_and_http_delivery() { ))], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; assert_eq!(response.status(), axum::http::StatusCode::OK); assert_eq!(trace.client_statuses(), vec![200]); @@ -378,6 +502,108 @@ async fn streaming_encodes_first_frame_before_commit_and_http_delivery() { std::mem::forget(response); } +#[tokio::test] +async fn streaming_response_forwards_only_allowlisted_provider_headers() { + let trace = Arc::new(Trace::default()); + let headers = vec![ + ProviderResponseHeader::new( + "x-models-etag", + SafeUpstreamValue::new("models-v2").expect("etag"), + ) + .expect("safe header"), + ProviderResponseHeader::new( + "authorization", + SafeUpstreamValue::new("should-not-cross-boundary").expect("header value"), + ) + .expect("syntactically safe header"), + ]; + let session = FakeSession::streaming( + Arc::clone(&trace), + vec![NextStep::Event(delivery( + started(), + CommitRequirement::CommitBeforeDelivery, + ))], + ) + .with_response_headers(headers); + + let response = stream_execution_response(Box::new(session), 1, None).await; + + assert_eq!( + response + .headers() + .get("x-models-etag") + .and_then(|value| value.to_str().ok()), + Some("models-v2") + ); + assert!(response.headers().get("authorization").is_none()); + std::mem::forget(response); +} + +#[tokio::test] +async fn streaming_commit_batch_encodes_pre_identity_wire_before_committing() { + let trace = Arc::new(Trace::default()); + let upstream_id = SafeUpstreamValue::new("resp_upstream").expect("upstream response ID"); + let metadata = + ResponseMeta::new("resp_gateway", "public-model").with_upstream_response_id(upstream_id); + let events = vec![ + ProviderEvent::wire( + ProtocolWireEvent::json( + "openai", + Some("response.future_metadata".to_owned()), + json!({"type":"response.future_metadata","opaque":true}), + ) + .expect("future wire event"), + ), + ProviderEvent::canonical_with_wire( + vec![GatewayEvent::Started(metadata.clone())], + ProtocolWireEvent::json( + "openai", + Some("response.created".to_owned()), + json!({ + "type":"response.created", + "response":{"id":"resp_upstream","model":"public-model","status":"in_progress","output":[]} + }), + ) + .expect("created wire event"), + ), + ProviderEvent::canonical_with_wire( + vec![GatewayEvent::Completed(metadata)], + ProtocolWireEvent::json( + "openai", + Some("response.completed".to_owned()), + json!({ + "type":"response.completed", + "response":{"id":"resp_upstream","model":"public-model","status":"completed","output":[]} + }), + ) + .expect("completed wire event"), + ), + ]; + let batch = CoordinatedEvent::try_batch(events, CommitRequirement::CommitBeforeDelivery) + .expect("non-empty commit batch"); + let session = FakeSession::streaming( + Arc::clone(&trace), + vec![NextStep::Event(batch), NextStep::FinalizeSuccess], + ); + + let response = stream_execution_response(Box::new(session), 1, None).await; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("read SSE body"); + let body = String::from_utf8(body.to_vec()).expect("SSE is UTF-8"); + + let future = body.find("response.future_metadata").expect("future event"); + let created = body.find("response.created").expect("created event"); + let completed = body.find("response.completed").expect("completed event"); + assert!(future < created && created < completed); + assert!(body.contains("resp_gateway")); + assert!(!body.contains("resp_upstream")); + assert!(body.ends_with("data: [DONE]\n\n")); + assert_eq!(trace.client_statuses(), vec![200]); + assert_eq!(trace.snapshot(), vec!["next_event", "commit", "next_end"]); + assert!(!trace.is_cancelled()); +} + #[tokio::test] async fn streaming_first_frame_encode_failure_cancels_before_commit() { let trace = Arc::new(Trace::default()); @@ -392,7 +618,7 @@ async fn streaming_first_frame_encode_failure_cancels_before_commit() { ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; assert_eq!( response.status(), @@ -414,7 +640,7 @@ async fn streaming_rate_limit_before_first_frame_should_persist_the_returned_429 ))], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); assert_eq!(trace.client_statuses(), vec![429]); @@ -433,7 +659,7 @@ async fn streaming_postcommit_encode_failure_cancels_without_second_commit() { ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read SSE body"); @@ -470,7 +696,7 @@ async fn streaming_success_should_emit_terminal_event_and_done_marker() { ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read complete SSE body"); @@ -505,7 +731,7 @@ async fn streaming_completed_event_without_finalized_execution_should_fail_close ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read failed SSE body"); @@ -540,7 +766,7 @@ async fn streaming_empty_terminal_should_emit_failure_done_and_cancel_execution( ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read failed SSE body"); @@ -571,7 +797,7 @@ async fn streaming_second_commit_request_should_fail_and_finalize_once() { ], ); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read failed SSE body"); @@ -598,7 +824,7 @@ async fn streaming_commit_failure_should_not_deliver_the_prepared_first_frame() ) .with_commit_failure(); - let response = stream_execution_response(session, 1).await; + let response = stream_execution_response(Box::new(session), 1, None).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read error response"); @@ -620,7 +846,7 @@ async fn buffered_response_commits_only_after_complete_json_is_encoded() { let trace = Arc::new(Trace::default()); let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]); - let response = collect_execution_response(session, 1).await; + let response = collect_execution_response(Box::new(session), 1).await; let body = to_bytes(response.into_body(), usize::MAX) .await .expect("read JSON body"); @@ -632,12 +858,34 @@ async fn buffered_response_commits_only_after_complete_json_is_encoded() { assert!(!trace.is_cancelled()); } +#[tokio::test] +async fn buffered_response_forwards_allowlisted_provider_headers() { + let trace = Arc::new(Trace::default()); + let header = ProviderResponseHeader::new( + "openai-model", + SafeUpstreamValue::new("gpt-5.5-effective").expect("model header"), + ) + .expect("safe header"); + let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]) + .with_response_headers(vec![header]); + + let response = collect_execution_response(Box::new(session), 1).await; + + assert_eq!( + response + .headers() + .get("openai-model") + .and_then(|value| value.to_str().ok()), + Some("gpt-5.5-effective") + ); +} + #[tokio::test] async fn buffered_encode_failure_cancels_without_commit() { let trace = Arc::new(Trace::default()); let session = FakeSession::buffered(Arc::clone(&trace), vec![completed()]); - let response = collect_execution_response(session, 1).await; + let response = collect_execution_response(Box::new(session), 1).await; assert_eq!( response.status(), @@ -654,7 +902,7 @@ async fn buffered_commit_failure_should_cancel_after_encoding_without_returning_ let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()]) .with_commit_failure(); - let response = collect_execution_response(session, 1).await; + let response = collect_execution_response(Box::new(session), 1).await; let status = response.status(); assert_eq!(status, axum::http::StatusCode::INTERNAL_SERVER_ERROR); @@ -670,7 +918,7 @@ async fn buffered_commit_failure_should_cancel_after_encoding_without_returning_ async fn dropping_buffered_handler_before_commit_cancels_execution() { let trace = Arc::new(Trace::default()); let session = FakeSession::pending_buffered(Arc::clone(&trace)); - let task = tokio::spawn(async move { collect_execution_response(session, 1).await }); + let task = tokio::spawn(async move { collect_execution_response(Box::new(session), 1).await }); tokio::task::yield_now().await; task.abort(); @@ -691,7 +939,7 @@ async fn buffered_rate_limit_should_persist_the_returned_429_status() { ), ); - let response = collect_execution_response(session, 1).await; + let response = collect_execution_response(Box::new(session), 1).await; assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); assert_eq!(trace.client_statuses(), vec![429]); diff --git a/backend/crates/gateway-api/tests/openai/responses/mod.rs b/backend/crates/gateway-api/tests/openai/responses/mod.rs index 0602091f4..3dcc283c8 100644 --- a/backend/crates/gateway-api/tests/openai/responses/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/mod.rs @@ -2,13 +2,15 @@ mod http; mod websocket; use gateway_api::openai::responses::{ - ContinuationIntent, RequestDecodeError, ResponseEncodeError, ResponsesCollector, decode_request, + ContinuationIntent, OpenAiResponsesEncoder, RequestDecodeError, ResponseEncodeError, + ResponsesCollector, decode_request, decode_response_create, }; use gateway_core::{ accounting::{CalculatedCost, ProviderReportedCost, Usage}, + error::SafeUpstreamValue, event::{ - ContentItem, ContentKind, EventSequenceError, FinishReason, GatewayEvent, ReasoningDelta, - ResponseMeta, TextDelta, ToolCallDelta, + ContentItem, ContentKind, EventSequenceError, FinishReason, GatewayEvent, + ProtocolWireEvent, ProviderEvent, ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta, }, operation::{ ContentPart, Feature, MessageRole, Operation, OutputFormat, ReasoningEffort, @@ -31,6 +33,15 @@ fn generate_operation( } } +fn openai_wire_body( + decoded: &gateway_api::openai::responses::DecodedResponsesRequest, +) -> &serde_json::Map { + generate_operation(decoded) + .protocol_payload() + .expect("OpenAI request should retain wire payload") + .body() +} + fn response_meta() -> ResponseMeta { ResponseMeta::new("resp_gateway_contract", "smart-code") } @@ -149,6 +160,19 @@ fn decoder_should_preserve_roles_tools_reasoning_schema_and_provider_options() { ); } +#[test] +fn decoder_should_preserve_a_bounded_prompt_cache_key_without_debug_exposure() { + let decoded = generate_request(json!({ + "model": "smart-code", + "input": "cache this prefix", + "prompt_cache_key": "private-cache-route" + })); + let request = generate_operation(&decoded); + + assert_eq!(request.prompt_cache_key(), Some("private-cache-route")); + assert!(!format!("{decoded:?}").contains("private-cache-route")); +} + #[test] fn decoder_should_map_string_input_to_user_text() { let decoded = generate_request(json!({"model": "smart-code", "input": "hello"})); @@ -161,7 +185,7 @@ fn decoder_should_map_string_input_to_user_text() { } #[test] -fn decoder_should_default_to_buffered_stored_response() { +fn decoder_should_default_to_streaming_non_stored_response() { let decoded = generate_request(json!({"model": "smart-code", "input": "hello"})); assert_eq!( @@ -170,10 +194,98 @@ fn decoder_should_default_to_buffered_stored_response() { decoded.metadata().store(), generate_operation(&decoded).response_persistence(), ), - (false, true, ResponsePersistence::Store) + (true, false, ResponsePersistence::DoNotStore) ); } +#[test] +fn decoder_should_project_boolean_transport_override_without_wire_leak() { + for (use_websocket, expected) in [(true, "websocket"), (false, "http_sse")] { + let decoded = generate_request(json!({ + "model": "smart-code", + "input": "hello", + "use_websocket": use_websocket + })); + let request = generate_operation(&decoded); + + assert!(openai_wire_body(&decoded).get("use_websocket").is_none()); + assert_eq!( + request + .provider_options() + .get("openai") + .and_then(|options| options.get("transport")), + Some(&json!(expected)) + ); + } +} + +#[test] +fn decoder_should_reject_non_boolean_transport_override_without_disclosure() { + let secret = "private-transport-marker"; + let error = decode_request( + json!({"model": "smart-code", "input": "hello", "use_websocket": secret}) + .to_string() + .as_bytes(), + ) + .expect_err("transport override must be a boolean"); + + assert_eq!( + error, + RequestDecodeError::InvalidType { + field: "use_websocket".to_owned(), + expected: "a boolean" + } + ); + assert!(!format!("{error:?}\n{error}").contains(secret)); +} + +#[test] +fn decoder_should_prefer_explicit_provider_transport_over_body_override() { + let decoded = generate_request(json!({ + "model": "smart-code", + "input": "hello", + "use_websocket": true, + "provider_options": { + "version": "v1", + "providers": { + "openai": {"schema_version": 1, "transport": "http_sse"} + } + } + })); + + assert_eq!( + generate_operation(&decoded) + .provider_options() + .get("openai") + .and_then(|options| options.get("transport")), + Some(&json!("http_sse")) + ); + assert!(openai_wire_body(&decoded).get("use_websocket").is_none()); +} + +#[test] +fn response_create_should_share_transport_override_projection() { + let decoded = decode_response_create( + &json!({ + "type": "response.create", + "model": "smart-code", + "input": "hello", + "use_websocket": false + }) + .to_string(), + ) + .expect("WebSocket response.create should use the shared request decoder"); + + assert_eq!( + generate_operation(&decoded) + .provider_options() + .get("openai") + .and_then(|options| options.get("transport")), + Some(&json!("http_sse")) + ); + assert!(openai_wire_body(&decoded).get("use_websocket").is_none()); +} + #[test] fn connection_context_should_attach_without_entering_request_debug() { let user_agent = "Codex-CLI/private-client-context"; @@ -327,43 +439,41 @@ fn decoder_should_require_input() { } #[test] -fn decoder_should_reject_zero_max_output_tokens() { - let error = decode_request( +fn decoder_should_preserve_zero_max_output_tokens_for_upstream_validation() { + let decoded = decode_request( json!({"model": "smart-code", "input": "hello", "max_output_tokens": 0}) .to_string() .as_bytes(), ) - .expect_err("zero output limit must fail"); + .expect("opaque OpenAI field should pass through"); assert_eq!( - error, - RequestDecodeError::InvalidValue { - field: "max_output_tokens".to_owned() - } + ( + generate_operation(&decoded).max_output_tokens(), + openai_wire_body(&decoded).get("max_output_tokens"), + ), + (None, Some(&json!(0))) ); } #[test] -fn decoder_should_reject_fractional_max_output_tokens() { - let error = decode_request( +fn decoder_should_preserve_fractional_max_output_tokens_for_upstream_validation() { + let decoded = decode_request( json!({"model": "smart-code", "input": "hello", "max_output_tokens": 1.5}) .to_string() .as_bytes(), ) - .expect_err("fractional output limit must fail"); + .expect("opaque OpenAI field should pass through"); assert_eq!( - error, - RequestDecodeError::InvalidType { - field: "max_output_tokens".to_owned(), - expected: "a non-negative integer" - } + openai_wire_body(&decoded).get("max_output_tokens"), + Some(&json!(1.5)) ); } #[test] -fn decoder_should_reject_an_unscoped_previous_response_id() { - let error = decode_request( +fn decoder_should_accept_an_external_previous_response_id() { + let decoded = decode_request( json!({ "model": "smart-code", "input": "continue", @@ -372,14 +482,12 @@ fn decoder_should_reject_an_unscoped_previous_response_id() { .to_string() .as_bytes(), ) - .expect_err("non-gateway response ID must fail"); + .expect("external response ID should pass through"); - assert_eq!( - error, - RequestDecodeError::InvalidValue { - field: "previous_response_id".to_owned() - } - ); + assert!(matches!( + decoded.metadata().continuation(), + ContinuationIntent::PreviousResponseId(value) if value == "upstream-response-id" + )); } #[test] @@ -426,8 +534,8 @@ fn decoder_should_reject_an_oversized_gateway_response_id() { } #[test] -fn decoder_should_reject_inline_image_data() { - let error = decode_request( +fn decoder_should_preserve_inline_image_data_without_decoding_it() { + let decoded = decode_request( json!({ "model": "smart-vision", "input": [{ @@ -441,20 +549,27 @@ fn decoder_should_reject_inline_image_data() { .to_string() .as_bytes(), ) - .expect_err("inline image data must fail"); + .expect("Provider should own OpenAI image validation"); + assert!( + decoded + .operation() + .capability_requirements() + .features() + .contains(&Feature::Vision) + ); assert_eq!( - error, - RequestDecodeError::InvalidValue { - field: "input[0].content[0].image_url".to_owned() - } + openai_wire_body(&decoded) + .get("input") + .and_then(|input| input.pointer("/0/content/0/image_url")), + Some(&json!("data:image/png;base64,cHJpdmF0ZQ==")) ); } #[test] -fn decoder_should_report_nested_unknown_field_path_without_its_value() { +fn decoder_should_preserve_nested_unknown_fields_without_debug_exposure() { let secret = "nested-private-value"; - let error = decode_request( + let decoded = decode_request( json!({ "model": "smart-code", "input": [{ @@ -469,49 +584,45 @@ fn decoder_should_report_nested_unknown_field_path_without_its_value() { .to_string() .as_bytes(), ) - .expect_err("nested unknown field must fail"); - let rendered = format!("{error:?}\n{error}\n{}", error.protocol_body().into_value()); + .expect("unknown OpenAI fields should pass through"); + let rendered = format!("{decoded:?}"); assert_eq!( - error, - RequestDecodeError::UnknownField { - field: "input[0].content[0].secret_value".to_owned() - } + openai_wire_body(&decoded) + .get("input") + .and_then(|input| input.pointer("/0/content/0/secret_value")), + Some(&json!(secret)) ); assert!(!rendered.contains(secret)); } #[test] -fn decoder_should_reject_unknown_top_level_field() { - let error = decode_request( +fn decoder_should_preserve_unknown_top_level_field() { + let decoded = decode_request( json!({"model": "smart-code", "input": "hello", "mystery": true}) .to_string() .as_bytes(), ) - .expect_err("unknown field must fail"); + .expect("unknown OpenAI fields should pass through"); assert_eq!( - error, - RequestDecodeError::UnknownField { - field: "mystery".to_owned() - } + openai_wire_body(&decoded).get("mystery"), + Some(&json!(true)) ); } #[test] -fn decoder_should_reject_known_unsupported_semantic() { - let error = decode_request( +fn decoder_should_preserve_new_official_semantic_without_gateway_enumeration() { + let decoded = decode_request( json!({"model": "smart-code", "input": "hello", "background": true}) .to_string() .as_bytes(), ) - .expect_err("known unsupported field must fail"); + .expect("Provider should own OpenAI request evolution"); assert_eq!( - error, - RequestDecodeError::UnsupportedField { - field: "background".to_owned() - } + openai_wire_body(&decoded).get("background"), + Some(&json!(true)) ); } @@ -535,11 +646,11 @@ fn decoder_should_reject_unknown_provider_options_version() { fn request_errors_should_not_disclose_prompt_or_raw_body() { let secret = "private prompt that must not appear"; let error = decode_request( - json!({"model": "smart-code", "input": secret, "background": true}) + json!({"model": "smart-code", "input": secret, "stream": "private-invalid-value"}) .to_string() .as_bytes(), ) - .expect_err("unsupported field must fail"); + .expect_err("wire routing field type must fail"); let rendered = format!("{error:?}\n{error}\n{}", error.protocol_body().into_value()); assert!(!rendered.contains(secret) && !rendered.contains("private prompt")); @@ -838,3 +949,75 @@ fn collector_should_reject_unrepresentable_media_output() { ResponseEncodeError::UnsupportedContentKind { kind: "image" } ); } + +#[test] +fn openai_wire_encoder_should_preserve_unknown_media_events_and_hide_upstream_response_id() { + let upstream_id = "resp_upstream_private"; + let gateway_id = "resp_gateway_public"; + let started = ResponseMeta::new(gateway_id, "gpt-test").with_upstream_response_id( + SafeUpstreamValue::new(upstream_id).expect("safe upstream response ID"), + ); + let completed = ResponseMeta::new(gateway_id, "gpt-test") + .with_finish_reason(FinishReason::Stop) + .with_upstream_response_id( + SafeUpstreamValue::new(upstream_id).expect("safe upstream response ID"), + ); + let events = [ + openai_wire_event( + vec![GatewayEvent::Started(started)], + "response.created", + json!({ + "type": "response.created", + "response": {"id": upstream_id, "status": "in_progress"} + }), + ), + openai_wire_event( + Vec::new(), + "response.image_generation_call.partial_image", + json!({ + "type": "response.image_generation_call.partial_image", + "response_id": upstream_id, + "partial_image_b64": "opaque-image-fragment" + }), + ), + openai_wire_event( + vec![GatewayEvent::Completed(completed)], + "response.completed", + json!({ + "type": "response.completed", + "response": { + "id": upstream_id, + "status": "completed", + "output": [{"type": "image_generation_call", "result": "opaque-image"}] + } + }), + ), + ]; + let mut encoder = OpenAiResponsesEncoder::new(1_700_000_000); + let frames = events + .iter() + .flat_map(|event| encoder.push_sse(event).expect("wire event should encode")) + .collect::>(); + let response = encoder.finish().expect("wire response should finish"); + let parsed = parse_sse_events(&frames.join("")).expect("wire SSE should parse"); + + assert_eq!(parsed.len(), 3); + assert_eq!( + serde_json::from_str::(&parsed[1].data) + .expect("unknown event JSON") + .get("partial_image_b64"), + Some(&json!("opaque-image-fragment")) + ); + assert_eq!(response.get("id"), Some(&json!(gateway_id))); + assert!(!frames.join("").contains(upstream_id)); +} + +fn openai_wire_event(canonical: Vec, event_type: &str, data: Value) -> ProviderEvent { + let wire = ProtocolWireEvent::json("openai", Some(event_type.to_owned()), data) + .expect("valid OpenAI wire event"); + if canonical.is_empty() { + ProviderEvent::wire(wire) + } else { + ProviderEvent::canonical_with_wire(canonical, wire) + } +} diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket.rs b/backend/crates/gateway-api/tests/openai/responses/websocket.rs index f17a4350c..9592b62c3 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket.rs @@ -93,11 +93,12 @@ fn response_create_errors_should_not_retain_prompt_or_unknown_field_values() { "type": "response.create", "model": "smart-code", "input": prompt, - "metadata": {"secret": secret} + "metadata": {"secret": secret}, + "stream": secret }) .to_string(), ) - .expect_err("unsupported field"); + .expect_err("invalid stream field"); let rendered = format!("{error:?}\n{error}"); assert!(!rendered.contains(prompt)); diff --git a/backend/crates/gateway-core/src/accounting/mod.rs b/backend/crates/gateway-core/src/accounting/mod.rs index 87f6ca6b9..41c2ea72c 100644 --- a/backend/crates/gateway-core/src/accounting/mod.rs +++ b/backend/crates/gateway-core/src/accounting/mod.rs @@ -277,34 +277,78 @@ pub struct CalculatedCostBreakdown { multiplier_percent: u32, } -impl CalculatedCostBreakdown { - #[allow(clippy::too_many_arguments)] +/// 一次请求的费用组成，全部使用同一币种。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CalculatedCostAmounts { + input: Money, + output: Money, + cache_read: Money, + cache_write: Money, + standard: Money, + total: Money, +} + +impl CalculatedCostAmounts { #[must_use] pub const fn new( - input_amount: Money, - output_amount: Money, - cache_read_amount: Money, - cache_write_amount: Money, - standard_amount: Money, - total_amount: Money, - input_price_per_million: Money, - output_price_per_million: Money, - cache_read_price_per_million: Money, - cache_write_price_per_million: Money, + input: Money, + output: Money, + cache_read: Money, + cache_write: Money, + standard: Money, + total: Money, + ) -> Self { + Self { + input, + output, + cache_read, + cache_write, + standard, + total, + } + } +} + +/// 每百万 Token 的费率组成，全部使用同一币种。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CalculatedCostRates { + input: Money, + output: Money, + cache_read: Money, + cache_write: Money, +} + +impl CalculatedCostRates { + #[must_use] + pub const fn new(input: Money, output: Money, cache_read: Money, cache_write: Money) -> Self { + Self { + input, + output, + cache_read, + cache_write, + } + } +} + +impl CalculatedCostBreakdown { + #[must_use] + pub const fn new( + amounts: CalculatedCostAmounts, + rates: CalculatedCostRates, service_tier: Option, multiplier_percent: u32, ) -> Self { Self { - input_amount, - output_amount, - cache_read_amount, - cache_write_amount, - standard_amount, - total_amount, - input_price_per_million, - output_price_per_million, - cache_read_price_per_million, - cache_write_price_per_million, + input_amount: amounts.input, + output_amount: amounts.output, + cache_read_amount: amounts.cache_read, + cache_write_amount: amounts.cache_write, + standard_amount: amounts.standard, + total_amount: amounts.total, + input_price_per_million: rates.input, + output_price_per_million: rates.output, + cache_read_price_per_million: rates.cache_read, + cache_write_price_per_million: rates.cache_write, service_tier, multiplier_percent, } diff --git a/backend/crates/gateway-core/src/engine/admission.rs b/backend/crates/gateway-core/src/engine/admission.rs new file mode 100644 index 000000000..715469843 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/admission.rs @@ -0,0 +1,126 @@ +//! Client Key 准入、释放与进程启动恢复的中立契约。 + +use std::time::{Duration, SystemTime}; + +use futures::future::BoxFuture; + +use crate::policy::{ClientApiKeyId, RateLimits}; + +use super::{ExecutionStore, ModelRequestId}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientAdmissionRequest { + pub model_request_id: ModelRequestId, + pub client_api_key_id: ClientApiKeyId, + pub input_token_estimate: u64, + pub lease_ttl: Duration, + pub limits: RateLimits, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientAdmissionRejection { + RateLimited, + ConcurrencyLimited, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientAdmissionDecision { + Granted, + Rejected(ClientAdmissionRejection), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RecentAdmissionFact { + pub model_request_id: ModelRequestId, + pub started_at: SystemTime, + pub input_token_estimate: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RunningAdmissionFact { + pub model_request_id: ModelRequestId, + pub expires_at: SystemTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientAdmissionRecovery { + pub client_api_key_id: ClientApiKeyId, + pub recent_requests: Vec, + pub running_requests: Vec, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ClientAdmissionRestoreResult { + pub restored_recent_requests: u64, + pub restored_running_requests: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("client admission store is unavailable")] +pub struct ClientAdmissionError; + +pub trait ClientAdmissionPort: Send + Sync { + fn admit( + &self, + request: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result>; + + fn release<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + model_request_id: &'a ModelRequestId, + ) -> BoxFuture<'a, Result>; + + fn restore( + &self, + recovery: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result>; +} + +pub trait ClientAdmissionRecoveryPort: Send + Sync { + fn load_recovery( + &self, + since: SystemTime, + ) -> BoxFuture<'_, Result, ClientAdmissionError>>; +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ClientAdmissionStartupRecoveryReport { + pub expired_model_requests: u64, + pub restored_clients: u64, + pub restored_recent_requests: u64, + pub restored_running_requests: u64, +} + +/// 监听端口前完成 PostgreSQL 终态收敛与 Redis 热状态恢复。 +pub async fn restore_client_admission_startup( + execution: &dyn ExecutionStore, + recovery: &dyn ClientAdmissionRecoveryPort, + admissions: &dyn ClientAdmissionPort, + now: SystemTime, +) -> Result { + let expired = execution + .recover_expired(now) + .await + .map_err(|_| ClientAdmissionError)?; + let since = now + .checked_sub(Duration::from_secs(61)) + .ok_or(ClientAdmissionError)?; + let recoveries = recovery.load_recovery(since).await?; + let restored_clients = u64::try_from(recoveries.len()).unwrap_or(u64::MAX); + let mut restored_recent_requests = 0_u64; + let mut restored_running_requests = 0_u64; + for recovery in recoveries { + let restored = admissions.restore(recovery).await?; + restored_recent_requests = + restored_recent_requests.saturating_add(restored.restored_recent_requests); + restored_running_requests = + restored_running_requests.saturating_add(restored.restored_running_requests); + } + Ok(ClientAdmissionStartupRecoveryReport { + expired_model_requests: expired.requests, + restored_clients, + restored_recent_requests, + restored_running_requests, + }) +} diff --git a/backend/crates/gateway-core/src/engine/continuation.rs b/backend/crates/gateway-core/src/engine/continuation.rs index a0433d138..cd103a4f8 100644 --- a/backend/crates/gateway-core/src/engine/continuation.rs +++ b/backend/crates/gateway-core/src/engine/continuation.rs @@ -1,14 +1,16 @@ -//! 客户端历史与可选 Provider 原生 previous-response 约束。 +//! Provider 原生 previous-response 的调用方隔离、账号绑定与复用约束。 //! -//! Core 不保存 transcript，也不提供 portable history repository。普通连续对话已经 -//! 作为 [`Operation`](crate::operation::Operation) 的完整消息输入进入本次请求。 +//! Core 不解释 Provider transcript；同一客户端连接需要的可携带状态由 +//! [`ProviderSessionState`](crate::operation::ProviderSessionState) 不透明承载。 use std::fmt; +use futures::future::BoxFuture; use thiserror::Error; use crate::engine::credential::ProviderAccountId; use crate::error::{IdentifierError, SafeUpstreamValue, validate_text}; +use crate::policy::ClientApiKeyId; use crate::routing::{ProviderInstanceId, ProviderKind}; /// 客户端传入的 previous response ID。 @@ -40,6 +42,15 @@ pub enum NativeContinuationReuse { SingleUse, } +/// Provider 原生 response handle 的续接范围。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NativeContinuationScope { + /// 上游已持久化，可由同账号的新连接继续。 + Persisted, + /// 只存在于完成上一轮的 WebSocket。 + ConnectionLocal, +} + /// 外层已解析并认证的 native previous-response pin。 /// /// 该值不代表数据库 transcript；它只阻止 native handle 被发送到错误的 @@ -54,6 +65,7 @@ pub struct NativeContinuationPin { instance: ProviderInstanceId, account: ProviderAccountId, reuse: NativeContinuationReuse, + scope: NativeContinuationScope, } impl NativeContinuationPin { @@ -73,9 +85,17 @@ impl NativeContinuationPin { instance, account, reuse, + scope: NativeContinuationScope::ConnectionLocal, } } + /// 设置 Store 已确认的原生续接范围。 + #[must_use] + pub const fn with_scope(mut self, scope: NativeContinuationScope) -> Self { + self.scope = scope; + self + } + #[must_use] pub const fn previous_response_id(&self) -> &PreviousResponseId { &self.previous_response_id @@ -107,6 +127,11 @@ impl NativeContinuationPin { self.reuse } + #[must_use] + pub const fn scope(&self) -> NativeContinuationScope { + self.scope + } + /// 校验本次 route/account 选择没有破坏 native pin。 #[must_use] pub fn matches( @@ -129,10 +154,51 @@ impl fmt::Debug for NativeContinuationPin { .field("instance", &self.instance) .field("account", &self.account) .field("reuse", &self.reuse) + .field("scope", &self.scope) .finish() } } +/// 一次请求最终采用的 previous-response 绑定方式。 +/// +/// 已命中网关历史的 handle 携带完整账号 pin；未命中历史的外部 handle 只保留 +/// 客户端提交的 opaque ID，由目标 Provider 在首次且唯一一次 attempt 中解释。 +#[derive(Clone, PartialEq, Eq)] +pub enum ContinuationBinding { + Pinned(NativeContinuationPin), + External(PreviousResponseId), +} + +impl ContinuationBinding { + #[must_use] + pub const fn previous_response_id(&self) -> &PreviousResponseId { + match self { + Self::Pinned(pin) => pin.previous_response_id(), + Self::External(previous_response_id) => previous_response_id, + } + } + + #[must_use] + pub const fn pinned(&self) -> Option<&NativeContinuationPin> { + match self { + Self::Pinned(pin) => Some(pin), + Self::External(_) => None, + } + } +} + +impl fmt::Debug for ContinuationBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Pinned(pin) => formatter.debug_tuple("Pinned").field(pin).finish(), + Self::External(_) => formatter + .debug_tuple("External") + .field(&"") + .finish(), + } + } +} + /// Request-local 的 single-use claim 状态。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum NativeClaimState { @@ -208,3 +274,16 @@ pub enum ContinuationError { #[error("native continuation handle is already claimed or terminal")] AlreadyClaimed, } + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("native continuation store is unavailable")] +pub struct NativeContinuationStoreError; + +/// 调用方隔离后的 previous-response 解析端口。 +pub trait NativeContinuationPort: Send + Sync { + fn resolve<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + previous_response_id: &'a PreviousResponseId, + ) -> BoxFuture<'a, Result, NativeContinuationStoreError>>; +} diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index d96f48740..e602a363c 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -1,6 +1,6 @@ //! 唯一的账号重试、Provider instance 切换、发送与下游 commit barrier owner。 -use std::collections::BTreeSet; +use std::collections::{BTreeSet, VecDeque}; use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, SystemTime}; @@ -10,25 +10,28 @@ use futures_timer::Delay; use uuid::Uuid; use crate::accounting::{CostEstimate, CostSource, Usage}; -use crate::engine::continuation::NativeContinuationPin; +use crate::engine::continuation::ContinuationBinding; use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream}; use crate::engine::{ - AccountSelectionConstraints, AttemptContext, AttemptRecord, AttemptTrigger, CancellationToken, - CommitRequirement, CoordinatedEvent, EngineError, ExecutionOutcome, ExecutionStore, - GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, - ModelRequestTimings, NewModelRequest, ProviderAttemptOutcome, UpstreamSendState, + AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTrigger, CancellationToken, + CommitRequirement, ContinuationAttempt, CoordinatedEvent, EngineError, ExecutionOutcome, + ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, + ModelRequestTimings, NewModelRequest, ProviderAccountStateOwner, ProviderAttemptOutcome, + UpstreamSendState, }; use crate::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind}; -use crate::event::GatewayEvent; +use crate::event::{ + GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, +}; use crate::operation::{Operation, RetrySafety}; use crate::routing::RoutingPlan; /// Request 级协调器；不会创建或写入 `request_attempts`。 -pub struct AttemptCoordinator { +pub struct AttemptCoordinator { engine: Arc>, } -impl AttemptCoordinator +impl AttemptCoordinator where S: ExecutionStore, { @@ -50,7 +53,7 @@ where operation: Operation, plan: RoutingPlan, required_account: Option, - continuation: Option, + continuation: Option, cancellation: CancellationToken, ) -> Result, EngineError> { let request_id = request.id.clone(); @@ -58,6 +61,17 @@ where let deadline = request.deadline_at; self.engine.store().create_model_request(request).await?; let gateway_response_id = format!("resp_{}", Uuid::now_v7().simple()); + let account_state_owner = continuation + .as_ref() + .and_then(ContinuationBinding::pinned) + .map(ProviderAccountStateOwner::from_continuation); + let continuation_attempt = + initial_continuation_attempt(&operation, &plan, continuation.as_ref()); + let commit_policy = if continuation_attempt == ContinuationAttempt::None { + StreamCommitPolicy::FirstCanonicalEvent + } else { + StreamCommitPolicy::UntilOutputOrTerminal + }; let mut session = ResponseExecutionSession { engine: Arc::clone(&self.engine), @@ -68,6 +82,9 @@ where plan, required_account, continuation, + continuation_attempt, + commit_policy, + account_state_owner, cancellation, candidate_index: 0, last_attempt_candidate_index: None, @@ -86,6 +103,7 @@ where client_response_id: None, upstream_response_id: None, provider_attempt_outcomes: Vec::new(), + before_commit: VecDeque::new(), }; if session.cancellation.is_cancelled() { @@ -107,6 +125,7 @@ struct CurrentAttempt { index: NonZeroU32, started_at: SystemTime, send_observed: bool, + response_observation: Option, } struct FailureFinalization { @@ -119,11 +138,17 @@ struct FailureFinalization { provider_response_id: Option, } +#[derive(Clone, Copy)] +enum StreamCommitPolicy { + FirstCanonicalEvent, + UntilOutputOrTerminal, +} + /// API 可逐事件消费的 Core 执行会话。 /// /// API 只能提交下游 delivery 边界；账号重试、Provider instance 切换、断流终结与 /// `model_requests` 写回均留在本类型内。 -pub struct ResponseExecutionSession { +pub struct ResponseExecutionSession { engine: Arc>, request_id: ModelRequestId, request_started_at: SystemTime, @@ -131,7 +156,10 @@ pub struct ResponseExecutionSession { operation: Operation, plan: RoutingPlan, required_account: Option, - continuation: Option, + continuation: Option, + continuation_attempt: ContinuationAttempt, + commit_policy: StreamCommitPolicy, + account_state_owner: Option, cancellation: CancellationToken, candidate_index: usize, last_attempt_candidate_index: Option, @@ -150,9 +178,10 @@ pub struct ResponseExecutionSession { client_response_id: Option, upstream_response_id: Option, provider_attempt_outcomes: Vec, + before_commit: VecDeque, } -impl ResponseExecutionSession +impl ResponseExecutionSession where S: ExecutionStore, { @@ -168,19 +197,43 @@ where if self.finalized { return Ok(None); } + if self.downstream_committed_at.is_some() + && let Some(event) = self.before_commit.pop_front() + { + return Ok(Some(CoordinatedEvent::single( + event, + CommitRequirement::AlreadyCommitted, + ))); + } loop { match self.pull().await? { PullOutcome::Event(event) => { - let requirement = if self.downstream_committed_at.is_some() { - CommitRequirement::AlreadyCommitted - } else { - self.delivery_pending = true; - CommitRequirement::CommitBeforeDelivery + if self.downstream_committed_at.is_some() { + return Ok(Some(CoordinatedEvent::single( + event, + CommitRequirement::AlreadyCommitted, + ))); + } + let commit_significant = match self.commit_policy { + StreamCommitPolicy::FirstCanonicalEvent => event.has_canonical_facts(), + StreamCommitPolicy::UntilOutputOrTerminal => event.is_commit_significant(), }; - return Ok(Some(CoordinatedEvent::new(event, requirement))); + self.before_commit.push_back(event); + if !commit_significant { + continue; + } + self.delivery_pending = true; + let events = self.before_commit.drain(..).collect(); + return CoordinatedEvent::try_batch( + events, + CommitRequirement::CommitBeforeDelivery, + ) + .map(Some); + } + PullOutcome::AttemptDiscarded => { + self.before_commit.clear(); } - PullOutcome::AttemptDiscarded => {} PullOutcome::End => { if self.downstream_committed_at.is_some() { self.finish_success().await?; @@ -196,7 +249,7 @@ where /// # Errors /// /// 会话已提交、已有待提交结果或执行失败时返回错误。 - pub async fn collect_uncommitted(&mut self) -> Result, EngineError> { + pub async fn collect_uncommitted(&mut self) -> Result, EngineError> { if self.downstream_committed_at.is_some() || self.delivery_pending { return Err(EngineError::InvalidDeliveryState); } @@ -279,6 +332,16 @@ where &self.provider_attempt_outcomes } + /// 返回最终选中 attempt 已公开给协议层的安全响应头。 + #[must_use] + pub fn response_headers(&self) -> &[ProviderResponseHeader] { + self.current + .as_ref() + .and_then(|current| current.response_observation.as_ref()) + .map(ProviderResponseObservation::client_headers) + .unwrap_or_default() + } + /// 请求取消；实际终态在下一次会话 poll 时由 Core 持久化。 pub fn cancel(&self) { self.cancellation.cancel(); @@ -325,17 +388,37 @@ where return Err(EngineError::Deadline); } PollBoundary::Item(Some(Ok(mut event))) => { - match self.freeze_response_identity(&mut event) { - Ok(()) => { - self.observe_event(&event).await?; - return Ok(PullOutcome::Event(event)); + if let Some(observation) = event.take_observation() + && let Err(error) = self.observe_response(observation) + { + if self.handle_stream_error(error).await? { + return Ok(PullOutcome::AttemptDiscarded); } - Err(error) => { - if self.handle_stream_error(error).await? { - return Ok(PullOutcome::AttemptDiscarded); - } + continue; + } + let mut identity_error = None; + for fact in event.canonical_facts_mut() { + if let Err(error) = self.freeze_response_identity(fact) { + identity_error = Some(error); + break; } } + if let Some(error) = identity_error { + if self.handle_stream_error(error).await? { + return Ok(PullOutcome::AttemptDiscarded); + } + continue; + } + for fact in event.canonical_facts() { + self.observe_event(fact).await?; + } + if event.wire_event().is_some() && !event.has_canonical_facts() { + self.observe_wire_event().await?; + } + if !event.has_client_event() { + continue; + } + return Ok(PullOutcome::Event(event)); } PollBoundary::Item(Some(Err(error))) => { if self.handle_stream_error(error).await? { @@ -383,13 +466,15 @@ where next_attempt, self.deadline, self.plan.account_selection_policy(), - AccountSelectionConstraints::new( + AccountAttemptContext::new( self.excluded_accounts.clone(), self.required_account.clone(), + self.account_state_owner.clone(), ), self.continuation.clone(), self.cancellation.clone(), - ); + ) + .with_continuation_attempt(self.continuation_attempt); let trigger = if self.attempts == 0 { AttemptTrigger::Initial } else if self.last_attempt_candidate_index == Some(self.candidate_index) { @@ -493,7 +578,11 @@ where .await?; return Err(EngineError::RequiredAccountMismatch); } - if let Some(pin) = self.continuation.as_ref() + if let Some(pin) = self + .continuation + .as_ref() + .and_then(ContinuationBinding::pinned) + && self.continuation_attempt == ContinuationAttempt::Native && !metadata.provider_account_id().is_some_and(|account| { pin.matches(metadata.provider(), metadata.instance(), account) }) @@ -514,6 +603,15 @@ where .await?; return Err(EngineError::ContinuationPinMismatch); } + if self.account_state_owner.is_none() + && let Some(account) = metadata.provider_account_id() + { + self.account_state_owner = Some(ProviderAccountStateOwner::new( + metadata.provider().clone(), + metadata.instance().clone(), + account.clone(), + )); + } self.engine .store() .record_attempt(AttemptRecord { @@ -538,6 +636,7 @@ where index: next_attempt, started_at: SystemTime::now(), send_observed: false, + response_observation: None, }); return Ok(()); } @@ -569,6 +668,51 @@ where Ok(()) } + async fn observe_wire_event(&mut self) -> Result<(), EngineError> { + let current = self.current.as_mut().ok_or(EngineError::EmptyRoutingPlan)?; + if !current.send_observed { + self.engine + .store() + .mark_send_state(&self.request_id, UpstreamSendState::Sent) + .await?; + current.send_observed = true; + } + Ok(()) + } + + fn observe_response( + &mut self, + observation: ProviderResponseObservation, + ) -> Result<(), ProviderError> { + let send_state = self.current_send_state(); + let current = self + .current + .as_mut() + .ok_or_else(|| ProviderError::new(ProviderErrorKind::Protocol, send_state))?; + if current + .response_observation + .as_ref() + .is_some_and(|existing| existing.transport() != observation.transport()) + { + return Err(ProviderError::new(ProviderErrorKind::Protocol, send_state)); + } + let observed = observation.timings(); + if let Some(value) = observed.transport_decision_wait_ms { + self.timings.transport_decision_wait_ms = Some(value); + } + if let Some(value) = observed.connect_ms { + self.timings.connect_ms = Some(value); + } + if let Some(value) = observed.headers_ms { + self.timings.headers_ms = Some(value); + } + if let Some(value) = observed.first_event_ms { + self.timings.first_event_ms = Some(value); + } + current.response_observation = Some(observation); + Ok(()) + } + fn freeze_response_identity(&mut self, event: &mut GatewayEvent) -> Result<(), ProviderError> { let upstream_response_id = event .freeze_gateway_response_id(&self.gateway_response_id) @@ -606,19 +750,25 @@ where .store() .mark_send_state(&self.request_id, send_state) .await?; - let provider_proved_replay_safe = - !current.send_observed && provider_proved_replay_safe(&error); - let retryable = self.required_account.is_none() - && self.continuation.is_none() + let provider_proved_replay_safe = provider_proved_replay_safe(&error); + let continuation_retry = self.prepare_continuation_retry( + ¤t, + &error, + send_state, + provider_proved_replay_safe, + ); + let ordinary_retry = self.required_account.is_none() + && self.continuation_attempt == ContinuationAttempt::None && self.downstream_committed_at.is_none() && !self.delivery_pending && send_state != UpstreamSendState::Ambiguous && (self.operation.retry_safety() == RetrySafety::Idempotent || provider_proved_replay_safe) && self.attempts < self.plan.max_attempts().get(); + let retryable = continuation_retry || ordinary_retry; if retryable { - if let Some(account) = current.metadata.provider_account_id() { + if !continuation_retry && let Some(account) = current.metadata.provider_account_id() { self.excluded_accounts.insert(account.clone()); } if error.kind() == ProviderErrorKind::Unsupported @@ -636,6 +786,15 @@ where provider_kind: current.metadata.provider().clone(), account_id: current.metadata.provider_account_id().cloned(), upstream_model_id: current.metadata.upstream_model().clone(), + upstream_status_code: current + .response_observation + .as_ref() + .and_then(ProviderResponseObservation::status_code), + upstream_request_id: current + .response_observation + .as_ref() + .and_then(ProviderResponseObservation::request_id) + .map(|value| value.as_str().to_owned()), latency: current.started_at.elapsed().unwrap_or_default(), error, }) @@ -651,24 +810,85 @@ where Err(engine_error) } + fn prepare_continuation_retry( + &mut self, + current: &CurrentAttempt, + error: &ProviderError, + send_state: UpstreamSendState, + provider_proved_replay_safe: bool, + ) -> bool { + if self.required_account.is_some() + || self.continuation_attempt == ContinuationAttempt::None + || self.downstream_committed_at.is_some() + || self.delivery_pending + || send_state == UpstreamSendState::Ambiguous + || !provider_proved_replay_safe + || self.attempts >= self.plan.max_attempts().get() + || self + .operation + .provider_session_state(current.metadata.provider().as_str()) + .is_none() + { + return false; + } + + match self.continuation_attempt { + ContinuationAttempt::Native if error.continuation_failure().is_some() => { + self.continuation_attempt = ContinuationAttempt::ReplayOwner; + } + ContinuationAttempt::Native | ContinuationAttempt::ReplayOwner => { + self.continuation_attempt = ContinuationAttempt::ReplayAny; + if let Some(account) = current.metadata.provider_account_id() { + self.excluded_accounts.insert(account.clone()); + } + } + ContinuationAttempt::ReplayAny => { + if let Some(account) = current.metadata.provider_account_id() { + self.excluded_accounts.insert(account.clone()); + } + } + ContinuationAttempt::None => return false, + } + true + } + fn reset_uncommitted_observations(&mut self) { self.usage = Usage::new(); self.cost = CostEstimate::unavailable(); self.client_response_id = None; self.upstream_response_id = None; + self.timings.transport_decision_wait_ms = None; + self.timings.connect_ms = None; + self.timings.headers_ms = None; self.timings.first_event_ms = None; self.timings.first_reasoning_ms = None; self.timings.first_text_ms = None; self.timings.first_token_ms = None; self.upstream_complete = false; + self.before_commit.clear(); } fn can_fallback_instance(&self, provider_proved_replay_safe: bool) -> bool { - self.required_account.is_none() - && self.continuation.is_none() - && (self.operation.retry_safety() == RetrySafety::Idempotent - || provider_proved_replay_safe) - && self.candidate_index + 1 < self.plan.candidates().len() + if self.required_account.is_some() + || self.candidate_index + 1 >= self.plan.candidates().len() + { + return false; + } + let next = &self.plan.candidates()[self.candidate_index + 1]; + match self.continuation_attempt { + ContinuationAttempt::None => { + self.operation.retry_safety() == RetrySafety::Idempotent + || provider_proved_replay_safe + } + ContinuationAttempt::ReplayAny => { + provider_proved_replay_safe + && self + .operation + .provider_session_state(next.provider().as_str()) + .is_some() + } + ContinuationAttempt::Native | ContinuationAttempt::ReplayOwner => false, + } } async fn finish_success(&mut self) -> Result<(), EngineError> { @@ -679,10 +899,27 @@ where self.timings.latency_ms = Some(elapsed_ms(self.request_started_at, completed_at)); let upstream_request_id = self.current.as_ref().and_then(|current| { current - .metadata - .upstream_request_id() - .map(|value| value.as_str().to_owned()) + .response_observation + .as_ref() + .and_then(|observation| { + observation + .request_id() + .map(|value| value.as_str().to_owned()) + }) + .or_else(|| { + current + .metadata + .upstream_request_id() + .map(|value| value.as_str().to_owned()) + }) }); + let upstream_status_code = self + .current + .as_ref() + .and_then(|current| current.response_observation.as_ref()) + .and_then(ProviderResponseObservation::status_code) + .or(Some(200)); + let (upstream_transport, http_version) = self.current_transport_observation(); self.engine .store() .finalize_model_request(ModelRequestFinalization { @@ -692,10 +929,12 @@ where attempt_count: self.attempts, downstream_committed_at: self.downstream_committed_at, client_status_code: self.client_status_code, - upstream_status_code: Some(200), + upstream_status_code, client_response_id: self.client_response_id.clone(), upstream_request_id, upstream_response_id: self.upstream_response_id.clone(), + upstream_transport, + http_version, error: None, provider_error_code: None, retry_after_ms: None, @@ -800,10 +1039,26 @@ where self.timings.latency_ms = Some(elapsed_ms(self.request_started_at, completed_at)); let upstream_request_id = self.current.as_ref().and_then(|current| { current - .metadata - .upstream_request_id() - .map(|value| value.as_str().to_owned()) + .response_observation + .as_ref() + .and_then(|observation| { + observation + .request_id() + .map(|value| value.as_str().to_owned()) + }) + .or_else(|| { + current + .metadata + .upstream_request_id() + .map(|value| value.as_str().to_owned()) + }) }); + let observed_status_code = self + .current + .as_ref() + .and_then(|current| current.response_observation.as_ref()) + .and_then(ProviderResponseObservation::status_code); + let (upstream_transport, http_version) = self.current_transport_observation(); self.engine .store() .finalize_model_request(ModelRequestFinalization { @@ -813,12 +1068,14 @@ where attempt_count: self.attempts, downstream_committed_at: self.downstream_committed_at, client_status_code: self.client_status_code, - upstream_status_code: finalization.upstream_status_code, + upstream_status_code: finalization.upstream_status_code.or(observed_status_code), client_response_id: self.client_response_id.clone(), upstream_request_id, upstream_response_id: finalization .provider_response_id .or_else(|| self.upstream_response_id.clone()), + upstream_transport, + http_version, error: Some(finalization.error), provider_error_code: finalization.provider_error_code, retry_after_ms: finalization.retry_after_ms, @@ -832,6 +1089,21 @@ where Ok(()) } + fn current_transport_observation(&self) -> (Option, Option) { + let Some(current) = self.current.as_ref() else { + return (None, None); + }; + let Some(observation) = current.response_observation.as_ref() else { + return (None, None); + }; + ( + Some(observation.transport().as_str().to_owned()), + observation + .http_version() + .map(|version| version.as_str().to_owned()), + ) + } + fn current_send_state(&self) -> UpstreamSendState { if self.attempts == 0 { UpstreamSendState::NotSent @@ -883,13 +1155,13 @@ where } enum PullOutcome { - Event(GatewayEvent), + Event(ProviderEvent), AttemptDiscarded, End, } enum PollBoundary { - Item(Option>), + Item(Option>), Cancelled, Deadline, } @@ -940,6 +1212,26 @@ async fn poll_provider( } } +fn initial_continuation_attempt( + operation: &Operation, + plan: &RoutingPlan, + continuation: Option<&ContinuationBinding>, +) -> ContinuationAttempt { + match continuation { + None => ContinuationAttempt::None, + Some(ContinuationBinding::External(_)) + if plan.candidates().first().is_some_and(|candidate| { + operation + .provider_session_state(candidate.provider().as_str()) + .is_some() + }) => + { + ContinuationAttempt::ReplayAny + } + Some(_) => ContinuationAttempt::Native, + } +} + fn provider_engine_error(error: &ProviderError) -> EngineError { if error.kind() == ProviderErrorKind::Cancelled { EngineError::Cancelled diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index a82782c41..22490efb8 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -354,10 +354,16 @@ impl ProviderAccount { /// 判断账号当前能否进入同 target 的候选池。 #[must_use] pub fn is_schedulable(&self, now: SystemTime) -> bool { - self.enabled - && self.availability == AccountAvailability::Ready - && self.cooldown_until.is_none_or(|until| until <= now) - && self.access_token_expires_at > now + let available = match self.availability { + AccountAvailability::Ready => true, + AccountAvailability::Cooldown => self.cooldown_until.is_some_and(|until| until <= now), + AccountAvailability::Unknown + | AccountAvailability::QuotaExhausted + | AccountAvailability::Expired + | AccountAvailability::Banned + | AccountAvailability::Invalid => false, + }; + self.enabled && available && self.access_token_expires_at > now } } @@ -595,10 +601,10 @@ pub trait ProviderAccountStore: Send + Sync { update: CredentialCasUpdate, ) -> Result; - async fn get_quota( + async fn get_quotas( &self, - account: &ProviderAccountId, - ) -> Result, StoreError>; + accounts: &[ProviderAccountId], + ) -> Result, StoreError>; async fn compare_and_swap_quota( &self, @@ -677,7 +683,7 @@ impl AccountSelectionPolicy { } } -/// Redis/进程内 owner 提供的可失效调度信号。 +/// Store 提供并发事实，Provider 叠加自己解释的额度事实；全部信号均可重建。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AccountRuntimeSignals { pub in_flight: u32, @@ -686,6 +692,44 @@ pub struct AccountRuntimeSignals { pub quota_remaining_rank: Option, } +impl AccountRuntimeSignals { + #[must_use] + pub fn with_provider_quota(mut self, quota: Option) -> Self { + if let Some(quota) = quota { + self.quota_reset_at = quota.reset_at; + self.quota_remaining_rank = quota.remaining_rank; + } + self + } +} + +/// Provider 从私有 quota JSON 投影出的中立调度事实。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountQuotaSignals { + reset_at: Option, + remaining_rank: Option, +} + +impl AccountQuotaSignals { + #[must_use] + pub const fn new(reset_at: Option, remaining_rank: Option) -> Self { + Self { + reset_at, + remaining_rank, + } + } + + #[must_use] + pub const fn reset_at(self) -> Option { + self.reset_at + } + + #[must_use] + pub const fn remaining_rank(self) -> Option { + self.remaining_rank + } +} + /// 账号持久事实与可重建运行信号的请求级组合。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AccountCandidate { @@ -708,7 +752,7 @@ pub struct AccountSelectionContext { pub struct AccountSelector; impl AccountSelector { - /// 从可调度账号中确定一个候选，不读取 Provider quota JSON。 + /// 从可调度账号中确定一个候选；这里只消费 Provider 已解析的额度投影。 #[must_use] pub fn select<'a>( &self, @@ -735,6 +779,7 @@ impl AccountSelector { match context.policy.strategy() { RotationStrategy::QuotaResetPriority => eligible.sort_by_key(|candidate| { ( + candidate.signals.quota_reset_at.is_none(), candidate.signals.quota_reset_at, candidate.signals.in_flight, candidate.signals.last_started_at, diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs new file mode 100644 index 000000000..8209f5522 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -0,0 +1,844 @@ +//! 数据面执行用例：认证、准入、路由、continuation、circuit 与会话生命周期。 + +use std::collections::BTreeSet; +use std::fmt; +use std::net::IpAddr; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; + +use futures::future::BoxFuture; +use uuid::Uuid; + +use crate::engine::admission::{ + ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, +}; +use crate::engine::continuation::{ + ContinuationBinding, NativeContinuationPort, PreviousResponseId, +}; +use crate::engine::coordinator::ResponseExecutionSession; +use crate::engine::probe::{AccountProbe, AccountProbeRequest, AccountProbeResult}; +use crate::engine::provider::ProviderRegistry; +use crate::engine::{ + AttemptCoordinator, CancellationToken, CommitRequirement, CoordinatedEvent, EngineError, + ExecutionStore, GatewayEngine, ModelRequestId, NewModelRequest, ProviderAttemptOutcome, +}; +use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind}; +use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader}; +use crate::operation::{ + ContentPart, GenerateRequest, Message, MessageRole, Operation, ReasoningEffort, +}; +use crate::policy::{ClientApiKeyId, ClientPolicy}; +use crate::routing::snapshot::RuntimeSnapshotHandle; +use crate::routing::{ProviderInstanceId, PublicModelId, RoutingContext, RuntimeSnapshot}; + +const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientTransport { + HttpJson, + HttpSse, + WebSocket, + InternalProbe, +} + +impl ClientTransport { + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::HttpJson => "http_json", + Self::HttpSse => "http_sse", + Self::WebSocket => "websocket", + Self::InternalProbe => "internal", + } + } +} + +/// API 解码后交给 Core 的稳定请求元数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutionRequestMetadata { + pub protocol: String, + pub endpoint: String, + pub transport: ClientTransport, + pub stream: bool, + pub client_ip: Option, + pub user_agent: Option, + pub previous_response_id: Option, +} + +#[derive(Clone)] +pub struct AuthenticatedClient { + snapshot: Arc, + policy: ClientPolicy, +} + +impl AuthenticatedClient { + #[must_use] + pub const fn snapshot(&self) -> &Arc { + &self.snapshot + } + + #[must_use] + pub const fn policy(&self) -> &ClientPolicy { + &self.policy + } +} + +impl fmt::Debug for AuthenticatedClient { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AuthenticatedClient") + .field("key_id", &self.policy.key_id()) + .field("revision", &self.snapshot.revision()) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum ClientAuthenticationError { + #[error("client API key is invalid")] + InvalidKey, + #[error("runtime snapshot is unavailable")] + SnapshotUnavailable, +} + +pub struct StartExecution { + pub client: AuthenticatedClient, + pub public_model: PublicModelId, + pub operation: Operation, + pub metadata: ExecutionRequestMetadata, +} + +pub struct StartedExecution { + pub request_id: ModelRequestId, + pub created_at: SystemTime, + pub stream: bool, + pub session: Box, +} + +pub trait ExecutionSession: Send { + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>>; + fn response_headers(&self) -> &[ProviderResponseHeader]; + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>>; + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>>; + fn is_finalized(&self) -> bool; + fn cancel(&self); + fn detach_finalize(self: Box) -> BoxFuture<'static, ()>; +} + +pub trait ExecutionService: Send + Sync { + fn authenticate( + &self, + plaintext: &str, + ) -> Result; + fn public_models(&self, client: &AuthenticatedClient) -> Vec; + fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool; + fn start( + &self, + request: StartExecution, + ) -> BoxFuture<'_, Result>; +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderCircuitDecision { + Allow, + BlockedUntil(SystemTime), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("provider circuit store is unavailable")] +pub struct ProviderCircuitError; + +pub trait ProviderCircuitPort: Send + Sync { + fn decision<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result>; + fn observe_failure<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>>; + fn observe_success<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> BoxFuture<'a, Result<(), ProviderCircuitError>>; +} + +pub struct DefaultExecutionService { + snapshots: RuntimeSnapshotHandle, + coordinator: Arc>, + providers: ProviderRegistry, + admissions: Arc, + circuits: Arc, + continuation: Arc, +} + +impl DefaultExecutionService { + #[must_use] + pub fn new( + snapshots: RuntimeSnapshotHandle, + execution: Arc, + providers: ProviderRegistry, + admissions: Arc, + circuits: Arc, + continuation: Arc, + ) -> Self { + let engine = GatewayEngine::::new(execution, providers.clone()); + Self { + snapshots, + coordinator: Arc::new(AttemptCoordinator::new(engine)), + providers, + admissions, + circuits, + continuation, + } + } + + async fn start_inner(&self, request: StartExecution) -> Result { + request.client.policy.authorize().map_err(|_| { + GatewayError::new(GatewayErrorKind::PolicyDenied, "client API key is disabled") + })?; + let started_at = SystemTime::now(); + let deadline_at = started_at + .checked_add(MODEL_REQUEST_DEADLINE) + .ok_or_else(|| { + GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") + })?; + let request_id = new_request_id()?; + let routing_context = self + .route_context( + &request.client.snapshot, + request.client.policy.provider_kind(), + &request.public_model, + &request.operation, + ) + .await?; + let plan = request + .client + .snapshot + .plan(&request.public_model, &request.operation, &routing_context) + .map_err(map_routing_error)?; + let continuation = match request.metadata.previous_response_id.as_ref() { + Some(previous) => match self + .continuation + .resolve(request.client.policy.key_id(), previous) + .await + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::Internal, + "previous response history is temporarily unavailable", + ) + })? { + Some(pin) => ContinuationBinding::Pinned(pin), + None => ContinuationBinding::External(previous.clone()), + }, + None => { + return self + .start_without_continuation( + request, + request_id, + started_at, + deadline_at, + plan, + None, + ) + .await; + } + }; + self.start_without_continuation( + request, + request_id, + started_at, + deadline_at, + plan, + Some(continuation), + ) + .await + } + + async fn start_without_continuation( + &self, + request: StartExecution, + request_id: ModelRequestId, + started_at: SystemTime, + deadline_at: SystemTime, + plan: crate::routing::RoutingPlan, + continuation: Option, + ) -> Result { + let input_token_estimate = request + .operation + .capability_requirements() + .minimum_context_tokens(); + let admission_request = ClientAdmissionRequest { + model_request_id: request_id.clone(), + client_api_key_id: request.client.policy.key_id().clone(), + input_token_estimate, + lease_ttl: MODEL_REQUEST_DEADLINE, + limits: request.client.policy.limits(), + }; + match self + .admissions + .admit(admission_request) + .await + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "request admission is temporarily unavailable", + ) + })? { + ClientAdmissionDecision::Granted => {} + ClientAdmissionDecision::Rejected( + ClientAdmissionRejection::RateLimited + | ClientAdmissionRejection::ConcurrencyLimited, + ) => { + return Err(GatewayError::new( + GatewayErrorKind::RateLimited, + "request exceeds client API key limits", + )); + } + } + let admission = AdmissionLease { + port: Arc::clone(&self.admissions), + client_api_key_id: request.client.policy.key_id().clone(), + model_request_id: request_id.clone(), + }; + let observation = self + .providers + .request_observation(request.client.policy.provider_kind(), &request.operation); + let reasoning_effort = match &request.operation { + Operation::Generate(generate) => generate + .reasoning() + .and_then(|reasoning| reasoning.effort) + .map(reasoning_effort_name) + .map(str::to_owned), + _ => None, + }; + let new_request = NewModelRequest { + id: request_id.clone(), + client_api_key_id: Some(request.client.policy.key_id().clone()), + client_api_key_ref: request.client.policy.key_id().clone(), + config_revision: plan.config_revision(), + protocol: request.metadata.protocol, + operation: request.operation.kind(), + endpoint: request.metadata.endpoint, + client_transport: request.metadata.transport.as_str().to_owned(), + requested_model: request.public_model, + input_token_estimate, + client_ip: request.metadata.client_ip, + user_agent: request.metadata.user_agent, + reasoning_effort, + reasoning_preset: observation.reasoning_preset, + request_kind: observation.request_kind, + subagent_kind: observation.subagent_kind, + compact: observation.compact, + started_at, + deadline_at, + }; + let core = match self + .coordinator + .start( + new_request, + request.operation, + plan, + None, + continuation, + CancellationToken::new(), + ) + .await + { + Ok(core) => core, + Err(error) => { + admission.release().await; + return Err(gateway_error_from_engine(&error)); + } + }; + Ok(StartedExecution { + request_id, + created_at: started_at, + stream: request.metadata.stream, + session: Box::new(DefaultExecutionSession::new( + core, + admission, + Arc::clone(&self.circuits), + )), + }) + } + + async fn route_context( + &self, + snapshot: &RuntimeSnapshot, + provider_kind: &crate::routing::ProviderKind, + model: &PublicModelId, + operation: &Operation, + ) -> Result { + let allowed_instances = snapshot.instance_ids_for_provider(provider_kind); + let preliminary = snapshot + .plan( + model, + operation, + &RoutingContext { + provider_kind: Some(provider_kind.clone()), + allowed_instances: Some(allowed_instances.clone()), + ..RoutingContext::default() + }, + ) + .map_err(map_routing_error)?; + let mut blocked_instances = BTreeSet::new(); + let mut checked = BTreeSet::new(); + for candidate in preliminary.candidates() { + if !checked.insert(candidate.instance().clone()) { + continue; + } + match self + .circuits + .decision(candidate.instance()) + .await + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "provider health state is temporarily unavailable", + ) + })? { + ProviderCircuitDecision::Allow => {} + ProviderCircuitDecision::BlockedUntil(_) => { + blocked_instances.insert(candidate.instance().clone()); + } + } + } + Ok(RoutingContext { + provider_kind: Some(provider_kind.clone()), + allowed_instances: Some(allowed_instances), + blocked_instances, + ..RoutingContext::default() + }) + } + + async fn probe_inner( + &self, + request: AccountProbeRequest, + ) -> Result { + let snapshot = self.snapshots.acquire().map_err(|_| { + GatewayError::new( + GatewayErrorKind::Internal, + "runtime snapshot is unavailable", + ) + })?; + let provider_kind = snapshot + .provider_for_instance(&request.provider_instance_id) + .cloned() + .ok_or_else(|| { + GatewayError::new( + GatewayErrorKind::Unsupported, + "Provider instance is unavailable", + ) + })?; + let message = Message::new( + MessageRole::User, + vec![ContentPart::Text("Reply with exactly OK.".to_owned())], + ) + .map_err(|_| { + GatewayError::new( + GatewayErrorKind::Internal, + "connection test request is invalid", + ) + })?; + let operation = Operation::Generate(GenerateRequest::new(vec![message]).map_err(|_| { + GatewayError::new( + GatewayErrorKind::Internal, + "connection test request is invalid", + ) + })?); + let public_model = + PublicModelId::new(request.upstream_model.as_str().to_owned()).map_err(|_| { + GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid") + })?; + let routing_context = RoutingContext { + provider_kind: Some(provider_kind), + allowed_instances: Some(BTreeSet::from([request.provider_instance_id.clone()])), + ..RoutingContext::default() + }; + let plan = snapshot + .plan(&public_model, &operation, &routing_context) + .map_err(map_routing_error)?; + let started_at = SystemTime::now(); + let deadline_at = started_at + .checked_add(MODEL_REQUEST_DEADLINE) + .ok_or_else(|| { + GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") + })?; + let request_id = new_request_id()?; + let actor = ClientApiKeyId::new("admin_connection_test") + .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "invalid admin actor"))?; + let input_token_estimate = operation.capability_requirements().minimum_context_tokens(); + let new_request = NewModelRequest { + id: request_id, + client_api_key_id: None, + client_api_key_ref: actor, + config_revision: plan.config_revision(), + protocol: "admin_connection_test".to_owned(), + operation: operation.kind(), + endpoint: "/api/admin/accounts/test".to_owned(), + client_transport: ClientTransport::InternalProbe.as_str().to_owned(), + requested_model: public_model, + input_token_estimate, + client_ip: None, + user_agent: None, + reasoning_effort: None, + reasoning_preset: None, + request_kind: Some("account_connection_test".to_owned()), + subagent_kind: None, + compact: false, + started_at, + deadline_at, + }; + let mut session = self + .coordinator + .start( + new_request, + operation, + plan, + Some(request.account_id), + None, + CancellationToken::new(), + ) + .await + .map_err(|error| gateway_error_from_engine(&error))?; + let events = session.collect_uncommitted().await; + publish_provider_attempt_outcomes( + self.circuits.as_ref(), + session.provider_attempt_outcomes(), + ) + .await; + let events = events.map_err(|error| gateway_error_from_engine(&error))?; + session + .commit_downstream(Some(200)) + .await + .map_err(|error| gateway_error_from_engine(&error))?; + Ok(AccountProbeResult { + text: events + .into_iter() + .flat_map(|event| event.into_parts().0) + .filter_map(|fact| match fact { + GatewayEvent::TextDelta(delta) => Some(delta.text), + _ => None, + }) + .collect(), + }) + } +} + +impl ExecutionService for DefaultExecutionService { + fn authenticate( + &self, + plaintext: &str, + ) -> Result { + let snapshot = self + .snapshots + .acquire() + .map_err(|_| ClientAuthenticationError::SnapshotUnavailable)?; + let policy = snapshot + .client_policies() + .filter(|policy| { + constant_time_equal(plaintext, policy.plaintext_key().expose_for_auth()) + }) + .find(|policy| policy.authorize().is_ok()) + .cloned() + .ok_or(ClientAuthenticationError::InvalidKey)?; + Ok(AuthenticatedClient { snapshot, policy }) + } + + fn public_models(&self, client: &AuthenticatedClient) -> Vec { + client + .snapshot + .public_models_for_provider(client.policy.provider_kind()) + } + + fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool { + client + .snapshot + .contains_public_model_for_provider(model, client.policy.provider_kind()) + } + + fn start( + &self, + request: StartExecution, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { self.start_inner(request).await }) + } +} + +impl AccountProbe for DefaultExecutionService { + fn probe( + &self, + request: AccountProbeRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { self.probe_inner(request).await }) + } +} + +struct AdmissionLease { + port: Arc, + client_api_key_id: ClientApiKeyId, + model_request_id: ModelRequestId, +} + +impl AdmissionLease { + async fn release(self) { + let _ = self + .port + .release(&self.client_api_key_id, &self.model_request_id) + .await; + } +} + +struct DefaultExecutionSession { + core: Option>, + admission: Option, + circuits: Arc, + observed_provider_outcomes: usize, +} + +impl DefaultExecutionSession { + fn new( + core: ResponseExecutionSession, + admission: AdmissionLease, + circuits: Arc, + ) -> Self { + Self { + core: Some(core), + admission: Some(admission), + circuits, + observed_provider_outcomes: 0, + } + } + + fn core_mut( + &mut self, + ) -> Result<&mut ResponseExecutionSession, EngineError> { + self.core.as_mut().ok_or(EngineError::InvalidDeliveryState) + } + + async fn settle_if_finalized(&mut self) { + if self + .core + .as_ref() + .is_some_and(ResponseExecutionSession::is_finalized) + && let Some(admission) = self.admission.take() + { + admission.release().await; + } + } + + async fn observe_provider_outcomes(&mut self) { + let Some(core) = self.core.as_ref() else { + return; + }; + let outcomes = core.provider_attempt_outcomes(); + let new_outcomes = outcomes + .get(self.observed_provider_outcomes..) + .unwrap_or_default() + .to_vec(); + self.observed_provider_outcomes = outcomes.len(); + publish_provider_attempt_outcomes(self.circuits.as_ref(), &new_outcomes).await; + } + + async fn finalize_detached(&mut self) { + let Some(mut core) = self.core.take() else { + return; + }; + core.cancel(); + if !core.is_finalized() { + let _ = core.cancel_and_finalize().await; + } + let pending = core + .provider_attempt_outcomes() + .get(self.observed_provider_outcomes..) + .unwrap_or_default(); + publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await; + if let Some(admission) = self.admission.take() { + admission.release().await; + } + } +} + +impl Drop for DefaultExecutionSession { + fn drop(&mut self) { + if let Some(core) = &self.core { + core.cancel(); + } + } +} + +impl ExecutionSession for DefaultExecutionSession { + fn next_event(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + let result = self.core_mut()?.next_event().await; + self.observe_provider_outcomes().await; + self.settle_if_finalized().await; + result + }) + } + + fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result, EngineError>> { + Box::pin(async move { + let result = self.core_mut()?.collect_uncommitted().await; + self.observe_provider_outcomes().await; + self.settle_if_finalized().await; + result + }) + } + + fn response_headers(&self) -> &[ProviderResponseHeader] { + self.core + .as_ref() + .map(ResponseExecutionSession::response_headers) + .unwrap_or_default() + } + + fn commit_downstream( + &mut self, + client_status_code: Option, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self.core_mut()?.commit_downstream(client_status_code).await; + self.observe_provider_outcomes().await; + self.settle_if_finalized().await; + result + }) + } + + fn record_client_status( + &mut self, + client_status_code: u16, + ) -> BoxFuture<'_, Result<(), EngineError>> { + Box::pin(async move { + let result = self + .core_mut()? + .record_client_status(client_status_code) + .await; + self.observe_provider_outcomes().await; + self.settle_if_finalized().await; + result + }) + } + + fn is_finalized(&self) -> bool { + self.core + .as_ref() + .is_none_or(ResponseExecutionSession::is_finalized) + } + + fn cancel(&self) { + if let Some(core) = &self.core { + core.cancel(); + } + } + + fn detach_finalize(mut self: Box) -> BoxFuture<'static, ()> { + Box::pin(async move { self.finalize_detached().await }) + } +} + +#[must_use] +pub const fn provider_failure_affects_circuit(error_kind: ProviderErrorKind) -> bool { + matches!( + error_kind, + ProviderErrorKind::Timeout + | ProviderErrorKind::Transport + | ProviderErrorKind::Protocol + | ProviderErrorKind::Unavailable + ) +} + +async fn publish_provider_attempt_outcomes( + circuits: &dyn ProviderCircuitPort, + outcomes: &[ProviderAttemptOutcome], +) { + for outcome in outcomes { + let result = match outcome.error_kind() { + None => { + circuits + .observe_success(outcome.provider_instance_id()) + .await + } + Some(kind) if provider_failure_affects_circuit(kind) => { + circuits + .observe_failure(outcome.provider_instance_id()) + .await + } + Some(_) => continue, + }; + let _ = result; + } +} + +fn constant_time_equal(left: &str, right: &str) -> bool { + if left.len() != right.len() { + return false; + } + left.as_bytes() + .iter() + .zip(right.as_bytes()) + .fold(0_u8, |difference, (left, right)| { + difference | (left ^ right) + }) + == 0 +} + +fn new_request_id() -> Result { + ModelRequestId::new(format!("req_{}", Uuid::now_v7().simple())) + .map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "failed to allocate request ID")) +} + +const fn reasoning_effort_name(effort: ReasoningEffort) -> &'static str { + match effort { + ReasoningEffort::Minimal => "minimal", + ReasoningEffort::Low => "low", + ReasoningEffort::Medium => "medium", + ReasoningEffort::High => "high", + ReasoningEffort::ExtraHigh => "xhigh", + } +} + +fn map_routing_error(error: crate::error::RoutingError) -> GatewayError { + match error { + crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no Provider instance can execute this request", + ), + _ => GatewayError::new( + GatewayErrorKind::Internal, + "runtime routing configuration is invalid", + ), + } +} + +pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { + match error { + EngineError::Cancelled => { + GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled") + } + EngineError::Deadline => { + GatewayError::new(GatewayErrorKind::Timeout, "request deadline elapsed") + } + EngineError::Provider(provider) => GatewayError::from_provider(provider), + EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => { + GatewayError::new( + GatewayErrorKind::NoAvailableProvider, + "no Provider instance is available", + ) + } + _ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"), + } +} + +#[must_use] +pub const fn commit_requirement(event: &CoordinatedEvent) -> CommitRequirement { + event.commit_requirement() +} diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 9bdca3cfe..e88a1a99b 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -1,8 +1,11 @@ //! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。 +pub mod admission; pub mod continuation; pub mod coordinator; pub mod credential; +pub mod execution; +pub mod probe; pub mod provider; pub use coordinator::{AttemptCoordinator, ResponseExecutionSession}; @@ -20,12 +23,12 @@ use futures::channel::oneshot; use thiserror::Error; use crate::accounting::{CostEstimate, Usage}; -use crate::engine::continuation::NativeContinuationPin; +use crate::engine::continuation::{ContinuationBinding, NativeContinuationPin}; use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId}; use crate::error::{ GatewayError, IdentifierError, ProviderError, ProviderErrorKind, StoreError, validate_text, }; -use crate::event::GatewayEvent; +use crate::event::ProviderEvent; use crate::operation::OperationKind; use crate::policy::ClientApiKeyId; use crate::routing::{ @@ -230,25 +233,28 @@ fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { .unwrap_or_else(std::sync::PoisonError::into_inner) } -/// 单次 attempt 的账号选择硬约束。 +/// 单次 attempt 的账号选择与账号绑定状态事实。 /// /// `required_account` 只用于管理端 connection test 等必须命中唯一账号的内部请求； /// 一旦设置，Core 与 Provider 都不得换号或切换 target。 #[derive(Debug, Clone, Default)] -pub struct AccountSelectionConstraints { +pub struct AccountAttemptContext { excluded_accounts: BTreeSet, required_account: Option, + state_owner: Option, } -impl AccountSelectionConstraints { +impl AccountAttemptContext { #[must_use] pub const fn new( excluded_accounts: BTreeSet, required_account: Option, + state_owner: Option, ) -> Self { Self { excluded_accounts, required_account, + state_owner, } } @@ -261,6 +267,84 @@ impl AccountSelectionConstraints { pub const fn required_account(&self) -> Option<&ProviderAccountId> { self.required_account.as_ref() } + + #[must_use] + pub const fn state_owner(&self) -> Option<&ProviderAccountStateOwner> { + self.state_owner.as_ref() + } +} + +/// 请求中 Provider 账号绑定状态的唯一归属。 +/// +/// `turn_state` 等 opaque 状态只能发送给创建它的 Provider、instance 与账号； +/// Core 在首次真实选号后冻结该事实，Provider 据此决定是否清理跨账号状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderAccountStateOwner { + provider: ProviderKind, + instance: ProviderInstanceId, + account: ProviderAccountId, +} + +impl ProviderAccountStateOwner { + #[must_use] + pub const fn new( + provider: ProviderKind, + instance: ProviderInstanceId, + account: ProviderAccountId, + ) -> Self { + Self { + provider, + instance, + account, + } + } + + #[must_use] + pub fn from_continuation(pin: &NativeContinuationPin) -> Self { + Self::new( + pin.provider().clone(), + pin.instance().clone(), + pin.account().clone(), + ) + } + + #[must_use] + pub fn matches( + &self, + provider: &ProviderKind, + instance: &ProviderInstanceId, + account: &ProviderAccountId, + ) -> bool { + self.provider == *provider && self.instance == *instance && self.account == *account + } + + #[must_use] + pub const fn provider(&self) -> &ProviderKind { + &self.provider + } + + #[must_use] + pub const fn instance(&self) -> &ProviderInstanceId { + &self.instance + } + + #[must_use] + pub const fn account(&self) -> &ProviderAccountId { + &self.account + } +} + +/// 当前 attempt 对 previous-response 的唯一处理方式。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ContinuationAttempt { + /// 没有 previous-response。 + None, + /// 使用 Store 解析出的原生 handle 与账号绑定。 + Native, + /// 在原账号上用完整连接内 transcript 重放。 + ReplayOwner, + /// transcript 已可携带，允许选择其他账号。 + ReplayAny, } /// Provider 每次执行可见的 request-local context。 @@ -270,8 +354,9 @@ pub struct AttemptContext { attempt_index: NonZeroU32, deadline: SystemTime, account_selection_policy: AccountSelectionPolicy, - account_constraints: AccountSelectionConstraints, - continuation: Option, + account: AccountAttemptContext, + continuation: Option, + continuation_attempt: ContinuationAttempt, cancellation: CancellationToken, } @@ -282,21 +367,37 @@ impl AttemptContext { attempt_index: NonZeroU32, deadline: SystemTime, account_selection_policy: AccountSelectionPolicy, - account_constraints: AccountSelectionConstraints, - continuation: Option, + account: AccountAttemptContext, + continuation: Option, cancellation: CancellationToken, ) -> Self { + let continuation_attempt = if continuation.is_some() { + ContinuationAttempt::Native + } else { + ContinuationAttempt::None + }; Self { request_id, attempt_index, deadline, account_selection_policy, - account_constraints, + account, continuation, + continuation_attempt, cancellation, } } + /// 覆盖本次 attempt 的 continuation 恢复方式。 + #[must_use] + pub const fn with_continuation_attempt( + mut self, + continuation_attempt: ContinuationAttempt, + ) -> Self { + self.continuation_attempt = continuation_attempt; + self + } + #[must_use] pub const fn request_id(&self) -> &ModelRequestId { &self.request_id @@ -319,20 +420,30 @@ impl AttemptContext { #[must_use] pub const fn excluded_accounts(&self) -> &BTreeSet { - self.account_constraints.excluded_accounts() + self.account.excluded_accounts() } /// 管理端 connection test 等内部请求强制使用的唯一账号。 #[must_use] pub const fn required_account(&self) -> Option<&ProviderAccountId> { - self.account_constraints.required_account() + self.account.required_account() } #[must_use] - pub const fn continuation(&self) -> Option<&NativeContinuationPin> { + pub const fn account_state_owner(&self) -> Option<&ProviderAccountStateOwner> { + self.account.state_owner() + } + + #[must_use] + pub const fn continuation(&self) -> Option<&ContinuationBinding> { self.continuation.as_ref() } + #[must_use] + pub const fn continuation_attempt(&self) -> ContinuationAttempt { + self.continuation_attempt + } + #[must_use] pub const fn cancellation(&self) -> &CancellationToken { &self.cancellation @@ -388,6 +499,8 @@ pub struct IntermediateFailure { pub provider_kind: ProviderKind, pub account_id: Option, pub upstream_model_id: UpstreamModelId, + pub upstream_status_code: Option, + pub upstream_request_id: Option, pub error: ProviderError, pub latency: Duration, } @@ -419,6 +532,8 @@ pub struct ModelRequestFinalization { pub upstream_status_code: Option, pub upstream_request_id: Option, pub upstream_response_id: Option, + pub upstream_transport: Option, + pub http_version: Option, pub error: Option, pub provider_error_code: Option, pub retry_after_ms: Option, @@ -476,27 +591,45 @@ pub enum CommitRequirement { #[derive(Debug)] pub struct CoordinatedEvent { - event: GatewayEvent, + events: Vec, commit_requirement: CommitRequirement, } impl CoordinatedEvent { #[must_use] - pub const fn new(event: GatewayEvent, commit_requirement: CommitRequirement) -> Self { + pub fn single(event: ProviderEvent, commit_requirement: CommitRequirement) -> Self { Self { - event, + events: vec![event], commit_requirement, } } + /// 创建一个必须在同一下游提交边界内按序编码的事件批次。 + /// + /// # Errors + /// + /// 空批次无法建立交付边界。 + pub fn try_batch( + events: Vec, + commit_requirement: CommitRequirement, + ) -> Result { + if events.is_empty() { + return Err(EngineError::InvalidDeliveryState); + } + Ok(Self { + events, + commit_requirement, + }) + } + #[must_use] pub const fn commit_requirement(&self) -> CommitRequirement { self.commit_requirement } #[must_use] - pub fn into_event(self) -> GatewayEvent { - self.event + pub fn into_provider_events(self) -> Vec { + self.events } } @@ -527,12 +660,12 @@ pub enum EngineError { } /// Engine 只组合 Store 与 Provider Registry；重试算法完全位于 coordinator。 -pub struct GatewayEngine { +pub struct GatewayEngine { store: Arc, providers: provider::ProviderRegistry, } -impl GatewayEngine { +impl GatewayEngine { #[must_use] pub const fn new(store: Arc, providers: provider::ProviderRegistry) -> Self { Self { store, providers } diff --git a/backend/crates/gateway-core/src/engine/probe.rs b/backend/crates/gateway-core/src/engine/probe.rs new file mode 100644 index 000000000..904878854 --- /dev/null +++ b/backend/crates/gateway-core/src/engine/probe.rs @@ -0,0 +1,26 @@ +//! 管理端账号连通性测试的真实执行链端口。 + +use futures::future::BoxFuture; + +use crate::engine::credential::ProviderAccountId; +use crate::error::GatewayError; +use crate::routing::{ProviderInstanceId, UpstreamModelId}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountProbeRequest { + pub account_id: ProviderAccountId, + pub provider_instance_id: ProviderInstanceId, + pub upstream_model: UpstreamModelId, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountProbeResult { + pub text: Vec, +} + +pub trait AccountProbe: Send + Sync { + fn probe( + &self, + request: AccountProbeRequest, + ) -> BoxFuture<'_, Result>; +} diff --git a/backend/crates/gateway-core/src/engine/provider.rs b/backend/crates/gateway-core/src/engine/provider.rs index 443102400..54dcec7e6 100644 --- a/backend/crates/gateway-core/src/engine/provider.rs +++ b/backend/crates/gateway-core/src/engine/provider.rs @@ -15,16 +15,16 @@ use crate::engine::{AttemptContext, UpstreamSendState}; use crate::error::{ IdentifierError, ProviderError, ProviderErrorKind, SafeUpstreamValue, validate_text, }; -use crate::event::{EventSequenceValidator, GatewayEvent}; +use crate::event::{EventSequenceValidator, ProviderEvent}; use crate::operation::Operation; use crate::routing::{ ModelCapabilities, ProviderCandidate, ProviderInstance, ProviderInstanceId, ProviderKind, UpstreamModelId, }; -/// Box 只出现在 Provider Registry 的 event stream 边界。 +/// Box 只出现在 Provider Registry 的统一 event envelope 边界。 pub type EventStream = - Pin> + Send + 'static>>; + Pin> + Send + 'static>>; /// 匿名 credential/resource 引用。 /// @@ -235,14 +235,13 @@ pub struct ProviderStream { impl ProviderStream { /// 组装一次、且仅一次可见上游调用。 #[must_use] - pub fn new( - metadata: ProviderCallMetadata, - events: EventStream, - lease: impl ResourceLease, - ) -> Self { + pub fn new(metadata: ProviderCallMetadata, events: S, lease: impl ResourceLease) -> Self + where + S: Stream> + Send + 'static, + { Self { metadata, - events, + events: Box::pin(events), _lease: Box::new(lease), validator: EventSequenceValidator::new(), terminated: false, @@ -257,7 +256,7 @@ impl ProviderStream { } impl Stream for ProviderStream { - type Item = Result; + type Item = Result; fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { let this = self.get_mut(); @@ -267,16 +266,18 @@ impl Stream for ProviderStream { match this.events.as_mut().poll_next(context) { Poll::Pending => Poll::Pending, - Poll::Ready(Some(Ok(event))) => match this.validator.observe(&event) { - Ok(()) => Poll::Ready(Some(Ok(event))), - Err(_) => { - this.terminated = true; - Poll::Ready(Some(Err(ProviderError::new( - ProviderErrorKind::Protocol, - UpstreamSendState::Sent, - )))) + Poll::Ready(Some(Ok(event))) => { + for fact in event.canonical_facts() { + if this.validator.observe(fact).is_err() { + this.terminated = true; + return Poll::Ready(Some(Err(ProviderError::new( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )))); + } } - }, + Poll::Ready(Some(Ok(event))) + } Poll::Ready(Some(Err(error))) => { this.terminated = true; Poll::Ready(Some(Err(error))) @@ -302,6 +303,15 @@ pub struct ProviderRequest { candidate: ProviderCandidate, } +/// Provider 对公共观测表可解释的请求语义；未知字段保持空值。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ProviderRequestObservation { + pub reasoning_preset: Option, + pub request_kind: Option, + pub subagent_kind: Option, + pub compact: bool, +} + /// Provider 实时目录编译后的单模型能力。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProviderModelCapabilities { @@ -309,6 +319,24 @@ pub struct ProviderModelCapabilities { capabilities: ModelCapabilities, } +/// Provider 实时目录成功发布后的进程内单调代次。 +/// +/// 代次只表达“目录内容已经变化”，不承载模型、ETag 或 Provider 私有数据。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderCatalogGeneration(u64); + +impl ProviderCatalogGeneration { + #[must_use] + pub const fn new(value: u64) -> Self { + Self(value) + } + + #[must_use] + pub const fn get(self) -> u64 { + self.0 + } +} + impl ProviderModelCapabilities { #[must_use] pub const fn new(upstream_model: UpstreamModelId, capabilities: ModelCapabilities) -> Self { @@ -374,6 +402,14 @@ pub trait Provider: Send + Sync { /// 返回编译期注册名称。 fn name(&self) -> &'static str; + /// 返回当前进程已经成功发布的目录代次。 + fn catalog_generation(&self) -> ProviderCatalogGeneration; + + /// 解释 Provider 差异化观测字段；不参与路由和传输。 + fn request_observation(&self, _operation: &Operation) -> ProviderRequestObservation { + ProviderRequestObservation::default() + } + /// 查询 instance 的实时模型目录，并由 Provider 自己编译能力事实。 /// /// # Errors @@ -466,6 +502,17 @@ pub struct ProviderRegistry { } impl ProviderRegistry { + /// 从一组异构 Provider 构造冻结注册表。 + pub fn new( + providers: impl IntoIterator>, + ) -> Result { + let mut builder = ProviderRegistryBuilder::new(); + for provider in providers { + builder.register(provider)?; + } + Ok(builder.build()) + } + /// 创建 builder。 #[must_use] pub const fn builder() -> ProviderRegistryBuilder { @@ -478,6 +525,28 @@ impl ProviderRegistry { self.providers.get(provider) } + #[must_use] + pub fn request_observation( + &self, + provider: &ProviderKind, + operation: &Operation, + ) -> ProviderRequestObservation { + self.providers + .get(provider) + .map_or_else(ProviderRequestObservation::default, |registered| { + registered.request_observation(operation) + }) + } + + /// 读取全部 Provider 的目录代次，用于 Core 快照稳定性校验与对账。 + #[must_use] + pub fn catalog_generations(&self) -> BTreeMap { + self.providers + .iter() + .map(|(kind, provider)| (kind.clone(), provider.catalog_generation())) + .collect() + } + /// 通过编译期 Provider adapter 查询并编译 instance 的实时能力目录。 /// /// # Errors diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index 4282e1b8d..6d0b869b5 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -174,6 +174,15 @@ impl ProviderErrorKind { } } +/// Provider 对 previous-response 失败给出的稳定恢复语义。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ContinuationFailure { + /// 原生 continuation 正被同账号的另一请求占用。 + Busy, + /// 原生 history 在当前连接或账号不可用，可由完整输入重放恢复。 + HistoryUnavailable, +} + /// Adapter 已明确分类为非 bearer、可用于诊断的上游值。 #[derive(Clone, PartialEq, Eq, Hash)] pub struct SafeUpstreamValue(String); @@ -218,6 +227,7 @@ pub struct ProviderError { upstream_request_id: Option, upstream_response_id: Option, retry_after: Option, + continuation_failure: Option, replay_safe: bool, sensitive_context_redacted: bool, } @@ -234,6 +244,7 @@ impl ProviderError { upstream_request_id: None, upstream_response_id: None, retry_after: None, + continuation_failure: None, replay_safe: false, sensitive_context_redacted: false, } @@ -276,6 +287,13 @@ impl ProviderError { self } + /// 附加 previous-response 的恢复语义。 + #[must_use] + pub const fn with_continuation_failure(mut self, failure: ContinuationFailure) -> Self { + self.continuation_failure = Some(failure); + self + } + /// 标记 Provider 已证明本次拒绝没有执行生成，可在下游提交前重放。 #[must_use] pub const fn with_replay_safe(mut self) -> Self { @@ -332,6 +350,12 @@ impl ProviderError { self.retry_after } + /// 返回 previous-response 的恢复语义。 + #[must_use] + pub const fn continuation_failure(&self) -> Option { + self.continuation_failure + } + /// 返回 Provider 是否已证明本次失败可安全重放。 #[must_use] pub const fn replay_is_safe(&self) -> bool { @@ -371,6 +395,7 @@ impl fmt::Debug for ProviderError { .map(|_| ""), ) .field("retry_after", &self.retry_after) + .field("continuation_failure", &self.continuation_failure) .field("replay_safe", &self.replay_safe) .field( "sensitive_context", diff --git a/backend/crates/gateway-core/src/event.rs b/backend/crates/gateway-core/src/event.rs index cb4e6ab5c..1bd83b971 100644 --- a/backend/crates/gateway-core/src/event.rs +++ b/backend/crates/gateway-core/src/event.rs @@ -3,10 +3,13 @@ use std::collections::BTreeMap; use std::fmt; +use serde_json::Value; use thiserror::Error; use crate::accounting::{CalculatedCost, ProviderReportedCost, Usage}; -use crate::error::SafeUpstreamValue; +use crate::engine::provider::UpstreamTransport; +use crate::error::{IdentifierError, SafeUpstreamValue, validate_text}; +use crate::operation::ProviderSessionState; /// 一次响应的稳定元数据。 #[derive(Debug, Clone, PartialEq, Eq)] @@ -200,6 +203,250 @@ impl fmt::Debug for ToolCallDelta { } } +/// Provider 返回的协议原生 JSON event。 +/// +/// Core 不解释 `data`，只把它与同一上游事件产生的 canonical facts 一起 +/// 交给客户端协议 adapter。Debug 永不输出正文。 +#[derive(Clone, PartialEq)] +pub struct ProtocolWireEvent { + protocol: String, + event_type: Option, + data: Value, +} + +impl ProtocolWireEvent { + /// 创建协议原生 JSON event。 + /// + /// # Errors + /// + /// 协议名或显式事件名为空、过长或含控制字符时返回错误。 + pub fn json( + protocol: impl Into, + event_type: Option, + data: Value, + ) -> Result { + let protocol = protocol.into(); + validate_text(&protocol, 64, true, None)?; + if let Some(event_type) = event_type.as_deref() { + validate_text(event_type, 256, true, None)?; + } + Ok(Self { + protocol, + event_type, + data, + }) + } + + /// 返回客户端协议名称。 + #[must_use] + pub fn protocol(&self) -> &str { + &self.protocol + } + + /// 返回可选的协议事件名称。 + #[must_use] + pub fn event_type(&self) -> Option<&str> { + self.event_type.as_deref() + } + + /// 返回协议原生 JSON 数据。 + #[must_use] + pub const fn data(&self) -> &Value { + &self.data + } + + /// 拆出协议原生 JSON 数据。 + #[must_use] + pub fn into_data(self) -> Value { + self.data + } +} + +impl fmt::Debug for ProtocolWireEvent { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProtocolWireEvent") + .field("protocol", &self.protocol) + .field("event_type", &self.event_type) + .field("data", &"") + .finish() + } +} + +/// 上游响应使用的 HTTP 协议版本。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UpstreamHttpVersion { + Unknown, + Http09, + Http10, + Http11, + Http2, + Http3, +} + +impl UpstreamHttpVersion { + /// 解析 transport 已规范化的协议版本。 + #[must_use] + pub fn parse(value: &str) -> Option { + match value.trim().to_ascii_uppercase().as_str() { + "HTTP/0.9" | "0.9" => Some(Self::Http09), + "HTTP/1.0" | "HTTP/1" | "1.0" => Some(Self::Http10), + "HTTP/1.1" | "1.1" => Some(Self::Http11), + "HTTP/2" | "HTTP/2.0" | "2" | "2.0" => Some(Self::Http2), + "HTTP/3" | "HTTP/3.0" | "3" | "3.0" => Some(Self::Http3), + _ => None, + } + } + + /// 返回数据库使用的稳定名称。 + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Unknown => "unknown", + Self::Http09 => "HTTP/0.9", + Self::Http10 => "HTTP/1.0", + Self::Http11 => "HTTP/1.1", + Self::Http2 => "HTTP/2", + Self::Http3 => "HTTP/3", + } + } +} + +/// Provider transport 边界测得的阶段耗时。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ProviderResponseTimings { + pub transport_decision_wait_ms: Option, + pub connect_ms: Option, + pub headers_ms: Option, + pub first_event_ms: Option, +} + +/// Provider 已筛选、可由协议 adapter 再次按白名单表达的响应头。 +#[derive(Clone, PartialEq, Eq)] +pub struct ProviderResponseHeader { + name: String, + value: SafeUpstreamValue, +} + +impl ProviderResponseHeader { + /// 创建规范化的小写响应头；非 ASCII token 名称直接拒绝。 + #[must_use] + pub fn new(name: impl Into, value: SafeUpstreamValue) -> Option { + let name = name.into().trim().to_ascii_lowercase(); + let valid = !name.is_empty() + && name.len() <= 64 + && name + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-'); + valid.then_some(Self { name, value }) + } + + #[must_use] + pub fn name(&self) -> &str { + &self.name + } + + #[must_use] + pub const fn value(&self) -> &SafeUpstreamValue { + &self.value + } +} + +impl fmt::Debug for ProviderResponseHeader { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderResponseHeader") + .field("name", &self.name) + .field("value", &"") + .finish() + } +} + +/// Core 消费的实际上游响应事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderResponseObservation { + transport: UpstreamTransport, + http_version: Option, + status_code: Option, + request_id: Option, + timings: ProviderResponseTimings, + client_headers: Vec, +} + +impl ProviderResponseObservation { + #[must_use] + pub fn new(transport: UpstreamTransport) -> Self { + Self { + transport, + http_version: None, + status_code: None, + request_id: None, + timings: ProviderResponseTimings::default(), + client_headers: Vec::new(), + } + } + + #[must_use] + pub const fn with_http_version(mut self, version: UpstreamHttpVersion) -> Self { + self.http_version = Some(version); + self + } + + #[must_use] + pub fn with_status_code(mut self, status_code: u16) -> Self { + self.status_code = (100..=599).contains(&status_code).then_some(status_code); + self + } + + #[must_use] + pub fn with_request_id(mut self, request_id: SafeUpstreamValue) -> Self { + self.request_id = Some(request_id); + self + } + + #[must_use] + pub const fn with_timings(mut self, timings: ProviderResponseTimings) -> Self { + self.timings = timings; + self + } + + #[must_use] + pub fn with_client_headers(mut self, client_headers: Vec) -> Self { + self.client_headers = client_headers; + self + } + + #[must_use] + pub const fn transport(&self) -> &UpstreamTransport { + &self.transport + } + + #[must_use] + pub const fn http_version(&self) -> Option { + self.http_version + } + + #[must_use] + pub const fn status_code(&self) -> Option { + self.status_code + } + + #[must_use] + pub const fn request_id(&self) -> Option<&SafeUpstreamValue> { + self.request_id.as_ref() + } + + #[must_use] + pub const fn timings(&self) -> ProviderResponseTimings { + self.timings + } + + #[must_use] + pub fn client_headers(&self) -> &[ProviderResponseHeader] { + &self.client_headers + } +} + /// 所有 Provider 都必须输出的稳定事件集合。 #[derive(Debug, Clone, PartialEq)] #[non_exhaustive] @@ -224,6 +471,154 @@ pub enum GatewayEvent { Completed(ResponseMeta), } +/// Provider 单个上游事件产生的事实与可选协议原生表达。 +/// +/// 一个值至少包含一个 canonical fact 或一条 wire event。把同一 wire event +/// 产生的多个 canonical facts 放在同一封套，可避免客户端重复收到该事件。 +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderEvent { + canonical: Vec, + wire: Option>, + observation: Option>, + session_update: Option>, +} + +impl ProviderEvent { + /// 创建只有一个 canonical fact 的事件。 + #[must_use] + pub fn canonical(event: GatewayEvent) -> Self { + Self { + canonical: vec![event], + wire: None, + observation: None, + session_update: None, + } + } + + /// 创建只有协议原生表达的事件。 + #[must_use] + pub fn wire(wire: ProtocolWireEvent) -> Self { + Self { + canonical: Vec::new(), + wire: Some(Box::new(wire)), + observation: None, + session_update: None, + } + } + + /// 创建同一上游事件的 canonical facts 与协议原生表达。 + #[must_use] + pub fn canonical_with_wire(canonical: Vec, wire: ProtocolWireEvent) -> Self { + Self { + canonical, + wire: Some(Box::new(wire)), + observation: None, + session_update: None, + } + } + + /// 创建仅供 Core 消费的上游响应观察；该事件不会进入客户端 adapter。 + #[must_use] + pub fn observation(observation: ProviderResponseObservation) -> Self { + Self { + canonical: Vec::new(), + wire: None, + observation: Some(Box::new(observation)), + session_update: None, + } + } + + /// 附着只由同协议客户端连接保存、并在下一轮原样交还 Provider 的状态。 + #[must_use] + pub fn with_session_update(mut self, state: ProviderSessionState) -> Self { + self.session_update = Some(Box::new(state)); + self + } + + /// 把本事件标记为 Provider 连接内状态的提交边界。 + pub fn attach_session_update(&mut self, state: ProviderSessionState) { + self.session_update = Some(Box::new(state)); + } + + /// 返回 Provider 私有的连接内状态更新。 + #[must_use] + pub fn session_update(&self) -> Option<&ProviderSessionState> { + self.session_update.as_deref() + } + + /// 取出 Provider 连接内状态更新，交给协议连接持有。 + #[must_use] + pub fn take_session_update(&mut self) -> Option { + self.session_update.take().map(|state| *state) + } + + /// 返回 Core 可解释的全部 facts。 + #[must_use] + pub fn canonical_facts(&self) -> &[GatewayEvent] { + &self.canonical + } + + /// 返回 Core 可改写客户端响应身份的全部 facts。 + #[must_use] + pub fn canonical_facts_mut(&mut self) -> &mut [GatewayEvent] { + &mut self.canonical + } + + /// 返回协议原生表达。 + #[must_use] + pub fn wire_event(&self) -> Option<&ProtocolWireEvent> { + self.wire.as_deref() + } + + /// 取出仅供 Core 持久化的响应观察。 + #[must_use] + pub fn take_observation(&mut self) -> Option { + self.observation.take().map(|observation| *observation) + } + + /// 返回仅供 Core 持久化的响应观察。 + #[must_use] + pub fn response_observation(&self) -> Option<&ProviderResponseObservation> { + self.observation.as_deref() + } + + /// 拆分 canonical facts 与协议原生表达。 + #[must_use] + pub fn into_parts(self) -> (Vec, Option) { + (self.canonical, self.wire.map(|wire| *wire)) + } + + /// 返回是否含有可用于 commit barrier 的 canonical fact。 + #[must_use] + pub fn has_canonical_facts(&self) -> bool { + !self.canonical.is_empty() + } + + /// 返回该封套是否仍包含可交付客户端的表达。 + #[must_use] + pub fn has_client_event(&self) -> bool { + !self.canonical.is_empty() || self.wire.is_some() + } + + /// 返回是否包含足以冻结下游 commit barrier 的 canonical fact。 + /// + /// 单独的 wire event 不能冻结边界：协议层在看到 `Started` + /// 携带的网关响应身份前无法安全编码它。Core 会保留这些 + /// 事件，直到首个 canonical 可见增量或终态再一并交付。 + #[must_use] + pub fn is_commit_significant(&self) -> bool { + self.canonical + .iter() + .any(GatewayEvent::is_commit_significant) + } +} + +impl From for ProviderEvent { + fn from(event: GatewayEvent) -> Self { + Self::canonical(event) + } +} + impl GatewayEvent { /// Canonical event 均为客户端 encoder 可交付事件。 /// @@ -234,6 +629,21 @@ impl GatewayEvent { true } + /// 返回该 fact 是否足以冻结下游 commit barrier。 + /// + /// 生命周期、结构与结算 facts 可以在换号前安全丢弃；首个可见增量或 + /// Provider 正常终态才会使客户端输出不可撤回。 + #[must_use] + pub const fn is_commit_significant(&self) -> bool { + matches!( + self, + Self::TextDelta(_) + | Self::ReasoningDelta(_) + | Self::ToolCallDelta(_) + | Self::Completed(_) + ) + } + /// 冻结客户端可见的网关 response ID，并返回 Provider 原生 response ID。 /// /// 只有 Coordinator 会调用该边界；Provider 产生的 ID 永远不会直接下发。 diff --git a/backend/crates/gateway-core/src/health.rs b/backend/crates/gateway-core/src/health.rs new file mode 100644 index 000000000..030fb058c --- /dev/null +++ b/backend/crates/gateway-core/src/health.rs @@ -0,0 +1,60 @@ +//! Store、Core 与 Host 向 HTTP 健康探针暴露的中立契约。 + +use std::time::SystemTime; + +use futures::future::BoxFuture; + +use crate::task::{WorkerFencingToken, WorkerId, WorkerKind}; + +/// 一个健康探针的语义状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HealthState { + Healthy, + Degraded(String), + Unhealthy(String), +} + +/// 可被 API 异质聚合的异步健康探针。 +pub trait HealthProbe: Send + Sync { + fn name(&self) -> &'static str; + + fn check(&self) -> BoxFuture<'_, HealthState>; +} + +/// Host 监督器对一个 worker 观测到的当前状态。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerRuntimeState { + Starting, + AcquiringLease, + Standby, + Running, + Idle, + BackingOff, + Stopped, + Disabled, +} + +/// worker 健康快照的稳定 key。 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum WorkerHealthKey { + Task(WorkerId), + Disabled(WorkerKind), +} + +/// 一个 worker 在某一时刻的完整可观测快照。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerHealthSnapshot { + pub key: WorkerHealthKey, + pub state: WorkerRuntimeState, + pub consecutive_failures: u32, + pub completed_cycles: u64, + pub last_fencing_token: Option, + pub last_success_at: Option, + pub last_failure_at: Option, + pub last_error: Option, +} + +/// Host 提供、API 消费的 worker 健康快照来源。 +pub trait WorkerHealthSource: Send + Sync { + fn snapshot(&self) -> Vec; +} diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index c8ca9f60d..36f8a252d 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -3,12 +3,160 @@ //! 本 crate 只描述协议与 Provider 无关的业务语义。HTTP、数据库、Redis、 //! 具体客户端协议和具体 Provider 都通过外层 adapter 接入。 -#![forbid(unsafe_code)] - pub mod accounting; pub mod engine; pub mod error; pub mod event; +pub mod health; +pub mod lifecycle; pub mod operation; pub mod policy; +pub mod provider_ports; pub mod routing; +pub mod task; + +use std::sync::Arc; +use std::time::SystemTime; + +use engine::ExecutionStore; +use engine::admission::{ + ClientAdmissionPort, ClientAdmissionRecoveryPort, restore_client_admission_startup, +}; +use engine::continuation::NativeContinuationPort; +use engine::execution::{DefaultExecutionService, ExecutionService, ProviderCircuitPort}; +use engine::probe::AccountProbe; +use engine::provider::ProviderRegistry; +use health::HealthProbe; +use routing::snapshot::{ + RuntimeSnapshotCompiler, RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, + SnapshotStorePort, SnapshotSubscriptionPort, +}; +use task::{WorkerContribution, WorkerDefinitionError}; + +/// Store 提供给数据面 Core 的封闭能力集合。 +#[derive(Clone)] +pub struct CoreStorePorts { + execution: Arc, + admissions: Arc, + admission_recovery: Arc, + circuits: Arc, + continuation: Arc, + snapshots: Arc, + snapshot_subscriptions: Arc, +} + +impl CoreStorePorts { + #[must_use] + pub fn new( + execution: Arc, + admissions: Arc, + admission_recovery: Arc, + circuits: Arc, + continuation: Arc, + snapshots: Arc, + snapshot_subscriptions: Arc, + ) -> Self { + Self { + execution, + admissions, + admission_recovery, + circuits, + continuation, + snapshots, + snapshot_subscriptions, + } + } +} + +pub struct CoreBundle { + execution: Arc, + snapshot_control: Arc, + account_probe: Arc, + health_probes: Vec>, + worker_contributions: Vec, +} + +impl CoreBundle { + #[must_use] + pub fn execution_service(&self) -> Arc { + Arc::clone(&self.execution) + } + + #[must_use] + pub fn snapshot_control(&self) -> Arc { + Arc::clone(&self.snapshot_control) + } + + #[must_use] + pub fn account_probe(&self) -> Arc { + Arc::clone(&self.account_probe) + } + + #[must_use] + pub fn health_probes(&self) -> Vec> { + self.health_probes.clone() + } + + pub fn take_worker_contributions(&mut self) -> Vec { + std::mem::take(&mut self.worker_contributions) + } +} + +/// 首个快照与准入恢复均为监听前 fail-closed 屏障。 +pub async fn initialize( + ports: CoreStorePorts, + providers: ProviderRegistry, +) -> Result { + restore_client_admission_startup( + ports.execution.as_ref(), + ports.admission_recovery.as_ref(), + ports.admissions.as_ref(), + SystemTime::now(), + ) + .await + .map_err(|_| CoreError::AdmissionRecoveryUnavailable)?; + let compiler = Arc::new(RuntimeSnapshotCompiler::new( + Arc::clone(&ports.snapshots), + providers.clone(), + )); + let initial = compiler + .compile() + .await + .map_err(|_| CoreError::SnapshotUnavailable)?; + let snapshots = RuntimeSnapshotHandle::new(initial); + let publisher = Arc::new(RuntimeSnapshotPublisher::new( + compiler, + snapshots.clone(), + Arc::clone(&ports.snapshot_subscriptions), + )); + let worker_contributions = publisher.worker_contributions()?; + let service = Arc::new(DefaultExecutionService::new( + snapshots.clone(), + ports.execution, + providers, + ports.admissions, + ports.circuits, + ports.continuation, + )); + let execution: Arc = service.clone(); + let account_probe: Arc = service; + let snapshot_control: Arc = publisher; + let health_probes: Vec> = vec![Arc::new(snapshots)]; + Ok(CoreBundle { + execution, + snapshot_control, + account_probe, + health_probes, + worker_contributions, + }) +} + +#[derive(Debug, thiserror::Error)] +pub enum CoreError { + #[error("initial runtime snapshot is unavailable")] + SnapshotUnavailable, + #[error("client admission startup recovery is unavailable")] + AdmissionRecoveryUnavailable, + #[error(transparent)] + WorkerDefinition(#[from] WorkerDefinitionError), +} diff --git a/backend/crates/gateway-core/src/lifecycle.rs b/backend/crates/gateway-core/src/lifecycle.rs new file mode 100644 index 000000000..f8a7c3a4e --- /dev/null +++ b/backend/crates/gateway-core/src/lifecycle.rs @@ -0,0 +1,35 @@ +//! Host 与 API 之间的连接注册、drain 与取消契约。 + +use std::fmt; + +use crate::engine::CancellationToken; + +/// 进程已进入 drain，新连接不得再注册。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ConnectionDraining; + +impl fmt::Display for ConnectionDraining { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("connection lifecycle is draining") + } +} + +impl std::error::Error for ConnectionDraining {} + +/// 一次成功的活跃连接注册。 +/// +/// 实现必须在 guard `Drop` 时原子减少活跃连接计数。 +pub trait ConnectionGuard: Send + 'static {} + +/// API 消费、Host 实现的进程连接生命周期。 +pub trait ConnectionLifecycle: Send + Sync { + /// 原子地检查 drain 状态并注册一个活跃连接。 + /// + /// 当本方法成功时，drain 必须等待返回的 guard 被释放； + /// 当 drain 已经线性化生效时，本方法必须返回 [`ConnectionDraining`]。 + fn try_register(&self) -> Result, ConnectionDraining>; + + fn cancellation(&self) -> CancellationToken; + + fn is_draining(&self) -> bool; +} diff --git a/backend/crates/gateway-core/src/operation/mod.rs b/backend/crates/gateway-core/src/operation/mod.rs index 518d7fd31..bb1f3f8b5 100644 --- a/backend/crates/gateway-core/src/operation/mod.rs +++ b/backend/crates/gateway-core/src/operation/mod.rs @@ -5,6 +5,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt; +use std::sync::Arc; use serde_json::{Map, Value}; @@ -201,6 +202,104 @@ impl fmt::Debug for ProviderOptions { } } +/// 同一客户端连接内由 Provider 生成并解释的不透明会话状态。 +/// +/// Core 只在重试和路由过程中保持该值；协议层只能把 Provider 返回的状态原样带入 +/// 下一轮，不能读取或改写正文。 +#[derive(Clone, PartialEq)] +pub struct ProviderSessionState { + provider: String, + payload: Map, +} + +impl ProviderSessionState { + /// 创建 Provider 私有会话状态。 + /// + /// # Errors + /// + /// Provider 名称无效时返回错误。 + pub fn new( + provider: impl Into, + payload: Map, + ) -> Result { + let provider = provider.into(); + validate_text(&provider, 64, true, None).map_err(|_| OperationError::EmptyField { + field: "provider_session_state provider", + })?; + Ok(Self { provider, payload }) + } + + #[must_use] + pub fn provider(&self) -> &str { + &self.provider + } + + #[must_use] + pub const fn payload(&self) -> &Map { + &self.payload + } +} + +impl fmt::Debug for ProviderSessionState { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProviderSessionState") + .field("provider", &self.provider) + .field("payload", &"") + .finish() + } +} + +/// 客户端协议交给 Provider 的不透明 JSON object。 +/// +/// Core 只负责在路由和重试过程中保持该值，不读取或改写正文。协议 adapter +/// 与对应 Provider 共同拥有其语义。 +#[derive(Clone, PartialEq)] +pub struct ProtocolPayload { + protocol: String, + body: Map, +} + +impl ProtocolPayload { + /// 创建协议不透明正文。 + /// + /// # Errors + /// + /// 协议名称为空、过长或含控制字符时返回错误。 + pub fn json_object( + protocol: impl Into, + body: Map, + ) -> Result { + let protocol = protocol.into(); + validate_text(&protocol, 64, true, None).map_err(|_| OperationError::EmptyField { + field: "protocol_payload protocol", + })?; + Ok(Self { protocol, body }) + } + + /// 返回协议名称。 + #[must_use] + pub fn protocol(&self) -> &str { + &self.protocol + } + + /// 返回仅供对应 Provider 解释的 JSON object。 + #[must_use] + pub const fn body(&self) -> &Map { + &self.body + } +} + +impl fmt::Debug for ProtocolPayload { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProtocolPayload") + .field("protocol", &self.protocol) + .field("body", &"") + .finish() + } +} + /// 生成输入的角色。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum MessageRole { @@ -549,16 +648,25 @@ pub enum ContinuationMode { /// 通用生成请求。 #[derive(Clone, PartialEq)] pub struct GenerateRequest { - messages: Vec, - tools: Vec, - output_format: OutputFormat, + payload: Arc, reasoning: Option, continuation: Option, max_output_tokens: Option, estimated_context_tokens: u64, retry_safety: RetrySafety, response_persistence: ResponsePersistence, +} + +#[derive(Clone, PartialEq)] +struct GeneratePayload { + messages: Vec, + tools: Vec, + output_format: OutputFormat, + prompt_cache_key: Option, provider_options: ProviderOptions, + required_features: BTreeSet, + protocol_payload: Option, + provider_session_state: Option, } impl GenerateRequest { @@ -571,31 +679,53 @@ impl GenerateRequest { if messages.is_empty() { return Err(OperationError::EmptyField { field: "messages" }); } - Ok(Self { - messages, - tools: Vec::new(), - output_format: OutputFormat::Text, + Ok(Self::from_parts(messages, None)) + } + + /// 创建携带协议不透明正文的生成请求。 + /// + /// `messages` 是 Router 使用的已知语义投影，允许为空；完整请求始终由 + /// `protocol_payload` 保留并交给对应 Provider 解释。 + #[must_use] + pub fn from_protocol_payload( + messages: Vec, + protocol_payload: ProtocolPayload, + ) -> Self { + Self::from_parts(messages, Some(protocol_payload)) + } + + fn from_parts(messages: Vec, protocol_payload: Option) -> Self { + Self { + payload: Arc::new(GeneratePayload { + messages, + tools: Vec::new(), + output_format: OutputFormat::Text, + prompt_cache_key: None, + provider_options: ProviderOptions::new(), + required_features: BTreeSet::new(), + protocol_payload, + provider_session_state: None, + }), reasoning: None, continuation: None, max_output_tokens: None, estimated_context_tokens: 0, retry_safety: RetrySafety::NonIdempotent, response_persistence: ResponsePersistence::Store, - provider_options: ProviderOptions::new(), - }) + } } /// 设置工具。 #[must_use] pub fn with_tools(mut self, tools: Vec) -> Self { - self.tools = tools; + Arc::make_mut(&mut self.payload).tools = tools; self } /// 设置输出格式。 #[must_use] pub fn with_output_format(mut self, output_format: OutputFormat) -> Self { - self.output_format = output_format; + Arc::make_mut(&mut self.payload).output_format = output_format; self } @@ -620,6 +750,13 @@ impl GenerateRequest { self } + /// 设置客户端提供的稳定 prompt cache 路由键。 + #[must_use] + pub fn with_prompt_cache_key(mut self, key: impl Into) -> Self { + Arc::make_mut(&mut self.payload).prompt_cache_key = Some(key.into()); + self + } + /// 设置协议层估算的 context token 数。 #[must_use] pub const fn with_estimated_context_tokens(mut self, tokens: u64) -> Self { @@ -647,26 +784,42 @@ impl GenerateRequest { /// 设置 Provider 专属参数。 #[must_use] pub fn with_provider_options(mut self, provider_options: ProviderOptions) -> Self { - self.provider_options = provider_options; + Arc::make_mut(&mut self.payload).provider_options = provider_options; + self + } + + /// 附着同一客户端连接上一轮由 Provider 返回的不透明状态。 + #[must_use] + pub fn with_provider_session_state(mut self, state: ProviderSessionState) -> Self { + Arc::make_mut(&mut self.payload).provider_session_state = Some(state); + self + } + + /// 增加从不透明协议正文识别出的稳定能力要求。 + #[must_use] + pub fn require_feature(mut self, feature: Feature) -> Self { + Arc::make_mut(&mut self.payload) + .required_features + .insert(feature); self } /// 返回消息。 #[must_use] pub fn messages(&self) -> &[Message] { - &self.messages + &self.payload.messages } /// 返回工具。 #[must_use] pub fn tools(&self) -> &[ToolDefinition] { - &self.tools + &self.payload.tools } /// 返回输出格式。 #[must_use] - pub const fn output_format(&self) -> &OutputFormat { - &self.output_format + pub fn output_format(&self) -> &OutputFormat { + &self.payload.output_format } /// 返回推理要求。 @@ -687,6 +840,12 @@ impl GenerateRequest { self.max_output_tokens } + /// 返回客户端提供的 prompt cache 路由键。 + #[must_use] + pub fn prompt_cache_key(&self) -> Option<&str> { + self.payload.prompt_cache_key.as_deref() + } + /// 返回协议层估算的 context token 数。 #[must_use] pub const fn estimated_context_tokens(&self) -> u64 { @@ -707,18 +866,36 @@ impl GenerateRequest { /// 返回 Provider 专属参数。 #[must_use] - pub const fn provider_options(&self) -> &ProviderOptions { - &self.provider_options + pub fn provider_options(&self) -> &ProviderOptions { + &self.payload.provider_options + } + + /// 返回指定 Provider 的连接内会话状态。 + #[must_use] + pub fn provider_session_state(&self, provider: &str) -> Option<&ProviderSessionState> { + self.payload + .provider_session_state + .as_ref() + .filter(|state| state.provider() == provider) + } + + /// 返回协议不透明正文。 + #[must_use] + pub fn protocol_payload(&self) -> Option<&ProtocolPayload> { + self.payload.protocol_payload.as_ref() } fn requirements(&self) -> CapabilityRequirements { let mut requirements = CapabilityRequirements::new(OperationKind::Generate) .with_minimum_context_tokens(self.estimated_context_tokens) .with_requested_output_tokens(self.max_output_tokens); - if !self.tools.is_empty() { + for feature in &self.payload.required_features { + requirements = requirements.require(*feature); + } + if !self.payload.tools.is_empty() { requirements = requirements.require(Feature::Tools); } - if self.messages.iter().any(|message| { + if self.payload.messages.iter().any(|message| { message .content() .iter() @@ -729,7 +906,7 @@ impl GenerateRequest { if self.reasoning.is_some() { requirements = requirements.require(Feature::Reasoning); } - if matches!(self.output_format, OutputFormat::JsonSchema(_)) { + if matches!(self.payload.output_format, OutputFormat::JsonSchema(_)) { requirements = requirements.require(Feature::JsonSchema); } if self.continuation.is_some() { @@ -744,16 +921,30 @@ impl fmt::Debug for GenerateRequest { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("GenerateRequest") - .field("messages", &self.messages.len()) - .field("tools", &self.tools.len()) - .field("output_format", &self.output_format) + .field("messages", &self.payload.messages.len()) + .field("tools", &self.payload.tools.len()) + .field("output_format", &self.payload.output_format) .field("reasoning", &self.reasoning) .field("continuation", &self.continuation) .field("max_output_tokens", &self.max_output_tokens) + .field( + "prompt_cache_key", + &self.payload.prompt_cache_key.as_ref().map(|_| ""), + ) .field("estimated_context_tokens", &self.estimated_context_tokens) .field("retry_safety", &self.retry_safety) .field("response_persistence", &self.response_persistence) - .field("provider_options", &self.provider_options) + .field("provider_options", &self.payload.provider_options) + .field( + "provider_session_state", + &self + .payload + .provider_session_state + .as_ref() + .map(|_| ""), + ) + .field("required_features", &self.payload.required_features) + .field("protocol_payload", &self.payload.protocol_payload) .finish() } } @@ -957,6 +1148,15 @@ pub enum Operation { } impl Operation { + /// 附着协议连接持有的 Provider 私有状态；非生成 operation 保持不变。 + #[must_use] + pub fn with_provider_session_state(self, state: ProviderSessionState) -> Self { + match self { + Self::Generate(request) => Self::Generate(request.with_provider_session_state(state)), + operation => operation, + } + } + /// 返回稳定 operation 分类。 #[must_use] pub const fn kind(&self) -> OperationKind { @@ -1000,6 +1200,15 @@ impl Operation { Self::Embed(_) | Self::Rerank(_) | Self::GenerateImage(_) | Self::Speech(_) => None, } } + + /// 返回当前 Provider 的连接内私有状态。 + #[must_use] + pub fn provider_session_state(&self, provider: &str) -> Option<&ProviderSessionState> { + match self { + Self::Generate(request) => request.provider_session_state(provider), + Self::Embed(_) | Self::Rerank(_) | Self::GenerateImage(_) | Self::Speech(_) => None, + } + } } impl fmt::Debug for Operation { diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs new file mode 100644 index 000000000..ce2abd90b --- /dev/null +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -0,0 +1,822 @@ +//! Provider 运行时所需的中立存储能力。 + +use std::collections::BTreeMap; +use std::fmt; +use std::num::NonZeroU32; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; + +use futures::future::BoxFuture; + +use crate::engine::credential::{ + AccountAvailability, AccountRuntimeSignals, CredentialRevision, OpaqueProviderData, + ProviderAccountId, ProviderAccountStore, +}; +use crate::routing::{ProviderInstanceId, ProviderKind}; + +const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); + +/// Provider 可据此决定是否重试，但看不到 SQL、Redis 或秘密原文。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderStoreErrorKind { + Unavailable, + InvalidData, + Conflict, +} + +/// Provider 存储端口的脱敏错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("provider store {operation} failed: {kind:?}")] +pub struct ProviderStoreError { + kind: ProviderStoreErrorKind, + operation: &'static str, +} + +impl ProviderStoreError { + #[must_use] + pub const fn new(kind: ProviderStoreErrorKind, operation: &'static str) -> Self { + Self { kind, operation } + } + + #[must_use] + pub const fn kind(&self) -> ProviderStoreErrorKind { + self.kind + } +} + +/// 一次 Provider instance 的完整可重建调度状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSchedulingState { + signals: BTreeMap, + sticky_account: Option, + round_robin_cursor: u64, +} + +impl ProviderSchedulingState { + #[must_use] + pub const fn new( + signals: BTreeMap, + sticky_account: Option, + round_robin_cursor: u64, + ) -> Self { + Self { + signals, + sticky_account, + round_robin_cursor, + } + } + + #[must_use] + pub const fn signals(&self) -> &BTreeMap { + &self.signals + } + + #[must_use] + pub const fn sticky_account(&self) -> Option<&ProviderAccountId> { + self.sticky_account.as_ref() + } + + #[must_use] + pub const fn round_robin_cursor(&self) -> u64 { + self.round_robin_cursor + } +} + +/// 请求级账号 lease 的全部中立事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderSchedulingLeaseRequest { + provider_instance_id: ProviderInstanceId, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: NonZeroU32, + request_interval: Duration, + deadline: SystemTime, +} + +impl ProviderSchedulingLeaseRequest { + #[must_use] + pub const fn new( + provider_instance_id: ProviderInstanceId, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + max_concurrent: NonZeroU32, + request_interval: Duration, + deadline: SystemTime, + ) -> Self { + Self { + provider_instance_id, + account_id, + credential_revision, + max_concurrent, + request_interval, + deadline, + } + } + + #[must_use] + pub const fn provider_instance_id(&self) -> &ProviderInstanceId { + &self.provider_instance_id + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn max_concurrent(&self) -> NonZeroU32 { + self.max_concurrent + } + + #[must_use] + pub const fn request_interval(&self) -> Duration { + self.request_interval + } + + #[must_use] + pub const fn deadline(&self) -> SystemTime { + self.deadline + } +} + +/// Lease 生命周期由具体 Store guard 管理，Provider 只能持有。 +pub trait ProviderLeaseGuard: Send + Sync + 'static {} + +impl ProviderLeaseGuard for T where T: Send + Sync + 'static {} + +pub enum ProviderLeaseAcquisition { + Acquired(Box), + Busy { retry_after: Option }, +} + +impl fmt::Debug for ProviderLeaseAcquisition { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), + Self::Busy { retry_after } => formatter + .debug_struct("Busy") + .field("retry_after", retry_after) + .finish(), + } + } +} + +/// Provider 运行时会持有的三类 lease；刷新必须同时持有全局容量与账号互斥 lease。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProviderLeaseRequest { + Scheduling(ProviderSchedulingLeaseRequest), + RefreshCapacity(ProviderRefreshCapacityRequest), + Refresh(ProviderRefreshLeaseRequest), +} + +pub trait ProviderLeasePort: Send + Sync { + fn load_state<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + accounts: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result>; + + fn try_acquire( + &self, + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result>; +} + +/// 所有 Provider、所有实例共享的 OAuth refresh 并发容量。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshCapacityRequest { + max_concurrent: NonZeroU32, +} + +impl ProviderRefreshCapacityRequest { + #[must_use] + pub const fn new(max_concurrent: NonZeroU32) -> Self { + Self { max_concurrent } + } + + #[must_use] + pub const fn max_concurrent(self) -> NonZeroU32 { + self.max_concurrent + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderRefreshLeaseRequest { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, +} + +impl ProviderRefreshLeaseRequest { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + ) -> Self { + Self { + account_id, + credential_revision, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } +} + +/// Opaque catalog cache 的平台、账号与 credential revision 作用域。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCatalogCacheKey { + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, +} + +impl ProviderCatalogCacheKey { + #[must_use] + pub const fn new( + provider_kind: ProviderKind, + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + ) -> Self { + Self { + provider_kind, + account_id, + credential_revision, + } + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } +} + +pub trait ProviderCatalogCachePort: Send + Sync { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} + +/// Provider worker 读取的中立 instance 事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInstanceConfig { + id: ProviderInstanceId, + provider_kind: ProviderKind, + base_url: String, + enabled: bool, +} + +impl ProviderInstanceConfig { + #[must_use] + pub const fn new( + id: ProviderInstanceId, + provider_kind: ProviderKind, + base_url: String, + enabled: bool, + ) -> Self { + Self { + id, + provider_kind, + base_url, + enabled, + } + } + + #[must_use] + pub const fn id(&self) -> &ProviderInstanceId { + &self.id + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub fn base_url(&self) -> &str { + &self.base_url + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } +} + +pub trait ProviderInstanceCatalogPort: Send + Sync { + fn list_instances<'a>( + &'a self, + provider_kind: &'a ProviderKind, + include_disabled: bool, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; +} + +/// Provider 的两种目录能力：instance 事实与 provider-owned model cache。 +#[derive(Clone)] +pub struct ProviderCatalogPorts { + instances: Arc, + model_cache: Arc, +} + +impl ProviderCatalogPorts { + #[must_use] + pub fn new( + instances: Arc, + model_cache: Arc, + ) -> Self { + Self { + instances, + model_cache, + } + } + + #[must_use] + pub fn instances(&self) -> Arc { + Arc::clone(&self.instances) + } + + #[must_use] + pub fn model_cache(&self) -> Arc { + Arc::clone(&self.model_cache) + } +} + +impl fmt::Debug for ProviderCatalogPorts { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderCatalogPorts([CAPABILITIES])") + } +} + +/// Redis 中可重建的账号状态投影。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCredentialState { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + availability: AccountAvailability, + observed_at: SystemTime, +} + +impl ProviderCredentialState { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + enabled: bool, + availability: AccountAvailability, + observed_at: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + enabled, + availability, + observed_at, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn enabled(&self) -> bool { + self.enabled + } + + #[must_use] + pub const fn availability(&self) -> AccountAvailability { + self.availability + } + + #[must_use] + pub const fn observed_at(&self) -> SystemTime { + self.observed_at + } +} + +pub trait ProviderCredentialStatePort: Send + Sync { + fn replace( + &self, + state: ProviderCredentialState, + ) -> BoxFuture<'_, Result<(), ProviderStoreError>>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result>; +} + +/// 临时 cooldown 只保存调度截止时间；原因事实仍由 PostgreSQL 账号状态持有。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderCooldown { + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, +} + +impl ProviderCooldown { + #[must_use] + pub const fn new( + account_id: ProviderAccountId, + credential_revision: CredentialRevision, + until: SystemTime, + ) -> Self { + Self { + account_id, + credential_revision, + until, + } + } + + #[must_use] + pub const fn account_id(&self) -> &ProviderAccountId { + &self.account_id + } + + #[must_use] + pub const fn credential_revision(&self) -> CredentialRevision { + self.credential_revision + } + + #[must_use] + pub const fn until(&self) -> SystemTime { + self.until + } +} + +pub trait ProviderCooldownPort: Send + Sync { + fn put_if_later( + &self, + cooldown: ProviderCooldown, + ) -> BoxFuture<'_, Result>; + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>>; + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result>; +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ProviderRefreshPolicy { + margin: Duration, + concurrency: NonZeroU32, +} + +impl ProviderRefreshPolicy { + pub fn try_new(margin: Duration, concurrency: NonZeroU32) -> Result { + if margin.is_zero() { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate refresh policy", + )); + } + Ok(Self { + margin, + concurrency, + }) + } + + #[must_use] + pub const fn margin(self) -> Duration { + self.margin + } + + #[must_use] + pub const fn concurrency(self) -> NonZeroU32 { + self.concurrency + } + + /// 按账号稳定扰动刷新提前量；相同账号与过期时间跨进程得到同一执行时刻。 + pub fn next_attempt_at( + self, + account_id: &ProviderAccountId, + access_token_expires_at: SystemTime, + observed_at: SystemTime, + ) -> Result { + let remaining = access_token_expires_at + .duration_since(observed_at) + .map_err(|_| invalid_refresh_policy("schedule expired access token"))?; + let base_margin_seconds = self.margin.as_secs(); + if base_margin_seconds == 0 { + return Ok(observed_at); + } + + let factor = stable_factor(account_id.as_str(), "normal", 850, 1_150); + let jittered_seconds = base_margin_seconds + .saturating_mul(u64::from(factor)) + .saturating_add(500) + / 1_000; + if jittered_seconds >= remaining.as_secs() { + return Ok(observed_at); + } + let lead = jittered_seconds.max(1); + access_token_expires_at + .checked_sub(Duration::from_secs(lead)) + .ok_or_else(|| invalid_refresh_policy("schedule access token refresh")) + } +} + +/// 临时失败后的持久重试时刻；稳定扰动避免多实例同频重试。 +pub fn provider_refresh_retry_at( + account_id: &ProviderAccountId, + observed_at: SystemTime, + base_delay: Duration, + reason: &'static str, +) -> Result { + if base_delay.is_zero() { + return Err(invalid_refresh_policy("schedule refresh retry")); + } + let factor = stable_factor(account_id.as_str(), reason, 800, 1_200); + let millis = u64::try_from(base_delay.as_millis()) + .unwrap_or(u64::MAX) + .saturating_mul(u64::from(factor)) + .saturating_add(500) + / 1_000; + observed_at + .checked_add(Duration::from_millis(millis.max(1))) + .ok_or_else(|| invalid_refresh_policy("schedule refresh retry")) +} + +fn stable_factor(value: &str, salt: &str, minimum: u32, maximum: u32) -> u32 { + let mut hash = 0xcbf2_9ce4_8422_2325_u64; + for byte in value.bytes().chain([0]).chain(salt.bytes()) { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + let width = u64::from(maximum - minimum) + 1; + minimum + u32::try_from(hash % width).unwrap_or_default() +} + +fn invalid_refresh_policy(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} + +pub trait ProviderRuntimePolicyPort: Send + Sync { + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result>; +} + +/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在。 +#[derive(Clone, PartialEq, Eq)] +pub struct OAuthPendingBinding(String); + +impl OAuthPendingBinding { + pub fn try_new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() || value.len() > 512 { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending binding", + )); + } + Ok(Self(value)) + } + + #[must_use] + pub fn expose_to_store(&self) -> &str { + &self.0 + } +} + +impl fmt::Debug for OAuthPendingBinding { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("OAuthPendingBinding([REDACTED])") + } +} + +#[derive(Clone, PartialEq)] +pub struct NewOAuthPendingFlow { + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, +} + +impl NewOAuthPendingFlow { + pub fn try_new( + provider_kind: ProviderKind, + flow: OAuthPendingBinding, + owner: OAuthPendingBinding, + ttl: Duration, + payload: OpaqueProviderData, + ) -> Result { + if ttl.is_zero() || ttl > MAX_PENDING_FLOW_TTL { + return Err(ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "validate OAuth pending TTL", + )); + } + Ok(Self { + provider_kind, + flow, + owner, + ttl, + payload, + }) + } + + #[must_use] + pub const fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + #[must_use] + pub const fn flow(&self) -> &OAuthPendingBinding { + &self.flow + } + + #[must_use] + pub const fn owner(&self) -> &OAuthPendingBinding { + &self.owner + } + + #[must_use] + pub const fn ttl(&self) -> Duration { + self.ttl + } + + #[must_use] + pub const fn payload(&self) -> &OpaqueProviderData { + &self.payload + } +} + +impl fmt::Debug for NewOAuthPendingFlow { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("NewOAuthPendingFlow") + .field("provider_kind", &self.provider_kind) + .field("flow", &self.flow) + .field("owner", &self.owner) + .field("ttl", &self.ttl) + .field("payload", &"[PROVIDER-OWNED]") + .finish() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthPendingPutOutcome { + Stored, + AlreadyExists, +} + +#[derive(Clone, PartialEq)] +pub enum OAuthPendingTakeOutcome { + Taken(OpaqueProviderData), + NotFound, + OwnerMismatch, +} + +impl fmt::Debug for OAuthPendingTakeOutcome { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Taken(_) => formatter.write_str("Taken([PROVIDER-OWNED])"), + Self::NotFound => formatter.write_str("NotFound"), + Self::OwnerMismatch => formatter.write_str("OwnerMismatch"), + } + } +} + +pub trait OAuthPendingFlowPort: Send + Sync { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> BoxFuture<'_, Result>; + + fn take_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result>; +} + +/// Provider 只能按能力取用端口，无法取得 Redis client 或 repository 集合。 +#[derive(Clone)] +pub struct ProviderStorePorts { + accounts: Arc, + leases: Arc, + catalog: ProviderCatalogPorts, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, +} + +impl ProviderStorePorts { + #[must_use] + pub fn new( + accounts: Arc, + leases: Arc, + catalog: ProviderCatalogPorts, + credential_state: Arc, + cooldowns: Arc, + runtime_policy: Arc, + oauth_pending: Arc, + ) -> Self { + Self { + accounts, + leases, + catalog, + credential_state, + cooldowns, + runtime_policy, + oauth_pending, + } + } + + #[must_use] + pub fn accounts(&self) -> Arc { + Arc::clone(&self.accounts) + } + + #[must_use] + pub fn leases(&self) -> Arc { + Arc::clone(&self.leases) + } + + #[must_use] + pub fn instances(&self) -> Arc { + self.catalog.instances() + } + + #[must_use] + pub fn catalog_cache(&self) -> Arc { + self.catalog.model_cache() + } + + #[must_use] + pub fn credential_state(&self) -> Arc { + Arc::clone(&self.credential_state) + } + + #[must_use] + pub fn cooldowns(&self) -> Arc { + Arc::clone(&self.cooldowns) + } + + #[must_use] + pub fn runtime_policy(&self) -> Arc { + Arc::clone(&self.runtime_policy) + } + + #[must_use] + pub fn oauth_pending(&self) -> Arc { + Arc::clone(&self.oauth_pending) + } +} + +impl fmt::Debug for ProviderStorePorts { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("ProviderStorePorts([CAPABILITIES])") + } +} diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index d736a3f60..87859d6c9 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -1,5 +1,9 @@ //! Provider 平台、实例、模型目录、精确模型映射与请求级候选计划。 +pub mod snapshot; + +pub use snapshot::RuntimeSnapshot; + use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::num::{NonZeroU32, NonZeroU64}; @@ -7,8 +11,7 @@ use std::sync::Arc; use crate::engine::credential::AccountSelectionPolicy; use crate::error::{IdentifierError, RoutingError, validate_text}; -use crate::operation::{CapabilityRequirements, Feature, Operation, OperationKind}; -use crate::policy::{ClientApiKeyId, ClientPolicy}; +use crate::operation::{CapabilityRequirements, Feature, OperationKind}; const MAX_REQUEST_ATTEMPTS: u32 = 32; @@ -410,261 +413,3 @@ impl RoutingPlan { &self.candidates } } - -/// 数据面使用的不可变配置快照。 -#[derive(Debug, Clone)] -pub struct RuntimeSnapshot { - revision: ConfigRevision, - account_selection_policy: AccountSelectionPolicy, - instances: Arc>, - provider_models: - Arc>>, - provider_model_mappings: Arc>>, - client_policies: Arc>, -} - -impl RuntimeSnapshot { - /// 校验 instance、实时模型目录和 Client API Key，并构建快照。 - pub fn new( - revision: ConfigRevision, - account_selection_policy: AccountSelectionPolicy, - instances: Vec, - provider_models: Vec, - client_policies: Vec, - ) -> Result { - let mut instance_map = BTreeMap::new(); - for instance in instances { - let id = instance.id().clone(); - if instance_map.insert(id.clone(), instance).is_some() { - return Err(RoutingError::DuplicateEntity { - entity: "provider instance", - id: id.to_string(), - }); - } - } - - let mut model_map = - BTreeMap::>::new(); - for model in provider_models { - if !instance_map.contains_key(model.instance()) { - return Err(RoutingError::NotFound { - entity: "provider instance", - id: model.instance().to_string(), - }); - } - let models = model_map.entry(model.instance).or_default(); - let upstream_model = model.upstream_model; - if models - .insert(upstream_model.clone(), model.capabilities) - .is_some() - { - return Err(RoutingError::DuplicateEntity { - entity: "provider model", - id: upstream_model.to_string(), - }); - } - } - - let mut client_policy_map = BTreeMap::new(); - for policy in client_policies { - let id = policy.key_id().clone(); - if client_policy_map.insert(id.clone(), policy).is_some() { - return Err(RoutingError::DuplicateEntity { - entity: "client API key", - id: id.to_string(), - }); - } - } - client_policy_map.retain(|_, policy| policy.enabled()); - - Ok(Self { - revision, - account_selection_policy, - instances: Arc::new(instance_map), - provider_models: Arc::new(model_map), - provider_model_mappings: Arc::new(BTreeMap::new()), - client_policies: Arc::new(client_policy_map), - }) - } - - #[must_use] - pub fn with_provider_model_mappings( - mut self, - mappings: BTreeMap>, - ) -> Self { - self.provider_model_mappings = Arc::new(mappings); - self - } - - #[must_use] - pub const fn revision(&self) -> ConfigRevision { - self.revision - } - - /// 返回目录发现模型与设置别名的并集，仅用于 `/v1/models` 展示。 - #[must_use] - pub fn public_models_for_provider(&self, provider: &ProviderKind) -> Vec { - let mut models = BTreeSet::new(); - for (instance_id, discovered) in self.provider_models.iter() { - let Some(instance) = self.instances.get(instance_id) else { - continue; - }; - if !instance.enabled || instance.provider() != provider { - continue; - } - models.extend( - discovered - .keys() - .filter_map(|model| PublicModelId::new(model.as_str().to_owned()).ok()), - ); - } - if let Some(mappings) = self.provider_model_mappings.get(provider) { - models.extend( - mappings - .keys() - .filter_map(|model| PublicModelId::new(model.clone()).ok()), - ); - } - models.into_iter().collect() - } - - #[must_use] - pub fn public_models(&self) -> Vec { - let providers = self - .instances - .values() - .map(|instance| instance.provider().clone()) - .collect::>(); - providers - .iter() - .flat_map(|provider| self.public_models_for_provider(provider)) - .collect::>() - .into_iter() - .collect() - } - - /// 透明代理不以目录白名单拒绝模型；只要求该平台至少有一个启用 instance。 - #[must_use] - pub fn contains_public_model_for_provider( - &self, - _public_model: &PublicModelId, - provider: &ProviderKind, - ) -> bool { - self.instances - .values() - .any(|instance| instance.enabled && instance.provider() == provider) - } - - #[must_use] - pub fn instance_ids_for_provider( - &self, - provider: &ProviderKind, - ) -> BTreeSet { - self.instances - .iter() - .filter_map(|(id, instance)| { - (instance.enabled && instance.provider() == provider).then_some(id.clone()) - }) - .collect() - } - - #[must_use] - pub fn provider_for_instance(&self, instance_id: &ProviderInstanceId) -> Option<&ProviderKind> { - self.instances - .get(instance_id) - .map(ProviderInstance::provider) - } - - #[must_use] - pub fn mapped_model(&self, provider: &ProviderKind, requested: &str) -> String { - self.provider_model_mappings - .get(provider) - .and_then(|mapping| mapping.get(requested)) - .cloned() - .unwrap_or_else(|| requested.to_owned()) - } - - pub fn client_policies(&self) -> impl Iterator { - self.client_policies.values() - } - - #[must_use] - pub fn client_policy(&self, key_id: &ClientApiKeyId) -> Option<&ClientPolicy> { - self.client_policies.get(key_id) - } - - pub fn plan( - &self, - public_model: &PublicModelId, - operation: &Operation, - context: &RoutingContext, - ) -> Result { - let upstream_model = context.provider_kind.as_ref().map_or_else( - || public_model.as_str().to_owned(), - |provider| self.mapped_model(provider, public_model.as_str()), - ); - let upstream_model = - UpstreamModelId::new(upstream_model).map_err(|_| RoutingError::InvalidIdentifier)?; - let requirements = operation.capability_requirements(); - let mut candidates = Vec::new(); - - for instance in self.instances.values() { - if !self.instance_is_eligible(instance, context) { - continue; - } - let emulated_features = match self - .provider_models - .get(instance.id()) - .and_then(|models| models.get(&upstream_model)) - { - Some(capabilities) => { - let Some(emulated) = capabilities.match_requirements(&requirements) else { - continue; - }; - emulated - } - // 目录用于展示与能力提示，不是模型白名单；未知模型交给上游判断。 - None => BTreeSet::new(), - }; - candidates.push(ProviderCandidate { - instance: instance.clone(), - upstream_model: upstream_model.clone(), - emulated_features, - observation_token: context - .provider_observation_tokens - .get(instance.id()) - .copied(), - }); - } - - if candidates.is_empty() { - return Err(RoutingError::NoCapableProvider { - model: public_model.as_str().to_owned(), - }); - } - - Ok(RoutingPlan { - config_revision: self.revision, - account_selection_policy: self.account_selection_policy, - public_model: public_model.clone(), - operation: operation.kind(), - max_attempts: NonZeroU32::new(MAX_REQUEST_ATTEMPTS) - .expect("constant request attempt limit is non-zero"), - candidates: Arc::from(candidates), - }) - } - - fn instance_is_eligible(&self, instance: &ProviderInstance, context: &RoutingContext) -> bool { - instance.enabled - && instance.health.is_routable(context.allow_degraded) - && context - .provider_kind - .as_ref() - .is_none_or(|provider| instance.provider() == provider) - && !context.blocked_instances.contains(instance.id()) - && context - .allowed_instances - .as_ref() - .is_none_or(|allowed| allowed.contains(instance.id())) - } -} diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs new file mode 100644 index 000000000..5e81561c1 --- /dev/null +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -0,0 +1,908 @@ +//! RuntimeSnapshot 事实、编译、原子发布与版本收敛规则。 + +use std::collections::{BTreeMap, BTreeSet}; +use std::num::NonZeroU32; +use std::pin::Pin; +use std::sync::{Arc, RwLock}; +use std::time::Duration; + +use futures::future::BoxFuture; +use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased}; +use futures_timer::Delay; + +use crate::engine::CancellationToken; +use crate::engine::credential::{AccountSelectionPolicy, RotationStrategy}; +use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry}; +use crate::error::RoutingError; +use crate::health::{HealthProbe, HealthState}; +use crate::operation::Operation; +use crate::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits}; +use crate::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, + WorkerSchedule, WorkerTaskError, +}; + +use super::{ + ConfigRevision, InstanceHealth, ModelCapabilities, ProviderCandidate, ProviderInstance, + ProviderInstanceId, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan, + UpstreamModelId, +}; + +const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5); +const INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const MAXIMUM_BACKOFF: Duration = Duration::from_secs(30); +const UNUSED_LEASE_TTL: Duration = Duration::from_secs(30); +const UNUSED_LEASE_RENEWAL: Duration = Duration::from_secs(10); +const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4; + +/// Store 在一个一致性读取中提供的调度设置事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnapshotSettingsFacts { + max_concurrent_per_account: u32, + request_interval_ms: u64, + rotation_strategy: String, + provider_model_mappings: BTreeMap>, +} + +impl SnapshotSettingsFacts { + #[must_use] + pub fn new( + max_concurrent_per_account: u32, + request_interval_ms: u64, + rotation_strategy: impl Into, + provider_model_mappings: BTreeMap>, + ) -> Self { + Self { + max_concurrent_per_account, + request_interval_ms, + rotation_strategy: rotation_strategy.into(), + provider_model_mappings, + } + } +} + +/// Store 读取到的一个 Provider instance 配置事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnapshotProviderInstanceFacts { + id: String, + provider_kind: String, + base_url: String, + enabled: bool, +} + +impl SnapshotProviderInstanceFacts { + #[must_use] + pub fn new( + id: impl Into, + provider_kind: impl Into, + base_url: impl Into, + enabled: bool, + ) -> Self { + Self { + id: id.into(), + provider_kind: provider_kind.into(), + base_url: base_url.into(), + enabled, + } + } +} + +/// Store 读取到的一个启用 Client API Key 策略事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnapshotClientPolicyFacts { + key_id: ClientApiKeyId, + plaintext_key: PlaintextClientApiKey, + provider_kind: String, + limits: RateLimits, +} + +impl SnapshotClientPolicyFacts { + #[must_use] + pub fn new( + key_id: ClientApiKeyId, + plaintext_key: PlaintextClientApiKey, + provider_kind: impl Into, + limits: RateLimits, + ) -> Self { + Self { + key_id, + plaintext_key, + provider_kind: provider_kind.into(), + limits, + } + } +} + +/// 一次一致性读取产生的全部 RuntimeSnapshot 持久事实。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnapshotFacts { + config_revision: ConfigRevision, + observed_current_revision: ConfigRevision, + settings: SnapshotSettingsFacts, + provider_instances: Vec, + client_policies: Vec, +} + +impl SnapshotFacts { + #[must_use] + pub fn new( + config_revision: ConfigRevision, + observed_current_revision: ConfigRevision, + settings: SnapshotSettingsFacts, + provider_instances: Vec, + client_policies: Vec, + ) -> Self { + Self { + config_revision, + observed_current_revision, + settings, + provider_instances, + client_policies, + } + } + + #[must_use] + pub const fn config_revision(&self) -> ConfigRevision { + self.config_revision + } + + #[must_use] + pub const fn observed_current_revision(&self) -> ConfigRevision { + self.observed_current_revision + } +} + +/// 不泄漏持久化实现细节的 Snapshot store 错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot store is unavailable")] +pub struct SnapshotStoreError; + +impl SnapshotStoreError { + #[must_use] + pub const fn unavailable() -> Self { + Self + } +} + +/// 不泄漏订阅基础设施细节的通知错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot notification is unavailable")] +pub struct SnapshotSubscriptionError; + +impl SnapshotSubscriptionError { + #[must_use] + pub const fn unavailable() -> Self { + Self + } +} + +/// 可丢失的配置 revision 通知流；权威 revision 始终由 Store 端口读取。 +pub type SnapshotRevisionStream = + Pin> + Send + 'static>>; + +/// RuntimeSnapshot 持久事实的数据库中立端口。 +pub trait SnapshotStorePort: Send + Sync { + fn load_snapshot_facts(&self) -> BoxFuture<'_, Result>; + + fn current_config_revision(&self) -> BoxFuture<'_, Result>; +} + +/// 跨进程 revision 通知的基础设施中立端口。 +pub trait SnapshotSubscriptionPort: Send + Sync { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; + + fn subscribe_snapshot_revisions( + &self, + ) -> BoxFuture<'_, Result>; +} + +/// 快照未发布时可安全记录的稳定错误。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum RuntimeSnapshotCompileError { + #[error("runtime snapshot store is unavailable")] + StoreUnavailable, + #[error("runtime configuration changed while the snapshot was loading")] + RevisionChanged, + #[error("runtime snapshot contains invalid frozen data")] + InvalidData, + #[error("provider model catalog changed while the snapshot was compiling")] + CatalogChanged, +} + +/// Store 一致性事实与 Provider 实时目录的唯一快照编译器。 +#[derive(Clone)] +pub struct RuntimeSnapshotCompiler { + store: Arc, + providers: ProviderRegistry, +} + +impl RuntimeSnapshotCompiler { + #[must_use] + pub const fn new(store: Arc, providers: ProviderRegistry) -> Self { + Self { store, providers } + } + + /// 读取一个 revision，并为 instance 查询实时模型目录。 + pub async fn compile(&self) -> Result { + for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS { + let catalog_generations = self.providers.catalog_generations(); + let facts = self + .store + .load_snapshot_facts() + .await + .map_err(|_| RuntimeSnapshotCompileError::StoreUnavailable)?; + if facts.config_revision != facts.observed_current_revision { + return Err(RuntimeSnapshotCompileError::RevisionChanged); + } + let snapshot = compile_runtime_snapshot(facts, &self.providers).await?; + let observed_generations = self.providers.catalog_generations(); + if catalog_generations == observed_generations { + return Ok(snapshot.with_provider_catalog_generations(observed_generations)); + } + } + Err(RuntimeSnapshotCompileError::CatalogChanged) + } +} + +async fn compile_runtime_snapshot( + facts: SnapshotFacts, + providers: &ProviderRegistry, +) -> Result { + let mut instances_by_id = BTreeMap::new(); + for record in facts.provider_instances { + let id = ProviderInstanceId::new(record.id) + .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?; + let instance = ProviderInstance::new( + id.clone(), + ProviderKind::new(record.provider_kind) + .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?, + record.base_url, + record.enabled, + InstanceHealth::Healthy, + ); + if instances_by_id.insert(id, instance).is_some() { + return Err(RuntimeSnapshotCompileError::InvalidData); + } + } + + // 实时目录只提供公开模型与能力提示；查询失败时保留 instance 透传语义。 + let mut provider_models = Vec::new(); + for instance in instances_by_id.values() { + let Ok(models) = providers.query_model_capabilities(instance).await else { + continue; + }; + provider_models.extend(models.into_iter().map(|model| { + ProviderModel::new( + instance.id().clone(), + model.upstream_model().clone(), + model.capabilities().clone(), + ) + })); + } + + let provider_model_mappings = facts + .settings + .provider_model_mappings + .into_iter() + .map(|(provider, mappings)| { + ProviderKind::new(provider) + .map(|provider| (provider, mappings)) + .map_err(|_| RuntimeSnapshotCompileError::InvalidData) + }) + .collect::, _>>()?; + let rotation_strategy = match facts.settings.rotation_strategy.as_str() { + "smart" => RotationStrategy::Smart, + "quota_reset_priority" => RotationStrategy::QuotaResetPriority, + "round_robin" => RotationStrategy::RoundRobin, + "sticky" => RotationStrategy::Sticky, + _ => return Err(RuntimeSnapshotCompileError::InvalidData), + }; + let selection_policy = AccountSelectionPolicy::new( + rotation_strategy, + NonZeroU32::new(facts.settings.max_concurrent_per_account) + .ok_or(RuntimeSnapshotCompileError::InvalidData)?, + Duration::from_millis(facts.settings.request_interval_ms), + ); + let client_policies = facts + .client_policies + .into_iter() + .map(|policy| { + let provider = ProviderKind::new(policy.provider_kind) + .map_err(|_| RuntimeSnapshotCompileError::InvalidData)?; + Ok(ClientPolicy::new( + policy.key_id, + policy.plaintext_key, + provider, + true, + policy.limits, + )) + }) + .collect::, RuntimeSnapshotCompileError>>()?; + + RuntimeSnapshot::new( + facts.config_revision, + selection_policy, + instances_by_id.into_values().collect(), + provider_models, + client_policies, + ) + .map_err(|_| RuntimeSnapshotCompileError::InvalidData) + .map(|snapshot| snapshot.with_provider_model_mappings(provider_model_mappings)) +} + +/// 数据面使用的不可变配置快照。 +#[derive(Debug, Clone)] +pub struct RuntimeSnapshot { + revision: ConfigRevision, + account_selection_policy: AccountSelectionPolicy, + instances: Arc>, + provider_models: + Arc>>, + provider_model_mappings: Arc>>, + provider_catalog_generations: Arc>, + client_policies: Arc>, +} + +impl RuntimeSnapshot { + /// 校验 instance、实时模型目录和 Client API Key，并构建快照。 + pub fn new( + revision: ConfigRevision, + account_selection_policy: AccountSelectionPolicy, + instances: Vec, + provider_models: Vec, + client_policies: Vec, + ) -> Result { + let mut instance_map = BTreeMap::new(); + for instance in instances { + let id = instance.id().clone(); + if instance_map.insert(id.clone(), instance).is_some() { + return Err(RoutingError::DuplicateEntity { + entity: "provider instance", + id: id.to_string(), + }); + } + } + + let mut model_map = + BTreeMap::>::new(); + for model in provider_models { + if !instance_map.contains_key(model.instance()) { + return Err(RoutingError::NotFound { + entity: "provider instance", + id: model.instance().to_string(), + }); + } + let models = model_map.entry(model.instance).or_default(); + let upstream_model = model.upstream_model; + if models + .insert(upstream_model.clone(), model.capabilities) + .is_some() + { + return Err(RoutingError::DuplicateEntity { + entity: "provider model", + id: upstream_model.to_string(), + }); + } + } + + let mut client_policy_map = BTreeMap::new(); + for policy in client_policies { + let id = policy.key_id().clone(); + if client_policy_map.insert(id.clone(), policy).is_some() { + return Err(RoutingError::DuplicateEntity { + entity: "client API key", + id: id.to_string(), + }); + } + } + client_policy_map.retain(|_, policy| policy.enabled()); + + Ok(Self { + revision, + account_selection_policy, + instances: Arc::new(instance_map), + provider_models: Arc::new(model_map), + provider_model_mappings: Arc::new(BTreeMap::new()), + provider_catalog_generations: Arc::new(BTreeMap::new()), + client_policies: Arc::new(client_policy_map), + }) + } + + #[must_use] + pub fn with_provider_model_mappings( + mut self, + mappings: BTreeMap>, + ) -> Self { + self.provider_model_mappings = Arc::new(mappings); + self + } + + #[must_use] + fn with_provider_catalog_generations( + mut self, + generations: BTreeMap, + ) -> Self { + self.provider_catalog_generations = Arc::new(generations); + self + } + + #[must_use] + pub fn provider_catalog_generations( + &self, + ) -> &BTreeMap { + &self.provider_catalog_generations + } + + #[must_use] + pub const fn revision(&self) -> ConfigRevision { + self.revision + } + + /// 返回目录发现模型与设置映射的并集，仅用于公开模型展示。 + #[must_use] + pub fn public_models_for_provider(&self, provider: &ProviderKind) -> Vec { + let mut models = BTreeSet::new(); + for (instance_id, discovered) in self.provider_models.iter() { + let Some(instance) = self.instances.get(instance_id) else { + continue; + }; + if !instance.enabled || instance.provider() != provider { + continue; + } + models.extend( + discovered + .keys() + .filter_map(|model| PublicModelId::new(model.as_str().to_owned()).ok()), + ); + } + if let Some(mappings) = self.provider_model_mappings.get(provider) { + models.extend( + mappings + .keys() + .filter_map(|model| PublicModelId::new(model.clone()).ok()), + ); + } + models.into_iter().collect() + } + + #[must_use] + pub fn public_models(&self) -> Vec { + let providers = self + .instances + .values() + .map(|instance| instance.provider().clone()) + .collect::>(); + providers + .iter() + .flat_map(|provider| self.public_models_for_provider(provider)) + .collect::>() + .into_iter() + .collect() + } + + /// 透明代理不以目录白名单拒绝模型；只要求该平台至少有一个启用 instance。 + #[must_use] + pub fn contains_public_model_for_provider( + &self, + _public_model: &PublicModelId, + provider: &ProviderKind, + ) -> bool { + self.instances + .values() + .any(|instance| instance.enabled && instance.provider() == provider) + } + + #[must_use] + pub fn instance_ids_for_provider( + &self, + provider: &ProviderKind, + ) -> BTreeSet { + self.instances + .iter() + .filter_map(|(id, instance)| { + (instance.enabled && instance.provider() == provider).then_some(id.clone()) + }) + .collect() + } + + #[must_use] + pub fn provider_for_instance(&self, instance_id: &ProviderInstanceId) -> Option<&ProviderKind> { + self.instances + .get(instance_id) + .map(ProviderInstance::provider) + } + + #[must_use] + pub fn mapped_model(&self, provider: &ProviderKind, requested: &str) -> String { + self.provider_model_mappings + .get(provider) + .and_then(|mapping| mapping.get(requested)) + .cloned() + .unwrap_or_else(|| requested.to_owned()) + } + + pub fn client_policies(&self) -> impl Iterator { + self.client_policies.values() + } + + #[must_use] + pub fn client_policy(&self, key_id: &ClientApiKeyId) -> Option<&ClientPolicy> { + self.client_policies.get(key_id) + } + + pub fn plan( + &self, + public_model: &PublicModelId, + operation: &Operation, + context: &RoutingContext, + ) -> Result { + let upstream_model = context.provider_kind.as_ref().map_or_else( + || public_model.as_str().to_owned(), + |provider| self.mapped_model(provider, public_model.as_str()), + ); + let upstream_model = + UpstreamModelId::new(upstream_model).map_err(|_| RoutingError::InvalidIdentifier)?; + let requirements = operation.capability_requirements(); + let mut candidates = Vec::new(); + + for instance in self.instances.values() { + if !self.instance_is_eligible(instance, context) { + continue; + } + let emulated_features = match self + .provider_models + .get(instance.id()) + .and_then(|models| models.get(&upstream_model)) + { + Some(capabilities) => { + let Some(emulated) = capabilities.match_requirements(&requirements) else { + continue; + }; + emulated + } + None => BTreeSet::new(), + }; + candidates.push(ProviderCandidate { + instance: instance.clone(), + upstream_model: upstream_model.clone(), + emulated_features, + observation_token: context + .provider_observation_tokens + .get(instance.id()) + .copied(), + }); + } + + if candidates.is_empty() { + return Err(RoutingError::NoCapableProvider { + model: public_model.as_str().to_owned(), + }); + } + + Ok(RoutingPlan { + config_revision: self.revision, + account_selection_policy: self.account_selection_policy, + public_model: public_model.clone(), + operation: operation.kind(), + max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS) + .expect("constant request attempt limit is non-zero"), + candidates: Arc::from(candidates), + }) + } + + fn instance_is_eligible(&self, instance: &ProviderInstance, context: &RoutingContext) -> bool { + instance.enabled + && instance.health.is_routable(context.allow_degraded) + && context + .provider_kind + .as_ref() + .is_none_or(|provider| instance.provider() == provider) + && !context.blocked_instances.contains(instance.id()) + && context + .allowed_instances + .as_ref() + .is_none_or(|allowed| allowed.contains(instance.id())) + } +} + +/// 请求级冻结失败；此状态必须 fail closed。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("runtime snapshot is unavailable")] +pub struct RuntimeSnapshotUnavailable; + +/// RuntimeSnapshot 原子发布和请求级冻结句柄。 +#[derive(Clone, Default)] +pub struct RuntimeSnapshotHandle { + current: Arc>>>, +} + +impl RuntimeSnapshotHandle { + #[must_use] + pub fn new(initial: RuntimeSnapshot) -> Self { + Self { + current: Arc::new(RwLock::new(Some(Arc::new(initial)))), + } + } + + pub fn publish(&self, snapshot: RuntimeSnapshot) { + *write_unpoisoned(&self.current) = Some(Arc::new(snapshot)); + } + + pub fn suspend(&self) { + *write_unpoisoned(&self.current) = None; + } + + #[must_use] + pub fn revision(&self) -> Option { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.revision()) + } + + #[must_use] + pub fn provider_catalog_generations( + &self, + ) -> Option> { + read_unpoisoned(&self.current) + .as_ref() + .map(|snapshot| snapshot.provider_catalog_generations().clone()) + } + + /// 冻结当前 Arc；后续发布不改变已经开始的请求。 + pub fn acquire(&self) -> Result, RuntimeSnapshotUnavailable> { + read_unpoisoned(&self.current) + .clone() + .ok_or(RuntimeSnapshotUnavailable) + } +} + +impl HealthProbe for RuntimeSnapshotHandle { + fn name(&self) -> &'static str { + "runtime_snapshot" + } + + fn check(&self) -> BoxFuture<'_, HealthState> { + Box::pin(async move { + if self.revision().is_some() { + HealthState::Healthy + } else { + HealthState::Unhealthy("Runtime snapshot is unavailable".to_owned()) + } + }) + } +} + +/// Admin 提交配置后触发本进程刷新与跨进程通知的对象安全端口。 +pub trait SnapshotControl: Send + Sync { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()>; +} + +/// 配置提交后的本进程快照发布与跨进程失效通知。 +#[derive(Clone)] +pub struct RuntimeSnapshotPublisher { + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, +} + +impl RuntimeSnapshotPublisher { + #[must_use] + pub const fn new( + compiler: Arc, + snapshots: RuntimeSnapshotHandle, + subscriptions: Arc, + ) -> Self { + Self { + compiler, + snapshots, + subscriptions, + } + } + + /// 重新编译并原子替换本进程快照。 + pub async fn refresh(&self) -> Result { + let snapshot = self.compiler.compile().await?; + let revision = snapshot.revision(); + self.snapshots.publish(snapshot); + Ok(revision) + } + + pub fn suspend(&self) { + self.snapshots.suspend(); + } + + #[must_use] + pub fn published_revision(&self) -> Option { + self.snapshots.revision() + } + + #[must_use] + fn provider_catalogs_need_refresh(&self) -> bool { + self.snapshots.provider_catalog_generations().as_ref() + != Some(&self.compiler.providers.catalog_generations()) + } + + /// 数据库提交不能被目录或通知基础设施的暂时故障伪装成回滚。 + async fn publish_committed_inner(&self, committed_revision: ConfigRevision) { + if self.refresh().await.is_err() { + self.suspend(); + } + let _ = self + .subscriptions + .publish_snapshot_revision(committed_revision) + .await; + } + + /// 交给 Host 的周期对账与长驻订阅任务。 + pub fn worker_contributions(&self) -> Result, WorkerDefinitionError> { + let reconciliation_id = WorkerId::try_new( + WorkerKind::RuntimeSnapshotReconciliation, + "runtime_snapshot", + )?; + let schedule = WorkerSchedule::try_new( + RECONCILIATION_INTERVAL, + INITIAL_BACKOFF, + MAXIMUM_BACKOFF, + UNUSED_LEASE_TTL, + UNUSED_LEASE_RENEWAL, + )?; + let reconciliation = WorkerRegistration::try_new( + reconciliation_id, + WorkerRunnable::Scheduled { + schedule, + lease: None, + task: Box::new(RuntimeSnapshotReconciliationTask { + store: Arc::clone(&self.compiler.store), + publisher: self.clone(), + }), + }, + )?; + let subscription_id = + WorkerId::try_new(WorkerKind::RuntimeChangeSubscription, "runtime_snapshot")?; + let restart = DaemonRestartPolicy::try_new(INITIAL_BACKOFF, MAXIMUM_BACKOFF)?; + let subscription = WorkerRegistration::try_new( + subscription_id, + WorkerRunnable::Daemon { + restart, + task: Box::new(RuntimeSnapshotSubscriptionTask { + subscriptions: Arc::clone(&self.subscriptions), + publisher: self.clone(), + }), + }, + )?; + Ok(vec![ + WorkerContribution::Registration(reconciliation), + WorkerContribution::Registration(subscription), + ]) + } +} + +impl SnapshotControl for RuntimeSnapshotPublisher { + fn publish_committed(&self, committed_revision: ConfigRevision) -> BoxFuture<'_, ()> { + Box::pin(async move { + self.publish_committed_inner(committed_revision).await; + }) + } +} + +struct RuntimeSnapshotReconciliationTask { + store: Arc, + publisher: RuntimeSnapshotPublisher, +} + +impl ScheduledTask for RuntimeSnapshotReconciliationTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let persisted_revision = match self.store.current_config_revision().await { + Ok(revision) => revision, + Err(_) => { + self.publisher.suspend(); + return Err(WorkerTaskError::safe( + "runtime snapshot revision is unavailable", + )); + } + }; + let configuration_changed = runtime_revision_needs_refresh( + self.publisher.published_revision().map(ConfigRevision::get), + persisted_revision.get(), + ); + if !configuration_changed && !self.publisher.provider_catalogs_need_refresh() { + return Ok(()); + } + if self.publisher.refresh().await.is_err() { + // 已提交的配置无法确认时必须 fail closed；仅目录重编译失败时继续 + // 服务旧的不可变快照，单调代次会让下一周期再次尝试。 + if configuration_changed { + self.publisher.suspend(); + } + return Err(WorkerTaskError::safe( + "runtime snapshot reconciliation failed", + )); + } + Ok(()) + }) + } +} + +struct RuntimeSnapshotSubscriptionTask { + subscriptions: Arc, + publisher: RuntimeSnapshotPublisher, +} + +impl DaemonTask for RuntimeSnapshotSubscriptionTask { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let mut retry_delay = INITIAL_BACKOFF; + loop { + if cancellation.is_cancelled() { + return Ok(()); + } + let subscription = self.subscriptions.subscribe_snapshot_revisions().await; + let mut subscription = match subscription { + Ok(subscription) => { + retry_delay = INITIAL_BACKOFF; + subscription + } + Err(_) => { + wait_or_cancel(&cancellation, retry_delay).await; + retry_delay = (retry_delay * 2).min(MAXIMUM_BACKOFF); + continue; + } + }; + loop { + let cancelled = cancellation.cancelled().fuse(); + let next = subscription.next().fuse(); + pin_mut!(cancelled, next); + let notified = select_biased! { + _ = cancelled => return Ok(()), + next = next => next, + }; + match notified { + Some(Ok(_)) => { + if self.publisher.refresh().await.is_err() { + self.publisher.suspend(); + } + } + Some(Err(_)) | None => break, + } + } + } + }) + } +} + +async fn wait_or_cancel(cancellation: &CancellationToken, duration: Duration) { + let cancelled = cancellation.cancelled().fuse(); + let delay = Delay::new(duration).fuse(); + pin_mut!(cancelled, delay); + select_biased! { + _ = cancelled => {}, + _ = delay => {}, + } +} + +/// 当前发布版本与持久版本不一致时必须重载；缺失和回退同样 fail closed。 +#[must_use] +pub fn runtime_revision_needs_refresh( + published_revision: Option, + persisted_revision: u64, +) -> bool { + published_revision != Some(persisted_revision) +} + +fn read_unpoisoned(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { + lock.read() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn write_unpoisoned(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { + lock.write() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-core/src/task/mod.rs b/backend/crates/gateway-core/src/task/mod.rs new file mode 100644 index 000000000..4c82e34f9 --- /dev/null +++ b/backend/crates/gateway-core/src/task/mod.rs @@ -0,0 +1,537 @@ +//! 后台任务、leader lease 与运行计划的中立契约。 +//! +//! 本模块不执行任务，也不启动异步运行时。各业务 owner 只返回 +//! [`WorkerContribution`]，由 Host 负责监督、续租、重启与关闭。 + +use std::fmt; +use std::num::NonZeroU64; +use std::time::Duration; + +use futures::future::BoxFuture; + +use crate::engine::CancellationToken; + +/// 冻结架构中的全部后台任务类别。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum WorkerKind { + OAuthRefresh, + QuotaCatalogHealth, + RuntimeSnapshotReconciliation, + RuntimeChangeSubscription, + NativeClaimRecovery, + StaleModelRequestRecovery, + Retention, + OpsFlush, +} + +impl WorkerKind { + pub const ALL: [Self; 8] = [ + Self::OAuthRefresh, + Self::QuotaCatalogHealth, + Self::RuntimeSnapshotReconciliation, + Self::RuntimeChangeSubscription, + Self::NativeClaimRecovery, + Self::StaleModelRequestRecovery, + Self::Retention, + Self::OpsFlush, + ]; + + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::OAuthRefresh => "oauth_refresh", + Self::QuotaCatalogHealth => "quota_catalog_health", + Self::RuntimeSnapshotReconciliation => "runtime_snapshot_reconciliation", + Self::RuntimeChangeSubscription => "runtime_change_subscription", + Self::NativeClaimRecovery => "native_claim_recovery", + Self::StaleModelRequestRecovery => "stale_model_request_recovery", + Self::Retention => "retention", + Self::OpsFlush => "ops_flush", + } + } +} + +impl fmt::Display for WorkerKind { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +/// 一个后台任务的稳定身份；同类任务可由多个 owner 各自贡献。 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerId { + kind: WorkerKind, + owner: String, +} + +impl WorkerId { + /// 从任务类别和可用于持久化 key 的 owner 创建身份。 + /// + /// owner 长度为 1..=64，必须以小写 ASCII 字母或数字开头， + /// 其余字符仅允许小写 ASCII 字母、数字、`-`、`_` 和 `.`。 + pub fn try_new( + kind: WorkerKind, + owner: impl Into, + ) -> Result { + let owner = owner.into(); + if !valid_owner(&owner) { + return Err(WorkerDefinitionError::InvalidOwner); + } + Ok(Self { kind, owner }) + } + + #[must_use] + pub const fn kind(&self) -> WorkerKind { + self.kind + } + + #[must_use] + pub fn owner(&self) -> &str { + &self.owner + } +} + +impl fmt::Display for WorkerId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{}:{}", self.kind, self.owner) + } +} + +/// 正整数 fencing token；零值无法被构造。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct WorkerFencingToken(NonZeroU64); + +impl WorkerFencingToken { + #[must_use] + pub const fn new(value: NonZeroU64) -> Self { + Self(value) + } + + #[must_use] + pub const fn get(self) -> u64 { + self.0.get() + } +} + +impl From for WorkerFencingToken { + fn from(value: NonZeroU64) -> Self { + Self::new(value) + } +} + +/// 一个周期任务的完整监督时序。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WorkerSchedule { + interval: Duration, + initial_backoff: Duration, + maximum_backoff: Duration, + leader_lease_ttl: Duration, + leader_lease_renewal_interval: Duration, +} + +impl WorkerSchedule { + /// 构造已校验的时序。 + /// + /// 所有时长必须大于零；最大退避不得小于初始退避； + /// lease 续租间隔必须严格小于 lease TTL。 + pub fn try_new( + interval: Duration, + initial_backoff: Duration, + maximum_backoff: Duration, + leader_lease_ttl: Duration, + leader_lease_renewal_interval: Duration, + ) -> Result { + if interval.is_zero() + || initial_backoff.is_zero() + || maximum_backoff.is_zero() + || leader_lease_ttl.is_zero() + || leader_lease_renewal_interval.is_zero() + || maximum_backoff < initial_backoff + || leader_lease_renewal_interval >= leader_lease_ttl + { + return Err(WorkerDefinitionError::InvalidSchedule); + } + Ok(Self { + interval, + initial_backoff, + maximum_backoff, + leader_lease_ttl, + leader_lease_renewal_interval, + }) + } + + #[must_use] + pub const fn interval(self) -> Duration { + self.interval + } + + #[must_use] + pub const fn initial_backoff(self) -> Duration { + self.initial_backoff + } + + #[must_use] + pub const fn maximum_backoff(self) -> Duration { + self.maximum_backoff + } + + #[must_use] + pub const fn leader_lease_ttl(self) -> Duration { + self.leader_lease_ttl + } + + #[must_use] + pub const fn leader_lease_renewal_interval(self) -> Duration { + self.leader_lease_renewal_interval + } +} + +/// 长驻任务退出或 panic 后的 Host 重启策略。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DaemonRestartPolicy { + initial_backoff: Duration, + maximum_backoff: Duration, +} + +impl DaemonRestartPolicy { + pub fn try_new( + initial_backoff: Duration, + maximum_backoff: Duration, + ) -> Result { + if initial_backoff.is_zero() + || maximum_backoff.is_zero() + || maximum_backoff < initial_backoff + { + return Err(WorkerDefinitionError::InvalidDaemonRestartPolicy); + } + Ok(Self { + initial_backoff, + maximum_backoff, + }) + } + + #[must_use] + pub const fn initial_backoff(self) -> Duration { + self.initial_backoff + } + + #[must_use] + pub const fn maximum_backoff(self) -> Duration { + self.maximum_backoff + } +} + +/// 一次 leader lease 申请。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerLeaseRequest { + worker: WorkerId, + ttl: Duration, +} + +impl WorkerLeaseRequest { + pub fn try_new(worker: WorkerId, ttl: Duration) -> Result { + if ttl.is_zero() { + return Err(WorkerDefinitionError::InvalidLeaseTtl); + } + Ok(Self { worker, ttl }) + } + + #[must_use] + pub const fn worker(&self) -> &WorkerId { + &self.worker + } + + #[must_use] + pub const fn ttl(&self) -> Duration { + self.ttl + } +} + +/// 已获取 lease 的可续租句柄。 +/// +/// Host 在正常路径必须显式 `await` [`Self::release`]；异常退出只依赖 +/// Redis TTL 兜底，实现不得在 `Drop` 中启动异步任务。 +pub trait WorkerLeaderLeaseGuard: Send + Sync { + fn fencing_token(&self) -> WorkerFencingToken; + + fn renew(&mut self) -> BoxFuture<'_, Result<(), WorkerLeaseError>>; + + fn release(self: Box) -> BoxFuture<'static, Result<(), WorkerLeaseError>>; +} + +/// Leader lease 的稳定获取结果。 +pub enum WorkerLeaseAcquisition { + Acquired(Box), + Busy { retry_after: Option }, +} + +impl fmt::Debug for WorkerLeaseAcquisition { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Acquired(_) => formatter.write_str("Acquired([LEASE_GUARD])"), + Self::Busy { retry_after } => formatter + .debug_struct("Busy") + .field("retry_after", retry_after) + .finish(), + } + } +} + +/// 多实例 worker leader lease 的中立端口。 +pub trait WorkerLeaderLeasePort: Send + Sync { + fn try_acquire( + &self, + request: WorkerLeaseRequest, + ) -> BoxFuture<'_, Result>; +} + +/// 周期任务收到的单周期上下文。 +#[derive(Debug, Clone)] +pub struct WorkerCycleContext { + worker: WorkerId, + fencing_token: Option, + cancellation: CancellationToken, +} + +impl WorkerCycleContext { + #[must_use] + pub fn new( + worker: WorkerId, + fencing_token: Option, + cancellation: CancellationToken, + ) -> Self { + Self { + worker, + fencing_token, + cancellation, + } + } + + #[must_use] + pub const fn worker(&self) -> &WorkerId { + &self.worker + } + + #[must_use] + pub const fn fencing_token(&self) -> Option { + self.fencing_token + } + + #[must_use] + pub const fn cancellation(&self) -> &CancellationToken { + &self.cancellation + } +} + +/// 由 Host 周期调用的短生命任务。 +pub trait ScheduledTask: Send + Sync { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>>; +} + +/// 长驻任务。 +/// +/// 实现负责连接级重连退避；Host 负责任务退出或 panic 后的重启退避。 +pub trait DaemonTask: Send + Sync { + fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>>; +} + +/// Host 可执行的两种任务形态，无法组合出“守护任务 + 周期调度”等非法状态。 +pub enum WorkerRunnable { + Scheduled { + schedule: WorkerSchedule, + lease: Option, + task: Box, + }, + Daemon { + restart: DaemonRestartPolicy, + task: Box, + }, +} + +impl fmt::Debug for WorkerRunnable { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Scheduled { + schedule, lease, .. + } => formatter + .debug_struct("Scheduled") + .field("schedule", schedule) + .field("lease", lease) + .field("task", &"[SCHEDULED_TASK]") + .finish(), + Self::Daemon { restart, .. } => formatter + .debug_struct("Daemon") + .field("restart", restart) + .field("task", &"[DAEMON_TASK]") + .finish(), + } + } +} + +/// 一个 owner 交给 Host 的任务注册。 +#[derive(Debug)] +pub struct WorkerRegistration { + pub id: WorkerId, + pub runnable: WorkerRunnable, +} + +impl WorkerRegistration { + pub fn try_new(id: WorkerId, runnable: WorkerRunnable) -> Result { + let registration = Self { id, runnable }; + registration.validate()?; + Ok(registration) + } + + /// 校验公开字段构造出的注册是否保持身份与 lease 时序一致。 + pub fn validate(&self) -> Result<(), WorkerDefinitionError> { + let WorkerRunnable::Scheduled { + schedule, + lease: Some(lease), + .. + } = &self.runnable + else { + return Ok(()); + }; + if lease.worker() != &self.id { + return Err(WorkerDefinitionError::LeaseWorkerMismatch); + } + if lease.ttl() != schedule.leader_lease_ttl() { + return Err(WorkerDefinitionError::LeaseTtlMismatch); + } + Ok(()) + } +} + +/// Final DB 明确没有相应持久化状态时允许的两种禁用原因。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum WorkerDisabledReason { + NoPersistentNativeClaimState, + NoBufferedOpsEvents, +} + +impl WorkerDisabledReason { + #[must_use] + pub const fn kind(self) -> WorkerKind { + match self { + Self::NoPersistentNativeClaimState => WorkerKind::NativeClaimRecovery, + Self::NoBufferedOpsEvents => WorkerKind::OpsFlush, + } + } + + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::NoPersistentNativeClaimState => { + "native continuation has no persistent claim state" + } + Self::NoBufferedOpsEvents => "ops events have no flush buffer", + } + } +} + +impl fmt::Display for WorkerDisabledReason { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +/// 各包对 Host 计划的一项贡献。 +#[derive(Debug)] +pub enum WorkerContribution { + Registration(WorkerRegistration), + Disabled { + kind: WorkerKind, + reason: WorkerDisabledReason, + }, +} + +impl WorkerContribution { + #[must_use] + pub const fn kind(&self) -> WorkerKind { + match self { + Self::Registration(registration) => registration.id.kind(), + Self::Disabled { kind, .. } => *kind, + } + } + + pub fn validate(&self) -> Result<(), WorkerDefinitionError> { + match self { + Self::Registration(registration) => registration.validate(), + Self::Disabled { kind, reason } if *kind == reason.kind() => Ok(()), + Self::Disabled { .. } => Err(WorkerDefinitionError::DisabledReasonMismatch), + } + } +} + +/// 任务定义在交给 Host 前的稳定校验错误。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum WorkerDefinitionError { + #[error("worker owner is invalid")] + InvalidOwner, + #[error("worker schedule is invalid")] + InvalidSchedule, + #[error("daemon restart policy is invalid")] + InvalidDaemonRestartPolicy, + #[error("worker lease ttl is invalid")] + InvalidLeaseTtl, + #[error("worker registration lease belongs to a different worker")] + LeaseWorkerMismatch, + #[error("worker registration lease ttl differs from its schedule")] + LeaseTtlMismatch, + #[error("disabled worker kind does not match its reason")] + DisabledReasonMismatch, +} + +/// 不暴露基础设施原文或 lease resource 的错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{message}")] +pub struct WorkerLeaseError { + message: String, +} + +impl WorkerLeaseError { + #[must_use] + pub fn safe(message: impl Into) -> Self { + Self { + message: message.into(), + } + } + + #[must_use] + pub fn as_safe_str(&self) -> &str { + &self.message + } +} + +/// 不暴露 Provider 原文或存储细节的后台任务错误。 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{message}")] +pub struct WorkerTaskError { + message: String, +} + +impl WorkerTaskError { + #[must_use] + pub fn safe(message: impl Into) -> Self { + Self { + message: message.into(), + } + } + + #[must_use] + pub fn as_safe_str(&self) -> &str { + &self.message + } +} + +fn valid_owner(owner: &str) -> bool { + !owner.is_empty() + && owner.len() <= 64 + && owner.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"-_.".contains(&byte) + }) + && owner + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) +} diff --git a/backend/crates/gateway-core/tests/accounting/mod.rs b/backend/crates/gateway-core/tests/accounting/mod.rs index c71195806..6c6d5df1a 100644 --- a/backend/crates/gateway-core/tests/accounting/mod.rs +++ b/backend/crates/gateway-core/tests/accounting/mod.rs @@ -1,8 +1,9 @@ use std::str::FromStr; use gateway_core::accounting::{ - CalculatedCost, CalculatedCostBreakdown, CostEstimate, CostEstimateStatus, CostSource, - CurrencyCode, Decimal, Money, ProviderReportedCost, Usage, + CalculatedCost, CalculatedCostAmounts, CalculatedCostBreakdown, CalculatedCostRates, + CostEstimate, CostEstimateStatus, CostSource, CurrencyCode, Decimal, Money, + ProviderReportedCost, Usage, }; #[test] @@ -41,16 +42,20 @@ fn calculated_breakdown_should_preserve_runtime_components_without_changing_tota let usd = CurrencyCode::new("USD").expect("USD currency"); let money = |ticks| Money::new(Decimal::from_scaled(ticks).expect("valid ticks"), usd); let breakdown = CalculatedCostBreakdown::new( - money(100), - money(200), - money(30), - money(0), - money(330), - money(330), - money(10_000_000_000), - money(20_000_000_000), - money(3_000_000_000), - money(0), + CalculatedCostAmounts::new( + money(100), + money(200), + money(30), + money(0), + money(330), + money(330), + ), + CalculatedCostRates::new( + money(10_000_000_000), + money(20_000_000_000), + money(3_000_000_000), + money(0), + ), Some("default".to_owned()), 100, ); diff --git a/backend/crates/gateway-core/tests/engine/admission.rs b/backend/crates/gateway-core/tests/engine/admission.rs new file mode 100644 index 000000000..2a8ce4af5 --- /dev/null +++ b/backend/crates/gateway-core/tests/engine/admission.rs @@ -0,0 +1,214 @@ +use std::collections::VecDeque; +use std::sync::Mutex; +use std::time::{Duration, SystemTime}; + +use async_trait::async_trait; +use futures::future::BoxFuture; +use gateway_core::engine::admission::{ + ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery, + ClientAdmissionRecoveryPort, ClientAdmissionRequest, ClientAdmissionRestoreResult, + RecentAdmissionFact, RunningAdmissionFact, restore_client_admission_startup, +}; +use gateway_core::engine::{ + AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId, + NewModelRequest, RecoveryReport, UpstreamSendState, +}; +use gateway_core::error::{StoreError, StoreErrorKind}; +use gateway_core::policy::ClientApiKeyId; + +#[test] +fn client_admission_startup_recovery_should_preserve_order_and_exact_facts() { + futures::executor::block_on(async { + let now = SystemTime::now(); + let request = ModelRequestId::new("req_recovery").expect("request id"); + let recovery = ClientAdmissionRecovery { + client_api_key_id: ClientApiKeyId::new("key_recovery").expect("key id"), + recent_requests: vec![RecentAdmissionFact { + model_request_id: request.clone(), + started_at: now - Duration::from_secs(2), + input_token_estimate: 31, + }], + running_requests: vec![RunningAdmissionFact { + model_request_id: request, + expires_at: now + Duration::from_secs(30), + }], + }; + let recoveries = ScriptedRecoveries::new(Ok(vec![recovery.clone()])); + let admissions = RecordingAdmissions::default(); + let report = restore_client_admission_startup( + &RecoveryStore::success(2), + &recoveries, + &admissions, + now, + ) + .await + .expect("startup recovery"); + + assert_eq!(report.expired_model_requests, 2); + assert_eq!(report.restored_clients, 1); + assert_eq!( + admissions.restored.lock().expect("restored").as_slice(), + &[recovery] + ); + }); +} + +#[test] +fn client_admission_startup_recovery_should_fail_closed_at_each_boundary() { + futures::executor::block_on(async { + let now = SystemTime::now(); + assert!( + restore_client_admission_startup( + &RecoveryStore::failure(), + &ScriptedRecoveries::new(Ok(Vec::new())), + &RecordingAdmissions::default(), + now, + ) + .await + .is_err() + ); + assert!( + restore_client_admission_startup( + &RecoveryStore::success(0), + &ScriptedRecoveries::new(Err(ClientAdmissionError)), + &RecordingAdmissions::default(), + now, + ) + .await + .is_err() + ); + let admissions = RecordingAdmissions { + fail_restore: true, + ..RecordingAdmissions::default() + }; + assert!( + restore_client_admission_startup( + &RecoveryStore::success(0), + &ScriptedRecoveries::new(Ok(vec![empty_recovery()])), + &admissions, + now, + ) + .await + .is_err() + ); + }); +} + +fn empty_recovery() -> ClientAdmissionRecovery { + ClientAdmissionRecovery { + client_api_key_id: ClientApiKeyId::new("key_empty").expect("key id"), + recent_requests: Vec::new(), + running_requests: Vec::new(), + } +} + +struct RecoveryStore(Result); + +impl RecoveryStore { + fn success(requests: u64) -> Self { + Self(Ok(RecoveryReport { requests })) + } + + fn failure() -> Self { + Self(Err(StoreError::new(StoreErrorKind::Unavailable))) + } +} + +#[async_trait] +impl ExecutionStore for RecoveryStore { + async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> { + unreachable!() + } + async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> { + unreachable!() + } + async fn mark_send_state( + &self, + _: &ModelRequestId, + _: UpstreamSendState, + ) -> Result<(), StoreError> { + unreachable!() + } + async fn mark_downstream_committed( + &self, + _: &ModelRequestId, + _: SystemTime, + _: Option, + ) -> Result<(), StoreError> { + unreachable!() + } + async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> { + unreachable!() + } + async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> { + unreachable!() + } + async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> { + unreachable!() + } + async fn recover_expired(&self, _: SystemTime) -> Result { + self.0.clone() + } +} + +struct ScriptedRecoveries( + Mutex, ClientAdmissionError>>>, +); + +impl ScriptedRecoveries { + fn new(result: Result, ClientAdmissionError>) -> Self { + Self(Mutex::new(VecDeque::from([result]))) + } +} + +impl ClientAdmissionRecoveryPort for ScriptedRecoveries { + fn load_recovery( + &self, + _: SystemTime, + ) -> BoxFuture<'_, Result, ClientAdmissionError>> { + Box::pin(async move { + self.0 + .lock() + .expect("recovery script") + .pop_front() + .expect("script") + }) + } +} + +#[derive(Default)] +struct RecordingAdmissions { + restored: Mutex>, + fail_restore: bool, +} + +impl ClientAdmissionPort for RecordingAdmissions { + fn admit( + &self, + _: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + unreachable!() + } + fn release<'a>( + &'a self, + _: &'a ClientApiKeyId, + _: &'a ModelRequestId, + ) -> BoxFuture<'a, Result> { + unreachable!() + } + fn restore( + &self, + recovery: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + if self.fail_restore { + return Err(ClientAdmissionError); + } + self.restored.lock().expect("restored").push(recovery); + Ok(ClientAdmissionRestoreResult { + restored_recent_requests: 1, + restored_running_requests: 1, + }) + }) + } +} diff --git a/backend/crates/gateway-core/tests/engine/continuation.rs b/backend/crates/gateway-core/tests/engine/continuation.rs index 8d5bd6d95..7eabaab9c 100644 --- a/backend/crates/gateway-core/tests/engine/continuation.rs +++ b/backend/crates/gateway-core/tests/engine/continuation.rs @@ -1,6 +1,6 @@ use gateway_core::engine::continuation::{ - NativeClaim, NativeClaimState, NativeContinuationPin, NativeContinuationReuse, - PreviousResponseId, + ContinuationBinding, NativeClaim, NativeClaimState, NativeContinuationPin, + NativeContinuationReuse, PreviousResponseId, }; use gateway_core::engine::credential::ProviderAccountId; use gateway_core::error::SafeUpstreamValue; @@ -24,6 +24,15 @@ fn native_pin_debug_should_redact_previous_response_id() { assert!(!debug.contains("upstream-private")); } +#[test] +fn external_binding_debug_should_redact_previous_response_id() { + let binding = ContinuationBinding::External( + PreviousResponseId::new("external-private").expect("valid response"), + ); + + assert!(!format!("{binding:?}").contains("external-private")); +} + #[test] fn native_pin_should_reject_different_account() { assert!(!pin().matches( diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index afbc41e04..4b9594260 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -10,24 +10,31 @@ use futures::executor::block_on; use gateway_core::accounting::{CalculatedCost, CostSource, ProviderReportedCost, Usage}; use gateway_core::engine::continuation::{ - NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, }; use gateway_core::engine::credential::{ AccountSelectionPolicy, CredentialRevision, ProviderAccountId, RotationStrategy, }; use gateway_core::engine::provider::{ - Provider, ProviderCallMetadata, ProviderModelCapabilities, ProviderRegistry, ProviderRequest, - ProviderResource, ProviderStream, UpstreamTransport, + Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities, + ProviderRegistry, ProviderRequest, ProviderResource, ProviderStream, UpstreamTransport, }; use gateway_core::engine::{ AttemptContext, AttemptCoordinator, AttemptRecord, CancellationToken, CommitRequirement, - ExecutionOutcome, ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization, - ModelRequestId, NewModelRequest, ProviderAttemptOutcome, RecoveryReport, UpstreamSendState, + ContinuationAttempt, ExecutionOutcome, ExecutionStore, GatewayEngine, IntermediateFailure, + ModelRequestFinalization, ModelRequestId, NewModelRequest, ProviderAttemptOutcome, + RecoveryReport, UpstreamSendState, +}; +use gateway_core::error::{ + ContinuationFailure, ProviderError, ProviderErrorKind, SafeUpstreamValue, StoreError, +}; +use gateway_core::event::{ + GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseObservation, + ProviderResponseTimings, ResponseMeta, UpstreamHttpVersion, }; -use gateway_core::error::{ProviderError, ProviderErrorKind, SafeUpstreamValue, StoreError}; -use gateway_core::event::{GatewayEvent, ResponseMeta}; use gateway_core::operation::{ - ContentPart, GenerateRequest, Message, MessageRole, Operation, OperationKind, RetrySafety, + ContentPart, GenerateRequest, Message, MessageRole, Operation, OperationKind, + ProviderSessionState, RetrySafety, }; use gateway_core::policy::ClientApiKeyId; use gateway_core::routing::{ @@ -35,6 +42,7 @@ use gateway_core::routing::{ ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan, RuntimeSnapshot, UpstreamModelId, }; +use serde_json::{Map, Value}; #[derive(Debug, Clone, PartialEq, Eq)] struct FinalState { @@ -49,6 +57,14 @@ struct FinalState { latency_ms: Option, client_response_id: Option, upstream_response_id: Option, + upstream_transport: Option, + http_version: Option, + upstream_request_id: Option, + upstream_status_code: Option, + transport_decision_wait_ms: Option, + connect_ms: Option, + headers_ms: Option, + first_event_ms: Option, cost_source: CostSource, cost_ticks: Option, } @@ -62,6 +78,8 @@ struct StoreState { committed_statuses: Vec>, recorded_statuses: Vec, intermediate_failures: usize, + intermediate_status_codes: Vec>, + intermediate_request_ids: Vec>, finalizations: Vec, } @@ -126,9 +144,16 @@ impl ExecutionStore for FakeStore { async fn record_intermediate_failure( &self, - _failure: IntermediateFailure, + failure: IntermediateFailure, ) -> Result<(), StoreError> { - self.state.lock().expect("store lock").intermediate_failures += 1; + let mut state = self.state.lock().expect("store lock"); + state.intermediate_failures += 1; + state + .intermediate_status_codes + .push(failure.upstream_status_code); + state + .intermediate_request_ids + .push(failure.upstream_request_id); Ok(()) } @@ -152,6 +177,14 @@ impl ExecutionStore for FakeStore { latency_ms: finalization.timings.latency_ms, client_response_id: finalization.client_response_id, upstream_response_id: finalization.upstream_response_id, + upstream_transport: finalization.upstream_transport, + http_version: finalization.http_version, + upstream_request_id: finalization.upstream_request_id, + upstream_status_code: finalization.upstream_status_code, + transport_decision_wait_ms: finalization.timings.transport_decision_wait_ms, + connect_ms: finalization.timings.connect_ms, + headers_ms: finalization.timings.headers_ms, + first_event_ms: finalization.timings.first_event_ms, cost_source: finalization.cost.source(), cost_ticks: finalization .cost @@ -171,6 +204,10 @@ enum Script { account_id: &'static str, items: Vec>, }, + ObservedStream { + account_id: &'static str, + items: Vec>, + }, Error(ProviderError), } @@ -194,6 +231,10 @@ impl Provider for ScriptedProvider { "openai" } + fn catalog_generation(&self) -> ProviderCatalogGeneration { + ProviderCatalogGeneration::default() + } + async fn query_model_capabilities( &self, _instance: &ProviderInstance, @@ -230,6 +271,28 @@ impl Provider for ScriptedProvider { .with_upstream_request_id( SafeUpstreamValue::new("upstream-request").expect("safe request id"), ); + Ok(ProviderStream::new( + metadata, + Box::pin(futures::stream::iter( + items + .into_iter() + .map(|event| event.map(ProviderEvent::from)), + )), + (), + )) + } + Script::ObservedStream { account_id, items } => { + let candidate = request.candidate(); + let metadata = ProviderCallMetadata::new( + candidate.provider().clone(), + candidate.instance().clone(), + candidate.upstream_model().clone(), + ProviderResource::Account { + id: ProviderAccountId::new(account_id).expect("account id"), + revision: CredentialRevision::new(1).expect("revision"), + }, + UpstreamTransport::new("websocket").expect("planned transport"), + ); Ok(ProviderStream::new( metadata, Box::pin(futures::stream::iter(items)), @@ -366,7 +429,7 @@ fn coordinator( fn terminal_non_idempotent_failure( items: Vec>, - continuation: Option, + continuation: Option, ) -> (Arc, Arc) { let operation = operation(RetrySafety::NonIdempotent); let route_plan = plan(&operation, 2, 1); @@ -417,11 +480,11 @@ fn success_updates_one_model_request_and_persists_usage() { .expect("start execution"); let events = block_on(session.collect_uncommitted()).expect("collect response"); assert_eq!(events.len(), 3); - let started_id = match &events[0] { + let started_id = match &events[0].canonical_facts()[0] { GatewayEvent::Started(metadata) => metadata.response_id(), event => panic!("unexpected first event: {event:?}"), }; - let completed_id = match &events[2] { + let completed_id = match &events[2].canonical_facts()[0] { GatewayEvent::Completed(metadata) => metadata.response_id(), event => panic!("unexpected final event: {event:?}"), }; @@ -459,6 +522,188 @@ fn success_updates_one_model_request_and_persists_usage() { ); } +#[test] +fn response_observation_is_persisted_but_never_delivered() { + let operation = operation(RetrySafety::Idempotent); + let route_plan = plan(&operation, 1, 1); + let observation = ProviderResponseObservation::new( + UpstreamTransport::new("http_sse").expect("actual transport"), + ) + .with_http_version(UpstreamHttpVersion::Http2) + .with_status_code(200) + .with_request_id(SafeUpstreamValue::new("upstream-observed").expect("request id")) + .with_timings(ProviderResponseTimings { + transport_decision_wait_ms: Some(7), + connect_ms: Some(11), + headers_ms: Some(13), + first_event_ms: Some(17), + }); + let mut items = vec![Ok(ProviderEvent::observation(observation))]; + items.extend( + complete_stream(Some(21)) + .into_iter() + .map(|event| event.map(ProviderEvent::canonical)), + ); + let (coordinator, store, _) = coordinator(vec![Script::ObservedStream { + account_id: "acct_observed", + items, + }]); + let mut session = block_on(coordinator.start( + model_request(SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start execution"); + + let events = block_on(session.collect_uncommitted()).expect("collect response"); + assert_eq!(events.len(), 3); + block_on(session.commit_downstream(Some(200))).expect("commit response"); + + let state = store.state.lock().expect("store lock"); + let finalization = &state.finalizations[0]; + assert_eq!(finalization.upstream_transport.as_deref(), Some("http_sse")); + assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2")); + assert_eq!( + finalization.upstream_request_id.as_deref(), + Some("upstream-observed") + ); + assert_eq!(finalization.upstream_status_code, Some(200)); + assert_eq!(finalization.transport_decision_wait_ms, Some(7)); + assert_eq!(finalization.connect_ms, Some(11)); + assert_eq!(finalization.headers_ms, Some(13)); + assert_eq!(finalization.first_event_ms, Some(17)); +} + +#[test] +fn unknown_wire_event_before_response_identity_is_discarded_with_retried_attempt() { + let operation = operation(RetrySafety::Idempotent); + let route_plan = plan(&operation, 2, 1); + let unknown = ProtocolWireEvent::json( + "openai", + Some("response.future_event".to_owned()), + serde_json::json!({"type": "response.future_event", "opaque": true}), + ) + .expect("wire event"); + let (coordinator, store, provider) = coordinator(vec![ + Script::ObservedStream { + account_id: "acct_first", + items: vec![ + Ok(ProviderEvent::wire(unknown)), + Err(ProviderError::new( + ProviderErrorKind::Unavailable, + UpstreamSendState::Sent, + )), + ], + }, + Script::Stream { + account_id: "acct_second", + items: complete_stream(None), + }, + ]); + let mut session = block_on(coordinator.start( + model_request(SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start execution"); + + block_on(session.collect_uncommitted()).expect("retry response"); + block_on(session.commit_downstream(Some(200))).expect("commit winning response"); + + assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2); + assert!(provider.scripts.lock().expect("scripts lock").is_empty()); + let state = store.state.lock().expect("store lock"); + assert_eq!(state.intermediate_failures, 1); + assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); +} + +#[test] +fn discarded_attempt_observation_does_not_leak_into_retry_result() { + let operation = operation(RetrySafety::Idempotent); + let route_plan = plan(&operation, 2, 1); + let first_observation = ProviderResponseObservation::new( + UpstreamTransport::new("websocket").expect("first transport"), + ) + .with_http_version(UpstreamHttpVersion::Http11) + .with_status_code(503) + .with_request_id(SafeUpstreamValue::new("discarded-request").expect("request id")) + .with_timings(ProviderResponseTimings { + transport_decision_wait_ms: Some(7), + connect_ms: Some(11), + headers_ms: Some(13), + first_event_ms: Some(987_654), + }); + let second_observation = ProviderResponseObservation::new( + UpstreamTransport::new("http_sse").expect("second transport"), + ) + .with_http_version(UpstreamHttpVersion::Http2) + .with_status_code(200) + .with_request_id(SafeUpstreamValue::new("winning-request").expect("request id")); + let mut second_items = vec![Ok(ProviderEvent::observation(second_observation))]; + second_items.extend( + complete_stream(None) + .into_iter() + .map(|event| event.map(ProviderEvent::canonical)), + ); + let (coordinator, store, _) = coordinator(vec![ + Script::ObservedStream { + account_id: "acct_first", + items: vec![ + Ok(ProviderEvent::observation(first_observation)), + Err( + ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent) + .with_replay_safe(), + ), + ], + }, + Script::ObservedStream { + account_id: "acct_second", + items: second_items, + }, + ]); + let mut session = block_on(coordinator.start( + model_request(SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start execution"); + + block_on(session.collect_uncommitted()).expect("retry response"); + block_on(session.commit_downstream(Some(200))).expect("commit response"); + + let state = store.state.lock().expect("store lock"); + let finalization = &state.finalizations[0]; + assert_eq!(finalization.upstream_transport.as_deref(), Some("http_sse")); + assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2")); + assert_eq!( + finalization.upstream_request_id.as_deref(), + Some("winning-request") + ); + assert_eq!(finalization.upstream_status_code, Some(200)); + assert_eq!(finalization.transport_decision_wait_ms, None); + assert_eq!(finalization.connect_ms, None); + assert_eq!(finalization.headers_ms, None); + assert!( + finalization + .first_event_ms + .is_some_and(|elapsed| elapsed < 987_654) + ); + assert_eq!(state.intermediate_status_codes, vec![Some(503)]); + assert_eq!( + state.intermediate_request_ids, + vec![Some("discarded-request".to_owned())] + ); +} + #[test] fn websocket_success_keeps_client_http_status_absent() { let operation = operation(RetrySafety::Idempotent); @@ -507,7 +752,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() { operation, route_plan, None, - Some(continuation), + Some(ContinuationBinding::Pinned(continuation)), CancellationToken::new(), )) .expect("start execution"); @@ -519,6 +764,8 @@ fn authenticated_native_continuation_reaches_every_attempt_context() { contexts[0] .continuation() .expect("continuation") + .pinned() + .expect("pinned continuation") .previous_response_id() .as_str(), "previous-secret-id" @@ -527,10 +774,101 @@ fn authenticated_native_continuation_reaches_every_attempt_context() { contexts[0] .continuation() .expect("continuation") + .pinned() + .expect("pinned continuation") .upstream_response_id() .as_str(), "provider-native-id" ); + let owner = contexts[0] + .account_state_owner() + .expect("continuation account-state owner"); + assert!(owner.matches( + &ProviderKind::new("openai").expect("provider"), + &ProviderInstanceId::new("inst_openai_1").expect("instance"), + &ProviderAccountId::new("acct_one").expect("account"), + )); +} + +#[test] +fn native_continuation_replays_owner_before_safely_switching_account() { + let Operation::Generate(generate) = operation(RetrySafety::Idempotent) else { + panic!("generate operation"); + }; + let mut payload = Map::new(); + payload.insert("transcript".to_owned(), Value::Array(Vec::new())); + let operation = Operation::Generate(generate.with_provider_session_state( + ProviderSessionState::new("openai", payload).expect("provider session state"), + )); + let route_plan = plan(&operation, 3, 1); + let (coordinator, _, provider) = coordinator(vec![ + Script::Stream { + account_id: "acct_one", + items: vec![ + Ok(GatewayEvent::Started(ResponseMeta::new( + "response-native", + "gpt-5", + ))), + Err( + ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent) + .with_continuation_failure(ContinuationFailure::HistoryUnavailable) + .with_replay_safe(), + ), + ], + }, + Script::Stream { + account_id: "acct_one", + items: vec![Err(ProviderError::new( + ProviderErrorKind::Transport, + UpstreamSendState::NotSent, + ) + .with_replay_safe())], + }, + Script::Stream { + account_id: "acct_two", + items: complete_stream(None), + }, + ]); + let continuation = NativeContinuationPin::new( + PreviousResponseId::new("previous-secret-id").expect("previous response"), + SafeUpstreamValue::new("provider-native-id").expect("upstream response"), + ProviderKind::new("openai").expect("provider"), + ProviderInstanceId::new("inst_openai_1").expect("instance"), + ProviderAccountId::new("acct_one").expect("account"), + NativeContinuationReuse::Reusable, + ); + + let mut session = block_on(coordinator.start( + model_request(SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + Some(ContinuationBinding::Pinned(continuation)), + CancellationToken::new(), + )) + .expect("start execution"); + block_on(session.collect_uncommitted()).expect("portable replay succeeds"); + block_on(session.commit_downstream(Some(200))).expect("commit response"); + + let contexts = provider.contexts.lock().expect("contexts lock"); + assert_eq!(contexts.len(), 3); + assert_eq!( + contexts[0].continuation_attempt(), + ContinuationAttempt::Native + ); + assert_eq!( + contexts[1].continuation_attempt(), + ContinuationAttempt::ReplayOwner + ); + assert_eq!( + contexts[2].continuation_attempt(), + ContinuationAttempt::ReplayAny + ); + assert!( + contexts[2] + .excluded_accounts() + .contains(&ProviderAccountId::new("acct_one").expect("account")) + ); } #[test] @@ -941,6 +1279,16 @@ fn pre_commit_failure_excludes_account_and_retries_same_target() { let contexts = provider.contexts.lock().expect("contexts lock"); assert_eq!(contexts.len(), 2); + assert!(contexts[0].account_state_owner().is_none()); + assert!( + contexts[1] + .account_state_owner() + .is_some_and(|owner| owner.matches( + &ProviderKind::new("openai").expect("provider"), + &ProviderInstanceId::new("inst_openai_1").expect("instance"), + &ProviderAccountId::new("acct_first").expect("account"), + )) + ); assert!( contexts[1] .excluded_accounts() @@ -1038,25 +1386,47 @@ fn explicit_429_with_ambiguous_send_state_is_not_retried() { } #[test] -fn explicit_429_after_canonical_event_is_not_retried() { - let (store, provider) = terminal_non_idempotent_failure( - vec![ - Ok(GatewayEvent::Started(ResponseMeta::new( - "response-started", - "gpt-5", - ))), - Err( - ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent) - .with_status(429) - .with_replay_safe(), - ), - ], +fn explicit_429_after_structural_event_should_retry_before_commit() { + let operation = operation(RetrySafety::NonIdempotent); + let route_plan = plan(&operation, 2, 1); + let (coordinator, store, provider) = coordinator(vec![ + Script::Stream { + account_id: "acct_first", + items: vec![ + Ok(GatewayEvent::Started(ResponseMeta::new( + "response-started", + "gpt-5", + ))), + Err( + ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent) + .with_status(429) + .with_replay_safe(), + ), + ], + }, + Script::Stream { + account_id: "acct_second", + items: complete_stream(None), + }, + ]); + let mut session = block_on(coordinator.start( + model_request(SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, None, - ); - assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1); + None, + CancellationToken::new(), + )) + .expect("start execution"); + + block_on(session.collect_uncommitted()).expect("second account succeeds"); + block_on(session.commit_downstream(Some(200))).expect("commit response"); + + assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2); let state = store.state.lock().expect("store lock"); - assert_eq!(state.attempts.len(), 1); - assert_eq!(state.intermediate_failures, 0); + assert_eq!(state.attempts.len(), 2); + assert_eq!(state.intermediate_failures, 1); + assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); } #[test] @@ -1069,6 +1439,26 @@ fn native_continuation_explicit_429_is_not_retried() { ProviderAccountId::new("acct_first").expect("account"), NativeContinuationReuse::Reusable, ); + let (store, provider) = terminal_non_idempotent_failure( + vec![Err(ProviderError::new( + ProviderErrorKind::RateLimited, + UpstreamSendState::Sent, + ) + .with_status(429) + .with_replay_safe())], + Some(ContinuationBinding::Pinned(continuation)), + ); + assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1); + let state = store.state.lock().expect("store lock"); + assert_eq!(state.attempts.len(), 1); + assert_eq!(state.intermediate_failures, 0); +} + +#[test] +fn external_continuation_explicit_429_is_not_retried() { + let continuation = ContinuationBinding::External( + PreviousResponseId::new("external-provider-response").expect("previous response"), + ); let (store, provider) = terminal_non_idempotent_failure( vec![Err(ProviderError::new( ProviderErrorKind::RateLimited, @@ -1079,6 +1469,11 @@ fn native_continuation_explicit_429_is_not_retried() { Some(continuation), ); assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1); + assert!( + provider.contexts.lock().expect("contexts lock")[0] + .account_state_owner() + .is_none() + ); let state = store.state.lock().expect("store lock"); assert_eq!(state.attempts.len(), 1); assert_eq!(state.intermediate_failures, 0); @@ -1111,7 +1506,10 @@ fn non_idempotent_not_sent_failure_can_fallback_target() { block_on(session.collect_uncommitted()).expect("safe instance fallback succeeds"); block_on(session.commit_downstream(Some(200))).expect("commit response"); - assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2); + let contexts = provider.contexts.lock().expect("contexts lock"); + assert_eq!(contexts.len(), 2); + assert!(contexts[0].account_state_owner().is_none()); + assert!(contexts[1].account_state_owner().is_none()); let state = store.state.lock().expect("store lock"); assert_eq!(state.attempts.len(), 1); assert_eq!( @@ -1167,7 +1565,7 @@ fn ambiguous_send_state_stops_retry() { } #[test] -fn failure_after_first_event_is_incomplete_and_never_retried() { +fn structural_event_before_replay_safe_failure_should_switch_account_before_commit() { let operation = operation(RetrySafety::Idempotent); let route_plan = plan(&operation, 2, 1); let (coordinator, store, provider) = coordinator(vec![ @@ -1178,10 +1576,10 @@ fn failure_after_first_event_is_incomplete_and_never_retried() { "response-1", "gpt-5", ))), - Err(ProviderError::new( - ProviderErrorKind::Transport, - UpstreamSendState::Sent, - )), + Err( + ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Sent) + .with_replay_safe(), + ), ], }, Script::Stream { @@ -1206,13 +1604,9 @@ fn failure_after_first_event_is_incomplete_and_never_retried() { first.commit_requirement(), CommitRequirement::CommitBeforeDelivery ); + assert_eq!(first.into_provider_events().len(), 1); block_on(session.commit_downstream(Some(200))).expect("commit first event"); - let error = block_on(session.next_event()).expect_err("committed stream failure is terminal"); - - assert!(matches!( - error, - gateway_core::engine::EngineError::Provider(_) - )); + block_on(session.next_event()).expect_err("committed stream failure must not be replayed"); assert_eq!( session.provider_attempt_outcomes(), &[ProviderAttemptOutcome::Failed { diff --git a/backend/crates/gateway-core/tests/engine/credential.rs b/backend/crates/gateway-core/tests/engine/credential.rs index 4caae21a9..7e6c88ccc 100644 --- a/backend/crates/gateway-core/tests/engine/credential.rs +++ b/backend/crates/gateway-core/tests/engine/credential.rs @@ -5,10 +5,10 @@ use std::time::{Duration, SystemTime}; use serde_json::{Map, Value}; use gateway_core::engine::credential::{ - AccountAvailability, AccountCandidate, AccountRuntimeSignals, AccountSelectionContext, - AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, CredentialRevision, - OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId, - ProviderAccountUpdate, RotationStrategy, + AccountAvailability, AccountCandidate, AccountQuotaSignals, AccountRuntimeSignals, + AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate, + CredentialRevision, OpaqueProviderData, PlaintextCredential, ProviderAccount, + ProviderAccountId, ProviderAccountUpdate, RotationStrategy, }; use gateway_core::routing::{ProviderInstanceId, ProviderKind}; @@ -85,6 +85,27 @@ fn disabled_account_should_not_be_schedulable() { assert!(!account.is_schedulable(SystemTime::now())); } +#[test] +fn cooldown_account_should_become_schedulable_at_its_deadline() { + let now = SystemTime::now(); + let account = + account("acct_cooldown").with_runtime_state(true, AccountAvailability::Cooldown, Some(now)); + + assert!(account.is_schedulable(now)); +} + +#[test] +fn cooldown_account_should_remain_isolated_before_its_deadline() { + let now = SystemTime::now(); + let account = account("acct_cooldown").with_runtime_state( + true, + AccountAvailability::Cooldown, + Some(now + Duration::from_secs(1)), + ); + + assert!(!account.is_schedulable(now)); +} + #[test] fn smart_selector_should_prefer_lower_inflight_count() { let candidates = vec![ @@ -98,6 +119,145 @@ fn smart_selector_should_prefer_lower_inflight_count() { assert_eq!(selected.account.id().as_str(), "acct_idle"); } +#[test] +fn provider_quota_overlay_should_preserve_store_concurrency_facts() { + let reset_at = SystemTime::now() + Duration::from_secs(60); + let last_started_at = SystemTime::now(); + let signals = AccountRuntimeSignals { + in_flight: 2, + last_started_at: Some(last_started_at), + quota_reset_at: None, + quota_remaining_rank: None, + } + .with_provider_quota(Some(AccountQuotaSignals::new(Some(reset_at), Some(75)))); + + assert_eq!( + ( + signals.in_flight, + signals.last_started_at, + signals.quota_reset_at, + signals.quota_remaining_rank, + ), + (2, Some(last_started_at), Some(reset_at), Some(75)) + ); +} + +#[test] +fn smart_selector_should_use_provider_quota_rank_after_load_is_equal() { + let candidates = vec![ + candidate("acct_low_quota", 0, Some(20)), + candidate("acct_high_quota", 0, Some(80)), + ]; + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::Smart)) + .expect("candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_high_quota"); +} + +#[test] +fn quota_reset_selector_should_prefer_known_earliest_window() { + let now = SystemTime::now(); + let unknown = candidate("acct_unknown", 0, None); + let mut later = candidate("acct_later", 0, None); + later.signals.quota_reset_at = Some(now + Duration::from_secs(120)); + let mut earlier = candidate("acct_earlier", 0, None); + earlier.signals.quota_reset_at = Some(now + Duration::from_secs(60)); + let candidates = vec![unknown, later, earlier]; + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::QuotaResetPriority)) + .expect("candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_earlier"); +} + +#[test] +fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() { + let candidates = vec![ + candidate("acct_unknown", 0, None), + candidate("acct_known", 0, Some(1)), + ]; + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::Smart)) + .expect("candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_known"); +} + +#[test] +fn selector_should_reject_account_at_concurrency_limit() { + let candidates = vec![ + candidate("acct_full", 3, Some(100)), + candidate("acct_available", 2, Some(1)), + ]; + let selected = AccountSelector + .select(&candidates, &context(RotationStrategy::Smart)) + .expect("candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_available"); +} + +#[test] +fn selector_should_enforce_request_interval_until_boundary() { + let now = SystemTime::now(); + let mut cooling = candidate("acct_cooling", 0, Some(100)); + cooling.signals.last_started_at = Some(now - Duration::from_millis(9)); + let mut ready = candidate("acct_ready", 0, Some(1)); + ready.signals.last_started_at = Some(now - Duration::from_millis(10)); + let candidates = vec![cooling, ready]; + let mut context = context(RotationStrategy::Smart); + context.now = now; + context.policy = AccountSelectionPolicy::new( + RotationStrategy::Smart, + NonZeroU32::new(3).expect("positive"), + Duration::from_millis(10), + ); + let selected = AccountSelector + .select(&candidates, &context) + .expect("candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_ready"); +} + +#[test] +fn sticky_selector_should_fall_back_when_requested_account_is_excluded() { + let candidates = vec![ + candidate("acct_sticky", 0, None), + candidate("acct_fallback", 0, None), + ]; + let mut context = context(RotationStrategy::Sticky); + let sticky = ProviderAccountId::new("acct_sticky").expect("valid account"); + context.sticky_account = Some(sticky.clone()); + context.excluded_accounts.insert(sticky); + let selected = AccountSelector + .select(&candidates, &context) + .expect("fallback candidate available"); + + assert_eq!(selected.account.id().as_str(), "acct_fallback"); +} + +#[test] +fn round_robin_selector_should_be_stable_for_unsorted_candidates() { + let candidates = vec![ + candidate("acct_second", 0, None), + candidate("acct_first", 0, None), + ]; + let mut context = context(RotationStrategy::RoundRobin); + context.round_robin_cursor = 0; + let first = AccountSelector + .select(&candidates, &context) + .expect("first candidate available"); + context.round_robin_cursor = 1; + let second = AccountSelector + .select(&candidates, &context) + .expect("second candidate available"); + + assert_eq!( + (first.account.id().as_str(), second.account.id().as_str()), + ("acct_first", "acct_second") + ); +} + #[test] fn selector_should_honor_request_local_exclusion() { let candidates = vec![ diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs new file mode 100644 index 000000000..b0af2bd97 --- /dev/null +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -0,0 +1,16 @@ +use gateway_core::engine::execution::provider_failure_affects_circuit; +use gateway_core::error::ProviderErrorKind; + +#[test] +fn only_instance_attributable_failures_should_affect_circuit() { + assert!(provider_failure_affects_circuit(ProviderErrorKind::Timeout)); + assert!(provider_failure_affects_circuit( + ProviderErrorKind::Transport + )); + assert!(!provider_failure_affects_circuit( + ProviderErrorKind::RateLimited + )); + assert!(!provider_failure_affects_circuit( + ProviderErrorKind::InvalidRequest + )); +} diff --git a/backend/crates/gateway-core/tests/engine/mod.rs b/backend/crates/gateway-core/tests/engine/mod.rs index 05f844c93..4d741afad 100644 --- a/backend/crates/gateway-core/tests/engine/mod.rs +++ b/backend/crates/gateway-core/tests/engine/mod.rs @@ -1,6 +1,9 @@ +mod admission; mod continuation; mod coordinator; mod credential; +mod execution; +mod probe; mod provider; use gateway_core::engine::{AttemptTrigger, CancellationToken, UpstreamSendState}; diff --git a/backend/crates/gateway-core/tests/engine/probe.rs b/backend/crates/gateway-core/tests/engine/probe.rs new file mode 100644 index 000000000..f7a8c6c0f --- /dev/null +++ b/backend/crates/gateway-core/tests/engine/probe.rs @@ -0,0 +1,7 @@ +use gateway_core::engine::probe::AccountProbe; + +#[test] +fn account_probe_is_an_object_safe_execution_chain_port() { + fn accepts(_: Option<&dyn AccountProbe>) {} + accepts(None); +} diff --git a/backend/crates/gateway-core/tests/engine/provider.rs b/backend/crates/gateway-core/tests/engine/provider.rs index 2afefd7e9..ab493d357 100644 --- a/backend/crates/gateway-core/tests/engine/provider.rs +++ b/backend/crates/gateway-core/tests/engine/provider.rs @@ -6,9 +6,9 @@ use async_trait::async_trait; use futures::stream; use gateway_core::engine::provider::{ - EventStream, Provider, ProviderCallMetadata, ProviderModelCapabilities, ProviderRegistry, - ProviderRequest, ProviderResource, ProviderStream, RegistryError, ResourceId, - UpstreamTransport, + EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, + ProviderModelCapabilities, ProviderRegistry, ProviderRequest, ProviderResource, ProviderStream, + RegistryError, ResourceId, UpstreamTransport, }; use gateway_core::engine::{AttemptContext, UpstreamSendState}; use gateway_core::error::{IdentifierError, ProviderError, ProviderErrorKind}; @@ -26,6 +26,10 @@ impl Provider for NamedProvider { self.0 } + fn catalog_generation(&self) -> ProviderCatalogGeneration { + ProviderCatalogGeneration::default() + } + async fn query_model_capabilities( &self, _instance: &ProviderInstance, diff --git a/backend/crates/gateway-core/tests/event.rs b/backend/crates/gateway-core/tests/event.rs index b54a3cab7..95bb64256 100644 --- a/backend/crates/gateway-core/tests/event.rs +++ b/backend/crates/gateway-core/tests/event.rs @@ -1,7 +1,25 @@ +use gateway_core::engine::provider::UpstreamTransport; use gateway_core::event::{ ContentItem, ContentKind, EventSequenceError, EventSequenceValidator, GatewayEvent, - ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta, + ProtocolWireEvent, ProviderEvent, ProviderResponseObservation, ReasoningDelta, ResponseMeta, + TextDelta, ToolCallDelta, }; +use serde_json::json; + +#[test] +fn provider_event_keeps_large_wire_and_observation_payloads_behind_one_indirection() { + assert!(std::mem::size_of::() <= 64); + + let wire = ProviderEvent::wire( + ProtocolWireEvent::json("openai", None, json!({"type": "response.created"})) + .expect("wire event"), + ); + let observation = ProviderEvent::observation(ProviderResponseObservation::new( + UpstreamTransport::new("http_sse").expect("transport"), + )); + assert!(wire.has_client_event()); + assert!(!observation.has_client_event()); +} #[test] fn validator_should_accept_started_content_delta_completed_sequence() { diff --git a/backend/crates/gateway-core/tests/health.rs b/backend/crates/gateway-core/tests/health.rs new file mode 100644 index 000000000..9f176a761 --- /dev/null +++ b/backend/crates/gateway-core/tests/health.rs @@ -0,0 +1,58 @@ +use futures::future::BoxFuture; +use gateway_core::health::{ + HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, + WorkerRuntimeState, +}; +use gateway_core::task::{WorkerId, WorkerKind}; + +struct HealthyProbe; + +impl HealthProbe for HealthyProbe { + fn name(&self) -> &'static str { + "test" + } + + fn check(&self) -> BoxFuture<'_, HealthState> { + Box::pin(async { HealthState::Healthy }) + } +} + +struct SingleSnapshot; + +impl WorkerHealthSource for SingleSnapshot { + fn snapshot(&self) -> Vec { + vec![WorkerHealthSnapshot { + key: WorkerHealthKey::Task( + WorkerId::try_new(WorkerKind::Retention, "store").expect("valid worker"), + ), + state: WorkerRuntimeState::Idle, + consecutive_failures: 0, + completed_cycles: 3, + last_fencing_token: None, + last_success_at: None, + last_failure_at: None, + last_error: None, + }] + } +} + +#[test] +fn health_probe_contract_is_object_safe() { + let probe: &dyn HealthProbe = &HealthyProbe; + + assert_eq!(probe.name(), "test"); +} + +#[test] +fn worker_health_source_preserves_worker_identity() { + let source: &dyn WorkerHealthSource = &SingleSnapshot; + + assert!(matches!( + source.snapshot().as_slice(), + [WorkerHealthSnapshot { + key: WorkerHealthKey::Task(id), + completed_cycles: 3, + .. + }] if id.owner() == "store" + )); +} diff --git a/backend/crates/gateway-core/tests/lifecycle.rs b/backend/crates/gateway-core/tests/lifecycle.rs new file mode 100644 index 000000000..c3cf6a662 --- /dev/null +++ b/backend/crates/gateway-core/tests/lifecycle.rs @@ -0,0 +1,95 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use gateway_core::engine::CancellationToken; +use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle}; + +#[derive(Default)] +struct LifecycleState { + draining: AtomicBool, + active: AtomicUsize, +} + +struct TestGuard { + state: Arc, +} + +impl ConnectionGuard for TestGuard {} + +impl Drop for TestGuard { + fn drop(&mut self) { + self.state.active.fetch_sub(1, Ordering::AcqRel); + } +} + +struct TestLifecycle { + state: Arc, + cancellation: CancellationToken, +} + +impl TestLifecycle { + fn new() -> Self { + Self { + state: Arc::new(LifecycleState::default()), + cancellation: CancellationToken::new(), + } + } + + fn begin_draining(&self) { + self.state.draining.store(true, Ordering::Release); + self.cancellation.cancel(); + } + + fn active(&self) -> usize { + self.state.active.load(Ordering::Acquire) + } +} + +impl ConnectionLifecycle for TestLifecycle { + fn try_register(&self) -> Result, ConnectionDraining> { + if self.state.draining.load(Ordering::Acquire) { + return Err(ConnectionDraining); + } + self.state.active.fetch_add(1, Ordering::AcqRel); + if self.state.draining.load(Ordering::Acquire) { + self.state.active.fetch_sub(1, Ordering::AcqRel); + return Err(ConnectionDraining); + } + Ok(Box::new(TestGuard { + state: Arc::clone(&self.state), + })) + } + + fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + fn is_draining(&self) -> bool { + self.state.draining.load(Ordering::Acquire) + } +} + +#[test] +fn connection_guard_drop_releases_active_registration() { + let lifecycle = TestLifecycle::new(); + let guard = lifecycle.try_register().expect("registration before drain"); + drop(guard); + + assert_eq!(lifecycle.active(), 0); +} + +#[test] +fn connection_registration_rejects_after_drain_linearization() { + let lifecycle = TestLifecycle::new(); + lifecycle.begin_draining(); + + assert!(matches!(lifecycle.try_register(), Err(ConnectionDraining))); +} + +#[test] +fn connection_lifecycle_contract_is_object_safe() { + let lifecycle = TestLifecycle::new(); + let object: &dyn ConnectionLifecycle = &lifecycle; + + assert!(!object.is_draining()); +} diff --git a/backend/crates/gateway-core/tests/main.rs b/backend/crates/gateway-core/tests/main.rs index 316d450a7..2be2dc7bb 100644 --- a/backend/crates/gateway-core/tests/main.rs +++ b/backend/crates/gateway-core/tests/main.rs @@ -2,6 +2,10 @@ mod accounting; mod engine; mod error; mod event; +mod health; +mod lifecycle; mod operation; mod policy; +mod provider_ports; mod routing; +mod task; diff --git a/backend/crates/gateway-core/tests/operation/mod.rs b/backend/crates/gateway-core/tests/operation/mod.rs index 765595025..b8c662fb9 100644 --- a/backend/crates/gateway-core/tests/operation/mod.rs +++ b/backend/crates/gateway-core/tests/operation/mod.rs @@ -85,6 +85,16 @@ fn response_persistence_should_be_explicit_in_generate_request() { ); } +#[test] +fn prompt_cache_key_should_be_explicit_and_redacted_in_generate_request() { + let request = GenerateRequest::new(vec![text_message(MessageRole::User, "input")]) + .expect("test request is valid") + .with_prompt_cache_key("private-cache-route"); + + assert_eq!(request.prompt_cache_key(), Some("private-cache-route")); + assert!(!format!("{request:?}").contains("private-cache-route")); +} + #[test] fn message_should_reject_empty_content() { let error = Message::new(MessageRole::User, Vec::new()).expect_err("content is required"); diff --git a/backend/crates/gateway-core/tests/provider_ports/mod.rs b/backend/crates/gateway-core/tests/provider_ports/mod.rs new file mode 100644 index 000000000..b127ae642 --- /dev/null +++ b/backend/crates/gateway-core/tests/provider_ports/mod.rs @@ -0,0 +1,114 @@ +use std::collections::BTreeMap; +use std::num::NonZeroU32; +use std::time::{Duration, SystemTime}; + +use gateway_core::engine::credential::{ + AccountRuntimeSignals, CredentialRevision, OpaqueProviderData, ProviderAccountId, +}; +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, ProviderRefreshPolicy, ProviderSchedulingState, + ProviderStoreErrorKind, +}; +use gateway_core::routing::ProviderKind; + +#[test] +fn oauth_pending_binding_debug_redacts_raw_value() { + let binding = OAuthPendingBinding::try_new("must-not-appear").expect("valid binding"); + + assert_eq!(format!("{binding:?}"), "OAuthPendingBinding([REDACTED])"); +} + +#[test] +fn oauth_pending_ttl_rejects_zero_and_more_than_thirty_minutes() { + let provider = ProviderKind::new("fixture").expect("valid provider"); + let flow = OAuthPendingBinding::try_new("flow").expect("valid flow"); + let owner = OAuthPendingBinding::try_new("owner").expect("valid owner"); + let payload = OpaqueProviderData::new(serde_json::Map::new()); + + for ttl in [Duration::ZERO, Duration::from_secs(30 * 60 + 1)] { + let error = NewOAuthPendingFlow::try_new( + provider.clone(), + flow.clone(), + owner.clone(), + ttl, + payload.clone(), + ) + .expect_err("invalid TTL must fail"); + assert_eq!(error.kind(), ProviderStoreErrorKind::InvalidData); + } +} + +#[test] +fn refresh_policy_requires_a_positive_margin() { + let error = ProviderRefreshPolicy::try_new( + Duration::ZERO, + NonZeroU32::new(1).expect("positive concurrency"), + ) + .expect_err("zero margin must fail"); + + assert_eq!(error.kind(), ProviderStoreErrorKind::InvalidData); +} + +#[test] +fn refresh_schedule_is_stable_and_never_uses_query_time_margin() { + let account = ProviderAccountId::new("acct_stable").expect("valid account"); + let policy = ProviderRefreshPolicy::try_new( + Duration::from_secs(3_600), + NonZeroU32::new(2).expect("positive concurrency"), + ) + .expect("valid policy"); + let observed_at = SystemTime::UNIX_EPOCH + Duration::from_secs(10_000); + let expires_at = observed_at + Duration::from_secs(7_200); + + let first = policy + .next_attempt_at(&account, expires_at, observed_at) + .expect("schedule refresh"); + let second = policy + .next_attempt_at(&account, expires_at, observed_at) + .expect("schedule refresh again"); + + assert_eq!(first, second); + assert!(first > observed_at); + assert!(first < expires_at); +} + +#[test] +fn refresh_schedule_uses_observed_time_when_token_is_inside_safe_window() { + let account = ProviderAccountId::new("acct_boundary").expect("valid account"); + let policy = ProviderRefreshPolicy::try_new( + Duration::from_secs(3_600), + NonZeroU32::new(1).expect("positive concurrency"), + ) + .expect("valid policy"); + let observed_at = SystemTime::UNIX_EPOCH + Duration::from_secs(10_000); + + assert_eq!( + policy + .next_attempt_at(&account, observed_at + Duration::from_secs(1), observed_at,) + .expect("schedule immediate refresh"), + observed_at + ); +} + +#[test] +fn scheduling_state_preserves_provider_neutral_signals() { + let account = ProviderAccountId::new("acct_fixture").expect("valid account"); + let signals = BTreeMap::from([( + account.clone(), + AccountRuntimeSignals { + in_flight: 2, + last_started_at: None, + quota_reset_at: None, + quota_remaining_rank: Some(7), + }, + )]); + let state = ProviderSchedulingState::new(signals, Some(account.clone()), 9); + + assert_eq!(state.signals()[&account].in_flight, 2); + assert_eq!(state.sticky_account(), Some(&account)); + assert_eq!(state.round_robin_cursor(), 9); + assert_eq!( + CredentialRevision::new(1).expect("positive revision").get(), + 1 + ); +} diff --git a/backend/crates/gateway-core/tests/routing/mod.rs b/backend/crates/gateway-core/tests/routing/mod.rs index 34ddfc931..457e9b6a1 100644 --- a/backend/crates/gateway-core/tests/routing/mod.rs +++ b/backend/crates/gateway-core/tests/routing/mod.rs @@ -12,6 +12,8 @@ use gateway_core::routing::{ ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId, }; +mod snapshot; + fn scheduling() -> AccountSelectionPolicy { AccountSelectionPolicy::new( RotationStrategy::Smart, diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs new file mode 100644 index 000000000..2e0985346 --- /dev/null +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -0,0 +1,345 @@ +use std::collections::BTreeMap; +use std::num::NonZeroU32; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use futures::executor::block_on; +use futures::future::BoxFuture; + +use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy}; +use gateway_core::engine::provider::{ + Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry, + ProviderRequest, ProviderStream, +}; +use gateway_core::engine::{AttemptContext, UpstreamSendState}; +use gateway_core::error::{ProviderError, ProviderErrorKind}; +use gateway_core::operation::OperationKind; +use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::snapshot::{ + RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle, + RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, + SnapshotProviderInstanceFacts, SnapshotRevisionStream, SnapshotSettingsFacts, + SnapshotStoreError, SnapshotStorePort, SnapshotSubscriptionError, SnapshotSubscriptionPort, + runtime_revision_needs_refresh, +}; +use gateway_core::routing::{ + ConfigRevision, ModelCapabilities, ProviderInstance, ProviderKind, PublicModelId, + RuntimeSnapshot, UpstreamModelId, +}; +use gateway_core::task::WorkerKind; + +#[derive(Clone)] +struct TestSnapshotStore { + facts: Arc>>, + current_revision: Arc>>, +} + +impl TestSnapshotStore { + fn new(facts: Result) -> Self { + let current_revision = facts.as_ref().map(facts_revision).map_err(Clone::clone); + Self { + facts: Arc::new(Mutex::new(facts)), + current_revision: Arc::new(Mutex::new(current_revision)), + } + } +} + +impl SnapshotStorePort for TestSnapshotStore { + fn load_snapshot_facts(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { self.facts.lock().expect("facts lock").clone() }) + } + + fn current_config_revision(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { self.current_revision.lock().expect("revision lock").clone() }) + } +} + +#[derive(Default)] +struct TestSnapshotSubscriptions { + published: Mutex>, +} + +impl SnapshotSubscriptionPort for TestSnapshotSubscriptions { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>> { + Box::pin(async move { + self.published + .lock() + .expect("published lock") + .push(revision); + Ok(()) + }) + } + + fn subscribe_snapshot_revisions( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(Box::pin(futures::stream::empty()) as SnapshotRevisionStream) }) + } +} + +struct PublishingCatalogProvider { + generation: AtomicU64, + queries: AtomicUsize, +} + +#[async_trait] +impl Provider for PublishingCatalogProvider { + fn name(&self) -> &'static str { + "alpha" + } + + fn catalog_generation(&self) -> ProviderCatalogGeneration { + ProviderCatalogGeneration::new(self.generation.load(Ordering::SeqCst)) + } + + async fn query_model_capabilities( + &self, + _: &ProviderInstance, + ) -> Result, ProviderError> { + if self.queries.fetch_add(1, Ordering::SeqCst) == 0 { + self.generation.store(1, Ordering::SeqCst); + } + Ok(vec![ProviderModelCapabilities::new( + UpstreamModelId::new("upstream-model").expect("model"), + ModelCapabilities::new( + std::collections::BTreeSet::from([OperationKind::Generate]), + 128_000, + None, + ), + )]) + } + + async fn execute( + &self, + _: ProviderRequest, + _: AttemptContext, + ) -> Result { + Err(ProviderError::new( + ProviderErrorKind::Unavailable, + UpstreamSendState::NotSent, + )) + } +} + +#[test] +fn runtime_revision_reconciliation_should_refresh_missing_or_stale_snapshot() { + assert!(!runtime_revision_needs_refresh(Some(7), 7)); + assert!(runtime_revision_needs_refresh(Some(6), 7)); + assert!(runtime_revision_needs_refresh(Some(8), 7)); + assert!(runtime_revision_needs_refresh(None, 7)); +} + +#[test] +fn compiler_should_reject_revision_changed_during_consistent_read() { + let facts = facts(1, 2); + let compiler = compiler(Arc::new(TestSnapshotStore::new(Ok(facts)))); + + let error = block_on(compiler.compile()).expect_err("revision drift must fail closed"); + + assert_eq!(error, RuntimeSnapshotCompileError::RevisionChanged); +} + +#[test] +fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() { + let compiler = compiler(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3))))); + + let snapshot = block_on(compiler.compile()).expect("compile snapshot"); + let provider = ProviderKind::new("alpha").expect("provider"); + + assert_eq!(snapshot.revision().get(), 3); + assert!(snapshot.contains_public_model_for_provider( + &PublicModelId::new("unknown-upstream-model").expect("model"), + &provider, + )); + assert_eq!( + snapshot.mapped_model(&provider, "public-model"), + "upstream-model" + ); + assert_eq!(snapshot.client_policies().count(), 1); +} + +#[test] +fn compiler_retries_when_provider_publishes_catalog_during_compilation() { + let provider = Arc::new(PublishingCatalogProvider { + generation: AtomicU64::new(0), + queries: AtomicUsize::new(0), + }); + let providers = + ProviderRegistry::new([provider.clone() as Arc]).expect("provider registry"); + let compiler = + RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers); + + let snapshot = block_on(compiler.compile()).expect("stable catalog snapshot"); + + assert_eq!(provider.queries.load(Ordering::SeqCst), 2); + assert_eq!( + snapshot + .provider_catalog_generations() + .get(&ProviderKind::new("alpha").expect("provider")) + .map(|generation| generation.get()), + Some(1), + ); +} + +#[test] +fn handle_should_keep_request_snapshot_frozen_across_publish() { + let handle = RuntimeSnapshotHandle::new(empty_snapshot(1)); + let frozen = handle.acquire().expect("initial snapshot"); + + handle.publish(empty_snapshot(2)); + + assert_eq!(frozen.revision().get(), 1); + assert_eq!(handle.revision().map(ConfigRevision::get), Some(2)); +} + +#[test] +fn publisher_should_refresh_locally_and_notify_committed_revision() { + let store = Arc::new(TestSnapshotStore::new(Ok(facts(2, 2)))); + let compiler = Arc::new(compiler(store)); + let handle = RuntimeSnapshotHandle::new(empty_snapshot(1)); + let subscriptions = Arc::new(TestSnapshotSubscriptions::default()); + let publisher = RuntimeSnapshotPublisher::new(compiler, handle.clone(), subscriptions.clone()); + + block_on(publisher.publish_committed(revision(2))); + + assert_eq!(handle.revision().map(ConfigRevision::get), Some(2)); + assert_eq!( + subscriptions + .published + .lock() + .expect("published lock") + .as_slice(), + &[revision(2)], + ); +} + +#[test] +fn publisher_should_suspend_but_still_notify_after_committed_refresh_failure() { + let store = Arc::new(TestSnapshotStore::new(Err( + SnapshotStoreError::unavailable(), + ))); + let handle = RuntimeSnapshotHandle::new(empty_snapshot(1)); + let subscriptions = Arc::new(TestSnapshotSubscriptions::default()); + let publisher = RuntimeSnapshotPublisher::new( + Arc::new(compiler(store)), + handle.clone(), + subscriptions.clone(), + ); + + block_on(publisher.publish_committed(revision(2))); + + assert!(handle.acquire().is_err()); + assert_eq!( + subscriptions + .published + .lock() + .expect("published lock") + .as_slice(), + &[revision(2)], + ); +} + +#[test] +fn snapshot_ports_and_control_should_remain_object_safe() { + fn accept_store(_: &dyn SnapshotStorePort) {} + fn accept_subscriptions(_: &dyn SnapshotSubscriptionPort) {} + fn accept_control(_: &dyn SnapshotControl) {} + + let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1)))); + let subscriptions = Arc::new(TestSnapshotSubscriptions::default()); + let publisher = RuntimeSnapshotPublisher::new( + Arc::new(compiler(store.clone())), + RuntimeSnapshotHandle::new(empty_snapshot(1)), + subscriptions.clone(), + ); + + accept_store(store.as_ref()); + accept_subscriptions(subscriptions.as_ref()); + accept_control(&publisher); +} + +#[test] +fn publisher_should_contribute_reconciliation_and_subscription_workers() { + let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1)))); + let publisher = RuntimeSnapshotPublisher::new( + Arc::new(compiler(store)), + RuntimeSnapshotHandle::new(empty_snapshot(1)), + Arc::new(TestSnapshotSubscriptions::default()), + ); + + let contributions = publisher + .worker_contributions() + .expect("valid frozen worker definitions"); + let kinds = contributions + .iter() + .map(gateway_core::task::WorkerContribution::kind) + .collect::>(); + + assert_eq!( + kinds, + vec![ + WorkerKind::RuntimeSnapshotReconciliation, + WorkerKind::RuntimeChangeSubscription, + ], + ); +} + +fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts { + SnapshotFacts::new( + revision(config_revision), + revision(observed_current_revision), + SnapshotSettingsFacts::new( + 3, + 50, + "smart", + BTreeMap::from([( + "alpha".to_owned(), + BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]), + )]), + ), + vec![SnapshotProviderInstanceFacts::new( + "inst_alpha", + "alpha", + "https://alpha.example.invalid", + true, + )], + vec![SnapshotClientPolicyFacts::new( + ClientApiKeyId::new("key_one").expect("key ID"), + PlaintextClientApiKey::new("sk_test").expect("plaintext key"), + "alpha", + RateLimits::unlimited(), + )], + ) +} + +fn facts_revision(facts: &SnapshotFacts) -> ConfigRevision { + facts.config_revision() +} + +fn compiler(store: Arc) -> RuntimeSnapshotCompiler { + RuntimeSnapshotCompiler::new(store, ProviderRegistry::default()) +} + +fn empty_snapshot(value: u64) -> RuntimeSnapshot { + RuntimeSnapshot::new( + revision(value), + AccountSelectionPolicy::new( + RotationStrategy::Smart, + NonZeroU32::new(1).expect("positive concurrency"), + Duration::ZERO, + ), + Vec::new(), + Vec::new(), + Vec::new(), + ) + .expect("empty snapshot") +} + +fn revision(value: u64) -> ConfigRevision { + ConfigRevision::new(value).expect("positive revision") +} diff --git a/backend/crates/gateway-core/tests/task/mod.rs b/backend/crates/gateway-core/tests/task/mod.rs new file mode 100644 index 000000000..ec85a5e3c --- /dev/null +++ b/backend/crates/gateway-core/tests/task/mod.rs @@ -0,0 +1,171 @@ +use std::num::NonZeroU64; +use std::time::Duration; + +use futures::future::BoxFuture; +use gateway_core::engine::CancellationToken; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerDisabledReason, WorkerFencingToken, WorkerId, WorkerKind, + WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError, +}; + +struct NoopScheduledTask; + +impl ScheduledTask for NoopScheduledTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async { Ok(()) }) + } +} + +struct NoopDaemonTask; + +impl DaemonTask for NoopDaemonTask { + fn run(&self, _cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async { Ok(()) }) + } +} + +#[test] +fn schedule_rejects_every_invalid_duration_boundary() { + let one = Duration::from_nanos(1); + let two = Duration::from_nanos(2); + let three = Duration::from_nanos(3); + for values in [ + (Duration::ZERO, one, one, three, one), + (one, Duration::ZERO, one, three, one), + (one, one, Duration::ZERO, three, one), + (one, one, one, Duration::ZERO, one), + (one, one, one, three, Duration::ZERO), + (one, two, one, three, one), + (one, one, one, three, three), + (one, one, one, two, three), + ] { + assert_eq!( + WorkerSchedule::try_new(values.0, values.1, values.2, values.3, values.4), + Err(WorkerDefinitionError::InvalidSchedule), + ); + } +} + +#[test] +fn worker_owner_rejects_empty_unsafe_or_oversized_names() { + for owner in [ + String::new(), + "Codex".to_owned(), + "provider/codex".to_owned(), + "x".repeat(65), + ] { + assert_eq!( + WorkerId::try_new(WorkerKind::OAuthRefresh, owner), + Err(WorkerDefinitionError::InvalidOwner), + ); + } +} + +#[test] +fn schedule_accepts_renewal_strictly_before_lease_expiry() { + let schedule = WorkerSchedule::try_new( + Duration::from_secs(60), + Duration::from_secs(1), + Duration::from_secs(30), + Duration::from_secs(15), + Duration::from_secs(5), + ) + .expect("valid schedule"); + + assert_eq!( + schedule.leader_lease_renewal_interval(), + Duration::from_secs(5) + ); +} + +#[test] +fn registration_rejects_lease_for_a_different_worker() { + let registration_id = worker_id("openai"); + let lease = WorkerLeaseRequest::try_new(worker_id("xai"), Duration::from_secs(15)) + .expect("valid lease"); + let runnable = WorkerRunnable::Scheduled { + schedule: schedule(), + lease: Some(lease), + task: Box::new(NoopScheduledTask), + }; + + assert!(matches!( + WorkerRegistration::try_new(registration_id, runnable), + Err(WorkerDefinitionError::LeaseWorkerMismatch) + )); +} + +#[test] +fn registration_rejects_lease_ttl_that_differs_from_schedule() { + let id = worker_id("openai"); + let lease = + WorkerLeaseRequest::try_new(id.clone(), Duration::from_secs(16)).expect("valid lease"); + let runnable = WorkerRunnable::Scheduled { + schedule: schedule(), + lease: Some(lease), + task: Box::new(NoopScheduledTask), + }; + + assert!(matches!( + WorkerRegistration::try_new(id, runnable), + Err(WorkerDefinitionError::LeaseTtlMismatch) + )); +} + +#[test] +fn disabled_contribution_rejects_reason_for_another_kind() { + let contribution = WorkerContribution::Disabled { + kind: WorkerKind::OpsFlush, + reason: WorkerDisabledReason::NoPersistentNativeClaimState, + }; + + assert_eq!( + contribution.validate(), + Err(WorkerDefinitionError::DisabledReasonMismatch) + ); +} + +#[test] +fn daemon_restart_policy_rejects_reverse_backoff_range() { + assert_eq!( + DaemonRestartPolicy::try_new(Duration::from_secs(2), Duration::from_secs(1)), + Err(WorkerDefinitionError::InvalidDaemonRestartPolicy) + ); +} + +#[test] +fn worker_cycle_context_preserves_optional_fencing_token() { + let token = WorkerFencingToken::new(NonZeroU64::new(7).expect("non-zero fencing token")); + let context = + WorkerCycleContext::new(worker_id("openai"), Some(token), CancellationToken::new()); + + assert_eq!(context.fencing_token(), Some(token)); +} + +#[test] +fn scheduled_and_daemon_tasks_are_object_safe() { + fn accept_scheduled(_task: &dyn ScheduledTask) {} + fn accept_daemon(_task: &dyn DaemonTask) {} + + accept_scheduled(&NoopScheduledTask); + accept_daemon(&NoopDaemonTask); +} + +fn worker_id(owner: &str) -> WorkerId { + WorkerId::try_new(WorkerKind::OAuthRefresh, owner).expect("valid worker identity") +} + +fn schedule() -> WorkerSchedule { + WorkerSchedule::try_new( + Duration::from_secs(60), + Duration::from_secs(1), + Duration::from_secs(30), + Duration::from_secs(15), + Duration::from_secs(5), + ) + .expect("valid schedule") +} diff --git a/backend/crates/gateway-host/Cargo.toml b/backend/crates/gateway-host/Cargo.toml new file mode 100644 index 000000000..d4c1fee77 --- /dev/null +++ b/backend/crates/gateway-host/Cargo.toml @@ -0,0 +1,45 @@ +[package] +name = "gateway-host" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true +autotests = false + +[dependencies] +async-trait = "0.1.89" +axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] } +chrono = { version = "0.4.45", features = ["clock", "serde"] } +config = { version = "0.15.23", default-features = false, features = ["yaml"] } +flate2 = "1.1.9" +futures = "0.3.32" +gateway-admin = { path = "../gateway-admin" } +gateway-core = { path = "../gateway-core" } +hex = "0.4.3" +reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "rustls-tls-native-roots"] } +semver = "1.0.27" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.150" +sha2 = "0.11.0" +tar = "0.4.44" +thiserror = "2.0.18" +tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "signal", "time", "sync", "net"] } +tracing = "0.1.44" +tracing-appender = "0.2.5" +tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] } +uuid = { version = "1.23.3", features = ["v7"] } + +[dev-dependencies] +filetime = "0.2.27" +tempfile = "3.27.0" +tokio = { version = "1.52.3", default-features = false, features = ["rt-multi-thread", "test-util"] } +tower = { version = "0.5.3", features = ["util"] } +wiremock = "0.6.5" + +[[test]] +name = "main" +path = "tests/main.rs" + +[lints] +workspace = true diff --git a/backend/crates/gateway-host/src/config.rs b/backend/crates/gateway-host/src/config.rs new file mode 100644 index 000000000..b1096ebf3 --- /dev/null +++ b/backend/crates/gateway-host/src/config.rs @@ -0,0 +1,196 @@ +//! 配置文件发现、反序列化与 Host-owned 配置校验。 + +use std::env; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use serde::Deserialize; +use serde::de::DeserializeOwned; +use tracing_subscriber::EnvFilter; + +use crate::system_update::SystemUpdateConfig; + +const CONFIG_RELATIVE_PATH: &str = "deploy/config.yaml"; +const SERVER_HOST_ENV: &str = "CPR_SERVER_HOST"; +const SERVER_PORT_ENV: &str = "CPR_SERVER_PORT"; + +/// 由组装根实现的顶层配置契约。 +/// +/// Host 只负责找到和解析文件；每个包的字段解释与相对路径解析 +/// 由顶层配置委托给对应的包完成。 +pub trait LoadableConfig: DeserializeOwned { + fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError>; +} + +/// 从当前目录或父目录中的 `deploy/config.yaml` 加载顶层配置。 +pub fn load_config() -> Result { + let current = env::current_dir().map_err(|_| ConfigError::CurrentDirectory)?; + let path = discover_config_path(¤t)?; + let source_dir = path.parent().ok_or(ConfigError::InvalidConfigPath)?; + let mut value: T = config::Config::builder() + .add_source(config::File::from(path.as_path()).required(true)) + .build() + .and_then(config::Config::try_deserialize) + .map_err(|_| ConfigError::InvalidDocument { path: path.clone() })?; + value.resolve_and_validate(source_dir)?; + Ok(value) +} + +/// Host 唯一拥有的进程配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct HostConfig { + pub listen: ListenConfig, + pub logging: LoggingConfig, + #[serde(default)] + pub system_update: SystemUpdateConfig, + #[serde(default = "default_drain_timeout_seconds")] + pub drain_timeout_seconds: u64, + #[serde(default = "default_worker_shutdown_timeout_seconds")] + pub worker_shutdown_timeout_seconds: u64, +} + +impl HostConfig { + pub fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError> { + if let Some(host) = optional_environment_value(SERVER_HOST_ENV)? { + self.listen.host = host; + } + if let Some(port) = optional_environment_value(SERVER_PORT_ENV)? { + self.listen.port = port + .parse::() + .ok() + .filter(|port| *port > 0) + .ok_or(ConfigError::InvalidEnvironment(SERVER_PORT_ENV))?; + } + if self.listen.host.trim().is_empty() { + return Err(ConfigError::InvalidField("host.listen.host")); + } + if self.listen.port == 0 { + return Err(ConfigError::InvalidField("host.listen.port")); + } + if self.drain_timeout_seconds == 0 { + return Err(ConfigError::InvalidField("host.drain_timeout_seconds")); + } + if self.worker_shutdown_timeout_seconds == 0 { + return Err(ConfigError::InvalidField( + "host.worker_shutdown_timeout_seconds", + )); + } + self.logging.resolve_and_validate(source_dir)?; + self.system_update.resolve_and_validate(source_dir)?; + Ok(()) + } + + #[must_use] + pub const fn drain_timeout(&self) -> Duration { + Duration::from_secs(self.drain_timeout_seconds) + } + + #[must_use] + pub const fn worker_shutdown_timeout(&self) -> Duration { + Duration::from_secs(self.worker_shutdown_timeout_seconds) + } +} + +/// HTTP 监听地址。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct ListenConfig { + pub host: String, + pub port: u16, +} + +/// Host 结构化日志配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct LoggingConfig { + pub level: String, + pub stdout: bool, + pub file: FileLoggingConfig, +} + +impl LoggingConfig { + fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError> { + EnvFilter::try_new(&self.level) + .map_err(|_| ConfigError::InvalidField("host.logging.level"))?; + if !self.stdout && !self.file.enabled { + return Err(ConfigError::InvalidField("host.logging")); + } + if self.file.directory.as_os_str().is_empty() { + return Err(ConfigError::InvalidField("host.logging.file.directory")); + } + if self.file.retention_days == 0 { + return Err(ConfigError::InvalidField( + "host.logging.file.retention_days", + )); + } + if self.file.max_file_size_mb == 0 { + return Err(ConfigError::InvalidField( + "host.logging.file.max_file_size_mb", + )); + } + if self.file.max_files == 0 { + return Err(ConfigError::InvalidField("host.logging.file.max_files")); + } + resolve_relative_path(source_dir, &mut self.file.directory); + Ok(()) + } +} + +/// Host 文件日志配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct FileLoggingConfig { + pub enabled: bool, + pub directory: PathBuf, + pub retention_days: usize, + pub max_file_size_mb: u64, + pub max_files: usize, +} + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("current directory is unavailable")] + CurrentDirectory, + #[error("deploy/config.yaml was not found")] + ConfigFileNotFound, + #[error("configuration path is invalid")] + InvalidConfigPath, + #[error("configuration document is invalid: {path}")] + InvalidDocument { path: PathBuf }, + #[error("configuration field is invalid: {0}")] + InvalidField(&'static str), + #[error("environment variable is invalid: {0}")] + InvalidEnvironment(&'static str), +} + +fn optional_environment_value(name: &'static str) -> Result, ConfigError> { + match env::var(name) { + Ok(value) if value.trim().is_empty() => Err(ConfigError::InvalidEnvironment(name)), + Ok(value) => Ok(Some(value)), + Err(env::VarError::NotPresent) => Ok(None), + Err(env::VarError::NotUnicode(_)) => Err(ConfigError::InvalidEnvironment(name)), + } +} + +fn discover_config_path(start: &Path) -> Result { + start + .ancestors() + .map(|directory| directory.join(CONFIG_RELATIVE_PATH)) + .find(|candidate| candidate.is_file()) + .ok_or(ConfigError::ConfigFileNotFound) +} + +fn resolve_relative_path(base: &Path, path: &mut PathBuf) { + if path.is_relative() { + *path = base.join(&*path); + } +} + +const fn default_drain_timeout_seconds() -> u64 { + 30 +} + +const fn default_worker_shutdown_timeout_seconds() -> u64 { + 30 +} diff --git a/backend/crates/gateway-host/src/lib.rs b/backend/crates/gateway-host/src/lib.rs new file mode 100644 index 000000000..206044935 --- /dev/null +++ b/backend/crates/gateway-host/src/lib.rs @@ -0,0 +1,113 @@ +//! 网关进程与操作系统能力：配置发现、日志、任务、自更新与 serve/drain。 + +pub mod config; +mod logging; +mod serve; +pub mod system_update; +pub mod workers; + +use std::sync::Arc; + +use axum::Router; +use gateway_admin::ports::system::SystemOperations; +use gateway_core::engine::CancellationToken; +use gateway_core::health::WorkerHealthSource; +use gateway_core::lifecycle::ConnectionLifecycle; +use gateway_core::task::{WorkerContribution, WorkerLeaderLeasePort}; + +pub use config::{ConfigError, HostConfig, LoadableConfig, load_config}; + +use self::logging::{LogGuard, initialize_logging}; +use self::serve::{ConnectionTracker, serve_router}; +use self::system_update::ProcessSystemOperations; +use self::workers::WorkerSupervisor; + +/// Host 初始化的能力集；字段全部私有，不暴露内部监督器或进程状态。 +pub struct HostBundle { + config: HostConfig, + _log_guard: LogGuard, + cancellation: CancellationToken, + connections: Arc, + workers: WorkerSupervisor, + system: Arc, +} + +/// 在启动其他包之前初始化进程级能力。 +pub async fn initialize(config: HostConfig) -> Result { + let log_guard = initialize_logging(&config.logging)?; + let cancellation = CancellationToken::new(); + let connections = Arc::new(ConnectionTracker::new(cancellation.clone())); + let workers = WorkerSupervisor::new(cancellation.clone()); + let system = Arc::new(ProcessSystemOperations::new( + cancellation.clone(), + config.system_update.clone(), + )); + Ok(HostBundle { + config, + _log_guard: log_guard, + cancellation, + connections, + workers, + system, + }) +} + +impl HostBundle { + #[must_use] + pub fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + #[must_use] + pub fn system_operations(&self) -> Arc { + self.system.clone() + } + + #[must_use] + pub fn connection_lifecycle(&self) -> Arc { + self.connections.clone() + } + + #[must_use] + pub fn worker_health(&self) -> Arc { + self.workers.health_source() + } + + pub fn start_workers( + &self, + plan: Vec, + lease: Arc, + ) -> Result<(), HostError> { + self.workers.start(plan, lease)?; + Ok(()) + } + + /// 进程唯一阻塞点；返回前完成 HTTP drain 与 worker join。 + pub async fn serve(self, router: Router) -> Result<(), HostError> { + let result = serve_router( + router, + &self.config.listen.host, + self.config.listen.port, + self.cancellation.clone(), + Arc::clone(&self.connections), + self.config.drain_timeout(), + ) + .await; + self.cancellation.cancel(); + self.workers + .shutdown(self.config.worker_shutdown_timeout()) + .await; + result?; + Ok(()) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum HostError { + #[error(transparent)] + Logging(#[from] logging::LogError), + #[error(transparent)] + Workers(#[from] workers::WorkerStartError), + #[error(transparent)] + Serve(#[from] serve::ServeError), +} diff --git a/backend/crates/gateway-host/src/logging.rs b/backend/crates/gateway-host/src/logging.rs new file mode 100644 index 000000000..866d10858 --- /dev/null +++ b/backend/crates/gateway-host/src/logging.rs @@ -0,0 +1,274 @@ +//! 进程结构化日志与按日期/大小轮转的文件 writer。 + +use std::env; +use std::fs::{self, File, OpenOptions}; +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; + +use chrono::Utc; +use tracing_appender::non_blocking::{NonBlockingBuilder, WorkerGuard}; +use tracing_subscriber::EnvFilter; +use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt as _}; + +use crate::config::LoggingConfig; + +const LOG_FILE_PREFIX: &str = "codex-proxy-rs"; + +/// non-blocking 日志 writer 的进程级守卫。 +pub struct LogGuard { + _writers: Vec, +} + +#[derive(Debug, thiserror::Error)] +pub enum LogError { + #[error("log IO failed")] + Io(#[from] io::Error), + #[error("logging filter is invalid")] + InvalidFilter, + #[error("global tracing subscriber is already initialized")] + AlreadyInitialized, + #[error("logging size limit is too large")] + SizeOverflow, +} + +/// 按自然日、单文件大小、保留天数和文件总数初始化结构化日志。 +pub fn initialize_logging(config: &LoggingConfig) -> Result { + let directive = env::var("RUST_LOG").unwrap_or_else(|_| config.level.clone()); + let filter = EnvFilter::try_new(directive).map_err(|_| LogError::InvalidFilter)?; + let mut guards = Vec::new(); + + let stdout_writer = config.stdout.then(|| { + let (writer, guard) = NonBlockingBuilder::default() + .thread_name("gateway-log-stdout") + .finish(io::stdout()); + guards.push(guard); + writer + }); + let file_writer = if config.file.enabled { + let maximum_bytes = config + .file + .max_file_size_mb + .checked_mul(1024 * 1024) + .ok_or(LogError::SizeOverflow)?; + let writer = RotatingLogWriter::open( + config.file.directory.clone(), + maximum_bytes, + config.file.retention_days, + config.file.max_files, + )?; + let (writer, guard) = NonBlockingBuilder::default() + .thread_name("gateway-log-file") + .finish(writer); + guards.push(guard); + Some(writer) + } else { + None + }; + + let stdout_layer = stdout_writer.map(|writer| { + tracing_subscriber::fmt::layer() + .compact() + .with_writer(writer) + .with_target(true) + .with_ansi(false) + }); + let file_layer = file_writer.map(|writer| { + tracing_subscriber::fmt::layer() + .json() + .with_writer(writer) + .with_target(true) + .with_file(true) + .with_line_number(true) + .with_thread_ids(true) + .with_thread_names(true) + .with_current_span(true) + .with_span_list(true) + }); + tracing_subscriber::registry() + .with(filter) + .with(stdout_layer) + .with(file_layer) + .try_init() + .map_err(|_| LogError::AlreadyInitialized)?; + Ok(LogGuard { _writers: guards }) +} + +struct RotatingLogWriter { + directory: PathBuf, + maximum_bytes: u64, + retention_days: usize, + maximum_files: usize, + date: chrono::NaiveDate, + segment: usize, + bytes_written: u64, + file: File, +} + +impl RotatingLogWriter { + fn open( + directory: PathBuf, + maximum_bytes: u64, + retention_days: usize, + maximum_files: usize, + ) -> io::Result { + fs::create_dir_all(&directory)?; + let date = Utc::now().date_naive(); + cleanup_log_files(&directory, date, retention_days, maximum_files)?; + let (segment, bytes_written) = writable_log_segment(&directory, date, maximum_bytes)?; + let file = open_log_segment(&directory, date, segment)?; + Ok(Self { + directory, + maximum_bytes, + retention_days, + maximum_files, + date, + segment, + bytes_written, + file, + }) + } + + fn rotate_if_required(&mut self, incoming_bytes: usize) -> io::Result<()> { + let date = Utc::now().date_naive(); + let incoming_bytes = u64::try_from(incoming_bytes).unwrap_or(u64::MAX); + let day_changed = date != self.date; + let size_exceeded = !day_changed + && self.bytes_written > 0 + && self.bytes_written.saturating_add(incoming_bytes) > self.maximum_bytes; + if !day_changed && !size_exceeded { + return Ok(()); + } + self.file.flush()?; + if day_changed { + self.date = date; + self.segment = 0; + } else { + self.segment = self.segment.saturating_add(1); + } + self.file = open_log_segment(&self.directory, self.date, self.segment)?; + self.bytes_written = self.file.metadata()?.len(); + cleanup_log_files( + &self.directory, + self.date, + self.retention_days, + self.maximum_files, + ) + } +} + +impl Write for RotatingLogWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.rotate_if_required(buffer.len())?; + let written = self.file.write(buffer)?; + self.bytes_written = self + .bytes_written + .saturating_add(u64::try_from(written).unwrap_or(u64::MAX)); + Ok(written) + } + + fn flush(&mut self) -> io::Result<()> { + self.file.flush() + } +} + +#[derive(Debug)] +struct ManagedLogFile { + date: chrono::NaiveDate, + segment: usize, + path: PathBuf, +} + +fn writable_log_segment( + directory: &Path, + date: chrono::NaiveDate, + maximum_bytes: u64, +) -> io::Result<(usize, u64)> { + let latest = managed_log_files(directory)? + .into_iter() + .filter(|entry| entry.date == date) + .max_by_key(|entry| entry.segment); + let Some(latest) = latest else { + return Ok((0, 0)); + }; + let length = latest.path.metadata()?.len(); + if length >= maximum_bytes { + Ok((latest.segment.saturating_add(1), 0)) + } else { + Ok((latest.segment, length)) + } +} + +fn open_log_segment(directory: &Path, date: chrono::NaiveDate, segment: usize) -> io::Result { + OpenOptions::new() + .create(true) + .append(true) + .open(directory.join(log_file_name(date, segment))) +} + +fn log_file_name(date: chrono::NaiveDate, segment: usize) -> String { + if segment == 0 { + format!("{LOG_FILE_PREFIX}.{date}.log") + } else { + format!("{LOG_FILE_PREFIX}.{date}.{segment}.log") + } +} + +fn managed_log_files(directory: &Path) -> io::Result> { + Ok(fs::read_dir(directory)? + .filter_map(Result::ok) + .filter_map(|entry| { + parse_log_file_name(entry.file_name().to_string_lossy().as_ref()).map( + |(date, segment)| ManagedLogFile { + date, + segment, + path: entry.path(), + }, + ) + }) + .collect()) +} + +fn parse_log_file_name(name: &str) -> Option<(chrono::NaiveDate, usize)> { + let body = name + .strip_prefix(&format!("{LOG_FILE_PREFIX}."))? + .strip_suffix(".log")?; + let (date, segment) = body.split_once('.').map_or((body, 0), |(date, segment)| { + (date, segment.parse().ok().unwrap_or(usize::MAX)) + }); + if segment == usize::MAX { + return None; + } + chrono::NaiveDate::parse_from_str(date, "%Y-%m-%d") + .ok() + .map(|date| (date, segment)) +} + +fn cleanup_log_files( + directory: &Path, + today: chrono::NaiveDate, + retention_days: usize, + maximum_files: usize, +) -> io::Result<()> { + let retention_offset = i64::try_from(retention_days.saturating_sub(1)).unwrap_or(i64::MAX); + let cutoff = today + .checked_sub_signed(chrono::Duration::days(retention_offset)) + .unwrap_or(chrono::NaiveDate::MIN); + let mut retained = Vec::new(); + for entry in managed_log_files(directory)? { + if entry.date < cutoff { + fs::remove_file(entry.path)?; + } else { + retained.push(entry); + } + } + retained.sort_by(|left, right| { + right + .date + .cmp(&left.date) + .then_with(|| right.segment.cmp(&left.segment)) + }); + for entry in retained.into_iter().skip(maximum_files) { + fs::remove_file(entry.path)?; + } + Ok(()) +} diff --git a/backend/crates/gateway-host/src/serve.rs b/backend/crates/gateway-host/src/serve.rs new file mode 100644 index 000000000..10e33e5cb --- /dev/null +++ b/backend/crates/gateway-host/src/serve.rs @@ -0,0 +1,171 @@ +//! HTTP 监听、OS signal、原子连接注册与优雅 drain。 + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use axum::Router; +use gateway_core::engine::CancellationToken; +use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle}; +use tokio::sync::Notify; + +const DRAINING_BIT: usize = 1usize << (usize::BITS - 1); +const ACTIVE_MASK: usize = !DRAINING_BIT; + +pub(crate) struct ConnectionTracker { + state: Arc, + cancellation: CancellationToken, +} + +struct ConnectionState { + value: AtomicUsize, + idle: Notify, +} + +impl ConnectionTracker { + pub(crate) fn new(cancellation: CancellationToken) -> Self { + Self { + state: Arc::new(ConnectionState { + value: AtomicUsize::new(0), + idle: Notify::new(), + }), + cancellation, + } + } + + pub(crate) fn begin_draining(&self) { + let previous = self.state.value.fetch_or(DRAINING_BIT, Ordering::AcqRel); + self.cancellation.cancel(); + if previous & ACTIVE_MASK == 0 { + self.state.idle.notify_waiters(); + } + } + + pub(crate) async fn wait_until_idle(&self, timeout: Duration) { + let wait = async { + loop { + let notified = self.state.idle.notified(); + if self.state.value.load(Ordering::Acquire) & ACTIVE_MASK == 0 { + return; + } + notified.await; + } + }; + let _ = tokio::time::timeout(timeout, wait).await; + } +} + +impl ConnectionLifecycle for ConnectionTracker { + fn try_register(&self) -> Result, ConnectionDraining> { + let mut observed = self.state.value.load(Ordering::Acquire); + loop { + if observed & DRAINING_BIT != 0 || observed & ACTIVE_MASK == ACTIVE_MASK { + return Err(ConnectionDraining); + } + match self.state.value.compare_exchange_weak( + observed, + observed + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + return Ok(Box::new(ActiveConnection { + state: Arc::clone(&self.state), + })); + } + Err(actual) => observed = actual, + } + } + } + + fn cancellation(&self) -> CancellationToken { + self.cancellation.clone() + } + + fn is_draining(&self) -> bool { + self.state.value.load(Ordering::Acquire) & DRAINING_BIT != 0 + } +} + +struct ActiveConnection { + state: Arc, +} + +impl ConnectionGuard for ActiveConnection {} + +impl Drop for ActiveConnection { + fn drop(&mut self) { + let previous = self.state.value.fetch_sub(1, Ordering::AcqRel); + if previous & ACTIVE_MASK == 1 { + self.state.idle.notify_waiters(); + } + } +} + +pub(crate) async fn serve_router( + router: Router, + host: &str, + port: u16, + cancellation: CancellationToken, + connections: Arc, + drain_timeout: Duration, +) -> Result<(), ServeError> { + let listener = tokio::net::TcpListener::bind(format!("{host}:{port}")) + .await + .map_err(ServeError::Bind)?; + tracing::info!(host, port, "网关开始监听"); + + let shutdown_cancellation = cancellation.clone(); + let shutdown_connections = Arc::clone(&connections); + let shutdown = async move { + wait_for_shutdown(&shutdown_cancellation).await; + shutdown_connections.begin_draining(); + }; + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown) + .await + .map_err(ServeError::Serve)?; + + connections.begin_draining(); + connections.wait_until_idle(drain_timeout).await; + Ok(()) +} + +async fn wait_for_shutdown(cancellation: &CancellationToken) { + #[cfg(unix)] + { + let terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()); + match terminate { + Ok(mut terminate) => { + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = terminate.recv() => {} + () = cancellation.cancelled() => return, + } + } + Err(_) => { + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + () = cancellation.cancelled() => return, + } + } + } + } + #[cfg(not(unix))] + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + () = cancellation.cancelled() => return, + } + cancellation.cancel(); +} + +#[derive(Debug, thiserror::Error)] +pub enum ServeError { + #[error("failed to bind HTTP listener")] + Bind(std::io::Error), + #[error("HTTP server failed")] + Serve(std::io::Error), +} diff --git a/backend/crates/gateway-host/src/system_update/archive.rs b/backend/crates/gateway-host/src/system_update/archive.rs new file mode 100644 index 000000000..daed31847 --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/archive.rs @@ -0,0 +1,127 @@ +//! Release tar.gz 安全解包与制品归一化。 + +use std::fs; +use std::path::{Component, Path, PathBuf}; + +use super::{OperationError, internal, invalid}; + +const APP_BINARY_NAME: &str = "codex-proxy-rs"; +const MAX_EXTRACTED_SIZE: u64 = 1024 * 1024 * 1024; +const MAX_ARCHIVE_FILES: usize = 20_000; + +#[derive(Debug)] +pub(crate) struct ExtractedRelease { + pub(crate) binary_path: PathBuf, + pub(crate) web_dist_dir: Option, +} + +pub(crate) fn extract_release( + archive_path: &Path, + temp_dir: &Path, +) -> Result { + let file = fs::File::open(archive_path) + .map_err(|error| internal(format!("failed to open release archive: {error}")))?; + let decoder = flate2::read::GzDecoder::new(file); + let mut archive = tar::Archive::new(decoder); + let binary_path = temp_dir.join(APP_BINARY_NAME); + let web_dist_dir = temp_dir.join("web-dist"); + let mut found_binary = false; + let mut found_web = false; + let mut extracted_size = 0_u64; + let mut file_count = 0_usize; + + for entry in archive + .entries() + .map_err(|error| internal(format!("failed to read release archive: {error}")))? + { + let mut entry = + entry.map_err(|error| internal(format!("invalid archive entry: {error}")))?; + let path = entry + .path() + .map_err(|error| internal(format!("invalid archive path: {error}")))? + .into_owned(); + if unsafe_archive_path(&path) { + return Err(invalid("release archive contains an unsafe path")); + } + if !entry.header().entry_type().is_file() { + continue; + } + file_count = file_count.saturating_add(1); + extracted_size = extracted_size.saturating_add(entry.header().size().unwrap_or(u64::MAX)); + if file_count > MAX_ARCHIVE_FILES || extracted_size > MAX_EXTRACTED_SIZE { + return Err(invalid("release archive expands beyond safety limits")); + } + + if path.file_name().is_some_and(|name| name == APP_BINARY_NAME) { + if found_binary { + return Err(invalid("release archive contains duplicate binaries")); + } + entry + .unpack(&binary_path) + .map_err(|error| internal(format!("failed to extract binary: {error}")))?; + found_binary = true; + continue; + } + if let Some(relative) = web_dist_relative_path(&path) { + if relative.as_os_str().is_empty() { + continue; + } + let target = web_dist_dir.join(relative); + if let Some(parent) = target.parent() { + fs::create_dir_all(parent).map_err(|error| { + internal(format!("failed to create web asset dir: {error}")) + })?; + } + entry + .unpack(&target) + .map_err(|error| internal(format!("failed to extract web asset: {error}")))?; + found_web = true; + } + } + if !found_binary { + return Err(invalid("release archive does not contain codex-proxy-rs")); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + fs::set_permissions(&binary_path, fs::Permissions::from_mode(0o755)) + .map_err(|error| internal(format!("failed to chmod binary: {error}")))?; + } + Ok(ExtractedRelease { + binary_path, + web_dist_dir: found_web.then_some(web_dist_dir), + }) +} + +fn unsafe_archive_path(path: &Path) -> bool { + path.is_absolute() + || path.components().any(|component| { + matches!( + component, + Component::ParentDir | Component::Prefix(_) | Component::RootDir + ) + }) +} + +fn web_dist_relative_path(path: &Path) -> Option { + let components = path + .components() + .filter_map(|component| match component { + Component::Normal(value) => Some(value.to_os_string()), + _ => None, + }) + .collect::>(); + for index in 0..components.len() { + if components[index] == "web" + && components + .get(index + 1) + .is_some_and(|value| value == "dist") + { + return Some(components[index + 2..].iter().collect()); + } + if components[index] == "dist" { + return Some(components[index + 1..].iter().collect()); + } + } + None +} diff --git a/backend/crates/gateway-host/src/system_update/download.rs b/backend/crates/gateway-host/src/system_update/download.rs new file mode 100644 index 000000000..32b512728 --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/download.rs @@ -0,0 +1,157 @@ +//! 受限制的 Release 下载、重定向信任链和 SHA-256 校验。 + +use std::fs; +use std::io::{Read, Write}; +use std::path::Path; +use std::time::Duration; + +use futures::StreamExt as _; +use sha2::Digest as _; + +use super::release::{download_client, validate_download_url}; +use super::{OperationError, UpdateEvents, invalid, upstream}; + +pub(crate) const MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024; +pub(crate) const MAX_CHECKSUM_SIZE: u64 = 1024 * 1024; + +pub(crate) async fn download_file( + url: &str, + destination: &Path, + expected_size: u64, + api_base: &str, + operation_id: &str, + events: &UpdateEvents, +) -> Result<(), OperationError> { + validate_download_url(url, api_base)?; + if expected_size == 0 || expected_size > MAX_DOWNLOAD_SIZE { + return Err(invalid("release archive size is invalid")); + } + let client = download_client(api_base, Duration::from_secs(120))?; + let response = client + .get(url) + .header(reqwest::header::ACCEPT_ENCODING, "identity") + .send() + .await + .map_err(|error| upstream(format!("release download failed: {error}")))?; + if !response.status().is_success() { + return Err(upstream(format!( + "release download failed with {}", + response.status() + ))); + } + if response + .content_length() + .is_some_and(|size| size != expected_size) + { + return Err(upstream("downloaded size does not match release metadata")); + } + let mut file = fs::File::create(destination) + .map_err(|error| super::internal(format!("failed to create download file: {error}")))?; + let mut stream = response.bytes_stream(); + let mut downloaded = 0_u64; + let mut last_percent = 0_u8; + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|error| upstream(format!("download stream failed: {error}")))?; + downloaded = downloaded.saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX)); + if downloaded > expected_size { + return Err(invalid("download exceeds declared release size")); + } + file.write_all(&chunk) + .map_err(|error| super::internal(format!("failed to write download: {error}")))?; + let percent = download_percent(downloaded, expected_size); + if percent > last_percent { + last_percent = percent; + events.emit_progress(operation_id, "downloading release archive", percent); + } + } + file.sync_all() + .map_err(|error| super::internal(format!("failed to sync download: {error}")))?; + if downloaded != expected_size { + return Err(upstream("downloaded size does not match release metadata")); + } + Ok(()) +} + +pub(crate) async fn verify_checksum( + file_path: &Path, + file_name: &str, + checksum_url: &str, + checksum_size: u64, + api_base: &str, +) -> Result<(), OperationError> { + validate_download_url(checksum_url, api_base)?; + if checksum_size == 0 || checksum_size > MAX_CHECKSUM_SIZE { + return Err(invalid("release checksum size is invalid")); + } + let client = download_client(api_base, Duration::from_secs(30))?; + let response = client + .get(checksum_url) + .header(reqwest::header::ACCEPT_ENCODING, "identity") + .send() + .await + .map_err(|error| upstream(format!("checksum download failed: {error}")))?; + if !response.status().is_success() { + return Err(upstream(format!( + "checksum download failed with {}", + response.status() + ))); + } + let bytes = response + .bytes() + .await + .map_err(|error| upstream(format!("checksum read failed: {error}")))?; + if bytes.len() > MAX_CHECKSUM_SIZE as usize || bytes.len() as u64 != checksum_size { + return Err(invalid("checksum document size is invalid")); + } + let body = + std::str::from_utf8(&bytes).map_err(|_| invalid("checksum document is not UTF-8"))?; + let expected = checksum_for(body, file_name) + .ok_or_else(|| upstream("release checksum entry is missing"))?; + let actual = sha256_file(file_path)?; + if !expected.eq_ignore_ascii_case(&actual) { + return Err(upstream("release checksum mismatch")); + } + Ok(()) +} + +fn checksum_for(body: &str, file_name: &str) -> Option { + body.lines().find_map(|line| { + let mut fields = line.split_whitespace(); + let hash = fields.next()?; + let name = fields.next()?.trim_start_matches('*'); + (hash.len() == 64 + && hash.bytes().all(|byte| byte.is_ascii_hexdigit()) + && Path::new(name).file_name()?.to_string_lossy() == file_name) + .then(|| hash.to_owned()) + }) +} + +fn sha256_file(path: &Path) -> Result { + let mut file = fs::File::open(path) + .map_err(|error| super::internal(format!("failed to open checksum file: {error}")))?; + let mut hasher = sha2::Sha256::new(); + let mut buffer = [0_u8; 64 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|error| super::internal(format!("failed to read checksum file: {error}")))?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + Ok(hex::encode(hasher.finalize())) +} + +fn download_percent(downloaded: u64, total: u64) -> u8 { + if total == 0 { + return 0; + } + u8::try_from( + downloaded + .saturating_mul(100) + .saturating_div(total) + .min(100), + ) + .unwrap_or(100) +} diff --git a/backend/crates/gateway-host/src/system_update/mod.rs b/backend/crates/gateway-host/src/system_update/mod.rs new file mode 100644 index 000000000..134e23ffb --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/mod.rs @@ -0,0 +1,617 @@ +//! 版本、安全自更新、回滚与进程重启的 Host-owned 实现。 + +mod archive; +mod download; +mod process; +mod release; +mod state; +mod swap; + +use std::env; +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use chrono::Utc; +use gateway_admin::model::system::{ + SystemOperationAccepted, SystemOperationKind, SystemUpdateDetail, SystemUpdateEvent, + SystemUpdateEventLevel, SystemUpdateStatus, SystemVersion, +}; +use gateway_admin::ports::system::{ + SystemOperationError, SystemOperationErrorKind, SystemOperations, SystemUpdateEventStream, +}; +use gateway_core::engine::CancellationToken; +use serde::Deserialize; +use tokio::sync::{Mutex as AsyncMutex, broadcast}; + +use crate::config::ConfigError; + +use self::archive::extract_release; +use self::download::{MAX_CHECKSUM_SIZE, MAX_DOWNLOAD_SIZE, download_file, verify_checksum}; +use self::process::{environment_value, spawn_replacement}; +use self::release::{ + ReleaseCache, confirmed_target, detail_from_release, fetch_latest, select_archive, +}; +use self::state::{ + OperationFileLock, UpdateTempDir, finish, operation_id, read_status, set_running, +}; +use self::swap::{replace_release_files, rollback_release}; + +const APP_BINARY_NAME: &str = "codex-proxy-rs"; +const DEFAULT_WEB_DIST_DIR: &str = "/app/web/dist"; +const DEFAULT_GITHUB_API_BASE: &str = "https://api.github.com/repos"; + +type OperationError = SystemOperationError; + +/// 系统更新与重启配置；所有字段只由 Host 解释。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +pub struct SystemUpdateConfig { + pub version: String, + pub git_sha: String, + pub build_time: String, + pub deployment_mode: String, + pub build_type: String, + pub update_channel: String, + pub update_repository: Option, + pub github_api_base: String, + pub executable_path: Option, + pub web_dist_dir: PathBuf, + pub update_state_file: PathBuf, + pub update_lock_file: PathBuf, + pub update_temp_dir: PathBuf, + pub self_restart_enabled: bool, +} + +impl Default for SystemUpdateConfig { + fn default() -> Self { + let deployment_mode = + environment_value("CPR_DEPLOYMENT_MODE").unwrap_or_else(|| "source".to_owned()); + let executable_path = environment_value("CPR_UPDATE_EXE_PATH") + .map(PathBuf::from) + .or_else(|| { + (deployment_mode == "docker") + .then(|| PathBuf::from("/app/bin").join(APP_BINARY_NAME)) + }); + let update_state_file = environment_value("CPR_UPDATE_STATE_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/app/.runtime/data/update-state.json")); + let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE") + .map(PathBuf::from) + .unwrap_or_else(|| update_state_file.with_extension("lock")); + let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| state::default_temp_dir(&update_state_file)); + Self { + version: option_env!("CPR_VERSION") + .unwrap_or(env!("CARGO_PKG_VERSION")) + .to_owned(), + git_sha: option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_owned(), + build_time: option_env!("CPR_BUILD_TIME") + .unwrap_or("unknown") + .to_owned(), + deployment_mode, + build_type: option_env!("CPR_BUILD_TYPE").unwrap_or("source").to_owned(), + update_channel: environment_value("CPR_UPDATE_CHANNEL") + .unwrap_or_else(|| "stable".to_owned()), + update_repository: environment_value("CPR_UPDATE_REPOSITORY"), + github_api_base: environment_value("CPR_GITHUB_API_BASE") + .unwrap_or_else(|| DEFAULT_GITHUB_API_BASE.to_owned()), + executable_path, + web_dist_dir: environment_value("CPR_WEB_DIST_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(DEFAULT_WEB_DIST_DIR)), + update_state_file, + update_lock_file, + update_temp_dir, + self_restart_enabled: environment_value("CPR_ENABLE_SELF_RESTART").as_deref() + == Some("true"), + } + } +} + +impl SystemUpdateConfig { + pub(crate) fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError> { + for path in [ + self.executable_path.as_mut(), + Some(&mut self.web_dist_dir), + Some(&mut self.update_state_file), + Some(&mut self.update_lock_file), + Some(&mut self.update_temp_dir), + ] + .into_iter() + .flatten() + { + if path.is_relative() { + *path = source_dir.join(&*path); + } + } + if self.version.trim().is_empty() { + return Err(ConfigError::InvalidField("host.system_update.version")); + } + if !matches!( + self.deployment_mode.as_str(), + "source" | "binary" | "docker" + ) { + return Err(ConfigError::InvalidField( + "host.system_update.deployment_mode", + )); + } + if !matches!(self.update_channel.as_str(), "stable" | "preview") { + return Err(ConfigError::InvalidField( + "host.system_update.update_channel", + )); + } + Ok(()) + } + + fn update_support_error(&self) -> Option { + if self.build_type != "release" { + return Some("一键更新需要正式构建包".to_owned()); + } + let Some(repository) = self.update_repository.as_deref() else { + return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_owned()); + }; + if let Err(error) = release::validate_repository(repository) { + return Some(error.to_string()); + } + if let Err(error) = release::validate_api_base(&self.github_api_base) { + return Some(error); + } + None + } + + fn release_cache_key(&self) -> String { + format!( + "{}|{}|{}|{}|{}|{}", + self.update_repository.as_deref().unwrap_or_default(), + self.github_api_base, + self.version, + self.deployment_mode, + self.build_type, + self.update_channel, + ) + } + + pub(crate) fn executable_path(&self) -> Result { + if let Some(path) = &self.executable_path { + return Ok(path.clone()); + } + env::current_exe() + .and_then(fs::canonicalize) + .map_err(|error| internal(format!("failed to resolve executable: {error}"))) + } +} + +/// gateway-admin 消费的真实进程/文件系统操作实现。 +pub struct ProcessSystemOperations { + cancellation: CancellationToken, + events: UpdateEvents, + config: SystemUpdateConfig, + operation_lock: AsyncMutex<()>, + release_cache: ReleaseCache, +} + +impl ProcessSystemOperations { + #[must_use] + pub fn new(cancellation: CancellationToken, config: SystemUpdateConfig) -> Self { + Self { + cancellation, + events: UpdateEvents::default(), + config, + operation_lock: AsyncMutex::const_new(()), + release_cache: ReleaseCache::default(), + } + } + + async fn perform_update_inner( + &self, + target_version: Option, + ) -> Result { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if let Some(reason) = self.config.update_support_error() { + self.events + .error_terminal(None, Some("preflight"), reason.clone()); + return Err(conflict(reason)); + } + let target = confirmed_target(target_version)?; + let repository = self + .config + .update_repository + .as_deref() + .ok_or_else(|| conflict("update repository is not configured"))?; + self.events + .info(None, Some("release"), "fetching latest release"); + let release = fetch_latest(&self.config.github_api_base, repository).await?; + let detail = detail_from_release(&self.config, &release); + if detail.latest_version != target { + self.events.warning_terminal( + None, + Some("release"), + "remote latest version changed; confirm again", + ); + return Err(conflict("remote latest version changed")); + } + if !detail.has_update { + self.events + .warning_terminal(None, Some("release"), "already up to date"); + return Err(conflict("already up to date")); + } + + let operation_id = operation_id("update"); + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + set_running( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Update, + Some(&target), + &self.config.version, + )?; + let result = self.install_release(&release, &target, &operation_id).await; + match &result { + Ok(()) => self.events.success_terminal( + Some(&operation_id), + Some("done"), + "release files replaced", + ), + Err(error) => { + self.events + .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) + } + } + finish( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Update, + result.as_ref().ok().map(|()| target.clone()), + result.as_ref().err().map(ToString::to_string), + ); + drop(file_lock); + result?; + Ok(SystemOperationAccepted::Update { + operation_id, + deployment_mode: self.config.deployment_mode.clone(), + message: "更新完成，请重启服务。".to_owned(), + need_restart: true, + target_version: target, + }) + } + + async fn install_release( + &self, + release: &release::GitHubRelease, + version: &str, + operation_id: &str, + ) -> Result<(), OperationError> { + let archive = select_archive(release, version)?; + if archive.size == 0 || archive.size > MAX_DOWNLOAD_SIZE { + return Err(invalid("release archive size is invalid")); + } + let checksum = release + .assets + .iter() + .find(|asset| asset.name == "checksums.txt") + .ok_or_else(|| upstream("release checksums.txt is required"))?; + if checksum.size == 0 || checksum.size > MAX_CHECKSUM_SIZE { + return Err(invalid("release checksum size is invalid")); + } + fs::create_dir_all(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to prepare update temp dir: {error}")))?; + let temp_root = fs::canonicalize(&self.config.update_temp_dir) + .map_err(|error| internal(format!("failed to resolve update temp dir: {error}")))?; + let temp = UpdateTempDir::create(&temp_root)?; + let archive_path = temp.path().join(&archive.name); + download_file( + &archive.browser_download_url, + &archive_path, + archive.size, + &self.config.github_api_base, + operation_id, + &self.events, + ) + .await?; + verify_checksum( + &archive_path, + &archive.name, + &checksum.browser_download_url, + checksum.size, + &self.config.github_api_base, + ) + .await?; + let extracted = extract_release(&archive_path, temp.path())?; + replace_release_files( + &self.config.executable_path()?, + &self.config.web_dist_dir, + extracted, + ) + } +} + +#[async_trait] +impl SystemOperations for ProcessSystemOperations { + async fn version(&self) -> Result { + let detail = self.update_detail(false).await?; + Ok(SystemVersion { + version: self.config.version.clone(), + git_sha: self.config.git_sha.clone(), + build_time: self.config.build_time.clone(), + deployment_mode: self.config.deployment_mode.clone(), + update_channel: self.config.update_channel.clone(), + latest_version: detail.latest_version, + has_update: detail.has_update, + update_cached: detail.cached, + update_warning: detail.warning, + }) + } + + async fn update_detail(&self, refresh: bool) -> Result { + if let Some(reason) = self.config.update_support_error() { + return Ok(base_update_detail(&self.config, Some(reason), None)); + } + match self.release_cache.detail(&self.config, refresh).await { + Ok(detail) => Ok(detail), + Err(error) => Ok(base_update_detail( + &self.config, + self.config.update_support_error(), + Some(error.to_string()), + )), + } + } + + fn update_events(&self) -> SystemUpdateEventStream { + let receiver = self.events.subscribe(); + Box::pin(futures::stream::unfold( + (receiver, false), + |(mut receiver, close)| async move { + if close { + return None; + } + loop { + match receiver.recv().await { + Ok((event, terminal)) => return Some((event, (receiver, terminal))), + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => return None, + } + } + }, + )) + } + + async fn perform_update( + &self, + target_version: Option, + ) -> Result { + self.perform_update_inner(target_version).await + } + + async fn update_status(&self) -> Result { + read_status(&self.config.update_state_file) + } + + async fn rollback(&self) -> Result { + let _operation = self + .operation_lock + .try_lock() + .map_err(|_| conflict("system operation is already running"))?; + if let Some(reason) = self.config.update_support_error() { + return Err(conflict(reason)); + } + let operation_id = operation_id("rollback"); + let file_lock = OperationFileLock::acquire(&self.config.update_lock_file)?; + set_running( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + &self.config.version, + )?; + let result = rollback_release(&self.config); + match &result { + Ok(()) => self.events.success_terminal( + Some(&operation_id), + Some("done"), + "previous release restored", + ), + Err(error) => { + self.events + .error_terminal(Some(&operation_id), Some("failed"), error.to_string()) + } + } + finish( + &self.config.update_state_file, + &operation_id, + SystemOperationKind::Rollback, + None, + result.as_ref().err().map(ToString::to_string), + ); + drop(file_lock); + result?; + Ok(SystemOperationAccepted::Rollback { + operation_id, + message: "回滚完成，请重启服务。".to_owned(), + need_restart: true, + }) + } + + async fn restart(&self) -> Result { + if !self.config.self_restart_enabled { + return Err(conflict("self restart is disabled")); + } + let message = if self.config.deployment_mode == "docker" { + "已安排进程内重启" + } else { + spawn_replacement(&self.config)?; + "已安排自重启" + }; + let operation_id = operation_id("restart"); + let cancellation = self.cancellation.clone(); + drop(tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(500)).await; + cancellation.cancel(); + })); + Ok(SystemOperationAccepted::Restart { + operation_id, + message: message.to_owned(), + }) + } +} + +fn base_update_detail( + config: &SystemUpdateConfig, + unsupported_reason: Option, + warning: Option, +) -> SystemUpdateDetail { + SystemUpdateDetail { + current_version: config.version.clone(), + latest_version: config.version.clone(), + has_update: false, + deployment_mode: config.deployment_mode.clone(), + build_type: config.build_type.clone(), + release_url: None, + notes: None, + cached: false, + update_supported: unsupported_reason.is_none() && warning.is_none(), + unsupported_reason, + warning, + } +} + +pub(crate) struct UpdateEvents { + sender: broadcast::Sender<(SystemUpdateEvent, bool)>, + sequence: AtomicU64, +} + +impl Default for UpdateEvents { + fn default() -> Self { + let (sender, _) = broadcast::channel(256); + Self { + sender, + sequence: AtomicU64::new(0), + } + } +} + +impl UpdateEvents { + fn subscribe(&self) -> broadcast::Receiver<(SystemUpdateEvent, bool)> { + self.sender.subscribe() + } + + fn info(&self, operation_id: Option<&str>, step: Option<&str>, message: impl Into) { + self.emit( + SystemUpdateEventLevel::Info, + operation_id, + step, + message, + None, + false, + ); + } + + fn warning_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Warning, + operation_id, + step, + message, + None, + true, + ); + } + + fn success_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Success, + operation_id, + step, + message, + Some(100), + true, + ); + } + + fn error_terminal( + &self, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + ) { + self.emit( + SystemUpdateEventLevel::Error, + operation_id, + step, + message, + None, + true, + ); + } + + pub(crate) fn emit_progress(&self, operation_id: &str, message: &str, progress: u8) { + self.emit( + SystemUpdateEventLevel::Info, + Some(operation_id), + Some("download"), + message, + Some(progress), + false, + ); + } + + fn emit( + &self, + level: SystemUpdateEventLevel, + operation_id: Option<&str>, + step: Option<&str>, + message: impl Into, + progress_percent: Option, + terminal: bool, + ) { + let now = Utc::now(); + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); + let event = SystemUpdateEvent { + id: format!( + "update-log-{}-{sequence}", + now.timestamp_nanos_opt() + .unwrap_or_else(|| now.timestamp_millis()) + ), + operation_id: operation_id.map(str::to_owned), + level, + step: step.map(str::to_owned), + message: message.into(), + terminal, + progress_percent, + occurred_at: now, + }; + let _ = self.sender.send((event, terminal)); + } +} + +fn invalid(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Invalid, message) +} + +fn conflict(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Conflict, message) +} + +fn upstream(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Upstream, message) +} + +fn internal(message: impl Into) -> OperationError { + SystemOperationError::new(SystemOperationErrorKind::Internal, message) +} diff --git a/backend/crates/gateway-host/src/system_update/process.rs b/backend/crates/gateway-host/src/system_update/process.rs new file mode 100644 index 000000000..475cd4a36 --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/process.rs @@ -0,0 +1,68 @@ +//! 宿主进程替换与容器内关闭语义。 + +use std::env; +use std::fs; +use std::process::{Command, Stdio}; + +use super::{OperationError, SystemUpdateConfig, internal}; + +const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS"; +const DEFAULT_RESTART_DELAY_MS: u64 = 1_200; + +pub(crate) fn spawn_replacement(config: &SystemUpdateConfig) -> Result<(), OperationError> { + let executable = config.executable_path()?; + let metadata = fs::metadata(&executable) + .map_err(|error| internal(format!("failed to schedule replacement process: {error}")))?; + if !metadata.is_file() { + return Err(internal( + "failed to schedule replacement process: executable is not a file", + )); + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + if metadata.permissions().mode() & 0o111 == 0 { + return Err(internal( + "failed to schedule replacement process: executable is not executable", + )); + } + } + + let delay_ms = environment_value(RESTART_DELAY_ENV) + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_RESTART_DELAY_MS); + #[cfg(unix)] + let mut command = { + let delay = format!("{}.{:03}", delay_ms / 1_000, delay_ms % 1_000); + let mut command = Command::new("/bin/sh"); + command + .arg("-c") + .arg("sleep \"$1\"; shift; exec \"$@\"") + .arg("codex-proxy-rs-restart") + .arg(delay) + .arg(&executable) + .args(env::args_os().skip(1)); + command + }; + #[cfg(not(unix))] + let mut command = { + let mut command = Command::new(&executable); + command + .args(env::args_os().skip(1)) + .env(RESTART_DELAY_ENV, delay_ms.to_string()); + command + }; + command.stdin(Stdio::null()); + command + .spawn() + .map(|_| ()) + .map_err(|error| internal(format!("failed to schedule replacement process: {error}"))) +} + +pub(crate) fn environment_value(key: &str) -> Option { + env::var(key) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} diff --git a/backend/crates/gateway-host/src/system_update/release.rs b/backend/crates/gateway-host/src/system_update/release.rs new file mode 100644 index 000000000..9c701c7dd --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/release.rs @@ -0,0 +1,326 @@ +//! GitHub Release 发现、缓存、版本比较与下载信任边界。 + +use std::env; +use std::time::{Duration, Instant}; + +use gateway_admin::model::system::SystemUpdateDetail; +use serde::Deserialize; +use tokio::sync::Mutex; + +use super::{OperationError, SystemUpdateConfig, conflict, invalid, upstream}; + +const APP_BINARY_NAME: &str = "codex-proxy-rs"; +const CACHE_TTL: Duration = Duration::from_secs(20 * 60); + +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct GitHubRelease { + pub(crate) tag_name: String, + pub(crate) name: Option, + pub(crate) body: Option, + pub(crate) html_url: Option, + pub(crate) prerelease: bool, + #[serde(default)] + pub(crate) assets: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct GitHubAsset { + pub(crate) name: String, + pub(crate) browser_download_url: String, + pub(crate) size: u64, +} + +struct CachedRelease { + key: String, + detail: SystemUpdateDetail, + cached_at: Instant, +} + +#[derive(Default)] +pub(crate) struct ReleaseCache { + entry: Mutex>, +} + +impl ReleaseCache { + pub(crate) async fn detail( + &self, + config: &SystemUpdateConfig, + refresh: bool, + ) -> Result { + let repository = config + .update_repository + .as_deref() + .ok_or_else(|| conflict("update repository is not configured"))?; + validate_repository(repository)?; + validate_api_base(&config.github_api_base).map_err(conflict)?; + let key = config.release_cache_key(); + if !refresh && let Some(detail) = self.cached(&key).await { + return Ok(detail); + } + match fetch_latest(&config.github_api_base, repository).await { + Ok(release) => { + let detail = detail_from_release(config, &release); + *self.entry.lock().await = Some(CachedRelease { + key, + detail: detail.clone(), + cached_at: Instant::now(), + }); + Ok(detail) + } + Err(error) => self.cached(&key).await.ok_or(error), + } + } + + async fn cached(&self, key: &str) -> Option { + let entry = self.entry.lock().await; + let cached = entry.as_ref()?; + (cached.key == key && cached.cached_at.elapsed() <= CACHE_TTL).then(|| { + let mut detail = cached.detail.clone(); + detail.cached = true; + detail + }) + } +} + +pub(crate) async fn fetch_latest( + api_base: &str, + repository: &str, +) -> Result { + validate_api_base(api_base).map_err(conflict)?; + validate_repository(repository)?; + let url = format!( + "{}/{repository}/releases/latest", + api_base.trim_end_matches('/') + ); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|error| upstream(format!("failed to create release client: {error}")))?; + let response = client + .get(url) + .header(reqwest::header::ACCEPT, "application/vnd.github+json") + .header(reqwest::header::USER_AGENT, APP_BINARY_NAME) + .send() + .await + .map_err(|error| upstream(format!("GitHub release check failed: {error}")))?; + if !response.status().is_success() { + return Err(upstream(format!( + "GitHub release check failed with {}", + response.status() + ))); + } + response + .json::() + .await + .map_err(|error| upstream(format!("invalid GitHub release response: {error}"))) +} + +pub(crate) fn detail_from_release( + config: &SystemUpdateConfig, + release: &GitHubRelease, +) -> SystemUpdateDetail { + let latest = normalize_version(&release.tag_name); + let available = (config.update_channel != "stable" || !release.prerelease) + && version_is_newer(&config.version, &latest).unwrap_or(false); + SystemUpdateDetail { + current_version: config.version.clone(), + latest_version: latest, + has_update: available, + deployment_mode: config.deployment_mode.clone(), + build_type: config.build_type.clone(), + release_url: release.html_url.clone(), + notes: release.body.clone().or_else(|| release.name.clone()), + cached: false, + update_supported: config.update_support_error().is_none(), + unsupported_reason: config.update_support_error(), + warning: None, + } +} + +pub(crate) fn confirmed_target(target: Option) -> Result { + let target = target.ok_or_else(|| conflict("target version must be confirmed"))?; + let target = normalize_version(&target); + if target.is_empty() || semver::Version::parse(&target).is_err() { + return Err(invalid("target version is invalid")); + } + Ok(target) +} + +pub(crate) fn select_archive<'a>( + release: &'a GitHubRelease, + version: &str, +) -> Result<&'a GitHubAsset, OperationError> { + let os = platform_os_aliases(); + let arch = platform_arch_aliases(); + let normalized = normalize_version(version); + release + .assets + .iter() + .find(|asset| { + asset.name.contains(APP_BINARY_NAME) + && asset.name.contains(&normalized) + && asset_matches_platform(&asset.name, os, arch) + && !asset.name.ends_with(".txt") + }) + .or_else(|| { + release.assets.iter().find(|asset| { + asset_matches_platform(&asset.name, os, arch) && !asset.name.ends_with(".txt") + }) + }) + .ok_or_else(|| { + conflict(format!( + "no compatible release archive for {}/{}", + env::consts::OS, + env::consts::ARCH + )) + }) +} + +pub(crate) fn validate_download_url(raw: &str, api_base: &str) -> Result<(), OperationError> { + let url = reqwest::Url::parse(raw) + .map_err(|error| invalid(format!("invalid download URL: {error}")))?; + if local_download_allowed(&url, api_base) { + return Ok(()); + } + if url.scheme() != "https" { + return Err(invalid("only HTTPS release downloads are allowed")); + } + let host = url + .host_str() + .ok_or_else(|| invalid("download URL is missing host"))?; + if github_download_host_allowed(host) { + Ok(()) + } else { + Err(invalid("release download host is not trusted")) + } +} + +pub(crate) fn download_client( + api_base: &str, + timeout: Duration, +) -> Result { + let api_base = api_base.to_owned(); + reqwest::Client::builder() + .timeout(timeout) + .redirect(reqwest::redirect::Policy::custom(move |attempt| { + if download_url_allowed(attempt.url(), &api_base) { + attempt.follow() + } else { + attempt.error(std::io::Error::other( + "release redirect target is not trusted", + )) + } + })) + .build() + .map_err(|error| upstream(format!("failed to create download client: {error}"))) +} + +pub(crate) fn validate_repository(repository: &str) -> Result<(), OperationError> { + let mut segments = repository.split('/'); + let owner = segments.next().unwrap_or_default(); + let name = segments.next().unwrap_or_default(); + if owner.is_empty() + || name.is_empty() + || segments.next().is_some() + || !owner.chars().all(repository_character) + || !name.chars().all(repository_character) + { + return Err(conflict("update repository must use owner/repository")); + } + Ok(()) +} + +pub(crate) fn validate_api_base(raw: &str) -> Result<(), String> { + let url = reqwest::Url::parse(raw).map_err(|error| format!("invalid API base: {error}"))?; + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err("API base must not include credentials, query, or fragment".to_owned()); + } + if url.path().trim_end_matches('/') != "/repos" { + return Err("API base path must be /repos".to_owned()); + } + if url_host_is_loopback(&url) { + return matches!(url.scheme(), "http" | "https") + .then_some(()) + .ok_or_else(|| "loopback API base must use HTTP or HTTPS".to_owned()); + } + if url.scheme() != "https" || url.host_str() != Some("api.github.com") { + return Err("API base must be https://api.github.com/repos".to_owned()); + } + Ok(()) +} + +fn normalize_version(version: &str) -> String { + version.trim().trim_start_matches('v').to_owned() +} + +fn version_is_newer(current: &str, latest: &str) -> Option { + let current = semver::Version::parse(&normalize_version(current)).ok()?; + let latest = semver::Version::parse(&normalize_version(latest)).ok()?; + Some(latest > current) +} + +fn repository_character(character: char) -> bool { + character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.') +} + +fn platform_os_aliases() -> &'static [&'static str] { + match env::consts::OS { + "macos" => &["macos", "darwin"], + "windows" => &["windows", "win32"], + "linux" => &["linux"], + _ => &[env::consts::OS], + } +} + +fn platform_arch_aliases() -> &'static [&'static str] { + match env::consts::ARCH { + "x86_64" => &["x86_64", "amd64"], + "aarch64" => &["aarch64", "arm64"], + _ => &[env::consts::ARCH], + } +} + +fn asset_matches_platform(name: &str, os: &[&str], arch: &[&str]) -> bool { + os.iter().any(|alias| name.contains(alias)) && arch.iter().any(|alias| name.contains(alias)) +} + +fn github_download_host_allowed(host: &str) -> bool { + host == "github.com" + || host.ends_with(".github.com") + || host == "objects.githubusercontent.com" + || host.ends_with(".objects.githubusercontent.com") +} + +fn local_download_allowed(url: &reqwest::Url, api_base: &str) -> bool { + if !matches!(url.scheme(), "http" | "https") || !url_host_is_loopback(url) { + return false; + } + let Ok(base) = reqwest::Url::parse(api_base) else { + return false; + }; + url_host_is_loopback(&base) + && url.scheme() == base.scheme() + && url.host_str() == base.host_str() + && url.port_or_known_default() == base.port_or_known_default() +} + +fn download_url_allowed(url: &reqwest::Url, api_base: &str) -> bool { + local_download_allowed(url, api_base) + || (url.scheme() == "https" && url.host_str().is_some_and(github_download_host_allowed)) +} + +fn url_host_is_loopback(url: &reqwest::Url) -> bool { + url.host_str().is_some_and(|host| { + host == "localhost" + || host == "127.0.0.1" + || host == "::1" + || host + .parse::() + .is_ok_and(|address| address.is_loopback()) + }) +} diff --git a/backend/crates/gateway-host/src/system_update/state.rs b/backend/crates/gateway-host/src/system_update/state.rs new file mode 100644 index 000000000..b9cb9a876 --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/state.rs @@ -0,0 +1,347 @@ +//! 系统操作状态文件、跨进程锁与临时目录。 + +use std::fs; +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use gateway_admin::model::system::{ + SystemOperationKind, SystemOperationState, SystemOperationStatus, SystemUpdateStatus, +}; +use serde::{Deserialize, Serialize}; + +use super::{OperationError, conflict, internal}; + +#[derive(Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct PersistedState { + previous_version: Option, + current_version: Option, + #[serde(default)] + operation: PersistedOperation, +} + +#[derive(Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct PersistedOperation { + operation_id: Option, + kind: Option, + #[serde(default)] + status: PersistedStatus, + target_version: Option, + message: Option, + error: Option, + started_at: Option, + finished_at: Option, +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +enum PersistedKind { + Update, + Rollback, + Restart, +} + +#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +enum PersistedStatus { + #[default] + Idle, + Running, + Succeeded, + Failed, +} + +pub(crate) struct OperationFileLock { + path: PathBuf, +} + +impl OperationFileLock { + pub(crate) fn acquire(path: &Path) -> Result { + if let Some(parent) = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent).map_err(|error| { + internal(format!("failed to create update lock directory: {error}")) + })?; + } + match Self::try_create(path) { + Ok(()) => Ok(Self { path: path.into() }), + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => { + if !stale_lock(path)? { + return Err(conflict("system update already running")); + } + fs::remove_file(path).map_err(|error| { + internal(format!("failed to remove stale update lock: {error}")) + })?; + Self::try_create(path).map_err(|error| { + if error.kind() == io::ErrorKind::AlreadyExists { + conflict("system update already running") + } else { + internal(format!("failed to create update lock: {error}")) + } + })?; + Ok(Self { path: path.into() }) + } + Err(error) => Err(internal(format!("failed to create update lock: {error}"))), + } + } + + fn try_create(path: &Path) -> io::Result<()> { + let mut file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(path)?; + writeln!( + file, + "pid={}\ncreated_at={}", + std::process::id(), + Utc::now().to_rfc3339() + )?; + file.sync_all() + } +} + +impl Drop for OperationFileLock { + fn drop(&mut self) { + if let Err(error) = fs::remove_file(&self.path) + && error.kind() != io::ErrorKind::NotFound + { + tracing::warn!(path = %self.path.display(), error = %error, "清理系统更新锁失败"); + } + } +} + +pub(crate) struct UpdateTempDir { + path: PathBuf, +} + +impl UpdateTempDir { + pub(crate) fn create(parent: &Path) -> Result { + for attempt in 0..100_u8 { + let path = parent.join(format!( + ".codex-proxy-rs-update-{}-{attempt}", + Utc::now().timestamp_nanos_opt().unwrap_or_default() + )); + match fs::create_dir(&path) { + Ok(()) => return Ok(Self { path }), + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue, + Err(error) => { + return Err(internal(format!( + "failed to create update temp directory: {error}" + ))); + } + } + } + Err(internal("failed to create unique update temp directory")) + } + + pub(crate) fn path(&self) -> &Path { + &self.path + } +} + +impl Drop for UpdateTempDir { + fn drop(&mut self) { + if let Err(error) = fs::remove_dir_all(&self.path) + && error.kind() != io::ErrorKind::NotFound + { + tracing::warn!(path = %self.path.display(), error = %error, "清理系统更新临时目录失败"); + } + } +} + +pub(crate) fn operation_id(kind: &str) -> String { + format!("sysop-{kind}-{}", Utc::now().timestamp_millis()) +} + +pub(crate) fn set_running( + path: &Path, + operation_id: &str, + kind: SystemOperationKind, + target_version: Option<&str>, + current_version: &str, +) -> Result<(), OperationError> { + let mut state = read_persisted(path)?; + if state.current_version.is_none() { + state.current_version = Some(current_version.to_owned()); + } + state.operation = PersistedOperation { + operation_id: Some(operation_id.to_owned()), + kind: Some(kind.into()), + status: PersistedStatus::Running, + target_version: target_version.map(ToOwned::to_owned), + message: Some("operation running".to_owned()), + error: None, + started_at: Some(Utc::now().to_rfc3339()), + finished_at: None, + }; + write_persisted(path, &state) +} + +pub(crate) fn finish( + path: &Path, + operation_id: &str, + kind: SystemOperationKind, + version: Option, + error: Option, +) { + let mut state = match read_persisted(path) { + Ok(state) => state, + Err(error) => { + tracing::warn!(error = %error, "读取系统更新状态失败"); + return; + } + }; + if state.operation.operation_id.as_deref() != Some(operation_id) { + return; + } + if let Some(error) = error { + state.operation.status = PersistedStatus::Failed; + state.operation.message = Some(error.clone()); + state.operation.error = Some(error); + } else { + state.operation.status = PersistedStatus::Succeeded; + state.operation.message = Some("operation succeeded".to_owned()); + state.operation.error = None; + match kind { + SystemOperationKind::Update => { + state.previous_version = state.current_version.take(); + state.current_version = version.clone(); + } + SystemOperationKind::Rollback => { + let current = state.current_version.take(); + state.current_version = state.previous_version.take(); + state.previous_version = current; + } + SystemOperationKind::Restart => {} + } + state.operation.target_version = version; + } + state.operation.finished_at = Some(Utc::now().to_rfc3339()); + if let Err(error) = write_persisted(path, &state) { + tracing::warn!(error = %error, "写入系统更新状态失败"); + } +} + +pub(crate) fn read_status(path: &Path) -> Result { + let state = read_persisted(path)?; + let operation = state.operation; + Ok(SystemUpdateStatus { + previous_version: state.previous_version, + current_version: state.current_version, + operation: SystemOperationState { + operation_id: operation.operation_id, + kind: operation.kind.map(Into::into), + status: operation.status.into(), + target_version: operation.target_version, + message: operation.message, + error: operation.error, + started_at: parse_time(operation.started_at), + finished_at: parse_time(operation.finished_at), + }, + }) +} + +pub(crate) fn default_temp_dir(state_file: &Path) -> PathBuf { + state_file + .parent() + .map(|parent| parent.join("update-tmp")) + .unwrap_or_else(|| std::env::temp_dir().join("codex-proxy-rs-update")) +} + +fn stale_lock(path: &Path) -> Result { + let modified = fs::metadata(path) + .and_then(|metadata| metadata.modified()) + .map_err(|error| internal(format!("failed to read update lock timestamp: {error}")))?; + modified + .elapsed() + .map(|age| age > Duration::from_secs(30 * 60)) + .map_err(|error| internal(format!("failed to calculate update lock age: {error}"))) +} + +fn read_persisted(path: &Path) -> Result { + if !path.exists() { + return Ok(PersistedState::default()); + } + let data = fs::read_to_string(path) + .map_err(|error| internal(format!("failed to read update state: {error}")))?; + serde_json::from_str(&data).map_err(|error| internal(format!("invalid update state: {error}"))) +} + +fn write_persisted(path: &Path, state: &PersistedState) -> Result<(), OperationError> { + if let Some(parent) = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent).map_err(|error| { + internal(format!("failed to create update state directory: {error}")) + })?; + } + let data = serde_json::to_vec_pretty(state) + .map_err(|error| internal(format!("failed to encode update state: {error}")))?; + let temporary = temporary_path(path); + let result = (|| { + let mut file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary)?; + file.write_all(&data)?; + file.sync_all()?; + fs::rename(&temporary, path) + })(); + if result.is_err() { + let _ = fs::remove_file(&temporary); + } + result.map_err(|error| internal(format!("failed to write update state: {error}"))) +} + +fn temporary_path(path: &Path) -> PathBuf { + let mut temporary = path.as_os_str().to_os_string(); + temporary.push(format!( + ".tmp-{}-{}", + std::process::id(), + Utc::now().timestamp_nanos_opt().unwrap_or_default() + )); + PathBuf::from(temporary) +} + +fn parse_time(value: Option) -> Option> { + value + .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) + .map(|value| value.with_timezone(&Utc)) +} + +impl From for PersistedKind { + fn from(value: SystemOperationKind) -> Self { + match value { + SystemOperationKind::Update => Self::Update, + SystemOperationKind::Rollback => Self::Rollback, + SystemOperationKind::Restart => Self::Restart, + } + } +} + +impl From for SystemOperationKind { + fn from(value: PersistedKind) -> Self { + match value { + PersistedKind::Update => Self::Update, + PersistedKind::Rollback => Self::Rollback, + PersistedKind::Restart => Self::Restart, + } + } +} + +impl From for SystemOperationStatus { + fn from(value: PersistedStatus) -> Self { + match value { + PersistedStatus::Idle => Self::Idle, + PersistedStatus::Running => Self::Running, + PersistedStatus::Succeeded => Self::Succeeded, + PersistedStatus::Failed => Self::Failed, + } + } +} diff --git a/backend/crates/gateway-host/src/system_update/swap.rs b/backend/crates/gateway-host/src/system_update/swap.rs new file mode 100644 index 000000000..c41f49df6 --- /dev/null +++ b/backend/crates/gateway-host/src/system_update/swap.rs @@ -0,0 +1,273 @@ +//! 二进制与 Web 制品的交易式替换、跨文件系统移动与回滚。 + +use std::fmt; +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use super::archive::ExtractedRelease; +use super::{OperationError, SystemUpdateConfig, conflict, internal}; + +pub(crate) fn replace_release_files( + executable: &Path, + web_dist: &Path, + extracted: ExtractedRelease, +) -> Result<(), OperationError> { + let web_backup = backup_path_for(web_dist); + let web_replaced = if let Some(new_web) = extracted.web_dist_dir { + replace_dir(web_dist, &web_backup, &new_web)?; + true + } else { + false + }; + + let binary_backup = backup_path_for(executable); + if binary_backup.exists() + && let Err(error) = fs::remove_file(&binary_backup) + { + let mut rollback = Vec::new(); + if web_replaced { + collect_rollback_error( + &mut rollback, + "restore web assets", + restore_dir(web_dist, &web_backup), + ); + } + return Err(error_with_rollback( + "failed to remove old binary backup", + error, + rollback, + )); + } + if let Err(error) = move_file(executable, &binary_backup) { + let mut rollback = Vec::new(); + if web_replaced { + collect_rollback_error( + &mut rollback, + "restore web assets", + restore_dir(web_dist, &web_backup), + ); + } + return Err(error_with_rollback("binary backup failed", error, rollback)); + } + if let Err(error) = move_file(&extracted.binary_path, executable) { + let mut rollback = Vec::new(); + if executable.exists() { + collect_rollback_error( + &mut rollback, + "remove partial replacement binary", + fs::remove_file(executable), + ); + } + collect_rollback_error( + &mut rollback, + "restore previous binary", + move_file(&binary_backup, executable), + ); + if web_replaced { + collect_rollback_error( + &mut rollback, + "restore web assets", + restore_dir(web_dist, &web_backup), + ); + } + return Err(error_with_rollback( + "binary replace failed", + error, + rollback, + )); + } + Ok(()) +} + +pub(crate) fn rollback_release(config: &SystemUpdateConfig) -> Result<(), OperationError> { + let executable = config.executable_path()?; + let binary_backup = backup_path_for(&executable); + if !binary_backup.exists() { + return Err(conflict("no binary backup found for rollback")); + } + swap_file(&executable, &binary_backup) + .map_err(|error| internal(format!("binary rollback failed: {error}")))?; + + let web_backup = backup_path_for(&config.web_dist_dir); + if web_backup.exists() + && let Err(error) = swap_dir(&config.web_dist_dir, &web_backup) + { + let mut rollback = Vec::new(); + collect_rollback_error( + &mut rollback, + "restore binary after web rollback failure", + swap_file(&executable, &binary_backup), + ); + return Err(error_with_rollback("web rollback failed", error, rollback)); + } + Ok(()) +} + +fn replace_dir(current: &Path, backup: &Path, replacement: &Path) -> Result<(), OperationError> { + if backup.exists() { + fs::remove_dir_all(backup) + .map_err(|error| internal(format!("failed to remove old web backup: {error}")))?; + } + if current.exists() { + move_dir(current, backup) + .map_err(|error| internal(format!("failed to backup web assets: {error}")))?; + } + if let Err(error) = move_dir(replacement, current) { + let mut rollback = Vec::new(); + if backup.exists() { + collect_rollback_error( + &mut rollback, + "restore previous web assets", + restore_dir(current, backup), + ); + } + return Err(error_with_rollback( + "failed to replace web assets", + error, + rollback, + )); + } + Ok(()) +} + +fn swap_file(current: &Path, backup: &Path) -> io::Result<()> { + if !current.exists() { + return Err(io::Error::new( + io::ErrorKind::NotFound, + format!("current binary is missing: {}", current.display()), + )); + } + let swap = swap_path_for(current); + if swap.exists() { + fs::remove_file(&swap)?; + } + move_file(current, &swap)?; + if let Err(error) = move_file(backup, current) { + let _ = move_file(&swap, current); + return Err(error); + } + if let Err(error) = move_file(&swap, backup) { + let _ = move_file(current, &swap); + let _ = move_file(backup, current); + let _ = move_file(&swap, backup); + return Err(error); + } + Ok(()) +} + +fn swap_dir(current: &Path, backup: &Path) -> io::Result<()> { + if !current.exists() { + return Err(io::Error::new( + io::ErrorKind::NotFound, + format!("current web directory is missing: {}", current.display()), + )); + } + let swap = swap_path_for(current); + if swap.exists() { + fs::remove_dir_all(&swap)?; + } + move_dir(current, &swap)?; + if let Err(error) = move_dir(backup, current) { + let _ = move_dir(&swap, current); + return Err(error); + } + if let Err(error) = move_dir(&swap, backup) { + let _ = move_dir(current, &swap); + let _ = move_dir(backup, current); + let _ = move_dir(&swap, backup); + return Err(error); + } + Ok(()) +} + +fn restore_dir(current: &Path, backup: &Path) -> io::Result<()> { + if current.exists() { + fs::remove_dir_all(current)?; + } + if backup.exists() { + move_dir(backup, current)?; + } + Ok(()) +} + +fn move_file(from: &Path, to: &Path) -> io::Result<()> { + match fs::rename(from, to) { + Ok(()) => Ok(()), + Err(error) if error.kind() == io::ErrorKind::CrossesDevices => { + if let Err(copy_error) = fs::copy(from, to) { + let _ = fs::remove_file(to); + return Err(copy_error); + } + fs::remove_file(from) + } + Err(error) => Err(error), + } +} + +fn move_dir(from: &Path, to: &Path) -> io::Result<()> { + match fs::rename(from, to) { + Ok(()) => Ok(()), + Err(error) if error.kind() == io::ErrorKind::CrossesDevices => { + if let Err(copy_error) = copy_dir_all(from, to) { + let _ = fs::remove_dir_all(to); + return Err(copy_error); + } + fs::remove_dir_all(from) + } + Err(error) => Err(error), + } +} + +fn copy_dir_all(from: &Path, to: &Path) -> io::Result<()> { + fs::create_dir_all(to)?; + for entry in fs::read_dir(from)? { + let entry = entry?; + let file_type = entry.file_type()?; + let target = to.join(entry.file_name()); + if file_type.is_dir() { + copy_dir_all(&entry.path(), &target)?; + } else if file_type.is_file() { + fs::copy(entry.path(), &target)?; + fs::set_permissions(&target, entry.metadata()?.permissions())?; + } else { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + format!("unsupported file type in {}", entry.path().display()), + )); + } + } + Ok(()) +} + +fn backup_path_for(path: &Path) -> PathBuf { + let mut backup = path.as_os_str().to_os_string(); + backup.push(".backup"); + PathBuf::from(backup) +} + +fn swap_path_for(path: &Path) -> PathBuf { + let mut swap = path.as_os_str().to_os_string(); + swap.push(".rollback-swap"); + PathBuf::from(swap) +} + +fn collect_rollback_error(errors: &mut Vec, action: &'static str, result: io::Result<()>) { + if let Err(error) = result { + errors.push(format!("{action}: {error}")); + } +} + +fn error_with_rollback( + context: &'static str, + error: impl fmt::Display, + rollback_errors: Vec, +) -> OperationError { + if rollback_errors.is_empty() { + return internal(format!("{context}: {error}")); + } + internal(format!( + "{context}: {error}; rollback failed: {}", + rollback_errors.join("; ") + )) +} diff --git a/backend/crates/gateway-host/src/workers/mod.rs b/backend/crates/gateway-host/src/workers/mod.rs new file mode 100644 index 000000000..081e590ec --- /dev/null +++ b/backend/crates/gateway-host/src/workers/mod.rs @@ -0,0 +1,632 @@ +//! Core 任务计划的唯一运行时：注册校验、lease、退避、健康与关闭。 + +use std::collections::{BTreeMap, BTreeSet}; +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime}; + +use futures::FutureExt as _; +use gateway_core::engine::CancellationToken; +use gateway_core::health::{ + WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState, +}; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerDisabledReason, WorkerId, WorkerKind, WorkerLeaderLeaseGuard, + WorkerLeaderLeasePort, WorkerLeaseAcquisition, WorkerLeaseRequest, WorkerRegistration, + WorkerRunnable, WorkerSchedule, WorkerTaskError, +}; +use tokio::task::JoinHandle; + +/// 进程内所有后台任务的唯一监督器。 +pub struct WorkerSupervisor { + cancellation: CancellationToken, + health: Arc, + started: AtomicBool, + handles: Mutex)>>, +} + +impl WorkerSupervisor { + #[must_use] + pub fn new(cancellation: CancellationToken) -> Self { + Self { + cancellation, + health: Arc::new(WorkerHealthRegistry::default()), + started: AtomicBool::new(false), + handles: Mutex::new(Vec::new()), + } + } + + #[must_use] + pub fn health_source(&self) -> Arc { + self.health.clone() + } + + /// 校验完整计划并且仅启动一次。 + pub fn start( + &self, + plan: Vec, + leader_lease: Arc, + ) -> Result<(), WorkerStartError> { + if self + .started + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return Err(WorkerStartError::AlreadyStarted); + } + let result = self.start_inner(plan, leader_lease); + if result.is_err() { + self.started.store(false, Ordering::Release); + } + result + } + + fn start_inner( + &self, + plan: Vec, + leader_lease: Arc, + ) -> Result<(), WorkerStartError> { + tokio::runtime::Handle::try_current().map_err(|_| WorkerStartError::RuntimeUnavailable)?; + let WorkerPlan { + registrations, + disabled, + } = WorkerPlan::try_from(plan)?; + self.health.install(®istrations, &disabled); + + let handles = registrations + .into_values() + .map(|registration| { + let id = registration.id.clone(); + let handle = match registration.runnable { + WorkerRunnable::Scheduled { + schedule, + lease, + task, + } => tokio::spawn(supervise_scheduled( + id.clone(), + task, + schedule, + lease, + Arc::clone(&leader_lease), + self.cancellation.clone(), + Arc::clone(&self.health), + )), + WorkerRunnable::Daemon { restart, task } => tokio::spawn(supervise_daemon( + id.clone(), + task, + restart, + self.cancellation.clone(), + Arc::clone(&self.health), + )), + }; + (id, handle) + }) + .collect(); + *lock_unpoisoned(&self.handles) = handles; + Ok(()) + } + + /// 先协作取消；超时后 abort 不合作任务，lease 依赖 TTL 兜底。 + pub async fn shutdown(&self, timeout: Duration) { + self.cancellation.cancel(); + let mut handles = std::mem::take(&mut *lock_unpoisoned(&self.handles)); + let joined = tokio::time::timeout(timeout, async { + for (_, handle) in &mut handles { + let _ = handle.await; + } + }) + .await + .is_ok(); + if !joined { + for (_, handle) in &handles { + handle.abort(); + } + for (_, handle) in handles { + let _ = handle.await; + } + } + self.health.stop_all(); + } +} + +impl Drop for WorkerSupervisor { + fn drop(&mut self) { + self.cancellation.cancel(); + for (_, handle) in lock_unpoisoned(&self.handles).iter() { + handle.abort(); + } + self.health.stop_all(); + } +} + +struct WorkerPlan { + registrations: BTreeMap, + disabled: BTreeMap, +} + +impl TryFrom> for WorkerPlan { + type Error = WorkerStartError; + + fn try_from(plan: Vec) -> Result { + let mut registrations = BTreeMap::new(); + let mut disabled = BTreeMap::new(); + for contribution in plan { + contribution.validate()?; + match contribution { + WorkerContribution::Registration(registration) => { + let kind = registration.id.kind(); + if disabled.contains_key(&kind) { + return Err(WorkerStartError::KindDisabled(kind)); + } + let id = registration.id.clone(); + if registrations.insert(id.clone(), registration).is_some() { + return Err(WorkerStartError::Duplicate(id)); + } + } + WorkerContribution::Disabled { kind, reason } => { + if registrations.keys().any(|id| id.kind() == kind) { + return Err(WorkerStartError::KindHasTasks(kind)); + } + if disabled.insert(kind, reason).is_some() { + return Err(WorkerStartError::DuplicateDisabled(kind)); + } + } + } + } + let active = registrations + .keys() + .map(WorkerId::kind) + .collect::>(); + let missing = WorkerKind::ALL + .into_iter() + .filter(|kind| !active.contains(kind) && !disabled.contains_key(kind)) + .collect::>(); + if !missing.is_empty() { + return Err(WorkerStartError::Missing(missing)); + } + Ok(Self { + registrations, + disabled, + }) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum WorkerStartError { + #[error("worker supervisor was already started")] + AlreadyStarted, + #[error("Tokio runtime is unavailable")] + RuntimeUnavailable, + #[error(transparent)] + InvalidDefinition(#[from] WorkerDefinitionError), + #[error("worker `{0}` was registered more than once")] + Duplicate(WorkerId), + #[error("worker kind `{0}` is disabled")] + KindDisabled(WorkerKind), + #[error("worker kind `{0}` already has tasks")] + KindHasTasks(WorkerKind), + #[error("worker kind `{0}` was disabled more than once")] + DuplicateDisabled(WorkerKind), + #[error("required workers are missing: {0:?}")] + Missing(Vec), +} + +#[derive(Clone)] +struct HealthEntry { + snapshot: WorkerHealthSnapshot, + success_freshness: Duration, +} + +#[derive(Default)] +struct WorkerHealthRegistry { + states: Mutex>, +} + +impl WorkerHealthRegistry { + fn install( + &self, + registrations: &BTreeMap, + disabled: &BTreeMap, + ) { + let mut states = registrations + .values() + .map(|registration| { + let freshness = match ®istration.runnable { + WorkerRunnable::Scheduled { schedule, .. } => (*schedule) + .interval() + .saturating_mul(2) + .saturating_add(schedule.maximum_backoff()), + WorkerRunnable::Daemon { restart, .. } => restart.maximum_backoff(), + }; + let key = WorkerHealthKey::Task(registration.id.clone()); + ( + key.clone(), + HealthEntry { + snapshot: empty_snapshot(key, WorkerRuntimeState::Starting), + success_freshness: freshness, + }, + ) + }) + .collect::>(); + states.extend(disabled.keys().map(|kind| { + let key = WorkerHealthKey::Disabled(*kind); + ( + key.clone(), + HealthEntry { + snapshot: empty_snapshot(key, WorkerRuntimeState::Disabled), + success_freshness: Duration::ZERO, + }, + ) + })); + *lock_unpoisoned(&self.states) = states; + } + + fn update(&self, id: &WorkerId, update: impl FnOnce(&mut WorkerHealthSnapshot)) { + if let Some(entry) = + lock_unpoisoned(&self.states).get_mut(&WorkerHealthKey::Task(id.clone())) + { + update(&mut entry.snapshot); + } + } + + fn acquiring(&self, id: &WorkerId) { + self.update(id, |state| { + state.state = WorkerRuntimeState::AcquiringLease; + }); + } + + fn standby(&self, id: &WorkerId) { + self.update(id, |state| { + state.state = WorkerRuntimeState::Standby; + state.consecutive_failures = 0; + state.last_error = None; + }); + } + + fn running( + &self, + id: &WorkerId, + fencing_token: Option, + ) { + self.update(id, |state| { + state.state = WorkerRuntimeState::Running; + state.last_fencing_token = fencing_token; + }); + } + + fn succeeded(&self, id: &WorkerId) { + self.update(id, |state| { + state.state = WorkerRuntimeState::Idle; + state.consecutive_failures = 0; + state.completed_cycles = state.completed_cycles.saturating_add(1); + state.last_success_at = Some(SystemTime::now()); + state.last_error = None; + }); + } + + fn failed(&self, id: &WorkerId, error: String) -> u32 { + let mut failures = 0; + self.update(id, |state| { + state.state = WorkerRuntimeState::BackingOff; + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + state.last_failure_at = Some(SystemTime::now()); + state.last_error = Some(error); + failures = state.consecutive_failures; + }); + failures + } + + fn stopped(&self, id: &WorkerId) { + self.update(id, |state| state.state = WorkerRuntimeState::Stopped); + } + + fn stop_all(&self) { + for entry in lock_unpoisoned(&self.states).values_mut() { + if entry.snapshot.state != WorkerRuntimeState::Disabled { + entry.snapshot.state = WorkerRuntimeState::Stopped; + } + } + } +} + +impl WorkerHealthSource for WorkerHealthRegistry { + fn snapshot(&self) -> Vec { + let now = SystemTime::now(); + lock_unpoisoned(&self.states) + .values() + .map(|entry| { + let mut snapshot = entry.snapshot.clone(); + let stale = matches!( + snapshot.state, + WorkerRuntimeState::Running | WorkerRuntimeState::Idle + ) && snapshot.last_success_at.is_some_and(|last_success| { + now.duration_since(last_success) + .is_ok_and(|age| age > entry.success_freshness) + }); + if stale { + snapshot.state = WorkerRuntimeState::BackingOff; + snapshot.last_error = Some("worker success is stale".to_owned()); + } + snapshot + }) + .collect() + } +} + +async fn supervise_scheduled( + id: WorkerId, + task: Box, + schedule: WorkerSchedule, + lease: Option, + leader_lease: Arc, + cancellation: CancellationToken, + health: Arc, +) { + let mut backoff = schedule.initial_backoff(); + loop { + if cancellation.is_cancelled() { + break; + } + let outcome = if let Some(request) = lease.clone() { + health.acquiring(&id); + acquire_and_run( + &id, + task.as_ref(), + schedule, + request, + leader_lease.as_ref(), + &cancellation, + &health, + ) + .await + } else { + run_local_cycle(&id, task.as_ref(), &cancellation, &health).await + }; + + let delay = match outcome { + ScheduledOutcome::Succeeded => { + health.succeeded(&id); + backoff = schedule.initial_backoff(); + schedule.interval() + } + ScheduledOutcome::Busy(retry_after) => { + health.standby(&id); + retry_after + .filter(|delay| !delay.is_zero()) + .unwrap_or(schedule.interval()) + .min(schedule.maximum_backoff()) + } + ScheduledOutcome::Failed(error) => { + let failures = health.failed(&id, error); + tracing::warn!(worker = %id, failures, "后台任务周期失败"); + take_backoff(&mut backoff, schedule.maximum_backoff()) + } + ScheduledOutcome::ShuttingDown => break, + }; + tokio::select! { + () = cancellation.cancelled() => break, + () = tokio::time::sleep(delay) => {} + } + } + health.stopped(&id); +} + +async fn acquire_and_run( + id: &WorkerId, + task: &dyn ScheduledTask, + schedule: WorkerSchedule, + request: WorkerLeaseRequest, + leader_lease: &dyn WorkerLeaderLeasePort, + cancellation: &CancellationToken, + health: &WorkerHealthRegistry, +) -> ScheduledOutcome { + let acquisition = AssertUnwindSafe(leader_lease.try_acquire(request)).catch_unwind(); + tokio::pin!(acquisition); + let acquisition = tokio::select! { + () = cancellation.cancelled() => return ScheduledOutcome::ShuttingDown, + acquisition = &mut acquisition => acquisition, + }; + match acquisition { + Ok(Ok(WorkerLeaseAcquisition::Acquired(guard))) => { + run_leased_cycle(id, task, guard, schedule, cancellation, health).await + } + Ok(Ok(WorkerLeaseAcquisition::Busy { retry_after })) => ScheduledOutcome::Busy(retry_after), + Ok(Err(error)) => ScheduledOutcome::Failed(error.as_safe_str().to_owned()), + Err(_) => ScheduledOutcome::Failed("leader lease port panicked".to_owned()), + } +} + +async fn run_local_cycle( + id: &WorkerId, + task: &dyn ScheduledTask, + cancellation: &CancellationToken, + health: &WorkerHealthRegistry, +) -> ScheduledOutcome { + health.running(id, None); + let cycle_cancel = CancellationToken::new(); + let cycle = AssertUnwindSafe(task.run_cycle(WorkerCycleContext::new( + id.clone(), + None, + cycle_cancel.clone(), + ))) + .catch_unwind(); + tokio::pin!(cycle); + tokio::select! { + outcome = &mut cycle => task_result(outcome), + () = cancellation.cancelled() => { + cycle_cancel.cancel(); + let _ = cycle.await; + ScheduledOutcome::ShuttingDown + } + } +} + +async fn run_leased_cycle( + id: &WorkerId, + task: &dyn ScheduledTask, + mut guard: Box, + schedule: WorkerSchedule, + cancellation: &CancellationToken, + health: &WorkerHealthRegistry, +) -> ScheduledOutcome { + let fencing_token = match std::panic::catch_unwind(AssertUnwindSafe(|| guard.fencing_token())) { + Ok(token) => token, + Err(_) => return ScheduledOutcome::Failed("leader lease guard panicked".to_owned()), + }; + health.running(id, Some(fencing_token)); + let cycle_cancel = CancellationToken::new(); + let cycle = AssertUnwindSafe(task.run_cycle(WorkerCycleContext::new( + id.clone(), + Some(fencing_token), + cycle_cancel.clone(), + ))) + .catch_unwind(); + tokio::pin!(cycle); + let renewal = tokio::time::sleep(schedule.leader_lease_renewal_interval()); + tokio::pin!(renewal); + + let outcome = loop { + tokio::select! { + result = &mut cycle => break task_result(result), + () = cancellation.cancelled() => { + cycle_cancel.cancel(); + let _ = cycle.await; + break ScheduledOutcome::ShuttingDown; + } + () = &mut renewal => { + let renewed = renew_guard( + guard.as_mut(), + schedule.leader_lease_renewal_interval(), + ).await; + if let Err(error) = renewed { + cycle_cancel.cancel(); + let _ = cycle.await; + break ScheduledOutcome::Failed(error); + } + renewal.as_mut().reset( + tokio::time::Instant::now() + schedule.leader_lease_renewal_interval(), + ); + } + } + }; + match release_guard(guard, schedule.leader_lease_renewal_interval()).await { + Ok(()) => outcome, + Err(error) if matches!(outcome, ScheduledOutcome::ShuttingDown) => { + tracing::warn!(worker = %id, error, "leader lease 释放失败"); + outcome + } + Err(error) => ScheduledOutcome::Failed(error), + } +} + +async fn renew_guard( + guard: &mut dyn WorkerLeaderLeaseGuard, + timeout: Duration, +) -> Result<(), String> { + let future = AssertUnwindSafe(guard.renew()).catch_unwind(); + match tokio::time::timeout(timeout, future).await { + Ok(Ok(Ok(()))) => Ok(()), + Ok(Ok(Err(error))) => Err(error.as_safe_str().to_owned()), + Ok(Err(_)) => Err("leader lease renewal panicked".to_owned()), + Err(_) => Err("leader lease renewal timed out".to_owned()), + } +} + +async fn release_guard( + guard: Box, + timeout: Duration, +) -> Result<(), String> { + let future = std::panic::catch_unwind(AssertUnwindSafe(|| guard.release())) + .map_err(|_| "leader lease release panicked".to_owned())?; + match tokio::time::timeout(timeout, AssertUnwindSafe(future).catch_unwind()).await { + Ok(Ok(Ok(()))) => Ok(()), + Ok(Ok(Err(error))) => Err(error.as_safe_str().to_owned()), + Ok(Err(_)) => Err("leader lease release panicked".to_owned()), + Err(_) => Err("leader lease release timed out".to_owned()), + } +} + +async fn supervise_daemon( + id: WorkerId, + task: Box, + restart: DaemonRestartPolicy, + cancellation: CancellationToken, + health: Arc, +) { + let mut backoff = restart.initial_backoff(); + loop { + if cancellation.is_cancelled() { + break; + } + health.running(&id, None); + let daemon_cancel = CancellationToken::new(); + let run = AssertUnwindSafe(task.run(daemon_cancel.clone())).catch_unwind(); + tokio::pin!(run); + let result = tokio::select! { + result = &mut run => Some(result), + () = cancellation.cancelled() => { + daemon_cancel.cancel(); + let _ = run.await; + None + } + }; + let Some(result) = result else { break }; + let error = match result { + Ok(Ok(())) => "daemon exited unexpectedly".to_owned(), + Ok(Err(error)) => error.as_safe_str().to_owned(), + Err(_) => "daemon panicked".to_owned(), + }; + let failures = health.failed(&id, error); + tracing::warn!(worker = %id, failures, "长驻任务将重启"); + let delay = take_backoff(&mut backoff, restart.maximum_backoff()); + tokio::select! { + () = cancellation.cancelled() => break, + () = tokio::time::sleep(delay) => {} + } + } + health.stopped(&id); +} + +enum ScheduledOutcome { + Succeeded, + Busy(Option), + Failed(String), + ShuttingDown, +} + +fn task_result( + result: Result, Box>, +) -> ScheduledOutcome { + match result { + Ok(Ok(())) => ScheduledOutcome::Succeeded, + Ok(Err(error)) => ScheduledOutcome::Failed(error.as_safe_str().to_owned()), + Err(_) => ScheduledOutcome::Failed("worker panicked".to_owned()), + } +} + +fn empty_snapshot(key: WorkerHealthKey, state: WorkerRuntimeState) -> WorkerHealthSnapshot { + WorkerHealthSnapshot { + key, + state, + consecutive_failures: 0, + completed_cycles: 0, + last_fencing_token: None, + last_success_at: None, + last_failure_at: None, + last_error: None, + } +} + +fn take_backoff(backoff: &mut Duration, maximum: Duration) -> Duration { + let delay = *backoff; + *backoff = backoff.checked_mul(2).unwrap_or(maximum).min(maximum); + delay +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-host/tests/config.rs b/backend/crates/gateway-host/tests/config.rs new file mode 100644 index 000000000..7cd54d2aa --- /dev/null +++ b/backend/crates/gateway-host/tests/config.rs @@ -0,0 +1,58 @@ +use std::path::PathBuf; + +use gateway_host::config::{FileLoggingConfig, HostConfig, ListenConfig, LoggingConfig}; +use gateway_host::system_update::SystemUpdateConfig; + +#[test] +fn host_config_resolves_only_host_owned_relative_paths() { + let mut config = valid_config(); + config + .resolve_and_validate(std::path::Path::new("/srv/gateway")) + .expect("valid host config"); + + assert_eq!( + config.logging.file.directory, + PathBuf::from("/srv/gateway/.runtime/logs") + ); +} + +#[test] +fn host_config_rejects_zero_drain_window() { + let mut config = valid_config(); + config.drain_timeout_seconds = 0; + + assert!( + config + .resolve_and_validate(std::path::Path::new("/srv/gateway")) + .is_err() + ); +} + +fn valid_config() -> HostConfig { + let system_update = SystemUpdateConfig { + update_state_file: PathBuf::from(".runtime/update-state.json"), + update_lock_file: PathBuf::from(".runtime/update-state.lock"), + update_temp_dir: PathBuf::from(".runtime/update-tmp"), + ..SystemUpdateConfig::default() + }; + HostConfig { + listen: ListenConfig { + host: "127.0.0.1".to_owned(), + port: 8080, + }, + logging: LoggingConfig { + level: "info".to_owned(), + stdout: true, + file: FileLoggingConfig { + enabled: true, + directory: PathBuf::from(".runtime/logs"), + retention_days: 7, + max_file_size_mb: 100, + max_files: 30, + }, + }, + system_update, + drain_timeout_seconds: 30, + worker_shutdown_timeout_seconds: 30, + } +} diff --git a/backend/crates/gateway-host/tests/logging.rs b/backend/crates/gateway-host/tests/logging.rs new file mode 100644 index 000000000..6e3270a41 --- /dev/null +++ b/backend/crates/gateway-host/tests/logging.rs @@ -0,0 +1,34 @@ +use std::path::PathBuf; + +use gateway_host::config::{FileLoggingConfig, HostConfig, ListenConfig, LoggingConfig}; +use gateway_host::system_update::SystemUpdateConfig; + +#[test] +fn logging_requires_at_least_one_sink() { + let mut config = HostConfig { + listen: ListenConfig { + host: "127.0.0.1".to_owned(), + port: 8080, + }, + logging: LoggingConfig { + level: "info".to_owned(), + stdout: false, + file: FileLoggingConfig { + enabled: false, + directory: PathBuf::from("logs"), + retention_days: 7, + max_file_size_mb: 100, + max_files: 30, + }, + }, + system_update: SystemUpdateConfig::default(), + drain_timeout_seconds: 30, + worker_shutdown_timeout_seconds: 30, + }; + + assert!( + config + .resolve_and_validate(std::path::Path::new("/tmp")) + .is_err() + ); +} diff --git a/backend/crates/gateway-host/tests/main.rs b/backend/crates/gateway-host/tests/main.rs new file mode 100644 index 000000000..2de636f9e --- /dev/null +++ b/backend/crates/gateway-host/tests/main.rs @@ -0,0 +1,5 @@ +mod config; +mod logging; +mod serve; +mod system_update; +mod workers; diff --git a/backend/crates/gateway-host/tests/serve.rs b/backend/crates/gateway-host/tests/serve.rs new file mode 100644 index 000000000..e237ed757 --- /dev/null +++ b/backend/crates/gateway-host/tests/serve.rs @@ -0,0 +1,8 @@ +use gateway_host::HostBundle; + +#[test] +fn host_bundle_serve_is_a_consuming_process_entrypoint() { + let _serve = HostBundle::serve; + + assert_eq!(std::mem::size_of_val(&_serve), 0); +} diff --git a/backend/crates/gateway-host/tests/system_update/mod.rs b/backend/crates/gateway-host/tests/system_update/mod.rs new file mode 100644 index 000000000..73ebf2219 --- /dev/null +++ b/backend/crates/gateway-host/tests/system_update/mod.rs @@ -0,0 +1,697 @@ +use std::fs; +use std::path::PathBuf; +use std::time::Duration; + +use filetime::FileTime; +use flate2::{Compression, write::GzEncoder}; +use futures::StreamExt as _; +use gateway_admin::model::system::{ + SystemOperationKind, SystemOperationStatus, SystemUpdateEventLevel, +}; +use gateway_admin::ports::system::{SystemOperationErrorKind, SystemOperations}; +use gateway_core::engine::CancellationToken; +use gateway_host::system_update::{ProcessSystemOperations, SystemUpdateConfig}; +use sha2::{Digest as _, Sha256}; +use tar::{Builder, EntryType, Header}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const TARGET_VERSION: &str = "9.9.9"; + +#[tokio::test] +async fn restart_should_not_shutdown_when_replacement_spawn_fails() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.executable_path = Some(fixture.root.path().join("missing")); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert!(service.restart().await.is_err()); + assert!(!shutdown.is_cancelled()); +} + +#[tokio::test] +async fn restart_should_request_process_restart_inside_docker() { + let fixture = Fixture::new(); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + config.deployment_mode = "docker".to_owned(); + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + service.restart().await.expect("restart accepted"); + + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested"); +} + +#[tokio::test] +async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { + let fixture = Fixture::new(); + fixture.write_executable("#!/bin/sh\nexit 0\n"); + let mut config = fixture.config("http://127.0.0.1:1/repos"); + config.self_restart_enabled = true; + let shutdown = CancellationToken::new(); + let service = ProcessSystemOperations::new(shutdown.clone(), config); + + assert_eq!( + service + .restart() + .await + .expect("replacement scheduled") + .kind(), + SystemOperationKind::Restart + ); + tokio::time::timeout(Duration::from_secs(2), shutdown.cancelled()) + .await + .expect("shutdown requested after spawn"); +} + +#[tokio::test] +async fn rollback_should_restore_binary_web_and_version_state() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let service = fixture.service(&server); + service + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("update"); + service.rollback().await.expect("rollback"); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"old-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); + assert_eq!( + service + .update_status() + .await + .expect("status") + .operation + .status, + SystemOperationStatus::Succeeded + ); +} + +#[tokio::test] +async fn update_detail_should_fallback_to_cache_when_refresh_fetch_fails() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service.update_detail(true).await.expect("prime cache"); + + assert_eq!( + service + .update_detail(true) + .await + .expect("cached fallback") + .latest_version, + TARGET_VERSION + ); +} + +#[tokio::test] +async fn update_detail_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + + let detail = service.update_detail(true).await.expect("safe rejection"); + assert!(!detail.update_supported); + assert!(detail.warning.is_some() || detail.unsupported_reason.is_some()); +} + +#[tokio::test] +async fn update_detail_should_report_source_build_as_unsupported() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + config.update_repository = None; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + + let detail = service.update_detail(false).await.expect("detail"); + assert_eq!(detail.latest_version, "1.0.0"); + assert!(!detail.update_supported); +} + +#[tokio::test] +async fn update_detail_should_use_cached_release_when_not_refreshed() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture.mount_release_once(&server, TARGET_VERSION).await; + let service = fixture.service(&server); + service.update_detail(true).await.expect("prime cache"); + + assert!(service.update_detail(false).await.is_ok()); +} + +#[tokio::test] +async fn update_events_should_close_after_terminal_update_log() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut events = service.update_events(); + let _ = service + .perform_update(Some(TARGET_VERSION.to_owned())) + .await; + let first = events.next().await.expect("terminal event"); + + assert_eq!(first.level, SystemUpdateEventLevel::Error); + assert!(events.next().await.is_none()); +} + +#[tokio::test] +async fn update_events_should_open_authenticated_sse_stream() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.build_type = "source".to_owned(); + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let mut stream = service.update_events(); + let _ = service + .perform_update(Some(TARGET_VERSION.to_owned())) + .await; + + assert!(stream.next().await.is_some()); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_is_missing() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Missing, + ) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_fail_when_release_checksum_mismatches() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Mismatch, + ) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_reject_insecure_release_archive() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "http://github.com/archive", None) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_from_untrusted_host() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_custom_release(&server, "https://evil.example/archive", None) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_reject_release_archive_with_unsafe_path() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::UnsafePath, + ChecksumKind::Valid, + ) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_reject_untrusted_github_api_base() { + let fixture = Fixture::new(); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.example/repos"), + ); + let error = service + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .expect_err("untrusted API rejected"); + + assert_eq!(error.kind(), SystemOperationErrorKind::Conflict); +} + +#[tokio::test] +async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release(&server, "9.9.8", ArchiveKind::Safe, ChecksumKind::Valid) + .await; + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); +} + +#[tokio::test] +async fn update_should_remove_stale_file_lock_and_continue() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::write(fixture.lock(), "stale").expect("lock"); + filetime::set_file_mtime(fixture.lock(), FileTime::from_unix_time(1, 0)).expect("old mtime"); + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_ok() + ); +} + +#[tokio::test] +async fn update_should_replace_local_release_files_with_latest_asset() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("update"); + + assert_eq!( + fs::read(fixture.executable()).expect("binary"), + b"new-binary" + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"new-web" + ); +} + +#[tokio::test] +async fn update_should_replace_web_assets_across_filesystems() { + let Ok(external) = tempfile::tempdir_in("/dev/shm") else { + return; + }; + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + let mut config = fixture.config(&format!("{}/repos", server.uri())); + config.web_dist_dir = external.path().join("dist"); + fs::create_dir_all(&config.web_dist_dir).expect("web dir"); + fs::write(config.web_dist_dir.join("index.html"), "old-web").expect("web"); + let service = ProcessSystemOperations::new(CancellationToken::new(), config.clone()); + service + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .expect("update"); + + assert_eq!( + fs::read(config.web_dist_dir.join("index.html")).expect("web"), + b"new-web" + ); +} + +#[tokio::test] +async fn update_should_restore_web_assets_when_binary_backup_fails() { + let server = MockServer::start().await; + let fixture = Fixture::new(); + fixture + .mount_release( + &server, + TARGET_VERSION, + ArchiveKind::Safe, + ChecksumKind::Valid, + ) + .await; + fs::create_dir(fixture.executable().with_extension("backup")).expect("blocking backup dir"); + + assert!( + fixture + .service(&server) + .perform_update(Some(TARGET_VERSION.to_owned())) + .await + .is_err() + ); + assert_eq!( + fs::read(fixture.web().join("index.html")).expect("web"), + b"old-web" + ); +} + +#[tokio::test] +async fn update_status_should_read_local_update_state() { + let fixture = Fixture::new(); + fs::write( + fixture.state(), + r#"{"previousVersion":"1.0.0","currentVersion":"2.0.0","operation":{"operationId":"x","kind":"update","status":"succeeded","targetVersion":"2.0.0","message":"done","error":null,"startedAt":"2026-07-19T00:00:00Z","finishedAt":"2026-07-19T00:01:00Z"}}"#, + ) + .expect("state"); + let service = ProcessSystemOperations::new( + CancellationToken::new(), + fixture.config("https://api.github.com/repos"), + ); + + assert_eq!( + service + .update_status() + .await + .expect("status") + .previous_version + .as_deref(), + Some("1.0.0") + ); +} + +#[tokio::test] +async fn version_should_return_backend_build_metadata() { + let fixture = Fixture::new(); + let mut config = fixture.config("https://api.github.com/repos"); + config.version = "2.3.4".to_owned(); + config.git_sha = "abc123".to_owned(); + config.update_repository = None; + let service = ProcessSystemOperations::new(CancellationToken::new(), config); + let version = service.version().await.expect("version"); + + assert_eq!( + (version.version.as_str(), version.git_sha.as_str()), + ("2.3.4", "abc123") + ); +} + +struct Fixture { + root: tempfile::TempDir, +} + +impl Fixture { + fn new() -> Self { + let root = tempfile::tempdir().expect("system update root"); + let fixture = Self { root }; + fixture.write_executable("old-binary"); + fs::create_dir_all(fixture.web()).expect("web dir"); + fs::write(fixture.web().join("index.html"), "old-web").expect("web"); + fixture + } + + fn executable(&self) -> PathBuf { + self.root.path().join("codex-proxy-rs") + } + + fn web(&self) -> PathBuf { + self.root.path().join("web/dist") + } + + fn state(&self) -> PathBuf { + self.root.path().join("update-state.json") + } + + fn lock(&self) -> PathBuf { + self.root.path().join("update.lock") + } + + fn write_executable(&self, content: &str) { + fs::write(self.executable(), content).expect("binary"); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + fs::set_permissions(self.executable(), fs::Permissions::from_mode(0o755)) + .expect("permissions"); + } + } + + fn config(&self, api_base: &str) -> SystemUpdateConfig { + SystemUpdateConfig { + version: "1.0.0".to_owned(), + git_sha: "test-sha".to_owned(), + build_time: "2026-07-19T00:00:00Z".to_owned(), + deployment_mode: "binary".to_owned(), + build_type: "release".to_owned(), + update_channel: "stable".to_owned(), + update_repository: Some("owner/repository".to_owned()), + github_api_base: api_base.to_owned(), + executable_path: Some(self.executable()), + web_dist_dir: self.web(), + update_state_file: self.state(), + update_lock_file: self.lock(), + update_temp_dir: self.root.path().join("tmp"), + self_restart_enabled: false, + } + } + + fn service(&self, server: &MockServer) -> ProcessSystemOperations { + ProcessSystemOperations::new( + CancellationToken::new(), + self.config(&format!("{}/repos", server.uri())), + ) + } + + async fn mount_release( + &self, + server: &MockServer, + version: &str, + archive_kind: ArchiveKind, + checksum_kind: ChecksumKind, + ) { + let archive = release_archive(archive_kind); + let name = archive_name(version); + let checksum = match checksum_kind { + ChecksumKind::Valid => format!("{} {name}\n", hex::encode(Sha256::digest(&archive))), + ChecksumKind::Mismatch => format!("{} {name}\n", "0".repeat(64)), + ChecksumKind::Missing => String::new(), + }; + let checksum_asset = (!matches!(checksum_kind, ChecksumKind::Missing)).then(|| { + serde_json::json!({ + "name": "checksums.txt", + "browser_download_url": format!("{}/checksums", server.uri()), + "size": checksum.len(), + }) + }); + let mut assets = vec![serde_json::json!({ + "name": name, + "browser_download_url": format!("{}/archive", server.uri()), + "size": archive.len(), + })]; + assets.extend(checksum_asset); + mount_release_json(server, version, assets, None).await; + Mock::given(method("GET")) + .and(path("/archive")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(archive)) + .mount(server) + .await; + if !checksum.is_empty() { + Mock::given(method("GET")) + .and(path("/checksums")) + .respond_with(ResponseTemplate::new(200).set_body_string(checksum)) + .mount(server) + .await; + } + } + + async fn mount_release_once(&self, server: &MockServer, version: &str) { + let assets = Vec::new(); + Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases/latest")) + .respond_with(release_response(version, assets)) + .up_to_n_times(1) + .mount(server) + .await; + } + + async fn mount_custom_release( + &self, + server: &MockServer, + archive_url: &str, + checksum_url: Option<&str>, + ) { + let archive = release_archive(ArchiveKind::Safe); + let name = archive_name(TARGET_VERSION); + let mut assets = vec![serde_json::json!({ + "name": name, + "browser_download_url": archive_url, + "size": archive.len(), + })]; + if let Some(checksum_url) = checksum_url { + assets.push(serde_json::json!({ + "name": "checksums.txt", + "browser_download_url": checksum_url, + "size": 80, + })); + } else { + assets.push(serde_json::json!({ + "name": "checksums.txt", + "browser_download_url": format!("{}/checksums", server.uri()), + "size": 80, + })); + } + mount_release_json(server, TARGET_VERSION, assets, None).await; + } +} + +#[derive(Clone, Copy)] +enum ArchiveKind { + Safe, + UnsafePath, +} + +#[derive(Clone, Copy)] +enum ChecksumKind { + Valid, + Mismatch, + Missing, +} + +fn release_archive(kind: ArchiveKind) -> Vec { + let encoder = GzEncoder::new(Vec::new(), Compression::default()); + let mut tar = Builder::new(encoder); + append_file(&mut tar, "codex-proxy-rs", b"new-binary", false); + append_file(&mut tar, "web/dist/index.html", b"new-web", false); + if matches!(kind, ArchiveKind::UnsafePath) { + append_file(&mut tar, "safe", b"escape", true); + } + let encoder = tar.into_inner().expect("tar"); + encoder.finish().expect("gzip") +} + +fn append_file(tar: &mut Builder>>, name: &str, data: &[u8], unsafe_path: bool) { + let mut header = Header::new_gnu(); + header.set_entry_type(EntryType::Regular); + header.set_mode(0o755); + header.set_size(u64::try_from(data.len()).expect("size")); + header.set_path(name).expect("path"); + if unsafe_path { + let bytes = header.as_mut_bytes(); + bytes[..100].fill(0); + bytes[..9].copy_from_slice(b"../escape"); + } + header.set_cksum(); + tar.append(&header, data).expect("append"); +} + +fn archive_name(version: &str) -> String { + format!( + "codex-proxy-rs-{version}-{}-{}.tar.gz", + std::env::consts::OS, + std::env::consts::ARCH + ) +} + +async fn mount_release_json( + server: &MockServer, + version: &str, + assets: Vec, + times: Option, +) { + let mock = Mock::given(method("GET")) + .and(path("/repos/owner/repository/releases/latest")) + .respond_with(release_response(version, assets)); + match times { + Some(times) => mock.up_to_n_times(times).mount(server).await, + None => mock.mount(server).await, + } +} + +fn release_response(version: &str, assets: Vec) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "tag_name": format!("v{version}"), + "name": format!("Release {version}"), + "body": "notes", + "html_url": "https://github.com/owner/repository/releases/latest", + "prerelease": false, + "published_at": "2026-07-19T00:00:00Z", + "assets": assets, + })) +} diff --git a/backend/crates/gateway-host/tests/workers/mod.rs b/backend/crates/gateway-host/tests/workers/mod.rs new file mode 100644 index 000000000..83ee519f0 --- /dev/null +++ b/backend/crates/gateway-host/tests/workers/mod.rs @@ -0,0 +1,856 @@ +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::future::pending; +use std::num::NonZeroU64; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures::future::BoxFuture; +use gateway_core::engine::CancellationToken; +use gateway_core::health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerRuntimeState}; +use gateway_core::task::{ + ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason, + WorkerFencingToken, WorkerId, WorkerKind, WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, + WorkerLeaseAcquisition, WorkerLeaseError, WorkerLeaseRequest, WorkerRegistration, + WorkerRunnable, WorkerSchedule, WorkerTaskError, +}; +use gateway_host::workers::{WorkerStartError, WorkerSupervisor}; +use tokio::sync::Notify; + +const ACTIVE_KINDS: [WorkerKind; 6] = [ + WorkerKind::OAuthRefresh, + WorkerKind::QuotaCatalogHealth, + WorkerKind::RuntimeSnapshotReconciliation, + WorkerKind::RuntimeChangeSubscription, + WorkerKind::StaleModelRequestRecovery, + WorkerKind::Retention, +]; + +#[derive(Clone)] +struct CountingTask { + runs: Arc, + fail_until: usize, +} + +impl CountingTask { + fn new(fail_until: usize) -> Self { + Self { + runs: Arc::new(AtomicUsize::new(0)), + fail_until, + } + } +} + +impl ScheduledTask for CountingTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let run = self.runs.fetch_add(1, Ordering::SeqCst) + 1; + if run <= self.fail_until { + Err(WorkerTaskError::safe("expected test failure")) + } else { + Ok(()) + } + }) + } +} + +#[derive(Clone)] +struct PanicOnceTask { + runs: Arc, +} + +impl ScheduledTask for PanicOnceTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + if self.runs.fetch_add(1, Ordering::SeqCst) == 0 { + panic!("expected isolated worker panic"); + } + Ok(()) + }) + } +} + +#[derive(Clone)] +struct CancellationTask { + entered: Arc, + observed: Arc, +} + +impl ScheduledTask for CancellationTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.entered.fetch_add(1, Ordering::SeqCst); + context.cancellation().cancelled().await; + self.observed.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + } +} + +#[derive(Clone)] +struct HungTask { + entered: Arc, + notification: Arc, +} + +impl ScheduledTask for HungTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.entered.store(true, Ordering::SeqCst); + self.notification.notify_one(); + pending::<()>().await; + Ok(()) + }) + } +} + +#[derive(Clone)] +struct SucceedsThenHangsTask { + runs: Arc, + second_cycle: Arc, +} + +impl ScheduledTask for SucceedsThenHangsTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + if self.runs.fetch_add(1, Ordering::SeqCst) == 0 { + return Ok(()); + } + self.second_cycle.notify_one(); + pending::<()>().await; + Ok(()) + }) + } +} + +#[derive(Clone)] +struct LeaseAwareTask { + entered: Arc, + release: Arc, + completed: Arc, + saw_live_guard: Arc, + expected_worker: WorkerId, + active_guards: Arc>>, + fence: Arc, +} + +impl ScheduledTask for LeaseAwareTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.saw_live_guard.store( + lock_unpoisoned(&self.active_guards).contains(context.worker()), + Ordering::SeqCst, + ); + if context.worker() != &self.expected_worker { + return Err(WorkerTaskError::safe("unexpected worker identity")); + } + self.fence.store( + context + .fencing_token() + .expect("leased task fencing token") + .get(), + Ordering::SeqCst, + ); + self.entered.notify_one(); + self.release.notified().await; + self.completed.store(true, Ordering::SeqCst); + Ok(()) + }) + } +} + +#[derive(Clone)] +struct ConcurrentCycleTask { + starts: Arc, + active: Arc, + maximum_active: Arc, +} + +impl ScheduledTask for ConcurrentCycleTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + self.starts.fetch_add(1, Ordering::SeqCst); + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.maximum_active.fetch_max(active, Ordering::SeqCst); + let _active = ActiveCycle(Arc::clone(&self.active)); + context.cancellation().cancelled().await; + Ok(()) + }) + } +} + +struct ActiveCycle(Arc); + +impl Drop for ActiveCycle { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} + +#[derive(Debug, Clone)] +enum LeaseStep { + Acquired, + Busy(Option), + Error, + Panic, +} + +#[derive(Debug, Clone, Copy)] +enum RenewStep { + Success, + Error, +} + +#[derive(Default)] +struct FakeLeasePort { + scripts: Mutex>>, + renew_scripts: Arc>>>, + requests: Mutex>, + renewals: Arc>>, + active_guards: Arc>>, + next_fence: AtomicU64, +} + +impl FakeLeasePort { + fn script(&self, worker: WorkerId, steps: impl IntoIterator) { + lock_unpoisoned(&self.scripts).insert(worker, steps.into_iter().collect()); + } + + fn renew_script(&self, worker: WorkerId, steps: impl IntoIterator) { + lock_unpoisoned(&self.renew_scripts).insert(worker, steps.into_iter().collect()); + } + + fn requests(&self) -> Vec { + lock_unpoisoned(&self.requests).clone() + } + + fn renewal_count(&self, worker: &WorkerId) -> usize { + lock_unpoisoned(&self.renewals) + .iter() + .filter(|renewed| *renewed == worker) + .count() + } +} + +struct FakeLeaseGuard { + worker: WorkerId, + fencing_token: WorkerFencingToken, + active_guards: Arc>>, + renew_scripts: Arc>>>, + renewals: Arc>>, +} + +impl WorkerLeaderLeaseGuard for FakeLeaseGuard { + fn fencing_token(&self) -> WorkerFencingToken { + self.fencing_token + } + + fn renew(&mut self) -> BoxFuture<'_, Result<(), WorkerLeaseError>> { + Box::pin(async move { + lock_unpoisoned(&self.renewals).push(self.worker.clone()); + match lock_unpoisoned(&self.renew_scripts) + .get_mut(&self.worker) + .and_then(VecDeque::pop_front) + .unwrap_or(RenewStep::Success) + { + RenewStep::Success => Ok(()), + RenewStep::Error => Err(WorkerLeaseError::safe("expected renewal failure")), + } + }) + } + + fn release(self: Box) -> BoxFuture<'static, Result<(), WorkerLeaseError>> { + Box::pin(async move { + drop(self); + Ok(()) + }) + } +} + +impl Drop for FakeLeaseGuard { + fn drop(&mut self) { + lock_unpoisoned(&self.active_guards).remove(&self.worker); + } +} + +impl WorkerLeaderLeasePort for FakeLeasePort { + fn try_acquire( + &self, + request: WorkerLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + lock_unpoisoned(&self.requests).push(request.clone()); + let step = lock_unpoisoned(&self.scripts) + .get_mut(request.worker()) + .and_then(VecDeque::pop_front) + .unwrap_or(LeaseStep::Acquired); + match step { + LeaseStep::Acquired => { + let fence = self.next_fence.fetch_add(1, Ordering::SeqCst) + 1; + lock_unpoisoned(&self.active_guards).insert(request.worker().clone()); + Ok(WorkerLeaseAcquisition::Acquired(Box::new(FakeLeaseGuard { + worker: request.worker().clone(), + fencing_token: WorkerFencingToken::new( + NonZeroU64::new(fence).expect("positive fence"), + ), + active_guards: Arc::clone(&self.active_guards), + renew_scripts: Arc::clone(&self.renew_scripts), + renewals: Arc::clone(&self.renewals), + }))) + } + LeaseStep::Busy(retry_after) => Ok(WorkerLeaseAcquisition::Busy { retry_after }), + LeaseStep::Error => Err(WorkerLeaseError::safe("expected lease failure")), + LeaseStep::Panic => panic!("expected isolated lease panic"), + } + }) + } +} + +#[tokio::test] +async fn registry_composes_multiple_owners_but_rejects_duplicate_identity() { + let duplicate = registration( + WorkerKind::OAuthRefresh, + "openai", + CountingTask::new(0), + test_schedule(), + true, + ); + let mut plan = complete_plan(vec![registration( + WorkerKind::OAuthRefresh, + "openai", + CountingTask::new(0), + test_schedule(), + true, + )]); + plan.push(duplicate); + let supervisor = supervisor(); + + assert!(matches!( + supervisor.start(plan, Arc::new(FakeLeasePort::default())), + Err(WorkerStartError::Duplicate(id)) if id.owner() == "openai" + )); +} + +#[tokio::test] +async fn only_database_fact_reasons_can_disable_the_two_nonexistent_workers() { + let mut plan = complete_plan(Vec::new()); + plan.push(WorkerContribution::Registration(registration_value( + WorkerKind::NativeClaimRecovery, + "store", + CountingTask::new(0), + test_schedule(), + true, + ))); + + assert!(matches!( + supervisor().start(plan, Arc::new(FakeLeasePort::default())), + Err(WorkerStartError::KindDisabled( + WorkerKind::NativeClaimRecovery + )) + )); +} + +#[tokio::test] +async fn registry_refuses_start_when_any_real_owner_kind_is_missing() { + let plan = complete_plan(Vec::new()) + .into_iter() + .filter(|item| item.kind() != WorkerKind::Retention) + .collect(); + + assert!(matches!( + supervisor().start(plan, Arc::new(FakeLeasePort::default())), + Err(WorkerStartError::Missing(missing)) if missing == vec![WorkerKind::Retention] + )); +} + +#[tokio::test] +async fn disabled_database_facts_never_execute_or_acquire_a_lease() { + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor + .start(complete_plan(Vec::new()), leases.clone()) + .expect("complete plan"); + wait_until(|| leases.requests().len() >= ACTIVE_KINDS.len()).await; + let snapshot = supervisor.health_source().snapshot(); + + assert!(snapshot.iter().any(|state| { + state.key == WorkerHealthKey::Disabled(WorkerKind::NativeClaimRecovery) + && state.state == WorkerRuntimeState::Disabled + })); + assert!(leases.requests().iter().all(|request| { + !matches!( + request.worker().kind(), + WorkerKind::NativeClaimRecovery | WorkerKind::OpsFlush + ) + })); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test] +async fn one_kind_runs_all_registered_owner_tasks() { + let second = CountingTask::new(0); + let plan = complete_plan(vec![ + registration( + WorkerKind::OAuthRefresh, + "openai", + CountingTask::new(0), + test_schedule(), + true, + ), + registration( + WorkerKind::OAuthRefresh, + "xai", + second.clone(), + test_schedule(), + true, + ), + ]); + let supervisor = supervisor(); + supervisor + .start(plan, Arc::new(FakeLeasePort::default())) + .expect("complete plan"); + wait_until(|| second.runs.load(Ordering::SeqCst) >= 1).await; + + assert_eq!( + supervisor + .health_source() + .snapshot() + .iter() + .filter(|state| matches!(&state.key, WorkerHealthKey::Task(id) if id.kind() == WorkerKind::OAuthRefresh)) + .count(), + 2 + ); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test] +async fn leader_guard_lives_through_cycle_and_fence_reaches_owner() { + let worker = target_worker(); + let leases = Arc::new(FakeLeasePort::default()); + let task = LeaseAwareTask { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + completed: Arc::new(AtomicBool::new(false)), + saw_live_guard: Arc::new(AtomicBool::new(false)), + expected_worker: worker.clone(), + active_guards: Arc::clone(&leases.active_guards), + fence: Arc::new(AtomicU64::new(0)), + }; + let supervisor = supervisor(); + supervisor + .start( + target_plan(task.clone(), long_interval_schedule()), + leases.clone(), + ) + .expect("complete plan"); + task.entered.notified().await; + + assert!(task.saw_live_guard.load(Ordering::SeqCst)); + assert!(lock_unpoisoned(&leases.active_guards).contains(&worker)); + task.release.notify_one(); + wait_until(|| task.completed.load(Ordering::SeqCst)).await; + wait_until(|| !lock_unpoisoned(&leases.active_guards).contains(&worker)).await; + assert!(task.fence.load(Ordering::SeqCst) > 0); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test(start_paused = true)] +async fn long_cycle_renews_leader_lease_before_ttl() { + let leases = Arc::new(FakeLeasePort::default()); + let task = lease_aware_task(&leases); + let supervisor = supervisor(); + supervisor + .start( + target_plan(task.clone(), renewal_schedule()), + leases.clone(), + ) + .expect("complete plan"); + task.entered.notified().await; + tokio::time::advance(Duration::from_millis(10)).await; + yield_until(|| leases.renewal_count(&target_worker()) >= 1).await; + + assert!(lock_unpoisoned(&leases.active_guards).contains(&target_worker())); + task.release.notify_one(); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test(start_paused = true)] +async fn renewal_failure_cancels_old_cycle_before_another_leader_cycle_starts() { + let task = ConcurrentCycleTask { + starts: Arc::new(AtomicUsize::new(0)), + active: Arc::new(AtomicUsize::new(0)), + maximum_active: Arc::new(AtomicUsize::new(0)), + }; + let leases = Arc::new(FakeLeasePort::default()); + leases.renew_script(target_worker(), [RenewStep::Error]); + let supervisor = supervisor(); + supervisor + .start(target_plan(task.clone(), renewal_schedule()), leases) + .expect("complete plan"); + yield_until(|| task.starts.load(Ordering::SeqCst) == 1).await; + tokio::time::advance(Duration::from_millis(10)).await; + yield_until(|| task.active.load(Ordering::SeqCst) == 0).await; + tokio::time::advance(Duration::from_millis(5)).await; + yield_until(|| task.starts.load(Ordering::SeqCst) == 2).await; + + assert_eq!(task.maximum_active.load(Ordering::SeqCst), 1); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test(start_paused = true)] +async fn supervisor_uses_exponential_backoff_and_recovers_health() { + let task = CountingTask::new(2); + let supervisor = supervisor(); + supervisor + .start( + target_plan(task.clone(), backoff_schedule()), + Arc::new(FakeLeasePort::default()), + ) + .expect("complete plan"); + yield_until(|| task.runs.load(Ordering::SeqCst) == 1).await; + tokio::time::advance(Duration::from_millis(10)).await; + yield_until(|| task.runs.load(Ordering::SeqCst) == 2).await; + tokio::time::advance(Duration::from_millis(20)).await; + yield_until(|| task.runs.load(Ordering::SeqCst) == 3).await; + let health = task_health(&supervisor, &target_worker()); + + assert_eq!( + (health.consecutive_failures, health.completed_cycles), + (0, 1) + ); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test(start_paused = true)] +async fn task_and_lease_port_panics_are_isolated_and_retried() { + let task = PanicOnceTask { + runs: Arc::new(AtomicUsize::new(0)), + }; + let leases = Arc::new(FakeLeasePort::default()); + leases.script( + target_worker(), + [LeaseStep::Error, LeaseStep::Panic, LeaseStep::Acquired], + ); + let supervisor = supervisor(); + supervisor + .start(target_plan(task.clone(), backoff_schedule()), leases) + .expect("complete plan"); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_millis(10)).await; + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_millis(20)).await; + yield_until(|| task.runs.load(Ordering::SeqCst) == 1).await; + tokio::time::advance(Duration::from_millis(20)).await; + yield_until(|| task.runs.load(Ordering::SeqCst) >= 2).await; + + assert_eq!( + task_health(&supervisor, &target_worker()).consecutive_failures, + 0 + ); + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test(start_paused = true)] +async fn lease_busy_is_healthy_standby_and_retries_without_running_owner() { + let task = CountingTask::new(0); + let leases = Arc::new(FakeLeasePort::default()); + leases.script( + target_worker(), + [ + LeaseStep::Busy(Some(Duration::from_millis(7))), + LeaseStep::Acquired, + ], + ); + let supervisor = supervisor(); + supervisor + .start(target_plan(task.clone(), backoff_schedule()), leases) + .expect("complete plan"); + yield_until(|| task_health(&supervisor, &target_worker()).state == WorkerRuntimeState::Standby) + .await; + + assert_eq!(task.runs.load(Ordering::SeqCst), 0); + tokio::time::advance(Duration::from_millis(7)).await; + yield_until(|| task.runs.load(Ordering::SeqCst) == 1).await; + supervisor.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test] +async fn cooperative_shutdown_propagates_cancel_before_dropping_all_guards() { + let task = CancellationTask { + entered: Arc::new(AtomicUsize::new(0)), + observed: Arc::new(AtomicUsize::new(0)), + }; + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor + .start(all_active_plan(task.clone()), leases.clone()) + .expect("complete plan"); + wait_until(|| task.entered.load(Ordering::SeqCst) == ACTIVE_KINDS.len()).await; + supervisor.shutdown(Duration::from_secs(1)).await; + + assert_eq!(task.observed.load(Ordering::SeqCst), ACTIVE_KINDS.len()); + assert!(lock_unpoisoned(&leases.active_guards).is_empty()); +} + +#[tokio::test] +async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() { + let task = HungTask { + entered: Arc::new(AtomicBool::new(false)), + notification: Arc::new(Notify::new()), + }; + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor + .start( + target_plan(task.clone(), long_interval_schedule()), + leases.clone(), + ) + .expect("complete plan"); + task.notification.notified().await; + supervisor.shutdown(Duration::from_millis(10)).await; + + assert!(lock_unpoisoned(&leases.active_guards).is_empty()); + assert_eq!( + task_health(&supervisor, &target_worker()).state, + WorkerRuntimeState::Stopped + ); +} + +#[tokio::test] +async fn running_worker_becomes_unhealthy_when_last_success_is_stale() { + let task = SucceedsThenHangsTask { + runs: Arc::new(AtomicUsize::new(0)), + second_cycle: Arc::new(Notify::new()), + }; + let supervisor = supervisor(); + supervisor + .start( + target_plan(task.clone(), freshness_schedule()), + Arc::new(FakeLeasePort::default()), + ) + .expect("complete plan"); + task.second_cycle.notified().await; + tokio::time::sleep(Duration::from_millis(45)).await; + + assert_eq!( + task_health(&supervisor, &target_worker()).state, + WorkerRuntimeState::BackingOff + ); + supervisor.shutdown(Duration::from_millis(10)).await; +} + +fn supervisor() -> WorkerSupervisor { + WorkerSupervisor::new(CancellationToken::new()) +} + +fn registration( + kind: WorkerKind, + owner: &str, + task: impl ScheduledTask + 'static, + schedule: WorkerSchedule, + leased: bool, +) -> WorkerContribution { + WorkerContribution::Registration(registration_value(kind, owner, task, schedule, leased)) +} + +fn registration_value( + kind: WorkerKind, + owner: &str, + task: impl ScheduledTask + 'static, + schedule: WorkerSchedule, + leased: bool, +) -> WorkerRegistration { + let id = WorkerId::try_new(kind, owner).expect("valid worker id"); + let lease = leased + .then(|| WorkerLeaseRequest::try_new(id.clone(), schedule.leader_lease_ttl())) + .transpose() + .expect("valid lease"); + WorkerRegistration::try_new( + id, + WorkerRunnable::Scheduled { + schedule, + lease, + task: Box::new(task), + }, + ) + .expect("valid registration") +} + +fn complete_plan(mut custom: Vec) -> Vec { + let present = custom + .iter() + .map(WorkerContribution::kind) + .collect::>(); + for kind in ACTIVE_KINDS { + if !present.contains(&kind) { + custom.push(registration( + kind, + kind.as_str(), + CountingTask::new(0), + long_interval_schedule(), + true, + )); + } + } + custom.push(WorkerContribution::Disabled { + kind: WorkerKind::NativeClaimRecovery, + reason: WorkerDisabledReason::NoPersistentNativeClaimState, + }); + custom.push(WorkerContribution::Disabled { + kind: WorkerKind::OpsFlush, + reason: WorkerDisabledReason::NoBufferedOpsEvents, + }); + custom +} + +fn target_plan( + task: impl ScheduledTask + 'static, + schedule: WorkerSchedule, +) -> Vec { + complete_plan(vec![registration( + WorkerKind::OAuthRefresh, + "openai", + task, + schedule, + true, + )]) +} + +fn all_active_plan(task: impl ScheduledTask + Clone + 'static) -> Vec { + let mut plan = ACTIVE_KINDS + .into_iter() + .map(|kind| registration(kind, kind.as_str(), task.clone(), test_schedule(), true)) + .collect::>(); + plan.push(WorkerContribution::Disabled { + kind: WorkerKind::NativeClaimRecovery, + reason: WorkerDisabledReason::NoPersistentNativeClaimState, + }); + plan.push(WorkerContribution::Disabled { + kind: WorkerKind::OpsFlush, + reason: WorkerDisabledReason::NoBufferedOpsEvents, + }); + plan +} + +fn lease_aware_task(leases: &Arc) -> LeaseAwareTask { + LeaseAwareTask { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + completed: Arc::new(AtomicBool::new(false)), + saw_live_guard: Arc::new(AtomicBool::new(false)), + expected_worker: target_worker(), + active_guards: Arc::clone(&leases.active_guards), + fence: Arc::new(AtomicU64::new(0)), + } +} + +fn target_worker() -> WorkerId { + WorkerId::try_new(WorkerKind::OAuthRefresh, "openai").expect("valid target worker") +} + +fn task_health(supervisor: &WorkerSupervisor, id: &WorkerId) -> WorkerHealthSnapshot { + supervisor + .health_source() + .snapshot() + .into_iter() + .find(|state| state.key == WorkerHealthKey::Task(id.clone())) + .expect("worker health") +} + +fn schedule( + interval: Duration, + initial: Duration, + maximum: Duration, + ttl: Duration, + renewal: Duration, +) -> WorkerSchedule { + WorkerSchedule::try_new(interval, initial, maximum, ttl, renewal).expect("valid schedule") +} + +fn test_schedule() -> WorkerSchedule { + schedule( + Duration::from_millis(20), + Duration::from_millis(5), + Duration::from_millis(20), + Duration::from_secs(30), + Duration::from_secs(10), + ) +} + +fn long_interval_schedule() -> WorkerSchedule { + schedule( + Duration::from_secs(60), + Duration::from_millis(5), + Duration::from_millis(20), + Duration::from_secs(30), + Duration::from_secs(10), + ) +} + +fn backoff_schedule() -> WorkerSchedule { + schedule( + Duration::from_secs(60), + Duration::from_millis(10), + Duration::from_millis(20), + Duration::from_secs(30), + Duration::from_secs(10), + ) +} + +fn renewal_schedule() -> WorkerSchedule { + schedule( + Duration::from_secs(60), + Duration::from_millis(5), + Duration::from_millis(20), + Duration::from_millis(30), + Duration::from_millis(10), + ) +} + +fn freshness_schedule() -> WorkerSchedule { + schedule( + Duration::from_millis(10), + Duration::from_millis(5), + Duration::from_millis(10), + Duration::from_millis(30), + Duration::from_millis(10), + ) +} + +async fn wait_until(mut condition: impl FnMut() -> bool) { + tokio::time::timeout(Duration::from_secs(1), async { + while !condition() { + tokio::task::yield_now().await; + } + }) + .await + .expect("worker test timeout"); +} + +async fn yield_until(mut condition: impl FnMut() -> bool) { + for _ in 0..100 { + if condition() { + return; + } + tokio::task::yield_now().await; + } + assert!(condition(), "condition did not become true after yielding"); +} + +fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} diff --git a/backend/crates/gateway-protocol/src/openai/mod.rs b/backend/crates/gateway-protocol/src/openai/mod.rs index d98897eec..676b5ab82 100644 --- a/backend/crates/gateway-protocol/src/openai/mod.rs +++ b/backend/crates/gateway-protocol/src/openai/mod.rs @@ -1,5 +1,14 @@ //! OpenAI wire contract。 +/// Codex Responses Lite 的 HTTP 请求头。 +pub const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str = + "x-openai-internal-codex-responses-lite"; +/// Codex Responses Lite 在 WebSocket `client_metadata` 中的投影键。 +pub const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str = + "ws_request_header_x_openai_internal_codex_responses_lite"; +/// Codex memory consolidation 请求标记。 +pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request"; + /// OpenAI/Codex 事件语义、用量与限流字段编解码。 pub mod events; /// OpenAI/Codex JSON schema 纯转换。 diff --git a/backend/crates/gateway-store/Cargo.toml b/backend/crates/gateway-store/Cargo.toml index 2253389f7..fa690d4ca 100644 --- a/backend/crates/gateway-store/Cargo.toml +++ b/backend/crates/gateway-store/Cargo.toml @@ -15,6 +15,7 @@ path = "tests/main.rs" async-trait = "0.1.89" chrono = "0.4.45" futures = "0.3.32" +gateway-admin = { path = "../gateway-admin" } gateway-core = { path = "../gateway-core" } hex = "0.4.3" redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] } @@ -24,6 +25,8 @@ sha2 = "0.11.0" sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "json", "macros", "migrate"] } thiserror = "2.0.18" tokio = { version = "1.52.3", default-features = false, features = ["rt"] } +tracing = "0.1.44" +url = "2.5.8" uuid = { version = "1.23.3", features = ["v4", "v7"] } [dev-dependencies] diff --git a/backend/crates/gateway-store/src/lib.rs b/backend/crates/gateway-store/src/lib.rs index 987c4feb9..4cab7bf0e 100644 --- a/backend/crates/gateway-store/src/lib.rs +++ b/backend/crates/gateway-store/src/lib.rs @@ -1,20 +1,41 @@ //! 多 Provider 网关的 PostgreSQL 持久化与 Redis 协调 adapter。 //! -//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的十张业务表 +//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的八张业务表 //! 和可丢失 Redis 状态映射为明确的基础设施操作。 -#![forbid(unsafe_code)] - +use std::sync::Arc; +use std::time::{Duration, SystemTime}; use std::{fmt, num::NonZeroU64, str::FromStr}; +use gateway_admin::model::auth::{AdminAuditEvent as AdminAuditModel, AdminSession}; +use gateway_admin::model::settings::{ + AdminApiKey, AdminApiKeyMutation, ProviderModelMappings, ReplaceRuntimeSettings, + RotationStrategy as AdminRotationStrategy, RuntimeSettings as AdminRuntimeSettings, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision as AdminRevision}; +use gateway_admin::ports::store::{ + AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, AuthStore, + SettingsStore, +}; +use gateway_core::CoreStorePorts; +use gateway_core::health::{HealthProbe, HealthState}; +use gateway_core::provider_ports::{ProviderCatalogPorts, ProviderStorePorts}; +use gateway_core::task::{ + ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason, WorkerId, + WorkerKind, WorkerLeaderLeasePort, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, + WorkerSchedule, WorkerTaskError, +}; +use serde::Deserialize; use serde_json::{Map, Value}; +const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL"; +const REDIS_URL_ENV: &str = "CPR_REDIS_URL"; +const DATABASE_PASSWORD_ENV: &str = "CPR_DATABASE_PASSWORD"; +const REDIS_PASSWORD_ENV: &str = "CPR_REDIS_PASSWORD"; + pub mod postgres; pub mod redis; -/// 保持依赖方向为 `gateway-store -> gateway-core`。 -pub use gateway_core as core; - /// 发生错误的基础设施边界。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum StoreBackend { @@ -59,6 +80,884 @@ pub enum StoreError { pub type StoreResult = Result; +/// Store 自己拥有并校验的启动配置。 +#[derive(Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StoreConfig { + database: StoreConnectionConfig, + redis: StoreConnectionConfig, +} + +impl StoreConfig { + pub fn resolve_and_validate(&mut self, _source_dir: &std::path::Path) -> StoreResult<()> { + if let Some(url) = optional_environment_value(DATABASE_URL_ENV)? { + self.database.url = url; + } + if let Some(url) = optional_environment_value(REDIS_URL_ENV)? { + self.redis.url = url; + } + if let Some(password) = optional_environment_value(DATABASE_PASSWORD_ENV)? { + self.database.password = password; + } + if let Some(password) = optional_environment_value(REDIS_PASSWORD_ENV)? { + self.redis.password = password; + } + self.database.validate("database")?; + self.redis.validate("redis")?; + Ok(()) + } + + fn database_url(&self) -> StoreResult { + self.database.connection_url("database") + } + + fn redis_url(&self) -> StoreResult { + self.redis.connection_url("redis") + } +} + +fn optional_environment_value(name: &'static str) -> StoreResult> { + match std::env::var(name) { + Ok(value) if value.trim().is_empty() => Err(StoreError::InvalidData { + entity: "store config", + message: format!("environment variable {name} is empty"), + }), + Ok(value) => Ok(Some(value)), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(std::env::VarError::NotUnicode(_)) => Err(StoreError::InvalidData { + entity: "store config", + message: format!("environment variable {name} is not Unicode"), + }), + } +} + +impl fmt::Debug for StoreConfig { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("StoreConfig") + .field("database", &"[REDACTED]") + .field("redis", &"[REDACTED]") + .finish() + } +} + +#[derive(Clone, Deserialize)] +#[serde(deny_unknown_fields)] +struct StoreConnectionConfig { + url: String, + password: String, +} + +impl StoreConnectionConfig { + fn validate(&self, field: &'static str) -> StoreResult<()> { + require_nonempty("store config", field, &self.url)?; + if self.password.len() != 48 || !self.password.bytes().all(|byte| byte.is_ascii_hexdigit()) + { + return Err(StoreError::InvalidData { + entity: "store config", + message: format!("{field}.password must be exactly 48 hexadecimal characters"), + }); + } + self.connection_url(field).map(|_| ()) + } + + fn connection_url(&self, field: &'static str) -> StoreResult { + let mut url = url::Url::parse(&self.url).map_err(|_| StoreError::InvalidData { + entity: "store config", + message: format!("{field}.url is invalid"), + })?; + if url.password().is_some() { + return Err(StoreError::InvalidData { + entity: "store config", + message: format!("{field}.url must not contain a password"), + }); + } + url.set_password(Some(&self.password)) + .map_err(|()| StoreError::InvalidData { + entity: "store config", + message: format!("{field}.url cannot carry credentials"), + })?; + Ok(url.to_string()) + } +} + +/// 已完成连接、迁移与 cooldown hydration 的 Store 能力集合。 +pub struct StoreBundle { + admin_ports: AdminStorePorts, + core_ports: CoreStorePorts, + provider_ports: ProviderStorePorts, + worker_leader_lease: Arc, + health_probes: Vec>, + worker_contributions: Vec, +} + +impl StoreBundle { + #[must_use] + pub fn admin_ports(&self) -> AdminStorePorts { + self.admin_ports.clone() + } + + #[must_use] + pub fn core_ports(&self) -> CoreStorePorts { + self.core_ports.clone() + } + + #[must_use] + pub fn provider_ports(&self) -> ProviderStorePorts { + self.provider_ports.clone() + } + + #[must_use] + pub fn worker_leader_lease(&self) -> Arc { + Arc::clone(&self.worker_leader_lease) + } + + #[must_use] + pub fn health_probes(&self) -> Vec> { + self.health_probes.clone() + } + + pub fn take_worker_contributions(&mut self) -> Vec { + std::mem::take(&mut self.worker_contributions) + } +} + +/// 在返回 Bundle 前完成全部 Store 启动屏障。 +pub async fn initialize(mut config: StoreConfig) -> StoreResult { + const REDIS_NAMESPACE: &str = "codex-proxy-rs"; + + config.resolve_and_validate(std::path::Path::new("."))?; + let pool = postgres::connect_and_migrate(&config.database_url()?).await?; + let redis_client = ::redis::Client::open(config.redis_url()?) + .map_err(|_| redis_unavailable("create Redis client"))?; + let redis_connection = redis_client + .get_connection_manager() + .await + .map_err(|_| redis_unavailable("connect Redis manager"))?; + + let provider_accounts = Arc::new(postgres::PgProviderAccountRepository::new(pool.clone())); + let cooldowns = Arc::new(redis::RedisCredentialCooldownRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?); + let account_store = Arc::new(redis::CooldownCachingProviderAccountStore::new( + provider_accounts, + cooldowns.clone(), + )); + let hydrated = account_store.hydrate(std::time::SystemTime::now()).await; + tracing::info!(hydrated, "账号 cooldown Redis 热缓存重建完成"); + + let credential_leases = + redis::RedisCredentialLeaseRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?; + let provider_leases = Arc::new(redis::RedisProviderLeaseCoordinator::new( + credential_leases.clone(), + )); + let credential_state = Arc::new(redis::RedisCredentialStateRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?); + let provider_catalog = ProviderCatalogPorts::new( + Arc::new(postgres::PgConfigCatalogRepository::new(pool.clone())), + credential_state.clone(), + ); + let runtime_policy = Arc::new(postgres::PgRuntimeSettingsRepository::new(pool.clone())); + let oauth_pending = Arc::new(redis::RedisOAuthPendingFlowRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?); + + let admin_ports = AdminStorePorts::new( + Arc::new(postgres::PgAdminAccountStore::new(pool.clone())), + Arc::new(AdminAuthStoreAdapter { + security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()), + settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()), + state: redis::RedisAdminAuthStateRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?, + }), + Arc::new(postgres::PgAdminCatalogStore::new(pool.clone())), + Arc::new(postgres::PgAdminClientKeyStore::new(pool.clone())), + Arc::new(postgres::PgAdminObservabilityStore::new(pool.clone())), + Arc::new(AdminSettingsStoreAdapter { + control_plane: postgres::PgControlPlaneRepository::new(pool.clone()), + }), + ); + + let execution = Arc::new(postgres::PgExecutionStore::new(pool.clone())); + let retention = Arc::new(postgres::PgRetentionRepository::new(pool.clone())); + let core_ports = CoreStorePorts::new( + execution.clone(), + Arc::new(redis::RedisClientAdmissionRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + )?), + Arc::new(postgres::PgClientAdmissionRecoveryRepository::new( + pool.clone(), + )), + Arc::new(redis::RedisProviderCircuitRepository::new( + redis_connection.clone(), + REDIS_NAMESPACE, + redis::ProviderCircuitPolicy::default(), + )?), + Arc::new(postgres::PgHistoryRepository::new(pool.clone())), + Arc::new(postgres::PgRuntimeSnapshotRepository::new(pool.clone())), + Arc::new(redis::RedisRuntimeChangeRepository::new( + redis_client, + REDIS_NAMESPACE, + )?), + ); + + let provider_ports = ProviderStorePorts::new( + account_store, + provider_leases, + provider_catalog, + credential_state, + cooldowns, + runtime_policy, + oauth_pending, + ); + let worker_leader_lease = Arc::new(redis::worker_lease::RedisWorkerLeaderLeasePort::new( + credential_leases, + )); + let health_probes: Vec> = vec![ + Arc::new(PostgresHealthProbe { pool: pool.clone() }), + Arc::new(RedisHealthProbe { + connection: redis_connection, + }), + ]; + let worker_contributions = store_worker_contributions(execution, retention)?; + Ok(StoreBundle { + admin_ports, + core_ports, + provider_ports, + worker_leader_lease, + health_probes, + worker_contributions, + }) +} + +fn store_worker_contributions( + execution: Arc, + retention: Arc, +) -> StoreResult> { + let stale_id = WorkerId::try_new(WorkerKind::StaleModelRequestRecovery, "postgres") + .map_err(worker_definition_error)?; + let retention_id = + WorkerId::try_new(WorkerKind::Retention, "postgres").map_err(worker_definition_error)?; + Ok(vec![ + WorkerContribution::Disabled { + kind: WorkerKind::NativeClaimRecovery, + reason: WorkerDisabledReason::NoPersistentNativeClaimState, + }, + WorkerContribution::Registration(scheduled_worker( + stale_id, + Duration::from_secs(30), + Box::new(StaleModelRequestRecoveryTask { execution }), + )?), + WorkerContribution::Registration(scheduled_worker( + retention_id, + Duration::from_secs(60 * 60), + Box::new(RetentionTask { retention }), + )?), + WorkerContribution::Disabled { + kind: WorkerKind::OpsFlush, + reason: WorkerDisabledReason::NoBufferedOpsEvents, + }, + ]) +} + +fn scheduled_worker( + id: WorkerId, + interval: Duration, + task: Box, +) -> StoreResult { + let schedule = WorkerSchedule::try_new( + interval, + Duration::from_secs(1), + Duration::from_secs(60), + Duration::from_secs(15 * 60), + Duration::from_secs(5 * 60), + ) + .map_err(worker_definition_error)?; + let lease = WorkerLeaseRequest::try_new(id.clone(), schedule.leader_lease_ttl()) + .map_err(worker_definition_error)?; + WorkerRegistration::try_new( + id, + WorkerRunnable::Scheduled { + schedule, + lease: Some(lease), + task, + }, + ) + .map_err(worker_definition_error) +} + +fn worker_definition_error(error: gateway_core::task::WorkerDefinitionError) -> StoreError { + StoreError::InvalidData { + entity: "store worker plan", + message: error.to_string(), + } +} + +struct StaleModelRequestRecoveryTask { + execution: Arc, +} + +impl ScheduledTask for StaleModelRequestRecoveryTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> futures::future::BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + gateway_core::engine::ExecutionStore::recover_expired( + self.execution.as_ref(), + SystemTime::now(), + ) + .await + .map(|_| ()) + .map_err(|_| WorkerTaskError::safe("stale request recovery failed")) + }) + } +} + +struct RetentionTask { + retention: Arc, +} + +impl ScheduledTask for RetentionTask { + fn run_cycle( + &self, + _context: WorkerCycleContext, + ) -> futures::future::BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let settings = + postgres::RetentionRepository::load_retention_settings(self.retention.as_ref()) + .await + .map_err(|_| WorkerTaskError::safe("retention settings read failed"))?; + postgres::RetentionRepository::apply_retention( + self.retention.as_ref(), + chrono::Utc::now(), + settings, + ) + .await + .map(|_| ()) + .map_err(|_| WorkerTaskError::safe("retention cleanup failed")) + }) + } +} + +struct PostgresHealthProbe { + pool: sqlx::PgPool, +} + +impl HealthProbe for PostgresHealthProbe { + fn name(&self) -> &'static str { + "postgres" + } + + fn check(&self) -> futures::future::BoxFuture<'_, HealthState> { + Box::pin(async move { + match sqlx::query_scalar::<_, i32>("select 1") + .fetch_one(&self.pool) + .await + { + Ok(1) => HealthState::Healthy, + Ok(_) => HealthState::Unhealthy("PostgreSQL health result is invalid".to_owned()), + Err(_) => HealthState::Unhealthy("PostgreSQL is unavailable".to_owned()), + } + }) + } +} + +struct RedisHealthProbe { + connection: ::redis::aio::ConnectionManager, +} + +struct AdminAuthStoreAdapter { + security: postgres::PgAdminSecurityAuditRepository, + settings: postgres::PgRuntimeSettingsRepository, + state: redis::RedisAdminAuthStateRepository, +} + +struct AdminSettingsStoreAdapter { + control_plane: postgres::PgControlPlaneRepository, +} + +#[async_trait::async_trait] +impl SettingsStore for AdminSettingsStoreAdapter { + async fn load_runtime_settings(&self) -> AdminStoreResult { + let snapshot = postgres::ControlPlaneRepository::load_control_plane(&self.control_plane) + .await + .map_err(|error| admin_store_error("runtime settings", error))?; + admin_runtime_settings(snapshot.settings) + } + + async fn admin_api_key_exists(&self) -> AdminStoreResult { + postgres::ControlPlaneRepository::load_control_plane(&self.control_plane) + .await + .map(|snapshot| snapshot.settings.admin_api_key.is_some()) + .map_err(|error| admin_store_error("admin API key", error)) + } + + async fn replace_runtime_settings( + &self, + command: ReplaceRuntimeSettings, + context: &MutationContext, + ) -> AdminStoreResult { + let current = postgres::ControlPlaneRepository::load_control_plane(&self.control_plane) + .await + .map_err(|error| admin_store_error("runtime settings", error))?; + let replacement = postgres::ControlPlaneReplacement { + settings: postgres::RuntimeSettingsUpdate { + admin_api_key: current.settings.admin_api_key, + refresh_margin_seconds: command.refresh_margin_seconds, + refresh_concurrency: command.refresh_concurrency, + max_concurrent_per_account: command.max_concurrent_per_account, + request_interval_ms: command.request_interval_ms, + rotation_strategy: admin_rotation_name(command.rotation_strategy).to_owned(), + provider_model_mappings: store_model_mappings(command.provider_model_mappings), + usage_retention_days: command.usage_retention_days, + ops_event_retention_days: command.ops_event_retention_days, + audit_retention_days: command.audit_retention_days, + }, + audit: mutation_audit( + context, + "settings.replace", + "runtime_settings", + "1", + vec![ + "provider_model_mappings_json".to_owned(), + "refresh_margin_seconds".to_owned(), + "refresh_concurrency".to_owned(), + "max_concurrent_per_account".to_owned(), + "request_interval_ms".to_owned(), + "rotation_strategy".to_owned(), + "retention".to_owned(), + ], + ), + }; + let snapshot = postgres::ControlPlaneRepository::replace_control_plane( + &self.control_plane, + store_revision(command.expected_config_revision)?, + replacement, + ) + .await + .map_err(|error| admin_store_error("runtime settings", error))?; + admin_runtime_settings(snapshot.settings) + } + + async fn replace_admin_api_key( + &self, + expected_config_revision: AdminRevision, + key: AdminApiKey, + context: &MutationContext, + ) -> AdminStoreResult { + self.replace_admin_api_key_value( + expected_config_revision, + Some(key.expose_for_auth().to_owned()), + context, + ) + .await + } + + async fn delete_admin_api_key( + &self, + expected_config_revision: AdminRevision, + context: &MutationContext, + ) -> AdminStoreResult { + self.replace_admin_api_key_value(expected_config_revision, None, context) + .await + } +} + +impl AdminSettingsStoreAdapter { + async fn replace_admin_api_key_value( + &self, + expected_config_revision: AdminRevision, + admin_api_key: Option, + context: &MutationContext, + ) -> AdminStoreResult { + let current = postgres::ControlPlaneRepository::load_control_plane(&self.control_plane) + .await + .map_err(|error| admin_store_error("admin API key", error))?; + let exists = admin_api_key.is_some(); + let update = postgres::RuntimeSettingsUpdate { + admin_api_key, + refresh_margin_seconds: current.settings.refresh_margin_seconds, + refresh_concurrency: current.settings.refresh_concurrency, + max_concurrent_per_account: current.settings.max_concurrent_per_account, + request_interval_ms: current.settings.request_interval_ms, + rotation_strategy: current.settings.rotation_strategy, + provider_model_mappings: current.settings.provider_model_mappings, + usage_retention_days: current.settings.usage_retention_days, + ops_event_retention_days: current.settings.ops_event_retention_days, + audit_retention_days: current.settings.audit_retention_days, + }; + let snapshot = postgres::ControlPlaneRepository::replace_control_plane( + &self.control_plane, + store_revision(expected_config_revision)?, + postgres::ControlPlaneReplacement { + settings: update, + audit: mutation_audit( + context, + if exists { + "admin_api_key.replace" + } else { + "admin_api_key.delete" + }, + "runtime_settings", + "1", + vec!["admin_api_key".to_owned()], + ), + }, + ) + .await + .map_err(|error| admin_store_error("admin API key", error))?; + Ok(AdminApiKeyMutation { + config_revision: admin_revision(snapshot.settings.config_revision)?, + exists, + }) + } +} + +fn admin_runtime_settings( + settings: postgres::RuntimeSettings, +) -> AdminStoreResult { + let rotation_strategy = match settings.rotation_strategy.as_str() { + "smart" => AdminRotationStrategy::Smart, + "quota_reset_priority" => AdminRotationStrategy::QuotaResetPriority, + "round_robin" => AdminRotationStrategy::RoundRobin, + "sticky" => AdminRotationStrategy::Sticky, + _ => { + return Err(AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "runtime settings", + "rotation strategy is invalid", + )); + } + }; + let provider_model_mappings = settings + .provider_model_mappings + .into_iter() + .map(|(provider, mappings)| { + let provider = gateway_core::routing::ProviderKind::new(provider).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "runtime settings", + "provider model mapping is invalid", + ) + })?; + let mappings = mappings + .into_iter() + .map(|(public, upstream)| { + let public = + gateway_core::routing::PublicModelId::new(public).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "runtime settings", + "public model mapping is invalid", + ) + })?; + let upstream = + gateway_core::routing::UpstreamModelId::new(upstream).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "runtime settings", + "upstream model mapping is invalid", + ) + })?; + Ok((public, upstream)) + }) + .collect::>()?; + Ok((provider, mappings)) + }) + .collect::>()?; + Ok(AdminRuntimeSettings { + config_revision: admin_revision(settings.config_revision)?, + provider_model_mappings, + refresh_margin_seconds: settings.refresh_margin_seconds, + refresh_concurrency: settings.refresh_concurrency, + max_concurrent_per_account: settings.max_concurrent_per_account, + request_interval_ms: settings.request_interval_ms, + rotation_strategy, + usage_retention_days: settings.usage_retention_days, + ops_event_retention_days: settings.ops_event_retention_days, + audit_retention_days: settings.audit_retention_days, + updated_at: settings.updated_at, + }) +} + +fn store_model_mappings( + mappings: ProviderModelMappings, +) -> std::collections::BTreeMap> { + mappings + .into_iter() + .map(|(provider, mappings)| { + ( + provider.as_str().to_owned(), + mappings + .into_iter() + .map(|(public, upstream)| { + (public.as_str().to_owned(), upstream.as_str().to_owned()) + }) + .collect(), + ) + }) + .collect() +} + +const fn admin_rotation_name(strategy: AdminRotationStrategy) -> &'static str { + match strategy { + AdminRotationStrategy::Smart => "smart", + AdminRotationStrategy::QuotaResetPriority => "quota_reset_priority", + AdminRotationStrategy::RoundRobin => "round_robin", + AdminRotationStrategy::Sticky => "sticky", + } +} + +pub(crate) fn store_revision(revision: AdminRevision) -> AdminStoreResult { + Revision::new(revision.get()).map_err(|error| admin_store_error("config revision", error)) +} + +pub(crate) fn admin_revision(revision: Revision) -> AdminStoreResult { + AdminRevision::new(revision.get()).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "config revision", + "config revision is invalid", + ) + }) +} + +pub(crate) fn mutation_audit( + context: &MutationContext, + action: &str, + entity_kind: &str, + entity_ref: &str, + changed_fields: Vec, +) -> postgres::AdminAuditEvent { + let (actor_kind, actor_admin_user_id, actor_ref) = match &context.actor { + MutationActor::AdminSession { admin_user_id } => ( + postgres::AdminAuditActorKind::AdminSession, + Some(admin_user_id.clone()), + admin_user_id.clone(), + ), + MutationActor::AdminApiKey => ( + postgres::AdminAuditActorKind::AdminApiKey, + None, + "admin_api_key".to_owned(), + ), + MutationActor::System => ( + postgres::AdminAuditActorKind::System, + None, + "system".to_owned(), + ), + }; + postgres::AdminAuditEvent { + id: format!("audit_{}", uuid::Uuid::now_v7().simple()), + actor_kind, + actor_admin_user_id, + actor_ref, + admin_request_id: Some(context.request_id.clone()), + action: action.to_owned(), + entity_kind: entity_kind.to_owned(), + entity_ref: entity_ref.to_owned(), + config_revision: None, + changed_fields, + created_at: chrono::Utc::now(), + } +} + +#[async_trait::async_trait] +impl AuthStore for AdminAuthStoreAdapter { + async fn load_password_hash(&self, admin_user_id: &str) -> AdminStoreResult> { + postgres::AdminSecurityAuditRepository::password_hash(&self.security, admin_user_id) + .await + .map_err(|error| admin_store_error("admin authentication", error)) + } + + async fn create_password_hash_if_absent( + &self, + admin_user_id: &str, + password_hash: &str, + ) -> AdminStoreResult { + postgres::AdminSecurityAuditRepository::create_password_hash_if_absent( + &self.security, + admin_user_id, + password_hash, + ) + .await + .map_err(|error| admin_store_error("admin authentication", error)) + } + + async fn load_admin_api_key(&self) -> AdminStoreResult> { + postgres::RuntimeSettingsRepository::load_runtime_settings(&self.settings) + .await + .map(|settings| settings.admin_api_key.map(AdminApiKey::new)) + .map_err(|error| admin_store_error("admin API key", error)) + } + + async fn load_session(&self, session_id: &str) -> AdminStoreResult> { + redis::AdminAuthStateRepository::load_admin_session(&self.state, session_id) + .await + .map(|session| { + session.map(|record| AdminSession { + admin_user_id: record.admin_user_id, + expires_at: record.expires_at, + }) + }) + .map_err(|error| admin_store_error("admin session", error)) + } + + async fn store_session( + &self, + session_id: &str, + session: &AdminSession, + ) -> AdminStoreResult<()> { + redis::AdminAuthStateRepository::store_admin_session( + &self.state, + session_id, + &redis::AdminSessionRecord { + admin_user_id: session.admin_user_id.clone(), + expires_at: session.expires_at, + }, + ) + .await + .map_err(|error| admin_store_error("admin session", error)) + } + + async fn delete_session(&self, session_id: &str) -> AdminStoreResult> { + redis::AdminAuthStateRepository::delete_admin_session(&self.state, session_id) + .await + .map(|session| { + session.map(|record| AdminSession { + admin_user_id: record.admin_user_id, + expires_at: record.expires_at, + }) + }) + .map_err(|error| admin_store_error("admin session", error)) + } + + async fn login_source_is_throttled( + &self, + source: &str, + failure_limit: u32, + window_seconds: u64, + ) -> AdminStoreResult { + redis::AdminAuthStateRepository::login_source_is_throttled( + &self.state, + source, + failure_limit, + window_seconds, + ) + .await + .map_err(|error| admin_store_error("admin login throttle", error)) + } + + async fn record_login_failure( + &self, + source: &str, + failure_limit: u32, + window_seconds: u64, + ) -> AdminStoreResult { + redis::AdminAuthStateRepository::record_login_failure( + &self.state, + source, + failure_limit, + window_seconds, + ) + .await + .map_err(|error| admin_store_error("admin login throttle", error)) + } + + async fn clear_login_failures(&self, source: &str) -> AdminStoreResult<()> { + redis::AdminAuthStateRepository::clear_login_failures(&self.state, source) + .await + .map_err(|error| admin_store_error("admin login throttle", error)) + } + + async fn append_audit_event(&self, event: AdminAuditModel) -> AdminStoreResult<()> { + let config_revision = event + .config_revision + .map(|revision| i64::try_from(revision.get())) + .transpose() + .map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + "admin audit", + "config revision is outside the supported range", + ) + })?; + let actor_kind = match event.actor_kind { + gateway_admin::model::auth::AuditActorKind::AdminSession => { + postgres::AdminAuditActorKind::AdminSession + } + gateway_admin::model::auth::AuditActorKind::AdminApiKey => { + postgres::AdminAuditActorKind::AdminApiKey + } + gateway_admin::model::auth::AuditActorKind::System => { + postgres::AdminAuditActorKind::System + } + gateway_admin::model::auth::AuditActorKind::Anonymous => { + postgres::AdminAuditActorKind::Anonymous + } + }; + postgres::AdminSecurityAuditRepository::append_admin_audit_event( + &self.security, + postgres::AdminAuditEvent { + id: event.id, + actor_kind, + actor_admin_user_id: event.actor_admin_user_id, + actor_ref: event.actor_ref, + admin_request_id: event.request_id, + action: event.action, + entity_kind: event.entity_kind, + entity_ref: event.entity_ref, + config_revision, + changed_fields: event.changed_fields, + created_at: event.occurred_at, + }, + ) + .await + .map_err(|error| admin_store_error("admin audit", error)) + } +} + +pub(crate) fn admin_store_error(resource: &'static str, error: StoreError) -> AdminStoreError { + let kind = match error { + StoreError::NotFound { .. } => AdminStoreErrorKind::NotFound, + StoreError::Conflict { + kind: ConflictKind::StaleRevision, + .. + } => AdminStoreErrorKind::StaleRevision, + StoreError::Conflict { .. } => AdminStoreErrorKind::Conflict, + StoreError::InvalidData { .. } => AdminStoreErrorKind::Invalid, + StoreError::Unavailable { .. } => AdminStoreErrorKind::Unavailable, + }; + AdminStoreError::new(kind, resource, "store operation failed") +} + +impl HealthProbe for RedisHealthProbe { + fn name(&self) -> &'static str { + "redis" + } + + fn check(&self) -> futures::future::BoxFuture<'_, HealthState> { + Box::pin(async move { + let mut connection = self.connection.clone(); + match ::redis::cmd("PING") + .query_async::(&mut connection) + .await + { + Ok(response) if response == "PONG" => HealthState::Healthy, + Ok(_) => HealthState::Unhealthy("Redis health result is invalid".to_owned()), + Err(_) => HealthState::Unhealthy("Redis is unavailable".to_owned()), + } + }) + } +} + /// 正整数 revision 或 fencing token。 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Revision(NonZeroU64); diff --git a/backend/crates/gateway-store/src/postgres/admin_security_audit.rs b/backend/crates/gateway-store/src/postgres/admin_security_audit.rs index 6f00042ba..bf83f8394 100644 --- a/backend/crates/gateway-store/src/postgres/admin_security_audit.rs +++ b/backend/crates/gateway-store/src/postgres/admin_security_audit.rs @@ -71,11 +71,11 @@ impl AdminAuditEvent { pub trait AdminSecurityAuditRepository: Send + Sync { async fn password_hash(&self, admin_user_id: &str) -> StoreResult>; - async fn replace_password_hash( + async fn create_password_hash_if_absent( &self, admin_user_id: &str, password_hash: &str, - ) -> StoreResult<()>; + ) -> StoreResult; async fn append_admin_audit_event(&self, event: AdminAuditEvent) -> StoreResult<()>; } @@ -103,25 +103,24 @@ impl AdminSecurityAuditRepository for PgAdminSecurityAuditRepository { .map_err(|_| postgres_unavailable("read admin password hash")) } - async fn replace_password_hash( + async fn create_password_hash_if_absent( &self, admin_user_id: &str, password_hash: &str, - ) -> StoreResult<()> { + ) -> StoreResult { require_nonempty("admin user", "id", admin_user_id)?; require_nonempty("admin user", "password_hash", password_hash)?; - sqlx::query( + let result = sqlx::query( "insert into admin_users (id, password_hash, created_at, updated_at) values ($1, $2, now(), now()) - on conflict (id) do update - set password_hash = excluded.password_hash, updated_at = now()", + on conflict (id) do nothing", ) .bind(admin_user_id) .bind(password_hash) .execute(&self.pool) .await - .map_err(|_| postgres_unavailable("replace admin password hash"))?; - Ok(()) + .map_err(|_| postgres_unavailable("create admin password hash"))?; + Ok(result.rows_affected() == 1) } async fn append_admin_audit_event(&self, event: AdminAuditEvent) -> StoreResult<()> { diff --git a/backend/crates/gateway-store/src/postgres/admission_recovery.rs b/backend/crates/gateway-store/src/postgres/admission_recovery.rs index cef920524..92c91d7e9 100644 --- a/backend/crates/gateway-store/src/postgres/admission_recovery.rs +++ b/backend/crates/gateway-store/src/postgres/admission_recovery.rs @@ -4,6 +4,16 @@ use std::collections::BTreeMap; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_core::{ + engine::{ + ModelRequestId, + admission::{ + ClientAdmissionError, ClientAdmissionRecovery as CoreAdmissionRecovery, + ClientAdmissionRecoveryPort, RecentAdmissionFact, RunningAdmissionFact, + }, + }, + policy::ClientApiKeyId, +}; use sqlx::PgPool; use crate::{StoreError, StoreResult, postgres_unavailable}; @@ -103,6 +113,54 @@ impl ClientAdmissionRecoveryRepository for PgClientAdmissionRecoveryRepository { } } +impl ClientAdmissionRecoveryPort for PgClientAdmissionRecoveryRepository { + fn load_recovery( + &self, + since: std::time::SystemTime, + ) -> futures::future::BoxFuture<'_, Result, ClientAdmissionError>> + { + Box::pin(async move { + self.load_client_admission_recovery(DateTime::::from(since)) + .await + .map_err(|_| ClientAdmissionError)? + .into_iter() + .map(|recovery| { + let client_api_key_id = ClientApiKeyId::new(recovery.client_api_key_ref) + .map_err(|_| ClientAdmissionError)?; + let recent_requests = recovery + .recent_requests + .into_iter() + .map(|request| { + Ok(RecentAdmissionFact { + model_request_id: ModelRequestId::new(request.model_request_id) + .map_err(|_| ClientAdmissionError)?, + started_at: request.started_at.into(), + input_token_estimate: request.input_token_estimate, + }) + }) + .collect::, ClientAdmissionError>>()?; + let running_requests = recovery + .running_requests + .into_iter() + .map(|request| { + Ok(RunningAdmissionFact { + model_request_id: ModelRequestId::new(request.model_request_id) + .map_err(|_| ClientAdmissionError)?, + expires_at: request.deadline_at.into(), + }) + }) + .collect::, ClientAdmissionError>>()?; + Ok(CoreAdmissionRecovery { + client_api_key_id, + recent_requests, + running_requests, + }) + }) + .collect() + }) + } +} + fn to_u64(value: i64) -> StoreResult { u64::try_from(value).map_err(|_| StoreError::InvalidData { entity: "client admission recovery", diff --git a/backend/crates/gateway-store/src/postgres/client_keys.rs b/backend/crates/gateway-store/src/postgres/client_keys.rs index dba982ff0..0842bdd2d 100644 --- a/backend/crates/gateway-store/src/postgres/client_keys.rs +++ b/backend/crates/gateway-store/src/postgres/client_keys.rs @@ -4,10 +4,31 @@ use std::fmt; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_admin::{ + model::{ + MutationContext, + client_keys::{ + ClientKeyCursor as AdminClientKeyCursor, + ClientKeyCursorValue as AdminClientKeyCursorValue, + ClientKeyListQuery as AdminClientKeyListQuery, ClientKeyPage as AdminClientKeyPage, + ClientKeyRecord as AdminClientKeyRecord, ClientKeySecret as AdminClientKeySecret, + ClientKeySort as AdminClientKeySort, ClientKeySortField as AdminClientKeySortField, + DeleteClientKey, NewClientKey, SetClientKeyEnabled, + SortDirection as AdminSortDirection, UpdateClientKey as AdminUpdateClientKey, + }, + }, + ports::store::{AdminStoreResult, ClientKeyStore}, +}; use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::ProviderKind; use sqlx::{PgPool, Postgres, QueryBuilder, Transaction}; -use crate::{StoreError, StoreResult, postgres_unavailable, require_nonempty}; +use crate::{ + StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit, + postgres_unavailable, require_nonempty, store_revision, +}; + +use super::{ControlPlaneRepository, PgControlPlaneRepository}; const ENTITY: &str = "client API key"; const KEY_LENGTH: usize = 46; @@ -189,8 +210,8 @@ pub struct ClientApiKeyListQuery { impl ClientApiKeyListQuery { pub fn validate(&self) -> StoreResult<()> { - if self.page_size == 0 || self.page_size > 200 { - return Err(invalid("page size must be between 1 and 200")); + if self.page_size == 0 { + return Err(invalid("page size must be between 1 and 65535")); } if self.search.as_deref().is_some_and(|search| { search.trim().is_empty() || search.len() > 256 || search.chars().any(char::is_control) @@ -293,6 +314,7 @@ pub trait ClientApiKeyRepository: Send + Sync { &self, query: ClientApiKeyListQuery, ) -> StoreResult; + async fn get_client_api_key(&self, id: &str) -> StoreResult>; async fn reveal_client_api_key(&self, id: &str) -> StoreResult>; async fn insert_client_api_key(&self, key: NewClientApiKey) -> StoreResult<()>; async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult; @@ -388,6 +410,21 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository { .transpose() } + async fn get_client_api_key(&self, id: &str) -> StoreResult> { + require_nonempty(ENTITY, "id", id)?; + sqlx::query_as::<_, ClientRecordRow>( + "select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency, + requests_per_minute, tokens_per_minute, last_used_at, created_at, updated_at + from client_api_keys where id = $1", + ) + .bind(id) + .fetch_optional(&self.pool) + .await + .map_err(|_| postgres_unavailable("get client API key"))? + .map(client_record_from_row) + .transpose() + } + async fn insert_client_api_key(&self, key: NewClientApiKey) -> StoreResult<()> { key.validate()?; sqlx::query( @@ -449,6 +486,315 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository { } } +/// Admin 用例所需的 Client Key 事务能力。 +#[derive(Clone)] +pub struct PgAdminClientKeyStore { + keys: PgClientApiKeyRepository, + control_plane: PgControlPlaneRepository, +} + +impl PgAdminClientKeyStore { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + keys: PgClientApiKeyRepository::new(pool.clone()), + control_plane: PgControlPlaneRepository::new(pool), + } + } + + async fn revision(&self) -> AdminStoreResult { + self.control_plane + .load_control_plane() + .await + .map_err(|error| admin_store_error(ENTITY, error)) + .and_then(|snapshot| admin_revision(snapshot.settings.config_revision)) + } + + async fn required_record(&self, id: &ClientApiKeyId) -> AdminStoreResult { + self.keys + .get_client_api_key(id.as_str()) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .ok_or_else(|| { + admin_store_error( + ENTITY, + StoreError::NotFound { + entity: ENTITY, + id: id.as_str().to_owned(), + }, + ) + }) + .and_then(admin_client_key_record) + } +} + +#[async_trait] +impl ClientKeyStore for PgAdminClientKeyStore { + async fn list_client_keys( + &self, + query: AdminClientKeyListQuery, + ) -> AdminStoreResult { + let config_revision = self.revision().await?; + let page = self + .keys + .list_client_api_keys(store_client_key_query(query)?) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok(AdminClientKeyPage { + config_revision, + items: page + .items + .into_iter() + .map(admin_client_key_record) + .collect::>>()?, + total: page.total, + next_cursor: page.next_cursor.map(admin_client_key_cursor).transpose()?, + }) + } + + async fn reveal_client_key( + &self, + id: &ClientApiKeyId, + ) -> AdminStoreResult> { + let Some(secret) = self + .keys + .reveal_client_api_key(id.as_str()) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + else { + return Ok(None); + }; + let record = self.required_record(id).await?; + Ok(Some(AdminClientKeySecret::new(record, secret.key))) + } + + async fn create_client_key( + &self, + command: NewClientKey, + context: &MutationContext, + ) -> AdminStoreResult<(gateway_admin::model::Revision, AdminClientKeyRecord)> { + let id = command.id; + let revision = self + .control_plane + .create_client_api_key( + store_revision(command.expected_config_revision)?, + NewClientApiKey { + id: id.as_str().to_owned(), + name: command.name, + label: command.label, + provider_kind: command.provider_kind.as_str().to_owned(), + key: command.plaintext, + max_concurrency: command.limits.max_concurrency, + requests_per_minute: command.limits.requests_per_minute, + tokens_per_minute: command.limits.tokens_per_minute, + }, + mutation_audit( + context, + "create", + "client_api_key", + id.as_str(), + [ + "name", + "label", + "provider_kind", + "key", + "enabled", + "max_concurrency", + "requests_per_minute", + "tokens_per_minute", + ] + .into_iter() + .map(str::to_owned) + .collect(), + ), + ) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok((admin_revision(revision)?, self.required_record(&id).await?)) + } + + async fn update_client_key( + &self, + command: AdminUpdateClientKey, + context: &MutationContext, + ) -> AdminStoreResult<(gateway_admin::model::Revision, AdminClientKeyRecord)> { + let id = command.id; + let revision = self + .control_plane + .update_client_api_key( + store_revision(command.expected_config_revision)?, + UpdateClientApiKeyDetails { + id: id.as_str().to_owned(), + name: command.name, + label: command.label, + provider_kind: command.provider_kind.as_str().to_owned(), + max_concurrency: command.limits.max_concurrency, + requests_per_minute: command.limits.requests_per_minute, + tokens_per_minute: command.limits.tokens_per_minute, + }, + mutation_audit( + context, + "update", + "client_api_key", + id.as_str(), + [ + "name", + "label", + "provider_kind", + "max_concurrency", + "requests_per_minute", + "tokens_per_minute", + ] + .into_iter() + .map(str::to_owned) + .collect(), + ), + ) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok((admin_revision(revision)?, self.required_record(&id).await?)) + } + + async fn set_client_key_enabled( + &self, + command: SetClientKeyEnabled, + context: &MutationContext, + ) -> AdminStoreResult<(gateway_admin::model::Revision, AdminClientKeyRecord)> { + let id = command.id; + let revision = self + .control_plane + .set_client_api_key_enabled( + store_revision(command.expected_config_revision)?, + id.as_str(), + command.enabled, + mutation_audit( + context, + if command.enabled { "enable" } else { "disable" }, + "client_api_key", + id.as_str(), + vec!["enabled".to_owned()], + ), + ) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok((admin_revision(revision)?, self.required_record(&id).await?)) + } + + async fn delete_client_key( + &self, + command: DeleteClientKey, + context: &MutationContext, + ) -> AdminStoreResult { + self.control_plane + .delete_client_api_key( + store_revision(command.expected_config_revision)?, + command.id.as_str(), + mutation_audit( + context, + "delete", + "client_api_key", + command.id.as_str(), + Vec::new(), + ), + ) + .await + .map_err(|error| admin_store_error(ENTITY, error)) + .and_then(admin_revision) + } +} + +fn store_client_key_query( + query: AdminClientKeyListQuery, +) -> AdminStoreResult { + let sort = store_client_key_sort(query.sort); + Ok(ClientApiKeyListQuery { + cursor: query + .cursor + .map(|cursor| store_client_key_cursor(cursor, sort)) + .transpose()?, + page_size: query.page_size.get(), + search: query.search, + sort, + }) +} + +fn store_client_key_sort(sort: AdminClientKeySort) -> ClientApiKeySort { + ClientApiKeySort { + field: match sort.field { + AdminClientKeySortField::Name => ClientApiKeySortField::Name, + AdminClientKeySortField::Enabled => ClientApiKeySortField::Enabled, + AdminClientKeySortField::CreatedAt => ClientApiKeySortField::CreatedAt, + AdminClientKeySortField::LastUsedAt => ClientApiKeySortField::LastUsedAt, + }, + direction: match sort.direction { + AdminSortDirection::Asc => ClientApiKeySortDirection::Asc, + AdminSortDirection::Desc => ClientApiKeySortDirection::Desc, + }, + } +} + +fn store_client_key_cursor( + cursor: AdminClientKeyCursor, + sort: ClientApiKeySort, +) -> AdminStoreResult { + let value = match cursor.value { + AdminClientKeyCursorValue::Name(value) => ClientApiKeyCursorValue::Name(value), + AdminClientKeyCursorValue::Enabled(value) => ClientApiKeyCursorValue::Enabled(value), + AdminClientKeyCursorValue::CreatedAt(value) => ClientApiKeyCursorValue::CreatedAt(value), + AdminClientKeyCursorValue::LastUsedAt(value) => ClientApiKeyCursorValue::LastUsedAt(value), + }; + ClientApiKeyCursor::new(sort, value, cursor.id.as_str().to_owned()) + .map_err(|error| admin_store_error(ENTITY, error)) +} + +fn admin_client_key_cursor(cursor: ClientApiKeyCursor) -> AdminStoreResult { + let sort = AdminClientKeySort { + field: match cursor.sort.field { + ClientApiKeySortField::Name => AdminClientKeySortField::Name, + ClientApiKeySortField::Enabled => AdminClientKeySortField::Enabled, + ClientApiKeySortField::CreatedAt => AdminClientKeySortField::CreatedAt, + ClientApiKeySortField::LastUsedAt => AdminClientKeySortField::LastUsedAt, + }, + direction: match cursor.sort.direction { + ClientApiKeySortDirection::Asc => AdminSortDirection::Asc, + ClientApiKeySortDirection::Desc => AdminSortDirection::Desc, + }, + }; + let value = match cursor.value { + ClientApiKeyCursorValue::Name(value) => AdminClientKeyCursorValue::Name(value), + ClientApiKeyCursorValue::Enabled(value) => AdminClientKeyCursorValue::Enabled(value), + ClientApiKeyCursorValue::CreatedAt(value) => AdminClientKeyCursorValue::CreatedAt(value), + ClientApiKeyCursorValue::LastUsedAt(value) => AdminClientKeyCursorValue::LastUsedAt(value), + }; + Ok(AdminClientKeyCursor { + sort, + value, + id: ClientApiKeyId::new(cursor.id) + .map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?, + }) +} + +fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult { + Ok(AdminClientKeyRecord { + id: ClientApiKeyId::new(record.id) + .map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?, + name: record.name, + label: record.label, + provider_kind: ProviderKind::new(record.provider_kind) + .map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))?, + prefix: record.prefix, + enabled: record.enabled, + limits: RateLimits { + max_concurrency: record.max_concurrency, + requests_per_minute: record.requests_per_minute, + tokens_per_minute: record.tokens_per_minute, + }, + last_used_at: record.last_used_at, + created_at: record.created_at, + updated_at: record.updated_at, + }) +} + pub(crate) async fn insert_client_api_key_in_transaction( transaction: &mut Transaction<'_, Postgres>, key: &NewClientApiKey, diff --git a/backend/crates/gateway-store/src/postgres/config_catalog.rs b/backend/crates/gateway-store/src/postgres/config_catalog.rs index e51213310..6ee238d08 100644 --- a/backend/crates/gateway-store/src/postgres/config_catalog.rs +++ b/backend/crates/gateway-store/src/postgres/config_catalog.rs @@ -2,9 +2,31 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_admin::{ + model::{ + MutationContext, + catalog::{ + CreateProviderInstance as AdminCreateProviderInstance, + DeleteProviderInstance as AdminDeleteProviderInstance, + ProviderInstance as AdminProviderInstance, ProviderInstanceCatalog, + ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled, + UpdateProviderInstance as AdminUpdateProviderInstance, + }, + }, + ports::store::{AdminStoreResult, CatalogStore}, +}; +use gateway_core::provider_ports::{ + ProviderInstanceCatalogPort, ProviderInstanceConfig, ProviderStoreError, ProviderStoreErrorKind, +}; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use sqlx::{PgPool, Postgres, Transaction}; -use crate::{StoreError, StoreResult, postgres_unavailable, require_nonempty}; +use crate::{ + StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit, + postgres_unavailable, require_nonempty, store_revision, +}; + +use super::{ControlPlaneRepository, PgControlPlaneRepository}; const ENTITY: &str = "provider instance"; @@ -185,6 +207,271 @@ impl ConfigCatalogRepository for PgConfigCatalogRepository { } } +impl ProviderInstanceCatalogPort for PgConfigCatalogRepository { + fn list_instances<'a>( + &'a self, + provider_kind: &'a ProviderKind, + include_disabled: bool, + ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> + { + Box::pin(async move { + self.list_provider_instances(include_disabled) + .await + .map_err(|_| { + ProviderStoreError::new( + ProviderStoreErrorKind::Unavailable, + "list Provider instances", + ) + })? + .into_iter() + .filter(|instance| instance.provider_kind == provider_kind.as_str()) + .map(|instance| { + let id = ProviderInstanceId::new(instance.id).map_err(|_| { + ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "decode Provider instance ID", + ) + })?; + let kind = ProviderKind::new(instance.provider_kind).map_err(|_| { + ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "decode Provider kind", + ) + })?; + Ok(ProviderInstanceConfig::new( + id, + kind, + instance.base_url, + instance.enabled, + )) + }) + .collect() + }) + } +} + +/// Admin 用例所需的 Provider instance 事务能力。 +#[derive(Clone)] +pub struct PgAdminCatalogStore { + catalog: PgConfigCatalogRepository, + control_plane: PgControlPlaneRepository, +} + +impl PgAdminCatalogStore { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + catalog: PgConfigCatalogRepository::new(pool.clone()), + control_plane: PgControlPlaneRepository::new(pool), + } + } + + async fn revision(&self) -> AdminStoreResult { + self.control_plane + .load_control_plane() + .await + .map_err(|error| admin_store_error("provider instance", error)) + .and_then(|snapshot| admin_revision(snapshot.settings.config_revision)) + } + + async fn load_mutation( + &self, + revision: crate::Revision, + id: &ProviderInstanceId, + ) -> AdminStoreResult { + let instance = self + .catalog + .get_provider_instance(id.as_str()) + .await + .map_err(|error| admin_store_error("provider instance", error))? + .map(admin_provider_instance) + .transpose()?; + Ok(ProviderInstanceMutation { + config_revision: admin_revision(revision)?, + instance, + }) + } +} + +#[async_trait] +impl CatalogStore for PgAdminCatalogStore { + async fn list_provider_instances( + &self, + include_disabled: bool, + ) -> AdminStoreResult { + let config_revision = self.revision().await?; + let items = self + .catalog + .list_provider_instances(include_disabled) + .await + .map_err(|error| admin_store_error("provider instance", error))? + .into_iter() + .map(admin_provider_instance) + .collect::>>()?; + Ok(ProviderInstanceCatalog { + config_revision, + items, + }) + } + + async fn load_provider_instance( + &self, + id: &ProviderInstanceId, + ) -> AdminStoreResult> { + let (revision, instance) = futures::try_join!(self.revision(), async { + self.catalog + .get_provider_instance(id.as_str()) + .await + .map_err(|error| admin_store_error("provider instance", error))? + .map(admin_provider_instance) + .transpose() + },)?; + Ok(instance.map(|item| ProviderInstanceDetail { + config_revision: revision, + item, + })) + } + + async fn create_provider_instance( + &self, + command: AdminCreateProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult { + let id = command.id; + let revision = self + .control_plane + .create_provider_instance( + store_revision(command.expected_config_revision)?, + NewProviderInstance { + id: id.as_str().to_owned(), + provider_kind: command.provider_kind.as_str().to_owned(), + name: command.name, + base_url: command.base_url, + }, + mutation_audit( + context, + "create", + "provider_instance", + id.as_str(), + vec![ + "provider_kind".to_owned(), + "name".to_owned(), + "base_url".to_owned(), + "enabled".to_owned(), + ], + ), + ) + .await + .map_err(|error| admin_store_error("provider instance", error))?; + self.load_mutation(revision, &id).await + } + + async fn update_provider_instance( + &self, + command: AdminUpdateProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult { + let id = command.id; + let revision = self + .control_plane + .update_provider_instance( + store_revision(command.expected_config_revision)?, + UpdateProviderInstanceDetails { + id: id.as_str().to_owned(), + name: command.name, + base_url: command.base_url, + }, + mutation_audit( + context, + "update", + "provider_instance", + id.as_str(), + vec!["name".to_owned(), "base_url".to_owned()], + ), + ) + .await + .map_err(|error| admin_store_error("provider instance", error))?; + self.load_mutation(revision, &id).await + } + + async fn set_provider_instance_enabled( + &self, + command: SetProviderInstanceEnabled, + context: &MutationContext, + ) -> AdminStoreResult { + let id = command.id; + let revision = self + .control_plane + .set_provider_instance_enabled( + store_revision(command.expected_config_revision)?, + id.as_str(), + command.enabled, + mutation_audit( + context, + if command.enabled { "enable" } else { "disable" }, + "provider_instance", + id.as_str(), + vec!["enabled".to_owned()], + ), + ) + .await + .map_err(|error| admin_store_error("provider instance", error))?; + self.load_mutation(revision, &id).await + } + + async fn delete_provider_instance( + &self, + command: AdminDeleteProviderInstance, + context: &MutationContext, + ) -> AdminStoreResult { + self.control_plane + .delete_provider_instance( + store_revision(command.expected_config_revision)?, + command.id.as_str(), + mutation_audit( + context, + "delete", + "provider_instance", + command.id.as_str(), + Vec::new(), + ), + ) + .await + .map_err(|error| admin_store_error("provider instance", error)) + .and_then(admin_revision) + } +} + +fn admin_provider_instance( + record: ProviderInstanceRecord, +) -> AdminStoreResult { + Ok(AdminProviderInstance { + id: ProviderInstanceId::new(record.id).map_err(|_| { + admin_store_error( + "provider instance", + StoreError::InvalidData { + entity: ENTITY, + message: "persisted instance ID is invalid".to_owned(), + }, + ) + })?, + provider_kind: ProviderKind::new(record.provider_kind).map_err(|_| { + admin_store_error( + "provider instance", + StoreError::InvalidData { + entity: ENTITY, + message: "persisted provider kind is invalid".to_owned(), + }, + ) + })?, + name: record.name, + base_url: record.base_url, + enabled: record.enabled, + created_at: record.created_at, + updated_at: record.updated_at, + }) +} + pub(crate) async fn insert_provider_instance_in_transaction( transaction: &mut Transaction<'_, Postgres>, instance: &NewProviderInstance, diff --git a/backend/crates/gateway-store/src/postgres/execution.rs b/backend/crates/gateway-store/src/postgres/execution.rs index f30996830..d4c36e20e 100644 --- a/backend/crates/gateway-store/src/postgres/execution.rs +++ b/backend/crates/gateway-store/src/postgres/execution.rs @@ -220,6 +220,8 @@ pub struct ModelRequestFinalization { pub client_response_id: Option, pub upstream_request_id: Option, pub upstream_response_id: Option, + pub upstream_transport: Option, + pub http_version: Option, pub error_kind: Option, pub provider_error_code: Option, pub error_message: Option, @@ -255,6 +257,14 @@ impl ModelRequestFinalization { { return Err(invalid("cost currency must be three uppercase characters")); } + for (field, value) in [ + ("upstream_transport", self.upstream_transport.as_deref()), + ("http_version", self.http_version.as_deref()), + ] { + if let Some(value) = value { + require_nonempty(ENTITY, field, value)?; + } + } self.timings.validate() } } @@ -480,7 +490,9 @@ impl ModelRequestRepository for PgExecutionStore { cost_currency = $23, transport_decision_wait_ms = $24, connect_ms = $25, headers_ms = $26, first_event_ms = $27, first_reasoning_ms = $28, first_text_ms = $29, first_token_ms = $30, provider_processing_ms = $31, - latency_ms = $32, completed_at = $33 + latency_ms = $32, completed_at = $33, + upstream_transport = coalesce($34, upstream_transport), + http_version = coalesce($35, http_version) where id = $1 and outcome = 'running'", ) .bind(&finalization.model_request_id) @@ -555,6 +567,8 @@ impl ModelRequestRepository for PgExecutionStore { )?) .bind(optional_i64(finalization.timings.latency_ms, "latency_ms")?) .bind(finalization.completed_at) + .bind(finalization.upstream_transport) + .bind(finalization.http_version) .execute(&self.pool) .await .map_err(|_| postgres_unavailable("finalize model request"))?; @@ -713,10 +727,13 @@ impl ExecutionStore for PgExecutionStore { provider_account_ref: failure.account_id.as_ref().map(|id| id.as_str().to_owned()), upstream_model_id: Some(failure.upstream_model_id.as_str().to_owned()), failure_kind: error.kind().as_str().to_owned(), - status_code: error.upstream_status(), + status_code: error.upstream_status().or(failure.upstream_status_code), provider_error_code: error.upstream_code().map(|code| code.as_str().to_owned()), retry_after_ms, - upstream_request_id: error.upstream_request_id().map(|id| id.as_str().to_owned()), + upstream_request_id: error + .upstream_request_id() + .map(|id| id.as_str().to_owned()) + .or(failure.upstream_request_id), latency_ms: Some( u64::try_from(failure.latency.as_millis()) .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?, @@ -771,6 +788,8 @@ impl ExecutionStore for PgExecutionStore { client_response_id: finalization.client_response_id, upstream_request_id: finalization.upstream_request_id, upstream_response_id: finalization.upstream_response_id, + upstream_transport: finalization.upstream_transport, + http_version: finalization.http_version, error_kind, provider_error_code: finalization.provider_error_code, error_message, diff --git a/backend/crates/gateway-store/src/postgres/history.rs b/backend/crates/gateway-store/src/postgres/history.rs index 875ea2bee..7955c4dfb 100644 --- a/backend/crates/gateway-store/src/postgres/history.rs +++ b/backend/crates/gateway-store/src/postgres/history.rs @@ -2,9 +2,14 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_core::engine::continuation::{NativeContinuationPort, NativeContinuationStoreError}; +use gateway_core::policy::ClientApiKeyId; use gateway_core::{ engine::{ - continuation::{NativeContinuationPin, NativeContinuationReuse, PreviousResponseId}, + continuation::{ + NativeContinuationPin, NativeContinuationReuse, NativeContinuationScope, + PreviousResponseId, + }, credential::ProviderAccountId, }, error::SafeUpstreamValue, @@ -127,7 +132,6 @@ impl ModelRequestHistoryRepository for PgHistoryRepository { and mr.client_api_key_ref = $2 and mr.outcome = 'succeeded' and mr.downstream_committed_at is not null - and mr.upstream_transport = 'websocket' and mr.upstream_response_id is not null", ) .bind(client_response_id) @@ -140,6 +144,27 @@ impl ModelRequestHistoryRepository for PgHistoryRepository { } } +impl NativeContinuationPort for PgHistoryRepository { + fn resolve<'a>( + &'a self, + client_api_key_id: &'a ClientApiKeyId, + previous_response_id: &'a PreviousResponseId, + ) -> futures::future::BoxFuture< + 'a, + Result, NativeContinuationStoreError>, + > { + Box::pin(async move { + self.resolve_native_continuation_pin( + previous_response_id.as_str(), + client_api_key_id.as_str(), + NativeContinuationReuse::Reusable, + ) + .await + .map_err(|_| NativeContinuationStoreError) + }) + } +} + async fn fetch_history( pool: &PgPool, lookup: &'static str, @@ -239,7 +264,8 @@ fn native_pin_from_row( instance, account, reuse, - )) + ) + .with_scope(NativeContinuationScope::Persisted)) } fn invalid_native_pin(message: &str) -> StoreError { diff --git a/backend/crates/gateway-store/src/postgres/observability.rs b/backend/crates/gateway-store/src/postgres/observability.rs index 1254a7034..08109c60c 100644 --- a/backend/crates/gateway-store/src/postgres/observability.rs +++ b/backend/crates/gateway-store/src/postgres/observability.rs @@ -7,9 +7,16 @@ use std::{ use async_trait::async_trait; use chrono::{DateTime, TimeDelta, Utc}; +use gateway_admin::{ + model::observability as admin_observability, + ports::store::{AdminStoreResult, ObservabilityStore as AdminObservabilityStore}, +}; use sqlx::{PgPool, Postgres, QueryBuilder, Row}; -use crate::{DecimalAmount, StoreError, StoreResult, postgres_unavailable, require_nonempty}; +use crate::{ + DecimalAmount, StoreError, StoreResult, admin_store_error, postgres_unavailable, + require_nonempty, +}; const MAX_PAGE_SIZE: u16 = 100; const MAX_FILTER_BYTES: usize = 256; @@ -113,6 +120,7 @@ impl UsageRecordFilter { (self.operation.as_deref(), "operation filter"), (self.provider_kind.as_deref(), "provider filter"), (self.model.as_deref(), "model filter"), + (self.outcome.as_deref(), "outcome filter"), (self.transport.as_deref(), "transport filter"), (self.response_id.as_deref(), "response ID filter"), ( @@ -123,14 +131,6 @@ impl UsageRecordFilter { validate_optional_text(value, MAX_FILTER_BYTES, field)?; } validate_optional_text(self.search.as_deref(), MAX_SEARCH_BYTES, "search filter")?; - if self.outcome.as_deref().is_some_and(|outcome| { - !matches!( - outcome, - "running" | "succeeded" | "failed" | "cancelled" | "incomplete" - ) - }) { - return Err(invalid("unknown request outcome filter")); - } if self .status_code .is_some_and(|status| !(100..=599).contains(&status)) @@ -367,6 +367,11 @@ pub struct ProviderAccountMetrics { pub total: u64, pub enabled: u64, pub unavailable: u64, + pub active: u64, + pub expired: u64, + pub quota_exhausted: u64, + pub disabled: u64, + pub banned: u64, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -674,6 +679,24 @@ impl PgObservabilityRepository { } } +/// `gateway-admin` 观测端口的 PostgreSQL adapter。 +/// +/// SQL 查询与内部投影继续由 [`PgObservabilityRepository`] 唯一拥有；本类型只负责 +/// Admin UTC 领域模型与持久化投影之间的无格式化转换。 +#[derive(Clone)] +pub struct PgAdminObservabilityStore { + repository: PgObservabilityRepository, +} + +impl PgAdminObservabilityStore { + #[must_use] + pub const fn new(pool: PgPool) -> Self { + Self { + repository: PgObservabilityRepository::new(pool), + } + } +} + #[async_trait] impl ObservabilityRepository for PgObservabilityRepository { async fn dashboard_summary( @@ -683,7 +706,7 @@ impl ObservabilityRepository for PgObservabilityRepository { let filter = UsageRecordFilter::default(); let requests = request_metrics(&self.pool, range, &filter).await?; let attempts = attempt_metrics(&self.pool, range, &filter).await?; - let provider_accounts = provider_account_metrics(&self.pool).await?; + let provider_accounts = provider_account_metrics(&self.pool, range.end).await?; let trend = request_metric_series(&self.pool, range, &filter).await?; let account_usage = provider_account_usage( &self.pool, @@ -777,6 +800,689 @@ impl ObservabilityRepository for PgObservabilityRepository { } } +#[async_trait] +impl AdminObservabilityStore for PgAdminObservabilityStore { + async fn dashboard_summary( + &self, + range: admin_observability::TimeRange, + ) -> AdminStoreResult { + let observation = self + .repository + .dashboard_summary(store_range(range)?) + .await + .map_err(observability_error)?; + admin_dashboard_observation(observation) + } + + async fn dashboard_trend( + &self, + range: admin_observability::TimeRange, + ) -> AdminStoreResult> { + self.repository + .dashboard_trend(store_range(range)?) + .await + .map_err(observability_error)? + .into_iter() + .map(admin_request_metric_point) + .collect() + } + + async fn usage_trend( + &self, + range: admin_observability::TimeRange, + filter: admin_observability::UsageFilter, + ) -> AdminStoreResult> { + self.repository + .usage_trend(store_range(range)?, store_usage_filter(filter)) + .await + .map_err(observability_error)? + .into_iter() + .map(admin_request_metric_point) + .collect() + } + + async fn list_usage_records( + &self, + query: admin_observability::UsageQuery, + ) -> AdminStoreResult { + let page = self + .repository + .list_usage_records(store_usage_query(query)?) + .await + .map_err(observability_error)?; + admin_usage_page(page) + } + + async fn usage_record_detail( + &self, + request_id: &str, + ) -> AdminStoreResult { + let detail = self + .repository + .usage_record_detail(request_id) + .await + .map_err(observability_error)?; + admin_usage_detail(detail) + } + + async fn usage_summary( + &self, + range: admin_observability::TimeRange, + filter: admin_observability::UsageFilter, + ) -> AdminStoreResult { + let overview = self + .repository + .usage_summary(store_range(range)?, store_usage_filter(filter)) + .await + .map_err(observability_error)?; + admin_usage_overview(overview) + } + + async fn usage_diagnostics( + &self, + range: admin_observability::TimeRange, + filter: admin_observability::UsageFilter, + dimension: admin_observability::DiagnosticDimension, + ) -> AdminStoreResult> { + self.repository + .usage_diagnostics( + store_range(range)?, + store_usage_filter(filter), + store_diagnostic_dimension(dimension), + ) + .await + .map_err(observability_error)? + .into_iter() + .map(admin_diagnostic_observation) + .collect() + } + + async fn list_ops_errors( + &self, + query: admin_observability::OpsErrorQuery, + ) -> AdminStoreResult { + let page = self + .repository + .list_ops_errors(store_ops_error_query(query)?) + .await + .map_err(observability_error)?; + admin_ops_error_page(page) + } +} + +fn store_range(range: admin_observability::TimeRange) -> AdminStoreResult { + ObservabilityRange::new(range.start, range.end).map_err(observability_error) +} + +fn store_usage_filter(filter: admin_observability::UsageFilter) -> UsageRecordFilter { + UsageRecordFilter { + client_api_key_ref: filter.client_api_key_ref, + request_id: filter.request_id, + provider_account_ref: filter.provider_account_ref, + operation: filter.operation, + provider_kind: filter.provider_kind, + model: filter.model, + outcome: filter.outcome.map(store_request_outcome), + status_code: filter.status_code, + transport: filter.transport, + attempt_index: filter.attempt_index, + response_id: filter.response_id, + upstream_request_id: filter.upstream_request_id, + search: filter.search, + } +} + +fn store_request_outcome(outcome: admin_observability::RequestOutcome) -> String { + outcome.as_str().to_owned() +} + +fn store_cursor( + cursor: admin_observability::ObservabilityCursor, +) -> AdminStoreResult { + ObservabilityCursor::new(cursor.observed_at, cursor.stable_id).map_err(observability_error) +} + +fn store_usage_query(query: admin_observability::UsageQuery) -> AdminStoreResult { + Ok(UsageRecordQuery { + range: store_range(query.range)?, + filter: store_usage_filter(query.filter), + cursor: query.cursor.map(store_cursor).transpose()?, + page: ObservabilityPageNumber::new(query.page.get()).map_err(observability_error)?, + page_size: ObservabilityPageSize::new(query.page_size.get()) + .map_err(observability_error)?, + }) +} + +fn store_ops_error_filter(filter: admin_observability::OpsErrorFilter) -> OpsErrorFilter { + OpsErrorFilter { + client_api_key_ref: filter.client_api_key_ref, + request_id: filter.request_id, + provider_account_ref: filter.provider_account_ref, + provider_kind: filter.provider_kind, + operation: filter.operation, + model: filter.model, + transport: filter.transport, + attempt_index: filter.attempt_index, + response_id: filter.response_id, + upstream_request_id: filter.upstream_request_id, + failure_kind: filter.failure_kind, + status_code: filter.status_code, + search: filter.search, + } +} + +fn store_ops_error_query( + query: admin_observability::OpsErrorQuery, +) -> AdminStoreResult { + Ok(OpsErrorQuery { + range: store_range(query.range)?, + filter: store_ops_error_filter(query.filter), + cursor: query.cursor.map(store_cursor).transpose()?, + page: ObservabilityPageNumber::new(query.page.get()).map_err(observability_error)?, + page_size: ObservabilityPageSize::new(query.page_size.get()) + .map_err(observability_error)?, + }) +} + +const fn store_diagnostic_dimension( + dimension: admin_observability::DiagnosticDimension, +) -> DiagnosticDimension { + match dimension { + admin_observability::DiagnosticDimension::Provider => DiagnosticDimension::Provider, + admin_observability::DiagnosticDimension::Model => DiagnosticDimension::Model, + admin_observability::DiagnosticDimension::Account => DiagnosticDimension::Account, + admin_observability::DiagnosticDimension::ApiKey => DiagnosticDimension::ApiKey, + admin_observability::DiagnosticDimension::Transport => DiagnosticDimension::Transport, + admin_observability::DiagnosticDimension::Failure => DiagnosticDimension::Failure, + admin_observability::DiagnosticDimension::Status => DiagnosticDimension::Status, + } +} + +fn admin_dashboard_observation( + observation: DashboardObservation, +) -> AdminStoreResult { + let DashboardObservation { + range, + requests, + attempts, + provider_accounts, + trend, + account_usage, + recent_requests, + } = observation; + Ok(admin_observability::DashboardObservation { + range: admin_range(range), + requests: admin_request_metrics(requests)?, + attempts: admin_attempt_metrics(attempts)?, + provider_accounts: admin_account_pool_metrics(provider_accounts), + trend: trend + .into_iter() + .map(admin_request_metric_point) + .collect::>()?, + account_usage: account_usage + .into_iter() + .map(admin_dashboard_account_usage) + .collect::>()?, + recent_requests: recent_requests + .into_iter() + .map(admin_usage_record) + .collect::>()?, + }) +} + +const fn admin_range(range: ObservabilityRange) -> admin_observability::TimeRange { + admin_observability::TimeRange { + start: range.start, + end: range.end, + } +} + +fn admin_request_metrics( + metrics: RequestMetrics, +) -> AdminStoreResult { + Ok(admin_observability::RequestMetrics { + request_count: metrics.request_count, + success_count: metrics.success_count, + failure_count: metrics.failure_count, + cancelled_count: metrics.cancelled_count, + incomplete_count: metrics.incomplete_count, + caller_error_count: metrics.caller_error_count, + input_tokens: metrics.input_tokens, + output_tokens: metrics.output_tokens, + cached_tokens: metrics.cached_tokens, + cache_write_tokens: metrics.cache_write_tokens, + reasoning_tokens: metrics.reasoning_tokens, + total_tokens: metrics.total_tokens, + first_token_latency_sum_ms: metrics.first_token_latency_sum, + first_token_latency_count: metrics.first_token_latency_count, + latency_sum_ms: metrics.latency_sum, + latency_count: metrics.latency_count, + min_latency_ms: metrics.min_latency_ms, + max_latency_ms: metrics.max_latency_ms, + latency_percentiles: admin_latency_percentiles(metrics.latency_percentiles)?, + first_token_latency_percentiles: admin_latency_percentiles( + metrics.first_token_latency_percentiles, + )?, + cache_eligible_request_count: metrics.cache_eligible_request_count, + cache_hit_request_count: metrics.cache_hit_request_count, + }) +} + +fn admin_latency_percentiles( + percentiles: LatencyPercentiles, +) -> AdminStoreResult { + Ok(admin_observability::LatencyPercentiles { + p50_ms: percentiles + .p50_ms + .map(|value| admin_observability::PercentileMilliseconds::new(value.as_f64())) + .transpose() + .map_err(|_| invalid_admin_percentile())?, + p95_ms: percentiles + .p95_ms + .map(|value| admin_observability::PercentileMilliseconds::new(value.as_f64())) + .transpose() + .map_err(|_| invalid_admin_percentile())?, + p99_ms: percentiles + .p99_ms + .map(|value| admin_observability::PercentileMilliseconds::new(value.as_f64())) + .transpose() + .map_err(|_| invalid_admin_percentile())?, + }) +} + +fn invalid_admin_percentile() -> gateway_admin::ports::store::AdminStoreError { + observability_error(StoreError::InvalidData { + entity: "observability latency percentile", + message: "latency percentile is outside the admin contract".to_owned(), + }) +} + +fn admin_attempt_metrics( + metrics: AttemptMetrics, +) -> AdminStoreResult { + Ok(admin_observability::AttemptMetrics { + attempt_count: metrics.attempt_count, + success_count: metrics.success_count, + failure_count: metrics.failure_count, + cancelled_count: metrics.cancelled_count, + incomplete_count: metrics.incomplete_count, + rate_limited_count: metrics.rate_limited_count, + auth_failure_count: metrics.auth_failure_count, + provider_5xx_count: metrics.provider_5xx_count, + cost_coverage: admin_cost_coverage(metrics.cost_coverage), + costs: admin_currency_costs(metrics.costs)?, + }) +} + +const fn admin_cost_coverage(coverage: CostCoverage) -> admin_observability::CostCoverage { + admin_observability::CostCoverage { + provider_reported_count: coverage.provider_reported_count, + calculated_count: coverage.calculated_count, + partial_count: 0, + unavailable_count: coverage.unavailable_count, + not_billable_count: 0, + } +} + +fn admin_currency_costs( + costs: Vec, +) -> AdminStoreResult> { + costs + .into_iter() + .map(|cost| { + Ok(admin_observability::CurrencyCost { + currency: cost.currency, + amount: admin_decimal_amount(cost.amount)?, + }) + }) + .collect() +} + +fn admin_decimal_amount( + amount: DecimalAmount, +) -> AdminStoreResult { + admin_observability::DecimalAmount::from_str(amount.as_str()).map_err(|_| { + observability_error(StoreError::InvalidData { + entity: "observability decimal amount", + message: "amount is outside the admin numeric contract".to_owned(), + }) + }) +} + +fn admin_optional_decimal_amount( + amount: Option, +) -> AdminStoreResult> { + amount.map(admin_decimal_amount).transpose() +} + +const fn admin_account_pool_metrics( + metrics: ProviderAccountMetrics, +) -> admin_observability::AccountPoolMetrics { + admin_observability::AccountPoolMetrics { + total: metrics.total, + enabled: metrics.enabled, + unavailable: metrics.unavailable, + active: metrics.active, + expired: metrics.expired, + quota_exhausted: metrics.quota_exhausted, + refreshing: None, + disabled: metrics.disabled, + banned: metrics.banned, + } +} + +fn admin_dashboard_account_usage( + usage: ProviderAccountUsageObservation, +) -> AdminStoreResult { + Ok(admin_observability::DashboardAccountUsage { + account_id: usage.account_id, + provider_instance_id: usage.provider_instance_id, + provider_kind: usage.provider_kind, + name: usage.name, + email: usage.email, + plan_type: usage.plan_type, + enabled: usage.enabled, + availability: usage.availability, + request_count: usage.request_count, + success_count: usage.success_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + cache_write_tokens: usage.cache_write_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + cost_coverage: admin_cost_coverage(usage.cost_coverage), + costs: admin_currency_costs(usage.costs)?, + last_used_at: usage.last_used_at, + models: usage + .models + .into_iter() + .map(admin_dashboard_account_model_usage) + .collect::>()?, + }) +} + +fn admin_dashboard_account_model_usage( + usage: ProviderAccountModelUsageObservation, +) -> AdminStoreResult { + Ok(admin_observability::DashboardAccountModelUsage { + model: usage.model, + request_count: usage.request_count, + success_count: usage.success_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + cache_write_tokens: usage.cache_write_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + cost_coverage: admin_cost_coverage(usage.cost_coverage), + costs: admin_currency_costs(usage.costs)?, + last_used_at: usage.last_used_at, + }) +} + +fn admin_request_metric_point( + point: RequestMetricPoint, +) -> AdminStoreResult { + Ok(admin_observability::RequestMetricPoint { + bucket_start: point.bucket_start, + granularity: admin_granularity(point.granularity), + metrics: admin_request_metrics(point.metrics)?, + cost_coverage: admin_cost_coverage(point.cost_coverage), + costs: admin_currency_costs(point.costs)?, + }) +} + +const fn admin_granularity( + granularity: ObservationGranularity, +) -> admin_observability::Granularity { + match granularity { + ObservationGranularity::FifteenMinutes => admin_observability::Granularity::FifteenMinutes, + ObservationGranularity::Hour => admin_observability::Granularity::Hour, + ObservationGranularity::Day => admin_observability::Granularity::Day, + } +} + +fn admin_usage_page(page: UsageRecordPage) -> AdminStoreResult { + Ok(admin_observability::UsagePage { + items: page + .items + .into_iter() + .map(admin_usage_record) + .collect::>()?, + total: page.total, + next_cursor: page.next_cursor.map(admin_cursor), + }) +} + +fn admin_cursor(cursor: ObservabilityCursor) -> admin_observability::ObservabilityCursor { + admin_observability::ObservabilityCursor { + observed_at: cursor.observed_at, + stable_id: cursor.stable_id, + } +} + +fn admin_usage_record(record: UsageRecord) -> AdminStoreResult { + let billing = match (&record.cost_amount, &record.cost_currency) { + (Some(amount), Some(currency)) => Some(admin_observability::UsageBilling::Total { + source: record.cost_source.clone(), + total: admin_observability::CurrencyCost { + currency: currency.clone(), + amount: admin_decimal_amount(amount.clone())?, + }, + }), + (None, None) => None, + _ => { + return Err(observability_error(StoreError::InvalidData { + entity: "observability request billing", + message: "cost amount and currency must be present together".to_owned(), + })); + } + }; + Ok(admin_observability::UsageRecord { + id: record.id, + client_api_key_ref: record.client_api_key_ref, + config_revision: record.config_revision, + protocol: record.protocol, + operation: record.operation, + endpoint: record.endpoint, + client_transport: record.client_transport, + requested_model_id: record.requested_model_id, + input_token_estimate: record.input_token_estimate, + provider_instance_id: record.provider_instance_id, + provider_kind: record.provider_kind, + provider_account_ref: record.provider_account_ref, + provider_account_name: record.provider_account_name, + provider_account_email: record.provider_account_email, + upstream_model_id: record.upstream_model_id, + upstream_transport: record.upstream_transport, + http_version: record.http_version, + attempt_count: record.attempt_count, + upstream_send_state: record.upstream_send_state, + downstream_committed_at: record.downstream_committed_at, + outcome: admin_request_outcome(&record.outcome)?, + client_status_code: record.client_status_code, + upstream_status_code: record.upstream_status_code, + client_response_id: record.client_response_id, + upstream_request_id: record.upstream_request_id, + upstream_response_id: record.upstream_response_id, + error_kind: record.error_kind, + provider_error_code: record.provider_error_code, + error_message: record.error_message, + retry_after_ms: record.retry_after_ms, + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cached_tokens: record.cached_tokens, + cache_write_tokens: record.cache_write_tokens, + reasoning_tokens: record.reasoning_tokens, + total_tokens: record.total_tokens, + cost_source: record.cost_source, + cost_amount: admin_optional_decimal_amount(record.cost_amount)?, + cost_currency: record.cost_currency, + billing, + transport_decision_wait_ms: record.transport_decision_wait_ms, + connect_ms: record.connect_ms, + headers_ms: record.headers_ms, + first_event_ms: record.first_event_ms, + first_reasoning_ms: record.first_reasoning_ms, + first_text_ms: record.first_text_ms, + first_token_ms: record.first_token_ms, + provider_processing_ms: record.provider_processing_ms, + latency_ms: record.latency_ms, + client_ip: record.client_ip, + user_agent: record.user_agent, + reasoning_effort: record.reasoning_effort, + reasoning_preset: record.reasoning_preset, + request_kind: record.request_kind, + subagent_kind: record.subagent_kind, + compact: record.compact, + started_at: record.started_at, + deadline_at: record.deadline_at, + completed_at: record.completed_at, + }) +} + +fn admin_request_outcome(outcome: &str) -> AdminStoreResult { + admin_observability::RequestOutcome::new(outcome.to_owned()).map_err(|_| { + observability_error(StoreError::InvalidData { + entity: "observability request outcome", + message: "invalid request outcome".to_owned(), + }) + }) +} + +fn admin_usage_detail( + detail: UsageRecordDetail, +) -> AdminStoreResult { + Ok(admin_observability::UsageDetail { + request: admin_usage_record(detail.request)?, + attempts: detail + .attempts + .into_iter() + .map(admin_usage_attempt) + .collect::>()?, + }) +} + +fn admin_usage_attempt( + attempt: UsageAttemptObservation, +) -> AdminStoreResult { + Ok(admin_observability::UsageAttempt { + source: attempt.source, + id: attempt.id, + attempt_index: attempt.attempt_index, + component: attempt.component, + operation: attempt.operation, + provider_instance_id: attempt.provider_instance_id, + provider_kind: attempt.provider_kind, + provider_account_ref: attempt.provider_account_ref, + upstream_model_id: attempt.upstream_model_id, + upstream_transport: attempt.upstream_transport, + upstream_send_state: attempt.upstream_send_state, + outcome: admin_request_outcome(&attempt.outcome)?, + downstream_committed: attempt.downstream_committed, + status_code: attempt.status_code, + provider_error_code: attempt.provider_error_code, + failure_kind: attempt.failure_kind, + retry_after_ms: attempt.retry_after_ms, + upstream_request_id: attempt.upstream_request_id, + latency_ms: attempt.latency_ms, + message: attempt.message, + input_tokens: attempt.input_tokens, + output_tokens: attempt.output_tokens, + cached_tokens: attempt.cached_tokens, + cache_write_tokens: attempt.cache_write_tokens, + reasoning_tokens: attempt.reasoning_tokens, + total_tokens: attempt.total_tokens, + cost_source: attempt.cost_source, + cost_amount: admin_optional_decimal_amount(attempt.cost_amount)?, + cost_currency: attempt.cost_currency, + occurred_at: attempt.occurred_at, + }) +} + +fn admin_usage_overview( + overview: UsageOverview, +) -> AdminStoreResult { + Ok(admin_observability::UsageOverview { + range: admin_range(overview.range), + requests: admin_request_metrics(overview.requests)?, + attempts: admin_attempt_metrics(overview.attempts)?, + providers: overview + .providers + .into_iter() + .map(admin_provider_observation) + .collect(), + }) +} + +fn admin_provider_observation( + observation: ProviderObservation, +) -> admin_observability::ProviderObservation { + admin_observability::ProviderObservation { + provider_kind: observation.provider_kind, + request_count: observation.request_count, + attempt_count: observation.attempt_count, + failure_count: observation.failure_count, + total_tokens: observation.total_tokens, + } +} + +fn admin_diagnostic_observation( + observation: DiagnosticObservation, +) -> AdminStoreResult { + Ok(admin_observability::DiagnosticObservation { + name: observation.name, + request_count: observation.request_count, + success_count: observation.success_count, + failure_count: observation.failure_count, + attempt_count: observation.attempt_count, + total_tokens: observation.total_tokens, + average_latency_ms: observation.average_latency_ms, + cost_coverage: admin_cost_coverage(observation.cost_coverage), + costs: admin_currency_costs(observation.costs)?, + }) +} + +fn admin_ops_error_page(page: OpsErrorPage) -> AdminStoreResult { + Ok(admin_observability::OpsErrorPage { + items: page.items.into_iter().map(admin_ops_error).collect(), + total: page.total, + next_cursor: page.next_cursor.map(admin_cursor), + }) +} + +fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::OpsError { + admin_observability::OpsError { + source: error.source, + event_id: error.event_id, + request_id: error.request_id, + attempt_index: error.attempt_index, + client_api_key_ref: error.client_api_key_ref, + component: error.component, + operation: error.operation, + provider_instance_id: error.provider_instance_id, + provider_kind: error.provider_kind, + provider_account_ref: error.provider_account_ref, + upstream_model_id: error.upstream_model_id, + upstream_transport: error.upstream_transport, + failure_kind: error.failure_kind, + status_code: error.status_code, + provider_error_code: error.provider_error_code, + client_response_id: error.client_response_id, + upstream_request_id: error.upstream_request_id, + latency_ms: error.latency_ms, + message: error.message, + occurrence_count: error.occurrence_count, + occurred_at: error.occurred_at, + stable_sort_id: error.stable_sort_id, + } +} + +fn observability_error(error: StoreError) -> gateway_admin::ports::store::AdminStoreError { + admin_store_error("observability", error) +} + async fn request_metrics( pool: &PgPool, range: ObservabilityRange, @@ -1058,14 +1764,28 @@ async fn request_costs( .collect() } -async fn provider_account_metrics(pool: &PgPool) -> StoreResult { +async fn provider_account_metrics( + pool: &PgPool, + observed_at: DateTime, +) -> StoreResult { let row = sqlx::query( "select count(*)::bigint as total, count(*) filter (where enabled)::bigint as enabled, count(*) filter (where not enabled or availability <> 'ready')::bigint - as unavailable + as unavailable, + count(*) filter ( + where enabled and availability = 'ready' and access_token_expires_at > $1 + )::bigint as active, + count(*) filter ( + where availability = 'expired' or access_token_expires_at <= $1 + )::bigint as expired, + count(*) filter (where availability = 'quota_exhausted')::bigint + as quota_exhausted, + count(*) filter (where not enabled)::bigint as disabled, + count(*) filter (where availability = 'banned')::bigint as banned from provider_accounts", ) + .bind(observed_at) .fetch_one(pool) .await .map_err(|_| postgres_unavailable("load provider account metrics"))?; @@ -1073,6 +1793,11 @@ async fn provider_account_metrics(pool: &PgPool) -> StoreResult>, pub enabled: bool, pub availability: ProviderAccountAvailability, + pub availability_reason: Option, pub cooldown_until: Option>, pub availability_observed_at: DateTime, } @@ -853,6 +887,943 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository { } } +/// Admin 账号用例所需的公共账号、留存观测与 revision 事务能力。 +/// +/// 三个 PostgreSQL adapter 都保持私有，调用方只能取得 [`AccountStore`] 暴露的领域能力。 +#[derive(Clone)] +pub struct PgAdminAccountStore { + accounts: PgProviderAccountRepository, + observability: PgObservabilityRepository, + control_plane: PgControlPlaneRepository, +} + +impl PgAdminAccountStore { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + accounts: PgProviderAccountRepository::new(pool.clone()), + observability: PgObservabilityRepository::new(pool.clone()), + control_plane: PgControlPlaneRepository::new(pool), + } + } + + async fn usage_observations( + &self, + range: ObservabilityRange, + account_ids: &[String], + ) -> AdminStoreResult> { + if account_ids.is_empty() { + return Ok(Vec::new()); + } + let mut observations = Vec::with_capacity(account_ids.len()); + for account_ids in account_ids.chunks(ADMIN_USAGE_CHUNK_SIZE) { + let query = ProviderAccountUsageQuery::for_accounts(range, account_ids.to_vec()) + .map_err(|error| admin_store_error(ENTITY, error))?; + observations.extend( + self.observability + .provider_account_usage(query) + .await + .map_err(|error| admin_store_error(ENTITY, error))?, + ); + } + Ok(observations) + } + + async fn required_scope( + &self, + account_id: &str, + ) -> AdminStoreResult { + let record = self + .accounts + .load_provider_account(account_id) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .ok_or_else(|| { + admin_store_error( + ENTITY, + StoreError::NotFound { + entity: ENTITY, + id: account_id.to_owned(), + }, + ) + })?; + Ok(ProviderAccountAdminScope { + provider_kind: record.summary.provider_kind, + provider_instance_id: record.summary.provider_instance_id, + }) + } + + async fn commit_prepared_import( + &self, + expected_config_revision: AdminRevision, + prepared: PreparedCredentialImport, + context: &MutationContext, + action: &str, + ) -> AdminStoreResult { + let provider_kind = prepared.provider_kind.as_str().to_owned(); + let provider_instance_id = prepared.provider_instance_id.as_str().to_owned(); + let credential_ids = prepared + .credentials + .iter() + .map(|credential| credential.account_id.clone()) + .collect::>(); + let accounts = prepared + .credentials + .into_iter() + .map(prepared_account) + .collect::>>() + .map_err(|error| admin_store_error(ENTITY, error))?; + let revision = self + .accounts + .import_provider_accounts( + store_revision(expected_config_revision)?, + ImportProviderAccounts { + scope: ProviderAccountAdminScope { + provider_kind, + provider_instance_id: provider_instance_id.clone(), + }, + accounts, + audit: mutation_audit( + context, + action, + "provider_account", + &provider_instance_id, + vec!["credentials".to_owned()], + ), + }, + ) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok(CredentialImportResult { + config_revision: admin_revision(revision)?, + credential_ids, + }) + } + + async fn commit_prepared_rotation( + &self, + expected_config_revision: AdminRevision, + prepared: PreparedCredentialRotationFacts, + context: &MutationContext, + action: &str, + ) -> AdminStoreResult { + let account_id = prepared.account_id.clone(); + let scope = ProviderAccountAdminScope { + provider_kind: prepared.provider_kind.as_str().to_owned(), + provider_instance_id: prepared.provider_instance_id.as_str().to_owned(), + }; + let rotation = self + .accounts + .rotate_provider_account( + store_revision(expected_config_revision)?, + RotateProviderAccount { + scope, + profile: UpdateProviderAccount { + id: account_id.as_str().to_owned(), + name: prepared.name, + email: prepared.email, + plan_type: prepared.plan_type, + }, + credential: ProviderCredentialUpdate { + account_id: account_id.as_str().to_owned(), + expected_revision: store_revision(prepared.expected_credential_revision)?, + provider_credentials_json: provider_document_json( + prepared.provider_material, + ) + .map_err(|error| admin_store_error(ENTITY, error))?, + has_refresh_token: prepared.has_refresh_token, + access_token_expires_at: prepared.access_token_expires_at, + next_refresh_at: prepared.next_refresh_at, + }, + audit: mutation_audit( + context, + action, + "provider_account", + account_id.as_str(), + vec!["credentials".to_owned()], + ), + }, + ) + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + Ok(CredentialMutationResult { + config_revision: admin_revision(rotation.config_revision)?, + account_id, + credential_revision: Some(admin_revision(rotation.credential_revision)?), + }) + } +} + +#[async_trait] +impl AccountStore for PgAdminAccountStore { + async fn list_accounts(&self, query: AdminAccountListQuery) -> AdminStoreResult { + if query.page == 0 { + return Err(AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "page number must be positive", + )); + } + let (control_plane, accounts) = futures::try_join!( + self.control_plane.load_control_plane(), + self.accounts.list_provider_accounts(None, true), + ) + .map_err(|error| admin_store_error(ENTITY, error))?; + let now = Utc::now(); + let summary = admin_account_summary(&accounts, now); + let mut items = accounts + .into_iter() + .filter_map(|account| { + let status = admin_account_status(&account, now); + account_matches_admin_query(&account, status, &query).then_some( + AdminAccountListItem { + account, + status, + usage: None, + }, + ) + }) + .collect::>(); + + if query.sort.is_some_and(|sort| { + matches!( + sort.field, + AdminAccountSortField::Usage | AdminAccountSortField::LastUsedAt + ) + }) { + let range = retained_usage_range(control_plane.settings.usage_retention_days, now)?; + let account_ids = items + .iter() + .map(|item| item.account.id.clone()) + .collect::>(); + let mut usage_by_account = self + .usage_observations(range, &account_ids) + .await? + .into_iter() + .map(|usage| (usage.account_id.clone(), usage)) + .collect::>(); + for item in &mut items { + item.usage = usage_by_account.remove(&item.account.id); + } + } + sort_admin_account_items(&mut items, query.sort); + + let total = u64::try_from(items.len()).unwrap_or(u64::MAX); + let page_size = usize::from(query.page_size.get()); + let offset = u64::from(query.page - 1).saturating_mul(u64::from(query.page_size.get())); + let offset = usize::try_from(offset).unwrap_or(usize::MAX); + let items = items + .into_iter() + .skip(offset) + .take(page_size) + .map(|item| admin_account_record(item.account)) + .collect::>>()?; + Ok(AccountPage { + config_revision: admin_revision(control_plane.settings.config_revision)?, + items, + total, + summary, + }) + } + + async fn load_account(&self, account_id: &str) -> AdminStoreResult> { + self.accounts + .load_provider_account(account_id) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .map(|record| admin_account_record(record.summary)) + .transpose() + } + + async fn load_account_usage( + &self, + range: TimeRange, + account_ids: &[String], + ) -> AdminStoreResult> { + let range = ObservabilityRange::new(range.start, range.end) + .map_err(|error| admin_store_error(ENTITY, error))?; + self.usage_observations(range, account_ids) + .await? + .into_iter() + .map(admin_account_usage) + .collect() + } + + async fn list_credentials( + &self, + provider_kind: &ProviderKind, + query: CredentialListQuery, + ) -> AdminStoreResult { + let provider_instance_id = query + .provider_instance_id + .as_ref() + .map(ProviderInstanceId::as_str); + let (control_plane, accounts) = futures::try_join!( + self.control_plane.load_control_plane(), + self.accounts + .list_provider_accounts(provider_instance_id, true), + ) + .map_err(|error| admin_store_error(ENTITY, error))?; + let mut accounts = accounts + .into_iter() + .filter(|account| account.provider_kind == provider_kind.as_str()) + .filter(|account| { + query.availability.as_ref().is_none_or(|expected| { + expected.matches(admin_account_availability(account.availability)) + }) + }) + .filter(|account| { + query + .enabled + .is_none_or(|enabled| account.enabled == enabled) + }) + .filter(|account| { + let CredentialListWindow::Page { + cursor: Some(cursor), + .. + } = &query.window + else { + return true; + }; + account.created_at > cursor.created_at + || (account.created_at == cursor.created_at + && account.id.as_str() > cursor.account_id.as_str()) + }) + .collect::>(); + if matches!(&query.window, CredentialListWindow::Page { .. }) { + accounts.sort_by(|left, right| { + left.created_at + .cmp(&right.created_at) + .then_with(|| left.id.cmp(&right.id)) + }); + } + let next_cursor = match query.window { + CredentialListWindow::All => None, + CredentialListWindow::Page { page_size, .. } => { + let page_size = usize::from(page_size.get()); + let has_more = accounts.len() > page_size; + accounts.truncate(page_size); + has_more + .then(|| accounts.last()) + .flatten() + .map(|account| { + Ok(CredentialCursor { + created_at: account.created_at, + account_id: CoreProviderAccountId::new(account.id.clone()).map_err( + |_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "persisted Provider account ID is invalid", + ) + }, + )?, + }) + }) + .transpose()? + } + }; + Ok(CredentialPage { + config_revision: admin_revision(control_plane.settings.config_revision)?, + items: accounts + .into_iter() + .map(admin_account_record) + .collect::>>()?, + next_cursor, + }) + } + + async fn credential_details( + &self, + provider_kind: &ProviderKind, + account_id: &CoreProviderAccountId, + ) -> AdminStoreResult> { + let (control_plane, account) = futures::try_join!( + self.control_plane.load_control_plane(), + self.accounts.load_provider_account(account_id.as_str()), + ) + .map_err(|error| admin_store_error(ENTITY, error))?; + account + .filter(|record| record.summary.provider_kind == provider_kind.as_str()) + .map(|record| { + Ok(CredentialDetails { + config_revision: admin_revision(control_plane.settings.config_revision)?, + credential: admin_account_record(record.summary)?, + }) + }) + .transpose() + } + + async fn load_credentials_for_export( + &self, + provider_kind: &ProviderKind, + account_ids: &[CoreProviderAccountId], + ) -> AdminStoreResult> { + let ids = account_ids + .iter() + .map(|id| id.as_str().to_owned()) + .collect::>(); + validate_admin_account_ids(&ids).map_err(|error| admin_store_error(ENTITY, error))?; + let mut credentials = Vec::with_capacity(account_ids.len()); + for account_id in account_ids { + let record = self + .accounts + .load_provider_account(account_id.as_str()) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .ok_or_else(|| { + AdminStoreError::new( + AdminStoreErrorKind::NotFound, + ENTITY, + "one or more exported credentials do not exist", + ) + })?; + if record.summary.provider_kind != provider_kind.as_str() { + return Err(AdminStoreError::new( + AdminStoreErrorKind::NotFound, + ENTITY, + "one or more exported credentials belong to another Provider", + )); + } + credentials.push(ProviderExportCredentialInput { + account: admin_account_record(record.summary)?, + provider_material: ProviderDocument::new(OpaqueProviderData::new( + record.provider_credentials_json.fields().clone(), + )), + }); + } + Ok(credentials) + } + + async fn commit_credential_import( + &self, + command: CredentialImportCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.commit_prepared_import( + command.expected_config_revision, + command.prepared, + context, + "import_document", + ) + .await + } + + async fn commit_authorization( + &self, + command: AuthorizationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + let expected_config_revision = command.pending.expected_config_revision(); + match command.credential { + AuthorizationCredentialCommit::Create(credential) => { + let account_id = credential.account_id.clone(); + let result = self + .commit_prepared_import( + expected_config_revision, + PreparedCredentialImport { + provider_kind: credential.provider_kind.clone(), + provider_instance_id: credential.provider_instance_id.clone(), + credentials: vec![credential], + }, + context, + "authorize", + ) + .await?; + let details = self + .accounts + .load_provider_account(account_id.as_str()) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .ok_or_else(|| { + AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + ENTITY, + "authorized credential was not visible after commit", + ) + })?; + Ok(CredentialMutationResult { + config_revision: result.config_revision, + account_id, + credential_revision: Some(admin_revision(details.summary.credential_revision)?), + }) + } + AuthorizationCredentialCommit::Reauthorize(prepared) => { + self.commit_prepared_rotation( + expected_config_revision, + prepared, + context, + "reauthorize", + ) + .await + } + } + } + + async fn commit_credential_rotation( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.commit_prepared_rotation( + command.expected_config_revision, + command.prepared, + context, + "rotate_credential", + ) + .await + } + + async fn commit_credential_refresh( + &self, + command: CredentialRotationCommit, + context: &MutationContext, + ) -> AdminStoreResult { + self.commit_prepared_rotation( + command.expected_config_revision, + command.prepared, + context, + "refresh_credential", + ) + .await + } + + async fn set_account_enabled( + &self, + command: SetAccountEnabled, + context: &MutationContext, + ) -> AdminStoreResult { + let scope = self.required_scope(&command.account_id).await?; + let action = if command.enabled { "enable" } else { "disable" }; + self.accounts + .set_provider_account_enabled_admin( + store_revision(command.expected_config_revision)?, + SetProviderAccountEnabled { + scope, + account_id: command.account_id.clone(), + enabled: command.enabled, + audit: mutation_audit( + context, + action, + "provider_account", + &command.account_id, + vec!["enabled".to_owned()], + ), + }, + ) + .await + .map_err(|error| admin_store_error(ENTITY, error)) + .and_then(admin_revision) + } + + async fn delete_account( + &self, + command: DeleteAccount, + context: &MutationContext, + ) -> AdminStoreResult { + let scope = self.required_scope(&command.account_id).await?; + self.accounts + .delete_provider_accounts_admin( + store_revision(command.expected_config_revision)?, + DeleteProviderAccounts { + scope, + account_ids: vec![command.account_id.clone()], + audit: mutation_audit( + context, + "delete", + "provider_account", + &command.account_id, + Vec::new(), + ), + }, + ) + .await + .map_err(|error| admin_store_error(ENTITY, error)) + .and_then(admin_revision) + } + + async fn record_credential_export( + &self, + account_ids: &[CoreProviderAccountId], + context: &MutationContext, + ) -> AdminStoreResult<()> { + let ids = account_ids + .iter() + .map(|id| id.as_str().to_owned()) + .collect::>(); + validate_admin_account_ids(&ids).map_err(|error| admin_store_error(ENTITY, error))?; + for account_id in &ids { + if self + .accounts + .load_provider_account(account_id) + .await + .map_err(|error| admin_store_error(ENTITY, error))? + .is_none() + { + return Err(AdminStoreError::new( + AdminStoreErrorKind::NotFound, + ENTITY, + "one or more exported credentials do not exist", + )); + } + } + let control_plane = self + .control_plane + .load_control_plane() + .await + .map_err(|error| admin_store_error(ENTITY, error))?; + let revision = control_plane.settings.config_revision; + let mut transaction = self.accounts.pool.begin().await.map_err(|_| { + admin_store_error( + ENTITY, + postgres_unavailable("begin credential export audit"), + ) + })?; + let result = async { + for account_id in &ids { + append_admin_audit_event_in_transaction( + &mut transaction, + mutation_audit( + context, + "export_credentials", + "provider_account", + account_id, + Vec::new(), + ), + revision, + ) + .await?; + } + Ok(()) + } + .await; + finish_admin_transaction(transaction, result, "credential export audit") + .await + .map_err(|error| admin_store_error(ENTITY, error)) + } +} + +struct AdminAccountListItem { + account: ProviderAccountSummary, + status: AdminAccountStatus, + usage: Option, +} + +fn retained_usage_range( + retention_days: u32, + now: DateTime, +) -> AdminStoreResult { + let duration = TimeDelta::try_days(i64::from(retention_days)).ok_or_else(|| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "usage retention is outside the supported range", + ) + })?; + let start = now.checked_sub_signed(duration).ok_or_else(|| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "usage retention start is outside the supported range", + ) + })?; + ObservabilityRange::new(start, now).map_err(|error| admin_store_error(ENTITY, error)) +} + +fn account_matches_admin_query( + account: &ProviderAccountSummary, + status: AdminAccountStatus, + query: &AdminAccountListQuery, +) -> bool { + let provider_matches = query + .provider_kind + .as_ref() + .is_none_or(|provider| account.provider_kind == provider.as_str()); + let search_matches = query.search.as_ref().is_none_or(|search| { + let search = search.to_lowercase(); + [ + Some(account.id.as_str()), + Some(account.name.as_str()), + account.email.as_deref(), + account.upstream_account_id.as_deref(), + Some(account.upstream_user_id.as_str()), + ] + .into_iter() + .flatten() + .any(|value| value.to_lowercase().contains(&search)) + }); + let status_matches = query.status.is_none_or(|expected| expected == status); + provider_matches && search_matches && status_matches +} + +fn admin_account_status( + account: &ProviderAccountSummary, + now: DateTime, +) -> AdminAccountStatus { + if !account.enabled { + AdminAccountStatus::Disabled + } else if account.availability == ProviderAccountAvailability::Banned { + AdminAccountStatus::Banned + } else if account.availability == ProviderAccountAvailability::QuotaExhausted { + AdminAccountStatus::QuotaExhausted + } else if account.availability == ProviderAccountAvailability::Expired + || account.access_token_expires_at <= now + { + AdminAccountStatus::Expired + } else if account.availability == ProviderAccountAvailability::Ready { + AdminAccountStatus::Active + } else { + AdminAccountStatus::Attention + } +} + +fn admin_account_summary( + accounts: &[ProviderAccountSummary], + now: DateTime, +) -> AccountSummary { + let total = u64::try_from(accounts.len()).unwrap_or(u64::MAX); + let active = u64::try_from( + accounts + .iter() + .filter(|account| admin_account_status(account, now) == AdminAccountStatus::Active) + .count(), + ) + .unwrap_or(u64::MAX); + let quota_exhausted = u64::try_from( + accounts + .iter() + .filter(|account| { + admin_account_status(account, now) == AdminAccountStatus::QuotaExhausted + }) + .count(), + ) + .unwrap_or(u64::MAX); + AccountSummary { + total, + active, + quota_exhausted, + attention: total.saturating_sub(active), + } +} + +fn sort_admin_account_items(items: &mut [AdminAccountListItem], sort: Option) { + let Some(sort) = sort else { + items.sort_by(|left, right| left.account.id.cmp(&right.account.id)); + return; + }; + items.sort_by(|left, right| { + let ordering = match sort.field { + AdminAccountSortField::Email => left.account.email.cmp(&right.account.email), + AdminAccountSortField::Status => { + admin_account_status_name(left.status).cmp(admin_account_status_name(right.status)) + } + AdminAccountSortField::PlanType => left.account.plan_type.cmp(&right.account.plan_type), + AdminAccountSortField::Usage => left + .usage + .as_ref() + .and_then(|usage| usage.total_tokens) + .cmp(&right.usage.as_ref().and_then(|usage| usage.total_tokens)), + AdminAccountSortField::LastUsedAt => left + .usage + .as_ref() + .and_then(|usage| usage.last_used_at.as_ref()) + .cmp( + &right + .usage + .as_ref() + .and_then(|usage| usage.last_used_at.as_ref()), + ), + AdminAccountSortField::ExpiresAt => left + .account + .access_token_expires_at + .cmp(&right.account.access_token_expires_at), + } + .then_with(|| left.account.id.cmp(&right.account.id)); + match sort.direction { + AdminSortDirection::Asc => ordering, + AdminSortDirection::Desc => ordering.reverse(), + } + }); +} + +const fn admin_account_status_name(status: AdminAccountStatus) -> &'static str { + match status { + AdminAccountStatus::Active => "active", + AdminAccountStatus::Expired => "expired", + AdminAccountStatus::QuotaExhausted => "quota_exhausted", + AdminAccountStatus::Disabled => "disabled", + AdminAccountStatus::Banned => "banned", + AdminAccountStatus::Attention => "attention", + } +} + +fn admin_account_record(summary: ProviderAccountSummary) -> AdminStoreResult { + Ok(AccountRecord { + id: summary.id, + provider_instance_id: ProviderInstanceId::new(summary.provider_instance_id).map_err( + |_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "persisted Provider instance ID is invalid", + ) + }, + )?, + provider_kind: ProviderKind::new(summary.provider_kind).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "persisted Provider kind is invalid", + ) + })?, + name: summary.name, + email: summary.email, + upstream_user_id: summary.upstream_user_id, + upstream_account_id: summary.upstream_account_id, + plan_type: summary.plan_type, + credential_revision: admin_revision(summary.credential_revision)?, + has_refresh_token: summary.has_refresh_token, + access_token_expires_at: summary.access_token_expires_at, + next_refresh_at: summary.next_refresh_at, + enabled: summary.enabled, + availability: admin_account_availability(summary.availability), + availability_reason: summary.availability_reason, + cooldown_until: summary.cooldown_until, + availability_observed_at: summary.availability_observed_at, + quota_observed_at: summary.quota_observed_at, + created_at: summary.created_at, + updated_at: summary.updated_at, + }) +} + +fn prepared_account(credential: PreparedCredentialCreate) -> StoreResult { + Ok(NewProviderAccount { + id: credential.account_id.as_str().to_owned(), + provider_instance_id: credential.provider_instance_id.as_str().to_owned(), + provider_kind: credential.provider_kind.as_str().to_owned(), + name: credential.name, + email: credential.email, + upstream_user_id: credential.upstream_user_id, + upstream_account_id: credential.upstream_account_id, + plan_type: credential.plan_type, + provider_credentials_json: provider_document_json(credential.provider_material)?, + has_refresh_token: credential.has_refresh_token, + access_token_expires_at: credential.access_token_expires_at, + next_refresh_at: credential.next_refresh_at, + enabled: credential.enabled, + availability: provider_account_availability(credential.availability), + availability_reason: credential.availability_reason, + cooldown_until: credential.cooldown_until, + availability_observed_at: credential.availability_observed_at, + }) +} + +fn provider_document_json(document: ProviderDocument) -> StoreResult { + JsonObject::try_from_value( + ENTITY, + serde_json::Value::Object(document.into_provider_data().into_inner()), + CREDENTIALS_MAX_BYTES, + ) +} + +const fn provider_account_availability( + availability: AdminAccountAvailability, +) -> ProviderAccountAvailability { + match availability { + AdminAccountAvailability::Unknown => ProviderAccountAvailability::Unknown, + AdminAccountAvailability::Ready => ProviderAccountAvailability::Ready, + AdminAccountAvailability::Cooldown => ProviderAccountAvailability::Cooldown, + AdminAccountAvailability::QuotaExhausted => ProviderAccountAvailability::QuotaExhausted, + AdminAccountAvailability::Expired => ProviderAccountAvailability::Expired, + AdminAccountAvailability::Banned => ProviderAccountAvailability::Banned, + AdminAccountAvailability::Invalid => ProviderAccountAvailability::Invalid, + } +} + +const fn admin_account_availability( + availability: ProviderAccountAvailability, +) -> AdminAccountAvailability { + match availability { + ProviderAccountAvailability::Unknown => AdminAccountAvailability::Unknown, + ProviderAccountAvailability::Ready => AdminAccountAvailability::Ready, + ProviderAccountAvailability::Cooldown => AdminAccountAvailability::Cooldown, + ProviderAccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted, + ProviderAccountAvailability::Expired => AdminAccountAvailability::Expired, + ProviderAccountAvailability::Banned => AdminAccountAvailability::Banned, + ProviderAccountAvailability::Invalid => AdminAccountAvailability::Invalid, + } +} + +fn admin_account_usage(usage: ProviderAccountUsageObservation) -> AdminStoreResult { + Ok(AccountUsage { + account_id: usage.account_id, + request_count: usage.request_count, + success_count: usage.success_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + cache_write_tokens: usage.cache_write_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + cost_coverage: admin_account_cost_coverage(usage.cost_coverage), + costs: admin_account_costs(usage.costs)?, + last_used_at: usage.last_used_at, + models: usage + .models + .into_iter() + .map(admin_account_model_usage) + .collect::>>()?, + }) +} + +fn admin_account_model_usage( + usage: ProviderAccountModelUsageObservation, +) -> AdminStoreResult { + Ok(AccountModelUsage { + model: usage.model, + request_count: usage.request_count, + success_count: usage.success_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + cache_write_tokens: usage.cache_write_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + cost_coverage: admin_account_cost_coverage(usage.cost_coverage), + costs: admin_account_costs(usage.costs)?, + last_used_at: usage.last_used_at, + }) +} + +const fn admin_account_cost_coverage(coverage: super::CostCoverage) -> AdminCostCoverage { + AdminCostCoverage { + provider_reported_count: coverage.provider_reported_count, + calculated_count: coverage.calculated_count, + partial_count: 0, + unavailable_count: coverage.unavailable_count, + not_billable_count: 0, + } +} + +fn admin_account_costs(costs: Vec) -> AdminStoreResult> { + costs + .into_iter() + .map(|cost| { + let amount = AdminDecimalAmount::from_str(cost.amount.as_str()).map_err(|_| { + AdminStoreError::new( + AdminStoreErrorKind::Invalid, + ENTITY, + "persisted account cost is invalid", + ) + })?; + Ok(AccountCost { + currency: cost.currency, + amount, + }) + }) + .collect() +} + async fn upsert_provider_account_in_transaction( transaction: &mut Transaction<'_, Postgres>, account: &NewProviderAccount, @@ -867,7 +1838,7 @@ async fn upsert_provider_account_in_transaction( availability_observed_at, quota_observed_at, created_at, updated_at ) values ( $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13, - $14, null, $15, null, $16, null, now(), greatest(now(), $16) + $14, $15, $16, null, $17, null, now(), greatest(now(), $17) ) on conflict (id) do update set name = excluded.name, @@ -880,7 +1851,7 @@ async fn upsert_provider_account_in_transaction( next_refresh_at = excluded.next_refresh_at, enabled = excluded.enabled, availability = excluded.availability, - availability_reason = null, + availability_reason = excluded.availability_reason, cooldown_until = excluded.cooldown_until, provider_quota_json = null, availability_observed_at = excluded.availability_observed_at, @@ -906,6 +1877,7 @@ async fn upsert_provider_account_in_transaction( .bind(account.next_refresh_at) .bind(account.enabled) .bind(account.availability.as_str()) + .bind(&account.availability_reason) .bind(account.cooldown_until) .bind(account.availability_observed_at) .fetch_optional(&mut **transaction) @@ -1116,6 +2088,7 @@ impl ProviderAccountStore for PgProviderAccountRepository { next_refresh_at: account.account.next_refresh_at().map(DateTime::::from), enabled: account.account.enabled(), availability: availability_from_core(account.account.availability()), + availability_reason: None, cooldown_until: account.account.cooldown_until().map(DateTime::::from), availability_observed_at: Utc::now(), }) @@ -1228,31 +2201,55 @@ impl ProviderAccountStore for PgProviderAccountRepository { } } - async fn get_quota( + async fn get_quotas( &self, - account: &CoreProviderAccountId, - ) -> Result, CoreStoreError> { - let Some(record) = self - .load_provider_account(account.as_str()) - .await - .map_err(core_store_error)? - else { - return Ok(None); - }; - if record.provider_quota_json.is_none() { - return Ok(None); + accounts: &[CoreProviderAccountId], + ) -> Result, CoreStoreError> { + if accounts.is_empty() { + return Ok(Vec::new()); } - Ok(Some(QuotaObservation { - account_id: account.clone(), - expected_revision: CoreCredentialRevision::new( - record.summary.credential_revision.get(), - ) - .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?, - quota: record - .provider_quota_json - .map(|value| OpaqueProviderData::new(value.fields().clone())), - observed_at: record.summary.quota_observed_at.map(Into::into), - })) + let account_ids = accounts + .iter() + .map(|account| account.as_str().to_owned()) + .collect::>(); + let rows = sqlx::query( + "select id, credential_revision, provider_quota_json, quota_observed_at \ + from provider_accounts \ + where id = any($1) and provider_quota_json is not null", + ) + .bind(account_ids) + .fetch_all(&self.pool) + .await + .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?; + rows.into_iter() + .map(|row| { + let account_id = row + .try_get::("id") + .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?; + let revision = row + .try_get::("credential_revision") + .ok() + .and_then(|value| u64::try_from(value).ok()) + .and_then(|value| CoreCredentialRevision::new(value).ok()) + .ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?; + let quota = row + .try_get::("provider_quota_json") + .ok() + .and_then(|value| value.as_object().cloned()) + .map(OpaqueProviderData::new) + .ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?; + let observed_at = row + .try_get::, _>("quota_observed_at") + .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?; + Ok(QuotaObservation { + account_id: CoreProviderAccountId::new(account_id) + .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?, + expected_revision: revision, + quota: Some(quota), + observed_at: Some(observed_at.into()), + }) + }) + .collect() } async fn compare_and_swap_quota( diff --git a/backend/crates/gateway-store/src/postgres/runtime_settings.rs b/backend/crates/gateway-store/src/postgres/runtime_settings.rs index 904332ae1..5f3b913b1 100644 --- a/backend/crates/gateway-store/src/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/src/postgres/runtime_settings.rs @@ -1,11 +1,17 @@ //! `runtime_settings` 单例与 config revision 的 PostgreSQL owner。 +use std::num::NonZeroU32; +use std::time::Duration; use std::{collections::BTreeMap, fmt}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use sqlx::{PgPool, Postgres, Transaction}; +use gateway_core::provider_ports::{ + ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind, +}; + use crate::{ConflictKind, Revision, StoreError, StoreResult, postgres_unavailable}; #[derive(Clone, PartialEq, Eq)] @@ -182,6 +188,24 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository { } } +impl ProviderRuntimePolicyPort for PgRuntimeSettingsRepository { + fn load_refresh_policy( + &self, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let settings = RuntimeSettingsRepository::load_runtime_settings(self) + .await + .map_err(|_| provider_unavailable("load refresh policy"))?; + let concurrency = NonZeroU32::new(settings.refresh_concurrency) + .ok_or_else(|| provider_invalid("decode refresh policy"))?; + ProviderRefreshPolicy::try_new( + Duration::from_secs(settings.refresh_margin_seconds), + concurrency, + ) + }) + } +} + pub(crate) async fn load_runtime_settings_in_transaction( transaction: &mut Transaction<'_, Postgres>, ) -> StoreResult { @@ -316,6 +340,14 @@ fn invalid_numeric() -> StoreError { } } +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + +fn provider_invalid(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} + fn valid_model_mappings(mappings: &BTreeMap>) -> bool { mappings.len() <= 32 && mappings.iter().all(|(provider, entries)| { diff --git a/backend/crates/gateway-store/src/postgres/snapshot.rs b/backend/crates/gateway-store/src/postgres/snapshot.rs index b6f51e7c2..dc38418ae 100644 --- a/backend/crates/gateway-store/src/postgres/snapshot.rs +++ b/backend/crates/gateway-store/src/postgres/snapshot.rs @@ -4,6 +4,13 @@ use std::collections::BTreeMap; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_core::routing::{ + ConfigRevision, + snapshot::{ + SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderInstanceFacts, + SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort, + }, +}; use sqlx::{PgPool, Postgres, Row, Transaction}; use crate::{Revision, StoreError, StoreResult, postgres_unavailable}; @@ -82,7 +89,8 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository { .await .map_err(|_| postgres_unavailable("commit runtime snapshot"))?; - let observed_current_revision = self.current_config_revision().await?; + let observed_current_revision = + RuntimeSnapshotRepository::current_config_revision(self).await?; Ok(RuntimeSnapshotData { config_revision, observed_current_revision, @@ -128,6 +136,73 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository { } } +impl SnapshotStorePort for PgRuntimeSnapshotRepository { + fn load_snapshot_facts( + &self, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let data = self + .load_runtime_snapshot() + .await + .map_err(|_| SnapshotStoreError::unavailable())?; + let config_revision = core_revision(data.config_revision)?; + let observed_current_revision = core_revision(data.observed_current_revision)?; + let settings = SnapshotSettingsFacts::new( + data.settings.max_concurrent_per_account, + data.settings.request_interval_ms, + data.settings.rotation_strategy, + data.settings.provider_model_mappings, + ); + let provider_instances = data + .provider_instances + .into_iter() + .map(|instance| { + SnapshotProviderInstanceFacts::new( + instance.id, + instance.provider_kind, + instance.base_url, + instance.enabled, + ) + }) + .collect(); + let client_policies = data + .client_api_keys + .into_iter() + .map(|key| { + SnapshotClientPolicyFacts::new( + key.id, + key.plaintext_key, + key.provider_kind, + key.limits, + ) + }) + .collect(); + Ok(SnapshotFacts::new( + config_revision, + observed_current_revision, + settings, + provider_instances, + client_policies, + )) + }) + } + + fn current_config_revision( + &self, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + RuntimeSnapshotRepository::current_config_revision(self) + .await + .map_err(|_| SnapshotStoreError::unavailable()) + .and_then(core_revision) + }) + } +} + +fn core_revision(revision: Revision) -> Result { + ConfigRevision::new(revision.get()).map_err(|_| SnapshotStoreError::unavailable()) +} + async fn load_settings( transaction: &mut Transaction<'_, Postgres>, ) -> StoreResult<(Revision, SnapshotRuntimeSettings)> { diff --git a/backend/crates/gateway-store/src/redis/client_admission.rs b/backend/crates/gateway-store/src/redis/client_admission.rs index 99f7547eb..f1d9c8f32 100644 --- a/backend/crates/gateway-store/src/redis/client_admission.rs +++ b/backend/crates/gateway-store/src/redis/client_admission.rs @@ -4,6 +4,13 @@ use std::{collections::HashSet, time::Duration}; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_core::engine::admission::{ + ClientAdmissionDecision as CoreAdmissionDecision, ClientAdmissionError as CoreAdmissionError, + ClientAdmissionPort, ClientAdmissionRecovery as CoreAdmissionRecovery, + ClientAdmissionRejection as CoreAdmissionRejection, + ClientAdmissionRequest as CoreAdmissionRequest, + ClientAdmissionRestoreResult as CoreAdmissionRestoreResult, +}; use redis::{Script, aio::ConnectionManager}; use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty}; @@ -397,6 +404,85 @@ impl ClientAdmissionRepository for RedisClientAdmissionRepository { } } +impl ClientAdmissionPort for RedisClientAdmissionRepository { + fn admit( + &self, + request: CoreAdmissionRequest, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + self.admit_client_request(&ClientAdmissionRequest { + model_request_id: request.model_request_id.as_str().to_owned(), + client_api_key_ref: request.client_api_key_id.as_str().to_owned(), + input_token_estimate: request.input_token_estimate, + lease_ttl: request.lease_ttl, + limits: ClientAdmissionLimits { + max_concurrency: request.limits.max_concurrency, + requests_per_minute: request.limits.requests_per_minute, + tokens_per_minute: request.limits.tokens_per_minute, + }, + }) + .await + .map(|decision| match decision { + ClientAdmissionDecision::Granted => CoreAdmissionDecision::Granted, + ClientAdmissionDecision::Rejected(ClientAdmissionRejection::RateLimited) => { + CoreAdmissionDecision::Rejected(CoreAdmissionRejection::RateLimited) + } + ClientAdmissionDecision::Rejected(ClientAdmissionRejection::ConcurrencyLimited) => { + CoreAdmissionDecision::Rejected(CoreAdmissionRejection::ConcurrencyLimited) + } + }) + .map_err(|_| CoreAdmissionError) + }) + } + + fn release<'a>( + &'a self, + client_api_key_id: &'a gateway_core::policy::ClientApiKeyId, + model_request_id: &'a gateway_core::engine::ModelRequestId, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + self.release_client_request(client_api_key_id.as_str(), model_request_id.as_str()) + .await + .map_err(|_| CoreAdmissionError) + }) + } + + fn restore( + &self, + recovery: CoreAdmissionRecovery, + ) -> futures::future::BoxFuture<'_, Result> + { + Box::pin(async move { + self.restore_client_admission(&ClientAdmissionRestore { + client_api_key_ref: recovery.client_api_key_id.as_str().to_owned(), + recent_requests: recovery + .recent_requests + .into_iter() + .map(|request| ClientAdmissionRecentRequest { + model_request_id: request.model_request_id.as_str().to_owned(), + started_at: DateTime::::from(request.started_at), + input_token_estimate: request.input_token_estimate, + }) + .collect(), + running_requests: recovery + .running_requests + .into_iter() + .map(|request| ClientAdmissionRunningRequest { + model_request_id: request.model_request_id.as_str().to_owned(), + expires_at: DateTime::::from(request.expires_at), + }) + .collect(), + }) + .await + .map(|restored| CoreAdmissionRestoreResult { + restored_recent_requests: restored.restored_recent_requests, + restored_running_requests: restored.restored_running_requests, + }) + .map_err(|_| CoreAdmissionError) + }) + } +} + fn invalid(message: &str) -> StoreError { StoreError::InvalidData { entity: "client admission", diff --git a/backend/crates/gateway-store/src/redis/credential_cooldown.rs b/backend/crates/gateway-store/src/redis/credential_cooldown.rs index f1f81cfda..dcaac8a21 100644 --- a/backend/crates/gateway-store/src/redis/credential_cooldown.rs +++ b/backend/crates/gateway-store/src/redis/credential_cooldown.rs @@ -12,6 +12,9 @@ use gateway_core::{ QuotaWriteOutcome, }, error::StoreError as CoreStoreError, + provider_ports::{ + ProviderCooldown, ProviderCooldownPort, ProviderStoreError, ProviderStoreErrorKind, + }, routing::ProviderInstanceId, }; use redis::{Script, aio::ConnectionManager}; @@ -180,6 +183,66 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository { } } +impl ProviderCooldownPort for RedisCredentialCooldownRepository { + fn put_if_later( + &self, + cooldown: ProviderCooldown, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let record = CredentialCooldown { + provider_account_id: cooldown.account_id().as_str().to_owned(), + credential_revision: Revision::new(cooldown.credential_revision().get()) + .map_err(|_| provider_invalid("encode credential cooldown"))?, + cooldown_until: cooldown.until().into(), + }; + CredentialCooldownRepository::cache_credential_cooldown(self, &record) + .await + .map_err(|_| provider_unavailable("cache credential cooldown")) + }) + } + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + CredentialCooldownRepository::read_credential_cooldown(self, account_id.as_str()) + .await + .map_err(|_| provider_unavailable("read credential cooldown"))? + .map(|record| { + let account_id = ProviderAccountId::new(record.provider_account_id) + .map_err(|_| provider_invalid("decode credential cooldown"))?; + let revision = CredentialRevision::new(record.credential_revision.get()) + .map_err(|_| provider_invalid("decode credential cooldown"))?; + Ok(ProviderCooldown::new( + account_id, + revision, + record.cooldown_until.into(), + )) + }) + .transpose() + }) + } + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + through_revision: CredentialRevision, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + let revision = Revision::new(through_revision.get()) + .map_err(|_| provider_invalid("encode credential cooldown revision"))?; + CredentialCooldownRepository::invalidate_credential_cooldown( + self, + account_id.as_str(), + revision, + ) + .await + .map_err(|_| provider_unavailable("clear credential cooldown")) + }) + } +} + /// PostgreSQL 账号事实到可丢失 cooldown cache 的 write-through adapter。 pub struct CooldownCachingProviderAccountStore { authoritative: Arc, @@ -305,11 +368,11 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore { Ok(outcome) } - async fn get_quota( + async fn get_quotas( &self, - account: &ProviderAccountId, - ) -> Result, CoreStoreError> { - self.authoritative.get_quota(account).await + accounts: &[ProviderAccountId], + ) -> Result, CoreStoreError> { + self.authoritative.get_quotas(accounts).await } async fn compare_and_swap_quota( @@ -383,3 +446,11 @@ fn invalid(message: &str) -> StoreError { message: message.to_owned(), } } + +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + +fn provider_invalid(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-store/src/redis/credential_leases.rs b/backend/crates/gateway-store/src/redis/credential_leases.rs index 51acb0ae8..6782e759a 100644 --- a/backend/crates/gateway-store/src/redis/credential_leases.rs +++ b/backend/crates/gateway-store/src/redis/credential_leases.rs @@ -1,10 +1,20 @@ //! Provider instance、账号与 OAuth refresh 的 Redis lease/fencing。 -use std::{fmt, time::Duration}; +use std::collections::BTreeMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, SystemTime}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use futures::future::join_all; +use gateway_core::engine::credential::{AccountRuntimeSignals, ProviderAccountId}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState, + ProviderStoreError, ProviderStoreErrorKind, +}; +use gateway_core::routing::ProviderInstanceId; use redis::{Script, aio::ConnectionManager}; use uuid::Uuid; @@ -13,6 +23,9 @@ use crate::{Revision, StoreError, StoreResult, redis_unavailable, require_nonemp use super::{MAX_REDIS_EXACT_INTEGER, namespace, resource_fingerprint}; const SIGNAL_TTL_MILLIS: u64 = 24 * 60 * 60 * 1_000; +const PROVIDER_ACCOUNT_LEASE_TTL: Duration = Duration::from_secs(10 * 60); +const OAUTH_REFRESH_LEASE_TTL: Duration = Duration::from_secs(5 * 60); +const OAUTH_REFRESH_CAPACITY_RESOURCE: &str = "oauth-refresh-global"; const ACQUIRE_SCRIPT: &str = r#" local clock = redis.call('TIME') @@ -99,10 +112,22 @@ local last_started = redis.call('GET', KEYS[2]) or '0' return {tostring(in_flight), tostring(last_started)} "#; +const ADVANCE_SCHEDULING_CURSOR_SCRIPT: &str = r#" +local cursor = redis.call('INCR', KEYS[1]) +if cursor > tonumber(ARGV[1]) then + redis.call('SET', KEYS[1], '1', 'PX', ARGV[2]) + cursor = 1 +else + redis.call('PEXPIRE', KEYS[1], ARGV[2]) +end +return tostring(cursor - 1) +"#; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CredentialLeaseScope { ProviderInstance, ProviderAccount, + OAuthRefreshCapacity, OAuthRefresh, ProviderTask, } @@ -112,6 +137,7 @@ impl CredentialLeaseScope { match self { Self::ProviderInstance => "instance", Self::ProviderAccount => "account", + Self::OAuthRefreshCapacity => "refresh-capacity", Self::OAuthRefresh => "refresh", Self::ProviderTask => "task", } @@ -135,9 +161,10 @@ impl CredentialLeaseRequest { } } -/// Provider account 的计数 lease 请求。 +/// 对任意明确资源施加并发数和启动间隔约束的计数 lease。 #[derive(Debug, Clone, PartialEq, Eq)] -pub struct CredentialSchedulingLeaseRequest { +pub struct CredentialBoundedLeaseRequest { + pub scope: CredentialLeaseScope, pub resource_id: String, pub owner_id: String, pub max_concurrent: u32, @@ -145,14 +172,10 @@ pub struct CredentialSchedulingLeaseRequest { pub ttl: Duration, } -impl CredentialSchedulingLeaseRequest { +impl CredentialBoundedLeaseRequest { pub fn validate(&self) -> StoreResult<()> { - require_nonempty( - "credential scheduling lease", - "resource_id", - &self.resource_id, - )?; - require_nonempty("credential scheduling lease", "owner_id", &self.owner_id)?; + require_nonempty("credential bounded lease", "resource_id", &self.resource_id)?; + require_nonempty("credential bounded lease", "owner_id", &self.owner_id)?; if self.max_concurrent == 0 { return Err(invalid("max_concurrent must be positive")); } @@ -163,7 +186,7 @@ impl CredentialSchedulingLeaseRequest { fn lease_request(&self) -> CredentialLeaseRequest { CredentialLeaseRequest { - scope: CredentialLeaseScope::ProviderAccount, + scope: self.scope, resource_id: self.resource_id.clone(), owner_id: self.owner_id.clone(), ttl: self.ttl, @@ -186,12 +209,12 @@ pub struct CredentialRuntimeSignal { pub last_started_at: Option>, } -pub enum CredentialSchedulingLeaseAcquisition { +pub enum CredentialBoundedLeaseAcquisition { Acquired(CredentialLeaseGuard), Busy { retry_after: Option }, } -impl fmt::Debug for CredentialSchedulingLeaseAcquisition { +impl fmt::Debug for CredentialBoundedLeaseAcquisition { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Acquired(_) => formatter.write_str("Acquired([LEASE_GUARD])"), @@ -273,10 +296,10 @@ pub trait CredentialLeaseRepository: Send + Sync { &self, resource_ids: &[String], ) -> StoreResult>; - async fn try_acquire_scheduling_lease( + async fn try_acquire_bounded_lease( &self, - request: &CredentialSchedulingLeaseRequest, - ) -> StoreResult; + request: &CredentialBoundedLeaseRequest, + ) -> StoreResult; } #[derive(Clone)] @@ -306,6 +329,27 @@ impl RedisCredentialLeaseRepository { })) } + /// 原子推进跨进程共享的 Provider instance 调度游标。 + pub async fn advance_scheduling_cursor(&self, provider_instance_id: &str) -> StoreResult { + require_nonempty( + "provider scheduling cursor", + "provider_instance_id", + provider_instance_id, + )?; + let key = self.scheduling_cursor_key(provider_instance_id)?; + let mut connection = self.connection.clone(); + let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT) + .key(key) + .arg(MAX_REDIS_EXACT_INTEGER) + .arg(SIGNAL_TTL_MILLIS) + .invoke_async::(&mut connection) + .await + .map_err(|_| redis_unavailable("advance provider scheduling cursor"))?; + cursor + .parse::() + .map_err(|_| invalid("Redis returned an invalid scheduling cursor")) + } + fn keys(&self, request: &CredentialLeaseRequest) -> StoreResult<[String; 3]> { let fingerprint = resource_fingerprint("credential lease", &request.resource_id)?; let tag = format!("{{{fingerprint}}}"); @@ -329,6 +373,14 @@ impl RedisCredentialLeaseRepository { )) } + fn scheduling_cursor_key(&self, provider_instance_id: &str) -> StoreResult { + let fingerprint = resource_fingerprint("provider scheduling cursor", provider_instance_id)?; + Ok(format!( + "{}:scheduler:instance:{{{fingerprint}}}:cursor", + self.namespace + )) + } + async fn acquire_with_limits( &self, request: &CredentialLeaseRequest, @@ -403,6 +455,192 @@ impl RedisCredentialLeaseRepository { } } +/// Store-owned 的通用 Provider lease 能力；具体 Provider 不感知 Redis。 +pub(crate) struct RedisProviderLeaseCoordinator { + repository: RedisCredentialLeaseRepository, + process_id: String, + sequence: AtomicU64, +} + +impl RedisProviderLeaseCoordinator { + #[must_use] + pub(crate) fn new(repository: RedisCredentialLeaseRepository) -> Self { + Self { + repository, + process_id: format!("gateway_{}", Uuid::now_v7().simple()), + sequence: AtomicU64::new(0), + } + } + + fn owner_id(&self, operation: &str) -> String { + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed); + format!("{}:{operation}:{sequence}", self.process_id) + } + + async fn next_scheduling_cursor( + &self, + provider_instance_id: &ProviderInstanceId, + ) -> Result { + self.repository + .advance_scheduling_cursor(provider_instance_id.as_str()) + .await + .map_err(|_| provider_unavailable("advance scheduling cursor")) + } + + async fn load_signals( + &self, + accounts: &[ProviderAccountId], + ) -> Result, ProviderStoreError> { + let ids = accounts + .iter() + .map(|account| account.as_str().to_owned()) + .collect::>(); + let signals = self + .repository + .credential_runtime_signals(&ids) + .await + .map_err(|_| provider_unavailable("load scheduling signals"))?; + signals + .into_iter() + .map(|signal| { + let account = ProviderAccountId::new(signal.resource_id).map_err(|_| { + ProviderStoreError::new( + ProviderStoreErrorKind::InvalidData, + "decode scheduling signals", + ) + })?; + Ok(( + account, + AccountRuntimeSignals { + in_flight: signal.in_flight, + last_started_at: signal.last_started_at.map(Into::into), + quota_reset_at: None, + quota_remaining_rank: None, + }, + )) + }) + .collect() + } + + async fn acquire_scheduling( + &self, + request: &ProviderSchedulingLeaseRequest, + ) -> Result { + let ttl = request + .deadline() + .duration_since(SystemTime::now()) + .ok() + .filter(|remaining| !remaining.is_zero()) + .map(|remaining| remaining.min(PROVIDER_ACCOUNT_LEASE_TTL)) + .ok_or_else(|| { + ProviderStoreError::new( + ProviderStoreErrorKind::Unavailable, + "acquire expired scheduling lease", + ) + })?; + let acquisition = self + .repository + .try_acquire_bounded_lease(&CredentialBoundedLeaseRequest { + scope: CredentialLeaseScope::ProviderAccount, + resource_id: request.account_id().as_str().to_owned(), + owner_id: self.owner_id("request"), + max_concurrent: request.max_concurrent().get(), + request_interval: request.request_interval(), + ttl, + }) + .await + .map_err(|_| provider_unavailable("acquire scheduling lease"))?; + Ok(match acquisition { + CredentialBoundedLeaseAcquisition::Acquired(guard) => { + ProviderLeaseAcquisition::Acquired(Box::new(guard)) + } + CredentialBoundedLeaseAcquisition::Busy { retry_after } => { + ProviderLeaseAcquisition::Busy { retry_after } + } + }) + } + + async fn acquire_refresh_capacity( + &self, + request: ProviderRefreshCapacityRequest, + ) -> Result { + let acquisition = self + .repository + .try_acquire_bounded_lease(&CredentialBoundedLeaseRequest { + scope: CredentialLeaseScope::OAuthRefreshCapacity, + resource_id: OAUTH_REFRESH_CAPACITY_RESOURCE.to_owned(), + owner_id: self.owner_id("refresh-capacity"), + max_concurrent: request.max_concurrent().get(), + request_interval: Duration::ZERO, + ttl: OAUTH_REFRESH_LEASE_TTL, + }) + .await + .map_err(|_| provider_unavailable("acquire refresh capacity"))?; + Ok(match acquisition { + CredentialBoundedLeaseAcquisition::Acquired(guard) => { + ProviderLeaseAcquisition::Acquired(Box::new(guard)) + } + CredentialBoundedLeaseAcquisition::Busy { retry_after } => { + ProviderLeaseAcquisition::Busy { retry_after } + } + }) + } +} + +impl ProviderLeasePort for RedisProviderLeaseCoordinator { + fn load_state<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + accounts: &'a [ProviderAccountId], + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + let signals = self.load_signals(accounts).await?; + let round_robin_cursor = self.next_scheduling_cursor(provider_instance_id).await?; + Ok(ProviderSchedulingState::new( + signals, + None, + round_robin_cursor, + )) + }) + } + + fn try_acquire( + &self, + request: ProviderLeaseRequest, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + match request { + ProviderLeaseRequest::Scheduling(request) => { + self.acquire_scheduling(&request).await + } + ProviderLeaseRequest::RefreshCapacity(request) => { + self.acquire_refresh_capacity(request).await + } + ProviderLeaseRequest::Refresh(request) => { + let lease = CredentialLeaseRequest { + scope: CredentialLeaseScope::OAuthRefresh, + resource_id: request.account_id().as_str().to_owned(), + owner_id: self.owner_id("refresh"), + ttl: OAUTH_REFRESH_LEASE_TTL, + }; + self.repository + .try_acquire_guard(lease) + .await + .map(|guard| match guard { + Some(guard) => ProviderLeaseAcquisition::Acquired(Box::new(guard)), + None => ProviderLeaseAcquisition::Busy { retry_after: None }, + }) + .map_err(|_| provider_unavailable("acquire refresh lease")) + } + } + }) + } +} + +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + #[async_trait] impl CredentialLeaseRepository for RedisCredentialLeaseRepository { async fn acquire_credential_lease( @@ -475,10 +713,10 @@ impl CredentialLeaseRepository for RedisCredentialLeaseRepository { signals.into_iter().collect() } - async fn try_acquire_scheduling_lease( + async fn try_acquire_bounded_lease( &self, - request: &CredentialSchedulingLeaseRequest, - ) -> StoreResult { + request: &CredentialBoundedLeaseRequest, + ) -> StoreResult { request.validate()?; let lease_request = request.lease_request(); let attempt = self @@ -489,14 +727,14 @@ impl CredentialLeaseRepository for RedisCredentialLeaseRepository { ) .await?; match attempt.grant { - Some(grant) => Ok(CredentialSchedulingLeaseAcquisition::Acquired( + Some(grant) => Ok(CredentialBoundedLeaseAcquisition::Acquired( CredentialLeaseGuard { repository: self.clone(), request: lease_request, grant: Some(grant), }, )), - None => Ok(CredentialSchedulingLeaseAcquisition::Busy { + None => Ok(CredentialBoundedLeaseAcquisition::Busy { retry_after: attempt.retry_after, }), } diff --git a/backend/crates/gateway-store/src/redis/credential_state.rs b/backend/crates/gateway-store/src/redis/credential_state.rs index b1273aaf6..7129416a8 100644 --- a/backend/crates/gateway-store/src/redis/credential_state.rs +++ b/backend/crates/gateway-store/src/redis/credential_state.rs @@ -4,7 +4,13 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use redis::{Script, aio::ConnectionManager}; -use gateway_core::engine::credential::OpaqueProviderData; +use gateway_core::engine::credential::{ + AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccountId, +}; +use gateway_core::provider_ports::{ + ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCredentialState, + ProviderCredentialStatePort, ProviderStoreError, ProviderStoreErrorKind, +}; use crate::{Revision, StoreError, StoreResult, redis_unavailable, require_nonempty}; @@ -23,6 +29,7 @@ return 1 "#; const MAX_PROVIDER_CATALOG_BYTES: usize = 1024 * 1024; +const CREDENTIAL_STATE_TTL_SECONDS: u64 = 24 * 60 * 60; #[derive(Debug, Clone, PartialEq, Eq)] pub struct CredentialStateCache { @@ -271,6 +278,122 @@ impl ProviderAccountCatalogCacheRepository for RedisCredentialStateRepository { } } +impl ProviderCredentialStatePort for RedisCredentialStateRepository { + fn replace( + &self, + state: ProviderCredentialState, + ) -> futures::future::BoxFuture<'_, Result<(), ProviderStoreError>> { + Box::pin(async move { + let observed_at: DateTime = state.observed_at().into(); + let cached = CredentialStateCache { + provider_account_id: state.account_id().as_str().to_owned(), + revision: Revision::new(state.credential_revision().get()) + .map_err(|_| provider_invalid("encode credential state"))?, + enabled: state.enabled(), + availability: state.availability().as_str().to_owned(), + observed_at, + }; + CredentialStateRepository::cache_credential_state( + self, + &cached, + CREDENTIAL_STATE_TTL_SECONDS, + ) + .await + .map(|_| ()) + .map_err(|_| provider_unavailable("replace credential state")) + }) + } + + fn read<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> + { + Box::pin(async move { + let cached = + CredentialStateRepository::read_credential_state(self, account_id.as_str()) + .await + .map_err(|_| provider_unavailable("read credential state"))?; + cached + .map(|state| { + let account_id = ProviderAccountId::new(state.provider_account_id) + .map_err(|_| provider_invalid("decode credential state"))?; + let revision = CredentialRevision::new(state.revision.get()) + .map_err(|_| provider_invalid("decode credential state"))?; + let availability = AccountAvailability::parse(&state.availability) + .ok_or_else(|| provider_invalid("decode credential state"))?; + Ok(ProviderCredentialState::new( + account_id, + revision, + state.enabled, + availability, + state.observed_at.into(), + )) + }) + .transpose() + }) + } + + fn clear<'a>( + &'a self, + account_id: &'a ProviderAccountId, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + CredentialStateRepository::clear_credential_state(self, account_id.as_str()) + .await + .map_err(|_| provider_unavailable("clear credential state")) + }) + } +} + +impl ProviderCatalogCachePort for RedisCredentialStateRepository { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + ttl: std::time::Duration, + ) -> futures::future::BoxFuture<'a, Result<(), ProviderStoreError>> { + Box::pin(async move { + let ttl_seconds = ttl.as_secs(); + if ttl_seconds == 0 || ttl.subsec_nanos() != 0 { + return Err(provider_invalid("validate catalog cache TTL")); + } + let key = ProviderAccountCatalogCacheKey { + provider_kind: key.provider_kind().as_str().to_owned(), + provider_account_id: key.account_id().as_str().to_owned(), + credential_revision: Revision::new(key.credential_revision().get()) + .map_err(|_| provider_invalid("encode catalog cache key"))?, + }; + ProviderAccountCatalogCacheRepository::replace_provider_account_catalog( + self, + &key, + catalog, + ttl_seconds, + ) + .await + .map_err(|_| provider_unavailable("replace catalog cache")) + }) + } + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> + { + Box::pin(async move { + let key = ProviderAccountCatalogCacheKey { + provider_kind: key.provider_kind().as_str().to_owned(), + provider_account_id: key.account_id().as_str().to_owned(), + credential_revision: Revision::new(key.credential_revision().get()) + .map_err(|_| provider_invalid("encode catalog cache key"))?, + }; + ProviderAccountCatalogCacheRepository::get_provider_account_catalog(self, &key) + .await + .map_err(|_| provider_unavailable("read catalog cache")) + }) + } +} + fn invalid(message: &str) -> StoreError { StoreError::InvalidData { entity: "credential state cache", @@ -284,3 +407,11 @@ fn catalog_invalid(message: &str) -> StoreError { message: message.to_owned(), } } + +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + +fn provider_invalid(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-store/src/redis/mod.rs b/backend/crates/gateway-store/src/redis/mod.rs index bd4b70a15..df311e07f 100644 --- a/backend/crates/gateway-store/src/redis/mod.rs +++ b/backend/crates/gateway-store/src/redis/mod.rs @@ -11,20 +11,22 @@ mod client_admission; mod credential_cooldown; mod credential_leases; mod credential_state; +mod oauth_pending; mod provider_circuit; mod runtime_change; +pub(crate) mod worker_lease; pub use client_admission::*; pub use credential_cooldown::*; pub use credential_leases::*; pub use credential_state::*; +pub use oauth_pending::*; pub use provider_circuit::*; pub use runtime_change::*; use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty}; pub(crate) const MAX_REDIS_EXACT_INTEGER: u64 = (1_u64 << 53) - 1; - const RECORD_ADMIN_LOGIN_FAILURE_SCRIPT: &str = r#" local count = redis.call('INCR', KEYS[1]) local ttl = redis.call('PTTL', KEYS[1]) diff --git a/backend/crates/gateway-store/src/redis/oauth_pending.rs b/backend/crates/gateway-store/src/redis/oauth_pending.rs new file mode 100644 index 000000000..b9802abe7 --- /dev/null +++ b/backend/crates/gateway-store/src/redis/oauth_pending.rs @@ -0,0 +1,164 @@ +//! OAuth 授权临时状态的 Redis 原子持久化。 + +use std::time::Duration; + +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome, + OAuthPendingTakeOutcome, ProviderStoreError, ProviderStoreErrorKind, +}; +use gateway_core::routing::ProviderKind; +use redis::{Script, aio::ConnectionManager}; + +use super::{MAX_REDIS_EXACT_INTEGER, namespace, resource_fingerprint}; +use crate::{StoreError, StoreResult}; + +const MAX_PAYLOAD_BYTES: usize = 1024 * 1024; + +const CREATE_SCRIPT: &str = r#" +if redis.call('EXISTS', KEYS[1]) == 1 then + return 0 +end +local clock = redis.call('TIME') +local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000) +local expires_at_ms = now_ms + tonumber(ARGV[3]) +redis.call('HSET', KEYS[1], + 'owner_fingerprint', ARGV[1], + 'expires_at', expires_at_ms, + 'provider_payload', ARGV[2]) +redis.call('PEXPIREAT', KEYS[1], expires_at_ms) +return 1 +"#; + +const TAKE_SCRIPT: &str = r#" +local owner = redis.call('HGET', KEYS[1], 'owner_fingerprint') +if owner == false then + return {0, ''} +end +if owner ~= ARGV[1] then + return {-1, ''} +end +local payload = redis.call('HGET', KEYS[1], 'provider_payload') +if payload == false then + return redis.error_reply('OAuth pending flow payload is missing') +end +redis.call('DEL', KEYS[1]) +return {1, payload} +"#; + +/// OAuth pending flow 的 Redis 原子持久化能力。 +#[derive(Clone)] +pub struct RedisOAuthPendingFlowRepository { + connection: ConnectionManager, + namespace: String, +} + +impl RedisOAuthPendingFlowRepository { + pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult { + Ok(Self { + connection, + namespace: format!("{}:oauth-pending:v1", namespace(key_namespace)?), + }) + } + + fn key(&self, provider_kind: &ProviderKind, flow: &OAuthPendingBinding) -> StoreResult { + let fingerprint = resource_fingerprint("OAuth pending flow", flow.expose_to_store())?; + Ok(format!( + "{}:{}:{fingerprint}", + self.namespace, + provider_kind.as_str() + )) + } +} + +impl OAuthPendingFlowPort for RedisOAuthPendingFlowRepository { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let key = self + .key(flow.provider_kind(), flow.flow()) + .map_err(|_| provider_invalid("encode OAuth pending key"))?; + let owner = resource_fingerprint("OAuth pending owner", flow.owner().expose_to_store()) + .map_err(|_| provider_invalid("encode OAuth pending owner"))?; + let payload = serde_json::to_vec(&serde_json::Value::Object( + flow.payload().expose_to_provider().clone(), + )) + .map_err(|_| provider_invalid("encode OAuth pending payload"))?; + if payload.len() > MAX_PAYLOAD_BYTES { + return Err(provider_invalid("validate OAuth pending payload")); + } + let ttl_millis = + ttl_millis(flow.ttl()).map_err(|_| provider_invalid("encode OAuth pending TTL"))?; + let mut connection = self.connection.clone(); + let stored = Script::new(CREATE_SCRIPT) + .key(key) + .arg(owner) + .arg(payload) + .arg(ttl_millis) + .invoke_async::(&mut connection) + .await + .map_err(|_| provider_unavailable("create OAuth pending flow"))?; + Ok(if stored == 1 { + OAuthPendingPutOutcome::Stored + } else { + OAuthPendingPutOutcome::AlreadyExists + }) + }) + } + + fn take_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + let key = self + .key(provider_kind, flow) + .map_err(|_| provider_invalid("encode OAuth pending key"))?; + let owner = resource_fingerprint("OAuth pending owner", owner.expose_to_store()) + .map_err(|_| provider_invalid("encode OAuth pending owner"))?; + let mut connection = self.connection.clone(); + let (status, payload): (i64, Vec) = Script::new(TAKE_SCRIPT) + .key(key) + .arg(owner) + .invoke_async(&mut connection) + .await + .map_err(|_| provider_unavailable("take OAuth pending flow"))?; + match status { + 0 => Ok(OAuthPendingTakeOutcome::NotFound), + -1 => Ok(OAuthPendingTakeOutcome::OwnerMismatch), + 1 => { + let value: serde_json::Value = serde_json::from_slice(&payload) + .map_err(|_| provider_invalid("decode OAuth pending payload"))?; + let serde_json::Value::Object(fields) = value else { + return Err(provider_invalid("decode OAuth pending payload")); + }; + Ok(OAuthPendingTakeOutcome::Taken( + gateway_core::engine::credential::OpaqueProviderData::new(fields), + )) + } + _ => Err(provider_invalid("decode OAuth pending outcome")), + } + }) + } +} + +fn ttl_millis(ttl: Duration) -> StoreResult { + u64::try_from(ttl.as_millis()) + .ok() + .filter(|value| *value > 0 && *value <= MAX_REDIS_EXACT_INTEGER) + .ok_or_else(|| StoreError::InvalidData { + entity: "OAuth pending flow", + message: "TTL is outside the supported range".to_owned(), + }) +} + +fn provider_unavailable(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) +} + +fn provider_invalid(operation: &'static str) -> ProviderStoreError { + ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation) +} diff --git a/backend/crates/gateway-store/src/redis/provider_circuit.rs b/backend/crates/gateway-store/src/redis/provider_circuit.rs index e8d62b321..806d0be01 100644 --- a/backend/crates/gateway-store/src/redis/provider_circuit.rs +++ b/backend/crates/gateway-store/src/redis/provider_circuit.rs @@ -5,6 +5,10 @@ use std::time::Duration; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use gateway_core::engine::execution::{ + ProviderCircuitDecision as CoreCircuitDecision, ProviderCircuitError, ProviderCircuitPort, +}; +use gateway_core::routing::ProviderInstanceId; use redis::{Script, aio::ConnectionManager}; use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty}; @@ -175,6 +179,48 @@ impl ProviderCircuitRepository for RedisProviderCircuitRepository { } } +impl ProviderCircuitPort for RedisProviderCircuitRepository { + fn decision<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + self.provider_circuit_decision(provider_instance_id.as_str()) + .await + .map(|decision| match decision { + ProviderCircuitDecision::Allow => CoreCircuitDecision::Allow, + ProviderCircuitDecision::BlockedUntil(until) => { + CoreCircuitDecision::BlockedUntil(until.into()) + } + }) + .map_err(|_| ProviderCircuitError) + }) + } + + fn observe_failure<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async move { + self.observe_provider_failure(provider_instance_id.as_str()) + .await + .map(|_| ()) + .map_err(|_| ProviderCircuitError) + }) + } + + fn observe_success<'a>( + &'a self, + provider_instance_id: &'a ProviderInstanceId, + ) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> { + Box::pin(async move { + self.observe_provider_success(provider_instance_id.as_str()) + .await + .map_err(|_| ProviderCircuitError) + }) + } +} + fn timestamp(value: &str) -> StoreResult> { let value = value .parse::() diff --git a/backend/crates/gateway-store/src/redis/runtime_change.rs b/backend/crates/gateway-store/src/redis/runtime_change.rs index 7be04d6c5..ff212d74b 100644 --- a/backend/crates/gateway-store/src/redis/runtime_change.rs +++ b/backend/crates/gateway-store/src/redis/runtime_change.rs @@ -4,6 +4,10 @@ use std::pin::Pin; use async_trait::async_trait; use futures::{Stream, StreamExt}; +use gateway_core::routing::{ + ConfigRevision, + snapshot::{SnapshotRevisionStream, SnapshotSubscriptionError, SnapshotSubscriptionPort}, +}; use redis::AsyncCommands; use serde::{Deserialize, Serialize}; @@ -101,6 +105,44 @@ impl RuntimeChangeRepository for RedisRuntimeChangeRepository { } } +impl SnapshotSubscriptionPort for RedisRuntimeChangeRepository { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> futures::future::BoxFuture<'_, Result<(), SnapshotSubscriptionError>> { + Box::pin(async move { + let config_revision = Revision::new(revision.get()) + .map_err(|_| SnapshotSubscriptionError::unavailable())?; + self.publish_runtime_change(&RuntimeChange::SnapshotPublished { config_revision }) + .await + .map_err(|_| SnapshotSubscriptionError::unavailable()) + }) + } + + fn subscribe_snapshot_revisions( + &self, + ) -> futures::future::BoxFuture<'_, Result> + { + Box::pin(async move { + let stream = self + .subscribe_runtime_changes() + .await + .map_err(|_| SnapshotSubscriptionError::unavailable())? + .filter_map(|change| async move { + match change { + Ok(RuntimeChange::SnapshotPublished { config_revision }) => Some( + ConfigRevision::new(config_revision.get()) + .map_err(|_| SnapshotSubscriptionError::unavailable()), + ), + Ok(RuntimeChange::ProviderAccountChanged { .. }) => None, + Err(_) => Some(Err(SnapshotSubscriptionError::unavailable())), + } + }); + Ok(Box::pin(stream) as SnapshotRevisionStream) + }) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] struct RuntimeChangeWire { kind: String, diff --git a/backend/crates/gateway-store/src/redis/worker_lease.rs b/backend/crates/gateway-store/src/redis/worker_lease.rs new file mode 100644 index 000000000..5f5868126 --- /dev/null +++ b/backend/crates/gateway-store/src/redis/worker_lease.rs @@ -0,0 +1,114 @@ +//! 多实例 worker leader lease 的 Redis 实现。 + +use std::num::NonZeroU64; + +use gateway_core::task::{ + WorkerFencingToken, WorkerLeaderLeaseGuard, WorkerLeaderLeasePort, WorkerLeaseAcquisition, + WorkerLeaseError, WorkerLeaseRequest, +}; +use uuid::Uuid; + +use super::{ + CredentialLeaseGrant, CredentialLeaseRepository, CredentialLeaseRequest, CredentialLeaseScope, + RedisCredentialLeaseRepository, +}; + +pub(crate) struct RedisWorkerLeaderLeasePort { + repository: RedisCredentialLeaseRepository, + owner_id: String, +} + +impl RedisWorkerLeaderLeasePort { + #[must_use] + pub(crate) fn new(repository: RedisCredentialLeaseRepository) -> Self { + Self { + repository, + owner_id: format!("worker-process-{}", Uuid::now_v7().simple()), + } + } +} + +struct RedisWorkerLeaderLeaseGuard { + repository: RedisCredentialLeaseRepository, + request: CredentialLeaseRequest, + grant: Option, + token: WorkerFencingToken, +} + +impl WorkerLeaderLeaseGuard for RedisWorkerLeaderLeaseGuard { + fn fencing_token(&self) -> WorkerFencingToken { + self.token + } + + fn renew(&mut self) -> futures::future::BoxFuture<'_, Result<(), WorkerLeaseError>> { + Box::pin(async move { + let current = self + .grant + .as_ref() + .ok_or_else(|| WorkerLeaseError::safe("worker lease is released"))?; + let renewed = CredentialLeaseRepository::renew_credential_lease( + &self.repository, + &self.request, + current, + ) + .await + .map_err(|_| WorkerLeaseError::safe("worker lease renewal failed"))? + .ok_or_else(|| WorkerLeaseError::safe("worker lease was lost"))?; + self.grant = Some(renewed); + Ok(()) + }) + } + + fn release( + mut self: Box, + ) -> futures::future::BoxFuture<'static, Result<(), WorkerLeaseError>> { + let repository = self.repository.clone(); + let request = self.request.clone(); + let grant = self.grant.take(); + Box::pin(async move { + let Some(grant) = grant else { + return Ok(()); + }; + CredentialLeaseRepository::release_credential_lease(&repository, &request, &grant) + .await + .map(|_| ()) + .map_err(|_| WorkerLeaseError::safe("worker lease release failed")) + }) + } +} + +impl WorkerLeaderLeasePort for RedisWorkerLeaderLeasePort { + fn try_acquire( + &self, + request: WorkerLeaseRequest, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let lease_request = CredentialLeaseRequest { + scope: CredentialLeaseScope::ProviderTask, + resource_id: request.worker().to_string(), + owner_id: self.owner_id.clone(), + ttl: request.ttl(), + }; + let grant = CredentialLeaseRepository::acquire_credential_lease( + &self.repository, + &lease_request, + ) + .await + .map_err(|_| WorkerLeaseError::safe("worker lease acquisition failed"))?; + let Some(grant) = grant else { + return Ok(WorkerLeaseAcquisition::Busy { retry_after: None }); + }; + let token = NonZeroU64::new(grant.fencing_token.get()) + .map(WorkerFencingToken::new) + .ok_or_else(|| WorkerLeaseError::safe("worker fencing token is invalid"))?; + Ok(WorkerLeaseAcquisition::Acquired(Box::new( + RedisWorkerLeaderLeaseGuard { + repository: self.repository.clone(), + request: lease_request, + grant: Some(grant), + token, + }, + ))) + }) + } +} diff --git a/backend/crates/gateway-store/tests/postgres/client_keys.rs b/backend/crates/gateway-store/tests/postgres/client_keys.rs index 05fe730ea..bbdeee5fb 100644 --- a/backend/crates/gateway-store/tests/postgres/client_keys.rs +++ b/backend/crates/gateway-store/tests/postgres/client_keys.rs @@ -1,7 +1,15 @@ +use gateway_admin::{ + model::client_keys::{ + ClientKeyListQuery as AdminClientKeyListQuery, ClientKeyPageSize, + ClientKeySort as AdminClientKeySort, ClientKeySortField as AdminClientKeySortField, + SortDirection as AdminSortDirection, + }, + ports::store::ClientKeyStore as _, +}; use gateway_store::postgres::{ ClientApiKeyCursor, ClientApiKeyCursorValue, ClientApiKeyListQuery, ClientApiKeyRepository, ClientApiKeySort, ClientApiKeySortDirection, ClientApiKeySortField, NewClientApiKey, - PgClientApiKeyRepository, + PgAdminClientKeyStore, PgClientApiKeyRepository, }; use super::TestDatabase; @@ -90,6 +98,16 @@ async fn client_key_list_uses_safe_keyset_search_and_filtered_total() { #[test] fn client_key_cursor_is_bound_to_one_sort_contract() { let created_sort = ClientApiKeySort::default(); + assert!( + ClientApiKeyListQuery { + cursor: None, + page_size: u16::MAX, + search: None, + sort: created_sort, + } + .validate() + .is_ok() + ); let cursor = ClientApiKeyCursor::new( created_sort, ClientApiKeyCursorValue::CreatedAt(chrono::Utc::now()), @@ -116,6 +134,29 @@ fn client_key_cursor_is_bound_to_one_sort_contract() { ); } +#[tokio::test] +async fn admin_client_key_adapter_should_preserve_the_full_nonzero_u16_page_size() { + let Some(database) = TestDatabase::create("admin_client_key_max_page").await else { + return; + }; + let page = PgAdminClientKeyStore::new(database.pool.clone()) + .list_client_keys(AdminClientKeyListQuery { + cursor: None, + page_size: ClientKeyPageSize::new(u16::MAX).expect("maximum page size"), + search: None, + sort: AdminClientKeySort { + field: AdminClientKeySortField::CreatedAt, + direction: AdminSortDirection::Desc, + }, + }) + .await + .expect("maximum Client Key page size"); + + assert_eq!(page.total, 0); + assert!(page.items.is_empty()); + database.close().await; +} + #[tokio::test] async fn client_key_database_sort_is_stable_and_keeps_null_last_used_at_last() { let Some(database) = TestDatabase::create("client_key_sort").await else { diff --git a/backend/crates/gateway-store/tests/postgres/execution.rs b/backend/crates/gateway-store/tests/postgres/execution.rs index c8a1c9ee4..068a44001 100644 --- a/backend/crates/gateway-store/tests/postgres/execution.rs +++ b/backend/crates/gateway-store/tests/postgres/execution.rs @@ -149,6 +149,8 @@ async fn core_adapter_should_persist_calculated_cost_exactly() { upstream_status_code: Some(200), upstream_request_id: None, upstream_response_id: None, + upstream_transport: Some("websocket".to_owned()), + http_version: Some("HTTP/2".to_owned()), error: None, provider_error_code: None, retry_after_ms: None, @@ -162,8 +164,8 @@ async fn core_adapter_should_persist_calculated_cost_exactly() { ) .await .expect("persist calculated cost"); - let persisted: (String, String, String) = sqlx::query_as( - "select cost_source, cost_amount::text, cost_currency + let persisted: (String, String, String, String, String) = sqlx::query_as( + "select cost_source, cost_amount::text, cost_currency, upstream_transport, http_version from model_requests where id = 'req_calculated_cost'", ) .fetch_one(&database.pool) @@ -176,6 +178,8 @@ async fn core_adapter_should_persist_calculated_cost_exactly() { "calculated".to_owned(), "0.0000012345".to_owned(), "USD".to_owned(), + "websocket".to_owned(), + "HTTP/2".to_owned(), ) ); database.close().await; diff --git a/backend/crates/gateway-store/tests/postgres/history.rs b/backend/crates/gateway-store/tests/postgres/history.rs index e1e26067a..7fcf895ab 100644 --- a/backend/crates/gateway-store/tests/postgres/history.rs +++ b/backend/crates/gateway-store/tests/postgres/history.rs @@ -86,7 +86,7 @@ async fn native_pin_is_caller_scoped_and_uses_live_account_and_upstream_handle() } #[tokio::test] -async fn native_pin_requires_committed_success_and_websocket_transport() { +async fn native_pin_requires_committed_success_but_not_a_specific_transport() { let Some(database) = TestDatabase::create("native_pin_boundaries").await else { return; }; @@ -164,23 +164,6 @@ async fn native_pin_requires_committed_success_and_websocket_transport() { ) .await .expect("resolve HTTP response") - .is_none() - ); - sqlx::query( - "update model_requests set upstream_transport = 'websocket' where id = 'req_history'", - ) - .execute(&database.pool) - .await - .expect("restore WebSocket response"); - assert!( - repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) - .await - .expect("resolve restored WebSocket response") .is_some() ); diff --git a/backend/crates/gateway-store/tests/postgres/observability.rs b/backend/crates/gateway-store/tests/postgres/observability.rs index 25af82686..1a47cb157 100644 --- a/backend/crates/gateway-store/tests/postgres/observability.rs +++ b/backend/crates/gateway-store/tests/postgres/observability.rs @@ -1,8 +1,14 @@ +use std::num::NonZeroU32; + use chrono::{TimeDelta, Utc}; +use gateway_admin::{ + model::{PageSize, observability as admin_observability}, + ports::store::ObservabilityStore as AdminObservabilityStore, +}; use gateway_store::postgres::{ DiagnosticDimension, ObservabilityPageNumber, ObservabilityPageSize, ObservabilityRange, - ObservabilityRepository, OpsErrorFilter, OpsErrorQuery, PgObservabilityRepository, - ProviderAccountUsageQuery, UsageRecordFilter, UsageRecordQuery, + ObservabilityRepository, OpsErrorFilter, OpsErrorQuery, PgAdminObservabilityStore, + PgObservabilityRepository, ProviderAccountUsageQuery, UsageRecordFilter, UsageRecordQuery, }; use sqlx::PgPool; @@ -24,12 +30,301 @@ fn observability_range_accepts_full_configured_retention_window() { assert_eq!(range.end, now); } +#[test] +fn usage_outcome_filter_should_accept_bounded_unknown_values() { + assert!( + UsageRecordFilter { + outcome: Some("provider_future_state".to_owned()), + ..UsageRecordFilter::default() + } + .validate() + .is_ok() + ); + assert!( + UsageRecordFilter { + outcome: Some("a".repeat(257)), + ..UsageRecordFilter::default() + } + .validate() + .is_err() + ); +} + #[test] fn postgres_observability_adapter_implements_query_port() { fn assert_port() {} assert_port::(); } +#[test] +fn postgres_admin_observability_adapter_implements_terminal_port() { + fn assert_port() {} + assert_port::(); +} + +#[tokio::test] +async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_details() { + let Some(database) = TestDatabase::create("admin_observability").await else { + return; + }; + let now = Utc::now(); + seed_observability_facts(&database.pool, now) + .await + .expect("seed observability facts"); + let range = + admin_observability::TimeRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1)) + .expect("admin observability range"); + let store = PgAdminObservabilityStore::new(database.pool.clone()); + + let dashboard = store + .dashboard_summary(range) + .await + .expect("admin dashboard summary"); + assert_eq!(dashboard.range, range); + assert_eq!(dashboard.requests.request_count, 2); + assert_eq!(dashboard.requests.first_token_latency_sum_ms, 120); + assert_eq!(dashboard.requests.latency_sum_ms, 1_600); + assert_eq!(dashboard.requests.min_latency_ms, Some(700)); + assert_eq!(dashboard.requests.max_latency_ms, Some(900)); + assert_eq!(dashboard.attempts.attempt_count, 3); + assert_eq!(dashboard.attempts.cost_coverage.provider_reported_count, 1); + assert_eq!(dashboard.attempts.cost_coverage.unavailable_count, 1); + assert_eq!(dashboard.attempts.costs[0].amount.as_str(), "1.25"); + assert_eq!(dashboard.provider_accounts.total, 1); + assert_eq!(dashboard.recent_requests.len(), 1); + assert_eq!( + dashboard.recent_requests[0].outcome, + admin_observability::RequestOutcome::Succeeded, + ); + assert_eq!( + dashboard.recent_requests[0] + .cost_amount + .as_ref() + .expect("dashboard request cost") + .as_str(), + "1.25", + ); + + let dashboard_trend = store + .dashboard_trend(range) + .await + .expect("admin dashboard trend"); + assert_eq!( + dashboard_trend + .iter() + .map(|point| point.metrics.request_count) + .sum::(), + 2, + ); + + let trend = store + .usage_trend( + range, + admin_observability::UsageFilter { + outcome: Some(admin_observability::RequestOutcome::Succeeded), + ..admin_observability::UsageFilter::default() + }, + ) + .await + .expect("admin usage trend"); + assert_eq!( + trend + .iter() + .map(|point| point.metrics.request_count) + .sum::(), + 1, + ); + assert_eq!( + trend + .iter() + .flat_map(|point| &point.costs) + .next() + .expect("trend cost") + .amount + .as_str(), + "1.25", + ); + + let first_page = store + .list_usage_records(admin_observability::UsageQuery { + range, + filter: admin_observability::UsageFilter::default(), + cursor: None, + page: admin_page(1), + page_size: PageSize::new(1).expect("page size"), + }) + .await + .expect("first usage page"); + assert_eq!(first_page.total, 2); + assert_eq!(first_page.items.len(), 1); + let second_page = store + .list_usage_records(admin_observability::UsageQuery { + range, + filter: admin_observability::UsageFilter::default(), + cursor: first_page.next_cursor, + page: admin_page(1), + page_size: PageSize::new(1).expect("page size"), + }) + .await + .expect("second usage page"); + assert_eq!(second_page.items.len(), 1); + assert_ne!(first_page.items[0].id, second_page.items[0].id); + + let filtered = store + .list_usage_records(admin_observability::UsageQuery { + range, + filter: admin_observability::UsageFilter { + client_api_key_ref: Some("key_observe".to_owned()), + request_id: Some("req_observe_success".to_owned()), + provider_account_ref: Some("acct_observe".to_owned()), + operation: Some("responses".to_owned()), + provider_kind: Some("openai".to_owned()), + model: Some("upstream-model".to_owned()), + outcome: Some(admin_observability::RequestOutcome::Succeeded), + status_code: Some(200), + transport: Some("http_sse".to_owned()), + attempt_index: Some(1), + response_id: Some("resp_observe_success".to_owned()), + upstream_request_id: Some("upstream_req_success".to_owned()), + search: Some("req_observe_success".to_owned()), + }, + cursor: None, + page: admin_page(1), + page_size: PageSize::new(10).expect("page size"), + }) + .await + .expect("fully filtered usage page"); + assert_eq!(filtered.total, 1); + assert_eq!( + filtered.items[0].upstream_request_id.as_deref(), + Some("upstream_req_success") + ); + + let other_outcome = store + .list_usage_records(admin_observability::UsageQuery { + range, + filter: admin_observability::UsageFilter { + outcome: Some( + admin_observability::RequestOutcome::new("provider_future_state") + .expect("bounded other outcome"), + ), + ..admin_observability::UsageFilter::default() + }, + cursor: None, + page: admin_page(1), + page_size: PageSize::new(10).expect("page size"), + }) + .await + .expect("other outcome filter should reach PostgreSQL"); + assert_eq!(other_outcome.total, 0); + assert!(other_outcome.items.is_empty()); + + let detail = store + .usage_record_detail("req_observe_failed") + .await + .expect("admin usage detail"); + assert_eq!( + detail.request.outcome, + admin_observability::RequestOutcome::Failed + ); + assert_eq!(detail.attempts.len(), 2); + assert_eq!( + detail.attempts[0].outcome, + admin_observability::RequestOutcome::Failed + ); + assert_eq!( + detail.attempts[1].outcome, + admin_observability::RequestOutcome::Failed + ); + + let overview = store + .usage_summary(range, admin_observability::UsageFilter::default()) + .await + .expect("admin usage overview"); + assert_eq!(overview.range, range); + assert_eq!(overview.providers[0].provider_kind, "openai"); + assert_eq!(overview.attempts.costs[0].amount.as_str(), "1.25"); + + let diagnostics = store + .usage_diagnostics( + range, + admin_observability::UsageFilter::default(), + admin_observability::DiagnosticDimension::Account, + ) + .await + .expect("admin diagnostics"); + assert_eq!(diagnostics[0].name, "acct_observe"); + assert_eq!(diagnostics[0].cost_coverage.provider_reported_count, 1); + assert_eq!(diagnostics[0].costs[0].amount.as_str(), "1.25"); + + let errors = store + .list_ops_errors(admin_observability::OpsErrorQuery { + range, + filter: admin_observability::OpsErrorFilter { + request_id: Some("req_observe_failed".to_owned()), + provider_kind: Some("openai".to_owned()), + provider_account_ref: Some("acct_observe".to_owned()), + model: Some("upstream-model".to_owned()), + failure_kind: Some("rate_limited".to_owned()), + status_code: Some(429), + search: Some("limited".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + cursor: None, + page: admin_page(1), + page_size: PageSize::new(10).expect("page size"), + }) + .await + .expect("admin ops errors"); + assert_eq!(errors.total, 2); + assert!(errors.items.iter().all(|item| item.occurred_at <= now)); + + for filter in [ + admin_observability::OpsErrorFilter { + client_api_key_ref: Some("missing-key".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + admin_observability::OpsErrorFilter { + operation: Some("missing-operation".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + admin_observability::OpsErrorFilter { + transport: Some("missing-transport".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + admin_observability::OpsErrorFilter { + attempt_index: Some(99), + ..admin_observability::OpsErrorFilter::default() + }, + admin_observability::OpsErrorFilter { + response_id: Some("missing-response".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + admin_observability::OpsErrorFilter { + upstream_request_id: Some("missing-upstream-request".to_owned()), + ..admin_observability::OpsErrorFilter::default() + }, + ] { + let page = store + .list_ops_errors(admin_observability::OpsErrorQuery { + range, + filter, + cursor: None, + page: admin_page(1), + page_size: PageSize::new(10).expect("page size"), + }) + .await + .expect("fully forwarded ops filter"); + assert_eq!(page.total, 0); + } + + database.close().await; +} + +fn admin_page(value: u32) -> admin_observability::PageNumber { + admin_observability::PageNumber::new(NonZeroU32::new(value).expect("positive page")) +} + #[tokio::test] async fn observability_queries_preserve_request_account_cost_and_diagnostic_facts() { let Some(database) = TestDatabase::create("observability").await else { @@ -229,7 +524,7 @@ async fn seed_observability_facts( provider_instance_id, provider_kind, provider_account_id, provider_account_ref, upstream_model_id, upstream_transport, attempt_count, upstream_send_state, downstream_committed_at, outcome, client_status_code, - upstream_status_code, client_response_id, upstream_response_id, + upstream_status_code, client_response_id, upstream_request_id, upstream_response_id, input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens, total_tokens, cost_source, cost_amount, cost_currency, first_token_ms, latency_ms, started_at, deadline_at, completed_at @@ -238,7 +533,8 @@ async fn seed_observability_facts( 'http_sse', 'public-model', 100, 'inst_observe', 'openai', 'acct_observe', 'acct_observe', 'upstream-model', 'http_sse', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200, - 'resp_observe_success', 'upstream_resp_success', 100, 20, 40, 3, 5, 120, + 'resp_observe_success', 'upstream_req_success', 'upstream_resp_success', + 100, 20, 40, 3, 5, 120, 'provider_reported', 1.25, 'USD', 120, 900, $1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes' )", diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs index 31b1389ea..ff58e02eb 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs @@ -1,18 +1,34 @@ use std::time::{Duration, SystemTime}; use chrono::{TimeDelta, Utc}; -use gateway_core::engine::credential::{ - CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, PlaintextCredential, - ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, +use gateway_admin::{ + model::{ + MutationActor, MutationContext, PageSize, Revision as AdminRevision, + accounts::{ + AccountListQuery, AccountSort, AccountSortField, AccountStatus, DeleteAccount, + SetAccountEnabled as AdminSetAccountEnabled, SortDirection, + }, + observability::TimeRange, + provider_credentials::{ + CredentialAvailabilityFilter, CredentialListQuery, CredentialListWindow, + }, + }, + ports::store::{AccountStore, AdminStoreErrorKind}, }; +use gateway_core::engine::credential::{ + CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, OpaqueProviderData, + PlaintextCredential, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, + QuotaObservation, QuotaWriteOutcome, +}; +use gateway_core::routing::ProviderKind; use gateway_store::{ ConflictKind, JsonObject, Revision, StoreError, postgres::{ AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts, - NewProviderAccount, PgProviderAccountRepository, ProviderAccountAdminRepository, - ProviderAccountAdminScope, ProviderAccountAvailability, ProviderAccountObservation, - ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount, - SetProviderAccountEnabled, UpdateProviderAccount, + NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository, + ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountAvailability, + ProviderAccountObservation, ProviderAccountRepository, ProviderCredentialUpdate, + RotateProviderAccount, SetProviderAccountEnabled, UpdateProviderAccount, }, }; use serde_json::json; @@ -26,6 +42,491 @@ fn postgres_provider_account_adapter_implements_core_port() { fn assert_admin_port() {} assert_admin_port::(); + + fn assert_terminal_admin_port() {} + assert_terminal_admin_port::(); +} + +#[tokio::test] +async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() { + let Some(database) = TestDatabase::create("provider_account_quota_batch").await else { + return; + }; + seed_instance(&database.pool, "inst_openai_admin", "openai") + .await + .expect("seed OpenAI instance"); + let repository = PgProviderAccountRepository::new(database.pool); + for id in ["acct_quota_a", "acct_quota_b", "acct_quota_empty"] { + repository + .insert_provider_account(account(id, &format!("user-{id}"))) + .await + .expect("insert quota fixture"); + } + let revision = CredentialRevision::new(1).expect("revision"); + for (id, remaining) in [("acct_quota_a", 20), ("acct_quota_b", 80)] { + let outcome = repository + .compare_and_swap_quota(QuotaObservation { + account_id: ProviderAccountId::new(id).expect("account id"), + expected_revision: revision, + quota: Some(OpaqueProviderData::new( + json!({"remaining": remaining}) + .as_object() + .expect("quota object") + .clone(), + )), + observed_at: Some(SystemTime::now()), + }) + .await + .expect("persist quota"); + assert_eq!(outcome, QuotaWriteOutcome::Updated); + } + + let mut observations = repository + .get_quotas(&[ + ProviderAccountId::new("acct_quota_b").expect("account id"), + ProviderAccountId::new("acct_quota_empty").expect("account id"), + ProviderAccountId::new("acct_quota_a").expect("account id"), + ]) + .await + .expect("read quota batch"); + observations.sort_by(|left, right| left.account_id.cmp(&right.account_id)); + + assert_eq!(observations.len(), 2); + assert_eq!(observations[0].account_id.as_str(), "acct_quota_a"); + assert_eq!(observations[0].expected_revision, revision); + assert_eq!( + observations[0] + .quota + .as_ref() + .expect("quota") + .expose_to_provider()["remaining"], + 20 + ); + assert!(observations.iter().all(|value| value.observed_at.is_some())); +} + +#[tokio::test] +async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_usage() { + let Some(database) = TestDatabase::create("provider_account_terminal_list").await else { + return; + }; + seed_instance(&database.pool, "inst_openai_admin", "openai") + .await + .expect("seed OpenAI instance"); + seed_instance(&database.pool, "inst_xai_admin", "xai") + .await + .expect("seed xAI instance"); + let repository = PgProviderAccountRepository::new(database.pool.clone()); + let now = Utc::now(); + + let mut alpha = account("acct_alpha", "user-alpha"); + alpha.email = Some("alpha@example.invalid".to_owned()); + let mut beta = account("acct_beta", "user-beta"); + beta.provider_instance_id = "inst_xai_admin".to_owned(); + beta.provider_kind = "xai".to_owned(); + beta.email = Some("beta@example.invalid".to_owned()); + beta.availability = ProviderAccountAvailability::Banned; + let mut charlie = account("acct_charlie", "user-charlie"); + charlie.email = Some("charlie@example.invalid".to_owned()); + let mut attention = account("acct_attention", "user-attention"); + attention.email = Some("attention@example.invalid".to_owned()); + attention.availability = ProviderAccountAvailability::Invalid; + for account in [alpha, beta, charlie, attention] { + repository + .insert_provider_account(account) + .await + .expect("insert account list fixture"); + } + + seed_model_request( + &database.pool, + ModelRequestSeed { + request_id: "req_alpha_recent", + account_id: "acct_alpha", + provider_instance_id: "inst_openai_admin", + provider_kind: "openai", + model: "gpt-list", + total_tokens: 10, + cost_amount: "0.10", + started_at: now - TimeDelta::minutes(20), + }, + ) + .await + .expect("seed alpha usage"); + seed_model_request( + &database.pool, + ModelRequestSeed { + request_id: "req_beta_recent", + account_id: "acct_beta", + provider_instance_id: "inst_xai_admin", + provider_kind: "xai", + model: "grok-list", + total_tokens: 50, + cost_amount: "0.50", + started_at: now - TimeDelta::minutes(5), + }, + ) + .await + .expect("seed beta usage"); + seed_model_request( + &database.pool, + ModelRequestSeed { + request_id: "req_beta_expired_retention", + account_id: "acct_beta", + provider_instance_id: "inst_xai_admin", + provider_kind: "xai", + model: "grok-list", + total_tokens: 500, + cost_amount: "5.00", + started_at: now - TimeDelta::days(40), + }, + ) + .await + .expect("seed expired beta usage"); + seed_model_request( + &database.pool, + ModelRequestSeed { + request_id: "req_charlie_recent", + account_id: "acct_charlie", + provider_instance_id: "inst_openai_admin", + provider_kind: "openai", + model: "gpt-list", + total_tokens: 60, + cost_amount: "0.60", + started_at: now - TimeDelta::minutes(10), + }, + ) + .await + .expect("seed charlie usage"); + + let store = PgAdminAccountStore::new(database.pool.clone()); + let usage_page = store + .list_accounts(AccountListQuery { + page: 1, + page_size: PageSize::new(2).expect("page size"), + provider_kind: None, + search: None, + status: None, + sort: Some(AccountSort { + field: AccountSortField::Usage, + direction: SortDirection::Desc, + }), + }) + .await + .expect("sort accounts by retained usage"); + assert_eq!(usage_page.config_revision.get(), 1); + assert_eq!(usage_page.total, 4); + assert_eq!(usage_page.summary.total, 4); + assert_eq!(usage_page.summary.active, 2); + assert_eq!(usage_page.summary.quota_exhausted, 0); + assert_eq!(usage_page.summary.attention, 2); + assert_eq!( + usage_page + .items + .iter() + .map(|account| account.id.as_str()) + .collect::>(), + ["acct_charlie", "acct_beta"] + ); + + let last_used_page = store + .list_accounts(AccountListQuery { + page: 1, + page_size: PageSize::new(2).expect("page size"), + provider_kind: None, + search: None, + status: None, + sort: Some(AccountSort { + field: AccountSortField::LastUsedAt, + direction: SortDirection::Desc, + }), + }) + .await + .expect("sort accounts by retained last use"); + assert_eq!( + last_used_page + .items + .iter() + .map(|account| account.id.as_str()) + .collect::>(), + ["acct_beta", "acct_charlie"] + ); + + let filtered = store + .list_accounts(AccountListQuery { + page: 1, + page_size: PageSize::new(10).expect("page size"), + provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")), + search: Some("ALPHA@EXAMPLE".to_owned()), + status: Some(AccountStatus::Active), + sort: None, + }) + .await + .expect("filter account directory"); + assert_eq!(filtered.total, 1); + assert_eq!(filtered.summary, usage_page.summary); + assert_eq!(filtered.items[0].id, "acct_alpha"); + assert_eq!(filtered.items[0].provider_kind.as_str(), "openai"); + + let banned = store + .list_accounts(AccountListQuery { + page: 1, + page_size: PageSize::new(10).expect("page size"), + provider_kind: None, + search: None, + status: Some(AccountStatus::Banned), + sort: None, + }) + .await + .expect("filter banned accounts"); + assert_eq!(banned.items[0].id, "acct_beta"); + let attention = store + .list_accounts(AccountListQuery { + page: 1, + page_size: PageSize::new(10).expect("page size"), + provider_kind: None, + search: None, + status: Some(AccountStatus::Attention), + sort: None, + }) + .await + .expect("filter attention accounts"); + assert_eq!(attention.items[0].id, "acct_attention"); + + database.close().await; +} + +#[tokio::test] +async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collections() { + let Some(database) = TestDatabase::create("provider_credential_terminal_windows").await else { + return; + }; + seed_instance(&database.pool, "inst_xai_admin", "xai") + .await + .expect("seed xAI instance"); + let repository = PgProviderAccountRepository::new(database.pool.clone()); + for index in 0..205_u16 { + let mut credential = account( + &format!("acct_xai_{index:03}"), + &format!("user-xai-{index:03}"), + ); + credential.provider_instance_id = "inst_xai_admin".to_owned(); + credential.provider_kind = "xai".to_owned(); + credential.availability = match index { + 0 => ProviderAccountAvailability::Expired, + 1 => ProviderAccountAvailability::Banned, + 2 => ProviderAccountAvailability::Invalid, + _ => ProviderAccountAvailability::Ready, + }; + repository + .insert_provider_account(credential) + .await + .expect("insert xAI credential fixture"); + } + + let store = PgAdminAccountStore::new(database.pool.clone()); + let provider = ProviderKind::new("xai").expect("xAI Provider kind"); + let complete = store + .list_credentials( + &provider, + CredentialListQuery { + provider_instance_id: None, + availability: None, + enabled: None, + window: CredentialListWindow::All, + }, + ) + .await + .expect("unpaged credential collection"); + assert_eq!(complete.items.len(), 205); + assert!(complete.next_cursor.is_none()); + + let page = store + .list_credentials( + &provider, + CredentialListQuery { + provider_instance_id: None, + availability: None, + enabled: None, + window: CredentialListWindow::Page { + cursor: None, + page_size: PageSize::new(200).expect("page size"), + }, + }, + ) + .await + .expect("paged credential collection"); + assert_eq!(page.items.len(), 200); + assert!(page.next_cursor.is_some()); + + let invalid_group = store + .list_credentials( + &provider, + CredentialListQuery { + provider_instance_id: None, + availability: Some(CredentialAvailabilityFilter::AnyOf(vec![ + gateway_admin::model::accounts::AccountAvailability::Expired, + gateway_admin::model::accounts::AccountAvailability::Banned, + gateway_admin::model::accounts::AccountAvailability::Invalid, + ])), + enabled: None, + window: CredentialListWindow::All, + }, + ) + .await + .expect("grouped invalid credential filter"); + assert_eq!(invalid_group.items.len(), 3); + + database.close().await; +} + +#[tokio::test] +async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs() { + let Some(database) = TestDatabase::create("provider_account_terminal_usage").await else { + return; + }; + seed_instance(&database.pool, "inst_openai_admin", "openai") + .await + .expect("seed OpenAI instance"); + let repository = PgProviderAccountRepository::new(database.pool.clone()); + repository + .insert_provider_account(account("acct_usage_exact", "user-usage-exact")) + .await + .expect("insert exact usage account"); + let now = Utc::now(); + seed_model_request( + &database.pool, + ModelRequestSeed { + request_id: "req_usage_exact", + account_id: "acct_usage_exact", + provider_instance_id: "inst_openai_admin", + provider_kind: "openai", + model: "gpt-exact", + total_tokens: 18, + cost_amount: "1.2345678901", + started_at: now - TimeDelta::minutes(1), + }, + ) + .await + .expect("seed exact usage request"); + let mut account_ids = vec!["acct_usage_exact".to_owned()]; + account_ids.extend((0..200).map(|index| format!("missing_account_{index}"))); + + let usage = PgAdminAccountStore::new(database.pool.clone()) + .load_account_usage( + TimeRange { + start: now - TimeDelta::hours(1), + end: now + TimeDelta::hours(1), + }, + &account_ids, + ) + .await + .expect("load account usage in bounded chunks"); + assert_eq!(usage.len(), 1); + assert_eq!(usage[0].account_id, "acct_usage_exact"); + assert_eq!(usage[0].request_count, 1); + assert_eq!(usage[0].success_count, 1); + assert_eq!(usage[0].total_tokens, Some(18)); + assert_eq!(usage[0].costs[0].currency, "USD"); + assert_eq!(usage[0].costs[0].amount.as_str(), "1.2345678901"); + assert_eq!(usage[0].models.len(), 1); + assert_eq!(usage[0].models[0].model, "gpt-exact"); + assert_eq!(usage[0].models[0].costs[0].amount.as_str(), "1.2345678901"); + + database.close().await; +} + +#[tokio::test] +async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() { + let Some(database) = TestDatabase::create("provider_account_terminal_mutation").await else { + return; + }; + seed_instance(&database.pool, "inst_openai_admin", "openai") + .await + .expect("seed OpenAI instance"); + PgProviderAccountRepository::new(database.pool.clone()) + .insert_provider_account(account("acct_terminal_mutation", "user-terminal-mutation")) + .await + .expect("insert mutation account"); + let store = PgAdminAccountStore::new(database.pool.clone()); + let context = MutationContext { + actor: MutationActor::System, + request_id: "request_terminal_mutation".to_owned(), + }; + + let revision = store + .set_account_enabled( + AdminSetAccountEnabled { + expected_config_revision: AdminRevision::new(1).expect("initial revision"), + account_id: "acct_terminal_mutation".to_owned(), + enabled: false, + }, + &context, + ) + .await + .expect("disable account atomically"); + assert_eq!(revision.get(), 2); + let enabled: bool = sqlx::query_scalar( + "select enabled from provider_accounts where id = 'acct_terminal_mutation'", + ) + .fetch_one(&database.pool) + .await + .expect("read disabled state"); + assert!(!enabled); + + let error = store + .delete_account( + DeleteAccount { + expected_config_revision: AdminRevision::new(1).expect("stale revision"), + account_id: "acct_terminal_mutation".to_owned(), + }, + &context, + ) + .await + .expect_err("stale delete must roll back"); + assert_eq!(error.kind(), AdminStoreErrorKind::StaleRevision); + assert_eq!(current_revision(&database.pool).await, 2); + assert_eq!( + account_count(&database.pool, "acct_terminal_mutation").await, + 1 + ); + let audit_count: i64 = sqlx::query_scalar("select count(*) from admin_audit_events") + .fetch_one(&database.pool) + .await + .expect("count audit rows after rollback"); + assert_eq!(audit_count, 1); + + let revision = store + .delete_account( + DeleteAccount { + expected_config_revision: revision, + account_id: "acct_terminal_mutation".to_owned(), + }, + &context, + ) + .await + .expect("delete disabled account atomically"); + assert_eq!(revision.get(), 3); + assert_eq!( + account_count(&database.pool, "acct_terminal_mutation").await, + 0 + ); + let audit_rows: Vec<(String, i64, Vec)> = sqlx::query_as( + "select action, config_revision, changed_fields + from admin_audit_events order by config_revision", + ) + .fetch_all(&database.pool) + .await + .expect("load terminal account audits"); + assert_eq!( + audit_rows, + vec![ + ("disable".to_owned(), 2, vec!["enabled".to_owned()]), + ("delete".to_owned(), 3, Vec::new()), + ] + ); + + database.close().await; } #[tokio::test] @@ -64,6 +565,8 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it() let mut updated = account("acct_admin_upsert", "user-admin-upsert"); updated.name = "updated import".to_owned(); updated.provider_credentials_json = credential_json("updated-import-secret"); + updated.availability = ProviderAccountAvailability::Banned; + updated.availability_reason = Some("upstream_account_deactivated".to_owned()); let revision = repository .import_provider_accounts( revision, @@ -76,8 +579,16 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it() .await .expect("update the same imported identity"); assert_eq!(revision.get(), 3); - let row: (String, serde_json::Value, i64, Option) = sqlx::query_as( - "select name, provider_credentials_json, credential_revision, provider_quota_json + let row: ( + String, + serde_json::Value, + i64, + Option, + String, + Option, + ) = sqlx::query_as( + "select name, provider_credentials_json, credential_revision, provider_quota_json, + availability, availability_reason from provider_accounts where id = 'acct_admin_upsert'", ) .fetch_one(&database.pool) @@ -90,6 +601,8 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it() row.3.is_none(), "credential replacement must clear stale quota" ); + assert_eq!(row.4, "banned"); + assert_eq!(row.5.as_deref(), Some("upstream_account_deactivated")); repository .import_provider_accounts( @@ -132,6 +645,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() { next_refresh_at: None, enabled: true, availability: ProviderAccountAvailability::Ready, + availability_reason: None, cooldown_until: None, availability_observed_at: Utc::now(), }) @@ -423,6 +937,7 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount { next_refresh_at: None, enabled: true, availability: ProviderAccountAvailability::Ready, + availability_reason: None, cooldown_until: None, availability_observed_at: Utc::now(), } @@ -481,6 +996,52 @@ fn audit(id: &str, action: &str, entity_ref: &str) -> AdminAuditEvent { } } +struct ModelRequestSeed<'a> { + request_id: &'a str, + account_id: &'a str, + provider_instance_id: &'a str, + provider_kind: &'a str, + model: &'a str, + total_tokens: i64, + cost_amount: &'a str, + started_at: chrono::DateTime, +} + +async fn seed_model_request( + pool: &sqlx::PgPool, + seed: ModelRequestSeed<'_>, +) -> Result<(), sqlx::Error> { + sqlx::query( + "insert into model_requests ( + id, client_api_key_ref, config_revision, protocol, operation, endpoint, + client_transport, requested_model_id, input_token_estimate, + provider_instance_id, provider_kind, provider_account_id, + provider_account_ref, upstream_model_id, upstream_transport, attempt_count, + upstream_send_state, outcome, client_status_code, upstream_status_code, + input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens, + total_tokens, cost_source, cost_amount, cost_currency, + started_at, deadline_at, completed_at + ) values ( + $1, 'key-provider-account-test', 1, 'openai', 'responses', '/v1/responses', + 'http_sse', $5, 1, $3, $4, $2, $2, $5, 'http_sse', 1, + 'sent', 'succeeded', 200, 200, $6, 0, 0, 0, 0, + $6, 'provider_reported', $7::numeric, 'USD', $8, + $8 + interval '5 minutes', $8 + interval '1 second' + )", + ) + .bind(seed.request_id) + .bind(seed.account_id) + .bind(seed.provider_instance_id) + .bind(seed.provider_kind) + .bind(seed.model) + .bind(seed.total_tokens) + .bind(seed.cost_amount) + .bind(seed.started_at) + .execute(pool) + .await?; + Ok(()) +} + async fn seed_instance( pool: &sqlx::PgPool, id: &str, diff --git a/backend/crates/gateway-store/tests/redis/credential_cooldown.rs b/backend/crates/gateway-store/tests/redis/credential_cooldown.rs index 7a85126ac..e6434a68e 100644 --- a/backend/crates/gateway-store/tests/redis/credential_cooldown.rs +++ b/backend/crates/gateway-store/tests/redis/credential_cooldown.rs @@ -83,10 +83,10 @@ impl ProviderAccountStore for CooldownTestAccountStore { )) } - async fn get_quota( + async fn get_quotas( &self, - _account: &ProviderAccountId, - ) -> Result, CoreStoreError> { + _accounts: &[ProviderAccountId], + ) -> Result, CoreStoreError> { unreachable!("cooldown tests do not read quota") } diff --git a/backend/crates/gateway-store/tests/redis/credential_leases.rs b/backend/crates/gateway-store/tests/redis/credential_leases.rs index 066555ae8..48397f5e3 100644 --- a/backend/crates/gateway-store/tests/redis/credential_leases.rs +++ b/backend/crates/gateway-store/tests/redis/credential_leases.rs @@ -1,9 +1,8 @@ use std::time::Duration; use gateway_store::redis::{ - CredentialLeaseRepository, CredentialLeaseRequest, CredentialLeaseScope, - CredentialSchedulingLeaseAcquisition, CredentialSchedulingLeaseRequest, - RedisCredentialLeaseRepository, + CredentialBoundedLeaseAcquisition, CredentialBoundedLeaseRequest, CredentialLeaseRepository, + CredentialLeaseRequest, CredentialLeaseScope, RedisCredentialLeaseRepository, }; use redis::aio::ConnectionManager; use uuid::Uuid; @@ -36,14 +35,14 @@ async fn scheduling_lease_counts_concurrency_interval_and_drop_release() { let first = acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("acquire first slot"), ); let first_fence = first.grant().expect("live grant").fencing_token.get(); let second = acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("acquire second slot"), ); @@ -55,12 +54,12 @@ async fn scheduling_lease_counts_concurrency_interval_and_drop_release() { assert!(signals[0].last_started_at.is_some()); let third = repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("capacity decision"); assert!(matches!( third, - CredentialSchedulingLeaseAcquisition::Busy { + CredentialBoundedLeaseAcquisition::Busy { retry_after: Some(_) } )); @@ -68,7 +67,7 @@ async fn scheduling_lease_counts_concurrency_interval_and_drop_release() { assert!(first.release().await.expect("release first slot")); let replacement = acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("reuse released slot"), ); @@ -110,7 +109,7 @@ async fn scheduling_lease_counts_concurrency_interval_and_drop_release() { ); acquired( repository - .try_acquire_scheduling_lease(&interval_request) + .try_acquire_bounded_lease(&interval_request) .await .expect("acquire interval slot"), ) @@ -118,12 +117,12 @@ async fn scheduling_lease_counts_concurrency_interval_and_drop_release() { .await .expect("release interval slot"); let interval_denied = repository - .try_acquire_scheduling_lease(&interval_request) + .try_acquire_bounded_lease(&interval_request) .await .expect("interval decision"); assert!(matches!( interval_denied, - CredentialSchedulingLeaseAcquisition::Busy { + CredentialBoundedLeaseAcquisition::Busy { retry_after: Some(value) } if value > Duration::ZERO && value <= Duration::from_secs(30) )); @@ -154,7 +153,7 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() { request.ttl = Duration::from_millis(25); let crashed = acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("acquire crash lease"), ); @@ -168,7 +167,7 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() { assert_eq!(recovered[0].in_flight, 0); acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("reacquire after TTL"), ) @@ -178,7 +177,7 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() { let cached = acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("acquire before cache loss"), ); @@ -202,7 +201,7 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() { assert_eq!(rebuilt[0].last_started_at, None); acquired( repository - .try_acquire_scheduling_lease(&request) + .try_acquire_bounded_lease(&request) .await .expect("acquire after cache rebuild"), ) @@ -211,6 +210,52 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() { .expect("release rebuilt lease"); } +#[tokio::test] +async fn scheduling_cursor_is_shared_across_gateway_processes() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + let peer = RedisCredentialLeaseRepository::new(connection.clone(), &namespace) + .expect("peer repository"); + + assert_eq!( + repository + .advance_scheduling_cursor("route_openai_real") + .await + .expect("first cursor"), + 0 + ); + assert_eq!( + peer.advance_scheduling_cursor("route_openai_real") + .await + .expect("shared cursor"), + 1 + ); + assert_eq!( + peer.advance_scheduling_cursor("route_openai_backup") + .await + .expect("independent cursor"), + 0 + ); + + let keys = redis::cmd("KEYS") + .arg(format!("{namespace}:*")) + .query_async::>(&mut connection) + .await + .expect("list cursor keys"); + assert_eq!(keys.len(), 2); + assert!( + keys.iter().all(|key| { + !key.contains("route_openai_real") && !key.contains("route_openai_backup") + }) + ); + redis::cmd("DEL") + .arg(keys) + .query_async::(&mut connection) + .await + .expect("clean cursor keys"); +} + #[tokio::test] async fn generic_refresh_guard_is_exclusive_and_released_on_drop() { let Some((repository, mut connection, namespace)) = repository().await else { @@ -271,13 +316,74 @@ async fn generic_refresh_guard_is_exclusive_and_released_on_drop() { } } +#[tokio::test] +async fn refresh_capacity_is_shared_across_provider_callers() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + let request = |owner_id: &str| CredentialBoundedLeaseRequest { + scope: CredentialLeaseScope::OAuthRefreshCapacity, + resource_id: "oauth-refresh-global".to_owned(), + owner_id: owner_id.to_owned(), + max_concurrent: 2, + request_interval: Duration::ZERO, + ttl: Duration::from_secs(60), + }; + + let openai = acquired( + repository + .try_acquire_bounded_lease(&request("openai-worker")) + .await + .expect("acquire OpenAI capacity"), + ); + let xai = acquired( + repository + .try_acquire_bounded_lease(&request("xai-worker")) + .await + .expect("acquire xAI capacity"), + ); + assert!(matches!( + repository + .try_acquire_bounded_lease(&request("second-openai-worker")) + .await + .expect("read shared capacity"), + CredentialBoundedLeaseAcquisition::Busy { .. } + )); + + openai.release().await.expect("release OpenAI capacity"); + acquired( + repository + .try_acquire_bounded_lease(&request("replacement-xai-worker")) + .await + .expect("reuse shared capacity"), + ) + .release() + .await + .expect("release replacement capacity"); + xai.release().await.expect("release xAI capacity"); + + let keys = redis::cmd("KEYS") + .arg(format!("{namespace}:*")) + .query_async::>(&mut connection) + .await + .expect("list isolated capacity keys"); + if !keys.is_empty() { + redis::cmd("DEL") + .arg(keys) + .query_async::(&mut connection) + .await + .expect("clean isolated capacity keys"); + } +} + fn scheduling_request( resource_id: &str, owner_id: &str, max_concurrent: u32, request_interval: Duration, -) -> CredentialSchedulingLeaseRequest { - CredentialSchedulingLeaseRequest { +) -> CredentialBoundedLeaseRequest { + CredentialBoundedLeaseRequest { + scope: CredentialLeaseScope::ProviderAccount, resource_id: resource_id.to_owned(), owner_id: owner_id.to_owned(), max_concurrent, @@ -287,11 +393,11 @@ fn scheduling_request( } fn acquired( - acquisition: CredentialSchedulingLeaseAcquisition, + acquisition: CredentialBoundedLeaseAcquisition, ) -> gateway_store::redis::CredentialLeaseGuard { match acquisition { - CredentialSchedulingLeaseAcquisition::Acquired(guard) => guard, - CredentialSchedulingLeaseAcquisition::Busy { retry_after } => { + CredentialBoundedLeaseAcquisition::Acquired(guard) => guard, + CredentialBoundedLeaseAcquisition::Busy { retry_after } => { panic!("expected acquired lease, retry after {retry_after:?}") } } diff --git a/backend/crates/gateway-store/tests/redis/mod.rs b/backend/crates/gateway-store/tests/redis/mod.rs index 23c989ed2..bf9c76cba 100644 --- a/backend/crates/gateway-store/tests/redis/mod.rs +++ b/backend/crates/gateway-store/tests/redis/mod.rs @@ -2,8 +2,10 @@ mod client_admission; mod credential_cooldown; mod credential_leases; mod credential_state; +mod oauth_pending; mod provider_circuit; mod runtime_change; +mod worker_lease; use chrono::Utc; use gateway_store::redis::{ diff --git a/backend/crates/gateway-store/tests/redis/oauth_pending.rs b/backend/crates/gateway-store/tests/redis/oauth_pending.rs new file mode 100644 index 000000000..b9933ae93 --- /dev/null +++ b/backend/crates/gateway-store/tests/redis/oauth_pending.rs @@ -0,0 +1,81 @@ +use std::{sync::Arc, time::Duration}; + +use gateway_core::{ + engine::credential::OpaqueProviderData, + provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome, + OAuthPendingTakeOutcome, + }, + routing::ProviderKind, +}; +use gateway_store::redis::RedisOAuthPendingFlowRepository; +use redis::aio::ConnectionManager; +use uuid::Uuid; + +#[tokio::test] +async fn codex_pending_flow_wrong_owner_does_not_consume_it() { + let Some(repository) = repository().await else { + return; + }; + let provider = ProviderKind::new("openai").expect("provider kind"); + let flow = + OAuthPendingBinding::try_new(format!("flow-{}", Uuid::new_v4())).expect("flow binding"); + let owner = OAuthPendingBinding::try_new("owner-a").expect("owner binding"); + let other_owner = OAuthPendingBinding::try_new("owner-b").expect("other owner binding"); + let mut payload = serde_json::Map::new(); + payload.insert( + "state".to_owned(), + serde_json::Value::String("opaque".to_owned()), + ); + let payload = OpaqueProviderData::new(payload); + + assert_eq!( + repository + .put_if_absent( + NewOAuthPendingFlow::try_new( + provider.clone(), + flow.clone(), + owner.clone(), + Duration::from_secs(60), + payload, + ) + .expect("pending flow"), + ) + .await + .expect("store pending flow"), + OAuthPendingPutOutcome::Stored + ); + assert_eq!( + repository + .take_if_owner(&provider, &flow, &other_owner) + .await + .expect("reject wrong owner"), + OAuthPendingTakeOutcome::OwnerMismatch + ); + assert!(matches!( + repository + .take_if_owner(&provider, &flow, &owner) + .await + .expect("consume right owner"), + OAuthPendingTakeOutcome::Taken(_) + )); + assert_eq!( + repository + .take_if_owner(&provider, &flow, &owner) + .await + .expect("flow is one shot"), + OAuthPendingTakeOutcome::NotFound + ); +} + +async fn repository() -> Option> { + let redis_url = std::env::var("CPR_TEST_REDIS_URL").ok()?; + let client = redis::Client::open(redis_url).ok()?; + let connection: ConnectionManager = client.get_connection_manager().await.ok()?; + RedisOAuthPendingFlowRepository::new( + connection, + &format!("gateway-store-test-{}", Uuid::new_v4()), + ) + .ok() + .map(Arc::new) +} diff --git a/backend/crates/gateway-store/tests/redis/provider_circuit.rs b/backend/crates/gateway-store/tests/redis/provider_circuit.rs index 8de890676..bc673f3e1 100644 --- a/backend/crates/gateway-store/tests/redis/provider_circuit.rs +++ b/backend/crates/gateway-store/tests/redis/provider_circuit.rs @@ -1,5 +1,6 @@ use std::{num::NonZeroU32, time::Duration}; +use gateway_core::{engine::execution::provider_failure_affects_circuit, error::ProviderErrorKind}; use gateway_store::redis::{ ProviderCircuitDecision, ProviderCircuitPolicy, ProviderCircuitRepository, RedisProviderCircuitRepository, @@ -12,6 +13,30 @@ fn provider_circuit_default_has_positive_threshold() { assert!(ProviderCircuitPolicy::default().failure_threshold.get() > 0); } +#[test] +fn provider_circuit_should_only_count_instance_attributable_failures() { + for error_kind in [ + ProviderErrorKind::Timeout, + ProviderErrorKind::Transport, + ProviderErrorKind::Protocol, + ProviderErrorKind::Unavailable, + ] { + assert!(provider_failure_affects_circuit(error_kind)); + } + for error_kind in [ + ProviderErrorKind::InvalidRequest, + ProviderErrorKind::Unsupported, + ProviderErrorKind::Unauthorized, + ProviderErrorKind::PermissionDenied, + ProviderErrorKind::RateLimited, + ProviderErrorKind::QuotaExhausted, + ProviderErrorKind::Cancelled, + ProviderErrorKind::ProcessTerminated, + ] { + assert!(!provider_failure_affects_circuit(error_kind)); + } +} + #[tokio::test] async fn provider_circuit_opens_at_threshold_and_success_resets_it() { let Some((repository, mut connection, namespace)) = repository(2).await else { diff --git a/backend/crates/gateway-store/tests/redis/worker_lease.rs b/backend/crates/gateway-store/tests/redis/worker_lease.rs new file mode 100644 index 000000000..c8b877f60 --- /dev/null +++ b/backend/crates/gateway-store/tests/redis/worker_lease.rs @@ -0,0 +1,96 @@ +//! Worker leader lease 通过 StoreBundle 的中立能力进行集成验证。 + +use std::collections::BTreeSet; +use std::time::Duration; + +use gateway_core::task::{WorkerId, WorkerKind, WorkerLeaseAcquisition, WorkerLeaseRequest}; +use gateway_store::{StoreConfig, initialize}; +use uuid::Uuid; + +#[tokio::test] +async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() { + let (Some(database_url), Some(redis_url)) = ( + std::env::var("CPR_TEST_DATABASE_URL").ok(), + std::env::var("CPR_TEST_REDIS_URL").ok(), + ) else { + return; + }; + let config = store_config(&database_url, &redis_url); + let mut first = initialize(config.clone()) + .await + .expect("first Store bundle"); + let second = initialize(config).await.expect("second Store bundle"); + + let kinds = first + .take_worker_contributions() + .into_iter() + .map(|contribution| contribution.kind()) + .collect::>(); + assert_eq!( + kinds, + BTreeSet::from([ + WorkerKind::NativeClaimRecovery, + WorkerKind::StaleModelRequestRecovery, + WorkerKind::Retention, + WorkerKind::OpsFlush, + ]) + ); + assert!(first.take_worker_contributions().is_empty()); + + let worker = WorkerId::try_new( + WorkerKind::Retention, + format!("test-{}", Uuid::new_v4().simple()), + ) + .expect("worker ID"); + let request = + WorkerLeaseRequest::try_new(worker, Duration::from_secs(5)).expect("worker lease request"); + let first_port = first.worker_leader_lease(); + let second_port = second.worker_leader_lease(); + let mut first_guard = match first_port + .try_acquire(request.clone()) + .await + .expect("first acquisition") + { + WorkerLeaseAcquisition::Acquired(guard) => guard, + WorkerLeaseAcquisition::Busy { .. } => panic!("fresh worker lease must be available"), + }; + let first_token = first_guard.fencing_token(); + assert!(matches!( + second_port + .try_acquire(request.clone()) + .await + .expect("contended acquisition"), + WorkerLeaseAcquisition::Busy { .. } + )); + first_guard.renew().await.expect("leader lease renewal"); + first_guard.release().await.expect("leader lease release"); + + let second_guard = match second_port + .try_acquire(request) + .await + .expect("acquisition after release") + { + WorkerLeaseAcquisition::Acquired(guard) => guard, + WorkerLeaseAcquisition::Busy { .. } => panic!("released worker lease must be reusable"), + }; + assert!(second_guard.fencing_token() > first_token); + second_guard.release().await.expect("second lease release"); +} + +fn store_config(database_url: &str, redis_url: &str) -> StoreConfig { + let (database_url, database_password) = split_connection_url(database_url); + let (redis_url, redis_password) = split_connection_url(redis_url); + serde_json::from_value(serde_json::json!({ + "database": { "url": database_url, "password": database_password }, + "redis": { "url": redis_url, "password": redis_password }, + })) + .expect("test Store config") +} + +fn split_connection_url(value: &str) -> (String, String) { + let mut url = url::Url::parse(value).expect("test connection URL"); + let password = url.password().expect("test connection password").to_owned(); + url.set_password(None) + .expect("connection URL supports credentials"); + (url.to_string(), password) +} diff --git a/backend/crates/providers/openai/Cargo.toml b/backend/crates/providers/openai/Cargo.toml index f09210de2..a6f00fde5 100644 --- a/backend/crates/providers/openai/Cargo.toml +++ b/backend/crates/providers/openai/Cargo.toml @@ -15,12 +15,15 @@ bytes = "1.11.1" chrono = { version = "0.4.45", features = ["clock", "serde"] } cookie = "0.18.1" futures = "0.3.32" +gateway-admin = { path = "../../gateway-admin" } gateway-core = { path = "../../gateway-core" } gateway-protocol = { path = "../../gateway-protocol" } getrandom = "0.3.4" hex = "0.4.3" jsonwebtoken = { version = "=10.3.0", default-features = false, features = ["rust_crypto"] } reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] } +roxmltree = "0.21.1" +semver = "1.0.27" rustls = { version = "=0.23.42", default-features = false, features = ["ring", "std", "tls12"] } rustls-native-certs = "0.8.4" rustls-pki-types = "1.14.1" diff --git a/backend/crates/providers/openai/src/admin.rs b/backend/crates/providers/openai/src/admin.rs new file mode 100644 index 000000000..8b057006a --- /dev/null +++ b/backend/crates/providers/openai/src/admin.rs @@ -0,0 +1,1268 @@ +//! OpenAI 管理边界：Provider preparation 与 Redis OAuth pending 适配。 + +use std::sync::Arc; +use std::time::SystemTime; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use gateway_admin::model::accounts::{ + AccountAvailability as AdminAccountAvailability, AccountRecord, +}; +use gateway_admin::model::observability::{ + CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireProfile, + DashboardWireTarget, DecimalAmount, DesktopReleaseStatus, ProviderBillingInput, +}; +use gateway_admin::model::provider_credentials::{ + AuthorizationMutationTarget, AuthorizationOwner, AuthorizationOwnerBinding, + AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedAuthorizationCredential, + PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation, + PreparedCredentialRotationFacts, ProviderDocument, ProviderExport, + ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota, + ProviderQuotaWindow, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision}; +use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind}; +use gateway_core::accounting::Money; +use gateway_core::engine::credential::{ + AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount, + OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId, + ProviderAccountStore, +}; +use gateway_core::error::StoreErrorKind; +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome, + OAuthPendingTakeOutcome, ProviderInstanceCatalogPort, ProviderRuntimePolicyPort, + ProviderStoreError, ProviderStoreErrorKind, +}; +use gateway_core::routing::{ + InstanceHealth, ProviderInstance, ProviderInstanceId, ProviderKind, UpstreamModelId, +}; +use secrecy::{ExposeSecret as _, SecretString}; +use serde::Deserialize; +use serde_json::{Map, Number, Value}; + +use crate::credential::{ + CodexAccountIdentityVerifier, CodexCredentialCodec, CodexIdentityExpectation, CodexOAuthSecret, + oauth_owner_ref, +}; +use crate::credential::{ + CodexAccountQuotaSnapshot, CodexCredentialAdmin, CodexCredentialAdminError, + CodexCredentialAdminService, CodexCredentialCatalogError, CodexCredentialCatalogService, + CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError, + CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexPendingAuthorization, + CodexQuotaWindowKind, CodexQuotaWindowRole, CompleteCodexOAuthAuthorization, + CompletedCodexOAuthCredential, ExportManagedCodexCredential, RotateManagedCodexCredential, + StartCodexOAuthAuthorization, StoredCodexPendingAuthorization, refresh_time, +}; +use crate::transport::CodexWebSocketPool; +use crate::transport::openai_billing_breakdown; +use crate::transport::profile::{ + CodexDesktopReleaseSnapshot, CodexDesktopReleaseStatus, CodexWireProfile, CodexWireProfileState, +}; + +const PROVIDER_NAME: &str = "openai"; +const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 1; +const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024; +const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024; + +/// OpenAI 对终态 Admin port 的唯一实现。 +pub(crate) struct OpenAiAdminProvider { + provider_kind: ProviderKind, + profile: CodexWireProfileState, + accounts: Arc, + instances: Arc, + credentials: Arc, + verifier: Arc, + oauth: Arc, + quota: Arc, + catalog: Arc, + runtime_policy: Arc, + websocket_pool: Arc, + desktop_release: CodexDesktopReleaseStatus, +} + +pub(crate) struct OpenAiAdminServices { + pub(crate) credentials: Arc, + pub(crate) verifier: Arc, + pub(crate) oauth: Arc, + pub(crate) quota: Arc, + pub(crate) catalog: Arc, + pub(crate) runtime_policy: Arc, +} + +impl OpenAiAdminProvider { + #[must_use] + pub(crate) fn new( + provider_kind: ProviderKind, + profile: CodexWireProfileState, + accounts: Arc, + instances: Arc, + services: OpenAiAdminServices, + websocket_pool: Arc, + desktop_release: CodexDesktopReleaseStatus, + ) -> Self { + Self { + provider_kind, + profile, + accounts, + instances, + credentials: services.credentials, + verifier: services.verifier, + oauth: services.oauth, + quota: services.quota, + catalog: services.catalog, + runtime_policy: services.runtime_policy, + websocket_pool, + desktop_release, + } + } + + async fn account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.accounts + .get_account(account_id) + .await + .map_err(map_store_error)? + .filter(|account| account.provider() == &self.provider_kind) + .ok_or_else(|| provider_admin_error(ProviderAdminErrorKind::NotFound)) + } + + async fn instance( + &self, + account: &ProviderAccount, + ) -> Result { + let instance = self + .instances + .list_instances(&self.provider_kind, true) + .await + .map_err(map_provider_store_error)? + .into_iter() + .find(|instance| instance.id() == account.instance()) + .ok_or_else(|| provider_admin_error(ProviderAdminErrorKind::NotFound))?; + if instance.provider_kind() != &self.provider_kind { + return Err(provider_admin_error(ProviderAdminErrorKind::Internal)); + } + Ok(ProviderInstance::new( + instance.id().clone(), + instance.provider_kind().clone(), + instance.base_url().to_owned(), + instance.enabled(), + InstanceHealth::Healthy, + )) + } + + async fn preserve_existing_installation_id( + &self, + mut incoming: NewProviderAccount, + ) -> Result { + let Some(existing) = self + .accounts + .get_account(incoming.account.id()) + .await + .map_err(map_store_error)? + else { + return Ok(incoming); + }; + if existing.provider() != incoming.account.provider() + || existing.instance() != incoming.account.instance() + || existing.upstream_user_id() != incoming.account.upstream_user_id() + || existing.upstream_account_id() != incoming.account.upstream_account_id() + { + return Ok(incoming); + } + let current = self + .accounts + .load_credential(existing.id(), existing.revision()) + .await + .map_err(map_store_error)?; + incoming.credential = CodexCredentialCodec::preserve_installation_id( + &incoming.credential, + ¤t.credential, + ) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + Ok(incoming) + } +} + +#[async_trait] +impl ProviderAdmin for OpenAiAdminProvider { + fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + async fn account_unavailable(&self, account_id: &ProviderAccountId) { + self.websocket_pool.evict_account(account_id.as_str()).await; + } + + fn dashboard_wire_profile(&self) -> Option { + let profile = self.profile.snapshot(); + let release = self.desktop_release.snapshot(); + let user_agent = profile.user_agent(); + let release = dashboard_desktop_release(&profile, release); + Some(DashboardWireProfile { + originator: profile.originator, + codex_version: profile.codex_version, + desktop_version: profile.desktop_version, + desktop_build: profile.desktop_build, + target: DashboardWireTarget { + os_type: profile.os_type, + os_version: profile.os_version, + arch: profile.arch, + terminal: profile.terminal, + }, + user_agent, + verified_at: profile.verified_at, + release, + }) + } + + fn calculated_billing( + &self, + input: &ProviderBillingInput, + ) -> Result, ProviderAdminError> { + let (Some(input_tokens), Some(output_tokens)) = (input.input_tokens, input.output_tokens) + else { + return Ok(None); + }; + let Some(breakdown) = openai_billing_breakdown( + &input.upstream_model_id, + input_tokens, + output_tokens, + input.cached_tokens.unwrap_or_default(), + input.cache_write_tokens.unwrap_or_default(), + None, + ) else { + return Ok(None); + }; + let total_amount = currency_cost(breakdown.total_amount())?; + if total_amount != input.total { + return Ok(None); + } + Ok(Some(CalculatedBillingBreakdown { + input_amount: currency_cost(breakdown.input_amount())?, + output_amount: currency_cost(breakdown.output_amount())?, + cache_read_amount: currency_cost(breakdown.cache_read_amount())?, + cache_write_amount: currency_cost(breakdown.cache_write_amount())?, + standard_amount: currency_cost(breakdown.standard_amount())?, + total_amount, + input_price_per_million: currency_cost(breakdown.input_price_per_million())?, + output_price_per_million: currency_cost(breakdown.output_price_per_million())?, + cache_read_price_per_million: currency_cost(breakdown.cache_read_price_per_million())?, + cache_write_price_per_million: currency_cost( + breakdown.cache_write_price_per_million(), + )?, + service_tier: breakdown.service_tier().map(str::to_owned), + multiplier_percent: breakdown.multiplier_percent(), + })) + } + + async fn prepare_import( + &self, + command: PrepareCredentialImport, + ) -> Result { + let prepared = self + .credentials + .prepare_import_document( + command.provider_instance_id.clone(), + Value::Object(command.document.into_provider_data().into_inner()), + ) + .await + .map_err(map_credential_admin_error)?; + let observed_at = Utc::now(); + let mut credentials = Vec::with_capacity(prepared.accounts().len()); + for account in prepared.into_accounts() { + let account = self.preserve_existing_installation_id(account).await?; + credentials.push(prepared_create(account, observed_at)?); + } + Ok(PreparedCredentialImport { + provider_kind: self.provider_kind.clone(), + provider_instance_id: command.provider_instance_id, + credentials, + }) + } + + async fn start_authorization( + &self, + pending: PendingAuthorizationMutation, + ) -> Result { + if pending.provider_kind() != &self.provider_kind { + return Err(provider_admin_error(ProviderAdminErrorKind::Invalid)); + } + let started = self + .oauth + .start_authorization(StartCodexOAuthAuthorization { mutation: pending }) + .await + .map_err(map_oauth_error)?; + Ok(AuthorizationStarted { + flow_id: started.flow_id, + authorization_url: started.authorization_url, + expires_at: started.expires_at, + }) + } + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result { + let binding = AuthorizationOwnerBinding::from_context(&command.context); + let completed = self + .oauth + .complete_authorization(CompleteCodexOAuthAuthorization { + owner_ref: oauth_owner_ref(binding.owner()), + flow_id: command.flow_id, + callback_url: SecretString::from(command.callback_url), + }) + .await + .map_err(map_oauth_error)?; + let credential = match completed.credential { + CompletedCodexOAuthCredential::Create(credential) => { + PreparedAuthorizationCredential::Create(prepared_create(credential, Utc::now())?) + } + CompletedCodexOAuthCredential::Reauthorize(credential) => { + let (provider_instance_id, provider_kind) = match completed.mutation.target() { + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + .. + } => ( + provider_instance_id.clone(), + completed.mutation.provider_kind().clone(), + ), + AuthorizationMutationTarget::Create { .. } => { + return Err(provider_admin_error(ProviderAdminErrorKind::Internal)); + } + }; + PreparedAuthorizationCredential::Reauthorize(prepared_rotation( + credential, + provider_instance_id, + provider_kind, + )?) + } + }; + Ok(PreparedAuthorizationCommit { + pending: completed.mutation, + credential, + }) + } + + async fn prepare_rotation( + &self, + command: PrepareCredentialRotation, + ) -> Result { + validate_account_record(&command.account, &self.provider_kind)?; + if command.account.credential_revision != command.expected_credential_revision { + return Err(provider_admin_error(ProviderAdminErrorKind::Conflict)); + } + let account_id = ProviderAccountId::new(command.account.id.clone()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let expected_revision = CredentialRevision::new(command.expected_credential_revision.get()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let current = self + .accounts + .load_credential(&account_id, expected_revision) + .await + .map_err(map_store_error)?; + if !account_matches_record(¤t.account, &command.account) { + return Err(provider_admin_error(ProviderAdminErrorKind::Conflict)); + } + let runtime = CodexCredentialCodec::decode(¤t.credential) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let account_id = current + .account + .upstream_account_id() + .ok_or_else(|| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let expectation = CodexIdentityExpectation::current( + runtime.principal.oauth_subject, + runtime.principal.poid, + account_id.to_owned(), + current.account.upstream_user_id().to_owned(), + runtime.installation_id, + ) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let secret = rotation_secret(command.provider_material)?; + let verified_account = self + .verifier + .verify(&secret, &expectation) + .await + .and_then(crate::credential::CodexIdentityVerification::into_complete) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(map_provider_store_error)?; + let next_refresh_at = refresh_time( + policy, + current.account.id(), + verified_account.access_token_expires_at, + secret.refresh_token.is_some(), + ) + .map_err(map_credential_admin_error)?; + let prepared = CodexCredentialAdmin + .prepare_rotation(RotateManagedCodexCredential { + current, + secret, + verified_account, + next_refresh_at, + }) + .map_err(map_credential_admin_error)?; + prepared_rotation( + prepared, + command.account.provider_instance_id, + command.account.provider_kind, + ) + } + + async fn prepare_refresh( + &self, + command: PrepareCredentialRefresh, + ) -> Result { + validate_account_record(&command.account, &self.provider_kind)?; + let account_id = ProviderAccountId::new(command.account.id.clone()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let revision = CredentialRevision::new(command.account.credential_revision.get()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let prepared = self + .credentials + .manual_refresh(account_id, revision) + .await + .map_err(map_credential_admin_error)?; + prepared_rotation( + prepared, + command.account.provider_instance_id, + command.account.provider_kind, + ) + } + + async fn quota( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + let account = self.account(account_id).await?; + let snapshot = if refresh { + let instance = self.instance(&account).await?; + Some( + self.quota + .refresh_account(&instance, account_id) + .await + .map_err(map_quota_error)?, + ) + } else { + self.quota + .read_account(account_id) + .await + .map_err(map_quota_error)? + }; + Ok(snapshot.map_or_else(empty_quota, project_quota)) + } + + async fn models( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + let account = self.account(account_id).await?; + let instance = self.instance(&account).await?; + let models = if refresh { + self.catalog + .synchronize_account(&instance, account_id) + .await + .map_err(map_catalog_error)? + } else { + self.catalog + .cached_account_models(instance.id(), account_id, account.revision()) + .map_err(map_catalog_error)? + .unwrap_or_default() + }; + let observed_at = self + .catalog + .cached(instance.id()) + .map_err(map_catalog_error)? + .map(|snapshot| DateTime::::from(snapshot.observed_at())); + let models = models + .into_iter() + .map(|model| { + let id = UpstreamModelId::new(model.clone()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Internal))?; + Ok(ProviderModel { id, name: model }) + }) + .collect::, ProviderAdminError>>()?; + Ok(ProviderModels { + models, + observed_at, + }) + } + + async fn export_credentials( + &self, + credentials: Vec, + ) -> Result { + let mut account_ids = Vec::with_capacity(credentials.len()); + let mut items = Vec::with_capacity(credentials.len()); + for input in credentials { + validate_account_record(&input.account, &self.provider_kind)?; + let current = LoadedCredential { + account: account_from_record(&input.account)?, + credential: PlaintextCredential::new( + input.provider_material.into_provider_data().into_inner(), + ), + }; + account_ids.push(current.account.id().clone()); + items.push(ExportManagedCodexCredential { + current, + added_at: input.account.created_at, + updated_at: input.account.updated_at, + }); + } + let document = CodexCredentialAdmin + .format_cpr_export(items) + .and_then(|document| document.into_json()) + .map_err(map_credential_admin_error)?; + let Value::Object(document) = document else { + return Err(provider_admin_error(ProviderAdminErrorKind::Internal)); + }; + Ok(ProviderExport { + provider_kind: self.provider_kind.clone(), + account_ids, + document: ProviderDocument::new(OpaqueProviderData::new(document)), + }) + } +} + +struct OpenAiCredentialCommitGuard { + _guard: crate::credential::PreparedCodexCredentialRotationGuard, +} + +impl CredentialCommitGuard for OpenAiCredentialCommitGuard { + fn finish(self: Box) {} +} + +fn prepared_create( + prepared: NewProviderAccount, + observed_at: DateTime, +) -> Result { + let NewProviderAccount { + account, + credential, + } = prepared; + let availability = admin_availability(account.availability()); + Ok(PreparedCredentialCreate { + account_id: account.id().clone(), + provider_instance_id: account.instance().clone(), + provider_kind: account.provider().clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + upstream_user_id: account.upstream_user_id().to_owned(), + upstream_account_id: account.upstream_account_id().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + provider_material: ProviderDocument::new(OpaqueProviderData::new(credential.into_inner())), + has_refresh_token: account.has_refresh_token(), + access_token_expires_at: DateTime::::from(account.access_token_expires_at()), + next_refresh_at: account.next_refresh_at().map(DateTime::::from), + enabled: account.enabled(), + availability, + availability_reason: None, + cooldown_until: account.cooldown_until().map(DateTime::::from), + availability_observed_at: observed_at, + }) +} + +fn prepared_rotation( + prepared: crate::credential::PreparedCodexCredentialRotation, + provider_instance_id: ProviderInstanceId, + provider_kind: ProviderKind, +) -> Result { + let (profile, credential, guard) = prepared.into_parts(); + let ( + account_id, + expected_revision, + credential_profile, + credential, + has_refresh_token, + access_token_expires_at, + next_refresh_at, + ) = credential.into_parts(); + if profile != credential_profile || profile.account_id != account_id { + return Err(provider_admin_error(ProviderAdminErrorKind::Internal)); + } + let expected_credential_revision = Revision::new(expected_revision.get()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Internal))?; + Ok(PreparedCredentialRotation::new( + PreparedCredentialRotationFacts { + account_id, + provider_instance_id, + provider_kind, + expected_credential_revision, + name: profile.name, + email: profile.email, + plan_type: profile.plan_type, + provider_material: ProviderDocument::new(OpaqueProviderData::new( + credential.into_inner(), + )), + has_refresh_token, + access_token_expires_at: DateTime::::from(access_token_expires_at), + next_refresh_at: next_refresh_at.map(DateTime::::from), + }, + Box::new(OpenAiCredentialCommitGuard { _guard: guard }), + )) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RotationDocument { + access_token: String, + refresh_token: Option, +} + +fn rotation_secret(document: ProviderDocument) -> Result { + let document: RotationDocument = + serde_json::from_value(Value::Object(document.into_provider_data().into_inner())) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + if document.access_token.is_empty() + || document.access_token.len() > MAX_ACCESS_TOKEN_BYTES + || document + .refresh_token + .as_deref() + .is_some_and(|token| token.is_empty() || token.len() > MAX_REFRESH_TOKEN_BYTES) + { + return Err(provider_admin_error(ProviderAdminErrorKind::Invalid)); + } + Ok(CodexOAuthSecret { + access_token: SecretString::from(document.access_token), + refresh_token: document.refresh_token.map(SecretString::from), + id_token: None, + }) +} + +fn validate_account_record( + account: &AccountRecord, + provider_kind: &ProviderKind, +) -> Result<(), ProviderAdminError> { + if &account.provider_kind != provider_kind || provider_kind.as_str() != PROVIDER_NAME { + return Err(provider_admin_error(ProviderAdminErrorKind::Invalid)); + } + ProviderAccountId::new(account.id.clone()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + Ok(()) +} + +fn account_from_record(account: &AccountRecord) -> Result { + let account_id = ProviderAccountId::new(account.id.clone()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + let revision = CredentialRevision::new(account.credential_revision.get()) + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?; + Ok(ProviderAccount::new( + account_id, + account.provider_instance_id.clone(), + account.provider_kind.clone(), + account.name.clone(), + account.upstream_user_id.clone(), + revision, + SystemTime::from(account.access_token_expires_at), + ) + .with_profile( + account.email.clone(), + account.upstream_account_id.clone(), + account.plan_type.clone(), + ) + .with_runtime_state( + account.enabled, + core_availability(account.availability), + account.cooldown_until.map(SystemTime::from), + ) + .with_refresh_schedule( + account.has_refresh_token, + account.next_refresh_at.map(SystemTime::from), + )) +} + +fn dashboard_desktop_release( + profile: &CodexWireProfile, + snapshot: CodexDesktopReleaseSnapshot, +) -> DashboardDesktopRelease { + let status = if snapshot.checked_at.is_none() { + DesktopReleaseStatus::Unchecked + } else if snapshot.last_error.is_some() { + DesktopReleaseStatus::Failed + } else if snapshot.latest.as_ref().is_some_and(|latest| { + latest.version == profile.desktop_version && latest.build == profile.desktop_build + }) { + DesktopReleaseStatus::Current + } else if snapshot.latest.is_some() { + DesktopReleaseStatus::UpdateAvailable + } else { + DesktopReleaseStatus::Failed + }; + let latest = snapshot.latest; + DashboardDesktopRelease { + status, + checked_at: snapshot.checked_at, + latest_version: latest.as_ref().map(|release| release.version.clone()), + latest_build: latest.as_ref().map(|release| release.build.clone()), + published_at: latest.as_ref().and_then(|release| release.published_at), + minimum_system_version: latest + .as_ref() + .and_then(|release| release.minimum_system_version.clone()), + hardware_requirements: latest + .as_ref() + .and_then(|release| release.hardware_requirements.clone()), + download_url: latest + .as_ref() + .and_then(|release| release.download_url.clone()), + download_size: latest.as_ref().and_then(|release| release.download_size), + signature_present: latest.as_ref().map(|release| release.signature_present), + error: snapshot.last_error, + } +} + +fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool { + account.id().as_str() == record.id + && account.instance() == &record.provider_instance_id + && account.provider() == &record.provider_kind + && account.revision().get() == record.credential_revision.get() + && account.name() == record.name + && account.email() == record.email.as_deref() + && account.upstream_user_id() == record.upstream_user_id + && account.upstream_account_id() == record.upstream_account_id.as_deref() + && account.plan_type() == record.plan_type.as_deref() + && account.enabled() == record.enabled + && admin_availability(account.availability()) == record.availability + && account.cooldown_until().map(DateTime::::from) == record.cooldown_until + && DateTime::::from(account.access_token_expires_at()) + == record.access_token_expires_at + && account.next_refresh_at().map(DateTime::::from) == record.next_refresh_at + && account.has_refresh_token() == record.has_refresh_token +} + +const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability { + match value { + AccountAvailability::Unknown => AdminAccountAvailability::Unknown, + AccountAvailability::Ready => AdminAccountAvailability::Ready, + AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown, + AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted, + AccountAvailability::Expired => AdminAccountAvailability::Expired, + AccountAvailability::Banned => AdminAccountAvailability::Banned, + AccountAvailability::Invalid => AdminAccountAvailability::Invalid, + } +} + +const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability { + match value { + AdminAccountAvailability::Unknown => AccountAvailability::Unknown, + AdminAccountAvailability::Ready => AccountAvailability::Ready, + AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown, + AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted, + AdminAccountAvailability::Expired => AccountAvailability::Expired, + AdminAccountAvailability::Banned => AccountAvailability::Banned, + AdminAccountAvailability::Invalid => AccountAvailability::Invalid, + } +} + +fn empty_quota() -> ProviderQuota { + ProviderQuota { + observed_at: None, + refresh_token_expires_at: None, + windows: Vec::new(), + provider_data: None, + } +} + +fn project_quota(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota { + let mut provider_data = Map::new(); + provider_data.insert( + "remaining_percent".to_owned(), + snapshot + .fact() + .remaining_percent() + .map_or(Value::Null, |value| Value::Number(Number::from(value))), + ); + provider_data.insert( + "exhausted".to_owned(), + Value::Bool(snapshot.fact().exhausted()), + ); + let windows = snapshot + .windows() + .iter() + .map(|window| { + let mut data = Map::new(); + data.insert( + "source".to_owned(), + Value::String(window.source().to_owned()), + ); + data.insert( + "role".to_owned(), + Value::String(quota_role(window.role()).to_owned()), + ); + ProviderQuotaWindow { + key: window.key().to_owned(), + group: quota_group(window.kind()).to_owned(), + source: Some(window.source().to_owned()), + window_seconds: window.window_seconds(), + used_percent: window.used_percent(), + reset_at: window.reset_at(), + provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))), + } + }) + .collect(); + ProviderQuota { + observed_at: Some(DateTime::::from(snapshot.observed_at())), + refresh_token_expires_at: None, + windows, + provider_data: Some(ProviderDocument::new(OpaqueProviderData::new( + provider_data, + ))), + } +} + +const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str { + match kind { + CodexQuotaWindowKind::Monthly => "monthly", + CodexQuotaWindowKind::ShortTerm | CodexQuotaWindowKind::Weekly => "shortTerm", + CodexQuotaWindowKind::Other => "other", + } +} + +const fn quota_role(role: CodexQuotaWindowRole) -> &'static str { + match role { + CodexQuotaWindowRole::Primary => "primary", + CodexQuotaWindowRole::Secondary => "secondary", + CodexQuotaWindowRole::Monthly => "monthly", + } +} + +/// 将 Provider-owned PKCE/OIDC 状态保存到 Store 提供的 Redis 原子端口。 +pub(crate) struct OpenAiOAuthPendingStore { + port: Arc, + provider_kind: ProviderKind, +} + +impl OpenAiOAuthPendingStore { + pub(crate) const fn new( + port: Arc, + provider_kind: ProviderKind, + ) -> Self { + Self { + port, + provider_kind, + } + } +} + +#[async_trait] +impl CodexOAuthPendingStore for OpenAiOAuthPendingStore { + async fn create( + &self, + pending: &CodexPendingAuthorization, + ) -> Result<(), CodexOAuthPendingStoreError> { + let now = Utc::now(); + let ttl = (pending.expires_at() - now) + .to_std() + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?; + let flow = NewOAuthPendingFlow::try_new( + self.provider_kind.clone(), + binding(pending.flow_id())?, + binding(pending.owner_ref())?, + ttl, + OpaqueProviderData::new(encode_pending(pending)), + ) + .map_err(map_pending_store_error)?; + match self + .port + .put_if_absent(flow) + .await + .map_err(map_pending_store_error)? + { + OAuthPendingPutOutcome::Stored => Ok(()), + OAuthPendingPutOutcome::AlreadyExists => Err(CodexOAuthPendingStoreError::Conflict), + } + } + + async fn take( + &self, + owner_ref: &str, + flow_id: &str, + ) -> Result, CodexOAuthPendingStoreError> { + let flow = binding(flow_id)?; + let owner = binding(owner_ref)?; + match self + .port + .take_if_owner(&self.provider_kind, &flow, &owner) + .await + .map_err(map_pending_store_error)? + { + OAuthPendingTakeOutcome::Taken(payload) => decode_pending(payload).map(Some), + OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => Ok(None), + } + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PendingDocument { + flow_id: String, + owner_ref: String, + started_request_ref: String, + provider_instance_id: String, + name: String, + expires_at: DateTime, + state: String, + nonce: String, + code_verifier: String, + reauthorization_account_id: Option, + reauthorization_credential_revision: Option, + mutation: PendingMutationDocument, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PendingMutationDocument { + schema_version: u64, + expected_config_revision: u64, + provider_kind: String, + target: PendingTargetDocument, + owner: PendingOwnerDocument, + started_request_id: String, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum PendingTargetDocument { + Create { + provider_instance_id: String, + name: String, + }, + Reauthorize { + provider_instance_id: String, + account_id: String, + expected_credential_revision: u64, + }, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum PendingOwnerDocument { + AdminSession { admin_user_id: String }, + AdminApiKey, + System, +} + +fn encode_pending(pending: &CodexPendingAuthorization) -> Map { + let mut document = Map::new(); + document.insert( + "flow_id".to_owned(), + Value::String(pending.flow_id().to_owned()), + ); + document.insert( + "owner_ref".to_owned(), + Value::String(pending.owner_ref().to_owned()), + ); + document.insert( + "started_request_ref".to_owned(), + Value::String(pending.started_request_ref().to_owned()), + ); + document.insert( + "provider_instance_id".to_owned(), + Value::String(pending.provider_instance_id().to_owned()), + ); + document.insert("name".to_owned(), Value::String(pending.name().to_owned())); + document.insert( + "expires_at".to_owned(), + Value::String(pending.expires_at().to_rfc3339()), + ); + document.insert( + "state".to_owned(), + Value::String(pending.state().expose_secret().to_owned()), + ); + document.insert( + "nonce".to_owned(), + Value::String(pending.nonce().expose_secret().to_owned()), + ); + document.insert( + "code_verifier".to_owned(), + Value::String(pending.code_verifier().expose_secret().to_owned()), + ); + document.insert( + "reauthorization_account_id".to_owned(), + pending.reauthorization().map_or(Value::Null, |target| { + Value::String(target.account_id().to_string()) + }), + ); + document.insert( + "reauthorization_credential_revision".to_owned(), + pending.reauthorization().map_or(Value::Null, |target| { + Value::Number(Number::from(target.credential_revision().get())) + }), + ); + document.insert( + "mutation".to_owned(), + Value::Object(encode_mutation(pending.mutation())), + ); + document +} + +fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map { + let mut document = Map::new(); + document.insert( + "schema_version".to_owned(), + Value::Number(Number::from(PENDING_DOCUMENT_SCHEMA_VERSION)), + ); + document.insert( + "expected_config_revision".to_owned(), + Value::Number(Number::from(mutation.expected_config_revision().get())), + ); + document.insert( + "provider_kind".to_owned(), + Value::String(mutation.provider_kind().as_str().to_owned()), + ); + document.insert( + "target".to_owned(), + Value::Object(encode_target(mutation.target())), + ); + document.insert( + "owner".to_owned(), + Value::Object(encode_owner(mutation.owner_binding().owner())), + ); + document.insert( + "started_request_id".to_owned(), + Value::String(mutation.owner_binding().started_request_id().to_owned()), + ); + document +} + +fn encode_target(target: &AuthorizationMutationTarget) -> Map { + let mut document = Map::new(); + match target { + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + } => { + document.insert("kind".to_owned(), Value::String("create".to_owned())); + document.insert( + "provider_instance_id".to_owned(), + Value::String(provider_instance_id.to_string()), + ); + document.insert("name".to_owned(), Value::String(name.clone())); + } + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => { + document.insert("kind".to_owned(), Value::String("reauthorize".to_owned())); + document.insert( + "provider_instance_id".to_owned(), + Value::String(provider_instance_id.to_string()), + ); + document.insert( + "account_id".to_owned(), + Value::String(account_id.to_string()), + ); + document.insert( + "expected_credential_revision".to_owned(), + Value::Number(Number::from(expected_credential_revision.get())), + ); + } + } + document +} + +fn encode_owner(owner: &AuthorizationOwner) -> Map { + let mut document = Map::new(); + match owner { + AuthorizationOwner::AdminSession { admin_user_id } => { + document.insert("kind".to_owned(), Value::String("admin_session".to_owned())); + document.insert( + "admin_user_id".to_owned(), + Value::String(admin_user_id.clone()), + ); + } + AuthorizationOwner::AdminApiKey => { + document.insert("kind".to_owned(), Value::String("admin_api_key".to_owned())); + } + AuthorizationOwner::System => { + document.insert("kind".to_owned(), Value::String("system".to_owned())); + } + } + document +} + +fn decode_pending( + payload: OpaqueProviderData, +) -> Result { + let document: PendingDocument = serde_json::from_value(Value::Object(payload.into_inner())) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?; + CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization { + flow_id: document.flow_id, + owner_ref: document.owner_ref, + started_request_ref: document.started_request_ref, + provider_instance_id: document.provider_instance_id, + name: document.name, + expires_at: document.expires_at, + state: SecretString::from(document.state), + nonce: SecretString::from(document.nonce), + code_verifier: SecretString::from(document.code_verifier), + reauthorization_account_id: document.reauthorization_account_id, + reauthorization_credential_revision: document.reauthorization_credential_revision, + mutation: decode_mutation(document.mutation)?, + }) +} + +fn decode_mutation( + document: PendingMutationDocument, +) -> Result { + if document.schema_version != PENDING_DOCUMENT_SCHEMA_VERSION { + return Err(CodexOAuthPendingStoreError::InvalidValue); + } + let expected_config_revision = Revision::new(document.expected_config_revision) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?; + let provider_kind = ProviderKind::new(document.provider_kind) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?; + let target = match document.target { + PendingTargetDocument::Create { + provider_instance_id, + name, + } => AuthorizationMutationTarget::Create { + provider_instance_id: ProviderInstanceId::new(provider_instance_id) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?, + name, + }, + PendingTargetDocument::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => AuthorizationMutationTarget::Reauthorize { + provider_instance_id: ProviderInstanceId::new(provider_instance_id) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?, + account_id: ProviderAccountId::new(account_id) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?, + expected_credential_revision: Revision::new(expected_credential_revision) + .map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?, + }, + }; + let actor = match document.owner { + PendingOwnerDocument::AdminSession { admin_user_id } => { + MutationActor::AdminSession { admin_user_id } + } + PendingOwnerDocument::AdminApiKey => MutationActor::AdminApiKey, + PendingOwnerDocument::System => MutationActor::System, + }; + let context = MutationContext { + actor, + request_id: document.started_request_id, + }; + Ok(PendingAuthorizationMutation::new( + expected_config_revision, + provider_kind, + target, + AuthorizationOwnerBinding::from_context(&context), + )) +} + +fn binding(value: &str) -> Result { + OAuthPendingBinding::try_new(value.to_owned()).map_err(map_pending_store_error) +} + +fn provider_admin_error(kind: ProviderAdminErrorKind) -> ProviderAdminError { + ProviderAdminError::new(kind) +} + +fn currency_cost(money: Money) -> Result { + Ok(CurrencyCost { + currency: money.currency().as_str().to_owned(), + amount: money + .amount() + .to_string() + .parse::() + .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Internal))?, + }) +} + +fn map_pending_store_error(error: ProviderStoreError) -> CodexOAuthPendingStoreError { + match error.kind() { + ProviderStoreErrorKind::InvalidData => CodexOAuthPendingStoreError::InvalidValue, + ProviderStoreErrorKind::Conflict => CodexOAuthPendingStoreError::Conflict, + ProviderStoreErrorKind::Unavailable => CodexOAuthPendingStoreError::Unavailable, + } +} + +fn map_provider_store_error(error: ProviderStoreError) -> ProviderAdminError { + provider_admin_error(match error.kind() { + ProviderStoreErrorKind::InvalidData => ProviderAdminErrorKind::Invalid, + ProviderStoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict, + ProviderStoreErrorKind::Unavailable => ProviderAdminErrorKind::Unavailable, + }) +} + +fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError { + provider_admin_error(match error.kind() { + StoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict, + StoreErrorKind::InvalidData | StoreErrorKind::InvalidState => { + ProviderAdminErrorKind::NotFound + } + StoreErrorKind::Unavailable => ProviderAdminErrorKind::Unavailable, + _ => ProviderAdminErrorKind::Internal, + }) +} + +fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdminError { + use CodexCredentialAdminError as Error; + provider_admin_error(match error { + Error::InvalidInput + | Error::InvalidCredential + | Error::IdentityMismatch + | Error::MissingRefreshToken + | Error::RefreshRejected + | Error::AccountBanned + | Error::IdentityRejected => ProviderAdminErrorKind::Invalid, + Error::NotFound => ProviderAdminErrorKind::NotFound, + Error::RevisionConflict | Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => { + ProviderAdminErrorKind::Conflict + } + Error::StoreUnavailable | Error::RefreshUnavailable | Error::IdentityUnavailable => { + ProviderAdminErrorKind::Unavailable + } + }) +} + +fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError { + use CodexOAuthAdminError as Error; + provider_admin_error(match error { + Error::InvalidInput | Error::UpstreamRejected | Error::Credential => { + ProviderAdminErrorKind::Invalid + } + Error::NotFound | Error::FlowExpired => ProviderAdminErrorKind::NotFound, + Error::Conflict | Error::Ambiguous => ProviderAdminErrorKind::Conflict, + Error::UpstreamUnavailable | Error::StorageUnavailable => { + ProviderAdminErrorKind::Unavailable + } + }) +} + +fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError { + use CodexCredentialQuotaError as Error; + provider_admin_error(match error { + Error::InvalidInstance | Error::InvalidCredentialData => ProviderAdminErrorKind::Invalid, + Error::NotFound => ProviderAdminErrorKind::NotFound, + Error::RevisionConflict => ProviderAdminErrorKind::Conflict, + Error::Repository(_) | Error::Store | Error::Upstream => { + ProviderAdminErrorKind::Unavailable + } + }) +} + +fn map_catalog_error(error: CodexCredentialCatalogError) -> ProviderAdminError { + use CodexCredentialCatalogError as Error; + provider_admin_error(match error { + Error::InvalidInstance + | Error::InvalidCredentialData + | Error::ConflictingModelFacts + | Error::InvalidEtag => ProviderAdminErrorKind::Invalid, + Error::NoEligibleCredential => ProviderAdminErrorKind::NotFound, + Error::ConcurrentUpdate => ProviderAdminErrorKind::Conflict, + Error::Upstream | Error::Cache => ProviderAdminErrorKind::Unavailable, + }) +} diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs new file mode 100644 index 000000000..d955c3f3f --- /dev/null +++ b/backend/crates/providers/openai/src/config.rs @@ -0,0 +1,117 @@ +//! OpenAI Provider 启动配置与 Codex Desktop 请求画像校验。 + +use std::path::Path; + +use chrono::{DateTime, Utc}; +use serde::Deserialize; + +use crate::transport::profile::{CodexWireProfile, CodexWireProfileState}; + +/// OpenAI Provider 唯一启动配置。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct OpenAiConfig { + pub wire_profile: CodexWireProfileConfig, +} + +impl OpenAiConfig { + /// 校验 Provider-owned 字段；当前配置不含文件系统路径。 + pub fn resolve_and_validate(&mut self, _source_dir: &Path) -> Result<(), OpenAiConfigError> { + self.wire_profile.validate() + } + + #[must_use] + pub fn wire_profile_state(&self) -> CodexWireProfileState { + CodexWireProfileState::new(self.wire_profile.clone().into()) + } +} + +/// 经审计固定的 Codex Desktop 上游请求画像。 +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexWireProfileConfig { + pub originator: String, + pub codex_version: String, + pub desktop_version: String, + pub desktop_build: String, + pub os_type: String, + pub os_version: String, + pub arch: String, + pub terminal: String, + pub verified_at: DateTime, +} + +impl CodexWireProfileConfig { + fn validate(&self) -> Result<(), OpenAiConfigError> { + for (field, value) in [ + ("openai.wire_profile.originator", self.originator.as_str()), + ( + "openai.wire_profile.codex_version", + self.codex_version.as_str(), + ), + ( + "openai.wire_profile.desktop_version", + self.desktop_version.as_str(), + ), + ( + "openai.wire_profile.desktop_build", + self.desktop_build.as_str(), + ), + ("openai.wire_profile.os_type", self.os_type.as_str()), + ("openai.wire_profile.os_version", self.os_version.as_str()), + ("openai.wire_profile.arch", self.arch.as_str()), + ("openai.wire_profile.terminal", self.terminal.as_str()), + ] { + if value.trim().is_empty() { + return Err(OpenAiConfigError::InvalidField(field)); + } + } + if semver::Version::parse(&self.codex_version).is_err() { + return Err(OpenAiConfigError::InvalidField( + "openai.wire_profile.codex_version", + )); + } + if !numeric_dotted_version(&self.desktop_version) { + return Err(OpenAiConfigError::InvalidField( + "openai.wire_profile.desktop_version", + )); + } + if !self.desktop_build.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(OpenAiConfigError::InvalidField( + "openai.wire_profile.desktop_build", + )); + } + Ok(()) + } +} + +impl From for CodexWireProfile { + fn from(value: CodexWireProfileConfig) -> Self { + Self { + originator: value.originator, + codex_version: value.codex_version, + desktop_version: value.desktop_version, + desktop_build: value.desktop_build, + os_type: value.os_type, + os_version: value.os_version, + arch: value.arch, + terminal: value.terminal, + verified_at: value.verified_at, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum OpenAiConfigError { + #[error("OpenAI configuration field is invalid: {0}")] + InvalidField(&'static str), +} + +fn numeric_dotted_version(value: &str) -> bool { + let mut parts = value.split('.'); + let valid_parts = parts + .by_ref() + .filter(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit())) + .count(); + valid_parts >= 2 && valid_parts == value.split('.').count() +} diff --git a/backend/crates/providers/openai/src/credential/admin.rs b/backend/crates/providers/openai/src/credential/admin.rs index f1e5803ad..322d54bb5 100644 --- a/backend/crates/providers/openai/src/credential/admin.rs +++ b/backend/crates/providers/openai/src/credential/admin.rs @@ -6,7 +6,7 @@ use std::collections::BTreeSet; use std::fmt; use std::sync::Arc; -use std::time::{Duration, SystemTime}; +use std::time::SystemTime; use chrono::{DateTime, FixedOffset, Utc}; use gateway_core::engine::credential::{ @@ -14,16 +14,19 @@ use gateway_core::engine::credential::{ NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate, }; use gateway_core::error::StoreErrorKind; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, +}; use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use secrecy::{ExposeSecret, SecretString}; use serde::Serialize; use serde_json::Value; use thiserror::Error; -use super::identity::{CodexIdentityVerificationError, CodexTokenIdentityVerifier}; -use super::refresh::{ - CodexRefreshLeaseAcquisition, CodexRefreshLeaseCoordinator, CodexRefreshLeaseGuard, - CodexRefreshLeaseRequest, +use super::identity::{ + CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError, }; use super::repository::CodexCredentialRepository; use super::security::CodexCredentialCodec; @@ -66,6 +69,7 @@ pub struct ImportCodexOAuthCredential { pub name: String, pub secret: CodexOAuthSecret, pub verified_account: CodexAccountProfile, + pub next_refresh_at: Option>, pub enabled: bool, } @@ -78,6 +82,7 @@ impl std::fmt::Debug for ImportCodexOAuthCredential { .field("name", &self.name) .field("secret", &"") .field("verified_account", &self.verified_account) + .field("next_refresh_at", &self.next_refresh_at) .field("enabled", &self.enabled) .finish() } @@ -117,6 +122,8 @@ impl fmt::Debug for PreparedCodexAccountImport { struct ParsedCodexImportAccount { id: Option, name: Option, + chatgpt_account_id: Option, + chatgpt_user_id: Option, access_token: Option, refresh_token: Option, id_token: Option, @@ -129,6 +136,14 @@ impl fmt::Debug for ParsedCodexImportAccount { .debug_struct("ParsedCodexImportAccount") .field("id", &self.id) .field("name", &self.name) + .field( + "chatgpt_account_id", + &self.chatgpt_account_id.as_ref().map(|_| ""), + ) + .field( + "chatgpt_user_id", + &self.chatgpt_user_id.as_ref().map(|_| ""), + ) .field( "access_token", &self.access_token.as_ref().map(|_| ""), @@ -221,6 +236,7 @@ pub struct RotateManagedCodexCredential { pub current: LoadedCredential, pub secret: CodexOAuthSecret, pub verified_account: CodexAccountProfile, + pub next_refresh_at: Option>, } impl std::fmt::Debug for RotateManagedCodexCredential { @@ -230,6 +246,7 @@ impl std::fmt::Debug for RotateManagedCodexCredential { .field("current", &self.current) .field("secret", &"") .field("verified_account", &self.verified_account) + .field("next_refresh_at", &self.next_refresh_at) .finish() } } @@ -238,7 +255,12 @@ impl std::fmt::Debug for RotateManagedCodexCredential { pub struct PreparedCodexCredentialRotation { pub profile: ProviderAccountUpdate, pub credential: CredentialCasUpdate, - refresh_guard: Option>, + refresh_guards: Option, +} + +struct ProviderRefreshGuards { + _capacity: Box, + _account: Box, } impl std::fmt::Debug for PreparedCodexCredentialRotation { @@ -248,8 +270,8 @@ impl std::fmt::Debug for PreparedCodexCredentialRotation { .field("profile", &self.profile) .field("credential", &self.credential) .field( - "refresh_guard", - &self.refresh_guard.as_ref().map(|_| ""), + "refresh_guards", + &self.refresh_guards.as_ref().map(|_| ""), ) .finish() } @@ -258,7 +280,7 @@ impl std::fmt::Debug for PreparedCodexCredentialRotation { impl PreparedCodexCredentialRotation { #[must_use] pub const fn holds_refresh_lease(&self) -> bool { - self.refresh_guard.is_some() + self.refresh_guards.is_some() } /// 将 command 与 lease 一起交给 App；App 必须让返回的 guard 活到 CAS 提交结束。 @@ -273,13 +295,13 @@ impl PreparedCodexCredentialRotation { ( self.profile, self.credential, - PreparedCodexCredentialRotationGuard(self.refresh_guard), + PreparedCodexCredentialRotationGuard(self.refresh_guards), ) } } /// 手工刷新从 token exchange 到数据库 CAS 完成期间持有的 Redis lease。 -pub struct PreparedCodexCredentialRotationGuard(Option>); +pub struct PreparedCodexCredentialRotationGuard(Option); impl fmt::Debug for PreparedCodexCredentialRotationGuard { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -344,13 +366,10 @@ impl CodexCredentialAdmin { required_time(input.verified_account.access_token_expires_at)?; let revision = CredentialRevision::new(1).map_err(|_| CodexCredentialAdminError::InvalidCredential)?; - let upstream_user_id = input - .verified_account - .chatgpt_user_id - .clone() - .unwrap_or_else(|| input.verified_account.chatgpt_account_id.clone()); - let credential = CodexCredentialCodec::encode(&input.secret, Vec::new()) - .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; + let upstream_user_id = input.verified_account.chatgpt_user_id.clone(); + let credential = + CodexCredentialCodec::encode_new(&input.secret, &input.verified_account, Vec::new()) + .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; let account = ProviderAccount::new( account_id, instance, @@ -368,7 +387,7 @@ impl CodexCredentialAdmin { .with_runtime_state(input.enabled, AccountAvailability::Ready, None) .with_refresh_schedule( input.secret.refresh_token.is_some(), - optional_time(input.verified_account.next_refresh_at), + optional_time(input.next_refresh_at), ); Ok(NewProviderAccount { account, @@ -445,21 +464,19 @@ impl CodexCredentialAdmin { &self, input: RotateManagedCodexCredential, ) -> Result { - if input.current.account.provider().as_str() != PROVIDER_NAME - || input.current.account.upstream_account_id() - != Some(input.verified_account.chatgpt_account_id.as_str()) - || input - .verified_account - .chatgpt_user_id - .as_deref() - .is_some_and(|user| user != input.current.account.upstream_user_id()) - { - return Err(CodexCredentialAdminError::IdentityMismatch); - } let access_token_expires_at = required_time(input.verified_account.access_token_expires_at)?; let mut data = CodexCredentialCodec::decode_complete(&input.current.credential) .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; + if input.current.account.provider().as_str() != PROVIDER_NAME + || input.current.account.upstream_account_id() + != Some(input.verified_account.chatgpt_account_id.as_str()) + || input.current.account.upstream_user_id() != input.verified_account.chatgpt_user_id + || data.principal.oauth_subject != input.verified_account.oauth_subject + || data.principal.poid != input.verified_account.poid + { + return Err(CodexCredentialAdminError::IdentityMismatch); + } data.access_token = input.secret.access_token.expose_secret().to_owned(); data.refresh_token = input .secret @@ -486,13 +503,13 @@ impl CodexCredentialAdmin { credential, input.secret.refresh_token.is_some(), access_token_expires_at, - optional_time(input.verified_account.next_refresh_at), + optional_time(input.next_refresh_at), ) .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; Ok(PreparedCodexCredentialRotation { profile, credential, - refresh_guard: None, + refresh_guards: None, }) } } @@ -501,9 +518,9 @@ impl CodexCredentialAdmin { pub struct CodexCredentialAdminService { repository: CodexCredentialRepository, refresher: Arc, - verifier: Arc, - leases: Arc, - refresh_margin: Duration, + verifier: Arc, + leases: Arc, + runtime_policy: Arc, } impl fmt::Debug for CodexCredentialAdminService { @@ -512,9 +529,9 @@ impl fmt::Debug for CodexCredentialAdminService { .debug_struct("CodexCredentialAdminService") .field("repository", &"CodexCredentialRepository") .field("refresher", &"TokenRefresher") - .field("verifier", &"CodexTokenIdentityVerifier") - .field("leases", &"CodexRefreshLeaseCoordinator") - .field("refresh_margin", &self.refresh_margin) + .field("verifier", &"CodexAccountIdentityVerifier") + .field("leases", &"ProviderLeasePort") + .field("runtime_policy", &"ProviderRuntimePolicyPort") .finish() } } @@ -523,20 +540,17 @@ impl CodexCredentialAdminService { pub fn new( repository: CodexCredentialRepository, refresher: Arc, - verifier: Arc, - leases: Arc, - refresh_margin: Duration, - ) -> Result { - if refresh_margin.is_zero() { - return Err(CodexCredentialAdminError::InvalidInput); - } - Ok(Self { + verifier: Arc, + leases: Arc, + runtime_policy: Arc, + ) -> Self { + Self { repository, refresher, verifier, leases, - refresh_margin, - }) + runtime_policy, + } } /// 官方 RT exchange + AT 身份验证；结果由 App 在同一 revision/audit 事务中提交。 @@ -559,22 +573,40 @@ impl CodexCredentialAdminService { } let runtime = CodexCredentialCodec::decode(¤t.credential) .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; + let expectation = current_identity_expectation(¤t.account, &runtime)?; let refresh_token = runtime .secret .refresh_token .as_ref() .ok_or(CodexCredentialAdminError::MissingRefreshToken)?; - let guard = match self + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(|_| CodexCredentialAdminError::RefreshUnavailable)?; + let capacity_guard = match self .leases - .try_acquire(&CodexRefreshLeaseRequest { - account_id: account_id.clone(), - credential_revision: expected_revision, - }) + .try_acquire(ProviderLeaseRequest::RefreshCapacity( + ProviderRefreshCapacityRequest::new(policy.concurrency()), + )) .await .map_err(|_| CodexCredentialAdminError::RefreshUnavailable)? { - CodexRefreshLeaseAcquisition::Acquired(guard) => guard, - CodexRefreshLeaseAcquisition::Unavailable => { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { + return Err(CodexCredentialAdminError::RefreshLeaseUnavailable); + } + }; + let account_guard = match self + .leases + .try_acquire(ProviderLeaseRequest::Refresh( + ProviderRefreshLeaseRequest::new(account_id.clone(), expected_revision), + )) + .await + .map_err(|_| CodexCredentialAdminError::RefreshUnavailable)? + { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { return Err(CodexCredentialAdminError::RefreshLeaseUnavailable); } }; @@ -583,7 +615,7 @@ impl CodexCredentialAdminService { .refresh(refresh_token.expose_secret()) .await .map_err(map_refresh_failure)?; - if tokens.access_token.trim().is_empty() || tokens.expires_in <= self.refresh_margin { + if tokens.access_token.trim().is_empty() || tokens.expires_in.is_zero() { return Err(CodexCredentialAdminError::InvalidCredential); } let secret = CodexOAuthSecret { @@ -596,15 +628,26 @@ impl CodexCredentialAdminService { }; let verified_account = self .verifier - .verify(&secret) + .verify(&secret, &expectation) .await + .and_then(super::identity::CodexIdentityVerification::into_complete) .map_err(map_identity_error)?; + let next_refresh_at = refresh_time( + policy, + &account_id, + verified_account.access_token_expires_at, + secret.refresh_token.is_some(), + )?; let mut prepared = CodexCredentialAdmin.prepare_rotation(RotateManagedCodexCredential { current, secret, verified_account, + next_refresh_at, })?; - prepared.refresh_guard = Some(guard); + prepared.refresh_guards = Some(ProviderRefreshGuards { + _capacity: capacity_guard, + _account: account_guard, + }); Ok(prepared) } @@ -628,12 +671,17 @@ impl CodexCredentialAdminService { let mut account_ids = BTreeSet::new(); let mut upstream_identities = BTreeSet::new(); let mut accounts = Vec::with_capacity(candidates.len()); + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(|_| CodexCredentialAdminError::RefreshUnavailable)?; for candidate in candidates { let (secret, verified_account) = self.verify_import_candidate(&candidate).await?; - if let Some(user_id) = verified_account.chatgpt_user_id.as_ref() - && !upstream_identities - .insert((user_id.clone(), verified_account.chatgpt_account_id.clone())) - { + if !upstream_identities.insert(( + verified_account.chatgpt_user_id.clone(), + verified_account.chatgpt_account_id.clone(), + )) { return Err(CodexCredentialAdminError::InvalidInput); } let account_id = candidate @@ -644,6 +692,13 @@ impl CodexCredentialAdminService { return Err(CodexCredentialAdminError::InvalidInput); } let (enabled, availability) = import_runtime_state(candidate.status.as_deref())?; + let next_refresh_at = refresh_time( + policy, + &ProviderAccountId::new(account_id.clone()) + .map_err(|_| CodexCredentialAdminError::InvalidInput)?, + verified_account.access_token_expires_at, + secret.refresh_token.is_some(), + )?; let mut prepared = CodexCredentialAdmin.prepare_import(ImportCodexOAuthCredential { account_id, provider_instance_id: provider_instance_id.to_string(), @@ -653,6 +708,7 @@ impl CodexCredentialAdminService { .unwrap_or_else(|| "Codex OAuth".to_owned()), secret, verified_account, + next_refresh_at, enabled, })?; prepared.account = prepared @@ -682,14 +738,22 @@ impl CodexCredentialAdminService { .as_deref() .map(normalize_bearer) .filter(|token| !token.is_empty()); + let expectation = CodexIdentityExpectation::imported( + candidate.chatgpt_account_id.clone(), + candidate.chatgpt_user_id.clone(), + ) + .map_err(|_| CodexCredentialAdminError::InvalidInput)?; if let Some(access_token) = access_token { let secret = CodexOAuthSecret { access_token: SecretString::from(access_token), refresh_token: refresh_token.clone().map(SecretString::from), id_token: id_token.clone().map(SecretString::from), }; - match self.verifier.verify(&secret).await { - Ok(profile) => return Ok((secret, profile)), + match self.verifier.verify(&secret, &expectation).await { + Ok(verification) => { + let profile = verification.into_complete().map_err(map_identity_error)?; + return Ok((secret, profile)); + } Err(CodexIdentityVerificationError::Unavailable) => { return Err(CodexCredentialAdminError::IdentityUnavailable); } @@ -705,7 +769,7 @@ impl CodexCredentialAdminService { .refresh(&refresh_token) .await .map_err(map_refresh_failure)?; - if tokens.access_token.trim().is_empty() || tokens.expires_in <= self.refresh_margin { + if tokens.access_token.trim().is_empty() || tokens.expires_in.is_zero() { return Err(CodexCredentialAdminError::InvalidCredential); } // 导入时已经消费来源 RT；只有上游明确轮换返回的新 RT 才能继续保存。 @@ -716,13 +780,31 @@ impl CodexCredentialAdminService { }; let profile = self .verifier - .verify(&secret) + .verify(&secret, &expectation) .await + .and_then(super::identity::CodexIdentityVerification::into_complete) .map_err(map_identity_error)?; Ok((secret, profile)) } } +fn current_identity_expectation( + account: &ProviderAccount, + credential: &super::security::CodexRuntimeCredential, +) -> Result { + let account_id = account + .upstream_account_id() + .ok_or(CodexCredentialAdminError::InvalidCredential)?; + CodexIdentityExpectation::current( + credential.principal.oauth_subject.clone(), + credential.principal.poid.clone(), + account_id.to_owned(), + account.upstream_user_id().to_owned(), + credential.installation_id.clone(), + ) + .map_err(|_| CodexCredentialAdminError::InvalidCredential) +} + fn required_time( value: Option>, ) -> Result { @@ -735,6 +817,23 @@ fn optional_time(value: Option>) -> Option>, + has_refresh_token: bool, +) -> Result>, CodexCredentialAdminError> { + if !has_refresh_token { + return Ok(None); + } + let expires_at = required_time(access_token_expires_at)?; + policy + .next_attempt_at(account_id, expires_at, SystemTime::now()) + .map(DateTime::::from) + .map(Some) + .map_err(|_| CodexCredentialAdminError::InvalidCredential) +} + fn cpr_status(account: &ProviderAccount) -> &'static str { if !account.enabled() { return "disabled"; @@ -862,6 +961,8 @@ fn parse_native_account( Ok(ParsedCodexImportAccount { id: first_string(value, &["id"]), name: first_string(value, &["label", "email"]), + chatgpt_account_id: first_string(value, &["accountId", "account_id"]), + chatgpt_user_id: first_string(value, &["userId", "user_id"]), access_token, refresh_token, id_token: None, @@ -915,6 +1016,29 @@ fn parse_credential_bundle_account( return Ok(Some(ParsedCodexImportAccount { id: first_string(value, &["id"]), name: first_string(value, &["label", "name"]), + chatgpt_account_id: first_path_string( + credentials, + &[ + &["chatgpt_account_id"], + &["chatgptAccountId"], + &["account_id"], + &["accountId"], + &["account", "id"], + &["account", "account_id"], + &["account", "chatgpt_account_id"], + ], + ), + chatgpt_user_id: first_path_string( + credentials, + &[ + &["chatgpt_user_id"], + &["chatgptUserId"], + &["user_id"], + &["userId"], + &["user", "id"], + &["account", "user_id"], + ], + ), access_token, refresh_token, id_token: first_path_string(credentials, &[&["id_token"], &["idToken"]]), @@ -950,6 +1074,29 @@ fn parse_credential_bundle_account( Ok(Some(ParsedCodexImportAccount { id: first_string(value, &["id"]), name: first_path_string(value, &[&["label"], &["name"], &["user", "name"]]), + chatgpt_account_id: first_path_string( + value, + &[ + &["chatgpt_account_id"], + &["chatgptAccountId"], + &["account_id"], + &["accountId"], + &["account", "id"], + &["account", "account_id"], + &["account", "chatgpt_account_id"], + ], + ), + chatgpt_user_id: first_path_string( + value, + &[ + &["chatgpt_user_id"], + &["chatgptUserId"], + &["user_id"], + &["userId"], + &["user", "id"], + &["account", "user_id"], + ], + ), access_token, refresh_token, id_token: first_path_string( @@ -998,6 +1145,19 @@ fn parse_auth_document_account( Ok(Some(ParsedCodexImportAccount { id: None, name: first_string(value, &["label", "name", "email"]), + chatgpt_account_id: first_string( + value, + &[ + "chatgpt_account_id", + "chatgptAccountId", + "account_id", + "accountId", + ], + ), + chatgpt_user_id: first_string( + value, + &["chatgpt_user_id", "chatgptUserId", "user_id", "userId"], + ), access_token, refresh_token, id_token: first_string(value, &["id_token", "idToken"]), diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index d5add69db..ec93514b4 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -1,32 +1,39 @@ -//! Codex model entitlement 的实时查询与进程内 TTL cache；不落 PostgreSQL。 +//! Codex model entitlement 的实时查询与可重建进程内 cache；不落 PostgreSQL。 use std::collections::{BTreeMap, btree_map::Entry}; use std::fmt; -use std::sync::{Arc, RwLock}; -use std::time::{Duration, Instant, SystemTime}; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::SystemTime; -use gateway_core::engine::credential::{AccountAvailability, ProviderAccount, ProviderAccountId}; +use gateway_core::engine::credential::{ + AccountAvailability, CredentialRevision, ProviderAccount, ProviderAccountId, +}; +use gateway_core::engine::provider::ProviderCatalogGeneration; use gateway_core::routing::{ProviderInstance, ProviderInstanceId}; use secrecy::ExposeSecret; use thiserror::Error; +use tokio::sync::Notify; use uuid::Uuid; use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; -use crate::provider::{CodexEndpointPolicy, CodexProviderInstanceConfig}; +use crate::provider::{CodexOriginPolicy, CodexProviderInstanceConfig}; use crate::transport::profile::CodexWireProfileState; -use crate::transport::{ - CodexBackendClient, CodexCatalogModel, CodexRequestContext, build_reqwest_client, -}; +use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext}; -const DEFAULT_CATALOG_TTL: Duration = Duration::from_secs(5 * 60); +const MAX_RESPONSE_ETAG_BYTES: usize = 256; #[derive(Clone)] pub struct CodexCredentialCatalogSnapshot { provider_instance_id: ProviderInstanceId, observed_at: SystemTime, - expires_at: Instant, models: Vec, - account_models: BTreeMap>, + account_models: BTreeMap, +} + +#[derive(Clone, PartialEq, Eq)] +struct CodexAccountEntitlement { + revision: CredentialRevision, + models: Vec, } impl CodexCredentialCatalogSnapshot { @@ -46,13 +53,15 @@ impl CodexCredentialCatalogSnapshot { } #[must_use] - pub fn account_models(&self, account_id: &str) -> Option<&[String]> { - self.account_models.get(account_id).map(Vec::as_slice) - } - - #[must_use] - pub fn is_fresh(&self) -> bool { - Instant::now() < self.expires_at + pub fn account_models( + &self, + account_id: &ProviderAccountId, + revision: CredentialRevision, + ) -> Option<&[String]> { + self.account_models + .get(account_id) + .filter(|entry| entry.revision == revision) + .map(|entry| entry.models.as_slice()) } } @@ -82,6 +91,34 @@ pub enum CodexCredentialCatalogError { ConflictingModelFacts, #[error("Codex model catalog cache is unavailable")] Cache, + #[error("Codex model catalog ETag is invalid")] + InvalidEtag, + #[error("Codex model catalog changed during refresh")] + ConcurrentUpdate, +} + +#[derive(Default)] +struct CatalogEtagState { + applied: BTreeMap, + pending: BTreeMap, + inflight: BTreeMap, +} + +#[derive(Default)] +struct CatalogCacheState { + revision: u64, + generation: u64, + snapshots: BTreeMap, +} + +struct FetchedAccountModels { + models: Vec, + etag: Option, +} + +struct FetchedInstanceCatalog { + snapshot: CodexCredentialCatalogSnapshot, + etags: Vec, } impl From for CodexCredentialCatalogError { @@ -95,55 +132,49 @@ pub struct CodexCredentialCatalogService { repository: CodexCredentialRepository, profile: CodexWireProfileState, http: reqwest::Client, - endpoint_policy: CodexEndpointPolicy, - ttl: Duration, - cache: Arc>>, + origin_policy: Arc, + cache: Arc>, + etags: Arc>, + etag_notification: Arc, } impl CodexCredentialCatalogService { pub fn new( repository: CodexCredentialRepository, profile: CodexWireProfileState, - ) -> Result { - let http = build_reqwest_client().map_err(|_| CodexCredentialCatalogError::Upstream)?; - Ok(Self::new_with_endpoint_policy( - repository, - profile, - http, - CodexEndpointPolicy::Official, - DEFAULT_CATALOG_TTL, - )) - } - - #[must_use] - pub fn new_with_endpoint_policy( - repository: CodexCredentialRepository, - profile: CodexWireProfileState, http: reqwest::Client, - endpoint_policy: CodexEndpointPolicy, - ttl: Duration, + origin_policy: Arc, ) -> Self { Self { repository, profile, http, - endpoint_policy, - ttl: ttl.max(Duration::from_secs(1)), - cache: Arc::new(RwLock::new(BTreeMap::new())), + origin_policy, + cache: Arc::new(RwLock::new(CatalogCacheState::default())), + etags: Arc::new(Mutex::new(CatalogEtagState::default())), + etag_notification: Arc::new(Notify::new()), } } + #[must_use] + pub fn catalog_generation(&self) -> ProviderCatalogGeneration { + let cache = self + .cache + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + ProviderCatalogGeneration::new(cache.generation) + } + pub fn cached( &self, instance: &ProviderInstanceId, ) -> Result, CodexCredentialCatalogError> { - let guard = self + Ok(self .cache .read() - .map_err(|_| CodexCredentialCatalogError::Cache)?; - Ok(guard + .map_err(|_| CodexCredentialCatalogError::Cache)? + .snapshots .get(instance) - .filter(|value| value.is_fresh()) .cloned()) } @@ -152,22 +183,37 @@ impl CodexCredentialCatalogService { &self, instance: &ProviderInstanceId, account_id: &ProviderAccountId, + revision: CredentialRevision, ) -> Result>, CodexCredentialCatalogError> { Ok(self.cached(instance)?.and_then(|snapshot| { snapshot - .account_models(account_id.as_str()) + .account_models(account_id, revision) .map(<[String]>::to_vec) })) } - /// 只刷新指定账号的 realtime catalog，并原子替换其进程内 TTL entry。 + pub fn observed_model_support( + &self, + instance: &ProviderInstanceId, + account_id: &ProviderAccountId, + revision: CredentialRevision, + model: &str, + ) -> Result, CodexCredentialCatalogError> { + Ok(self.cached(instance)?.and_then(|snapshot| { + snapshot + .account_models(account_id, revision) + .map(|models| models.iter().any(|candidate| candidate == model)) + })) + } + + /// 只刷新指定账号的 realtime catalog，并以本地 revision 原子合并账号事实。 pub async fn synchronize_account( &self, instance: &ProviderInstance, account_id: &ProviderAccountId, ) -> Result, CodexCredentialCatalogError> { let config = - CodexProviderInstanceConfig::from_snapshot_with_policy(instance, self.endpoint_policy) + CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref()) .map_err(|_| CodexCredentialCatalogError::InvalidInstance)?; let account = self .repository @@ -184,12 +230,22 @@ impl CodexCredentialCatalogService { config.base_url().as_str(), self.profile.clone(), ); - let models = self.fetch_account_models(&client, &account).await?; - let entitlement = models + let cache_revision = self.cache_revision()?; + let fetched = self.fetch_account_models(&client, &account).await?; + let entitlement = fetched + .models .iter() .map(|model| model.request_model().as_str().to_owned()) .collect::>(); - self.replace_account_cache(config.id(), account_id, models, entitlement.clone())?; + self.replace_account_cache( + cache_revision, + config.id(), + account_id, + account.revision(), + fetched.models, + entitlement.clone(), + )?; + self.record_applied_catalog_etags(config.id(), fetched.etag)?; Ok(entitlement) } @@ -200,8 +256,17 @@ impl CodexCredentialCatalogService { if let Some(cached) = self.cached(instance.id())? { return Ok(cached); } + let cache_revision = self.cache_revision()?; + let fetched = self.fetch_instance_catalog(instance).await?; + self.commit_instance_catalog(cache_revision, fetched) + } + + async fn fetch_instance_catalog( + &self, + instance: &ProviderInstance, + ) -> Result { let config = - CodexProviderInstanceConfig::from_snapshot_with_policy(instance, self.endpoint_policy) + CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref()) .map_err(|_| CodexCredentialCatalogError::InvalidInstance)?; let accounts = self.repository.list_for_instance(config.id()).await?; let now = SystemTime::now(); @@ -220,10 +285,11 @@ impl CodexCredentialCatalogService { let mut union = BTreeMap::::new(); let mut union_order = Vec::new(); let mut account_models = BTreeMap::new(); + let mut etags = Vec::new(); for account in accounts { - let models = self.fetch_account_models(&client, &account).await?; - let mut entitlement = Vec::with_capacity(models.len()); - for model in &models { + let fetched = self.fetch_account_models(&client, &account).await?; + let mut entitlement = Vec::with_capacity(fetched.models.len()); + for model in &fetched.models { let id = model.request_model().as_str().to_owned(); entitlement.push(id.clone()); match union.entry(id) { @@ -237,31 +303,33 @@ impl CodexCredentialCatalogService { } } } - account_models.insert(account.id().to_string(), entitlement); + account_models.insert( + account.id().clone(), + CodexAccountEntitlement { + revision: account.revision(), + models: entitlement, + }, + ); + etags.extend(fetched.etag); } let observed_at = SystemTime::now(); let snapshot = CodexCredentialCatalogSnapshot { provider_instance_id: config.id().clone(), observed_at, - expires_at: Instant::now() + self.ttl, models: union_order .into_iter() .filter_map(|id| union.remove(&id)) .collect(), account_models, }; - self.cache - .write() - .map_err(|_| CodexCredentialCatalogError::Cache)? - .insert(config.id().clone(), snapshot.clone()); - Ok(snapshot) + Ok(FetchedInstanceCatalog { snapshot, etags }) } async fn fetch_account_models( &self, client: &CodexBackendClient, account: &ProviderAccount, - ) -> Result, CodexCredentialCatalogError> { + ) -> Result { let runtime = self.repository.load_runtime_credential(account).await?; let request_id = format!("catalog_{}", Uuid::now_v7().simple()); let context = CodexRequestContext { @@ -282,17 +350,22 @@ impl CodexCredentialCatalogService { client_request_id: None, turn_id: None, }; - client + let snapshot = client .fetch_models_with_context(context) .await - .map(|snapshot| snapshot.models().to_vec()) - .map_err(|_| CodexCredentialCatalogError::Upstream) + .map_err(|_| CodexCredentialCatalogError::Upstream)?; + Ok(FetchedAccountModels { + models: snapshot.models().to_vec(), + etag: snapshot.etag().map(str::to_owned), + }) } fn replace_account_cache( &self, + expected_cache_revision: u64, instance: &ProviderInstanceId, account_id: &ProviderAccountId, + revision: CredentialRevision, models: Vec, entitlement: Vec, ) -> Result<(), CodexCredentialCatalogError> { @@ -300,10 +373,13 @@ impl CodexCredentialCatalogService { .cache .write() .map_err(|_| CodexCredentialCatalogError::Cache)?; + if cache.revision != expected_cache_revision { + return Err(CodexCredentialCatalogError::ConcurrentUpdate); + } let mut union = BTreeMap::::new(); let mut union_order = Vec::new(); let mut account_models = BTreeMap::new(); - if let Some(existing) = cache.get(instance).filter(|snapshot| snapshot.is_fresh()) { + if let Some(existing) = cache.snapshots.get(instance) { for model in &existing.models { let id = model.request_model().as_str().to_owned(); union_order.push(id.clone()); @@ -324,26 +400,37 @@ impl CodexCredentialCatalogService { } } } - account_models.insert(account_id.to_string(), entitlement); + account_models.insert( + account_id.clone(), + CodexAccountEntitlement { + revision, + models: entitlement, + }, + ); union.retain(|model, _| { account_models .values() - .any(|entitlement| entitlement.iter().any(|item| item == model)) + .any(|entitlement| entitlement.models.iter().any(|item| item == model)) }); union_order.retain(|id| union.contains_key(id)); - cache.insert( - instance.clone(), - CodexCredentialCatalogSnapshot { - provider_instance_id: instance.clone(), - observed_at: SystemTime::now(), - expires_at: Instant::now() + self.ttl, - models: union_order - .into_iter() - .filter_map(|id| union.remove(&id)) - .collect(), - account_models, - }, - ); + let snapshot = CodexCredentialCatalogSnapshot { + provider_instance_id: instance.clone(), + observed_at: SystemTime::now(), + models: union_order + .into_iter() + .filter_map(|id| union.remove(&id)) + .collect(), + account_models, + }; + let changed = cache + .snapshots + .get(instance) + .is_none_or(|existing| !same_catalog(existing, &snapshot)); + cache.snapshots.insert(instance.clone(), snapshot); + cache.revision = cache.revision.saturating_add(1); + if changed { + cache.generation = cache.generation.saturating_add(1); + } Ok(()) } @@ -351,12 +438,193 @@ impl CodexCredentialCatalogService { &self, instance: &ProviderInstanceId, ) -> Result<(), CodexCredentialCatalogError> { - self.cache + let mut cache = self + .cache .write() - .map_err(|_| CodexCredentialCatalogError::Cache)? - .remove(instance); + .map_err(|_| CodexCredentialCatalogError::Cache)?; + if cache.snapshots.remove(instance).is_some() { + cache.revision = cache.revision.saturating_add(1); + cache.generation = cache.generation.saturating_add(1); + } Ok(()) } + + /// 记录普通 Responses 响应声明的目录版本；相同版本只触发一次。 + pub fn observe_response_etag( + &self, + instance: &ProviderInstanceId, + etag: &str, + ) -> Result { + validate_response_etag(etag)?; + let changed = { + let mut state = self + .etags + .lock() + .map_err(|_| CodexCredentialCatalogError::Cache)?; + let already_observed = state + .applied + .get(instance) + .is_some_and(|value| value == etag) + || state + .pending + .get(instance) + .is_some_and(|value| value == etag) + || state + .inflight + .get(instance) + .is_some_and(|value| value == etag); + if already_observed { + false + } else { + state.pending.insert(instance.clone(), etag.to_owned()); + true + } + }; + if changed { + self.etag_notification.notify_one(); + } + Ok(changed) + } + + /// 等待并合并一批需要强制刷新的 Provider instance。 + pub async fn wait_for_etag_refresh(&self) -> Vec { + loop { + let pending = self.begin_pending_etag_refreshes(); + if !pending.is_empty() { + return pending; + } + self.etag_notification.notified().await; + } + } + + /// 忽略当前 cache，按 ETag 变化强制生成一份完整新快照。 + pub async fn refresh_instance( + &self, + instance: &ProviderInstance, + ) -> Result { + let result = self.refresh_instance_inner(instance).await; + self.finish_etag_refresh(instance.id(), result.is_ok())?; + result + } + + async fn refresh_instance_inner( + &self, + instance: &ProviderInstance, + ) -> Result { + let cache_revision = self.cache_revision()?; + let fetched = self.fetch_instance_catalog(instance).await?; + self.commit_instance_catalog(cache_revision, fetched) + } + + fn record_applied_catalog_etags( + &self, + instance: &ProviderInstanceId, + etags: impl IntoIterator, + ) -> Result<(), CodexCredentialCatalogError> { + let mut distinct = etags.into_iter().collect::>(); + if distinct.len() != 1 { + return Ok(()); + } + let Some(etag) = distinct.pop_first() else { + return Ok(()); + }; + validate_response_etag(&etag)?; + self.etags + .lock() + .map_err(|_| CodexCredentialCatalogError::Cache)? + .applied + .insert(instance.clone(), etag); + Ok(()) + } + + fn begin_pending_etag_refreshes(&self) -> Vec { + let Ok(mut state) = self.etags.lock() else { + return Vec::new(); + }; + let pending = std::mem::take(&mut state.pending); + let instances = pending.keys().cloned().collect(); + state.inflight.extend(pending); + instances + } + + fn finish_etag_refresh( + &self, + instance: &ProviderInstanceId, + succeeded: bool, + ) -> Result<(), CodexCredentialCatalogError> { + let should_retry = { + let mut state = self + .etags + .lock() + .map_err(|_| CodexCredentialCatalogError::Cache)?; + let Some(etag) = state.inflight.remove(instance) else { + return Ok(()); + }; + if succeeded { + state.applied.insert(instance.clone(), etag); + false + } else { + state.pending.entry(instance.clone()).or_insert(etag); + true + } + }; + if should_retry { + self.etag_notification.notify_one(); + } + Ok(()) + } + + fn cache_revision(&self) -> Result { + self.cache + .read() + .map(|cache| cache.revision) + .map_err(|_| CodexCredentialCatalogError::Cache) + } + + fn commit_instance_catalog( + &self, + expected_cache_revision: u64, + fetched: FetchedInstanceCatalog, + ) -> Result { + let instance = fetched.snapshot.provider_instance_id.clone(); + let snapshot = fetched.snapshot; + let mut cache = self + .cache + .write() + .map_err(|_| CodexCredentialCatalogError::Cache)?; + if cache.revision != expected_cache_revision { + return Err(CodexCredentialCatalogError::ConcurrentUpdate); + } + let changed = cache + .snapshots + .get(&instance) + .is_none_or(|existing| !same_catalog(existing, &snapshot)); + cache.snapshots.insert(instance.clone(), snapshot.clone()); + cache.revision = cache.revision.saturating_add(1); + if changed { + cache.generation = cache.generation.saturating_add(1); + } + drop(cache); + self.record_applied_catalog_etags(&instance, fetched.etags)?; + Ok(snapshot) + } +} + +fn same_catalog( + left: &CodexCredentialCatalogSnapshot, + right: &CodexCredentialCatalogSnapshot, +) -> bool { + left.provider_instance_id == right.provider_instance_id + && left.models == right.models + && left.account_models == right.account_models +} + +fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> { + if etag.is_empty() || etag.len() > MAX_RESPONSE_ETAG_BYTES || etag.chars().any(char::is_control) + { + return Err(CodexCredentialCatalogError::InvalidEtag); + } + Ok(()) } fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool { @@ -365,10 +633,8 @@ fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool && match account.availability() { AccountAvailability::Unknown | AccountAvailability::Ready + | AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted => true, - AccountAvailability::Cooldown => { - account.cooldown_until().is_some_and(|until| until <= now) - } AccountAvailability::Expired | AccountAvailability::Banned | AccountAvailability::Invalid => false, diff --git a/backend/crates/providers/openai/src/credential/identity.rs b/backend/crates/providers/openai/src/credential/identity.rs index 250f01f5c..10b0ae7e7 100644 --- a/backend/crates/providers/openai/src/credential/identity.rs +++ b/backend/crates/providers/openai/src/credential/identity.rs @@ -2,21 +2,28 @@ use std::collections::HashMap; use std::fmt; +use std::sync::Arc; use std::time::{Duration, Instant}; use async_trait::async_trait; -use chrono::{DateTime, TimeDelta, Utc}; +use chrono::{DateTime, Utc}; use jsonwebtoken::{ Algorithm, DecodingKey, Validation, decode, decode_header, get_current_timestamp, }; -use reqwest::Client; use reqwest::header::{ACCEPT, CONTENT_TYPE}; +use reqwest::{Client, StatusCode}; use secrecy::{ExposeSecret, SecretString}; use serde::Deserialize; +use serde_json::Value; use tokio::sync::Mutex; +use uuid::Uuid; use super::types::{CodexAccountProfile, CodexOAuthSecret}; use crate::credential::token_client::OFFICIAL_CODEX_OAUTH_CLIENT_ID; +use crate::transport::profile::CodexWireProfileState; +use crate::transport::{ + CodexBackendClient, CodexClientError, CodexRequestContext, build_reqwest_client, +}; pub const OFFICIAL_OPENAI_ISSUER: &str = "https://auth.openai.com"; pub const OFFICIAL_OPENAI_API_AUDIENCE: &str = "https://api.openai.com/v1"; @@ -35,7 +42,7 @@ const MAX_TOKEN_LIFETIME_SECONDS: u64 = 366 * 24 * 60 * 60; const JWKS_CACHE_TTL: Duration = Duration::from_secs(60 * 60); const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); -const REFRESH_MARGIN: TimeDelta = TimeDelta::minutes(5); +const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api"; /// OAuth identity verification 的稳定失败分类。 #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] @@ -52,27 +59,211 @@ pub trait CodexJwksSource: Send + Sync { async fn fetch(&self) -> Result, CodexIdentityVerificationError>; } -/// 验证 Codex OAuth token 并返回由签名 claims 派生的账号事实。 -#[async_trait] -pub trait CodexTokenIdentityVerifier: Send + Sync { - async fn verify( - &self, - secret: &CodexOAuthSecret, - ) -> Result; +/// 已验签但尚未由账号接口补全的 OAuth 身份。 +#[derive(Clone)] +pub struct CodexSignedIdentity { + oauth_subject: String, + poid: Option, + claimed_account_id: Option, + claimed_user_id: Option, + email: Option, + plan_type: Option, + access_token_expires_at: DateTime, } -/// Authorization Code flow 的 OIDC token-set 校验端口。 -/// -/// 除 access token 的官方身份边界外，还必须验证 ID token 的签名、client audience、 -/// subject 与 server-side nonce；调用方不能用浏览器输入替代这些事实。 +impl fmt::Debug for CodexSignedIdentity { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CodexSignedIdentity") + .field("oauth_subject", &"") + .field("poid", &self.poid.as_ref().map(|_| "")) + .field( + "claimed_account_id", + &self.claimed_account_id.as_ref().map(|_| ""), + ) + .field( + "claimed_user_id", + &self.claimed_user_id.as_ref().map(|_| ""), + ) + .field("email", &self.email.as_ref().map(|_| "")) + .field("plan_type", &self.plan_type) + .field("access_token_expires_at", &self.access_token_expires_at) + .finish() + } +} + +impl CodexSignedIdentity { + #[must_use] + pub fn oauth_subject(&self) -> &str { + &self.oauth_subject + } + + #[must_use] + pub fn poid(&self) -> Option<&str> { + self.poid.as_deref() + } + + #[must_use] + pub fn claimed_account_id(&self) -> Option<&str> { + self.claimed_account_id.as_deref() + } + + #[must_use] + pub fn claimed_user_id(&self) -> Option<&str> { + self.claimed_user_id.as_deref() + } + + #[must_use] + pub const fn access_token_expires_at(&self) -> DateTime { + self.access_token_expires_at + } +} + +/// 导入文档或现有 credential 对认证结果施加的一致性约束。 +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CodexIdentityExpectation { + oauth_subject: Option, + poid: PoidExpectation, + chatgpt_account_id: Option, + chatgpt_user_id: Option, + installation_id: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +enum PoidExpectation { + #[default] + Unspecified, + Exact(Option), +} + +impl CodexIdentityExpectation { + pub fn imported( + chatgpt_account_id: Option, + chatgpt_user_id: Option, + ) -> Result { + if chatgpt_account_id + .as_deref() + .is_some_and(|value| !valid_identity(value)) + || chatgpt_user_id + .as_deref() + .is_some_and(|value| !valid_identity(value)) + { + return Err(CodexIdentityVerificationError::Rejected); + } + Ok(Self { + chatgpt_account_id, + chatgpt_user_id, + ..Self::default() + }) + } + + pub fn current( + oauth_subject: String, + poid: Option, + chatgpt_account_id: String, + chatgpt_user_id: String, + installation_id: String, + ) -> Result { + if !valid_identity(&oauth_subject) + || poid.as_deref().is_some_and(|value| !valid_identity(value)) + || !valid_identity(&chatgpt_account_id) + || !valid_identity(&chatgpt_user_id) + || !valid_installation_id(&installation_id) + { + return Err(CodexIdentityVerificationError::Rejected); + } + Ok(Self { + oauth_subject: Some(oauth_subject), + poid: PoidExpectation::Exact(poid), + chatgpt_account_id: Some(chatgpt_account_id), + chatgpt_user_id: Some(chatgpt_user_id), + installation_id: Some(installation_id), + }) + } + + #[must_use] + pub fn chatgpt_account_id(&self) -> Option<&str> { + self.chatgpt_account_id.as_deref() + } + + #[must_use] + pub fn chatgpt_user_id(&self) -> Option<&str> { + self.chatgpt_user_id.as_deref() + } + + #[must_use] + pub fn installation_id(&self) -> Option<&str> { + self.installation_id.as_deref() + } +} + +/// `/wham/usage` 认证后返回的账号身份事实。 +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CodexAuthenticatedAccount { + pub chatgpt_account_id: String, + pub chatgpt_user_id: String, + pub email: Option, + pub plan_type: Option, +} + +/// 完整认证成功，或仅签名已确认但账号接口暂时不可用。 +#[derive(Clone, Debug)] +pub enum CodexIdentityVerification { + Complete(CodexAccountProfile), + SignedOnly(CodexSignedIdentity), +} + +impl CodexIdentityVerification { + pub fn into_complete(self) -> Result { + match self { + Self::Complete(profile) => Ok(profile), + Self::SignedOnly(_) => Err(CodexIdentityVerificationError::Unavailable), + } + } +} + +/// JWT/OIDC 签名验证边界；结果不能直接持久化为 Ready 账号。 #[async_trait] -pub trait CodexAuthorizationTokenVerifier: Send + Sync { +pub trait CodexSignedIdentityVerifier: Send + Sync { + async fn verify_access( + &self, + secret: &CodexOAuthSecret, + ) -> Result; + async fn verify_authorization( &self, secret: &CodexOAuthSecret, id_token: &SecretString, expected_nonce: &SecretString, - ) -> Result; + ) -> Result; +} + +/// 用 access token 调用认证账号接口的窄端口。 +#[async_trait] +pub trait CodexAuthenticatedAccountSource: Send + Sync { + async fn fetch( + &self, + secret: &CodexOAuthSecret, + expectation: &CodexIdentityExpectation, + ) -> Result; +} + +/// OAuth、导入和刷新共享的完整身份认证边界。 +#[async_trait] +pub trait CodexAccountIdentityVerifier: Send + Sync { + async fn verify( + &self, + secret: &CodexOAuthSecret, + expectation: &CodexIdentityExpectation, + ) -> Result; + + async fn verify_authorization( + &self, + secret: &CodexOAuthSecret, + id_token: &SecretString, + expected_nonce: &SecretString, + expectation: &CodexIdentityExpectation, + ) -> Result; } /// 固定访问官方 OpenAI JWKS endpoint 的生产 source。 @@ -251,31 +442,150 @@ impl fmt::Debug for CodexJwtIdentityVerifier { } #[async_trait] -impl CodexTokenIdentityVerifier for CodexJwtIdentityVerifier { - async fn verify( +impl CodexSignedIdentityVerifier for CodexJwtIdentityVerifier { + async fn verify_access( &self, secret: &CodexOAuthSecret, - ) -> Result { + ) -> Result { self.decode_access_token(secret) .await? - .into_profile(secret.refresh_token.is_some()) + .into_signed_identity() } -} -#[async_trait] -impl CodexAuthorizationTokenVerifier for CodexJwtIdentityVerifier { async fn verify_authorization( &self, secret: &CodexOAuthSecret, id_token: &SecretString, expected_nonce: &SecretString, - ) -> Result { + ) -> Result { let access = self.decode_access_token(secret).await?; let id = self.decode_id_token(id_token).await?; if access.sub != id.sub || !id.valid_for_nonce(expected_nonce.expose_secret()) { return Err(CodexIdentityVerificationError::Rejected); } - access.into_profile(secret.refresh_token.is_some()) + access.into_signed_identity() + } +} + +/// 官方 `/wham/usage` 账号身份 source。 +pub struct ReqwestCodexAuthenticatedAccountSource { + client: CodexBackendClient, +} + +impl ReqwestCodexAuthenticatedAccountSource { + pub fn new(profile: CodexWireProfileState) -> Result { + let http = + build_reqwest_client().map_err(|_| CodexIdentityVerificationError::Unavailable)?; + Ok(Self { + client: CodexBackendClient::new(http, OFFICIAL_CODEX_BASE_URL, profile), + }) + } +} + +impl fmt::Debug for ReqwestCodexAuthenticatedAccountSource { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ReqwestCodexAuthenticatedAccountSource") + .field("endpoint", &OFFICIAL_CODEX_BASE_URL) + .finish_non_exhaustive() + } +} + +#[async_trait] +impl CodexAuthenticatedAccountSource for ReqwestCodexAuthenticatedAccountSource { + async fn fetch( + &self, + secret: &CodexOAuthSecret, + expectation: &CodexIdentityExpectation, + ) -> Result { + let request_id = format!("identity_{}", Uuid::now_v7().simple()); + let value = self + .client + .fetch_usage(CodexRequestContext { + access_token: secret.access_token.expose_secret(), + account_id: expectation.chatgpt_account_id(), + request_id: &request_id, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + installation_id: expectation.installation_id(), + session_id: None, + thread_id: None, + client_request_id: None, + turn_id: None, + }) + .await + .map_err(map_account_source_error)?; + authenticated_account(&value) + } +} + +/// 签名验证与认证账号接口的唯一组合服务。 +pub struct CodexAccountIdentityService { + signed: Arc, + accounts: Arc, +} + +impl CodexAccountIdentityService { + #[must_use] + pub fn new( + signed: Arc, + accounts: Arc, + ) -> Self { + Self { signed, accounts } + } + + async fn complete( + &self, + secret: &CodexOAuthSecret, + signed: CodexSignedIdentity, + expectation: &CodexIdentityExpectation, + ) -> Result { + match self.accounts.fetch(secret, expectation).await { + Ok(account) => complete_profile(signed, account, expectation) + .map(CodexIdentityVerification::Complete), + Err(CodexIdentityVerificationError::Unavailable) => { + Ok(CodexIdentityVerification::SignedOnly(signed)) + } + Err(error) => Err(error), + } + } +} + +impl fmt::Debug for CodexAccountIdentityService { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("CodexAccountIdentityService([VERIFIERS])") + } +} + +#[async_trait] +impl CodexAccountIdentityVerifier for CodexAccountIdentityService { + async fn verify( + &self, + secret: &CodexOAuthSecret, + expectation: &CodexIdentityExpectation, + ) -> Result { + let signed = self.signed.verify_access(secret).await?; + self.complete(secret, signed, expectation).await + } + + async fn verify_authorization( + &self, + secret: &CodexOAuthSecret, + id_token: &SecretString, + expected_nonce: &SecretString, + expectation: &CodexIdentityExpectation, + ) -> Result { + let signed = self + .signed + .verify_authorization(secret, id_token, expected_nonce) + .await?; + self.complete(secret, signed, expectation).await } } @@ -366,10 +676,7 @@ impl IdTokenClaims { } impl AccessTokenClaims { - fn into_profile( - self, - refreshable: bool, - ) -> Result { + fn into_signed_identity(self) -> Result { let now = get_current_timestamp(); if self.iss != OFFICIAL_OPENAI_ISSUER || !self.aud.is_exact(OFFICIAL_OPENAI_API_AUDIENCE) @@ -381,7 +688,16 @@ impl AccessTokenClaims { .nbf .is_some_and(|not_before| not_before > now || not_before >= self.exp) || !valid_identity(&self.sub) - || !valid_identity(&self.auth.chatgpt_account_id) + || self + .auth + .chatgpt_account_id + .as_deref() + .is_some_and(|value| !valid_identity(value)) + || self + .auth + .poid + .as_deref() + .is_some_and(|value| !valid_identity(value)) { return Err(CodexIdentityVerificationError::Rejected); } @@ -416,24 +732,24 @@ impl AccessTokenClaims { 0, ) .ok_or(CodexIdentityVerificationError::Rejected)?; - let next_refresh_at = refreshable.then(|| { - let planned = expires_at - REFRESH_MARGIN; - planned.max(Utc::now()) - }); - Ok(CodexAccountProfile { + Ok(CodexSignedIdentity { email: self.profile.and_then(|profile| profile.email), - chatgpt_account_id: self.auth.chatgpt_account_id, - chatgpt_user_id: self.auth.chatgpt_user_id.or(self.auth.user_id), + oauth_subject: self.sub, + poid: self.auth.poid, + claimed_account_id: self.auth.chatgpt_account_id, + claimed_user_id: self.auth.chatgpt_user_id.or(self.auth.user_id), plan_type: self.auth.chatgpt_plan_type, - access_token_expires_at: Some(expires_at), - next_refresh_at, + access_token_expires_at: expires_at, }) } } #[derive(Deserialize)] struct OpenAiAuthClaims { - chatgpt_account_id: String, + #[serde(default)] + chatgpt_account_id: Option, + #[serde(default)] + poid: Option, #[serde(default)] chatgpt_user_id: Option, #[serde(default)] @@ -464,6 +780,100 @@ impl AudienceClaim { } } +fn complete_profile( + signed: CodexSignedIdentity, + account: CodexAuthenticatedAccount, + expectation: &CodexIdentityExpectation, +) -> Result { + if expectation + .oauth_subject + .as_deref() + .is_some_and(|expected| expected != signed.oauth_subject) + || matches!(&expectation.poid, PoidExpectation::Exact(expected) if expected != &signed.poid) + || expectation + .chatgpt_account_id + .as_deref() + .is_some_and(|expected| expected != account.chatgpt_account_id) + || expectation + .chatgpt_user_id + .as_deref() + .is_some_and(|expected| expected != account.chatgpt_user_id) + || signed + .claimed_account_id + .as_deref() + .is_some_and(|claimed| claimed != account.chatgpt_account_id) + || signed + .claimed_user_id + .as_deref() + .is_some_and(|claimed| claimed != account.chatgpt_user_id) + { + return Err(CodexIdentityVerificationError::Rejected); + } + Ok(CodexAccountProfile { + email: account.email.or(signed.email), + oauth_subject: signed.oauth_subject, + poid: signed.poid, + chatgpt_account_id: account.chatgpt_account_id, + chatgpt_user_id: account.chatgpt_user_id, + plan_type: account.plan_type.or(signed.plan_type), + access_token_expires_at: Some(signed.access_token_expires_at), + }) +} + +fn authenticated_account( + value: &Value, +) -> Result { + let chatgpt_account_id = required_account_field(value, "account_id")?; + let chatgpt_user_id = required_account_field(value, "user_id")?; + let email = optional_account_field(value, "email", MAX_EMAIL_BYTES)?; + let plan_type = optional_account_field(value, "plan_type", MAX_PLAN_BYTES)?; + Ok(CodexAuthenticatedAccount { + chatgpt_account_id, + chatgpt_user_id, + email, + plan_type, + }) +} + +fn required_account_field( + value: &Value, + key: &str, +) -> Result { + let value = value + .get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| valid_identity(value)) + .ok_or(CodexIdentityVerificationError::Rejected)?; + Ok(value.to_owned()) +} + +fn optional_account_field( + value: &Value, + key: &str, + max_bytes: usize, +) -> Result, CodexIdentityVerificationError> { + let Some(value) = value.get(key).filter(|value| !value.is_null()) else { + return Ok(None); + }; + let value = value + .as_str() + .map(str::trim) + .filter(|value| valid_text(value, max_bytes)) + .ok_or(CodexIdentityVerificationError::Rejected)?; + Ok(Some(value.to_owned())) +} + +fn map_account_source_error(error: CodexClientError) -> CodexIdentityVerificationError { + match error { + CodexClientError::Upstream { + status: StatusCode::BAD_REQUEST | StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN, + .. + } => CodexIdentityVerificationError::Rejected, + _ => CodexIdentityVerificationError::Unavailable, + } +} + fn parse_jwks( body: &[u8], ) -> Result, CodexIdentityVerificationError> { @@ -571,6 +981,12 @@ fn valid_identity(value: &str) -> bool { valid_text(value, MAX_IDENTITY_BYTES) } +fn valid_installation_id(value: &str) -> bool { + Uuid::parse_str(value) + .ok() + .is_some_and(|uuid| uuid.get_version_num() == 4) +} + fn valid_text(value: &str, max_bytes: usize) -> bool { !value.is_empty() && value.len() <= max_bytes diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index f87e02550..3db07717f 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -13,6 +13,9 @@ mod selector; pub mod token_client; mod types; +pub(crate) use admin::refresh_time; +pub(crate) use oauth::oauth_owner_ref; + pub use admin::{ CodexCprExportDocument, CodexCredentialAdmin, CodexCredentialAdminError, CodexCredentialAdminService, ExportManagedCodexCredential, ImportCodexOAuthCredential, @@ -24,15 +27,17 @@ pub use catalog::{ }; pub use cookie::{CodexCookiePolicy, CookiePolicyError}; pub use identity::{ - CodexAuthorizationTokenVerifier, CodexIdentityVerificationError, CodexJwksSource, - CodexJwtIdentityVerifier, CodexTokenIdentityVerifier, OFFICIAL_OPENAI_API_AUDIENCE, - OFFICIAL_OPENAI_ISSUER, OFFICIAL_OPENAI_JWKS_URI, ReqwestOpenAiJwksSource, + CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexAuthenticatedAccount, + CodexAuthenticatedAccountSource, CodexIdentityExpectation, CodexIdentityVerification, + CodexIdentityVerificationError, CodexJwksSource, CodexJwtIdentityVerifier, CodexSignedIdentity, + CodexSignedIdentityVerifier, OFFICIAL_OPENAI_API_AUDIENCE, OFFICIAL_OPENAI_ISSUER, + OFFICIAL_OPENAI_JWKS_URI, ReqwestCodexAuthenticatedAccountSource, ReqwestOpenAiJwksSource, }; pub use oauth::{ CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthAuthorizationStarted, - CodexOAuthFlowBinding, CodexOAuthPendingStore, CodexOAuthPendingStoreError, - CodexOAuthReauthorizationTarget, CodexPendingAuthorization, CompleteCodexOAuthAuthorization, - StartCodexOAuthAuthorization, StartCodexOAuthReauthorization, StoredCodexPendingAuthorization, + CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthReauthorizationTarget, + CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthAuthorization, + CompletedCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization, }; pub use quota::{ CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService, @@ -41,18 +46,16 @@ pub use quota::{ }; pub use refresh::{ CodexCredentialRefreshError, CodexCredentialRefreshOutcome, CodexCredentialRefreshService, - CodexRefreshLeaseAcquisition, CodexRefreshLeaseCoordinator, CodexRefreshLeaseError, - CodexRefreshLeaseGuard, CodexRefreshLeaseRequest, DueCodexCredential, + DueCodexCredential, }; pub use repository::{CodexCredentialRepository, CredentialRepositoryError}; pub use security::{CodexCredentialCodec, CodexCredentialDataError, CodexRuntimeCredential}; pub use selector::{ - CodexCredentialLease, CodexCredentialSelector, CredentialLeaseCoordinator, - CredentialLeaseCoordinatorError, CredentialLeaseGuard, CredentialLeaseRequest, - CredentialSelectionError, LeaseAcquisition, SelectCodexCredential, + CodexAccountFailure, CodexCredentialLease, CodexCredentialSelector, CredentialSelectionError, + SelectCodexCredential, }; pub use types::{ CodexAccountProfile, CodexCookie, CodexCookieCaptureOutcome, CodexCredentialData, - CodexOAuthSecret, CreateCodexCredential, CredentialRecord, RotateCodexCredential, - RuntimeCodexCookie, UpsertCodexCookie, + CodexCredentialPrincipal, CodexOAuthSecret, CreateCodexCredential, CredentialRecord, + RotateCodexCredential, RuntimeCodexCookie, UpsertCodexCookie, }; diff --git a/backend/crates/providers/openai/src/credential/oauth.rs b/backend/crates/providers/openai/src/credential/oauth.rs index d145cfb62..ebee04837 100644 --- a/backend/crates/providers/openai/src/credential/oauth.rs +++ b/backend/crates/providers/openai/src/credential/oauth.rs @@ -6,9 +6,13 @@ use std::sync::Arc; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use chrono::{DateTime, TimeDelta, Utc}; +use gateway_admin::model::provider_credentials::{ + AuthorizationMutationTarget, AuthorizationOwner, PendingAuthorizationMutation, +}; use gateway_core::engine::credential::{ CredentialRevision, NewProviderAccount, ProviderAccountId, ProviderAccountStore, }; +use gateway_core::provider_ports::ProviderRuntimePolicyPort; use secrecy::{ExposeSecret, SecretString}; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq as _; @@ -17,9 +21,12 @@ use uuid::Uuid; use super::admin::{ CodexCredentialAdmin, CodexCredentialAdminError, ImportCodexOAuthCredential, - PreparedCodexCredentialRotation, RotateManagedCodexCredential, + PreparedCodexCredentialRotation, RotateManagedCodexCredential, refresh_time, }; -use super::identity::{CodexAuthorizationTokenVerifier, CodexIdentityVerificationError}; +use super::identity::{ + CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError, +}; +use super::security::CodexCredentialCodec; use super::token_client::{ AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant, OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI, @@ -31,50 +38,9 @@ const AUTHORIZATION_TTL: TimeDelta = TimeDelta::minutes(10); const MAX_CALLBACK_BYTES: usize = 16 * 1024; const MAX_TEXT_BYTES: usize = 512; -#[derive(Clone)] -pub struct CodexOAuthFlowBinding { - owner_ref: String, - started_request_ref: String, -} - -impl CodexOAuthFlowBinding { - pub fn new( - owner_ref: impl Into, - started_request_ref: impl Into, - ) -> Result { - let binding = Self { - owner_ref: owner_ref.into(), - started_request_ref: started_request_ref.into(), - }; - if !valid_text(&binding.owner_ref) || !valid_text(&binding.started_request_ref) { - return Err(CodexOAuthAdminError::InvalidInput); - } - Ok(binding) - } -} - -impl fmt::Debug for CodexOAuthFlowBinding { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("CodexOAuthFlowBinding") - .field("owner_ref", &"") - .field("started_request_ref", &"") - .finish() - } -} - #[derive(Clone, Debug)] pub struct StartCodexOAuthAuthorization { - pub binding: CodexOAuthFlowBinding, - pub provider_instance_id: String, - pub name: String, -} - -#[derive(Clone, Debug)] -pub struct StartCodexOAuthReauthorization { - pub binding: CodexOAuthFlowBinding, - pub account_id: ProviderAccountId, - pub expected_credential_revision: CredentialRevision, + pub mutation: PendingAuthorizationMutation, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -119,6 +85,37 @@ pub struct CompleteCodexOAuthAuthorization { pub callback_url: SecretString, } +/// OAuth exchange 后返回的 Provider prepared credential 及其原始事务信封。 +pub struct CompletedCodexOAuthAuthorization { + pub mutation: PendingAuthorizationMutation, + pub credential: T, +} + +/// OAuth exchange 后唯一的 credential preparation 结果。 +pub enum CompletedCodexOAuthCredential { + Create(NewProviderAccount), + Reauthorize(PreparedCodexCredentialRotation), +} + +impl fmt::Debug for CompletedCodexOAuthCredential { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Create(_) => formatter.write_str("Create([PREPARED])"), + Self::Reauthorize(_) => formatter.write_str("Reauthorize([PREPARED])"), + } + } +} + +impl fmt::Debug for CompletedCodexOAuthAuthorization { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CompletedCodexOAuthAuthorization") + .field("mutation", &self.mutation) + .field("credential", &"[PREPARED]") + .finish() + } +} + impl fmt::Debug for CompleteCodexOAuthAuthorization { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -141,6 +138,7 @@ pub struct CodexPendingAuthorization { nonce: SecretString, code_verifier: SecretString, reauthorization: Option, + mutation: PendingAuthorizationMutation, } pub struct StoredCodexPendingAuthorization { @@ -155,6 +153,7 @@ pub struct StoredCodexPendingAuthorization { pub code_verifier: SecretString, pub reauthorization_account_id: Option, pub reauthorization_credential_revision: Option, + pub mutation: PendingAuthorizationMutation, } impl CodexPendingAuthorization { @@ -185,6 +184,7 @@ impl CodexPendingAuthorization { nonce: input.nonce, code_verifier: input.code_verifier, reauthorization, + mutation: input.mutation, }; if !valid_text(&pending.flow_id) || !valid_text(&pending.owner_ref) @@ -195,6 +195,7 @@ impl CodexPendingAuthorization { || !valid_secret(pending.state.expose_secret()) || !valid_secret(pending.nonce.expose_secret()) || !valid_secret(pending.code_verifier.expose_secret()) + || !pending_mutation_matches(&pending) { return Err(CodexOAuthPendingStoreError::InvalidValue); } @@ -250,6 +251,11 @@ impl CodexPendingAuthorization { pub const fn reauthorization(&self) -> Option<&CodexOAuthReauthorizationTarget> { self.reauthorization.as_ref() } + + #[must_use] + pub const fn mutation(&self) -> &PendingAuthorizationMutation { + &self.mutation + } } impl fmt::Debug for CodexPendingAuthorization { @@ -266,6 +272,7 @@ impl fmt::Debug for CodexPendingAuthorization { .field("nonce", &"") .field("code_verifier", &"") .field("reauthorization", &self.reauthorization) + .field("mutation", &self.mutation) .finish() } } @@ -326,24 +333,15 @@ pub trait CodexOAuthAdmin: Send + Sync { async fn complete_authorization( &self, command: CompleteCodexOAuthAuthorization, - ) -> Result; - - async fn start_reauthorization( - &self, - command: StartCodexOAuthReauthorization, - ) -> Result; - - async fn complete_reauthorization( - &self, - command: CompleteCodexOAuthAuthorization, - ) -> Result; + ) -> Result, CodexOAuthAdminError>; } pub struct CodexOAuthAdminService { pending: Arc, exchanger: Arc, - verifier: Arc, + verifier: Arc, store: Arc, + runtime_policy: Arc, credentials: CodexCredentialAdmin, } @@ -352,8 +350,9 @@ impl CodexOAuthAdminService { pub const fn new( pending: Arc, exchanger: Arc, - verifier: Arc, + verifier: Arc, store: Arc, + runtime_policy: Arc, credentials: CodexCredentialAdmin, ) -> Self { Self { @@ -361,6 +360,7 @@ impl CodexOAuthAdminService { exchanger, verifier, store, + runtime_policy, credentials, } } @@ -372,11 +372,44 @@ impl CodexOAuthAdmin for CodexOAuthAdminService { &self, command: StartCodexOAuthAuthorization, ) -> Result { + let (provider_instance_id, name, reauthorization) = match command.mutation.target() { + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + } => (provider_instance_id.to_string(), name.clone(), None), + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => { + let expected_revision = CredentialRevision::new(expected_credential_revision.get()) + .map_err(|_| CodexOAuthAdminError::InvalidInput)?; + let current = self + .store + .load_credential(account_id, expected_revision) + .await + .map_err(map_store_error)?; + if current.account.provider().as_str() != "openai" + || current.account.id() != account_id + || current.account.instance() != provider_instance_id + { + return Err(CodexOAuthAdminError::NotFound); + } + ( + current.account.instance().to_string(), + current.account.name().to_owned(), + Some(CodexOAuthReauthorizationTarget { + account_id: account_id.clone(), + credential_revision: expected_revision, + }), + ) + } + }; self.start_pending( - command.binding, - command.provider_instance_id, - command.name, - None, + command.mutation, + provider_instance_id, + name, + reauthorization, ) .await } @@ -384,83 +417,97 @@ impl CodexOAuthAdmin for CodexOAuthAdminService { async fn complete_authorization( &self, command: CompleteCodexOAuthAuthorization, - ) -> Result { + ) -> Result, CodexOAuthAdminError> + { if !valid_text(&command.owner_ref) || !valid_text(&command.flow_id) || command.callback_url.expose_secret().len() > MAX_CALLBACK_BYTES { return Err(CodexOAuthAdminError::InvalidInput); } - let pending = self.exchange_pending(command).await?; - if pending.0.reauthorization().is_some() { - return Err(CodexOAuthAdminError::Conflict); - } - let (pending, secret, profile) = pending; - self.credentials - .prepare_import(ImportCodexOAuthCredential { - account_id: format!("acct_{}", Uuid::now_v7().simple()), - provider_instance_id: pending.provider_instance_id, - name: pending.name, - secret, - verified_account: profile, - enabled: true, - }) - .map_err(map_admin_error) - } - - async fn start_reauthorization( - &self, - command: StartCodexOAuthReauthorization, - ) -> Result { - let current = self - .store - .load_credential(&command.account_id, command.expected_credential_revision) + let (pending, mut secret, id_token) = self.exchange_pending(command).await?; + let mutation = pending.mutation.clone(); + let current = if let Some(target) = pending.reauthorization() { + let current = self + .store + .load_credential(target.account_id(), target.credential_revision()) + .await + .map_err(map_store_error)?; + Some(current) + } else { + None + }; + let expectation = current + .as_ref() + .map(oauth_identity_expectation) + .transpose()? + .unwrap_or_default(); + let profile = self + .verifier + .verify_authorization(&secret, &id_token, &pending.nonce, &expectation) .await - .map_err(map_store_error)?; - if current.account.provider().as_str() != "openai" - || current.account.id() != &command.account_id - { - return Err(CodexOAuthAdminError::NotFound); - } - self.start_pending( - command.binding, - current.account.instance().to_string(), - current.account.name().to_owned(), - Some(CodexOAuthReauthorizationTarget { - account_id: command.account_id, - credential_revision: command.expected_credential_revision, - }), - ) - .await - } - - async fn complete_reauthorization( - &self, - command: CompleteCodexOAuthAuthorization, - ) -> Result { - let (pending, secret, profile) = self.exchange_pending(command).await?; - let target = pending - .reauthorization() - .ok_or(CodexOAuthAdminError::Conflict)?; - let current = self - .store - .load_credential(target.account_id(), target.credential_revision()) + .and_then(super::identity::CodexIdentityVerification::into_complete) + .map_err(map_identity_error)?; + secret.id_token = Some(id_token); + let policy = self + .runtime_policy + .load_refresh_policy() .await - .map_err(map_store_error)?; - self.credentials - .prepare_rotation(RotateManagedCodexCredential { - current, - secret, - verified_account: profile, - }) - .map_err(map_admin_error) + .map_err(|_| CodexOAuthAdminError::StorageUnavailable)?; + let credential = if let Some(current) = current { + let next_refresh_at = refresh_time( + policy, + current.account.id(), + profile.access_token_expires_at, + secret.refresh_token.is_some(), + ) + .map_err(map_admin_error)?; + CompletedCodexOAuthCredential::Reauthorize( + self.credentials + .prepare_rotation(RotateManagedCodexCredential { + current, + secret, + verified_account: profile, + next_refresh_at, + }) + .map_err(map_admin_error)?, + ) + } else { + let account_id = format!("acct_{}", Uuid::now_v7().simple()); + let typed_account_id = ProviderAccountId::new(account_id.clone()) + .map_err(|_| CodexOAuthAdminError::Credential)?; + let next_refresh_at = refresh_time( + policy, + &typed_account_id, + profile.access_token_expires_at, + secret.refresh_token.is_some(), + ) + .map_err(map_admin_error)?; + CompletedCodexOAuthCredential::Create( + self.credentials + .prepare_import(ImportCodexOAuthCredential { + account_id, + provider_instance_id: pending.provider_instance_id, + name: pending.name, + secret, + verified_account: profile, + next_refresh_at, + enabled: true, + }) + .map_err(map_admin_error)?, + ) + }; + Ok(CompletedCodexOAuthAuthorization { + mutation, + credential, + }) } } impl CodexOAuthAdminService { async fn start_pending( &self, - binding: CodexOAuthFlowBinding, + mutation: PendingAuthorizationMutation, provider_instance_id: String, name: String, reauthorization: Option, @@ -471,10 +518,12 @@ impl CodexOAuthAdminService { let expires_at = Utc::now() .checked_add_signed(AUTHORIZATION_TTL) .ok_or(CodexOAuthAdminError::InvalidInput)?; + let owner_ref = oauth_owner_ref(mutation.owner_binding().owner()); + let started_request_ref = mutation.owner_binding().started_request_id().to_owned(); let pending = CodexPendingAuthorization::from_stored(StoredCodexPendingAuthorization { flow_id: random_secret()?, - owner_ref: binding.owner_ref, - started_request_ref: binding.started_request_ref, + owner_ref, + started_request_ref, provider_instance_id, name, expires_at, @@ -486,6 +535,7 @@ impl CodexOAuthAdminService { .map(|target| target.account_id().to_string()), reauthorization_credential_revision: reauthorization .map(|target| target.credential_revision().get()), + mutation, }) .map_err(map_pending_error)?; self.pending @@ -506,7 +556,7 @@ impl CodexOAuthAdminService { ( CodexPendingAuthorization, super::types::CodexOAuthSecret, - super::types::CodexAccountProfile, + SecretString, ), CodexOAuthAdminError, > { @@ -532,7 +582,7 @@ impl CodexOAuthAdminService { ) { return Err(CodexOAuthAdminError::UpstreamRejected); } - let mut tokens = self + let tokens = self .exchanger .exchange_authorization_code(AuthorizationCodeGrant { code, @@ -540,16 +590,29 @@ impl CodexOAuthAdminService { }) .await .map_err(map_exchange_error)?; - let profile = self - .verifier - .verify_authorization(&tokens.secret, &tokens.id_token, &pending.nonce) - .await - .map_err(map_identity_error)?; - tokens.secret.id_token = Some(tokens.id_token); - Ok((pending, tokens.secret, profile)) + Ok((pending, tokens.secret, tokens.id_token)) } } +fn oauth_identity_expectation( + current: &gateway_core::engine::credential::LoadedCredential, +) -> Result { + let runtime = CodexCredentialCodec::decode(¤t.credential) + .map_err(|_| CodexOAuthAdminError::Credential)?; + let account_id = current + .account + .upstream_account_id() + .ok_or(CodexOAuthAdminError::Credential)?; + CodexIdentityExpectation::current( + runtime.principal.oauth_subject, + runtime.principal.poid, + account_id.to_owned(), + current.account.upstream_user_id().to_owned(), + runtime.installation_id, + ) + .map_err(|_| CodexOAuthAdminError::Credential) +} + fn authorization_url(pending: &CodexPendingAuthorization) -> Result { let mut url = Url::parse(AUTHORIZATION_ENDPOINT).map_err(|_| CodexOAuthAdminError::InvalidInput)?; @@ -615,6 +678,50 @@ fn valid_secret(value: &str) -> bool { valid_text(value) && value.len() >= 16 } +pub(crate) fn oauth_owner_ref(owner: &AuthorizationOwner) -> String { + let mut digest = Sha256::new(); + match owner { + AuthorizationOwner::AdminSession { admin_user_id } => { + digest.update(b"admin-session\0"); + digest.update(admin_user_id.as_bytes()); + } + AuthorizationOwner::AdminApiKey => digest.update(b"admin-api-key"), + AuthorizationOwner::System => digest.update(b"system"), + } + URL_SAFE_NO_PAD.encode(digest.finalize()) +} + +fn pending_mutation_matches(pending: &CodexPendingAuthorization) -> bool { + if pending.mutation.provider_kind().as_str() != "openai" + || pending.mutation.owner_binding().started_request_id() != pending.started_request_ref + || oauth_owner_ref(pending.mutation.owner_binding().owner()) != pending.owner_ref + { + return false; + } + match (pending.mutation.target(), pending.reauthorization.as_ref()) { + ( + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + }, + None, + ) => provider_instance_id.as_str() == pending.provider_instance_id && name == &pending.name, + ( + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + }, + Some(target), + ) => { + provider_instance_id.as_str() == pending.provider_instance_id + && account_id == target.account_id() + && expected_credential_revision.get() == target.credential_revision().get() + } + _ => false, + } +} + fn map_pending_error(error: CodexOAuthPendingStoreError) -> CodexOAuthAdminError { match error { CodexOAuthPendingStoreError::InvalidValue => CodexOAuthAdminError::InvalidInput, diff --git a/backend/crates/providers/openai/src/credential/quota.rs b/backend/crates/providers/openai/src/credential/quota.rs index aab386d92..babcc8e13 100644 --- a/backend/crates/providers/openai/src/credential/quota.rs +++ b/backend/crates/providers/openai/src/credential/quota.rs @@ -1,30 +1,35 @@ //! Codex quota 原始观察、Provider-owned 解析与账号状态投影。 -use std::sync::Arc; -use std::time::{Duration, SystemTime}; +use std::collections::BTreeMap; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant, SystemTime}; use chrono::{DateTime, Utc}; use gateway_core::engine::credential::{ - AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccountId, - ProviderAccountStore, QuotaObservation, QuotaWriteOutcome, + AccountAvailability, AccountQuotaSignals, CredentialRevision, OpaqueProviderData, + ProviderAccount, ProviderAccountId, ProviderAccountStore, QuotaObservation, QuotaWriteOutcome, }; use gateway_core::routing::ProviderInstance; +use gateway_protocol::openai::events::{ + ParsedRateLimits, RateLimitDetails, RateLimitWindow, parse_rate_limit_headers, +}; use reqwest::{Client, StatusCode}; use secrecy::ExposeSecret; use serde_json::{Map, Value}; use thiserror::Error; +use tokio::sync::Mutex; use uuid::Uuid; -use crate::provider::{CodexEndpointPolicy, CodexProviderInstanceConfig}; +use crate::provider::{CodexOriginPolicy, CodexProviderInstanceConfig}; use crate::transport::profile::CodexWireProfileState; -use crate::transport::{ - CodexBackendClient, CodexClientError, CodexRequestContext, build_reqwest_client, -}; +use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext}; use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60); const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60); +const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); +const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct CodexQuotaFact { @@ -172,8 +177,6 @@ pub enum CodexCredentialQuotaError { InvalidInstance, #[error("Codex quota response is invalid")] InvalidCredentialData, - #[error("Codex quota transport could not initialize")] - TransportInitialization, #[error(transparent)] Repository(#[from] CredentialRepositoryError), #[error("provider account store is unavailable")] @@ -197,31 +200,237 @@ pub struct CodexCredentialQuotaService { store: Arc, profile: CodexWireProfileState, http: Client, - endpoint_policy: CodexEndpointPolicy, + origin_policy: Arc, + scheduling: CodexQuotaSchedulingProjection, +} + +#[derive(Clone, Default)] +struct CodexQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct CodexQuotaProjectionState { + next_version: u64, + entries: BTreeMap, +} + +#[derive(Debug, Clone, Copy)] +struct CodexQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct CodexQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl CodexQuotaSchedulingProjection { + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| CodexQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + fn observe( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + observed_at: SystemTime, + fact: CodexQuotaFact, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(observed_at) else { + return false; + }; + self.replace( + account_id, + revision, + remaining_ttl, + quota_scheduling_signals(fact), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &CodexQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn observe_if_unchanged( + &self, + target: &CodexQuotaHydrationTarget, + snapshot: &CodexAccountQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.fact()), + ); + true + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } +} + +fn insert_projection_entry( + state: &mut CodexQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + CodexQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); } impl CodexCredentialQuotaService { pub fn new( repository: CodexCredentialRepository, profile: CodexWireProfileState, - ) -> Result { - Self::new_with_endpoint_policy(repository, profile, CodexEndpointPolicy::Official) - } - - pub fn new_with_endpoint_policy( - repository: CodexCredentialRepository, - profile: CodexWireProfileState, - endpoint_policy: CodexEndpointPolicy, - ) -> Result { - let http = build_reqwest_client() - .map_err(|_| CodexCredentialQuotaError::TransportInitialization)?; - Ok(Self { + http: Client, + origin_policy: Arc, + ) -> Self { + Self { store: Arc::clone(repository.store()), repository, profile, http, - endpoint_policy, - }) + origin_policy, + scheduling: CodexQuotaSchedulingProjection::default(), + } + } + + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度。 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let account_ids = pending + .iter() + .map(|target| target.account.id().clone()) + .collect::>(); + let Ok(observations) = self.store.get_quotas(&account_ids).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) } pub async fn synchronize_instance( @@ -229,7 +438,7 @@ impl CodexCredentialQuotaService { instance: &gateway_core::routing::ProviderInstance, ) -> Result { let config = - CodexProviderInstanceConfig::from_snapshot_with_policy(instance, self.endpoint_policy) + CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref()) .map_err(|_| CodexCredentialQuotaError::InvalidInstance)?; let client = CodexBackendClient::new( self.http.clone(), @@ -238,7 +447,11 @@ impl CodexCredentialQuotaService { ); let accounts = self.repository.list_for_instance(config.id()).await?; let mut summary = CodexQuotaSyncSummary::default(); - for account in accounts.into_iter().filter(|account| account.enabled()) { + let now = SystemTime::now(); + for account in accounts + .into_iter() + .filter(|account| eligible_quota_sync_account(account, now)) + { let runtime = match self.repository.load_runtime_credential(&account).await { Ok(runtime) => runtime, Err(_) => { @@ -294,6 +507,12 @@ impl CodexCredentialQuotaService { summary.stale += 1; continue; } + self.scheduling.observe( + current.id().clone(), + current.revision(), + observed_at, + fact, + ); if fact.exhausted() { summary.exhausted += 1; } else { @@ -346,6 +565,69 @@ impl CodexCredentialQuotaService { Ok(summary) } + /// 把正常推理响应携带的限流事实合并进 Provider 原始 quota JSON。 + pub async fn synchronize_passive_headers( + &self, + account: &ProviderAccount, + headers: &[(String, String)], + ) -> Result { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return Ok(false); + }; + let existing = self + .store + .get_quotas(std::slice::from_ref(account.id())) + .await? + .into_iter() + .find(|observation| { + observation.account_id == *account.id() + && observation.expected_revision == account.revision() + }) + .and_then(|observation| observation.quota) + .map(OpaqueProviderData::into_inner) + .unwrap_or_default(); + let quota = merge_passive_quota(existing, &rate_limits); + let fact = parse_codex_quota_usage(&Value::Object(quota.clone()))?; + let observed_at = SystemTime::now(); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: Some(OpaqueProviderData::new(quota)), + observed_at: Some(observed_at), + }) + .await?; + if outcome == QuotaWriteOutcome::Conflict { + return Ok(false); + } + let Some(current) = self.store.get_account(account.id()).await? else { + return Ok(false); + }; + if current.revision() != account.revision() { + return Ok(false); + } + self.scheduling + .observe(current.id().clone(), current.revision(), observed_at, fact); + if let Some(availability) = quota_success_availability( + current.availability(), + current.cooldown_until(), + fact.exhausted(), + observed_at, + ) { + self.repository + .apply_state( + ¤t, + availability, + fact.exhausted().then_some("quota_exhausted".to_owned()), + None, + observed_at, + ) + .await?; + } + Ok(true) + } + /// 读取单账号最后一次落库的 Provider quota，并由 Codex 域解析展示窗口。 pub async fn read_account( &self, @@ -357,7 +639,13 @@ impl CodexCredentialQuotaService { .await? .filter(|account| account.provider().as_str() == "openai") .ok_or(CodexCredentialQuotaError::NotFound)?; - let Some(observation) = self.store.get_quota(account_id).await? else { + let Some(observation) = self + .store + .get_quotas(std::slice::from_ref(account_id)) + .await? + .into_iter() + .next() + else { return Ok(None); }; if observation.account_id != *account_id @@ -371,13 +659,19 @@ impl CodexCredentialQuotaService { let observed_at = observation .observed_at .ok_or(CodexCredentialQuotaError::InvalidCredentialData)?; - parse_account_quota_snapshot( + let snapshot = parse_account_quota_snapshot( account_id.clone(), account.revision(), observed_at, &Value::Object(data.expose_to_provider().clone()), - ) - .map(Some) + )?; + self.scheduling.observe( + snapshot.account_id().clone(), + snapshot.credential_revision(), + snapshot.observed_at(), + snapshot.fact(), + ); + Ok(Some(snapshot)) } /// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照。 @@ -387,7 +681,7 @@ impl CodexCredentialQuotaService { account_id: &ProviderAccountId, ) -> Result { let config = - CodexProviderInstanceConfig::from_snapshot_with_policy(instance, self.endpoint_policy) + CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref()) .map_err(|_| CodexCredentialQuotaError::InvalidInstance)?; let account = self .store @@ -473,6 +767,12 @@ impl CodexCredentialQuotaService { if current.revision() != account.revision() { return Err(CodexCredentialQuotaError::RevisionConflict); } + self.scheduling.observe( + snapshot.account_id().clone(), + snapshot.credential_revision(), + snapshot.observed_at(), + snapshot.fact(), + ); if let Some(availability) = quota_success_availability( current.availability(), current.cooldown_until(), @@ -496,6 +796,150 @@ impl CodexCredentialQuotaService { } } +fn eligible_quota_sync_account(account: &ProviderAccount, now: SystemTime) -> bool { + account.enabled() + && account.access_token_expires_at() > now + && match account.availability() { + AccountAvailability::Unknown + | AccountAvailability::Ready + | AccountAvailability::Cooldown + | AccountAvailability::QuotaExhausted => true, + AccountAvailability::Expired + | AccountAvailability::Banned + | AccountAvailability::Invalid => false, + } +} + +fn merge_passive_quota( + mut quota: Map, + rate_limits: &ParsedRateLimits, +) -> Map { + let active_limit = rate_limits + .active_limit + .as_deref() + .or_else(|| rate_limits.limits.contains_key("codex").then_some("codex")); + if let Some(active_limit) = active_limit { + quota.insert( + "active_limit".to_owned(), + Value::String(active_limit.to_owned()), + ); + } + if let Some(plan_type) = rate_limits.plan_type.as_ref() { + quota.insert("plan_type".to_owned(), Value::String(plan_type.clone())); + } + if let Some(promo_message) = rate_limits.promo_message.as_ref() { + quota.insert( + "promo_message".to_owned(), + Value::String(promo_message.clone()), + ); + } + if let Some(reached_type) = rate_limits.rate_limit_reached_type.as_ref() { + quota.insert( + "rate_limit_reached_type".to_owned(), + Value::String(reached_type.clone()), + ); + } + if let Some(credits) = rate_limits.credits.as_ref() { + let mut value = quota + .remove("credits") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + value.insert("has_credits".to_owned(), Value::Bool(credits.has_credits)); + value.insert("unlimited".to_owned(), Value::Bool(credits.unlimited)); + if let Some(balance) = credits.balance.as_ref() { + value.insert("balance".to_owned(), Value::String(balance.clone())); + } + quota.insert("credits".to_owned(), Value::Object(value)); + } + + let mut additional = quota + .remove("additional_rate_limits") + .and_then(|value| value.as_array().cloned()) + .unwrap_or_default(); + for (limit_id, details) in &rate_limits.limits { + if Some(limit_id.as_str()) == active_limit { + let mut value = quota + .remove("rate_limit") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + merge_rate_limit_details(&mut value, details); + quota.insert("rate_limit".to_owned(), Value::Object(value)); + continue; + } + let index = additional.iter().position(|item| { + item.get("metered_feature") + .or_else(|| item.get("limit_name")) + .and_then(Value::as_str) + .is_some_and(|existing| existing == limit_id) + }); + let mut item = index + .and_then(|index| additional.get(index)) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + item.insert( + "metered_feature".to_owned(), + Value::String(limit_id.clone()), + ); + if let Some(name) = details.limit_name.as_ref() { + item.insert("limit_name".to_owned(), Value::String(name.clone())); + } + let mut value = item + .remove("rate_limit") + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + merge_rate_limit_details(&mut value, details); + item.insert("rate_limit".to_owned(), Value::Object(value)); + match index { + Some(index) => additional[index] = Value::Object(item), + None => additional.push(Value::Object(item)), + } + } + quota.insert( + "additional_rate_limits".to_owned(), + Value::Array(additional), + ); + quota +} + +fn merge_rate_limit_details(target: &mut Map, details: &RateLimitDetails) { + if let Some(allowed) = details.allowed { + target.insert("allowed".to_owned(), Value::Bool(allowed)); + } + if let Some(limit_reached) = details.limit_reached { + target.insert("limit_reached".to_owned(), Value::Bool(limit_reached)); + } + for (field, window) in [ + ("primary_window", details.primary), + ("secondary_window", details.secondary), + ] { + let Some(window) = window else { + continue; + }; + let mut value = target + .remove(field) + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + merge_rate_limit_window(&mut value, window); + target.insert(field.to_owned(), Value::Object(value)); + } +} + +fn merge_rate_limit_window(target: &mut Map, window: RateLimitWindow) { + if let Some(number) = serde_json::Number::from_f64(window.used_percent) { + target.insert("used_percent".to_owned(), Value::Number(number)); + } + if let Some(seconds) = window + .window_minutes + .and_then(|minutes| minutes.checked_mul(60)) + { + target.insert("limit_window_seconds".to_owned(), Value::from(seconds)); + } + if let Some(reset_at) = window.reset_at { + target.insert("reset_at".to_owned(), Value::from(reset_at)); + } +} + fn quota_success_availability( current: AccountAvailability, cooldown_until: Option, @@ -634,6 +1078,40 @@ pub fn parse_codex_quota_usage(usage: &Value) -> Result Option { + let observed_at = observation.observed_at?; + let quota = observation.quota.as_ref()?; + parse_account_quota_snapshot( + observation.account_id.clone(), + observation.expected_revision, + observed_at, + &Value::Object(quota.expose_to_provider().clone()), + ) + .ok() +} + +fn quota_projection_ttl(observed_at: SystemTime) -> Option { + let age = SystemTime::now() + .duration_since(observed_at) + .unwrap_or(Duration::ZERO); + QUOTA_SCHEDULING_TTL + .checked_sub(age) + .filter(|remaining| !remaining.is_zero()) +} + +fn quota_scheduling_signals(fact: CodexQuotaFact) -> Option { + let remaining_rank = fact.remaining_percent().map(|value| u64::from(value) * 100); + let now = SystemTime::now(); + let reset_at = fact + .resets_at() + .map(SystemTime::from) + .filter(|reset_at| *reset_at > now); + (remaining_rank.is_some() || reset_at.is_some()) + .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) +} + fn parse_account_quota_snapshot( account_id: ProviderAccountId, credential_revision: CredentialRevision, diff --git a/backend/crates/providers/openai/src/credential/refresh.rs b/backend/crates/providers/openai/src/credential/refresh.rs index 838fcbe03..3c4cadf2d 100644 --- a/backend/crates/providers/openai/src/credential/refresh.rs +++ b/backend/crates/providers/openai/src/credential/refresh.rs @@ -4,25 +4,34 @@ use std::collections::BTreeSet; use std::sync::Arc; use std::time::{Duration, SystemTime}; -use async_trait::async_trait; use chrono::{DateTime, Utc}; -use gateway_core::engine::credential::{ - AccountAvailability, CredentialRevision, ProviderAccount, ProviderAccountId, +use futures::{StreamExt as _, stream}; +use gateway_core::engine::credential::{AccountAvailability, ProviderAccount, ProviderAccountId}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_retry_at, }; use secrecy::{ExposeSecret, SecretString}; use thiserror::Error; +use super::identity::{ + CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerification, + CodexIdentityVerificationError, +}; use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; use super::token_client::{RefreshFailure, TokenPair, TokenRefresher}; -use super::types::{CodexAccountProfile, CodexOAuthSecret, RotateCodexCredential}; +use super::types::{CodexCredentialPrincipal, CodexOAuthSecret, RotateCodexCredential}; const PROVIDER_NAME: &str = "openai"; const MAX_REFRESH_BATCH: u32 = 1_000; -const TRANSIENT_COOLDOWN: Duration = Duration::from_secs(30); +const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60); pub struct DueCodexCredential { pub account: ProviderAccount, pub secret: CodexOAuthSecret, + principal: CodexCredentialPrincipal, + installation_id: String, } impl std::fmt::Debug for DueCodexCredential { @@ -35,34 +44,6 @@ impl std::fmt::Debug for DueCodexCredential { } } -#[derive(Debug, Clone)] -pub struct CodexRefreshLeaseRequest { - pub account_id: ProviderAccountId, - pub credential_revision: CredentialRevision, -} - -pub trait CodexRefreshLeaseGuard: Send + Sync + 'static {} -impl CodexRefreshLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum CodexRefreshLeaseAcquisition { - Acquired(Box), - Unavailable, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] -pub enum CodexRefreshLeaseError { - #[error("Codex refresh lease runtime is unavailable")] - Unavailable, -} - -#[async_trait] -pub trait CodexRefreshLeaseCoordinator: Send + Sync { - async fn try_acquire( - &self, - request: &CodexRefreshLeaseRequest, - ) -> Result; -} - #[derive(Debug, Clone, PartialEq, Eq)] pub enum CodexCredentialRefreshOutcome { Refreshed { @@ -97,9 +78,7 @@ pub enum CodexCredentialRefreshError { #[error(transparent)] Repository(#[from] CredentialRepositoryError), #[error(transparent)] - Lease(#[from] CodexRefreshLeaseError), - #[error("Codex refresh configuration is invalid")] - InvalidConfiguration, + Lease(#[from] ProviderStoreError), #[error("Codex refresh response is invalid")] InvalidRefreshResponse, } @@ -107,48 +86,53 @@ pub enum CodexCredentialRefreshError { pub struct CodexCredentialRefreshService { repository: CodexCredentialRepository, refresher: Arc, - leases: Arc, - refresh_margin: Duration, + identity: Arc, + leases: Arc, + runtime_policy: Arc, } impl CodexCredentialRefreshService { pub fn new( repository: CodexCredentialRepository, refresher: Arc, - leases: Arc, - refresh_margin: Duration, - ) -> Result { - if refresh_margin.is_zero() { - return Err(CodexCredentialRefreshError::InvalidConfiguration); - } - Ok(Self { + identity: Arc, + leases: Arc, + runtime_policy: Arc, + ) -> Self { + Self { repository, refresher, + identity, leases, - refresh_margin, - }) + runtime_policy, + } } pub async fn refresh_due( &self, - limit: u32, ) -> Result, CodexCredentialRefreshError> { - self.refresh_due_excluding(limit, &BTreeSet::new()).await + self.refresh_due_excluding(&BTreeSet::new()).await } pub async fn refresh_due_excluding( &self, - limit: u32, excluded: &BTreeSet, ) -> Result, CodexCredentialRefreshError> { - if limit == 0 || limit > MAX_REFRESH_BATCH { - return Err(CodexCredentialRefreshError::InvalidConfiguration); - } - let (due, mut outcomes) = self.list_due_refresh(limit, excluded).await?; + let policy = self.runtime_policy.load_refresh_policy().await?; + let (due, mut outcomes) = self.list_due_refresh(excluded).await?; outcomes.reserve(due.len()); - for credential in due { + let refreshed = stream::iter(due.into_iter().map(|credential| async move { let account_id = credential.account.id().to_string(); - match self.refresh_one(credential).await { + ( + account_id, + self.refresh_one_with_policy(credential, policy).await, + ) + })) + .buffer_unordered(policy.concurrency().get() as usize) + .collect::>() + .await; + for (account_id, result) in refreshed { + match result { Ok(outcome) => outcomes.push(outcome), Err(_) => outcomes.push(CodexCredentialRefreshOutcome::Failed { account_id }), } @@ -159,18 +143,38 @@ impl CodexCredentialRefreshService { pub async fn refresh_one( &self, due: DueCodexCredential, + ) -> Result { + let policy = self.runtime_policy.load_refresh_policy().await?; + self.refresh_one_with_policy(due, policy).await + } + + async fn refresh_one_with_policy( + &self, + due: DueCodexCredential, + policy: ProviderRefreshPolicy, ) -> Result { let account_id = due.account.id().to_string(); + let capacity = self + .leases + .try_acquire(ProviderLeaseRequest::RefreshCapacity( + ProviderRefreshCapacityRequest::new(policy.concurrency()), + )) + .await?; + let _capacity_guard = match capacity { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { + return Ok(CodexCredentialRefreshOutcome::LeaseUnavailable { account_id }); + } + }; let acquisition = self .leases - .try_acquire(&CodexRefreshLeaseRequest { - account_id: due.account.id().clone(), - credential_revision: due.account.revision(), - }) + .try_acquire(ProviderLeaseRequest::Refresh( + ProviderRefreshLeaseRequest::new(due.account.id().clone(), due.account.revision()), + )) .await?; let _guard = match acquisition { - CodexRefreshLeaseAcquisition::Acquired(guard) => guard, - CodexRefreshLeaseAcquisition::Unavailable => { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { return Ok(CodexCredentialRefreshOutcome::LeaseUnavailable { account_id }); } }; @@ -180,37 +184,55 @@ impl CodexCredentialRefreshService { .as_ref() .ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?; match self.refresher.refresh(refresh_token.expose_secret()).await { - Ok(tokens) => self.persist_success(due, tokens).await, + Ok(tokens) => self.persist_success(due, tokens, policy).await, Err(RefreshFailure::InvalidGrant) => { - self.invalidate(&due.account, AccountAvailability::Expired, "invalid_grant") - .await?; - Ok(CodexCredentialRefreshOutcome::Invalidated { account_id }) + self.persist_terminal( + &due.account, + AccountAvailability::Expired, + "invalid_grant", + CodexCredentialRefreshOutcome::Invalidated { account_id }, + ) + .await } Err(RefreshFailure::Banned) => { - self.invalidate(&due.account, AccountAvailability::Banned, "account_banned") - .await?; - Ok(CodexCredentialRefreshOutcome::Banned { account_id }) + self.persist_terminal( + &due.account, + AccountAvailability::Banned, + "account_banned", + CodexCredentialRefreshOutcome::Banned { account_id }, + ) + .await } Err(RefreshFailure::RetryableTransport) => { - self.cooldown(&due.account).await?; - Ok(CodexCredentialRefreshOutcome::Transient { account_id }) + if self + .defer_refresh(&due.account, "transport-not-sent") + .await? + { + Ok(CodexCredentialRefreshOutcome::Transient { account_id }) + } else { + Ok(CodexCredentialRefreshOutcome::Stale { account_id }) + } } Err(RefreshFailure::Transport) => { - Ok(CodexCredentialRefreshOutcome::Ambiguous { account_id }) + self.persist_terminal( + &due.account, + AccountAvailability::Invalid, + "refresh_ambiguous", + CodexCredentialRefreshOutcome::Ambiguous { account_id }, + ) + .await } } } async fn list_due_refresh( &self, - limit: u32, excluded: &BTreeSet, ) -> Result< (Vec, Vec), CodexCredentialRefreshError, > { let now = SystemTime::now(); - let threshold = now.checked_add(self.refresh_margin).unwrap_or(now); let mut accounts = self .repository .store() @@ -222,12 +244,18 @@ impl CodexCredentialRefreshService { && account.enabled() && account.has_refresh_token() && !excluded.contains(account.id()) - && account - .next_refresh_at() - .is_some_and(|next| next <= threshold) + && !matches!( + account.availability(), + AccountAvailability::Expired + | AccountAvailability::Banned + | AccountAvailability::Invalid + ) + && (account.availability() != AccountAvailability::Cooldown + || account.cooldown_until().is_some_and(|until| until <= now)) + && account.next_refresh_at().is_some_and(|next| next <= now) }); accounts.sort_by_key(|account| (account.next_refresh_at(), account.id().clone())); - accounts.truncate(limit as usize); + accounts.truncate(MAX_REFRESH_BATCH as usize); let mut due = Vec::with_capacity(accounts.len()); let mut failures = Vec::new(); for account in accounts { @@ -237,6 +265,8 @@ impl CodexCredentialRefreshService { due.push(DueCodexCredential { account, secret: runtime.secret, + principal: runtime.principal, + installation_id: runtime.installation_id, }); } Ok(_) | Err(_) => { @@ -251,44 +281,106 @@ impl CodexCredentialRefreshService { &self, due: DueCodexCredential, tokens: TokenPair, + policy: ProviderRefreshPolicy, ) -> Result { - if tokens.access_token.is_empty() || tokens.expires_in <= self.refresh_margin { + if tokens.access_token.is_empty() || tokens.expires_in.is_zero() { return Err(CodexCredentialRefreshError::InvalidRefreshResponse); } let refresh_token = tokens .refresh_token .map(SecretString::from) .or(due.secret.refresh_token); - let now = SystemTime::now(); - let expires_at = now - .checked_add(tokens.expires_in) + let account_id = due + .account + .upstream_account_id() .ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?; - let next_refresh_at = expires_at - .checked_sub(self.refresh_margin) - .ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?; - let profile = CodexAccountProfile { - email: due.account.email().map(str::to_owned), - chatgpt_account_id: due - .account - .upstream_account_id() - .unwrap_or(due.account.upstream_user_id()) - .to_owned(), - chatgpt_user_id: Some(due.account.upstream_user_id().to_owned()), - plan_type: due.account.plan_type().map(str::to_owned), - access_token_expires_at: Some(DateTime::::from(expires_at)), - next_refresh_at: Some(DateTime::::from(next_refresh_at)), + let expectation = CodexIdentityExpectation::current( + due.principal.oauth_subject, + due.principal.poid, + account_id.to_owned(), + due.account.upstream_user_id().to_owned(), + due.installation_id, + ) + .map_err(|_| CodexCredentialRefreshError::InvalidRefreshResponse)?; + let secret = CodexOAuthSecret { + access_token: SecretString::from(tokens.access_token), + refresh_token, + id_token: due.secret.id_token, }; + let verification = match self.identity.verify(&secret, &expectation).await { + Ok(verification) => verification, + Err(CodexIdentityVerificationError::Rejected) => { + return self + .persist_terminal( + &due.account, + AccountAvailability::Invalid, + "refreshed_identity_rejected", + CodexCredentialRefreshOutcome::Invalidated { + account_id: due.account.id().to_string(), + }, + ) + .await; + } + Err(CodexIdentityVerificationError::Unavailable) => { + return self + .persist_terminal( + &due.account, + AccountAvailability::Invalid, + "refreshed_identity_unavailable", + CodexCredentialRefreshOutcome::Ambiguous { + account_id: due.account.id().to_string(), + }, + ) + .await; + } + }; + let profile = match verification { + CodexIdentityVerification::Complete(profile) => profile, + CodexIdentityVerification::SignedOnly(signed) => { + let retry_at = provider_refresh_retry_at( + due.account.id(), + SystemTime::now(), + REFRESH_RETRY_DELAY, + "identity-completion", + )?; + match self + .repository + .rotate_signed_secret(&due.account, secret, &signed, retry_at) + .await + { + Ok(_) => { + return Ok(CodexCredentialRefreshOutcome::Transient { + account_id: due.account.id().to_string(), + }); + } + Err(CredentialRepositoryError::RevisionConflict) => { + return Ok(CodexCredentialRefreshOutcome::Stale { + account_id: due.account.id().to_string(), + }); + } + Err(error) => return Err(error.into()), + } + } + }; + let next_refresh_at = policy + .next_attempt_at( + due.account.id(), + profile + .access_token_expires_at + .map(SystemTime::from) + .ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?, + SystemTime::now(), + ) + .map(DateTime::::from) + .map_err(CodexCredentialRefreshError::from)?; let result = self .repository .rotate_oauth_secret(RotateCodexCredential { account_id: due.account.id().to_string(), expected_credential_revision: due.account.revision().get(), - secret: CodexOAuthSecret { - access_token: SecretString::from(tokens.access_token), - refresh_token, - id_token: due.secret.id_token, - }, + secret, verified_account: profile, + next_refresh_at: Some(next_refresh_at), }) .await; match result { @@ -305,13 +397,24 @@ impl CodexCredentialRefreshService { } } - async fn invalidate( + async fn persist_terminal( &self, account: &ProviderAccount, availability: AccountAvailability, reason: &'static str, - ) -> Result<(), CodexCredentialRefreshError> { - self.repository + outcome: CodexCredentialRefreshOutcome, + ) -> Result { + match self.repository.load_runtime_credential(account).await { + Ok(_) => {} + Err(CredentialRepositoryError::RevisionConflict) => { + return Ok(CodexCredentialRefreshOutcome::Stale { + account_id: account.id().to_string(), + }); + } + Err(error) => return Err(error.into()), + } + match self + .repository .apply_state( account, availability, @@ -319,21 +422,33 @@ impl CodexCredentialRefreshService { None, SystemTime::now(), ) - .await?; - Ok(()) + .await + { + Ok(()) => Ok(outcome), + Err(CredentialRepositoryError::RevisionConflict) => { + Ok(CodexCredentialRefreshOutcome::Stale { + account_id: account.id().to_string(), + }) + } + Err(error) => Err(error.into()), + } } - async fn cooldown(&self, account: &ProviderAccount) -> Result<(), CodexCredentialRefreshError> { - let now = SystemTime::now(); - self.repository - .apply_state( - account, - AccountAvailability::Cooldown, - Some("refresh_transport".to_owned()), - now.checked_add(TRANSIENT_COOLDOWN), - now, - ) - .await?; - Ok(()) + async fn defer_refresh( + &self, + account: &ProviderAccount, + reason: &'static str, + ) -> Result { + let retry_at = provider_refresh_retry_at( + account.id(), + SystemTime::now(), + REFRESH_RETRY_DELAY, + reason, + )?; + match self.repository.defer_refresh(account, retry_at).await { + Ok(_) => Ok(true), + Err(CredentialRepositoryError::RevisionConflict) => Ok(false), + Err(error) => Err(error.into()), + } } } diff --git a/backend/crates/providers/openai/src/credential/repository.rs b/backend/crates/providers/openai/src/credential/repository.rs index 1ed97bca2..ce0270f1f 100644 --- a/backend/crates/providers/openai/src/credential/repository.rs +++ b/backend/crates/providers/openai/src/credential/repository.rs @@ -13,10 +13,11 @@ use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use secrecy::ExposeSecret; use thiserror::Error; +use super::identity::CodexSignedIdentity; use super::security::{CodexCredentialCodec, CodexCredentialDataError, CodexRuntimeCredential}; use super::types::{ - CodexAccountProfile, CodexCredentialData, CreateCodexCredential, CredentialRecord, - RotateCodexCredential, + CodexAccountProfile, CodexCredentialData, CodexOAuthSecret, CreateCodexCredential, + CredentialRecord, RotateCodexCredential, }; const PROVIDER_NAME: &str = "openai"; @@ -50,14 +51,11 @@ impl CodexCredentialRepository { let revision = CredentialRevision::new(1) .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; let access_token_expires_at = required_time(input.account.access_token_expires_at)?; - let upstream_user_id = input - .account - .chatgpt_user_id - .clone() - .unwrap_or_else(|| input.account.chatgpt_account_id.clone()); - let credential = CodexCredentialCodec::encode(&input.secret, Vec::new())?; + let upstream_user_id = input.account.chatgpt_user_id.clone(); + let credential = + CodexCredentialCodec::encode_new(&input.secret, &input.account, Vec::new())?; let has_refresh_token = input.secret.refresh_token.is_some(); - let next_refresh_at = optional_time(input.account.next_refresh_at); + let next_refresh_at = optional_time(input.next_refresh_at); let account = ProviderAccount::new( account_id.clone(), instance.clone(), @@ -96,8 +94,8 @@ impl CodexCredentialRepository { let expected = CredentialRevision::new(input.expected_credential_revision) .map_err(|_| CredentialRepositoryError::InvalidInput("credential_revision"))?; let current = self.store.load_credential(&account_id, expected).await?; - verify_identity(¤t, &input.verified_account)?; let mut data = CodexCredentialCodec::decode_complete(¤t.credential)?; + verify_identity(¤t, &data, &input.verified_account)?; data.access_token = input.secret.access_token.expose_secret().to_owned(); data.refresh_token = input .secret @@ -122,14 +120,89 @@ impl CodexCredentialRepository { credential, input.secret.refresh_token.is_some(), required_time(input.verified_account.access_token_expires_at)?, - optional_time(input.verified_account.next_refresh_at), + optional_time(input.next_refresh_at), ) .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; - let outcome = self.store.compare_and_swap_credential(update).await?; - match outcome { - CredentialCasOutcome::Updated(revision) => Ok(revision), - CredentialCasOutcome::Conflict => Err(CredentialRepositoryError::RevisionConflict), + cas_revision(self.store.compare_and_swap_credential(update).await?) + } + + /// usage 暂时不可用时，只在签名 principal 与现有绑定完全一致后保存已轮换 token。 + pub async fn rotate_signed_secret( + &self, + account: &ProviderAccount, + secret: CodexOAuthSecret, + signed: &CodexSignedIdentity, + next_refresh_at: SystemTime, + ) -> Result { + let current = self + .store + .load_credential(account.id(), account.revision()) + .await?; + if current.account != *account { + return Err(CredentialRepositoryError::RevisionConflict); } + let mut data = CodexCredentialCodec::decode_complete(¤t.credential)?; + if data.principal.oauth_subject != signed.oauth_subject() + || data.principal.poid.as_deref() != signed.poid() + || signed + .claimed_account_id() + .is_some_and(|value| account.upstream_account_id() != Some(value)) + || signed + .claimed_user_id() + .is_some_and(|value| account.upstream_user_id() != value) + { + return Err(CredentialRepositoryError::IdentityMismatch); + } + data.access_token = secret.access_token.expose_secret().to_owned(); + data.refresh_token = secret + .refresh_token + .as_ref() + .map(|value| value.expose_secret().to_owned()); + data.id_token = secret + .id_token + .as_ref() + .map(|value| value.expose_secret().to_owned()); + let credential = CodexCredentialCodec::encode_complete(data)?; + let update = CredentialCasUpdate::new( + account.id().clone(), + account.revision(), + unchanged_profile(account), + credential, + secret.refresh_token.is_some(), + SystemTime::from(signed.access_token_expires_at()), + Some(next_refresh_at), + ) + .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; + cas_revision(self.store.compare_and_swap_credential(update).await?) + } + + /// token 签名边界不可用时，以相同 credential 推进持久刷新退避。 + pub async fn defer_refresh( + &self, + account: &ProviderAccount, + next_refresh_at: SystemTime, + ) -> Result { + let current = self + .store + .load_credential(account.id(), account.revision()) + .await?; + if current.account != *account { + return Err(CredentialRepositoryError::RevisionConflict); + } + let data = CodexCredentialCodec::decode_complete(¤t.credential)?; + let has_refresh_token = data.refresh_token.is_some(); + let credential = CodexCredentialCodec::encode_complete(data)?; + let update = CredentialCasUpdate::new( + account.id().clone(), + account.revision(), + unchanged_profile(account), + credential, + has_refresh_token, + account.access_token_expires_at(), + Some(next_refresh_at), + ) + .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; + cas_revision(self.store.compare_and_swap_credential(update).await?) } pub async fn list_for_instance( @@ -192,11 +265,7 @@ impl CodexCredentialRepository { account.next_refresh_at(), ) .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; - let outcome = self.store.compare_and_swap_credential(update).await?; - match outcome { - CredentialCasOutcome::Updated(revision) => Ok(revision), - CredentialCasOutcome::Conflict => Err(CredentialRepositoryError::RevisionConflict), - } + cas_revision(self.store.compare_and_swap_credential(update).await?) } pub async fn apply_state( @@ -221,15 +290,34 @@ impl CodexCredentialRepository { } } +fn unchanged_profile(account: &ProviderAccount) -> ProviderAccountUpdate { + ProviderAccountUpdate { + account_id: account.id().clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + } +} + +fn cas_revision( + outcome: CredentialCasOutcome, +) -> Result { + match outcome { + CredentialCasOutcome::Updated(revision) => Ok(revision), + CredentialCasOutcome::Conflict => Err(CredentialRepositoryError::RevisionConflict), + } +} + fn verify_identity( current: &LoadedCredential, + credential: &CodexCredentialData, verified: &CodexAccountProfile, ) -> Result<(), CredentialRepositoryError> { - if current.account.upstream_account_id() != Some(verified.chatgpt_account_id.as_str()) - || verified - .chatgpt_user_id - .as_deref() - .is_some_and(|user| user != current.account.upstream_user_id()) + if current.account.provider().as_str() != PROVIDER_NAME + || current.account.upstream_account_id() != Some(verified.chatgpt_account_id.as_str()) + || verified.chatgpt_user_id != current.account.upstream_user_id() + || verified.oauth_subject != credential.principal.oauth_subject + || verified.poid != credential.principal.poid { return Err(CredentialRepositoryError::IdentityMismatch); } @@ -261,8 +349,14 @@ pub enum CredentialRepositoryError { } impl From for CredentialRepositoryError { - fn from(_: gateway_core::error::StoreError) -> Self { - Self::Store + fn from(error: gateway_core::error::StoreError) -> Self { + match error.kind() { + gateway_core::error::StoreErrorKind::Conflict => Self::RevisionConflict, + gateway_core::error::StoreErrorKind::Unavailable + | gateway_core::error::StoreErrorKind::InvalidState + | gateway_core::error::StoreErrorKind::InvalidData => Self::Store, + _ => Self::Store, + } } } diff --git a/backend/crates/providers/openai/src/credential/security.rs b/backend/crates/providers/openai/src/credential/security.rs index 284f7176d..1f6e67ceb 100644 --- a/backend/crates/providers/openai/src/credential/security.rs +++ b/backend/crates/providers/openai/src/credential/security.rs @@ -5,7 +5,10 @@ use secrecy::{ExposeSecret, SecretString}; use serde_json::Value; use thiserror::Error; -use super::types::{CodexCookie, CodexCredentialData, CodexOAuthSecret, RuntimeCodexCookie}; +use super::types::{ + CodexAccountProfile, CodexCookie, CodexCredentialData, CodexCredentialPrincipal, + CodexOAuthSecret, RuntimeCodexCookie, +}; const CODEX_CREDENTIAL_SCHEMA_VERSION: u32 = 1; const MAX_CREDENTIAL_BYTES: usize = 256 * 1024; @@ -15,6 +18,8 @@ const MAX_COOKIES: usize = 128; /// 已解析且只在 Provider 内可见的认证材料。 pub struct CodexRuntimeCredential { pub secret: CodexOAuthSecret, + pub principal: CodexCredentialPrincipal, + pub installation_id: String, pub cookies: Vec, pub oauth_client_id: Option, pub oauth_scope: Option, @@ -25,6 +30,8 @@ impl std::fmt::Debug for CodexRuntimeCredential { formatter .debug_struct("CodexRuntimeCredential") .field("secret", &self.secret) + .field("principal", &self.principal) + .field("installation_id", &"") .field("cookies", &self.cookies) .field("oauth_client_id", &self.oauth_client_id) .field("oauth_scope", &self.oauth_scope) @@ -45,12 +52,18 @@ pub enum CodexCredentialDataError { pub struct CodexCredentialCodec; impl CodexCredentialCodec { - pub fn encode( + pub fn encode_new( secret: &CodexOAuthSecret, + account: &CodexAccountProfile, cookies: Vec, ) -> Result { Self::encode_complete(CodexCredentialData { schema_version: CODEX_CREDENTIAL_SCHEMA_VERSION, + principal: CodexCredentialPrincipal { + oauth_subject: account.oauth_subject.clone(), + poid: account.poid.clone(), + }, + installation_id: uuid::Uuid::new_v4().to_string(), access_token: secret.access_token.expose_secret().to_owned(), refresh_token: secret .refresh_token @@ -105,6 +118,8 @@ impl CodexCredentialCodec { refresh_token: data.refresh_token.map(SecretString::from), id_token: data.id_token.map(SecretString::from), }, + principal: data.principal, + installation_id: data.installation_id, cookies: data .cookies .into_iter() @@ -132,10 +147,30 @@ impl CodexCredentialCodec { validate(&data)?; Ok(data) } + + pub fn preserve_installation_id( + incoming: &PlaintextCredential, + existing: &PlaintextCredential, + ) -> Result { + let mut incoming = Self::decode_complete(incoming)?; + let existing = Self::decode_complete(existing)?; + if incoming.principal != existing.principal { + return Err(CodexCredentialDataError::Invalid); + } + incoming.installation_id = existing.installation_id; + Self::encode_complete(incoming) + } } fn validate(data: &CodexCredentialData) -> Result<(), CodexCredentialDataError> { if data.schema_version != CODEX_CREDENTIAL_SCHEMA_VERSION + || !valid_identity(&data.principal.oauth_subject) + || data + .principal + .poid + .as_deref() + .is_some_and(|value| !valid_identity(value)) + || !valid_installation_id(&data.installation_id) || !valid_secret(&data.access_token) || data .refresh_token @@ -162,6 +197,16 @@ fn validate(data: &CodexCredentialData) -> Result<(), CodexCredentialDataError> Ok(()) } +fn valid_identity(value: &str) -> bool { + !value.is_empty() && value.len() <= 2_048 && !value.chars().any(char::is_control) +} + +fn valid_installation_id(value: &str) -> bool { + uuid::Uuid::parse_str(value) + .ok() + .is_some_and(|uuid| uuid.get_version_num() == 4) +} + fn valid_secret(value: &str) -> bool { !value.is_empty() && value.len() <= MAX_TOKEN_BYTES && !value.chars().any(char::is_control) } diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 2392e2f98..6a036f7b6 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -1,94 +1,112 @@ //! AttemptContext 驱动的 Codex 账号选择与 Redis lease port。 -use std::collections::BTreeMap; +use std::collections::HashMap; use std::fmt; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; -use async_trait::async_trait; use gateway_core::engine::credential::{ AccountAvailability, AccountCandidate, AccountRuntimeSignals, AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId, }; -use gateway_core::engine::{AttemptContext, UpstreamSendState}; -use gateway_core::error::ProviderErrorKind; +use gateway_core::engine::{AttemptContext, ContinuationAttempt}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest, + ProviderSchedulingLeaseRequest, ProviderStoreError, +}; use gateway_core::routing::ProviderInstanceId; use secrecy::ExposeSecret; -use sha2::{Digest, Sha256}; use thiserror::Error; use url::Url; +use super::catalog::CodexCredentialCatalogService; use super::cookie::CodexCookiePolicy; +use super::quota::CodexCredentialQuotaService; use super::repository::{CodexCredentialRepository, CredentialRepositoryError}; use super::types::{CodexCookie, CodexCookieCaptureOutcome, CodexOAuthSecret, RuntimeCodexCookie}; const PROVIDER_NAME: &str = "openai"; -const DEFAULT_COOLDOWN: Duration = Duration::from_secs(30); +const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60); +const CLOUDFLARE_PATH_BLOCK_COOLDOWN: Duration = Duration::from_secs(30); +const CLOUDFLARE_RECOVERY_STALE_AFTER: Duration = Duration::from_secs(60 * 60); +const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [ + Duration::from_secs(10), + Duration::from_secs(30), + Duration::from_secs(90), + Duration::from_secs(120), +]; +const CLOUDFLARE_PATH_BLOCK_THRESHOLD: u32 = 3; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexAccountFailure { + /// Access token 已被上游明确判定为过期或失效。 + CredentialExpired, + /// 账号需要完成身份验证后才能继续使用。 + IdentityVerificationRequired, + /// 账号、workspace 或 organization 已被封禁或停用。 + Banned, + /// 账号信用额度已耗尽。 + QuotaExhausted, + /// 账号触发临时用量限制。 + RateLimited { + /// 上游明确返回的最短冷却时长。 + retry_after: Option, + }, + /// Cloudflare challenge 要求账号进入递增冷却。 + CloudflareChallenge { + /// 上游明确返回的最短冷却时长。 + retry_after: Option, + }, + /// Cloudflare 对当前上游路径返回空 404。 + CloudflarePathBlocked, +} + +#[derive(Debug, Clone, Copy)] +struct RiskRecoveryState { + challenge_count: u32, + path_block_count: u32, + observed_at: SystemTime, +} + +#[derive(Debug, Clone, Copy)] +enum CookieRecovery { + None, + ExpireAt(SystemTime), + Clear, +} pub struct SelectCodexCredential<'a> { pub provider_instance_id: &'a ProviderInstanceId, + pub upstream_model: &'a str, pub request_url: &'a Url, pub attempt: &'a AttemptContext, } -#[derive(Debug, Clone)] -pub struct CredentialLeaseRequest { - pub account_id: ProviderAccountId, - pub max_concurrent: u32, - pub request_interval: Duration, - pub deadline: SystemTime, -} - -pub trait CredentialLeaseGuard: Send + Sync + 'static {} -impl CredentialLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum LeaseAcquisition { - Acquired(Box), - Busy { retry_after: Option }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] -pub enum CredentialLeaseCoordinatorError { - #[error("credential lease runtime is unavailable")] - Unavailable, -} - -/// Redis adapter 实现：信号可失效，lease guard 通过 Drop 释放。 -#[async_trait] -pub trait CredentialLeaseCoordinator: Send + Sync { - async fn runtime_signals( - &self, - accounts: &[ProviderAccountId], - ) -> Result, CredentialLeaseCoordinatorError>; - - fn next_round_robin_cursor( - &self, - provider_instance_id: &ProviderInstanceId, - ) -> Result; - - async fn try_acquire( - &self, - request: CredentialLeaseRequest, - ) -> Result; -} - pub struct CodexCredentialSelector { repository: CodexCredentialRepository, - leases: Arc, + leases: Arc, + catalog: Arc, + quota: Arc, cookie_policy: CodexCookiePolicy, + risk_recovery: Mutex>, } impl CodexCredentialSelector { #[must_use] - pub const fn new( + pub fn new( repository: CodexCredentialRepository, - leases: Arc, + leases: Arc, + catalog: Arc, + quota: Arc, cookie_policy: CodexCookiePolicy, ) -> Self { Self { repository, leases, + catalog, + quota, cookie_policy, + risk_recovery: Mutex::new(HashMap::new()), } } @@ -105,20 +123,32 @@ impl CodexCredentialSelector { .filter(|account| { account.provider().as_str() == PROVIDER_NAME && account.instance() == request.provider_instance_id + && !matches!( + self.catalog.observed_model_support( + request.provider_instance_id, + account.id(), + account.revision(), + request.upstream_model, + ), + Ok(Some(false)) + ) }) .collect::>(); + self.quota.prepare_scheduling(&accounts).await; let account_ids = accounts .iter() .map(|account| account.id().clone()) .collect::>(); - let signals = self.leases.runtime_signals(&account_ids).await?; - let round_robin_cursor = self + let scheduling = self .leases - .next_round_robin_cursor(request.provider_instance_id)?; + .load_state(request.provider_instance_id, &account_ids) + .await?; + let round_robin_cursor = scheduling.round_robin_cursor(); let candidates = accounts .into_iter() - .map(|account| AccountCandidate { - signals: signals + .map(|account| { + let signals = scheduling + .signals() .get(account.id()) .cloned() .unwrap_or(AccountRuntimeSignals { @@ -126,15 +156,28 @@ impl CodexCredentialSelector { last_started_at: None, quota_reset_at: None, quota_remaining_rank: None, - }), - account, + }) + .with_provider_quota(self.quota.scheduling_signals(&account)); + AccountCandidate { account, signals } }) .collect::>(); let mut excluded = request.attempt.excluded_accounts().clone(); - let continuation_account = request - .attempt - .continuation() - .map(|continuation| continuation.account().clone()); + let continuation_account = match request.attempt.continuation_attempt() { + ContinuationAttempt::Native => request + .attempt + .continuation() + .and_then(gateway_core::engine::continuation::ContinuationBinding::pinned) + .map(|continuation| continuation.account().clone()), + ContinuationAttempt::ReplayOwner => request + .attempt + .account_state_owner() + .filter(|owner| { + owner.provider().as_str() == PROVIDER_NAME + && owner.instance() == request.provider_instance_id + }) + .map(|owner| owner.account().clone()), + ContinuationAttempt::None | ContinuationAttempt::ReplayAny => None, + }; let required_account = request.attempt.required_account().cloned(); if required_account .as_ref() @@ -143,14 +186,15 @@ impl CodexCredentialSelector { { return Err(CredentialSelectionError::NoEligibleCredential); } - let sticky = required_account.or(continuation_account); - if let Some(required) = sticky.as_ref() { + let pinned_account = required_account.or(continuation_account); + if let Some(required) = pinned_account.as_ref() { for candidate in &candidates { if candidate.account.id() != required { excluded.insert(candidate.account.id().clone()); } } } + let sticky = pinned_account.or_else(|| scheduling.sticky_account().cloned()); let mut shortest_retry = None; loop { @@ -173,19 +217,23 @@ impl CodexCredentialSelector { let policy = request.attempt.account_selection_policy(); match self .leases - .try_acquire(CredentialLeaseRequest { - account_id: account.id().clone(), - max_concurrent: policy.max_concurrent_per_account().get(), - request_interval: policy.request_interval(), - deadline: request.attempt.deadline(), - }) + .try_acquire(ProviderLeaseRequest::Scheduling( + ProviderSchedulingLeaseRequest::new( + request.provider_instance_id.clone(), + account.id().clone(), + account.revision(), + policy.max_concurrent_per_account(), + policy.request_interval(), + request.attempt.deadline(), + ), + )) .await? { - LeaseAcquisition::Busy { retry_after } => { + ProviderLeaseAcquisition::Busy { retry_after } => { shortest_retry = minimum_duration(shortest_retry, retry_after); excluded.insert(account.id().clone()); } - LeaseAcquisition::Acquired(guard) => { + ProviderLeaseAcquisition::Acquired(guard) => { let runtime = self.repository.load_runtime_credential(&account).await?; let cookies = runtime .cookies @@ -204,11 +252,7 @@ impl CodexCredentialSelector { }) .collect(); return Ok(CodexCredentialLease { - installation_id: installation_id( - account - .upstream_account_id() - .unwrap_or(account.upstream_user_id()), - ), + installation_id: runtime.installation_id, account, secret: runtime.secret, cookies, @@ -222,36 +266,84 @@ impl CodexCredentialSelector { pub async fn record_failure( &self, lease: &CodexCredentialLease, - kind: ProviderErrorKind, - send_state: UpstreamSendState, - retry_after: Option, - ) { + failure: CodexAccountFailure, + ) -> Result<(), CredentialSelectionError> { let now = SystemTime::now(); - let (availability, reason, cooldown_until) = match kind { - ProviderErrorKind::Unauthorized => ( - AccountAvailability::Invalid, - Some("credential_rejected".to_owned()), + let (availability, reason, cooldown_until, cookie_recovery) = match failure { + CodexAccountFailure::CredentialExpired => ( + AccountAvailability::Expired, + Some("credential_expired".to_owned()), None, + CookieRecovery::None, ), - ProviderErrorKind::QuotaExhausted => ( + CodexAccountFailure::IdentityVerificationRequired => ( + AccountAvailability::Invalid, + Some("identity_verification_required".to_owned()), + None, + CookieRecovery::None, + ), + CodexAccountFailure::Banned => ( + AccountAvailability::Banned, + Some("account_banned".to_owned()), + None, + CookieRecovery::None, + ), + CodexAccountFailure::QuotaExhausted => ( AccountAvailability::QuotaExhausted, Some("quota_exhausted".to_owned()), None, + CookieRecovery::None, ), - ProviderErrorKind::RateLimited if send_state != UpstreamSendState::Ambiguous => { - let until = now.checked_add(retry_after.unwrap_or(DEFAULT_COOLDOWN)); + CodexAccountFailure::RateLimited { retry_after } => { + let until = now.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)); ( AccountAvailability::Cooldown, Some("rate_limited".to_owned()), until, + CookieRecovery::None, ) } - _ => return, + CodexAccountFailure::CloudflareChallenge { retry_after } => { + let delay = self.cloudflare_challenge_delay(lease.account_id(), now, retry_after); + let cooldown_until = now.checked_add(delay); + ( + AccountAvailability::Cooldown, + Some("cloudflare_challenge".to_owned()), + cooldown_until, + cooldown_until.map_or(CookieRecovery::Clear, CookieRecovery::ExpireAt), + ) + } + CodexAccountFailure::CloudflarePathBlocked => { + let blocked = self.record_cloudflare_path_block(lease.account_id(), now); + if blocked >= CLOUDFLARE_PATH_BLOCK_THRESHOLD { + ( + AccountAvailability::Invalid, + Some("cloudflare_path_blocked".to_owned()), + None, + CookieRecovery::Clear, + ) + } else { + ( + AccountAvailability::Cooldown, + Some("cloudflare_path_blocked".to_owned()), + now.checked_add(CLOUDFLARE_PATH_BLOCK_COOLDOWN), + CookieRecovery::Clear, + ) + } + } }; - let _ = self - .repository + self.repository .apply_state(&lease.account, availability, reason, cooldown_until, now) - .await; + .await?; + self.apply_cookie_recovery(lease, cookie_recovery).await?; + Ok(()) + } + + pub fn record_success(&self, lease: &CodexCredentialLease) { + self.risk_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(lease.account_id().as_str()); } pub async fn capture_response_cookies( @@ -307,6 +399,88 @@ impl CodexCredentialSelector { rejected: parsed.rejected, }) } + + fn cloudflare_challenge_delay( + &self, + account_id: &ProviderAccountId, + now: SystemTime, + retry_after: Option, + ) -> Duration { + let mut recovery = self + .risk_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let state = active_risk_recovery(&mut recovery, account_id.as_str(), now); + state.challenge_count = state.challenge_count.saturating_add(1); + state.observed_at = now; + let index = usize::try_from(state.challenge_count.saturating_sub(1)) + .unwrap_or(usize::MAX) + .min(CLOUDFLARE_CHALLENGE_BACKOFF.len() - 1); + retry_after + .unwrap_or_default() + .max(CLOUDFLARE_CHALLENGE_BACKOFF[index]) + } + + fn record_cloudflare_path_block(&self, account_id: &ProviderAccountId, now: SystemTime) -> u32 { + let mut recovery = self + .risk_recovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let state = active_risk_recovery(&mut recovery, account_id.as_str(), now); + state.path_block_count = state.path_block_count.saturating_add(1); + state.observed_at = now; + state.path_block_count + } + + async fn apply_cookie_recovery( + &self, + lease: &CodexCredentialLease, + recovery: CookieRecovery, + ) -> Result<(), CredentialSelectionError> { + if matches!(recovery, CookieRecovery::None) { + return Ok(()); + } + let mut data = self.repository.load_complete_data(&lease.account).await?; + if data.cookies.is_empty() { + return Ok(()); + } + match recovery { + CookieRecovery::None => return Ok(()), + CookieRecovery::ExpireAt(expires_at) => { + let expires_at = chrono::DateTime::::from(expires_at); + for cookie in &mut data.cookies { + cookie.expires_at = Some( + cookie + .expires_at + .map_or(expires_at, |current| current.min(expires_at)), + ); + } + } + CookieRecovery::Clear => data.cookies.clear(), + } + self.repository + .compare_and_swap_data(&lease.account, data) + .await?; + Ok(()) + } +} + +fn active_risk_recovery<'a>( + recovery: &'a mut HashMap, + account_id: &str, + now: SystemTime, +) -> &'a mut RiskRecoveryState { + recovery.retain(|_, state| match now.duration_since(state.observed_at) { + Ok(elapsed) => elapsed <= CLOUDFLARE_RECOVERY_STALE_AFTER, + Err(_) => true, + }); + recovery + .entry(account_id.to_owned()) + .or_insert(RiskRecoveryState { + challenge_count: 0, + path_block_count: 0, + observed_at: now, + }) } impl fmt::Debug for CodexCredentialSelector { @@ -314,7 +488,9 @@ impl fmt::Debug for CodexCredentialSelector { formatter .debug_struct("CodexCredentialSelector") .field("repository", &"ProviderAccountStore") - .field("leases", &"CredentialLeaseCoordinator") + .field("leases", &"ProviderLeasePort") + .field("catalog", &"CodexCredentialCatalogService") + .field("quota", &"CodexCredentialQuotaService") .field("cookie_policy", &self.cookie_policy) .finish() } @@ -325,7 +501,7 @@ pub struct CodexCredentialLease { secret: CodexOAuthSecret, cookies: Vec, installation_id: String, - _guard: Box, + _guard: Box, } impl CodexCredentialLease { @@ -396,8 +572,8 @@ impl From for CredentialSelectionError { } } -impl From for CredentialSelectionError { - fn from(_: CredentialLeaseCoordinatorError) -> Self { +impl From for CredentialSelectionError { + fn from(_: ProviderStoreError) -> Self { Self::Coordinator } } @@ -415,8 +591,3 @@ fn minimum_duration(current: Option, candidate: Option) -> O (None, candidate) => candidate, } } - -fn installation_id(identity: &str) -> String { - let digest = Sha256::digest(identity.as_bytes()); - format!("cpr-{}", hex::encode(&digest[..16])) -} diff --git a/backend/crates/providers/openai/src/credential/types.rs b/backend/crates/providers/openai/src/credential/types.rs index d2bdb5630..37b443337 100644 --- a/backend/crates/providers/openai/src/credential/types.rs +++ b/backend/crates/providers/openai/src/credential/types.rs @@ -29,15 +29,16 @@ impl fmt::Debug for CodexOAuthSecret { } } -/// 已由官方 JWT/OIDC 验证器确认的账号投影。 +/// 已由官方 JWT/OIDC 与认证 usage 响应共同确认的完整账号投影。 #[derive(Clone)] pub struct CodexAccountProfile { pub email: Option, + pub oauth_subject: String, + pub poid: Option, pub chatgpt_account_id: String, - pub chatgpt_user_id: Option, + pub chatgpt_user_id: String, pub plan_type: Option, pub access_token_expires_at: Option>, - pub next_refresh_at: Option>, } impl fmt::Debug for CodexAccountProfile { @@ -45,14 +46,30 @@ impl fmt::Debug for CodexAccountProfile { formatter .debug_struct("CodexAccountProfile") .field("email", &self.email.as_ref().map(|_| "")) + .field("oauth_subject", &"") + .field("poid", &self.poid.as_ref().map(|_| "")) .field("chatgpt_account_id", &"") - .field( - "chatgpt_user_id", - &self.chatgpt_user_id.as_ref().map(|_| ""), - ) + .field("chatgpt_user_id", &"") .field("plan_type", &self.plan_type) .field("access_token_expires_at", &self.access_token_expires_at) - .field("next_refresh_at", &self.next_refresh_at) + .finish() + } +} + +/// 持久化在 Provider credential JSON 中的签名认证主体。 +#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct CodexCredentialPrincipal { + pub oauth_subject: String, + pub poid: Option, +} + +impl fmt::Debug for CodexCredentialPrincipal { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CodexCredentialPrincipal") + .field("oauth_subject", &"") + .field("poid", &self.poid.as_ref().map(|_| "")) .finish() } } @@ -89,6 +106,8 @@ impl fmt::Debug for CodexCookie { #[serde(deny_unknown_fields)] pub struct CodexCredentialData { pub schema_version: u32, + pub principal: CodexCredentialPrincipal, + pub installation_id: String, pub access_token: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub refresh_token: Option, @@ -107,6 +126,8 @@ impl fmt::Debug for CodexCredentialData { formatter .debug_struct("CodexCredentialData") .field("schema_version", &self.schema_version) + .field("principal", &self.principal) + .field("installation_id", &"") .field("access_token", &"") .field( "refresh_token", @@ -127,6 +148,7 @@ pub struct CreateCodexCredential { pub name: String, pub secret: CodexOAuthSecret, pub account: CodexAccountProfile, + pub next_refresh_at: Option>, pub enabled: bool, } @@ -139,6 +161,7 @@ impl fmt::Debug for CreateCodexCredential { .field("name", &self.name) .field("secret", &self.secret) .field("account", &self.account) + .field("next_refresh_at", &self.next_refresh_at) .field("enabled", &self.enabled) .finish() } @@ -150,6 +173,7 @@ pub struct RotateCodexCredential { pub expected_credential_revision: u64, pub secret: CodexOAuthSecret, pub verified_account: CodexAccountProfile, + pub next_refresh_at: Option>, } impl fmt::Debug for RotateCodexCredential { @@ -163,6 +187,7 @@ impl fmt::Debug for RotateCodexCredential { ) .field("secret", &self.secret) .field("verified_account", &self.verified_account) + .field("next_refresh_at", &self.next_refresh_at) .finish() } } diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index b2e912c4b..a1b80ca0c 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -1,16 +1,223 @@ -//! Codex Provider 专属能力。 +//! OpenAI Provider 专属能力。 +mod admin; +pub mod config; mod provider; +use std::path::Path; +use std::sync::Arc; + +use gateway_admin::ports::provider::ProviderAdmin; +use gateway_core::engine::credential::ProviderAccountStore; +use gateway_core::engine::provider::Provider; +use gateway_core::provider_ports::ProviderStorePorts; +use gateway_core::routing::ProviderKind; +use gateway_core::task::WorkerContribution; + +use crate::admin::{OpenAiAdminProvider, OpenAiAdminServices, OpenAiOAuthPendingStore}; +use crate::credential::token_client::{AuthorizationCodeExchanger, TokenRefresher}; +use crate::credential::{ + CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexAuthenticatedAccountSource, + CodexCookiePolicy, CodexCredentialAdmin, CodexCredentialAdminService, + CodexCredentialCatalogService, CodexCredentialQuotaService, CodexCredentialRefreshService, + CodexCredentialRepository, CodexCredentialSelector, CodexJwtIdentityVerifier, CodexOAuthAdmin, + CodexOAuthAdminService, CodexSignedIdentityVerifier, ReqwestCodexAuthenticatedAccountSource, + ReqwestOpenAiJwksSource, +}; +use crate::transport::profile::{CodexDesktopReleaseService, OfficialCodexDesktopReleaseTransport}; +use crate::transport::{CodexWebSocketPool, build_reqwest_client}; + +pub use config::{CodexWireProfileConfig, OpenAiConfig, OpenAiConfigError}; pub use provider::{ - CodexEndpointPolicy, CodexProvider, CodexProviderConfigError, CodexProviderInstanceConfig, - CodexProviderTransport, OFFICIAL_CODEX_BASE_PATH, + CodexOriginPolicy, CodexProvider, CodexProviderConfigError, CodexProviderInstanceConfig, + CodexProviderTransport, OFFICIAL_CODEX_BASE_PATH, OfficialCodexOriginPolicy, }; pub mod credential; pub mod transport; pub use transport::{ - CodexCanonicalDecoder, CodexRequestEncodeError, codex_request_semantics, + CodexCanonicalDecoder, CodexCanonicalError, CodexRequestEncodeError, codex_request_semantics, encode_generate_request, openai_billing_breakdown, }; + +/// OpenAI 初始化后交给组装根的最小能力集。 +pub struct ProviderBundle { + core_provider: Arc, + admin_provider: Arc, + worker_contributions: Vec, +} + +/// 构造 OpenAI 数据面、Provider-owned 后台任务与 Redis OAuth pending owner。 +pub async fn initialize( + mut config: OpenAiConfig, + ports: ProviderStorePorts, +) -> Result { + config + .resolve_and_validate(Path::new(".")) + .map_err(OpenAiInitializeError::Config)?; + let provider_kind = + ProviderKind::new("openai").map_err(|_| OpenAiInitializeError::InvalidProviderKind)?; + let accounts: Arc = ports.accounts(); + let leases = ports.leases(); + let runtime_policy = ports.runtime_policy(); + let instances = ports.instances(); + let profile = config.wire_profile_state(); + let http = build_reqwest_client().map_err(|_| OpenAiInitializeError::Transport)?; + let origin_policy: Arc = Arc::new(OfficialCodexOriginPolicy); + let desktop_release = Arc::new(CodexDesktopReleaseService::new( + profile.clone(), + Arc::new( + OfficialCodexDesktopReleaseTransport::new() + .map_err(|_| OpenAiInitializeError::DesktopRelease)?, + ), + )); + let desktop_release_status = desktop_release.status(); + let repository = CodexCredentialRepository::new(Arc::clone(&accounts)); + let catalog = Arc::new(CodexCredentialCatalogService::new( + repository.clone(), + profile.clone(), + http.clone(), + Arc::clone(&origin_policy), + )); + let quota = Arc::new(CodexCredentialQuotaService::new( + repository.clone(), + profile.clone(), + http.clone(), + Arc::clone(&origin_policy), + )); + let selector = Arc::new(CodexCredentialSelector::new( + repository.clone(), + Arc::clone(&leases), + Arc::clone(&catalog), + Arc::clone("a), + CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?, + )); + let websocket_pool = Arc::new(CodexWebSocketPool::default()); + let core_provider: Arc = Arc::new(CodexProvider::new( + selector, + Arc::clone(&catalog), + Arc::clone("a), + http, + profile.clone(), + Arc::clone(&websocket_pool), + origin_policy, + )); + + let token_client = Arc::new( + credential::token_client::official_openai_token_client() + .map_err(|_| OpenAiInitializeError::TokenClient)?, + ); + let jwks = ReqwestOpenAiJwksSource::new().map_err(|_| OpenAiInitializeError::Identity)?; + let signed: Arc = + Arc::new(CodexJwtIdentityVerifier::new(Box::new(jwks))); + let account_source: Arc = Arc::new( + ReqwestCodexAuthenticatedAccountSource::new(profile.clone()) + .map_err(|_| OpenAiInitializeError::Identity)?, + ); + let identity: Arc = + Arc::new(CodexAccountIdentityService::new(signed, account_source)); + let refresher: Arc = token_client.clone(); + let exchanger: Arc = token_client; + let credential_admin = Arc::new(CodexCredentialAdminService::new( + repository.clone(), + Arc::clone(&refresher), + Arc::clone(&identity), + Arc::clone(&leases), + Arc::clone(&runtime_policy), + )); + let refresh = Arc::new(CodexCredentialRefreshService::new( + repository, + refresher, + Arc::clone(&identity), + Arc::clone(&leases), + Arc::clone(&runtime_policy), + )); + let pending = Arc::new(OpenAiOAuthPendingStore::new( + ports.oauth_pending(), + provider_kind.clone(), + )); + let oauth_admin: Arc = Arc::new(CodexOAuthAdminService::new( + pending, + exchanger, + Arc::clone(&identity), + Arc::clone(&accounts), + Arc::clone(&runtime_policy), + CodexCredentialAdmin, + )); + let admin_provider: Arc = Arc::new(OpenAiAdminProvider::new( + provider_kind.clone(), + profile, + accounts, + Arc::clone(&instances), + OpenAiAdminServices { + credentials: credential_admin, + verifier: identity, + oauth: oauth_admin, + quota: Arc::clone("a), + catalog: Arc::clone(&catalog), + runtime_policy, + }, + websocket_pool, + desktop_release_status, + )); + let worker_contributions = provider::worker_contributions( + refresh, + quota, + catalog, + desktop_release, + instances, + provider_kind, + ) + .map_err(|_| OpenAiInitializeError::Worker)?; + + Ok(ProviderBundle { + core_provider, + admin_provider, + worker_contributions, + }) +} + +impl ProviderBundle { + #[must_use] + pub fn core_provider(&self) -> Arc { + Arc::clone(&self.core_provider) + } + + #[must_use] + pub fn admin_provider(&self) -> Arc { + Arc::clone(&self.admin_provider) + } + + /// 一次性移交 Host 任务计划，防止同一 owner 被重复注册。 + pub fn take_worker_contributions(&mut self) -> Vec { + std::mem::take(&mut self.worker_contributions) + } +} + +/// OpenAI 初始化失败的脱敏分类。 +#[derive(Debug, thiserror::Error)] +pub enum OpenAiInitializeError { + #[error(transparent)] + Config(OpenAiConfigError), + #[error("OpenAI runtime policy is unavailable")] + RuntimePolicy, + #[error("OpenAI Provider kind is invalid")] + InvalidProviderKind, + #[error("OpenAI transport could not initialize")] + Transport, + #[error("OpenAI cookie policy could not initialize")] + CookiePolicy, + #[error("OpenAI token client could not initialize")] + TokenClient, + #[error("OpenAI identity verifier could not initialize")] + Identity, + #[error("OpenAI credential administration could not initialize")] + CredentialAdmin, + #[error("OpenAI credential refresh could not initialize")] + Refresh, + #[error("OpenAI Desktop release service could not initialize")] + DesktopRelease, + #[error("OpenAI worker plan is invalid")] + Worker, +} diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 12032451b..9b372b584 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -1,40 +1,70 @@ //! Codex 的 `gateway-core` Provider adapter。 -use std::collections::BTreeSet; +use std::collections::{BTreeSet, HashMap}; use std::fmt; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use async_trait::async_trait; -use futures::StreamExt; +use futures::{StreamExt, future::BoxFuture}; +use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope}; use gateway_core::engine::provider::{ - EventStream, Provider, ProviderCallMetadata, ProviderModelCapabilities, ProviderRequest, - ProviderResource, ProviderStream, UpstreamTransport, + EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, + ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, ProviderResource, + ProviderStream, UpstreamTransport, }; -use gateway_core::engine::{AttemptContext, UpstreamSendState}; -use gateway_core::error::{ProviderError, ProviderErrorKind}; -use gateway_core::event::GatewayEvent; -use gateway_core::operation::{Feature, Operation, OperationKind}; +use gateway_core::engine::{AttemptContext, ContinuationAttempt, UpstreamSendState}; +use gateway_core::error::{ + ContinuationFailure, ProviderError, ProviderErrorKind, SafeUpstreamValue, +}; +use gateway_core::event::{ + GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, + ProviderResponseTimings, UpstreamHttpVersion, +}; +use gateway_core::operation::{ + Feature, GenerateRequest, Operation, OperationKind, ProviderSessionState, +}; +use gateway_core::provider_ports::ProviderInstanceCatalogPort; use gateway_core::routing::{ - ModelCapabilities, ProviderInstance, ProviderInstanceId, ProviderKind, SupportLevel, - UpstreamModelId, + InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId, ProviderKind, + SupportLevel, UpstreamModelId, }; -use reqwest::{Client, StatusCode}; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration, + WorkerRunnable, WorkerSchedule, WorkerTaskError, +}; +use reqwest::Client; use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; use url::{Host, Url}; use crate::credential::{ - CodexCredentialCatalogService, CodexCredentialLease, CodexCredentialSelector, - CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, + CodexAccountFailure, CodexCredentialCatalogService, CodexCredentialLease, + CodexCredentialQuotaService, CodexCredentialRefreshOutcome, CodexCredentialRefreshService, + CodexCredentialSelector, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, +}; +use crate::transport::canonical::{ + CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, }; -use crate::transport::canonical::CodexCanonicalDecoder; use crate::transport::catalog::{CodexCatalogCapabilityEvidence, CodexCatalogModel}; -use crate::transport::profile::CodexWireProfileState; +use crate::transport::diagnostics::{ + CodexFailureCategory, CodexUpstreamFailure, CodexUpstreamSendPhase, +}; +use crate::transport::profile::{ + APPCAST_POLL_INTERVAL, CodexDesktopReleaseService, CodexWireProfileState, +}; use crate::transport::protocol::responses::{CodexResponsesRequest, PreviousResponseScope}; -use crate::transport::request::{CodexRequestEncodeError, encode_generate_request}; +use crate::transport::request::{ + CodexRequestEncodeError, encode_generate_request, sanitize_cross_account_item, + scope_request_to_account, +}; +use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason}; use crate::transport::{ CODEX_RESPONSES_PATH, CodexBackendClient, CodexBackendTransport, CodexClientError, - CodexRequestContext, CodexWebSocketPool, build_reqwest_client, endpoint_url, + CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics, + CodexWebSocketPool, endpoint_url, }; const PROVIDER_NAME: &str = "openai"; @@ -46,24 +76,37 @@ pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexProviderTransport { - HttpSse, - WebSocket, + HttpOnly, + PreferWebSocket, } impl CodexProviderTransport { - fn parse(value: &str) -> Option { + fn parse_explicit(value: &str) -> Option { match value { - HTTP_SSE_TRANSPORT => Some(Self::HttpSse), - WEBSOCKET_TRANSPORT => Some(Self::WebSocket), + HTTP_SSE_TRANSPORT => Some(Self::HttpOnly), + WEBSOCKET_TRANSPORT => Some(Self::PreferWebSocket), _ => None, } } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CodexEndpointPolicy { - Official, - Loopback, +/// Provider base URL 的显式信任边界。 +/// +/// 生产组装注入 [`OfficialCodexOriginPolicy`]；Provider 只依赖这个端口， +/// 不包含部署环境或地址类型分支。 +pub trait CodexOriginPolicy: Send + Sync { + fn allows(&self, url: &Url) -> bool; +} + +#[derive(Debug, Default)] +pub struct OfficialCodexOriginPolicy; + +impl CodexOriginPolicy for OfficialCodexOriginPolicy { + fn allows(&self, url: &Url) -> bool { + url.scheme() == "https" + && matches!(url.host(), Some(Host::Domain(host)) if host == OFFICIAL_CODEX_HOST) + && url.port().is_none() + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -74,24 +117,20 @@ pub struct CodexProviderInstanceConfig { } impl CodexProviderInstanceConfig { - pub fn from_snapshot(instance: &ProviderInstance) -> Result { - Self::from_snapshot_with_policy(instance, CodexEndpointPolicy::Official) - } - - pub fn from_snapshot_with_policy( + pub fn from_snapshot( instance: &ProviderInstance, - endpoint_policy: CodexEndpointPolicy, + origin_policy: &dyn CodexOriginPolicy, ) -> Result { if instance.provider().as_str() != PROVIDER_NAME { return Err(CodexProviderConfigError::ProviderMismatch); } let mut base_url = Url::parse(instance.base_url()) .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; - normalize_and_validate_base_url(&mut base_url, endpoint_policy)?; + normalize_and_validate_base_url(&mut base_url, origin_policy)?; Ok(Self { id: instance.id().clone(), base_url, - transport: CodexProviderTransport::HttpSse, + transport: CodexProviderTransport::PreferWebSocket, }) } @@ -119,8 +158,6 @@ pub enum CodexProviderConfigError { InvalidBaseUrl, #[error("Codex provider base URL is not allowed")] UnsafeBaseUrl, - #[error("Codex provider HTTP client initialization failed")] - TransportInitialization, } #[derive(Clone)] @@ -133,45 +170,40 @@ struct CompiledInstance { pub struct CodexProvider { selector: Arc, catalog: Arc, + quota: Arc, http: Client, profile: CodexWireProfileState, - endpoint_policy: CodexEndpointPolicy, + origin_policy: Arc, websocket_pool: Arc, + compiled_instances: Mutex>, } impl CodexProvider { pub fn new( selector: Arc, catalog: Arc, + quota: Arc, + http: Client, profile: CodexWireProfileState, - ) -> Result { - let http = build_reqwest_client() - .map_err(|_| CodexProviderConfigError::TransportInitialization)?; - Ok(Self { + websocket_pool: Arc, + origin_policy: Arc, + ) -> Self { + Self { selector, catalog, + quota, http, profile, - endpoint_policy: CodexEndpointPolicy::Official, - websocket_pool: Arc::new(CodexWebSocketPool::default()), - }) - } - - pub fn new_with_policy( - selector: Arc, - catalog: Arc, - profile: CodexWireProfileState, - endpoint_policy: CodexEndpointPolicy, - ) -> Result { - let mut provider = Self::new(selector, catalog, profile)?; - provider.endpoint_policy = endpoint_policy; - Ok(provider) + origin_policy, + websocket_pool, + compiled_instances: Mutex::new(HashMap::new()), + } } pub fn validate_instance( instance: &ProviderInstance, ) -> Result { - CodexProviderInstanceConfig::from_snapshot(instance) + CodexProviderInstanceConfig::from_snapshot(instance, &OfficialCodexOriginPolicy) } fn compile_instance( @@ -179,14 +211,23 @@ impl CodexProvider { snapshot: &ProviderInstance, ) -> Result { let config = - CodexProviderInstanceConfig::from_snapshot_with_policy(snapshot, self.endpoint_policy) + CodexProviderInstanceConfig::from_snapshot(snapshot, self.origin_policy.as_ref()) .map_err(map_instance_config_error)?; let responses_url = Url::parse(&endpoint_url( config.base_url.as_str(), CODEX_RESPONSES_PATH, )) .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - Ok(CompiledInstance { + let mut instances = self + .compiled_instances + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(compiled) = instances.get(config.id().as_str()) + && compiled.config == config + { + return Ok(compiled.clone()); + } + let compiled = CompiledInstance { client: CodexBackendClient::new( self.http.clone(), config.base_url.as_str(), @@ -195,7 +236,9 @@ impl CodexProvider { .with_websocket_pool(Arc::clone(&self.websocket_pool)), config, responses_url, - }) + }; + instances.insert(compiled.config.id().as_str().to_owned(), compiled.clone()); + Ok(compiled) } } @@ -205,7 +248,7 @@ impl fmt::Debug for CodexProvider { .debug_struct("CodexProvider") .field("selector", &"") .field("catalog", &"") - .field("endpoint_policy", &self.endpoint_policy) + .field("origin_policy", &"[ORIGIN_POLICY]") .finish() } } @@ -216,6 +259,37 @@ impl Provider for CodexProvider { PROVIDER_NAME } + fn catalog_generation(&self) -> ProviderCatalogGeneration { + self.catalog.catalog_generation() + } + + fn request_observation(&self, operation: &Operation) -> ProviderRequestObservation { + let Operation::Generate(request) = operation else { + return ProviderRequestObservation::default(); + }; + let semantics = encode_generate_request(request, "observability") + .map(|mut encoded| { + if let Ok(Some(previous)) = decode_openai_session_state(request) { + let mut input = previous + .transcript + .iter() + .map(OpenAiReplayItem::value) + .cloned() + .collect::>(); + input.extend(encoded.input().iter().cloned()); + encoded.set_input(input); + } + encoded.semantics() + }) + .unwrap_or_default(); + ProviderRequestObservation { + reasoning_preset: semantics.reasoning_preset.map(str::to_owned), + request_kind: semantics.request_kind, + subagent_kind: semantics.subagent_kind, + compact: semantics.compact, + } + } + async fn query_model_capabilities( &self, instance: &ProviderInstance, @@ -272,34 +346,109 @@ impl Provider for CodexProvider { UpstreamSendState::NotSent, )); }; + let previous_session = decode_openai_session_state(generate)?; + let continuation_requested = generate.continuation().is_some(); let mut upstream_request = encode_generate_request(generate, candidate.upstream_model().as_str()) .map_err(map_request_error)?; - let transport = - selected_transport(&request, &instance.config, context.continuation().is_some())?; - if let Some(continuation) = context.continuation() { - upstream_request.set_previous_response_id(Some( - continuation.upstream_response_id().as_str().to_owned(), - )); - // Codex 请求统一 store=false；原生 handle 只能在仍持有它的连接上续接。 - upstream_request.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); - } + let request_input = upstream_request.input().to_vec(); + let transport = selected_transport(&request, &instance.config)?; apply_transport(&mut upstream_request, transport); let lease = self .selector .select(&SelectCodexCredential { provider_instance_id: candidate.instance(), + upstream_model: candidate.upstream_model().as_str(), request_url: &instance.responses_url, attempt: &context, }) .await .map_err(map_selection_error)?; let lease = Arc::new(lease); + let provider_kind = ProviderKind::new(PROVIDER_NAME) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; + let cross_account = context.account_state_owner().is_some_and(|owner| { + !owner.matches(&provider_kind, candidate.instance(), lease.account_id()) + }) || previous_session + .as_ref() + .is_some_and(|state| state.account_id != lease.account_id().as_str()); + let replay_previous_response = matches!( + context.continuation_attempt(), + ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny + ) || previous_session.as_ref().is_some_and(|state| { + state.continuation_scope == OpenAiContinuationScope::ReplayRequired + }); + if replay_previous_response { + let state = previous_session.as_ref().ok_or_else(|| { + provider_error( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + ) + })?; + let mut input = replay_input_for_account(state, lease.account_id().as_str()); + input.reserve(request_input.len()); + input.extend(request_input.iter().cloned()); + upstream_request.set_input(input); + upstream_request.set_previous_response_id(None); + upstream_request.previous_response_scope = None; + upstream_request.turn_state = None; + } + scope_request_to_account( + &mut upstream_request, + lease.installation_id(), + cross_account, + ); + if let Some(conversation_id) = previous_session + .as_ref() + .and_then(|state| state.conversation_id.as_ref()) + { + upstream_request.local_conversation_id = Some(conversation_id.clone()); + } + if context.continuation_attempt() == ContinuationAttempt::Native + && !replay_previous_response + && let Some(continuation) = context.continuation() + { + match continuation { + ContinuationBinding::Pinned(continuation) => { + let previous_response_scope = match previous_session + .as_ref() + .map(|state| state.continuation_scope) + { + Some(OpenAiContinuationScope::Persisted) => { + PreviousResponseScope::Persisted + } + Some(OpenAiContinuationScope::ConnectionLocal) => { + PreviousResponseScope::ConnectionLocal + } + Some(OpenAiContinuationScope::ReplayRequired) => { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )); + } + None => match continuation.scope() { + NativeContinuationScope::Persisted => PreviousResponseScope::Persisted, + NativeContinuationScope::ConnectionLocal => { + PreviousResponseScope::ConnectionLocal + } + }, + }; + upstream_request.set_previous_response_id(Some( + continuation.upstream_response_id().as_str().to_owned(), + )); + upstream_request.previous_response_scope = Some(previous_response_scope); + } + ContinuationBinding::External(previous_response_id) => { + upstream_request + .set_previous_response_id(Some(previous_response_id.as_str().to_owned())); + upstream_request.previous_response_scope = + Some(PreviousResponseScope::ExternalUnknown); + } + } + } let metadata = ProviderCallMetadata::new( - ProviderKind::new(PROVIDER_NAME).map_err(|_| { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) - })?, + provider_kind, candidate.instance().clone(), candidate.upstream_model().clone(), ProviderResource::Account { @@ -310,15 +459,29 @@ impl Provider for CodexProvider { provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) })?, ); + let response_store = upstream_request.store(); + let session_capture = + (!continuation_requested || previous_session.is_some()).then(|| OpenAiSessionCapture { + previous: previous_session, + request_input, + account_id: lease.account_id().as_str().to_owned(), + conversation_id: upstream_request.local_conversation_id.clone(), + response_store, + continuation_scope: None, + }); let events = cold_response_stream(ColdResponse { client: instance.client, response_origin: instance.responses_url, request: upstream_request, upstream_model: candidate.upstream_model().clone(), - expected_transport: transport, + transport_policy: transport, context, selector: Arc::clone(&self.selector), + quota: Arc::clone(&self.quota), + catalog: Arc::clone(&self.catalog), + provider_instance_id: candidate.instance().clone(), lease: Arc::clone(&lease), + session_capture, }); Ok(ProviderStream::new(metadata, events, lease)) } @@ -329,10 +492,190 @@ struct ColdResponse { response_origin: Url, request: CodexResponsesRequest, upstream_model: UpstreamModelId, - expected_transport: CodexProviderTransport, + transport_policy: CodexProviderTransport, context: AttemptContext, selector: Arc, + quota: Arc, + catalog: Arc, + provider_instance_id: ProviderInstanceId, lease: Arc, + session_capture: Option, +} + +#[derive(Clone, Serialize, Deserialize)] +struct OpenAiSessionState { + account_id: String, + conversation_id: Option, + continuation_scope: OpenAiContinuationScope, + transcript: Vec, +} + +#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum OpenAiContinuationScope { + Persisted, + ConnectionLocal, + ReplayRequired, +} + +#[derive(Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum OpenAiReplayItem { + ClientInput(Value), + SanitizedOutput(Value), + AccountOutput { account_id: String, item: Value }, +} + +impl OpenAiReplayItem { + fn value(&self) -> &Value { + match self { + Self::ClientInput(value) + | Self::SanitizedOutput(value) + | Self::AccountOutput { item: value, .. } => value, + } + } +} + +struct OpenAiSessionCapture { + previous: Option, + request_input: Vec, + account_id: String, + conversation_id: Option, + response_store: bool, + continuation_scope: Option, +} + +fn decode_openai_session_state( + request: &GenerateRequest, +) -> Result, ProviderError> { + request + .provider_session_state(PROVIDER_NAME) + .map(|state| { + serde_json::from_value(Value::Object(state.payload().clone())).map_err(|_| { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) + }) + }) + .transpose() +} + +fn encode_openai_session_state( + state: OpenAiSessionState, +) -> Result { + let Value::Object(payload) = serde_json::to_value(state) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))? + else { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )); + }; + ProviderSessionState::new(PROVIDER_NAME, payload) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)) +} + +fn attach_openai_session_update( + events: &mut [ProviderEvent], + capture: &mut Option, +) -> Result<(), ProviderError> { + let Some((terminal_index, output)) = events.iter().enumerate().find_map(|(index, event)| { + terminal_response_output(event).map(|output| (index, output.to_vec())) + }) else { + return Ok(()); + }; + let Some(capture) = capture.take() else { + return Ok(()); + }; + let mut transcript = capture + .previous + .map(|state| state.transcript) + .unwrap_or_default(); + project_transcript_to_account(&mut transcript, &capture.account_id); + transcript.extend( + capture + .request_input + .into_iter() + .map(OpenAiReplayItem::ClientInput), + ); + transcript.extend( + output + .into_iter() + .map(|item| OpenAiReplayItem::AccountOutput { + account_id: capture.account_id.clone(), + item, + }), + ); + let update = encode_openai_session_state(OpenAiSessionState { + account_id: capture.account_id, + conversation_id: capture.conversation_id, + continuation_scope: capture + .continuation_scope + .ok_or_else(|| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?, + transcript, + })?; + events[terminal_index].attach_session_update(update); + Ok(()) +} + +fn terminal_response_output(event: &ProviderEvent) -> Option<&[Value]> { + let wire = event.wire_event()?; + if wire.protocol() != PROVIDER_NAME { + return None; + } + let event_type = wire + .event_type() + .or_else(|| wire.data().get("type").and_then(Value::as_str)); + matches!( + event_type, + Some("response.completed" | "response.incomplete") + ) + .then(|| { + wire.data() + .pointer("/response/output") + .and_then(Value::as_array) + .map(Vec::as_slice) + }) + .flatten() +} + +fn project_transcript_to_account(transcript: &mut Vec, account_id: &str) { + *transcript = transcript + .drain(..) + .filter_map(|item| match item { + OpenAiReplayItem::AccountOutput { + account_id: owner, + item, + } if owner != account_id => { + sanitize_cross_account_item(item).map(OpenAiReplayItem::SanitizedOutput) + } + item => Some(item), + }) + .collect(); +} + +fn replay_input_for_account(state: &OpenAiSessionState, account_id: &str) -> Vec { + state + .transcript + .iter() + .filter_map(|item| match item { + OpenAiReplayItem::ClientInput(value) | OpenAiReplayItem::SanitizedOutput(value) => { + Some(value.clone()) + } + OpenAiReplayItem::AccountOutput { + account_id: owner, + item, + } if owner == account_id => Some(without_output_id(item.clone())), + OpenAiReplayItem::AccountOutput { item, .. } => { + sanitize_cross_account_item(item.clone()) + } + }) + .collect() +} + +fn without_output_id(mut item: Value) -> Value { + if let Value::Object(object) = &mut item { + object.remove("id"); + } + item } fn cold_response_stream(response: ColdResponse) -> EventStream { @@ -341,10 +684,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { response_origin, request, upstream_model, - expected_transport, + transport_policy, context, selector, + quota, + catalog, + provider_instance_id, lease, + mut session_capture, } = response; Box::pin(async_stream::try_stream! { let cookie_header = build_cookie_header(lease.cookies())?; @@ -362,14 +709,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { ); let response = tokio::select! { biased; - _ = cancellation.cancelled() => Err(provider_error( + _ = cancellation.cancelled() => Err(MappedProviderFailure::plain(provider_error( ProviderErrorKind::Cancelled, UpstreamSendState::Ambiguous, - )), - _ = tokio::time::sleep(handshake_deadline) => Err(provider_error( + ))), + _ = tokio::time::sleep(handshake_deadline) => Err(MappedProviderFailure::plain(provider_error( ProviderErrorKind::Timeout, UpstreamSendState::Ambiguous, - )), + ))), response = client.create_response_stream_with_pool_account( &request, request_context, @@ -378,23 +725,63 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { }; let response = match response { Ok(response) => response, - Err(error) => { - selector - .record_failure(&lease, error.kind(), error.send_state(), error.retry_after()) - .await; - Err(error)?; + Err(mut failure) => { + if let Some(observation) = failure.observation.take() { + yield ProviderEvent::observation(observation); + } + apply_failure( + &client, + &selector, + "a, + &lease, + &response_origin, + &failure, + ) + .await; + Err(failure.error)?; return; } }; - if response.transport != backend_transport(expected_transport) { + if !accepts_backend_transport(transport_policy, response.transport) { Err(provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?; } + if let Some(capture) = session_capture.as_mut() { + capture.continuation_scope = Some(if capture.response_store { + OpenAiContinuationScope::Persisted + } else if response.transport == CodexBackendTransport::WebSocket + && response.connection_local_continuation + { + OpenAiContinuationScope::ConnectionLocal + } else { + OpenAiContinuationScope::ReplayRequired + }); + } + yield ProviderEvent::observation(codex_response_observation( + response.transport, + &response.diagnostics, + &response.response_metadata, + &response.transport_metrics, + )?); + synchronize_passive_quota("a, &lease, &response.rate_limit_headers).await; + if let Some(etag) = response.response_metadata.models_etag.as_deref() + && let Err(error) = catalog.observe_response_etag(&provider_instance_id, etag) + { + tracing::warn!( + provider_instance_id = %provider_instance_id, + error = %error, + "OpenAI model ETag observation was rejected" + ); + } if !response.set_cookie_headers.is_empty() { let _ = selector .capture_response_cookies(&lease, &response_origin, &response.set_cookie_headers) .await; } let mut body = response.body; + let failure_diagnostics = response.diagnostics.clone(); + let failure_set_cookie_headers = response.set_cookie_headers.clone(); + let failure_rate_limit_headers = response.rate_limit_headers.clone(); + let rate_limit_updates = response.rate_limit_header_updates; let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str()); loop { let Some(stream_deadline) = remaining(context.deadline()) else { @@ -403,14 +790,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { }; let next = tokio::select! { biased; - _ = cancellation.cancelled() => Err(provider_error( + _ = cancellation.cancelled() => Err(MappedProviderFailure::plain(provider_error( ProviderErrorKind::Cancelled, UpstreamSendState::Sent, - )), - _ = tokio::time::sleep(stream_deadline) => Err(provider_error( + ))), + _ = tokio::time::sleep(stream_deadline) => Err(MappedProviderFailure::plain(provider_error( ProviderErrorKind::Timeout, UpstreamSendState::Sent, - )), + ))), chunk = body.next() => match chunk { Some(Ok(chunk)) => Ok(Some(chunk)), Some(Err(error)) => Err(map_stream_error(error)), @@ -419,32 +806,220 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { }; let next = match next { Ok(next) => next, - Err(error) => { - selector - .record_failure(&lease, error.kind(), error.send_state(), error.retry_after()) - .await; - Err(error)?; + Err(failure) => { + apply_failure( + &client, + &selector, + "a, + &lease, + &response_origin, + &failure, + ) + .await; + Err(failure.error)?; return; } }; let Some(chunk) = next else { break; }; - let events = decoder.push(&chunk)?; + if let Some(updates) = rate_limit_updates.as_ref() { + let updates = std::mem::take(&mut *updates.lock().await); + synchronize_passive_quota("a, &lease, &updates).await; + } + let (mut events, terminal_error) = match decoder.push(&chunk) { + CodexCanonicalOutcome::Events(events) => (events, None), + CodexCanonicalOutcome::Failed(failure) => { + let (events, error, semantic_output_seen) = failure.into_parts(); + let failure = map_canonical_error( + error, + &failure_diagnostics, + &failure_set_cookie_headers, + &failure_rate_limit_headers, + ReplayBoundary::from_semantic_output(semantic_output_seen), + ); + apply_failure( + &client, + &selector, + "a, + &lease, + &response_origin, + &failure, + ) + .await; + (events, Some(failure.error)) + } + }; + attach_openai_session_update(&mut events, &mut session_capture)?; let completed = events .iter() + .flat_map(ProviderEvent::canonical_facts) .any(|event| matches!(event, GatewayEvent::Completed(_))); for event in events { yield event; } + if let Some(error) = terminal_error { + Err(error)?; + return; + } if completed { + selector.record_success(&lease); return; } } - for event in decoder.finish()? { + let (mut events, terminal_error) = match decoder.finish() { + CodexCanonicalOutcome::Events(events) => (events, None), + CodexCanonicalOutcome::Failed(failure) => { + let (events, error, semantic_output_seen) = failure.into_parts(); + let failure = map_canonical_error( + error, + &failure_diagnostics, + &failure_set_cookie_headers, + &failure_rate_limit_headers, + ReplayBoundary::from_semantic_output(semantic_output_seen), + ); + apply_failure( + &client, + &selector, + "a, + &lease, + &response_origin, + &failure, + ) + .await; + (events, Some(failure.error)) + } + }; + attach_openai_session_update(&mut events, &mut session_capture)?; + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .any(|event| matches!(event, GatewayEvent::Completed(_))); + for event in events { yield event; } + if let Some(error) = terminal_error { + Err(error)?; + return; + } + if completed { + selector.record_success(&lease); + } }) } +fn codex_response_observation( + transport: CodexBackendTransport, + diagnostics: &CodexUpstreamDiagnostics, + response_metadata: &CodexResponseMetadata, + metrics: &CodexTransportMetrics, +) -> Result { + let mut observation = ProviderResponseObservation::new( + UpstreamTransport::new(actual_transport_name(transport)) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent))?, + ) + .with_timings(ProviderResponseTimings { + transport_decision_wait_ms: nonnegative_millis(metrics.transport_decision_wait_ms), + connect_ms: nonnegative_millis(metrics.ws_connect_ms), + headers_ms: nonnegative_millis(metrics.upstream_headers_ms), + first_event_ms: nonnegative_millis(metrics.first_event_ms), + }); + if let Some(version) = metrics + .http_version + .as_deref() + .and_then(UpstreamHttpVersion::parse) + { + observation = observation.with_http_version(version); + } + if let Some(status_code) = diagnostics.status_code { + observation = observation.with_status_code(status_code); + } + if let Some(request_id) = diagnostics + .request_id + .as_deref() + .and_then(|request_id| SafeUpstreamValue::new(request_id.to_owned()).ok()) + { + observation = observation.with_request_id(request_id); + } + let client_headers = response_metadata + .client_headers + .iter() + .filter_map(|(name, value)| { + let value = SafeUpstreamValue::new(value.to_owned()).ok()?; + ProviderResponseHeader::new(name.to_owned(), value) + }) + .collect(); + observation = observation.with_client_headers(client_headers); + Ok(observation) +} + +fn codex_error_observation(error: &CodexClientError) -> Option { + let transport = error.transport()?; + let mut observation = ProviderResponseObservation::new( + UpstreamTransport::new(actual_transport_name(transport)).ok()?, + ); + match error { + CodexClientError::Upstream { + status, + diagnostics, + transport, + transport_metrics, + .. + } => { + observation = codex_response_observation( + *transport, + diagnostics, + &CodexResponseMetadata::default(), + transport_metrics, + ) + .ok()? + .with_status_code(status.as_u16()); + } + CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => { + observation = observation.with_status_code(upstream.status_code); + if let Some(request_id) = upstream + .diagnostics + .request_id + .as_deref() + .and_then(|value| SafeUpstreamValue::new(value.to_owned()).ok()) + { + observation = observation.with_request_id(request_id); + } + } + _ => {} + } + Some(observation) +} + +async fn synchronize_passive_quota( + quota: &CodexCredentialQuotaService, + lease: &CodexCredentialLease, + headers: &[(String, String)], +) { + if headers.is_empty() { + return; + } + if let Err(error) = quota + .synchronize_passive_headers(lease.account(), headers) + .await + { + tracing::warn!( + account_id = %lease.account_id(), + error = %error, + "OpenAI passive quota synchronization failed" + ); + } +} + +const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str { + match transport { + CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT, + CodexBackendTransport::WebSocket => WEBSOCKET_TRANSPORT, + } +} + +fn nonnegative_millis(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + fn compile_model_capabilities(model: &CodexCatalogModel) -> ProviderModelCapabilities { let evidence = model.capabilities(); let mut operations = BTreeSet::new(); @@ -479,11 +1054,7 @@ const fn support(evidence: CodexCatalogCapabilityEvidence) -> SupportLevel { fn selected_transport( request: &ProviderRequest, instance: &CodexProviderInstanceConfig, - has_continuation: bool, ) -> Result { - if has_continuation { - return Ok(CodexProviderTransport::WebSocket); - } let mut transport = instance.transport(); if let Some(value) = request .operation() @@ -492,7 +1063,7 @@ fn selected_transport( { transport = value .as_str() - .and_then(CodexProviderTransport::parse) + .and_then(CodexProviderTransport::parse_explicit) .ok_or_else(|| { provider_error( ProviderErrorKind::InvalidRequest, @@ -504,23 +1075,32 @@ fn selected_transport( } fn apply_transport(request: &mut CodexResponsesRequest, transport: CodexProviderTransport) { - let websocket = transport == CodexProviderTransport::WebSocket; - request.force_http_sse = !websocket; - request.force_websocket = websocket; - request.use_websocket = websocket; + match transport { + CodexProviderTransport::HttpOnly => { + request.force_http_sse = true; + request.use_websocket = false; + } + CodexProviderTransport::PreferWebSocket => { + request.force_http_sse = false; + request.use_websocket = true; + } + } } const fn transport_name(transport: CodexProviderTransport) -> &'static str { match transport { - CodexProviderTransport::HttpSse => HTTP_SSE_TRANSPORT, - CodexProviderTransport::WebSocket => WEBSOCKET_TRANSPORT, + CodexProviderTransport::HttpOnly => HTTP_SSE_TRANSPORT, + CodexProviderTransport::PreferWebSocket => WEBSOCKET_TRANSPORT, } } -const fn backend_transport(transport: CodexProviderTransport) -> CodexBackendTransport { +const fn accepts_backend_transport( + transport: CodexProviderTransport, + actual: CodexBackendTransport, +) -> bool { match transport { - CodexProviderTransport::HttpSse => CodexBackendTransport::HttpSse, - CodexProviderTransport::WebSocket => CodexBackendTransport::WebSocket, + CodexProviderTransport::HttpOnly => matches!(actual, CodexBackendTransport::HttpSse), + CodexProviderTransport::PreferWebSocket => true, } } @@ -622,9 +1202,9 @@ fn map_request_error(error: CodexRequestEncodeError) -> ProviderError { fn map_instance_config_error(error: CodexProviderConfigError) -> ProviderError { let kind = match error { CodexProviderConfigError::ProviderMismatch => ProviderErrorKind::InvalidRequest, - CodexProviderConfigError::InvalidBaseUrl - | CodexProviderConfigError::UnsafeBaseUrl - | CodexProviderConfigError::TransportInitialization => ProviderErrorKind::Protocol, + CodexProviderConfigError::InvalidBaseUrl | CodexProviderConfigError::UnsafeBaseUrl => { + ProviderErrorKind::Protocol + } }; provider_error(kind, UpstreamSendState::NotSent) } @@ -648,82 +1228,357 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError { } } -fn map_handshake_error(error: CodexClientError) -> ProviderError { - map_client_error(error, UpstreamSendState::Ambiguous) +struct MappedProviderFailure { + error: ProviderError, + account_failure: Option, + set_cookie_headers: Vec, + rate_limit_headers: Vec<(String, String)>, + observation: Option, + capture_response_cookies: bool, } -fn map_stream_error(error: CodexClientError) -> ProviderError { - map_client_error(error, UpstreamSendState::Sent) -} - -fn map_client_error(error: CodexClientError, uncertain_state: UpstreamSendState) -> ProviderError { - match error { - CodexClientError::Upstream { - status, - retry_after_seconds, - .. - } => { - let kind = http_error_kind(status); - let mut mapped = provider_error(kind, UpstreamSendState::Sent) - .with_status(status.as_u16()) - .redact_sensitive_context("upstream response body"); - if explicit_rejection_is_replay_safe(kind, status.as_u16()) { - mapped = mapped.with_replay_safe(); - } - match retry_after_seconds { - Some(seconds) => mapped.with_retry_after(Duration::from_secs(seconds)), - None => mapped, - } +impl MappedProviderFailure { + fn plain(error: ProviderError) -> Self { + Self { + error, + account_failure: None, + set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), + observation: None, + capture_response_cookies: false, } + } +} + +#[derive(Clone, Copy)] +enum ReplayBoundary { + BeforeSemanticOutput, + AfterSemanticOutput, +} + +impl ReplayBoundary { + const fn from_semantic_output(semantic_output_seen: bool) -> Self { + if semantic_output_seen { + Self::AfterSemanticOutput + } else { + Self::BeforeSemanticOutput + } + } + + const fn permits_provider_proof(self) -> bool { + matches!(self, Self::BeforeSemanticOutput) + } +} + +async fn apply_failure( + client: &CodexBackendClient, + selector: &CodexCredentialSelector, + quota: &CodexCredentialQuotaService, + lease: &CodexCredentialLease, + response_origin: &Url, + failure: &MappedProviderFailure, +) { + synchronize_passive_quota(quota, lease, &failure.rate_limit_headers).await; + if let Some(account_failure) = failure.account_failure { + client + .evict_websocket_account(lease.account_id().as_str()) + .await; + if let Err(error) = selector.record_failure(lease, account_failure).await { + tracing::warn!( + account_id = %lease.account_id(), + error = %error, + "Failed to persist OpenAI account failure state" + ); + } + } + if failure.capture_response_cookies + && !failure.set_cookie_headers.is_empty() + && let Err(error) = selector + .capture_response_cookies(lease, response_origin, &failure.set_cookie_headers) + .await + { + tracing::warn!( + account_id = %lease.account_id(), + error = %error, + "Failed to persist OpenAI response cookies" + ); + } +} + +fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { + map_client_error(error, UpstreamSendState::Ambiguous, true) +} + +fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { + map_client_error(error, UpstreamSendState::Sent, false) +} + +fn map_canonical_error( + error: CodexCanonicalError, + diagnostics: &CodexUpstreamDiagnostics, + set_cookie_headers: &[String], + rate_limit_headers: &[(String, String)], + replay_boundary: ReplayBoundary, +) -> MappedProviderFailure { + match error { + CodexCanonicalError::Protocol(error) => MappedProviderFailure::plain(error), + CodexCanonicalError::Upstream(failure) => map_upstream_failure( + CodexUpstreamFailure::from_sse_failure( + &failure, + diagnostics, + set_cookie_headers, + rate_limit_headers, + CodexUpstreamSendPhase::AfterPayload, + ), + None, + replay_boundary, + ), + } +} + +fn map_client_error( + error: CodexClientError, + uncertain_state: UpstreamSendState, + observe_transport: bool, +) -> MappedProviderFailure { + let continuation_failure = match &error { + CodexClientError::WebSocket(CodexWebSocketExchangeError::ContinuationUnavailable { + reason: PreviousResponseUnavailableReason::ConnectionBusy, + }) => Some(ContinuationFailure::Busy), + CodexClientError::WebSocket(CodexWebSocketExchangeError::ContinuationUnavailable { + .. + }) => Some(ContinuationFailure::HistoryUnavailable), + _ => None, + }; + let observation = observe_transport + .then(|| codex_error_observation(&error)) + .flatten(); + if let Some(failure) = error.upstream_failure() { + return map_upstream_failure(failure, observation, ReplayBoundary::BeforeSemanticOutput); + } + let mut failure = match error { + CodexClientError::Upstream { .. } => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )), CodexClientError::InvalidHeaderName(_) | CodexClientError::InvalidHeaderValue(_) | CodexClientError::WebSocketEncode(_) | CodexClientError::ModelCatalog(_) - | CodexClientError::CustomCa(_) => { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) - } - CodexClientError::StreamIdleTimeout { .. } => { - provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent) - } - CodexClientError::InvalidSse(_) => { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent) - } - CodexClientError::Http(error) => provider_error( - if error.is_timeout() { - ProviderErrorKind::Timeout + | CodexClientError::CustomCa(_) => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )), + CodexClientError::StreamIdleTimeout { .. } => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + )), + CodexClientError::InvalidSse(_) => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )), + CodexClientError::Http(error) => { + let send_state = if error.is_connect() { + UpstreamSendState::NotSent } else { - ProviderErrorKind::Transport - }, - uncertain_state, + uncertain_state + }; + MappedProviderFailure::plain(provider_error( + if error.is_timeout() { + ProviderErrorKind::Timeout + } else { + ProviderErrorKind::Transport + }, + send_state, + )) + } + CodexClientError::WebSocket(error) => MappedProviderFailure::plain(provider_error( + websocket_error_kind(&error), + websocket_send_state(&error), + )), + }; + if let Some(continuation_failure) = continuation_failure { + failure.error = failure + .error + .with_continuation_failure(continuation_failure); + } + failure.observation = observation; + failure +} + +fn map_upstream_failure( + failure: CodexUpstreamFailure, + observation: Option, + replay_boundary: ReplayBoundary, +) -> MappedProviderFailure { + let category = failure.category(); + let continuation_failure = failure + .persistable_code() + .filter(|code| is_history_failure_code(code)) + .map(|_| ContinuationFailure::HistoryUnavailable); + let send_state = upstream_send_state(failure.send_phase); + let mut error = provider_error(provider_error_kind(category), send_state) + .redact_sensitive_context("upstream response body"); + if let Some(status) = failure.status { + error = error.with_status(status.as_u16()); + } + if replay_boundary.permits_provider_proof() + && (failure.replay_is_safe() || continuation_failure.is_some()) + { + error = error.with_replay_safe(); + } + if let Some(continuation_failure) = continuation_failure { + error = error.with_continuation_failure(continuation_failure); + } + if let Some(retry_after) = failure.retry_after_seconds.map(Duration::from_secs) { + error = error.with_retry_after(retry_after); + } + if let Some(code) = failure + .persistable_code() + .and_then(|code| SafeUpstreamValue::new(code.to_owned()).ok()) + { + error = error.with_upstream_code(code); + } + if let Some(request_id) = failure + .request_id + .as_deref() + .and_then(|request_id| SafeUpstreamValue::new(request_id.to_owned()).ok()) + { + error = error.with_upstream_request_id(request_id); + } + MappedProviderFailure { + error, + account_failure: account_failure(category, failure.retry_after_seconds), + set_cookie_headers: failure.set_cookie_headers, + rate_limit_headers: failure.rate_limit_headers, + observation, + capture_response_cookies: !matches!( + category, + CodexFailureCategory::CloudflareChallenge | CodexFailureCategory::CloudflarePathBlocked ), - CodexClientError::WebSocket(_) => { - provider_error(ProviderErrorKind::Unavailable, uncertain_state) + } +} + +fn is_history_failure_code(code: &str) -> bool { + matches!( + code, + "previous_response_not_found" + | "invalid_encrypted_content" + | "missing_tool_output" + | "no_tool_output" + ) +} + +const fn provider_error_kind(category: CodexFailureCategory) -> ProviderErrorKind { + match category { + CodexFailureCategory::ModelUnsupported => ProviderErrorKind::Unsupported, + CodexFailureCategory::CredentialExpired => ProviderErrorKind::Unauthorized, + CodexFailureCategory::IdentityVerificationRequired | CodexFailureCategory::Banned => { + ProviderErrorKind::PermissionDenied + } + CodexFailureCategory::RateLimited => ProviderErrorKind::RateLimited, + CodexFailureCategory::QuotaExhausted => ProviderErrorKind::QuotaExhausted, + CodexFailureCategory::CloudflareChallenge + | CodexFailureCategory::CloudflarePathBlocked + | CodexFailureCategory::Unavailable => ProviderErrorKind::Unavailable, + CodexFailureCategory::InvalidRequest => ProviderErrorKind::InvalidRequest, + CodexFailureCategory::PermissionDenied => ProviderErrorKind::PermissionDenied, + CodexFailureCategory::Timeout => ProviderErrorKind::Timeout, + CodexFailureCategory::Transport => ProviderErrorKind::Transport, + } +} + +const fn upstream_send_state(phase: CodexUpstreamSendPhase) -> UpstreamSendState { + match phase { + CodexUpstreamSendPhase::BeforePayload => UpstreamSendState::NotSent, + CodexUpstreamSendPhase::AfterPayload => UpstreamSendState::Sent, + CodexUpstreamSendPhase::Ambiguous => UpstreamSendState::Ambiguous, + } +} + +fn account_failure( + category: CodexFailureCategory, + retry_after_seconds: Option, +) -> Option { + match category { + CodexFailureCategory::CredentialExpired => Some(CodexAccountFailure::CredentialExpired), + CodexFailureCategory::IdentityVerificationRequired => { + Some(CodexAccountFailure::IdentityVerificationRequired) + } + CodexFailureCategory::Banned => Some(CodexAccountFailure::Banned), + CodexFailureCategory::RateLimited => Some(CodexAccountFailure::RateLimited { + retry_after: retry_after_seconds.map(Duration::from_secs), + }), + CodexFailureCategory::QuotaExhausted => Some(CodexAccountFailure::QuotaExhausted), + CodexFailureCategory::CloudflareChallenge => { + Some(CodexAccountFailure::CloudflareChallenge { + retry_after: retry_after_seconds.map(Duration::from_secs), + }) + } + CodexFailureCategory::CloudflarePathBlocked => { + Some(CodexAccountFailure::CloudflarePathBlocked) + } + CodexFailureCategory::ModelUnsupported + | CodexFailureCategory::InvalidRequest + | CodexFailureCategory::PermissionDenied + | CodexFailureCategory::Timeout + | CodexFailureCategory::Unavailable + | CodexFailureCategory::Transport => None, + } +} + +const fn websocket_send_state(error: &CodexWebSocketExchangeError) -> UpstreamSendState { + match error { + CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::Connect(_) + | CodexWebSocketExchangeError::ConnectTimeout { .. } + | CodexWebSocketExchangeError::FastPathTimeout { .. } + | CodexWebSocketExchangeError::OriginCircuitOpen + | CodexWebSocketExchangeError::OriginHalfOpenBusy + | CodexWebSocketExchangeError::SharedConnectFailed + | CodexWebSocketExchangeError::ContinuationUnavailable { .. } => UpstreamSendState::NotSent, + CodexWebSocketExchangeError::Upstream(_) + | CodexWebSocketExchangeError::InvalidSse(_) + | CodexWebSocketExchangeError::InvalidCompletedResponse { .. } + | CodexWebSocketExchangeError::UnexpectedBinaryEvent => UpstreamSendState::Sent, + CodexWebSocketExchangeError::Transport(_) + | CodexWebSocketExchangeError::PostSendAmbiguous { .. } + | CodexWebSocketExchangeError::SendTimeout { .. } + | CodexWebSocketExchangeError::ClosedBeforeTerminal + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } + | CodexWebSocketExchangeError::InitialEventTimeout { .. } => UpstreamSendState::Ambiguous, + } +} + +const fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> ProviderErrorKind { + match error { + CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::InvalidSse(_) + | CodexWebSocketExchangeError::InvalidCompletedResponse { .. } + | CodexWebSocketExchangeError::UnexpectedBinaryEvent => ProviderErrorKind::Protocol, + CodexWebSocketExchangeError::ConnectTimeout { .. } + | CodexWebSocketExchangeError::FastPathTimeout { .. } + | CodexWebSocketExchangeError::SendTimeout { .. } + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } + | CodexWebSocketExchangeError::InitialEventTimeout { .. } => ProviderErrorKind::Timeout, + CodexWebSocketExchangeError::OriginCircuitOpen + | CodexWebSocketExchangeError::OriginHalfOpenBusy + | CodexWebSocketExchangeError::SharedConnectFailed + | CodexWebSocketExchangeError::ContinuationUnavailable { .. } => { + ProviderErrorKind::Unavailable + } + CodexWebSocketExchangeError::Upstream(_) => ProviderErrorKind::Unavailable, + CodexWebSocketExchangeError::Transport(_) + | CodexWebSocketExchangeError::Connect(_) + | CodexWebSocketExchangeError::PostSendAmbiguous { .. } + | CodexWebSocketExchangeError::ClosedBeforeTerminal + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { + ProviderErrorKind::Transport } } } -fn http_error_kind(status: StatusCode) -> ProviderErrorKind { - match status.as_u16() { - 400 | 404 | 409 | 422 => ProviderErrorKind::InvalidRequest, - 401 => ProviderErrorKind::Unauthorized, - 402 => ProviderErrorKind::QuotaExhausted, - 403 => ProviderErrorKind::PermissionDenied, - 408 | 504 => ProviderErrorKind::Timeout, - 429 => ProviderErrorKind::RateLimited, - 500..=599 => ProviderErrorKind::Unavailable, - _ => ProviderErrorKind::Transport, - } -} - -fn explicit_rejection_is_replay_safe(kind: ProviderErrorKind, status: u16) -> bool { - matches!( - (kind, status), - (ProviderErrorKind::Unauthorized, 401) - | (ProviderErrorKind::QuotaExhausted, 402) - | (ProviderErrorKind::RateLimited, 429) - ) -} - fn provider_error(kind: ProviderErrorKind, send_state: UpstreamSendState) -> ProviderError { ProviderError::new(kind, send_state) } @@ -737,7 +1592,7 @@ fn remaining(deadline: SystemTime) -> Option { fn normalize_and_validate_base_url( url: &mut Url, - endpoint_policy: CodexEndpointPolicy, + origin_policy: &dyn CodexOriginPolicy, ) -> Result<(), CodexProviderConfigError> { if url.cannot_be_a_base() || url.host().is_none() @@ -748,18 +1603,7 @@ fn normalize_and_validate_base_url( { return Err(CodexProviderConfigError::UnsafeBaseUrl); } - let allowed = match endpoint_policy { - CodexEndpointPolicy::Official => { - url.scheme() == "https" - && matches!(url.host(), Some(Host::Domain(host)) if host == OFFICIAL_CODEX_HOST) - && url.port().is_none() - } - CodexEndpointPolicy::Loopback => { - url.scheme() == "http" - && matches!(url.host(), Some(Host::Ipv4(host)) if host.is_loopback()) - } - }; - if !allowed { + if !origin_policy.allows(url) { return Err(CodexProviderConfigError::UnsafeBaseUrl); } let normalized_path = url.path().trim_end_matches('/'); @@ -769,3 +1613,259 @@ fn normalize_and_validate_base_url( url.set_path(OFFICIAL_CODEX_BASE_PATH); Ok(()) } + +const WORKER_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60); +const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); +const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); +const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); +const QUOTA_REFRESH_INTERVAL: Duration = Duration::from_secs(5 * 60); +const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release"; +const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag"; + +pub(crate) fn worker_contributions( + refresh: Arc, + quota: Arc, + catalog: Arc, + desktop_release: Arc, + instances: Arc, + provider_kind: ProviderKind, +) -> Result, WorkerDefinitionError> { + let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, PROVIDER_NAME)?; + let quota_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, PROVIDER_NAME)?; + let etag_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_ETAG_WORKER_OWNER)?; + let desktop_release_id = + WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?; + Ok(vec![ + WorkerContribution::Registration(scheduled_registration( + refresh_id, + OAUTH_REFRESH_INTERVAL, + Box::new(OpenAiOAuthRefreshTask { service: refresh }), + )?), + WorkerContribution::Registration(scheduled_registration( + quota_id, + QUOTA_REFRESH_INTERVAL, + Box::new(OpenAiQuotaTask { + instances: Arc::clone(&instances), + quota, + provider_kind: provider_kind.clone(), + }), + )?), + WorkerContribution::Registration(WorkerRegistration::try_new( + etag_id, + WorkerRunnable::Daemon { + restart: DaemonRestartPolicy::try_new( + WORKER_INITIAL_BACKOFF, + WORKER_MAXIMUM_BACKOFF, + )?, + task: Box::new(OpenAiCatalogEtagTask { + instances, + catalog, + provider_kind, + }), + }, + )?), + WorkerContribution::Registration(scheduled_registration( + desktop_release_id, + APPCAST_POLL_INTERVAL, + Box::new(OpenAiDesktopReleaseTask { + service: desktop_release, + }), + )?), + ]) +} + +fn scheduled_registration( + id: WorkerId, + interval: Duration, + task: Box, +) -> Result { + let schedule = WorkerSchedule::try_new( + interval, + WORKER_INITIAL_BACKOFF, + WORKER_MAXIMUM_BACKOFF, + WORKER_LEASE_TTL, + WORKER_LEASE_RENEWAL, + )?; + let lease = WorkerLeaseRequest::try_new(id.clone(), WORKER_LEASE_TTL)?; + WorkerRegistration::try_new( + id, + WorkerRunnable::Scheduled { + schedule, + lease: Some(lease), + task, + }, + ) +} + +struct OpenAiOAuthRefreshTask { + service: Arc, +} + +impl ScheduledTask for OpenAiOAuthRefreshTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + if context.cancellation().is_cancelled() { + return Ok(()); + } + let outcomes = self + .service + .refresh_due() + .await + .map_err(|_| WorkerTaskError::safe("OpenAI OAuth refresh failed"))?; + let operational_failures = outcomes + .iter() + .filter(|outcome| { + matches!( + outcome, + CodexCredentialRefreshOutcome::Transient { .. } + | CodexCredentialRefreshOutcome::Ambiguous { .. } + | CodexCredentialRefreshOutcome::Failed { .. } + ) + }) + .count(); + if operational_failures > 0 { + tracing::warn!( + operational_failures, + "OpenAI OAuth refresh cycle contained operational failures" + ); + } + Ok(()) + }) + } +} + +struct OpenAiQuotaTask { + instances: Arc, + quota: Arc, + provider_kind: ProviderKind, +} + +struct OpenAiCatalogEtagTask { + instances: Arc, + catalog: Arc, + provider_kind: ProviderKind, +} + +struct OpenAiDesktopReleaseTask { + service: Arc, +} + +impl ScheduledTask for OpenAiDesktopReleaseTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let refresh = self.service.refresh(); + tokio::pin!(refresh); + let result = tokio::select! { + () = context.cancellation().cancelled() => return Ok(()), + result = &mut refresh => result, + }; + if let Err(error) = result { + // 上游检查失败已经作为 Provider 观察事实保存；本周期本身正常完成， + // 避免 Host 的短退避持续请求固定官方 appcast。 + tracing::warn!(error = %error, "OpenAI Desktop release check failed"); + } + Ok(()) + }) + } +} + +impl ScheduledTask for OpenAiQuotaTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let instances = self + .instances + .list_instances(&self.provider_kind, false) + .await + .map_err(|_| WorkerTaskError::safe("OpenAI Provider instances unavailable"))?; + let mut failures = 0_u64; + for config in instances { + if context.cancellation().is_cancelled() { + return Ok(()); + } + if !config.enabled() || config.provider_kind() != &self.provider_kind { + failures = failures.saturating_add(1); + continue; + } + let instance = ProviderInstance::new( + config.id().clone(), + config.provider_kind().clone(), + config.base_url().to_owned(), + config.enabled(), + InstanceHealth::Healthy, + ); + match self.quota.synchronize_instance(&instance).await { + Ok(summary) if summary.has_operational_failures() => { + tracing::warn!( + provider_instance_id = %instance.id(), + "OpenAI quota cycle contained operational failures" + ); + } + Ok(_) => {} + Err(_) => { + failures = failures.saturating_add(1); + tracing::warn!( + provider_instance_id = %instance.id(), + "OpenAI quota synchronization failed" + ); + } + } + } + if failures == 0 { + Ok(()) + } else { + Err(WorkerTaskError::safe("OpenAI quota synchronization failed")) + } + }) + } +} + +impl DaemonTask for OpenAiCatalogEtagTask { + fn run( + &self, + cancellation: gateway_core::engine::CancellationToken, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + loop { + let changed = tokio::select! { + () = cancellation.cancelled() => return Ok(()), + changed = self.catalog.wait_for_etag_refresh() => changed, + }; + let instances = self + .instances + .list_instances(&self.provider_kind, false) + .await + .map_err(|_| { + WorkerTaskError::safe( + "OpenAI Provider instances unavailable for model ETag refresh", + ) + })?; + for config in instances { + if cancellation.is_cancelled() { + return Ok(()); + } + if !changed.contains(config.id()) + || !config.enabled() + || config.provider_kind() != &self.provider_kind + { + continue; + } + let instance = ProviderInstance::new( + config.id().clone(), + config.provider_kind().clone(), + config.base_url().to_owned(), + true, + InstanceHealth::Healthy, + ); + if let Err(error) = self.catalog.refresh_instance(&instance).await { + tracing::warn!( + provider_instance_id = %instance.id(), + error = %error, + "OpenAI model catalog ETag refresh failed" + ); + } + } + } + }) + } +} diff --git a/backend/crates/providers/openai/src/transport/canonical.rs b/backend/crates/providers/openai/src/transport/canonical.rs index a65b8f9f6..dc43fe416 100644 --- a/backend/crates/providers/openai/src/transport/canonical.rs +++ b/backend/crates/providers/openai/src/transport/canonical.rs @@ -1,18 +1,23 @@ //! Codex Responses SSE 到核心 canonical event 的单一解码边界。 -use std::collections::{BTreeMap, BTreeSet}; +use std::{ + collections::{BTreeMap, BTreeSet}, + fmt, +}; use gateway_core::accounting::Usage; use gateway_core::engine::UpstreamSendState; use gateway_core::error::{ProviderError, ProviderErrorKind}; use gateway_core::event::{ - ContentItem, ContentKind, FinishReason, GatewayEvent, ReasoningDelta, ResponseMeta, TextDelta, - ToolCallDelta, + ContentItem, ContentKind, FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, + ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta, }; use gateway_protocol::openai::events::{TokenUsage, extract_usage}; use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder}; use serde_json::Value; +use thiserror::Error; +use super::protocol::responses::{ResponsesSseFailure, response_event_signals}; use super::usage::openai_billing_breakdown; const CONTENTS_PER_OUTPUT: u32 = 1_024; @@ -30,6 +35,83 @@ pub struct CodexCanonicalDecoder { content: BTreeMap, tool_arguments_seen: BTreeSet, usage_emitted: bool, + semantic_output_seen: bool, +} + +/// 上游 Responses 事件的两类失败：协议损坏，或上游明确报告业务失败。 +#[derive(Error)] +pub enum CodexCanonicalError { + /// SSE/JSON 违反了已知协议不变量。 + #[error("invalid Codex Responses event")] + Protocol(#[source] ProviderError), + /// 上游在成功建立流后发送了明确的失败事件。 + #[error("Codex upstream reported a failed response")] + Upstream(ResponsesSseFailure), +} + +impl fmt::Debug for CodexCanonicalError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Protocol(error) => formatter.debug_tuple("Protocol").field(error).finish(), + Self::Upstream(_) => formatter.write_str("Upstream()"), + } + } +} + +impl From for CodexCanonicalError { + fn from(error: ProviderError) -> Self { + Self::Protocol(error) + } +} + +/// 单次增量解码的有序结果。 +/// +/// 上游失败不是解析异常：它与失败前已经产生的事件一起返回，Provider 因而可以 +/// 先保留真实输出，再把类型化失败交给 Core 收敛。 +pub enum CodexCanonicalOutcome { + /// 本批次只包含正常事件。 + Events(Vec), + /// 本批次在若干正常事件后到达失败边界。 + Failed(CodexCanonicalFailure), +} + +/// `response.failed` 或协议错误发生时的单一 typed outcome。 +pub struct CodexCanonicalFailure { + events: Vec, + error: CodexCanonicalError, + semantic_output_seen: bool, +} + +impl fmt::Debug for CodexCanonicalFailure { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CodexCanonicalFailure") + .field("events", &self.events) + .field("error", &self.error) + .field("semantic_output_seen", &self.semantic_output_seen) + .finish() + } +} + +impl CodexCanonicalFailure { + /// 失败前按上游顺序产生的事件。 + pub fn events(&self) -> &[ProviderEvent] { + &self.events + } + + /// 类型化失败；其 Debug/Display 不包含上游正文。 + pub const fn error(&self) -> &CodexCanonicalError { + &self.error + } + + /// 失败前是否已经出现客户端可消费的真实输出。 + pub const fn semantic_output_seen(&self) -> bool { + self.semantic_output_seen + } + + pub(crate) fn into_parts(self) -> (Vec, CodexCanonicalError, bool) { + (self.events, self.error, self.semantic_output_seen) + } } impl CodexCanonicalDecoder { @@ -43,37 +125,107 @@ impl CodexCanonicalDecoder { content: BTreeMap::new(), tool_arguments_seen: BTreeSet::new(), usage_emitted: false, + semantic_output_seen: false, } } - pub fn push(&mut self, chunk: &[u8]) -> Result, ProviderError> { - let events = self.decoder.push(chunk).map_err(protocol_error)?; + pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome { + let events = match self.decoder.push(chunk) { + Ok(events) => events, + Err(error) => { + return self.failure( + Vec::new(), + CodexCanonicalError::Protocol(protocol_error(error)), + ); + } + }; self.decode(events) } - pub fn finish(&mut self) -> Result, ProviderError> { - let events = self.decoder.finish().map_err(protocol_error)?; + pub fn finish(&mut self) -> CodexCanonicalOutcome { + let events = match self.decoder.finish() { + Ok(events) => events, + Err(error) => { + return self.failure( + Vec::new(), + CodexCanonicalError::Protocol(protocol_error(error)), + ); + } + }; self.decode(events) } - fn decode(&mut self, events: Vec) -> Result, ProviderError> { - let mut canonical = Vec::new(); + fn decode(&mut self, events: Vec) -> CodexCanonicalOutcome { + let mut output = Vec::new(); for event in events { if event.data.trim() == "[DONE]" { if !self.completed { - return Err(protocol_error_marker()); + return self.failure( + output, + CodexCanonicalError::Protocol(protocol_error_marker()), + ); } continue; } - let value = serde_json::from_str::(&event.data).map_err(protocol_error)?; - let event_type = event + let value = match serde_json::from_str::(&event.data) { + Ok(value) => value, + Err(error) => { + return self + .failure(output, CodexCanonicalError::Protocol(protocol_error(error))); + } + }; + let Some(event_type) = event .event .as_deref() .or_else(|| value.get("type").and_then(Value::as_str)) - .ok_or_else(protocol_error_marker)?; - self.decode_event(event_type, &value, &mut canonical)?; + else { + return self.failure( + output, + CodexCanonicalError::Protocol(protocol_error_marker()), + ); + }; + if matches!(event_type, "response.failed" | "error") { + return self.failure( + output, + CodexCanonicalError::Upstream(ResponsesSseFailure::from_event( + event_type, &value, + )), + ); + } + let mut canonical = Vec::new(); + if let Err(error) = self.decode_event(event_type, &value, &mut canonical) { + return self.failure(output, CodexCanonicalError::Protocol(error)); + } + let semantic_output = response_event_signals(Some(event_type), &value).semantic_output; + let wire = match ProtocolWireEvent::json("openai", event.event, value) { + Ok(wire) => wire, + Err(_) => { + return self.failure( + output, + CodexCanonicalError::Protocol(protocol_error_marker()), + ); + } + }; + output.push(if canonical.is_empty() { + ProviderEvent::wire(wire) + } else { + ProviderEvent::canonical_with_wire(canonical, wire) + }); + self.semantic_output_seen |= semantic_output; } - Ok(canonical) + CodexCanonicalOutcome::Events(output) + } + + fn failure( + &self, + events: Vec, + error: CodexCanonicalError, + ) -> CodexCanonicalOutcome { + CodexCanonicalOutcome::Failed(CodexCanonicalFailure { + events, + error, + semantic_output_seen: self.semantic_output_seen, + }) } fn decode_event( @@ -105,7 +257,7 @@ impl CodexCanonicalDecoder { "response.completed" | "response.incomplete" => { self.complete(event_type, value, output) } - "response.failed" | "error" => Err(upstream_event_error(value)), + "response.failed" | "error" => Err(protocol_error_marker()), "response.output_text.done" | "response.refusal.done" | "response.reasoning_summary_text.done" @@ -118,7 +270,7 @@ impl CodexCanonicalDecoder { | "response.rate_limits.updated" | "codex.rate_limits" | "response.metadata" => Ok(()), - _ => Err(unsupported_event_error()), + _ => Ok(()), } } @@ -181,10 +333,8 @@ impl CodexCanonicalDecoder { self.add_content(index, ContentKind::Reasoning, output) } Some("message") => Ok(()), - Some("image_generation_call" | "computer_call" | "web_search_call") => { - Err(unsupported_event_error()) - } - _ => Err(unsupported_event_error()), + Some("image_generation_call" | "computer_call" | "web_search_call") => Ok(()), + _ => Ok(()), } } @@ -252,7 +402,7 @@ impl CodexCanonicalDecoder { let kind = match part.get("type").and_then(Value::as_str) { Some("output_text" | "refusal") => ContentKind::Text, Some("summary_text" | "reasoning_text") => ContentKind::Reasoning, - _ => return Err(unsupported_event_error()), + _ => return Ok(()), }; self.add_content(content_index(output_index, part_index)?, kind, output) } @@ -501,33 +651,6 @@ fn incomplete_finish_reason(response: &Value) -> FinishReason { } } -fn upstream_event_error(value: &Value) -> ProviderError { - let kind = match value - .pointer("/error/code") - .or_else(|| value.pointer("/response/error/code")) - .and_then(Value::as_str) - { - Some("invalid_request" | "invalid_prompt" | "cyber_policy") => { - ProviderErrorKind::InvalidRequest - } - Some("unsupported" | "unsupported_feature" | "model_not_supported") => { - ProviderErrorKind::Unsupported - } - Some( - "unauthorized" - | "invalid_api_key" - | "token_invalid" - | "token_invalidated" - | "access_token_expired", - ) => ProviderErrorKind::Unauthorized, - Some("permission_denied") => ProviderErrorKind::PermissionDenied, - Some("rate_limit_exceeded") => ProviderErrorKind::RateLimited, - Some("quota_exceeded" | "insufficient_quota") => ProviderErrorKind::QuotaExhausted, - _ => ProviderErrorKind::Unavailable, - }; - ProviderError::new(kind, UpstreamSendState::Sent).redact_sensitive_context("upstream event") -} - fn protocol_error(_error: impl std::fmt::Debug) -> ProviderError { protocol_error_marker() } @@ -536,8 +659,3 @@ fn protocol_error_marker() -> ProviderError { ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::Sent) .redact_sensitive_context("invalid upstream event") } - -fn unsupported_event_error() -> ProviderError { - ProviderError::new(ProviderErrorKind::Unsupported, UpstreamSendState::Sent) - .redact_sensitive_context("unsupported upstream event") -} diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index 9a1c989a9..c2625c10a 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -12,6 +12,7 @@ use crate::transport::profile::CodexWireProfileState; use bytes::Bytes; use futures::{Stream, StreamExt}; use gateway_protocol::openai::{ + WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY, events::{TokenUsage, retry_after_seconds_from_body}, sse::SseError, }; @@ -22,12 +23,9 @@ use reqwest::{ use serde_json::{Value, map::Map}; use thiserror::Error; -use crate::transport::protocol::responses::{ - CodexResponsesRequest, TransportRequirement, - WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY, -}; +use crate::transport::protocol::responses::{CodexResponsesRequest, TransportRequirement}; -use super::diagnostics::CodexUpstreamDiagnostics; +use super::diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamFailure, CodexUpstreamSendPhase}; use super::response_meta::CodexResponseMetadata; use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key}; use super::websocket::{ @@ -137,6 +135,10 @@ pub enum CodexClientError { rate_limit_headers: Vec<(String, String)>, /// 实际收到该上游响应的 transport。 transport: CodexBackendTransport, + /// 错误响应前已经确认的 transport 与 HTTP 阶段事实。 + transport_metrics: Box, + /// 上游拒绝发生时业务 payload 的发送阶段。 + send_phase: CodexUpstreamSendPhase, }, } @@ -168,12 +170,14 @@ impl fmt::Debug for CodexClientError { status, retry_after_seconds, transport, + send_phase, .. } => formatter .debug_struct("CodexClientError::Upstream") .field("status", status) .field("retry_after_seconds", retry_after_seconds) .field("transport", transport) + .field("send_phase", send_phase) .field("body", &"[REDACTED]") .finish(), } @@ -196,6 +200,41 @@ impl CodexClientError { | Self::WebSocketEncode(_) => None, } } + + pub(crate) fn upstream_failure(&self) -> Option { + match self { + Self::Upstream { + status, + body, + retry_after_seconds, + diagnostics, + set_cookie_headers, + rate_limit_headers, + send_phase, + .. + } => Some(CodexUpstreamFailure::from_response( + *status, + body, + *retry_after_seconds, + diagnostics, + set_cookie_headers, + rate_limit_headers, + *send_phase, + )), + Self::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => { + Some(CodexUpstreamFailure::from_response( + StatusCode::from_u16(upstream.status_code).unwrap_or(StatusCode::BAD_GATEWAY), + &upstream.body, + upstream.retry_after_seconds, + &upstream.diagnostics, + &upstream.set_cookie_headers, + &[], + upstream.send_phase, + )) + } + _ => None, + } + } } /// Codex 客户端结果类型。 @@ -293,8 +332,8 @@ pub struct CodexBackendResponse { pub response_metadata: CodexResponseMetadata, /// 传输选择与低延迟阶段耗时。 pub transport_metrics: CodexTransportMetrics, - /// 池中 WebSocket 保留 connection-local continuation 的最晚时间。 - pub connection_local_continuation_expires_at: Option, + /// 当前响应是否由池中 WebSocket 保留 connection-local continuation。 + pub connection_local_continuation: bool, } /// Codex Responses 实际使用的上游传输。 @@ -377,8 +416,8 @@ pub struct CodexBackendStreamingResponse { pub response_metadata: CodexResponseMetadata, /// 传输选择与低延迟阶段耗时。 pub transport_metrics: CodexTransportMetrics, - /// terminal completed 后池中 WebSocket 保留 connection-local continuation 的最晚时间。 - pub connection_local_continuation_expires_at: Option, + /// terminal completed 后是否由池中 WebSocket 保留 connection-local continuation。 + pub connection_local_continuation: bool, } // --------------------------------------------------------------------------- @@ -612,6 +651,8 @@ pub(super) fn websocket_exchange_error_to_client_error( set_cookie_headers: upstream.set_cookie_headers, rate_limit_headers: Vec::new(), transport: CodexBackendTransport::WebSocket, + transport_metrics: Box::default(), + send_phase: upstream.send_phase, } } error => CodexClientError::WebSocket(error), @@ -624,9 +665,7 @@ pub(super) fn websocket_success_decision( ) -> CodexTransportDecision { match requirement { TransportRequirement::ExactWebSocketContinuation => CodexTransportDecision::ExactWebSocket, - TransportRequirement::WebSocketRequired | TransportRequirement::ExplicitWebSocketWarmup => { - CodexTransportDecision::RequiredWebSocket - } + TransportRequirement::ExplicitWebSocketWarmup => CodexTransportDecision::RequiredWebSocket, _ if prepared.reused() => CodexTransportDecision::ReusedWebSocket, _ => CodexTransportDecision::ConnectedWebSocket, } diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index c13358fbc..9749533e2 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -5,6 +5,7 @@ use std::{ use futures::{StreamExt, TryStreamExt}; use gateway_protocol::openai::{ + X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER, events::{extract_sse_usage, retry_after_seconds_from_body}, sse::SseEventDecoder, }; @@ -21,6 +22,7 @@ use crate::transport::{ CodexModelCatalogError, CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, catalog_etag, parse_codex_model_catalog, }, + diagnostics::CodexUpstreamSendPhase, endpoints::{CODEX_RESPONSES_PATH, endpoint_url}, headers::{build_codex_base_headers, insert_optional_header, websocket_header_pairs}, profile::{CodexWireProfile, CodexWireProfileState}, @@ -126,6 +128,12 @@ impl CodexBackendClient { set_cookie_headers, rate_limit_headers, transport: CodexBackendTransport::HttpSse, + transport_metrics: Box::new(CodexTransportMetrics { + upstream_headers_ms: Some(upstream_headers_ms), + http_version: Some(http_version), + ..CodexTransportMetrics::default() + }), + send_phase: CodexUpstreamSendPhase::AfterPayload, }); } @@ -169,7 +177,7 @@ impl CodexBackendClient { http_version: Some(http_version), ..CodexTransportMetrics::default() }, - connection_local_continuation_expires_at: None, + connection_local_continuation: false, }) } @@ -209,6 +217,12 @@ impl CodexBackendClient { set_cookie_headers, rate_limit_headers, transport: CodexBackendTransport::HttpSse, + transport_metrics: Box::new(CodexTransportMetrics { + upstream_headers_ms: Some(upstream_headers_ms), + http_version: Some(http_version), + ..CodexTransportMetrics::default() + }), + send_phase: CodexUpstreamSendPhase::AfterPayload, }); } @@ -228,7 +242,7 @@ impl CodexBackendClient { http_version: Some(http_version), ..CodexTransportMetrics::default() }, - connection_local_continuation_expires_at: None, + connection_local_continuation: false, }) } @@ -334,8 +348,7 @@ impl CodexBackendClient { TransportRequirement::PersistedContinuation | TransportRequirement::ExternalUnknown | TransportRequirement::NewChain => Some(self.websocket_fast_path_budget), - TransportRequirement::WebSocketRequired - | TransportRequirement::ExplicitWebSocketWarmup + TransportRequirement::ExplicitWebSocketWarmup | TransportRequirement::ExactWebSocketContinuation => None, TransportRequirement::HttpRequired => None, }; @@ -459,8 +472,7 @@ impl CodexBackendClient { first_event_ms: exchange.first_event_ms, ..metrics }, - connection_local_continuation_expires_at: exchange - .connection_local_continuation_expires_at, + connection_local_continuation: exchange.connection_local_continuation, }) } }; @@ -524,8 +536,7 @@ impl CodexBackendClient { diagnostics: exchange.diagnostics, response_metadata: exchange.response_metadata, transport_metrics: metrics, - connection_local_continuation_expires_at: exchange - .connection_local_continuation_expires_at, + connection_local_continuation: exchange.connection_local_continuation, } }) } @@ -595,6 +606,8 @@ impl CodexBackendClient { set_cookie_headers, rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, + transport_metrics: Box::default(), + send_phase: CodexUpstreamSendPhase::AfterPayload, }); } Ok(parse_codex_model_catalog(&body, etag.as_deref())?) @@ -615,12 +628,12 @@ impl CodexBackendClient { } insert_optional_header( &mut headers, - crate::transport::protocol::responses::X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, + X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, request.responses_lite.as_deref(), )?; insert_optional_header( &mut headers, - crate::transport::protocol::responses::X_OPENAI_MEMGEN_REQUEST_HEADER, + X_OPENAI_MEMGEN_REQUEST_HEADER, request.memgen_request.as_deref(), )?; Ok(headers) @@ -632,9 +645,7 @@ impl CodexBackendClient { context: CodexRequestContext<'_>, ) -> CodexClientResult { let mut headers = self.request_headers_for_http_response(request, context)?; - headers.remove( - crate::transport::protocol::responses::X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, - ); + headers.remove(X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER); Ok(headers) } @@ -725,30 +736,6 @@ impl CodexBackendClient { insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?; insert_optional_header(&mut headers, "cookie", context.cookie_header)?; headers.insert(ACCEPT, HeaderValue::from_static("application/json")); - headers.insert( - HeaderName::from_static("openai-beta"), - HeaderValue::from_static("codex-1"), - ); - headers.insert( - HeaderName::from_static("oai-language"), - HeaderValue::from_static("zh-CN"), - ); - headers.insert( - HeaderName::from_static("sec-fetch-site"), - HeaderValue::from_static("none"), - ); - headers.insert( - HeaderName::from_static("sec-fetch-mode"), - HeaderValue::from_static("no-cors"), - ); - headers.insert( - HeaderName::from_static("sec-fetch-dest"), - HeaderValue::from_static("empty"), - ); - headers.insert( - HeaderName::from_static("priority"), - HeaderValue::from_static("u=4, i"), - ); Ok(headers) } } @@ -776,18 +763,14 @@ async fn read_model_catalog_body(response: ReqwestResponse) -> CodexClientResult } fn websocket_connection_profile(headers: &HeaderMap) -> String { - [ - "originator", - "user-agent", - crate::transport::protocol::responses::X_OPENAI_MEMGEN_REQUEST_HEADER, - ] - .map(|name| { - headers - .get(name) - .and_then(|value| value.to_str().ok()) - .unwrap_or_default() - }) - .join("\0") + ["originator", "user-agent", X_OPENAI_MEMGEN_REQUEST_HEADER] + .map(|name| { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + }) + .join("\0") } fn http_sse_stream(response: ReqwestResponse) -> CodexBackendSseStream { diff --git a/backend/crates/providers/openai/src/transport/diagnostics.rs b/backend/crates/providers/openai/src/transport/diagnostics.rs index 4d211b5d1..b1c3eaf70 100644 --- a/backend/crates/providers/openai/src/transport/diagnostics.rs +++ b/backend/crates/providers/openai/src/transport/diagnostics.rs @@ -1,7 +1,8 @@ //! Upstream response diagnostics captured at the transport boundary. +use super::protocol::responses::ResponsesSseFailure; use base64::{Engine as _, engine::general_purpose::STANDARD}; -use reqwest::header::HeaderMap; +use reqwest::{StatusCode, header::HeaderMap}; use serde_json::Value; const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[ @@ -22,6 +23,238 @@ const UPSTREAM_TRACE_HEADERS: &[&str] = &[ ]; const IDENTITY_AUTHORIZATION_ERROR_HEADER: &str = "x-openai-authorization-error"; const IDENTITY_ERROR_JSON_HEADER: &str = "x-error-json"; +const PERSISTABLE_UPSTREAM_CODES: &[&str] = &[ + "access_token_expired", + "account_banned", + "account_deactivated", + "account_disabled", + "account_suspended", + "authentication_error", + "billing_limit", + "cyber_policy", + "deactivated_workspace", + "identity_verification_required", + "insufficient_quota", + "invalid_api_key", + "invalid_encrypted_content", + "invalid_prompt", + "invalid_request", + "missing_tool_output", + "model_not_available", + "model_not_supported", + "no_tool_output", + "organization_disabled", + "payment_required", + "permission_denied", + "previous_response_not_found", + "quota_exceeded", + "quota_exhausted", + "rate_limit_error", + "rate_limit_exceeded", + "rate_limit_reached", + "refresh_token_invalidated", + "server_is_overloaded", + "slow_down", + "token_expired", + "token_invalid", + "token_invalidated", + "token_revoked", + "unauthorized", + "unsupported", + "unsupported_feature", + "usage_limit_reached", + "verification_required", + "workspace_deactivated", + "workspace_member_credits_depleted", + "workspace_member_usage_limit_reached", + "workspace_owner_credits_depleted", + "workspace_owner_usage_limit_reached", +]; + +/// 上游拒绝相对业务 payload 的发送阶段。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexUpstreamSendPhase { + /// DNS/TCP/TLS/WS opening 阶段，业务 payload 尚未发送。 + BeforePayload, + /// 上游已收到业务 payload，并返回了明确拒绝。 + AfterPayload, + /// 无法证明 payload 是否到达上游。 + Ambiguous, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum CodexFailureCategory { + ModelUnsupported, + CredentialExpired, + IdentityVerificationRequired, + Banned, + RateLimited, + QuotaExhausted, + CloudflareChallenge, + CloudflarePathBlocked, + InvalidRequest, + PermissionDenied, + Timeout, + Unavailable, + Transport, +} + +pub(crate) struct CodexUpstreamFailure { + pub(crate) status: Option, + pub(crate) code: Option, + pub(crate) identity_error_code: Option, + pub(crate) retry_after_seconds: Option, + pub(crate) request_id: Option, + pub(crate) set_cookie_headers: Vec, + pub(crate) rate_limit_headers: Vec<(String, String)>, + pub(crate) send_phase: CodexUpstreamSendPhase, + category: CodexFailureCategory, +} + +impl CodexUpstreamFailure { + pub(crate) fn from_response( + status: StatusCode, + body: &str, + retry_after_seconds: Option, + diagnostics: &CodexUpstreamDiagnostics, + set_cookie_headers: &[String], + rate_limit_headers: &[(String, String)], + send_phase: CodexUpstreamSendPhase, + ) -> Self { + let fields = ParsedUpstreamError::from_body(body); + let category = classify_upstream_failure( + Some(status), + body, + &fields, + diagnostics.identity_authorization_error.as_deref(), + diagnostics.identity_error_code.as_deref(), + ); + Self { + status: Some(status), + code: fields.code, + identity_error_code: diagnostics.identity_error_code.clone(), + retry_after_seconds, + request_id: diagnostics.request_id.clone(), + set_cookie_headers: set_cookie_headers.to_vec(), + rate_limit_headers: rate_limit_headers.to_vec(), + send_phase, + category, + } + } + + pub(crate) fn from_sse_failure( + failure: &ResponsesSseFailure, + diagnostics: &CodexUpstreamDiagnostics, + set_cookie_headers: &[String], + rate_limit_headers: &[(String, String)], + send_phase: CodexUpstreamSendPhase, + ) -> Self { + let fields = ParsedUpstreamError { + code: failure.upstream_code.clone(), + error_type: failure.upstream_type.clone(), + message: failure.message.clone(), + }; + let status = failure + .explicit_status_code + .and_then(|code| StatusCode::from_u16(code).ok()); + let category = classify_upstream_failure( + status, + "", + &fields, + diagnostics.identity_authorization_error.as_deref(), + diagnostics.identity_error_code.as_deref(), + ); + Self { + status, + code: fields.code, + identity_error_code: diagnostics.identity_error_code.clone(), + retry_after_seconds: failure.retry_after_seconds, + request_id: diagnostics.request_id.clone(), + set_cookie_headers: set_cookie_headers.to_vec(), + rate_limit_headers: rate_limit_headers.to_vec(), + send_phase, + category, + } + } + + pub(crate) const fn category(&self) -> CodexFailureCategory { + self.category + } + + pub(crate) const fn replay_is_safe(&self) -> bool { + match self.send_phase { + CodexUpstreamSendPhase::BeforePayload => true, + CodexUpstreamSendPhase::Ambiguous => false, + CodexUpstreamSendPhase::AfterPayload => matches!( + self.category, + CodexFailureCategory::ModelUnsupported + | CodexFailureCategory::CredentialExpired + | CodexFailureCategory::IdentityVerificationRequired + | CodexFailureCategory::Banned + | CodexFailureCategory::RateLimited + | CodexFailureCategory::QuotaExhausted + | CodexFailureCategory::CloudflareChallenge + | CodexFailureCategory::CloudflarePathBlocked + ), + } + } + + pub(crate) fn persistable_code(&self) -> Option<&'static str> { + self.code + .as_deref() + .and_then(persistable_upstream_code) + .or_else(|| { + self.identity_error_code + .as_deref() + .and_then(persistable_upstream_code) + }) + } +} + +struct ParsedUpstreamError { + code: Option, + error_type: Option, + message: String, +} + +impl ParsedUpstreamError { + fn from_body(body: &str) -> Self { + let Ok(value) = serde_json::from_str::(body) else { + return Self { + code: None, + error_type: None, + message: body.to_owned(), + }; + }; + let error = value + .pointer("/response/error") + .or_else(|| value.get("error")) + .or_else(|| value.get("detail")) + .unwrap_or(&value); + let code = error + .get("code") + .or_else(|| value.get("code")) + .and_then(Value::as_str) + .and_then(non_empty_owned); + let error_type = error + .get("type") + .or_else(|| value.get("type")) + .and_then(Value::as_str) + .and_then(non_empty_owned); + let message = error + .get("message") + .or_else(|| value.get("message")) + .and_then(Value::as_str) + .or_else(|| error.as_str()) + .and_then(non_empty_owned) + .unwrap_or_else(|| body.to_owned()); + Self { + code, + error_type, + message, + } + } +} #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct CodexUpstreamDiagnostics { @@ -129,3 +362,204 @@ fn decode_identity_error_code(encoded: &str) -> Option { .filter(|code| !code.is_empty()) .map(ToString::to_string) } + +fn classify_upstream_failure( + status: Option, + body: &str, + fields: &ParsedUpstreamError, + identity_authorization_error: Option<&str>, + identity_error_code: Option<&str>, +) -> CodexFailureCategory { + let code = normalized(fields.code.as_deref()); + let error_type = normalized(fields.error_type.as_deref()); + let identity_code = normalized(identity_error_code); + let identity_authorization = normalized(identity_authorization_error); + let message = fields.message.to_ascii_lowercase(); + let body = body.to_ascii_lowercase(); + + if [code.as_str(), message.as_str(), body.as_str()] + .into_iter() + .any(is_model_unsupported) + { + return CodexFailureCategory::ModelUnsupported; + } + let identity_signals = [ + identity_code.as_str(), + identity_authorization.as_str(), + code.as_str(), + error_type.as_str(), + message.as_str(), + body.as_str(), + ]; + if identity_signals + .into_iter() + .any(is_identity_verification_required) + { + return CodexFailureCategory::IdentityVerificationRequired; + } + if identity_signals.into_iter().any(is_banned_account) { + return CodexFailureCategory::Banned; + } + if identity_signals.into_iter().any(is_expired_credential) + || status == Some(StatusCode::UNAUTHORIZED) + { + return CodexFailureCategory::CredentialExpired; + } + if status == Some(StatusCode::FORBIDDEN) && is_cloudflare_challenge(&body) { + return CodexFailureCategory::CloudflareChallenge; + } + if status == Some(StatusCode::NOT_FOUND) && body.trim().is_empty() { + return CodexFailureCategory::CloudflarePathBlocked; + } + if [code.as_str(), error_type.as_str()] + .into_iter() + .any(is_rate_limit_signal) + || is_rate_limit_message(&message) + || is_rate_limit_message(&body) + || status == Some(StatusCode::TOO_MANY_REQUESTS) + { + return CodexFailureCategory::RateLimited; + } + if [code.as_str(), error_type.as_str()] + .into_iter() + .any(is_quota_signal) + || is_quota_message(&message) + || is_quota_message(&body) + || status == Some(StatusCode::PAYMENT_REQUIRED) + { + return CodexFailureCategory::QuotaExhausted; + } + if is_upstream_overload(&code) || is_upstream_overload(&message) { + return CodexFailureCategory::Unavailable; + } + match status.map(|status| status.as_u16()) { + Some(status) => match status { + 400 | 404 | 409 | 422 => CodexFailureCategory::InvalidRequest, + 403 => CodexFailureCategory::PermissionDenied, + 408 | 504 => CodexFailureCategory::Timeout, + 500..=599 => CodexFailureCategory::Unavailable, + _ => CodexFailureCategory::Transport, + }, + None => CodexFailureCategory::Unavailable, + } +} + +fn normalized(value: Option<&str>) -> String { + value.unwrap_or_default().trim().to_ascii_lowercase() +} + +fn non_empty_owned(value: &str) -> Option { + let value = value.trim(); + (!value.is_empty()).then(|| value.to_owned()) +} + +fn persistable_upstream_code(value: &str) -> Option<&'static str> { + PERSISTABLE_UPSTREAM_CODES + .iter() + .copied() + .find(|known| value.trim().eq_ignore_ascii_case(known)) +} + +fn is_model_unsupported(value: &str) -> bool { + value.contains("model_not_supported") + || value.contains("model_not_available") + || (value.contains("model") + && (value.contains("not supported") + || value.contains("not available") + || value.contains("not_supported") + || value.contains("not_available"))) +} + +fn is_identity_verification_required(value: &str) -> bool { + matches!( + value, + "identity_verification_required" | "verification_required" + ) || value.contains("identity verification is required") +} + +fn is_banned_account(value: &str) -> bool { + matches!( + value, + "account_banned" + | "account_deactivated" + | "account_disabled" + | "account_suspended" + | "deactivated_workspace" + | "organization_disabled" + | "workspace_deactivated" + ) || value.contains("account is banned") + || value.contains("account has been banned") + || value.contains("account deactivated") + || value.contains("account has been deactivated") + || value.contains("account disabled") + || value.contains("account has been disabled") + || value.contains("account suspended") + || value.contains("organization has been disabled") + || value.contains("workspace has been deactivated") + || value.contains("deactivated_workspace") +} + +fn is_expired_credential(value: &str) -> bool { + matches!( + value, + "token_invalid" + | "token_invalidated" + | "token_expired" + | "token_revoked" + | "refresh_token_invalidated" + | "unauthorized" + | "invalid_api_key" + | "authentication_error" + | "access_token_expired" + ) || value.contains("token revoked") + || value.contains("token invalidated") + || value.contains("token invalid") + || value.contains("token expired") + || value.contains("unauthorized") + || value.contains("invalid api key") +} + +fn is_rate_limit_signal(value: &str) -> bool { + matches!( + value, + "usage_limit_reached" + | "rate_limit_exceeded" + | "rate_limit_reached" + | "rate_limit_error" + | "workspace_owner_usage_limit_reached" + | "workspace_member_usage_limit_reached" + ) +} + +fn is_quota_signal(value: &str) -> bool { + matches!( + value, + "quota_exhausted" + | "quota_exceeded" + | "payment_required" + | "insufficient_quota" + | "workspace_owner_credits_depleted" + | "workspace_member_credits_depleted" + ) || value.starts_with("billing_limit") +} + +fn is_rate_limit_message(value: &str) -> bool { + value.contains("rate limit") || value.contains("usage limit") +} + +fn is_quota_message(value: &str) -> bool { + value.contains("quota") || value.contains("payment required") || value.contains("billing limit") +} + +fn is_cloudflare_challenge(value: &str) -> bool { + value.contains("cf-mitigated") + || value.contains("cf-chl-bypass") + || value.contains("_cf_chl") + || value.contains("cf_chl") + || value.contains("attention required") + || value.contains("just a moment") +} + +fn is_upstream_overload(value: &str) -> bool { + matches!(value, "server_is_overloaded" | "slow_down") || value.contains("server_overloaded") +} diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index 176cc7e63..a747264a8 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -17,7 +17,7 @@ pub mod usage; pub mod websocket; pub use self::{ - canonical::CodexCanonicalDecoder, + canonical::{CodexCanonicalDecoder, CodexCanonicalError}, catalog::{ CodexCatalogCapabilities, CodexCatalogCapabilityEvidence, CodexCatalogLimits, CodexCatalogMetadata, CodexCatalogModel, CodexCatalogVisibility, CodexModelCatalogError, @@ -29,7 +29,7 @@ pub use self::{ CodexRateLimitHeaderUpdates, CodexRequestContext, CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client, }, - diagnostics::CodexUpstreamDiagnostics, + diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase}, endpoints::{ CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, WHAM_USAGE_PATH, endpoint_request_path, endpoint_url, usage_endpoint_urls, diff --git a/backend/crates/providers/openai/src/transport/profile.rs b/backend/crates/providers/openai/src/transport/profile.rs index 0a51e2bb1..f97fc242c 100644 --- a/backend/crates/providers/openai/src/transport/profile.rs +++ b/backend/crates/providers/openai/src/transport/profile.rs @@ -1,8 +1,23 @@ //! Codex Desktop 上游请求画像。 use std::sync::{Arc, RwLock}; +use std::time::Duration; use chrono::{DateTime, Utc}; +use futures::{StreamExt as _, future::BoxFuture}; +use reqwest::Client; +use reqwest::redirect::Policy; +use roxmltree::{Document, Node}; +use url::Url; + +/// Codex Desktop 官方 appcast 地址。 +pub const CODEX_DESKTOP_APPCAST_URL: &str = + "https://persistent.oaistatic.com/codex-app-prod/appcast.xml"; +/// 官方制品画像检查周期。 +pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60); + +const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); +const MAX_APPCAST_BYTES: usize = 1024 * 1024; /// Codex Desktop 上游请求身份。 /// @@ -68,4 +83,302 @@ impl CodexWireProfileState { .unwrap_or_else(std::sync::PoisonError::into_inner) .clone() } + + /// 使用官方 appcast 已验证的 Desktop 版本和构建号更新运行时画像。 + pub fn update_desktop_release(&self, desktop_version: &str, desktop_build: &str) { + let mut profile = self + .profile + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if profile.desktop_version == desktop_version && profile.desktop_build == desktop_build { + return; + } + profile.desktop_version = desktop_version.to_owned(); + profile.desktop_build = desktop_build.to_owned(); + } +} + +/// 官方 appcast 中按顺序出现的首个完整 Desktop 制品。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexDesktopRelease { + pub version: String, + pub build: String, + pub published_at: Option>, + pub minimum_system_version: Option, + pub hardware_requirements: Option, + pub download_url: Option, + pub download_size: Option, + pub signature_present: bool, +} + +/// 最近一次官方 Desktop 制品检查结果。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CodexDesktopReleaseSnapshot { + pub checked_at: Option>, + pub latest: Option, + pub last_error: Option, +} + +/// Provider 内共享的 appcast 观察状态。 +#[derive(Debug, Clone, Default)] +pub struct CodexDesktopReleaseStatus { + snapshot: Arc>, +} + +impl CodexDesktopReleaseStatus { + #[must_use] + pub fn snapshot(&self) -> CodexDesktopReleaseSnapshot { + self.snapshot + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + fn record_success(&self, checked_at: DateTime, latest: CodexDesktopRelease) { + *self + .snapshot + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = CodexDesktopReleaseSnapshot { + checked_at: Some(checked_at), + latest: Some(latest), + last_error: None, + }; + } + + fn record_failure(&self, checked_at: DateTime, error: &CodexDesktopReleaseError) { + let mut snapshot = self + .snapshot + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + snapshot.checked_at = Some(checked_at); + snapshot.last_error = Some(error.to_string()); + } +} + +/// Desktop appcast 获取边界。生产实现固定访问官方 HTTPS，测试实现只替换此边界。 +pub trait CodexDesktopReleaseTransport: Send + Sync { + fn fetch(&self) -> BoxFuture<'_, Result>; +} + +/// 固定访问官方 Desktop appcast 的生产 transport。 +#[derive(Clone)] +pub struct OfficialCodexDesktopReleaseTransport { + client: Client, + endpoint: Url, +} + +impl OfficialCodexDesktopReleaseTransport { + /// 构造禁用环境代理和 redirect 的官方 HTTPS transport。 + pub fn new() -> Result { + let endpoint = Url::parse(CODEX_DESKTOP_APPCAST_URL) + .map_err(|_| CodexDesktopReleaseError::InvalidEndpoint)?; + let client = Client::builder() + .https_only(true) + .no_proxy() + .redirect(Policy::none()) + .connect_timeout(Duration::from_secs(10)) + .timeout(APPCAST_TIMEOUT) + .build() + .map_err(|_| CodexDesktopReleaseError::ClientInitialization)?; + Ok(Self { client, endpoint }) + } +} + +impl CodexDesktopReleaseTransport for OfficialCodexDesktopReleaseTransport { + fn fetch(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { + let response = self.client.get(self.endpoint.clone()).send().await?; + if !response.status().is_success() { + return Err(CodexDesktopReleaseError::HttpStatus( + response.status().as_u16(), + )); + } + if response + .content_length() + .is_some_and(|length| length > MAX_APPCAST_BYTES as u64) + { + return Err(CodexDesktopReleaseError::ResponseTooLarge); + } + let mut bytes = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk?; + if bytes + .len() + .checked_add(chunk.len()) + .is_none_or(|length| length > MAX_APPCAST_BYTES) + { + return Err(CodexDesktopReleaseError::ResponseTooLarge); + } + bytes.extend_from_slice(&chunk); + } + let xml = std::str::from_utf8(&bytes) + .map_err(|_| CodexDesktopReleaseError::InvalidDocument)?; + parse_desktop_release(xml) + }) + } +} + +/// 拉取 appcast 并原子发布请求画像的 Provider 服务。 +#[derive(Clone)] +pub struct CodexDesktopReleaseService { + transport: Arc, + status: CodexDesktopReleaseStatus, + profile: CodexWireProfileState, +} + +impl CodexDesktopReleaseService { + /// 从正式 transport seam 构造服务；生产组装只注入官方 transport。 + #[must_use] + pub fn new( + profile: CodexWireProfileState, + transport: Arc, + ) -> Self { + Self { + transport, + status: CodexDesktopReleaseStatus::default(), + profile, + } + } + + #[must_use] + pub fn status(&self) -> CodexDesktopReleaseStatus { + self.status.clone() + } + + /// 执行一次有界检查；失败只更新观察状态，不修改上一份成功画像。 + pub async fn refresh(&self) -> Result { + let checked_at = Utc::now(); + let result = self.transport.fetch().await; + match result { + Ok(release) => { + self.profile + .update_desktop_release(&release.version, &release.build); + self.status.record_success(checked_at, release.clone()); + Ok(release) + } + Err(error) => { + self.status.record_failure(checked_at, &error); + Err(error) + } + } + } +} + +/// 解析 appcast 中按顺序出现的首个完整发布项。 +pub fn parse_desktop_release(xml: &str) -> Result { + let document = Document::parse(xml).map_err(|_| CodexDesktopReleaseError::InvalidDocument)?; + for item in document + .descendants() + .filter(|node| node.is_element() && node.tag_name().name() == "item") + { + if let Some(release) = parse_release_item(item) { + return release; + } + } + Err(CodexDesktopReleaseError::MissingItem) +} + +fn parse_release_item( + item: Node<'_, '_>, +) -> Option> { + let enclosure = item + .children() + .find(|node| node.is_element() && node.tag_name().name() == "enclosure"); + let version = child_text(item, "shortVersionString") + .or_else(|| local_attribute(enclosure, "shortVersionString"))?; + if !numeric_dotted_version(&version) { + return Some(Err(CodexDesktopReleaseError::InvalidVersion)); + } + let build = child_text(item, "version").or_else(|| local_attribute(enclosure, "version"))?; + if build.is_empty() || !build.bytes().all(|byte| byte.is_ascii_digit()) { + return Some(Err(CodexDesktopReleaseError::InvalidBuild)); + } + let published_at = child_text(item, "pubDate") + .map(|value| { + DateTime::parse_from_rfc2822(&value) + .map(|value| value.to_utc()) + .map_err(|_| CodexDesktopReleaseError::InvalidPublishedAt) + }) + .transpose(); + let published_at = match published_at { + Ok(value) => value, + Err(error) => return Some(Err(error)), + }; + let download_size = local_attribute(enclosure, "length") + .map(|value| { + value + .parse::() + .map_err(|_| CodexDesktopReleaseError::InvalidDownloadSize) + }) + .transpose(); + let download_size = match download_size { + Ok(value) => value, + Err(error) => return Some(Err(error)), + }; + Some(Ok(CodexDesktopRelease { + version, + build, + published_at, + minimum_system_version: child_text(item, "minimumSystemVersion"), + hardware_requirements: child_text(item, "hardwareRequirements"), + download_url: local_attribute(enclosure, "url"), + download_size, + signature_present: local_attribute(enclosure, "edSignature").is_some(), + })) +} + +fn child_text(parent: Node<'_, '_>, local_name: &str) -> Option { + parent + .children() + .find(|node| node.is_element() && node.tag_name().name() == local_name) + .and_then(|node| node.text()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn local_attribute(node: Option>, local_name: &str) -> Option { + node? + .attributes() + .find(|attribute| attribute.name() == local_name) + .map(|attribute| attribute.value().trim()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn numeric_dotted_version(value: &str) -> bool { + let mut parts = value.split('.'); + let valid = parts + .by_ref() + .filter(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit())) + .count(); + valid >= 2 && valid == value.split('.').count() +} + +/// Desktop appcast 的稳定、无响应正文错误分类。 +#[derive(Debug, thiserror::Error)] +pub enum CodexDesktopReleaseError { + #[error("Codex Desktop appcast client initialization failed")] + ClientInitialization, + #[error("Codex Desktop appcast endpoint is invalid")] + InvalidEndpoint, + #[error("Codex Desktop appcast request failed")] + Http(#[from] reqwest::Error), + #[error("Codex Desktop appcast returned HTTP {0}")] + HttpStatus(u16), + #[error("Codex Desktop appcast response exceeded the size limit")] + ResponseTooLarge, + #[error("Codex Desktop appcast document is invalid")] + InvalidDocument, + #[error("Codex Desktop appcast contains no complete release item")] + MissingItem, + #[error("Codex Desktop appcast version is invalid")] + InvalidVersion, + #[error("Codex Desktop appcast build is invalid")] + InvalidBuild, + #[error("Codex Desktop appcast publish time is invalid")] + InvalidPublishedAt, + #[error("Codex Desktop appcast download size is invalid")] + InvalidDownloadSize, } diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index c54ecac1d..4648ac025 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeSet; +use std::{collections::BTreeSet, fmt}; use gateway_protocol::openai::{ events, @@ -12,15 +12,6 @@ const MULTI_AGENT_MODE_OPEN_TAG: &str = ""; const MULTI_AGENT_MODE_CLOSE_TAG: &str = ""; const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active."; -/// Codex Responses Lite 的 HTTP 兼容头。 -pub const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str = - "x-openai-internal-codex-responses-lite"; -/// Codex Responses Lite 在 WebSocket `client_metadata` 中的官方投影键。 -pub const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str = - "ws_request_header_x_openai_internal_codex_responses_lite"; -/// Codex memory consolidation 请求标记。 -pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request"; - /// Codex Responses 上游请求体。 /// /// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object，逐字段（含顺序、 @@ -29,7 +20,7 @@ pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request"; /// 普通客户端请求不修改 body；模型路由只写入明确受控字段。 /// /// 其余字段是代理控制状态，不进上游 body（原 `#[serde(skip)]` 字段）。 -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct CodexResponsesRequest { /// 上游请求体（唯一真相源）。 body: Map, @@ -61,8 +52,6 @@ pub struct CodexResponsesRequest { pub use_websocket: bool, /// 是否强制 HTTP SSE。 pub force_http_sse: bool, - /// 是否强制 WebSocket，且禁止在同一 attempt 内降级到 HTTP。 - pub force_websocket: bool, /// turn state 透传头。 pub turn_state: Option, /// turn metadata 透传头。 @@ -85,6 +74,23 @@ pub struct CodexResponsesRequest { pub previous_response_scope: Option, } +impl fmt::Debug for CodexResponsesRequest { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CodexResponsesRequest") + .field("body", &"") + .field("has_tuple_schema", &self.tuple_schema.is_some()) + .field("explicit_prompt_cache_key", &self.explicit_prompt_cache_key) + .field( + "has_local_conversation_id", + &self.local_conversation_id.is_some(), + ) + .field("use_websocket", &self.use_websocket) + .field("force_http_sse", &self.force_http_sse) + .finish() + } +} + /// previous response 在上游的可续接范围。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] #[serde(rename_all = "snake_case")] @@ -122,8 +128,6 @@ impl Serialize for CodexResponsesRequest { pub enum TransportRequirement { /// 客户端显式要求 HTTP。 HttpRequired, - /// Provider instance 显式冻结 WebSocket transport。 - WebSocketRequired, /// `generate=false + store=false` 预热必须保留在同一条 WebSocket。 ExplicitWebSocketWarmup, /// 只能使用持有指定 connection-local response 的精确 WebSocket。 @@ -141,9 +145,7 @@ impl TransportRequirement { pub fn requires_websocket(self) -> bool { matches!( self, - Self::WebSocketRequired - | Self::ExplicitWebSocketWarmup - | Self::ExactWebSocketContinuation + Self::ExplicitWebSocketWarmup | Self::ExactWebSocketContinuation ) } @@ -159,7 +161,6 @@ impl TransportRequirement { pub fn as_str(self) -> &'static str { match self { Self::HttpRequired => "http_required", - Self::WebSocketRequired => "websocket_required", Self::ExplicitWebSocketWarmup => "explicit_websocket_warmup", Self::ExactWebSocketContinuation => "exact_websocket_continuation", Self::PersistedContinuation => "persisted_continuation", @@ -174,17 +175,6 @@ pub fn transport_requirement(request: &CodexResponsesRequest) -> TransportRequir if !request.generate() && !request.store() { return TransportRequirement::ExplicitWebSocketWarmup; } - if request.previous_response_id().is_some() - && matches!( - request.previous_response_scope, - Some(PreviousResponseScope::ConnectionLocal) - ) - { - return TransportRequirement::ExactWebSocketContinuation; - } - if request.force_websocket { - return TransportRequirement::WebSocketRequired; - } if request.force_http_sse { return TransportRequirement::HttpRequired; } @@ -436,7 +426,7 @@ pub enum CollectedResponse { } /// Codex Responses SSE 失败事件。 -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct ResponsesSseFailure { /// SSE event 名称。 pub event: String, @@ -452,6 +442,20 @@ pub struct ResponsesSseFailure { pub retry_after_seconds: Option, } +impl fmt::Debug for ResponsesSseFailure { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ResponsesSseFailure") + .field("event", &self.event) + .field("message", &"") + .field("has_upstream_code", &self.upstream_code.is_some()) + .field("has_upstream_type", &self.upstream_type.is_some()) + .field("explicit_status_code", &self.explicit_status_code) + .field("retry_after_seconds", &self.retry_after_seconds) + .finish() + } +} + impl ResponsesSseFailure { pub fn from_event(event: &str, value: &Value) -> Self { Self { @@ -687,7 +691,6 @@ impl CodexResponsesRequest { client_api_key_id: None, use_websocket: false, force_http_sse: false, - force_websocket: false, turn_state: None, turn_metadata: None, beta_features: None, @@ -968,10 +971,6 @@ pub(crate) fn codex_request_semantics_from_parts( } } -pub(crate) fn proactive_multi_agent_mode_from_text(text: &str) -> Option { - multi_agent_mode_from_text(text).map(|mode| mode.starts_with(PROACTIVE_MULTI_AGENT_MODE_PREFIX)) -} - fn latest_multi_agent_mode<'a>( input: impl DoubleEndedIterator, ) -> Option<&'a str> { diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 7d223fe8f..8ae1ae01b 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -3,14 +3,12 @@ use base64::Engine as _; use gateway_core::operation::{ ContentPart, GenerateRequest, ImageSource, MessageRole, OutputFormat, ReasoningEffort, - ReasoningSummary, + ReasoningSummary, ResponsePersistence, }; use serde_json::{Map, Value, json}; +use sha2::{Digest, Sha256}; -use crate::transport::protocol::responses::{ - CodexRequestSemantics, CodexResponsesRequest, codex_request_semantics_from_parts, - proactive_multi_agent_mode_from_text, -}; +use crate::transport::protocol::responses::{CodexRequestSemantics, CodexResponsesRequest}; const CODEX_OPTION_FIELDS: &[&str] = &[ "beta_features", @@ -33,6 +31,64 @@ const CODEX_OPTION_FIELDS: &[&str] = &[ "version", ]; +const CROSS_ACCOUNT_IDENTITY_KEYS: &[&str] = &[ + "authorization", + "Authorization", + "cookie", + "Cookie", + "chatgpt-account-id", + "chatgpt_account_id", + "chatgptAccountId", + "account_id", + "accountId", + "user_id", + "userId", + "chatgpt_user_id", + "chatgptUserId", + "access_token", + "accessToken", + "session_token", + "sessionToken", + "refresh_token", + "refreshToken", + "id_token", + "idToken", + "token", + "cookies", + "cookie_header", + "cookieHeader", + "cf_clearance", +]; + +const ACCOUNT_BOUND_STATE_KEYS: &[&str] = &[ + "turnState", + "turn_state", + "x-codex-turn-state", + "turnMetadata", + "turn_metadata", + "x-codex-turn-metadata", + "previous_response_id", + "previousResponseId", + "response_id", + "responseId", + "conversation", + "conversation_id", + "conversationId", +]; + +const UNTRUSTED_CONTINUATION_KEYS: &[&str] = &[ + "previous_response_id", + "previousResponseId", + "response_id", + "responseId", +]; + +const INSTALLATION_ID_KEYS: &[&str] = &[ + "installation_id", + "installationId", + "x-codex-installation-id", +]; + /// Provider 专属编码错误；不保存 prompt、schema 或 option 值。 #[derive(Debug, thiserror::Error, PartialEq, Eq)] pub enum CodexRequestEncodeError { @@ -48,16 +104,27 @@ pub fn encode_generate_request( request: &GenerateRequest, upstream_model: &str, ) -> Result { - let mut body = Map::new(); + let (mut body, protocol_payload) = match request.protocol_payload() { + Some(payload) if payload.protocol() == "openai" => (payload.body().clone(), true), + Some(_) => return Err(CodexRequestEncodeError::InvalidProviderOptions), + None => (Map::new(), false), + }; body.insert("model".to_owned(), Value::String(upstream_model.to_owned())); - body.insert("input".to_owned(), Value::Array(encode_messages(request)?)); - // ProviderStream 必须统一从 live canonical stream 解码；非流式仅由 client - // protocol collector 聚合，不能另开一条 Provider 路径。 - body.insert("stream".to_owned(), Value::Bool(true)); - // 客户端每轮携带完整历史；OAuth 上游不代存正文。 - body.insert("store".to_owned(), Value::Bool(false)); + if !protocol_payload { + body.insert("input".to_owned(), Value::Array(encode_messages(request)?)); + } + if !protocol_payload { + body.insert("stream".to_owned(), Value::Bool(true)); + body.insert( + "store".to_owned(), + Value::Bool(matches!( + request.response_persistence(), + ResponsePersistence::Store + )), + ); + } - if !request.tools().is_empty() { + if !protocol_payload && !request.tools().is_empty() { body.insert( "tools".to_owned(), Value::Array( @@ -86,13 +153,13 @@ pub fn encode_generate_request( ); } - if !matches!(request.output_format(), OutputFormat::Text) { + if !protocol_payload && !matches!(request.output_format(), OutputFormat::Text) { body.insert( "text".to_owned(), encode_output_format(request.output_format()), ); } - if let Some(reasoning) = request.reasoning() { + if !protocol_payload && let Some(reasoning) = request.reasoning() { let mut value = Map::new(); if let Some(effort) = reasoning.effort { value.insert( @@ -108,11 +175,16 @@ pub fn encode_generate_request( } body.insert("reasoning".to_owned(), Value::Object(value)); } - if let Some(tokens) = request.max_output_tokens() { + if !protocol_payload && let Some(tokens) = request.max_output_tokens() { body.insert("max_output_tokens".to_owned(), Value::from(tokens)); } + if !protocol_payload && let Some(key) = request.prompt_cache_key() { + body.insert("prompt_cache_key".to_owned(), Value::String(key.to_owned())); + } let mut encoded = CodexResponsesRequest::from_body(body); + encoded.explicit_prompt_cache_key = encoded.prompt_cache_key().is_some(); + extract_request_context(&mut encoded); if let Some(options) = request.provider_options().get("openai") { apply_codex_options(&mut encoded, options)?; } @@ -122,30 +194,9 @@ pub fn encode_generate_request( /// 复用正式 Codex 编码规则提取请求诊断语义；编码失败时不伪造 Provider 事实。 #[must_use] pub fn codex_request_semantics(request: &GenerateRequest) -> CodexRequestSemantics { - let turn_metadata = request - .provider_options() - .get("openai") - .and_then(|options| options.get("turn_metadata")) - .and_then(Value::as_str); - let effort = request - .reasoning() - .and_then(|reasoning| reasoning.effort) - .map(reasoning_effort); - let proactive_multi_agent = request - .messages() - .iter() - .rev() - .filter(|message| message.role() == MessageRole::Developer) - .find_map(|message| { - message.content().iter().rev().find_map(|content| { - let ContentPart::Text(text) = content else { - return None; - }; - proactive_multi_agent_mode_from_text(text) - }) - }) - .unwrap_or(false); - codex_request_semantics_from_parts(turn_metadata, effort, proactive_multi_agent, false) + encode_generate_request(request, "observability") + .map(|request| request.semantics()) + .unwrap_or_default() } fn encode_messages(request: &GenerateRequest) -> Result, CodexRequestEncodeError> { @@ -243,6 +294,462 @@ fn encode_output_format(format: &OutputFormat) -> Value { json!({"format": format}) } +fn extract_request_context(request: &mut CodexResponsesRequest) { + let context = ExtractedRequestContext::from_body(request.body()); + request.turn_state = context.turn_state; + request.turn_metadata = context.turn_metadata; + request.beta_features = context.beta_features; + request.version = context.version; + request.include_timing_metrics = context.include_timing_metrics; + request.codex_window_id = context.codex_window_id; + request.parent_thread_id = context.parent_thread_id; + request.client_conversation_id = context.conversation_id; + request.client_session_id = context.session_id; + request.client_thread_id = context.thread_id; + request.client_request_id = context.client_request_id; + request.client_turn_id = context.turn_id; + request.responses_lite = context.responses_lite; + request.memgen_request = context.memgen_request; +} + +struct ExtractedRequestContext { + turn_state: Option, + turn_metadata: Option, + beta_features: Option, + version: Option, + include_timing_metrics: Option, + codex_window_id: Option, + parent_thread_id: Option, + conversation_id: Option, + session_id: Option, + thread_id: Option, + client_request_id: Option, + turn_id: Option, + responses_lite: Option, + memgen_request: Option, +} + +impl ExtractedRequestContext { + fn from_body(body: &Map) -> Self { + Self { + turn_state: body_or_metadata_string( + body, + &["turnState", "turn_state", "x-codex-turn-state"], + ), + turn_metadata: body_or_metadata_string( + body, + &["turnMetadata", "turn_metadata", "x-codex-turn-metadata"], + ), + beta_features: body_or_metadata_string(body, &["betaFeatures", "beta_features"]), + version: body_or_metadata_string(body, &["version"]), + include_timing_metrics: body_or_metadata_string( + body, + &["includeTimingMetrics", "include_timing_metrics"], + ), + codex_window_id: body_or_metadata_string( + body, + &["codexWindowId", "codex_window_id", "x-codex-window-id"], + ), + parent_thread_id: body_or_metadata_string( + body, + &[ + "parentThreadId", + "parent_thread_id", + "x-codex-parent-thread-id", + ], + ), + conversation_id: body_or_metadata_string(body, &["conversation_id", "conversationId"]), + session_id: body_or_metadata_string(body, &["session_id", "sessionId"]), + thread_id: body_or_metadata_string(body, &["thread_id", "threadId"]), + client_request_id: body_or_metadata_string( + body, + &[ + "x-client-request-id", + "client_request_id", + "clientRequestId", + ], + ), + turn_id: body_or_metadata_string(body, &["turn_id", "turnId", "x-codex-turn-id"]), + responses_lite: body_or_metadata_string( + body, + &["ws_request_header_x_openai_internal_codex_responses_lite"], + ), + memgen_request: body_or_metadata_string( + body, + &["x-openai-memgen-request", "memgen_request"], + ), + } + } +} + +fn body_or_metadata_string(body: &Map, keys: &[&str]) -> Option { + keys.iter() + .find_map(|key| non_empty_string(body.get(*key))) + .or_else(|| { + let metadata = body.get("client_metadata")?.as_object()?; + keys.iter() + .find_map(|key| non_empty_string(metadata.get(*key))) + }) +} + +fn non_empty_string(value: Option<&Value>) -> Option { + value + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn derive_account_scoped_local_conversation_id( + request: &mut CodexResponsesRequest, + installation_id: &str, +) { + let anchor = request + .prompt_cache_key() + .map(|value| ("prompt-cache", value.to_owned())) + .or_else(|| { + request + .client_session_id + .as_deref() + .map(|value| ("session", value.to_owned())) + }) + .or_else(|| { + request + .client_thread_id + .as_deref() + .map(|value| ("thread", value.to_owned())) + }) + .or_else(|| { + request + .client_conversation_id + .as_deref() + .map(|value| ("conversation", value.to_owned())) + }) + .or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value))); + let Some((domain, value)) = anchor else { + return; + }; + let digest = hmac_sha256( + installation_id.as_bytes(), + &[ + b"codex-local-conversation-v2\0", + domain.as_bytes(), + b"\0", + value.as_bytes(), + ], + ); + request.local_conversation_id = Some(format!("lc_{}", hex::encode(digest))); +} + +const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; +const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; + +fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option { + let instructions = request + .instructions() + .chars() + .take(2_000) + .collect::(); + let first_user_text = first_user_text(request.input()); + let normalized = normalize_conversation_anchor_text(&first_user_text); + let first_user_text = if normalized.is_empty() { + first_user_text + } else { + normalized + }; + if instructions.is_empty() && first_user_text.is_empty() { + return None; + } + + let mut hasher = Sha256::new(); + hasher.update(request.model().as_bytes()); + hasher.update(b"\0"); + hasher.update(instructions.as_bytes()); + hasher.update(b"\0"); + hasher.update(first_user_text.as_bytes()); + let digest = hex::encode(hasher.finalize()); + Some(format!( + "{}-{}-{}-{}-{}", + &digest[0..8], + &digest[8..12], + &digest[12..16], + &digest[16..20], + &digest[20..32] + )) +} + +fn first_user_text(input: &[Value]) -> String { + for item in input { + if item.get("role").and_then(Value::as_str) != Some("user") { + continue; + } + let Some(content) = item.get("content") else { + return String::new(); + }; + if let Some(text) = content.as_str() { + return text.to_owned(); + } + if let Some(parts) = content.as_array() { + return parts + .iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("input_text")) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .collect(); + } + return String::new(); + } + String::new() +} + +fn normalize_conversation_anchor_text(text: &str) -> String { + let mut rest = text.trim_start(); + loop { + let lower = rest.to_ascii_lowercase(); + if !lower.starts_with(LEADING_SYSTEM_REMINDER_OPEN) { + break; + } + let Some(close_start) = lower.find(LEADING_SYSTEM_REMINDER_CLOSE) else { + break; + }; + rest = rest[close_start + LEADING_SYSTEM_REMINDER_CLOSE.len()..].trim_start(); + } + rest.to_owned() +} + +fn hmac_sha256(key: &[u8], message: &[&[u8]]) -> [u8; 32] { + const BLOCK_BYTES: usize = 64; + + let mut key_block = [0_u8; BLOCK_BYTES]; + if key.len() > BLOCK_BYTES { + key_block[..32].copy_from_slice(&Sha256::digest(key)); + } else { + key_block[..key.len()].copy_from_slice(key); + } + let mut inner_pad = key_block; + let mut outer_pad = key_block; + for byte in &mut inner_pad { + *byte ^= 0x36; + } + for byte in &mut outer_pad { + *byte ^= 0x5c; + } + + let mut inner = Sha256::new(); + inner.update(inner_pad); + for part in message { + inner.update(part); + } + let inner_digest = inner.finalize(); + + let mut outer = Sha256::new(); + outer.update(outer_pad); + outer.update(inner_digest); + outer.finalize().into() +} + +/// 把客户端正文收敛到当前 lease 的账号身份边界。 +/// +/// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入 +/// 请求头；正文只替换客户端原本声明过的 installation 字段，绝不接受客户端 +/// 提供的 token、cookie 或账号身份。 +pub(crate) fn scope_request_to_account( + request: &mut CodexResponsesRequest, + installation_id: &str, + cross_account: bool, +) { + for key in CROSS_ACCOUNT_IDENTITY_KEYS { + request.body_mut().remove(*key); + } + for key in UNTRUSTED_CONTINUATION_KEYS { + request.body_mut().remove(*key); + } + + let client_metadata_turn_state = metadata_string(request, "x-codex-turn-state"); + let preserve_turn_state = + !cross_account && (request.turn_state.is_some() || client_metadata_turn_state.is_some()); + let turn_state = preserve_turn_state + .then(|| request.turn_state.clone()) + .flatten(); + let client_metadata_turn_state = if preserve_turn_state { + client_metadata_turn_state + } else { + None + }; + let turn_metadata = request + .turn_metadata + .as_deref() + .and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account)); + let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata") + .and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account)); + + if cross_account { + sanitize_cross_account_input(request); + for key in ACCOUNT_BOUND_STATE_KEYS { + request.body_mut().remove(*key); + } + } + + for key in INSTALLATION_ID_KEYS { + request.replace_existing_identity_field(key, Some(installation_id)); + } + for (key, value) in [ + ("turnState", turn_state.as_deref()), + ("x-codex-turn-state", turn_state.as_deref()), + ("turnMetadata", turn_metadata.as_deref()), + ("x-codex-turn-metadata", turn_metadata.as_deref()), + ] { + request.replace_existing_identity_field(key, value); + } + + if let Some(client_metadata) = request.client_metadata().cloned() { + let scoped = match client_metadata { + Value::Object(mut metadata) => { + let scoped_turn_metadata = + ["turnMetadata", "turn_metadata", "x-codex-turn-metadata"].map(|key| { + ( + key, + metadata.get(key).and_then(Value::as_str).and_then(|value| { + scope_turn_metadata(value, installation_id, cross_account) + }), + ) + }); + for key in CROSS_ACCOUNT_IDENTITY_KEYS { + metadata.remove(*key); + } + if cross_account { + for key in ACCOUNT_BOUND_STATE_KEYS { + metadata.remove(*key); + } + } + metadata.insert( + "x-codex-installation-id".to_owned(), + Value::String(installation_id.to_owned()), + ); + replace_existing_metadata_field( + &mut metadata, + "installation_id", + Some(installation_id), + ); + replace_existing_metadata_field( + &mut metadata, + "installationId", + Some(installation_id), + ); + replace_metadata_field( + &mut metadata, + "x-codex-turn-state", + client_metadata_turn_state.as_deref(), + ); + replace_metadata_field( + &mut metadata, + "x-codex-turn-metadata", + client_metadata_turn_metadata.as_deref(), + ); + if !cross_account { + for (key, value) in scoped_turn_metadata { + replace_existing_metadata_field(&mut metadata, key, value.as_deref()); + } + } + (!metadata.is_empty()).then_some(Value::Object(metadata)) + } + value if !cross_account => Some(value), + Value::Null => None, + _ => None, + }; + request.set_client_metadata(scoped); + } + + request.turn_state = turn_state; + request.turn_metadata = turn_metadata; + derive_account_scoped_local_conversation_id(request, installation_id); +} + +fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) { + if request.input().is_empty() { + return; + } + let input = request + .input() + .iter() + .cloned() + .filter_map(sanitize_cross_account_item) + .collect(); + request.set_input(input); +} + +pub(crate) fn sanitize_cross_account_item(mut item: Value) -> Option { + if let Value::Object(object) = &mut item { + if matches!( + object.get("type").and_then(Value::as_str), + Some("compaction" | "compaction_summary" | "context_compaction") + ) { + return None; + } + object.remove("id"); + object.remove("encrypted_content"); + } + Some(item) +} + +fn metadata_string(request: &CodexResponsesRequest, key: &str) -> Option { + request + .client_metadata()? + .as_object()? + .get(key)? + .as_str() + .map(ToOwned::to_owned) +} + +fn scope_turn_metadata(raw: &str, installation_id: &str, cross_account: bool) -> Option { + let Ok(Value::Object(mut metadata)) = serde_json::from_str::(raw) else { + return (!cross_account).then(|| raw.to_owned()); + }; + let mut changed = false; + for key in CROSS_ACCOUNT_IDENTITY_KEYS { + changed |= metadata.remove(*key).is_some(); + } + if cross_account { + for key in ACCOUNT_BOUND_STATE_KEYS { + changed |= metadata.remove(*key).is_some(); + } + } + if !cross_account + && !changed + && !INSTALLATION_ID_KEYS + .iter() + .any(|key| metadata.contains_key(*key)) + { + return Some(raw.to_owned()); + } + for key in INSTALLATION_ID_KEYS { + if metadata.contains_key(*key) { + metadata.insert((*key).to_owned(), Value::String(installation_id.to_owned())); + } + } + serde_json::to_string(&metadata).ok() +} + +fn replace_existing_metadata_field( + metadata: &mut Map, + key: &str, + value: Option<&str>, +) { + if metadata.contains_key(key) { + replace_metadata_field(metadata, key, value); + } +} + +fn replace_metadata_field(metadata: &mut Map, key: &str, value: Option<&str>) { + match value.filter(|value| !value.trim().is_empty()) { + Some(value) => { + metadata.insert(key.to_owned(), Value::String(value.to_owned())); + } + None => { + metadata.remove(key); + } + } +} + fn apply_codex_options( request: &mut CodexResponsesRequest, options: &Map, @@ -257,21 +764,58 @@ fn apply_codex_options( return Err(CodexRequestEncodeError::InvalidProviderOptions); } - request.turn_state = optional_string(options, "turn_state")?; - request.turn_metadata = optional_string(options, "turn_metadata")?; - request.beta_features = optional_string(options, "beta_features")?; - request.version = optional_string(options, "version")?; - request.include_timing_metrics = optional_string(options, "include_timing_metrics")?; - request.codex_window_id = optional_string(options, "codex_window_id")?; - request.parent_thread_id = optional_string(options, "parent_thread_id")?; - request.client_conversation_id = optional_string(options, "conversation_id")?; - request.local_conversation_id = request.client_conversation_id.clone(); - request.client_session_id = optional_string(options, "session_id")?; - request.client_thread_id = optional_string(options, "thread_id")?; - request.client_request_id = optional_string(options, "client_request_id")?; - request.client_turn_id = optional_string(options, "turn_id")?; - request.responses_lite = optional_string(options, "responses_lite")?; - request.memgen_request = optional_string(options, "memgen_request")?; + request.turn_state = request + .turn_state + .take() + .or(optional_string(options, "turn_state")?); + request.turn_metadata = request + .turn_metadata + .take() + .or(optional_string(options, "turn_metadata")?); + request.beta_features = request + .beta_features + .take() + .or(optional_string(options, "beta_features")?); + request.version = request + .version + .take() + .or(optional_string(options, "version")?); + request.include_timing_metrics = request + .include_timing_metrics + .take() + .or(optional_string(options, "include_timing_metrics")?); + request.codex_window_id = request + .codex_window_id + .take() + .or(optional_string(options, "codex_window_id")?); + request.parent_thread_id = request + .parent_thread_id + .take() + .or(optional_string(options, "parent_thread_id")?); + request.client_conversation_id = request + .client_conversation_id + .take() + .or(optional_string(options, "conversation_id")?); + request.client_session_id = request + .client_session_id + .take() + .or(optional_string(options, "session_id")?); + request.client_thread_id = request + .client_thread_id + .take() + .or(optional_string(options, "thread_id")?); + request.client_request_id = request + .client_request_id + .take() + .or(optional_string(options, "client_request_id")?); + request.client_turn_id = request + .client_turn_id + .take() + .or(optional_string(options, "turn_id")?); + request.responses_lite = + optional_string(options, "responses_lite")?.or_else(|| request.responses_lite.take()); + request.memgen_request = + optional_string(options, "memgen_request")?.or_else(|| request.memgen_request.take()); match optional_string(options, "transport")?.as_deref() { Some("http_sse") => encoded_transport(request, false), Some("websocket") => encoded_transport(request, true), @@ -281,6 +825,9 @@ fn apply_codex_options( for field in ["prompt_cache_key", "service_tier"] { if let Some(value) = optional_string(options, field)? { + if field == "prompt_cache_key" && request.prompt_cache_key().is_some() { + return Err(CodexRequestEncodeError::InvalidProviderOptions); + } request .body_mut() .insert(field.to_owned(), Value::String(value)); @@ -291,7 +838,6 @@ fn apply_codex_options( fn encoded_transport(request: &mut CodexResponsesRequest, websocket: bool) { request.force_http_sse = !websocket; - request.force_websocket = websocket; request.use_websocket = websocket; } diff --git a/backend/crates/providers/openai/src/transport/usage.rs b/backend/crates/providers/openai/src/transport/usage.rs index e52ff756a..78b8d8cee 100644 --- a/backend/crates/providers/openai/src/transport/usage.rs +++ b/backend/crates/providers/openai/src/transport/usage.rs @@ -1,4 +1,7 @@ -use gateway_core::accounting::{CalculatedCostBreakdown, CurrencyCode, Decimal, Money}; +use gateway_core::accounting::{ + CalculatedCostAmounts, CalculatedCostBreakdown, CalculatedCostRates, CurrencyCode, Decimal, + Money, +}; use gateway_protocol::openai::events::retry_after_seconds_from_body; use reqwest::StatusCode; use serde_json::Value; @@ -248,16 +251,20 @@ pub fn openai_billing_breakdown( let total_amount = apply_percent(standard_amount, multiplier_percent)?; Some(CalculatedCostBreakdown::new( - usd_money(input_amount)?, - usd_money(output_amount)?, - usd_money(cache_read_amount)?, - usd_money(cache_write_amount)?, - usd_money(standard_amount)?, - usd_money(total_amount)?, - usd_price_per_million(rates.input_ticks)?, - usd_price_per_million(rates.output_ticks)?, - usd_price_per_million(rates.cache_read_ticks)?, - usd_price_per_million(cache_write_rate)?, + CalculatedCostAmounts::new( + usd_money(input_amount)?, + usd_money(output_amount)?, + usd_money(cache_read_amount)?, + usd_money(cache_write_amount)?, + usd_money(standard_amount)?, + usd_money(total_amount)?, + ), + CalculatedCostRates::new( + usd_price_per_million(rates.input_ticks)?, + usd_price_per_million(rates.output_ticks)?, + usd_price_per_million(rates.cache_read_ticks)?, + usd_price_per_million(cache_write_rate)?, + ), Some(normalized_tier.unwrap_or_else(|| "default".to_owned())), multiplier_percent, )) @@ -344,6 +351,8 @@ impl CodexBackendClient { set_cookie_headers: Vec::new(), rate_limit_headers: Vec::new(), transport: super::client::CodexBackendTransport::HttpSse, + transport_metrics: Box::default(), + send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload, }); } let body = body.into_string(); @@ -362,6 +371,8 @@ impl CodexBackendClient { set_cookie_headers: Vec::new(), rate_limit_headers: Vec::new(), transport: super::client::CodexBackendTransport::HttpSse, + transport_metrics: Box::default(), + send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload, }); } @@ -384,6 +395,8 @@ impl CodexBackendClient { set_cookie_headers: Vec::new(), rate_limit_headers: Vec::new(), transport: super::client::CodexBackendTransport::HttpSse, + transport_metrics: Box::default(), + send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload, }) } } diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 2e83ba74d..ef50fd9ab 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -1,6 +1,6 @@ //! Responses WebSocket pre-send/post-send 与 pool/breaker 编排。 -use std::time::{Duration, Instant, SystemTime}; +use std::time::{Duration, Instant}; use tokio::{sync::oneshot, time::timeout}; @@ -35,7 +35,6 @@ pub(crate) const WEBSOCKET_FAST_PATH_BUDGET: Duration = Duration::from_millis(80 pub(crate) struct PreparedWebSocket { connection: PooledWebSocketConnection, binding: PoolBinding, - connection_local_continuation_expires_at: Option, connect_elapsed: Option, decision_wait_elapsed: Duration, initial_event_timeout: Option, @@ -145,8 +144,6 @@ pub(crate) async fn prepare_response_create_request_with_pool( PreviousResponseUnavailableReason::LatestResponseMismatch, )); } - let connection_local_continuation_expires_at = - pool.connection_expires_at(connection.created_at); Ok(PreparedWebSocket { connection: *connection, binding: PoolBinding::Pooled { @@ -154,7 +151,6 @@ pub(crate) async fn prepare_response_create_request_with_pool( reused: true, decision: WebSocketPoolDecision::reuse(), }, - connection_local_continuation_expires_at, connect_elapsed: None, decision_wait_elapsed: decision_started_at.elapsed(), initial_event_timeout: pool.initial_event_timeout(), @@ -219,7 +215,6 @@ async fn prepare_unpooled_websocket( Ok(PreparedWebSocket { connection, binding: PoolBinding::Unpooled, - connection_local_continuation_expires_at: None, connect_elapsed: Some(connect_elapsed), decision_wait_elapsed: decision_started_at.elapsed(), initial_event_timeout, @@ -249,8 +244,6 @@ async fn prepare_pooled_websocket( permit, ); let handoff = waiter.wait(fast_path_budget).await?; - let connection_local_continuation_expires_at = - pool.connection_expires_at(handoff.connection.created_at); Ok(PreparedWebSocket { connection: *handoff.connection, binding: PoolBinding::Pooled { @@ -258,7 +251,6 @@ async fn prepare_pooled_websocket( reused: false, decision: WebSocketPoolDecision::new(), }, - connection_local_continuation_expires_at, connect_elapsed: Some(handoff.connect_elapsed), decision_wait_elapsed: decision_started_at.elapsed(), initial_event_timeout: pool.initial_event_timeout(), @@ -508,7 +500,6 @@ pub(crate) async fn execute_prepared_response_create_request( let PreparedWebSocket { connection, binding, - connection_local_continuation_expires_at, initial_event_timeout, .. } = prepared; @@ -523,10 +514,7 @@ pub(crate) async fn execute_prepared_response_create_request( discard_after_send(websocket, lease).await; return Err(post_send_ambiguous(error)); } - let connection_local_continuation_expires_at = lease - .is_some() - .then_some(connection_local_continuation_expires_at) - .flatten(); + let connection_local_available = lease.is_some(); let collected = collect_websocket_response( websocket, metadata, @@ -552,7 +540,7 @@ pub(crate) async fn execute_prepared_response_create_request( } }; exchange.pool_decision = pool_decision; - exchange.connection_local_continuation_expires_at = connection_local_continuation_expires_at; + exchange.connection_local_continuation = connection_local_available; match (terminal, lease) { (WebSocketTerminalKind::Completed, Some(lease)) => { lease @@ -580,7 +568,6 @@ pub(crate) async fn execute_prepared_response_create_request_stream( let PreparedWebSocket { connection, binding, - connection_local_continuation_expires_at, initial_event_timeout, .. } = prepared; @@ -595,10 +582,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream( discard_after_send(websocket, lease).await; return Err(post_send_ambiguous(error)); } - let connection_local_continuation_expires_at = lease - .is_some() - .then_some(connection_local_continuation_expires_at) - .flatten(); + let connection_local_available = lease.is_some(); let pool_return = lease.map(|lease| WebSocketStreamPoolReturn { lease, created_at, @@ -612,7 +596,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream( initial_event_timeout, ); exchange.pool_decision = pool_decision; - exchange.connection_local_continuation_expires_at = connection_local_continuation_expires_at; + exchange.connection_local_continuation = connection_local_available; Ok(exchange) } diff --git a/backend/crates/providers/openai/src/transport/websocket/error.rs b/backend/crates/providers/openai/src/transport/websocket/error.rs index ea0ee9879..86d8b401e 100644 --- a/backend/crates/providers/openai/src/transport/websocket/error.rs +++ b/backend/crates/providers/openai/src/transport/websocket/error.rs @@ -6,6 +6,7 @@ use gateway_protocol::openai::sse::SseError; use thiserror::Error; use crate::transport::diagnostics::CodexUpstreamDiagnostics; +use crate::transport::diagnostics::CodexUpstreamSendPhase; use super::PreviousResponseUnavailableReason; @@ -113,6 +114,8 @@ pub struct CodexWebSocketUpstreamError { pub set_cookie_headers: Vec, /// 上游诊断元数据。 pub diagnostics: CodexUpstreamDiagnostics, + /// 上游拒绝相对业务 payload 的发送阶段。 + pub send_phase: CodexUpstreamSendPhase, } impl fmt::Display for CodexWebSocketUpstreamError { @@ -132,6 +135,7 @@ impl CodexWebSocketExchangeError { body: String, set_cookie_headers: Vec, diagnostics: CodexUpstreamDiagnostics, + send_phase: CodexUpstreamSendPhase, ) -> Self { Self::Upstream(Box::new(CodexWebSocketUpstreamError { status_code, @@ -139,6 +143,7 @@ impl CodexWebSocketExchangeError { body, set_cookie_headers, diagnostics, + send_phase, })) } diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs index a359ab813..eb85031da 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs @@ -107,7 +107,7 @@ pub(in crate::transport::websocket) async fn collect_websocket_response( first_text_ms, first_event_ms, pool_decision: None, - connection_local_continuation_expires_at: None, + connection_local_continuation: false, diagnostics: metadata.diagnostics.clone(), response_metadata: metadata.response_metadata.clone(), }; diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs index 0790ff976..d938d3499 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs @@ -8,7 +8,7 @@ mod stream; use std::{ pin::Pin, sync::Arc, - time::{Duration, Instant, SystemTime}, + time::{Duration, Instant}, }; use bytes::Bytes; @@ -62,8 +62,8 @@ pub struct CodexWebSocketExchange { pub first_event_ms: Option, /// WebSocket 连接池决策。 pub pool_decision: Option, - /// terminal completed 后该 socket 保留 connection-local continuation 的最晚时间。 - pub connection_local_continuation_expires_at: Option, + /// terminal completed 后该 socket 是否会保留 connection-local continuation。 + pub connection_local_continuation: bool, /// 上游诊断元数据。 pub diagnostics: CodexUpstreamDiagnostics, /// 安全响应元数据。 @@ -88,8 +88,8 @@ pub struct CodexWebSocketStreamingExchange { pub turn_state_update: CodexWebSocketTurnStateUpdate, /// WebSocket 连接池决策。 pub pool_decision: Option, - /// terminal completed 后该 socket 保留 connection-local continuation 的最晚时间。 - pub connection_local_continuation_expires_at: Option, + /// terminal completed 后该 socket 是否会保留 connection-local continuation。 + pub connection_local_continuation: bool, /// 上游诊断元数据。 pub diagnostics: CodexUpstreamDiagnostics, /// 安全响应元数据。 @@ -125,7 +125,6 @@ pub async fn execute_response_create_request( pub(super) fn reusable_websocket_metadata( mut metadata: CodexWebSocketConnectionMetadata, ) -> CodexWebSocketConnectionMetadata { - metadata.set_cookie_headers.clear(); metadata.rate_limit_headers.clear(); metadata } diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index d356bdf46..5258b8b6b 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -110,7 +110,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response( rate_limit_header_updates, turn_state_update, pool_decision: None, - connection_local_continuation_expires_at: None, + connection_local_continuation: false, diagnostics: response_metadata.diagnostics, response_metadata: response_metadata.response_metadata, } diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index 638ffe370..e5661dd2c 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -17,7 +17,9 @@ use crate::{ transport::protocol::{ responses::CodexResponsesRequest, websocket::websocket_response_create_payload_text, }, - transport::{endpoints::CODEX_RESPONSES_PATH, response_meta, tls}, + transport::{ + diagnostics::CodexUpstreamSendPhase, endpoints::CODEX_RESPONSES_PATH, response_meta, tls, + }, }; use super::{ @@ -192,6 +194,7 @@ fn websocket_opening_error(response: &WsResponse>>) -> CodexWebSo body, response_meta::set_cookie_headers(response.headers()), response_meta::diagnostics(Some(status_code), response.headers()), + CodexUpstreamSendPhase::BeforePayload, ) } diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs index f0dda24af..1d0a3f155 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs @@ -7,7 +7,7 @@ mod supervisor; use std::{ future::Future, sync::{Arc, Mutex, MutexGuard, atomic::AtomicBool}, - time::{Duration, SystemTime}, + time::Duration, }; use serde_json::{Value, json}; @@ -150,15 +150,6 @@ impl CodexWebSocketPool { self.config.initial_event_timeout } - /// 返回指定连接在池策略下的绝对失效时间。 - pub(crate) fn connection_expires_at(&self, created_at: Instant) -> Option { - let remaining = self.config.max_age.checked_sub(created_at.elapsed())?; - if remaining.is_zero() { - return None; - } - SystemTime::now().checked_add(remaining) - } - /// 注册由连接池生命周期托管的 opening 任务。 pub(crate) fn spawn_connect_task(&self, future: impl Future + Send + 'static) { drop(self.tasks.spawn(future)); @@ -178,7 +169,7 @@ impl CodexWebSocketPool { } let key = if let Some(response_id) = required_response_id { let Some(key) = state.slots.iter().find_map(|(candidate, slot)| { - (candidate.same_upstream_owner(key) + (candidate.same_logical_connection(key) && slot.latest_response_id() == Some(response_id)) .then(|| candidate.clone()) }) else { diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs index 7563a7f4f..87c03c7ee 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs @@ -60,12 +60,10 @@ impl CodexWebSocketPoolKey { ]) } - /// 原生 continuation 不依赖客户端 conversation ID，只依赖上游 origin、 - /// 已冻结账号和握手画像；具体连接再由 latest_response_id 精确确认。 - pub(super) fn same_upstream_owner(&self, other: &Self) -> bool { + pub(super) fn same_logical_connection(&self, other: &Self) -> bool { self.base_url == other.base_url && self.account_id == other.account_id - && self.connection_profile == other.connection_profile + && self.conversation_id == other.conversation_id } } diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs new file mode 100644 index 000000000..7ece42150 --- /dev/null +++ b/backend/crates/providers/openai/tests/admin.rs @@ -0,0 +1,591 @@ +use std::collections::BTreeMap; +use std::num::NonZeroU32; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime}; + +use chrono::{DateTime, TimeZone as _, Utc}; +use futures::future::BoxFuture; +use gateway_admin::model::accounts::{ + AccountAvailability as AdminAccountAvailability, AccountRecord, +}; +use gateway_admin::model::observability::{ + CurrencyCost, DesktopReleaseStatus, ProviderBillingInput, +}; +use gateway_admin::model::provider_credentials::{ + AuthorizationMutationTarget, AuthorizationOwnerBinding, CompleteAuthorization, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision}; +use gateway_admin::ports::provider::ProviderAdminErrorKind; +use gateway_core::engine::credential::{ + CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId, + ProviderAccountStore, +}; +use gateway_core::operation::{ + ContentPart, GenerateRequest, Message, MessageRole, Operation, ProviderOptions, + ReasoningEffort, ReasoningRequirement, +}; +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingTakeOutcome, + ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogPorts, ProviderCooldown, + ProviderCooldownPort, ProviderCredentialState, ProviderCredentialStatePort, + ProviderInstanceCatalogPort, ProviderInstanceConfig, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts, +}; +use gateway_core::routing::{ProviderInstanceId, ProviderKind}; +use gateway_core::task::{WorkerContribution, WorkerKind, WorkerRunnable}; +use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig}; +use provider_openai::credential::CreateCodexCredential; +use provider_openai::transport::profile::APPCAST_POLL_INTERVAL; +use serde_json::{Map, Value, json}; + +use crate::support::{MemoryAccountStore, TestLeaseCoordinator, instance_id, profile, secret}; + +#[tokio::test] +async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions_once() { + let mut bundle = provider_openai::initialize(valid_config(), provider_ports()) + .await + .expect("OpenAI bundle"); + + assert_eq!(bundle.core_provider().name(), "openai"); + assert_eq!(bundle.admin_provider().provider_kind().as_str(), "openai"); + let contributions = bundle.take_worker_contributions(); + assert_eq!(contributions.len(), 4); + assert!( + contributions + .iter() + .any(|item| item.kind() == WorkerKind::OAuthRefresh) + ); + assert!( + contributions + .iter() + .any(|item| item.kind() == WorkerKind::QuotaCatalogHealth) + ); + let release_worker = contributions + .iter() + .find_map(|contribution| match contribution { + WorkerContribution::Registration(registration) + if registration.id.owner() == "openai-desktop-release" => + { + Some(registration) + } + WorkerContribution::Registration(_) | WorkerContribution::Disabled { .. } => None, + }) + .expect("Desktop release worker"); + assert_eq!(release_worker.id.kind(), WorkerKind::QuotaCatalogHealth); + let WorkerRunnable::Scheduled { schedule, .. } = &release_worker.runnable else { + panic!("Desktop release worker must be scheduled"); + }; + assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL); + assert!(contributions.iter().any(|contribution| { + matches!( + contribution, + WorkerContribution::Registration(registration) + if registration.id.owner() == "openai-model-etag" + && matches!(®istration.runnable, WorkerRunnable::Daemon { .. }) + ) + })); + assert!(bundle.take_worker_contributions().is_empty()); +} + +#[tokio::test] +async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing() { + let bundle = provider_openai::initialize(valid_config(), provider_ports()) + .await + .expect("OpenAI bundle"); + let admin = bundle.admin_provider(); + let profile = admin.dashboard_wire_profile().expect("wire profile"); + assert_eq!(profile.codex_version, "0.102.0"); + assert_eq!(profile.release.status, DesktopReleaseStatus::Unchecked); + let billing = admin + .calculated_billing(&ProviderBillingInput { + upstream_model_id: "gpt-4o".to_owned(), + input_tokens: Some(1_000_000), + output_tokens: Some(0), + cached_tokens: Some(0), + cache_write_tokens: Some(0), + total: CurrencyCost { + currency: "USD".to_owned(), + amount: "2.5".parse().expect("amount"), + }, + }) + .expect("billing") + .expect("known pricing"); + assert_eq!(billing.total_amount.amount.as_str(), "2.5"); + assert_eq!(billing.input_price_per_million.amount.as_str(), "2.5"); +} + +#[tokio::test] +async fn openai_core_provider_projects_codex_request_observation_without_routing_side_effects() { + let bundle = provider_openai::initialize(valid_config(), provider_ports()) + .await + .expect("OpenAI bundle"); + let message = Message::new( + MessageRole::User, + vec![ContentPart::Text("summarize".to_owned())], + ) + .expect("message"); + let mut options = ProviderOptions::new(); + options + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), json!(1)), + ( + "turn_metadata".to_owned(), + json!(r#"{"request_kind":"compaction","subagent_kind":"review"}"#), + ), + ]), + ) + .expect("provider options"); + let operation = Operation::Generate( + GenerateRequest::new(vec![message]) + .expect("generate") + .with_reasoning(ReasoningRequirement { + effort: Some(ReasoningEffort::High), + summary: None, + }) + .with_provider_options(options), + ); + + let observation = bundle.core_provider().request_observation(&operation); + + assert_eq!(observation.request_kind.as_deref(), Some("compaction")); + assert_eq!(observation.subagent_kind.as_deref(), Some("review")); + // Codex 当前只在特定多代理预设组合下给出 reasoning_preset；普通 high 保持空值。 + assert_eq!(observation.reasoning_preset, None); + assert!(observation.compact); +} + +#[tokio::test] +async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owner() { + let pending = Arc::new(TestOAuthPending::default()); + let bundle = provider_openai::initialize( + valid_config(), + provider_ports_with( + Arc::new(MemoryAccountStore::default()), + Arc::clone(&pending), + ), + ) + .await + .expect("OpenAI bundle"); + let start_context = MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "admin-owner".to_owned(), + }, + request_id: "request-start".to_owned(), + }; + let started = bundle + .admin_provider() + .start_authorization(PendingAuthorizationMutation::new( + Revision::new(7).expect("revision"), + ProviderKind::new("openai").expect("provider"), + AuthorizationMutationTarget::Create { + provider_instance_id: instance_id(), + name: "OAuth account".to_owned(), + }, + AuthorizationOwnerBinding::from_context(&start_context), + )) + .await + .expect("start authorization"); + { + let values = pending.values.lock().expect("OAuth pending"); + let (_, payload, _) = values.values().next().expect("stored pending"); + let mutation = payload + .expose_to_provider() + .get("mutation") + .and_then(Value::as_object) + .expect("pending mutation"); + assert_eq!( + mutation + .get("expected_config_revision") + .and_then(Value::as_u64), + Some(7) + ); + assert_eq!( + mutation.get("started_request_id").and_then(Value::as_str), + Some("request-start") + ); + } + let error = bundle + .admin_provider() + .complete_authorization(CompleteAuthorization { + context: MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "different-owner".to_owned(), + }, + request_id: "request-complete".to_owned(), + }, + flow_id: started.flow_id, + callback_url: "http://localhost:1455/auth/callback?code=unused&state=unused".to_owned(), + }) + .await + .expect_err("wrong owner"); + assert_eq!(error.kind(), ProviderAdminErrorKind::NotFound); + assert_eq!(pending.values.lock().expect("OAuth pending").len(), 1); +} + +#[tokio::test] +async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export() { + let store = Arc::new(MemoryAccountStore::default()); + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_admin_projection".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "admin projection".to_owned(), + secret: secret("admin-projection-access"), + account: profile("chatgpt-admin-projection"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await + .expect("create account"); + let account = store + .account("acct_admin_projection") + .expect("stored account"); + let record = account_record(&account); + let bundle = provider_openai::initialize( + valid_config(), + provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), + ) + .await + .expect("OpenAI bundle"); + let admin = bundle.admin_provider(); + let account_id = account.id().clone(); + let quota = admin.quota(&account_id, false).await.expect("cached quota"); + assert!(quota.windows.is_empty()); + let models = admin + .models(&account_id, false) + .await + .expect("cached models"); + assert!(models.models.is_empty()); + let loaded = store + .load_credential(account.id(), account.revision()) + .await + .expect("loaded credential"); + let exported = admin + .export_credentials(vec![ProviderExportCredentialInput { + account: record, + provider_material: ProviderDocument::new(OpaqueProviderData::new( + loaded.credential.into_inner(), + )), + }]) + .await + .expect("canonical export"); + assert_eq!(exported.account_ids, vec![account_id]); + assert_eq!( + exported + .document + .expose_to_provider() + .expose_to_provider() + .get("sourceFormat") + .and_then(Value::as_str), + Some("cpr") + ); +} + +#[tokio::test] +async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit() { + let store = Arc::new(MemoryAccountStore::default()); + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_admin_invalid".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "admin invalid".to_owned(), + secret: secret("admin-invalid-access"), + account: profile("chatgpt-admin-invalid"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await + .expect("create account"); + let account = store.account("acct_admin_invalid").expect("stored account"); + let record = account_record(&account); + let bundle = provider_openai::initialize( + valid_config(), + provider_ports_with(store, Arc::new(TestOAuthPending::default())), + ) + .await + .expect("OpenAI bundle"); + let admin = bundle.admin_provider(); + let import_error = admin + .prepare_import(PrepareCredentialImport { + provider_instance_id: instance_id(), + document: ProviderDocument::new(OpaqueProviderData::new(Map::new())), + }) + .await + .expect_err("invalid import"); + assert_eq!(import_error.kind(), ProviderAdminErrorKind::Invalid); + let rotation_error = admin + .prepare_rotation(PrepareCredentialRotation { + account: record.clone(), + expected_credential_revision: record.credential_revision, + provider_material: ProviderDocument::new(OpaqueProviderData::new(Map::new())), + }) + .await + .expect_err("invalid rotation"); + assert_eq!(rotation_error.kind(), ProviderAdminErrorKind::Invalid); + let mut missing = record; + missing.id = "acct_admin_missing".to_owned(); + let refresh_error = admin + .prepare_refresh(PrepareCredentialRefresh { account: missing }) + .await + .expect_err("missing refresh target"); + assert_eq!(refresh_error.kind(), ProviderAdminErrorKind::NotFound); +} + +fn provider_ports() -> ProviderStorePorts { + provider_ports_with( + Arc::new(MemoryAccountStore::default()), + Arc::new(TestOAuthPending::default()), + ) +} + +fn provider_ports_with( + accounts: Arc, + pending: Arc, +) -> ProviderStorePorts { + ProviderStorePorts::new( + accounts, + Arc::new(TestLeaseCoordinator::default()), + ProviderCatalogPorts::new( + Arc::new(TestInstanceCatalog), + Arc::new(TestCatalogCache::default()), + ), + Arc::new(TestCredentialState), + Arc::new(TestCooldown), + Arc::new(TestRuntimePolicy), + pending, + ) +} + +fn account_record(account: &ProviderAccount) -> AccountRecord { + let now = Utc::now(); + AccountRecord { + id: account.id().to_string(), + provider_instance_id: account.instance().clone(), + provider_kind: account.provider().clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + upstream_user_id: account.upstream_user_id().to_owned(), + upstream_account_id: account.upstream_account_id().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + credential_revision: Revision::new(account.revision().get()).expect("revision"), + has_refresh_token: account.has_refresh_token(), + access_token_expires_at: DateTime::::from(account.access_token_expires_at()), + next_refresh_at: account.next_refresh_at().map(DateTime::::from), + enabled: account.enabled(), + availability: AdminAccountAvailability::Ready, + availability_reason: None, + cooldown_until: account.cooldown_until().map(DateTime::::from), + availability_observed_at: now, + quota_observed_at: None, + created_at: now, + updated_at: now, + } +} + +fn valid_config() -> OpenAiConfig { + OpenAiConfig { + wire_profile: CodexWireProfileConfig { + originator: "Codex Desktop".to_owned(), + codex_version: "0.102.0".to_owned(), + desktop_version: "1.2026.190".to_owned(), + desktop_build: "19012345678".to_owned(), + os_type: "macOS".to_owned(), + os_version: "15.5.0".to_owned(), + arch: "arm64".to_owned(), + terminal: "xterm-256color".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) + .single() + .expect("valid test time"), + }, + } +} + +struct TestInstanceCatalog; + +impl ProviderInstanceCatalogPort for TestInstanceCatalog { + fn list_instances<'a>( + &'a self, + provider_kind: &'a ProviderKind, + _include_disabled: bool, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(vec![ProviderInstanceConfig::new( + ProviderInstanceId::new("inst_openai_primary").expect("instance"), + provider_kind.clone(), + "https://chatgpt.com/backend-api".to_owned(), + true, + )]) + }) + } +} + +#[derive(Default)] +struct TestCatalogCache { + values: Mutex>, +} + +impl ProviderCatalogCachePort for TestCatalogCache { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + _ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>> { + Box::pin(async move { + self.values + .lock() + .expect("catalog cache") + .insert(key.account_id().to_string(), catalog.clone()); + Ok(()) + }) + } + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(self + .values + .lock() + .expect("catalog cache") + .get(key.account_id().as_str()) + .cloned()) + }) + } +} + +struct TestCredentialState; + +impl ProviderCredentialStatePort for TestCredentialState { + fn replace( + &self, + _state: ProviderCredentialState, + ) -> BoxFuture<'_, Result<(), ProviderStoreError>> { + Box::pin(async { Ok(()) }) + } + + fn read<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async { Ok(None) }) + } + + fn clear<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(false) }) + } +} + +struct TestCooldown; + +impl ProviderCooldownPort for TestCooldown { + fn put_if_later( + &self, + _cooldown: ProviderCooldown, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(false) }) + } + + fn read<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async { Ok(None) }) + } + + fn clear<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + _through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(false) }) + } +} + +struct TestRuntimePolicy; + +impl ProviderRuntimePolicyPort for TestRuntimePolicy { + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + ProviderRefreshPolicy::try_new( + Duration::from_secs(300), + NonZeroU32::new(4).expect("nonzero concurrency"), + ) + }) + } +} + +#[derive(Default)] +struct TestOAuthPending { + values: Mutex>, +} + +type PendingKey = (String, String); +type PendingValue = (String, OpaqueProviderData, SystemTime); + +impl OAuthPendingFlowPort for TestOAuthPending { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let key = ( + flow.provider_kind().as_str().to_owned(), + flow.flow().expose_to_store().to_owned(), + ); + let mut values = self.values.lock().expect("OAuth pending"); + if values.contains_key(&key) { + return Ok(OAuthPendingPutOutcome::AlreadyExists); + } + values.insert( + key, + ( + flow.owner().expose_to_store().to_owned(), + flow.payload().clone(), + SystemTime::now() + flow.ttl(), + ), + ); + Ok(OAuthPendingPutOutcome::Stored) + }) + } + + fn take_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a gateway_core::provider_ports::OAuthPendingBinding, + owner: &'a gateway_core::provider_ports::OAuthPendingBinding, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let key = ( + provider_kind.as_str().to_owned(), + flow.expose_to_store().to_owned(), + ); + let mut values = self.values.lock().expect("OAuth pending"); + let Some((stored_owner, _, expires_at)) = values.get(&key) else { + return Ok(OAuthPendingTakeOutcome::NotFound); + }; + if expires_at <= &SystemTime::now() { + values.remove(&key); + return Ok(OAuthPendingTakeOutcome::NotFound); + } + if stored_owner != owner.expose_to_store() { + return Ok(OAuthPendingTakeOutcome::OwnerMismatch); + } + let (_, payload, _) = values.remove(&key).expect("checked pending"); + Ok(OAuthPendingTakeOutcome::Taken(payload)) + }) + } +} diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs new file mode 100644 index 000000000..5f82fbca6 --- /dev/null +++ b/backend/crates/providers/openai/tests/config.rs @@ -0,0 +1,65 @@ +use std::path::Path; + +use chrono::{TimeZone as _, Utc}; +use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig}; + +#[test] +fn openai_config_builds_the_audited_wire_profile() { + let mut config = valid_config(); + config + .resolve_and_validate(Path::new("/srv/gateway")) + .expect("valid OpenAI config"); + let profile = config.wire_profile_state().snapshot(); + + assert_eq!( + profile.user_agent(), + "Codex Desktop/0.102.0 (macOS 15.5.0; arm64) xterm-256color (Codex Desktop; 1.2026.190)" + ); + assert_eq!(profile.desktop_build, "19012345678"); +} + +#[test] +fn openai_config_rejects_noncanonical_versions_and_empty_fields() { + let mut config = valid_config(); + config.wire_profile.codex_version = "latest".to_owned(); + assert!( + config + .resolve_and_validate(Path::new("/srv/gateway")) + .is_err() + ); + + let mut config = valid_config(); + config.wire_profile.desktop_version = "1.preview".to_owned(); + assert!( + config + .resolve_and_validate(Path::new("/srv/gateway")) + .is_err() + ); + + let mut config = valid_config(); + config.wire_profile.originator.clear(); + assert!( + config + .resolve_and_validate(Path::new("/srv/gateway")) + .is_err() + ); +} + +fn valid_config() -> OpenAiConfig { + OpenAiConfig { + wire_profile: CodexWireProfileConfig { + originator: "Codex Desktop".to_owned(), + codex_version: "0.102.0".to_owned(), + desktop_version: "1.2026.190".to_owned(), + desktop_build: "19012345678".to_owned(), + os_type: "macOS".to_owned(), + os_version: "15.5.0".to_owned(), + arch: "arm64".to_owned(), + terminal: "xterm-256color".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 19, 0, 0, 0) + .single() + .expect("valid test time"), + }, + } +} diff --git a/backend/crates/providers/openai/tests/credential/admin.rs b/backend/crates/providers/openai/tests/credential/admin.rs index 318590d86..ed37d02ac 100644 --- a/backend/crates/providers/openai/tests/credential/admin.rs +++ b/backend/crates/providers/openai/tests/credential/admin.rs @@ -4,21 +4,29 @@ use std::time::Duration; use async_trait::async_trait; use chrono::{TimeZone, Utc}; +use futures::future::BoxFuture; use gateway_core::engine::credential::{ CredentialRevision, LoadedCredential, PlaintextCredential, ProviderAccountId, }; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest, + ProviderStoreError, +}; use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use provider_openai::credential::{ - CodexCredentialAdmin, CodexCredentialAdminError, CodexCredentialAdminService, - CodexCredentialCodec, CodexIdentityVerificationError, CodexJwtIdentityVerifier, - CodexRefreshLeaseAcquisition, CodexRefreshLeaseCoordinator, CodexRefreshLeaseError, - CodexRefreshLeaseRequest, CodexTokenIdentityVerifier, CreateCodexCredential, - ExportManagedCodexCredential, ImportCodexOAuthCredential, ImportCodexOAuthCredentialBatch, - ReqwestOpenAiJwksSource, RotateManagedCodexCredential, + CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexCredentialAdmin, + CodexCredentialAdminError, CodexCredentialAdminService, CodexCredentialCodec, + CodexIdentityExpectation, CodexIdentityVerification, CodexIdentityVerificationError, + CodexJwtIdentityVerifier, CreateCodexCredential, ExportManagedCodexCredential, + ImportCodexOAuthCredential, ImportCodexOAuthCredentialBatch, + ReqwestCodexAuthenticatedAccountSource, ReqwestOpenAiJwksSource, RotateManagedCodexCredential, }; -use secrecy::ExposeSecret; +use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; +use secrecy::{ExposeSecret, SecretString}; -use crate::support::{MemoryAccountStore, codex_account, instance_id, profile, secret}; +use crate::support::{ + MemoryAccountStore, codex_account, instance_id, profile, runtime_policy, secret, +}; fn import(id: &str, token: &str) -> ImportCodexOAuthCredential { ImportCodexOAuthCredential { @@ -27,10 +35,16 @@ fn import(id: &str, token: &str) -> ImportCodexOAuthCredential { name: format!("name-{id}"), secret: secret(token), verified_account: profile(&format!("chatgpt-{id}")), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, } } +fn encoded_credential(account_id: &str, token: &str) -> PlaintextCredential { + CodexCredentialCodec::encode_new(&secret(token), &profile(account_id), Vec::new()) + .expect("encode current credential") +} + #[test] fn prepare_import_returns_core_account_with_plaintext_json() { let prepared = CodexCredentialAdmin @@ -76,8 +90,7 @@ fn prepare_batch_is_all_or_nothing_validation() { #[test] fn prepare_rotation_preserves_provider_owned_cookie_data() { let account = codex_account("acct_rotate"); - let mut credential = CodexCredentialCodec::encode(&secret("old-access"), Vec::new()) - .expect("encode current credential"); + let mut credential = encoded_credential("chatgpt-acct_rotate", "old-access"); credential .expose_to_provider() .get("access_token") @@ -93,6 +106,7 @@ fn prepare_rotation_preserves_provider_owned_cookie_data() { }, secret: secret("new-access"), verified_account: profile("chatgpt-acct_rotate"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), }) .expect("prepare rotation"); let decoded = CodexCredentialCodec::decode(prepared.credential.credential()) @@ -107,11 +121,11 @@ fn prepare_rotation_rejects_account_rebinding() { .prepare_rotation(RotateManagedCodexCredential { current: LoadedCredential { account: codex_account("acct_rotate"), - credential: CodexCredentialCodec::encode(&secret("old-access"), Vec::new()) - .expect("encode current credential"), + credential: encoded_credential("chatgpt-acct_rotate", "old-access"), }, secret: secret("new-access"), verified_account: profile("chatgpt-other"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), }) .expect_err("identity rebinding must fail"); assert_eq!(error, CodexCredentialAdminError::IdentityMismatch); @@ -130,8 +144,7 @@ fn export_item(id: &str, token: &str) -> ExportManagedCodexCredential { ExportManagedCodexCredential { current: LoadedCredential { account: codex_account(id), - credential: CodexCredentialCodec::encode(&secret(token), Vec::new()) - .expect("credential"), + credential: encoded_credential(&format!("chatgpt-{id}"), token), }, added_at: Utc .with_ymd_and_hms(2026, 7, 18, 2, 47, 1) @@ -241,16 +254,28 @@ impl TokenRefresher for ManualRefresher { struct ManualVerifier; #[async_trait] -impl CodexTokenIdentityVerifier for ManualVerifier { +impl CodexAccountIdentityVerifier for ManualVerifier { async fn verify( &self, secret: &provider_openai::credential::CodexOAuthSecret, - ) -> Result - { + _expectation: &CodexIdentityExpectation, + ) -> Result { if secret.access_token.expose_secret() != "refreshed-access" { return Err(CodexIdentityVerificationError::Rejected); } - Ok(profile("chatgpt-acct_manual_refresh")) + Ok(CodexIdentityVerification::Complete(profile( + "chatgpt-acct_manual_refresh", + ))) + } + + async fn verify_authorization( + &self, + _secret: &provider_openai::credential::CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) } } @@ -263,27 +288,48 @@ impl Drop for DropCountGuard { } struct ManualLeases { - requests: Mutex>, + requests: Mutex>, drops: Arc, available: bool, } -#[async_trait] -impl CodexRefreshLeaseCoordinator for ManualLeases { - async fn try_acquire( +impl ProviderLeasePort for ManualLeases { + fn load_state<'a>( + &'a self, + _: &'a gateway_core::routing::ProviderInstanceId, + _: &'a [ProviderAccountId], + ) -> BoxFuture< + 'a, + Result, + > { + Box::pin(async { + Ok(gateway_core::provider_ports::ProviderSchedulingState::new( + Default::default(), + None, + 0, + )) + }) + } + + fn try_acquire( &self, - request: &CodexRefreshLeaseRequest, - ) -> Result { - self.requests - .lock() - .expect("lease requests") - .push(request.clone()); - Ok(if self.available { - CodexRefreshLeaseAcquisition::Acquired(Box::new(DropCountGuard(Arc::clone( - &self.drops, - )))) - } else { - CodexRefreshLeaseAcquisition::Unavailable + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + match request { + ProviderLeaseRequest::RefreshCapacity(_) => {} + ProviderLeaseRequest::Refresh(request) => { + self.requests.lock().expect("lease requests").push(request); + } + ProviderLeaseRequest::Scheduling(_) => panic!("unexpected scheduling lease"), + } + Ok(if self.available { + ProviderLeaseAcquisition::Acquired(Box::new(DropCountGuard(Arc::clone( + &self.drops, + )))) + } else { + ProviderLeaseAcquisition::Busy { retry_after: None } + }) }) } } @@ -306,6 +352,7 @@ async fn manual_refresh_fixture( name: "manual refresh".to_owned(), secret: secret("old-access"), account: profile("chatgpt-acct_manual_refresh"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await @@ -324,9 +371,8 @@ async fn manual_refresh_fixture( refresher.clone(), Arc::new(ManualVerifier), leases.clone(), - Duration::from_secs(60), - ) - .expect("admin service"); + runtime_policy(), + ); (store, refresher, leases, service) } @@ -383,7 +429,7 @@ async fn manual_refresh_prepares_revision_fenced_rotation_without_store_mutation ); assert_eq!(leases.drops.load(Ordering::SeqCst), 0); drop(prepared); - assert_eq!(leases.drops.load(Ordering::SeqCst), 1); + assert_eq!(leases.drops.load(Ordering::SeqCst), 2); } #[tokio::test] @@ -423,12 +469,12 @@ async fn manual_refresh_stale_revision_and_missing_lease_fail_before_exchange() struct ImportVerifier; #[async_trait] -impl CodexTokenIdentityVerifier for ImportVerifier { +impl CodexAccountIdentityVerifier for ImportVerifier { async fn verify( &self, secret: &provider_openai::credential::CodexOAuthSecret, - ) -> Result - { + expectation: &CodexIdentityExpectation, + ) -> Result { let suffix = secret .access_token .expose_secret() @@ -437,9 +483,28 @@ impl CodexTokenIdentityVerifier for ImportVerifier { let mut profile = profile(&format!("chatgpt-{suffix}")); if let Some(user) = suffix.strip_prefix("shared-") { profile.chatgpt_account_id = "chatgpt-shared".to_owned(); - profile.chatgpt_user_id = Some(format!("user-{user}")); + profile.chatgpt_user_id = format!("user-{user}"); } - Ok(profile) + if expectation + .chatgpt_account_id() + .is_some_and(|expected| expected != profile.chatgpt_account_id) + || expectation + .chatgpt_user_id() + .is_some_and(|expected| expected != profile.chatgpt_user_id) + { + return Err(CodexIdentityVerificationError::Rejected); + } + Ok(CodexIdentityVerification::Complete(profile)) + } + + async fn verify_authorization( + &self, + _secret: &provider_openai::credential::CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) } } @@ -456,9 +521,8 @@ fn import_service( drops: Arc::new(AtomicUsize::new(0)), available: true, }), - Duration::from_secs(60), + runtime_policy(), ) - .expect("import service") } fn unused_import_refresher() -> Arc { @@ -518,7 +582,7 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() { #[tokio::test] async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_projections() { - let prepared = import_service(unused_import_refresher()) + let error = import_service(unused_import_refresher()) .prepare_import_document( instance_id(), serde_json::json!({ @@ -532,11 +596,9 @@ async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_project }), ) .await - .expect("signed token identity is authoritative during import"); + .expect_err("stale document identity must not override authenticated identity"); - let account = &prepared.accounts()[0].account; - assert_eq!(account.upstream_account_id(), Some("chatgpt-cpr")); - assert_eq!(account.upstream_user_id(), "user-chatgpt-cpr"); + assert_eq!(error, CodexCredentialAdminError::IdentityRejected); } #[tokio::test] @@ -671,9 +733,28 @@ async fn real_cpr_fixture_import_contract() { provider_openai::credential::token_client::official_openai_token_client() .expect("official token client"), ); - let verifier = Arc::new(CodexJwtIdentityVerifier::new(Box::new( + let signed = Arc::new(CodexJwtIdentityVerifier::new(Box::new( ReqwestOpenAiJwksSource::new().expect("official JWKS source"), ))); + let profile = CodexWireProfileState::new(CodexWireProfile { + originator: "codex_cli_rs".to_owned(), + codex_version: "0.144.0".to_owned(), + desktop_version: "1.0.0".to_owned(), + desktop_build: "1".to_owned(), + os_type: "linux".to_owned(), + os_version: "6.8".to_owned(), + arch: "x86_64".to_owned(), + terminal: "xterm".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 18, 0, 0, 0) + .single() + .expect("fixture time"), + }); + let accounts = Arc::new( + ReqwestCodexAuthenticatedAccountSource::new(profile) + .expect("official authenticated account source"), + ); + let verifier = Arc::new(CodexAccountIdentityService::new(signed, accounts)); let service = CodexCredentialAdminService::new( store.repository(), refresher, @@ -683,9 +764,8 @@ async fn real_cpr_fixture_import_contract() { drops: Arc::new(AtomicUsize::new(0)), available: true, }), - Duration::from_secs(60), - ) - .expect("admin service"); + runtime_policy(), + ); let prepared = service .prepare_import_document(instance_id(), payload) .await @@ -695,6 +775,20 @@ async fn real_cpr_fixture_import_contract() { for account in prepared.accounts() { assert_eq!(account.account.provider().as_str(), "openai"); assert!(account.account.id().as_str().starts_with("acct_")); - assert!(CodexCredentialCodec::decode(&account.credential).is_ok()); + assert!(account.account.upstream_account_id().is_some()); + assert!(!account.account.upstream_user_id().is_empty()); + let runtime = + CodexCredentialCodec::decode(&account.credential).expect("new credential schema"); + assert!(runtime.principal.poid.is_some()); + assert_ne!( + runtime.principal.poid.as_deref(), + account.account.upstream_account_id() + ); + assert_eq!( + uuid::Uuid::parse_str(&runtime.installation_id) + .expect("random installation UUID") + .get_version_num(), + 4 + ); } } diff --git a/backend/crates/providers/openai/tests/credential/catalog.rs b/backend/crates/providers/openai/tests/credential/catalog.rs index 8a8f2f67d..faa6e6d38 100644 --- a/backend/crates/providers/openai/tests/credential/catalog.rs +++ b/backend/crates/providers/openai/tests/credential/catalog.rs @@ -1,13 +1,13 @@ use std::sync::Arc; -use std::time::{Duration, SystemTime}; +use std::time::SystemTime; use chrono::{TimeZone, Utc}; use futures::executor::block_on; use gateway_core::engine::credential::{ - AccountAvailability, AccountStateChange, ProviderAccountId, ProviderAccountStore as _, + AccountAvailability, AccountStateChange, CredentialRevision, ProviderAccountId, + ProviderAccountStore as _, }; use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderKind}; -use provider_openai::CodexEndpointPolicy; use provider_openai::credential::{ CodexCredentialCatalogError, CodexCredentialCatalogService, CreateCodexCredential, }; @@ -15,7 +15,7 @@ use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileStat use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use crate::support::{MemoryAccountStore, instance_id, profile, secret}; +use crate::support::{MemoryAccountStore, instance_id, loopback_origin_policy, profile, secret}; fn wire_profile() -> CodexWireProfileState { CodexWireProfileState::new(CodexWireProfile { @@ -53,6 +53,7 @@ async fn create(store: &Arc, id: &str, token: &str) { name: id.to_owned(), secret: secret(token), account: profile(&format!("chatgpt-{id}")), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await @@ -60,15 +61,14 @@ async fn create(store: &Arc, id: &str, token: &str) { } fn service(store: &Arc) -> CodexCredentialCatalogService { - CodexCredentialCatalogService::new_with_endpoint_policy( + CodexCredentialCatalogService::new( store.repository(), wire_profile(), reqwest::Client::builder() .no_proxy() .build() .expect("client"), - CodexEndpointPolicy::Loopback, - Duration::from_secs(60), + loopback_origin_policy(), ) } @@ -84,6 +84,14 @@ fn catalog(model: &str, display: &str) -> serde_json::Value { }) } +fn account_id(value: &str) -> ProviderAccountId { + ProviderAccountId::new(value).expect("account id") +} + +fn revision() -> CredentialRevision { + CredentialRevision::new(1).expect("revision") +} + #[tokio::test] async fn realtime_catalog_builds_union_and_per_account_entitlements() { let server = MockServer::start().await; @@ -105,11 +113,11 @@ async fn realtime_catalog_builds_union_and_per_account_entitlements() { .expect("synchronize catalog"); assert_eq!(snapshot.models().len(), 2); assert_eq!( - snapshot.account_models("acct_one"), + snapshot.account_models(&account_id("acct_one"), revision()), Some(&["gpt-5.4".to_owned()][..]) ); assert_eq!( - snapshot.account_models("acct_two"), + snapshot.account_models(&account_id("acct_two"), revision()), Some(&["gpt-5.5".to_owned()][..]) ); } @@ -153,13 +161,13 @@ async fn realtime_catalog_should_preserve_the_official_model_order() { official_order ); assert_eq!( - snapshot.account_models("acct_order"), + snapshot.account_models(&account_id("acct_order"), revision()), Some(official_order.map(str::to_owned).as_slice()) ); } #[tokio::test] -async fn fresh_catalog_is_served_from_ttl_cache_without_postgres_snapshot() { +async fn catalog_is_reused_until_explicit_invalidation_or_etag_change() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/codex/models")) @@ -206,16 +214,13 @@ async fn single_account_refresh_queries_only_that_account_and_updates_its_cache( assert_eq!(models, ["gpt-5.5"]); assert_eq!( service - .cached_account_models(instance.id(), &selected) + .cached_account_models(instance.id(), &selected, revision()) .expect("cache read"), Some(vec!["gpt-5.5".to_owned()]) ); assert_eq!( service - .cached_account_models( - instance.id(), - &ProviderAccountId::new("acct_one").expect("account id") - ) + .cached_account_models(instance.id(), &account_id("acct_one"), revision(),) .expect("other cache read"), None ); @@ -234,15 +239,62 @@ async fn invalidation_forces_a_new_realtime_query() { create(&store, "acct_one", "token-one").await; let service = service(&store); let instance = instance(&format!("{}/backend-api", server.uri())); + assert_eq!(service.catalog_generation().get(), 0); service .synchronize_instance(&instance) .await .expect("first query"); + assert_eq!(service.catalog_generation().get(), 1); service.invalidate(instance.id()).expect("invalidate"); + assert_eq!(service.catalog_generation().get(), 2); service .synchronize_instance(&instance) .await .expect("second query"); + assert_eq!(service.catalog_generation().get(), 3); +} + +#[tokio::test] +async fn response_etag_change_is_deduplicated_and_queued_for_refresh() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/codex/models")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("etag", "\"models-v1\"") + .set_body_json(catalog("gpt-5.4", "GPT-5.4")), + ) + .expect(1) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + create(&store, "acct_one", "token-one").await; + let service = service(&store); + let instance = instance(&format!("{}/backend-api", server.uri())); + service + .synchronize_instance(&instance) + .await + .expect("initial catalog"); + + assert!( + !service + .observe_response_etag(instance.id(), "\"models-v1\"") + .expect("same ETag") + ); + assert!( + service + .observe_response_etag(instance.id(), "\"models-v2\"") + .expect("changed ETag") + ); + assert!( + !service + .observe_response_etag(instance.id(), "\"models-v2\"") + .expect("duplicate ETag") + ); + assert_eq!( + service.wait_for_etag_refresh().await, + vec![instance.id().clone()] + ); } #[tokio::test] diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index 58e910f8f..47cc7650d 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -3,19 +3,21 @@ use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, SystemTime}; +use chrono::Utc; use futures::executor::block_on; use gateway_core::engine::credential::{ AccountAvailability, AccountSelectionPolicy, ProviderAccountId, RotationStrategy, }; use gateway_core::engine::{ - AccountSelectionConstraints, AttemptContext, CancellationToken, ModelRequestId, - UpstreamSendState, + AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, }; -use gateway_core::error::ProviderErrorKind; use provider_openai::credential::{ - CodexCookiePolicy, CodexCredentialCodec, CodexCredentialSelector, CreateCodexCredential, + CodexAccountFailure, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialCodec, + CodexCredentialQuotaService, CodexCredentialSelector, CreateCodexCredential, CredentialSelectionError, SelectCodexCredential, }; +use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; +use provider_openai::{CodexOriginPolicy, OfficialCodexOriginPolicy}; use secrecy::ExposeSecret; use url::Url; @@ -33,6 +35,7 @@ fn create_account(store: &Arc, id: &str, token: &str) { name: id.to_owned(), secret: secret(token), account: profile(&format!("chatgpt-{id}")), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }), ) @@ -52,7 +55,7 @@ fn attempt_with_required( NonZeroU32::new(1).expect("attempt"), SystemTime::now() + Duration::from_secs(30), account_policy(), - AccountSelectionConstraints::new(excluded_accounts, required_account), + AccountAttemptContext::new(excluded_accounts, required_account, None), None, CancellationToken::new(), ) @@ -68,7 +71,7 @@ fn round_robin_attempt() -> AttemptContext { NonZeroU32::new(2).expect("concurrency"), Duration::ZERO, ), - AccountSelectionConstraints::new(BTreeSet::new(), None), + AccountAttemptContext::new(BTreeSet::new(), None, None), None, CancellationToken::new(), ) @@ -78,17 +81,48 @@ fn selector( store: &Arc, leases: Arc, ) -> CodexCredentialSelector { + let profile = CodexWireProfileState::new(CodexWireProfile { + originator: "codex_cli_rs".to_owned(), + codex_version: "0.144.0".to_owned(), + desktop_version: "1.0.0".to_owned(), + desktop_build: "1".to_owned(), + os_type: "linux".to_owned(), + os_version: "6.8".to_owned(), + arch: "x86_64".to_owned(), + terminal: "selector-contract".to_owned(), + verified_at: Utc::now(), + }); + let http = reqwest::Client::builder().build().expect("HTTP client"); + let origin: Arc = Arc::new(OfficialCodexOriginPolicy); + let catalog = Arc::new(CodexCredentialCatalogService::new( + store.repository(), + profile.clone(), + http.clone(), + Arc::clone(&origin), + )); + let quota = Arc::new(CodexCredentialQuotaService::new( + store.repository(), + profile, + http, + origin, + )); CodexCredentialSelector::new( store.repository(), leases, + catalog, + quota, CodexCookiePolicy::official().expect("official cookie policy"), ) } #[test] fn codec_persists_tokens_as_plaintext_provider_json() { - let encoded = CodexCredentialCodec::encode(&secret("literal-access-token"), Vec::new()) - .expect("encode plaintext credential"); + let encoded = CodexCredentialCodec::encode_new( + &secret("literal-access-token"), + &profile("chatgpt-literal"), + Vec::new(), + ) + .expect("encode plaintext credential"); assert_eq!( encoded .expose_to_provider() @@ -111,10 +145,61 @@ fn codec_persists_tokens_as_plaintext_provider_json() { keys.sort_unstable(); assert_eq!( keys, - ["access_token", "cookies", "refresh_token", "schema_version"] + [ + "access_token", + "cookies", + "installation_id", + "principal", + "refresh_token", + "schema_version", + ] ); } +#[test] +fn codec_reimport_preserves_existing_installation_id_for_the_same_principal() { + let existing = CodexCredentialCodec::encode_new( + &secret("existing-access-token"), + &profile("chatgpt-stable-installation"), + Vec::new(), + ) + .expect("existing credential"); + let incoming = CodexCredentialCodec::encode_new( + &secret("incoming-access-token"), + &profile("chatgpt-stable-installation"), + Vec::new(), + ) + .expect("incoming credential"); + let existing_id = CodexCredentialCodec::decode_complete(&existing) + .expect("existing data") + .installation_id; + + let preserved = CodexCredentialCodec::preserve_installation_id(&incoming, &existing) + .expect("preserve installation ID"); + let preserved = CodexCredentialCodec::decode_complete(&preserved).expect("preserved data"); + + assert_eq!(preserved.installation_id, existing_id); + assert_eq!(preserved.access_token, "incoming-access-token"); +} + +#[test] +fn codec_reimport_rejects_installation_reuse_across_principals() { + let existing = CodexCredentialCodec::encode_new( + &secret("existing-access-token"), + &profile("chatgpt-existing-principal"), + Vec::new(), + ) + .expect("existing credential"); + let incoming = CodexCredentialCodec::encode_new( + &secret("incoming-access-token"), + &profile("chatgpt-incoming-principal"), + Vec::new(), + ) + .expect("incoming credential"); + + assert!(CodexCredentialCodec::preserve_installation_id(&incoming, &existing).is_err()); +} + #[test] fn repository_round_trips_plaintext_runtime_secret() { let store = Arc::new(MemoryAccountStore::default()); @@ -145,6 +230,7 @@ fn selector_uses_frozen_global_account_policy_for_lease() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -153,9 +239,26 @@ fn selector_uses_frozen_global_account_policy_for_lease() { .expect("select account"); assert_eq!(lease.account_id().as_str(), "acct_primary"); + let installation_id = lease.installation_id(); + assert_eq!( + uuid::Uuid::parse_str(installation_id) + .expect("installation UUID") + .get_version_num(), + 4 + ); + let runtime = block_on(store.repository().load_runtime_credential(lease.account())) + .expect("runtime credential"); + assert_eq!(runtime.installation_id, installation_id); let requests = leases.requests.lock().expect("lease requests lock"); - assert_eq!(requests[0].max_concurrent, 2); - assert_eq!(requests[0].request_interval, Duration::from_millis(10)); + assert_eq!(requests[0].provider_instance_id(), &instance_id()); + assert_eq!(requests[0].account_id(), lease.account_id()); + assert_eq!( + requests[0].credential_revision(), + lease.account().revision() + ); + assert_eq!(requests[0].max_concurrent().get(), 2); + assert_eq!(requests[0].request_interval(), Duration::from_millis(10)); + assert_eq!(requests[0].deadline(), attempt.deadline()); } #[test] @@ -172,6 +275,7 @@ fn selector_round_robin_cursor_advances_across_requests() { let attempt = round_robin_attempt(); let lease = block_on(selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &request_url, attempt: &attempt, })) @@ -198,6 +302,7 @@ fn selector_honors_attempt_local_account_exclusion() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -219,6 +324,7 @@ fn selector_uses_only_the_required_account() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -241,6 +347,7 @@ fn unavailable_required_account_never_falls_back() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -265,6 +372,7 @@ fn selector_returns_capacity_error_when_every_redis_lease_is_busy() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -280,7 +388,7 @@ fn selector_returns_capacity_error_when_every_redis_lease_is_busy() { } #[test] -fn unauthorized_failure_marks_unified_account_invalid() { +fn credential_expired_failure_marks_unified_account_expired() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_primary", "at-primary"); let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); @@ -289,29 +397,26 @@ fn unauthorized_failure_marks_unified_account_invalid() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, }), ) .expect("select account"); - block_on(selector.record_failure( - &lease, - ProviderErrorKind::Unauthorized, - UpstreamSendState::Sent, - None, - )); + block_on(selector.record_failure(&lease, CodexAccountFailure::CredentialExpired)) + .expect("record credential expiry"); assert_eq!( store .account("acct_primary") .expect("account") .availability(), - AccountAvailability::Invalid + AccountAvailability::Expired ); } #[test] -fn permission_denied_failure_keeps_every_account_ready() { +fn identity_verification_failure_isolates_only_selected_account() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_primary", "at-primary"); create_account(&store, "acct_other", "at-other"); @@ -321,6 +426,46 @@ fn permission_denied_failure_keeps_every_account_ready() { block_on( selector.select(&SelectCodexCredential { provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") + .expect("request URL"), + attempt: &attempt, + }), + ) + .expect("select account"); + + block_on(selector.record_failure(&lease, CodexAccountFailure::IdentityVerificationRequired)) + .expect("record identity verification failure"); + + assert_eq!( + store + .account(lease.account_id().as_str()) + .expect("selected account") + .availability(), + AccountAvailability::Invalid + ); + let other = if lease.account_id().as_str() == "acct_primary" { + "acct_other" + } else { + "acct_primary" + }; + assert_eq!( + store.account(other).expect("other account").availability(), + AccountAvailability::Ready + ); +} + +#[test] +fn cloudflare_challenge_backoff_escalates_and_success_resets_it() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "at-primary"); + let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); + let attempt = attempt(BTreeSet::new()); + let lease = + block_on( + selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") .expect("request URL"), attempt: &attempt, @@ -330,22 +475,174 @@ fn permission_denied_failure_keeps_every_account_ready() { block_on(selector.record_failure( &lease, - ProviderErrorKind::PermissionDenied, - UpstreamSendState::Sent, - None, - )); + CodexAccountFailure::CloudflareChallenge { retry_after: None }, + )) + .expect("record first challenge"); + let first = store + .account("acct_primary") + .expect("account") + .cooldown_until() + .expect("first cooldown") + .duration_since(SystemTime::now()) + .expect("future cooldown"); + assert!(first >= Duration::from_secs(9) && first <= Duration::from_secs(10)); + + block_on(selector.record_failure( + &lease, + CodexAccountFailure::CloudflareChallenge { retry_after: None }, + )) + .expect("record second challenge"); + let second = store + .account("acct_primary") + .expect("account") + .cooldown_until() + .expect("second cooldown") + .duration_since(SystemTime::now()) + .expect("future cooldown"); + assert!(second >= Duration::from_secs(29) && second <= Duration::from_secs(30)); + + selector.record_success(&lease); + block_on(selector.record_failure( + &lease, + CodexAccountFailure::CloudflareChallenge { retry_after: None }, + )) + .expect("record reset challenge"); + let reset = store + .account("acct_primary") + .expect("account") + .cooldown_until() + .expect("reset cooldown") + .duration_since(SystemTime::now()) + .expect("future cooldown"); + assert!(reset >= Duration::from_secs(9) && reset <= Duration::from_secs(10)); +} + +#[test] +fn repeated_cloudflare_path_block_marks_only_the_affected_account_invalid() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "at-primary"); + create_account(&store, "acct_other", "at-other"); + let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); + let attempt = attempt_with_required( + BTreeSet::new(), + Some(ProviderAccountId::new("acct_primary").expect("account id")), + ); + let lease = + block_on( + selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") + .expect("request URL"), + attempt: &attempt, + }), + ) + .expect("select account"); + + for _ in 0..3 { + block_on(selector.record_failure(&lease, CodexAccountFailure::CloudflarePathBlocked)) + .expect("record path block"); + } assert_eq!( - ( - store - .account("acct_primary") - .expect("primary account") - .availability(), - store - .account("acct_other") - .expect("other account") - .availability(), - ), - (AccountAvailability::Ready, AccountAvailability::Ready) + store + .account("acct_primary") + .expect("affected account") + .availability(), + AccountAvailability::Invalid + ); + assert_eq!( + store + .account("acct_other") + .expect("other account") + .availability(), + AccountAvailability::Ready ); } + +#[test] +fn cloudflare_challenge_expires_provider_owned_cookies_at_cooldown_boundary() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "at-primary"); + let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); + let required = ProviderAccountId::new("acct_primary").expect("account id"); + let request_url = + Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL"); + let first_attempt = attempt_with_required(BTreeSet::new(), Some(required.clone())); + let first = block_on(selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &first_attempt, + })) + .expect("select account"); + block_on(selector.capture_response_cookies( + &first, + &request_url, + &["cf_clearance=old; Path=/; Domain=chatgpt.com; Secure; Max-Age=3600".to_owned()], + )) + .expect("capture cookie"); + + let second_attempt = attempt_with_required(BTreeSet::new(), Some(required)); + let second = block_on(selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &second_attempt, + })) + .expect("select revised account"); + block_on(selector.record_failure( + &second, + CodexAccountFailure::CloudflareChallenge { retry_after: None }, + )) + .expect("record challenge"); + + let account = store.account("acct_primary").expect("account"); + let cooldown_until = account.cooldown_until().expect("cooldown"); + let data = block_on(store.repository().load_complete_data(&account)).expect("credential data"); + assert_eq!(data.cookies.len(), 1); + assert!( + data.cookies[0] + .expires_at + .is_some_and(|expires_at| SystemTime::from(expires_at) <= cooldown_until) + ); +} + +#[test] +fn cloudflare_path_block_deletes_provider_owned_cookies() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "at-primary"); + let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); + let required = ProviderAccountId::new("acct_primary").expect("account id"); + let request_url = + Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL"); + let first_attempt = attempt_with_required(BTreeSet::new(), Some(required.clone())); + let first = block_on(selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &first_attempt, + })) + .expect("select account"); + block_on(selector.capture_response_cookies( + &first, + &request_url, + &["__cf_bm=old; Path=/; Domain=chatgpt.com; Secure; Max-Age=3600".to_owned()], + )) + .expect("capture cookie"); + + let second_attempt = attempt_with_required(BTreeSet::new(), Some(required)); + let second = block_on(selector.select(&SelectCodexCredential { + provider_instance_id: &instance_id(), + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &second_attempt, + })) + .expect("select revised account"); + block_on(selector.record_failure(&second, CodexAccountFailure::CloudflarePathBlocked)) + .expect("record path block"); + + let account = store.account("acct_primary").expect("account"); + let data = block_on(store.repository().load_complete_data(&account)).expect("credential data"); + assert!(data.cookies.is_empty()); +} diff --git a/backend/crates/providers/openai/tests/credential/identity.rs b/backend/crates/providers/openai/tests/credential/identity.rs index f6b7daab1..5246132d0 100644 --- a/backend/crates/providers/openai/tests/credential/identity.rs +++ b/backend/crates/providers/openai/tests/credential/identity.rs @@ -6,9 +6,11 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use jsonwebtoken::{Algorithm, EncodingKey, Header, encode, get_current_timestamp}; use provider_openai::credential::token_client::OFFICIAL_CODEX_OAUTH_CLIENT_ID; use provider_openai::credential::{ - CodexAuthorizationTokenVerifier, CodexIdentityVerificationError, CodexJwksSource, - CodexJwtIdentityVerifier, CodexOAuthSecret, CodexTokenIdentityVerifier, - OFFICIAL_OPENAI_API_AUDIENCE, OFFICIAL_OPENAI_ISSUER, ReqwestOpenAiJwksSource, + CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexAuthenticatedAccount, + CodexAuthenticatedAccountSource, CodexIdentityExpectation, CodexIdentityVerification, + CodexIdentityVerificationError, CodexJwksSource, CodexJwtIdentityVerifier, CodexOAuthSecret, + CodexSignedIdentityVerifier, OFFICIAL_OPENAI_API_AUDIENCE, OFFICIAL_OPENAI_ISSUER, + ReqwestOpenAiJwksSource, }; use rsa::RsaPrivateKey; use rsa::pkcs1::EncodeRsaPrivateKey; @@ -65,6 +67,21 @@ impl CodexJwksSource for StaticJwksSource { } } +struct StaticAccountSource { + result: Result, +} + +#[async_trait] +impl CodexAuthenticatedAccountSource for StaticAccountSource { + async fn fetch( + &self, + _secret: &CodexOAuthSecret, + _expectation: &CodexIdentityExpectation, + ) -> Result { + self.result.clone() + } +} + fn verifier() -> (CodexJwtIdentityVerifier, Arc) { let calls = Arc::new(AtomicUsize::new(0)); let private = RsaPrivateKey::from_pkcs8_pem(PRIVATE_KEY).expect("parse test RSA key"); @@ -130,6 +147,17 @@ fn secret(claims: &Value, refresh_token: Option<&str>) -> CodexOAuthSecret { } } +pub(crate) async fn signed_identity_fixture() -> provider_openai::credential::CodexSignedIdentity { + let (verifier, _) = verifier(); + verifier + .verify_access(&secret( + &claims(get_current_timestamp() + 3_600), + Some("refresh-token"), + )) + .await + .expect("signed identity fixture") +} + fn id_token_claims(nonce: &str, subject: &str) -> Value { let now = get_current_timestamp(); json!({ @@ -147,20 +175,128 @@ fn id_token_claims(nonce: &str, subject: &str) -> Value { async fn verifier_should_derive_identity_only_from_valid_signed_claims() { let (verifier, calls) = verifier(); let profile = verifier - .verify(&secret( + .verify_access(&secret( &claims(get_current_timestamp() + 3_600), Some("refresh-token"), )) .await .expect("verify signed Codex token"); - assert_eq!(profile.chatgpt_account_id, "account-signed"); - assert_eq!(profile.chatgpt_user_id.as_deref(), Some("user-signed")); - assert_eq!(profile.email.as_deref(), Some("signed@example.invalid")); - assert!(profile.next_refresh_at.is_some()); + assert_eq!(profile.oauth_subject(), "subject-signed"); + assert_eq!(profile.poid(), None); + assert_eq!(profile.claimed_account_id(), Some("account-signed")); + assert_eq!(profile.claimed_user_id(), Some("user-signed")); assert_eq!(calls.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn verifier_should_accept_current_signed_poid_account_identity() { + let (verifier, _) = verifier(); + let mut current = claims(get_current_timestamp() + 3_600); + current["https://api.openai.com/auth"] + .as_object_mut() + .expect("auth claims") + .remove("chatgpt_account_id"); + current["https://api.openai.com/auth"]["poid"] = json!("organization-signed"); + + let profile = verifier + .verify_access(&secret(¤t, Some("refresh-token"))) + .await + .expect("verify current signed Codex token"); + + assert_eq!(profile.oauth_subject(), "subject-signed"); + assert_eq!(profile.poid(), Some("organization-signed")); + assert_eq!(profile.claimed_account_id(), None); +} + +#[tokio::test] +async fn account_identity_requires_matching_authenticated_usage_facts() { + let (signed, _) = verifier(); + let service = CodexAccountIdentityService::new( + Arc::new(signed), + Arc::new(StaticAccountSource { + result: Ok(CodexAuthenticatedAccount { + chatgpt_account_id: "account-signed".to_owned(), + chatgpt_user_id: "user-signed".to_owned(), + email: Some("usage@example.invalid".to_owned()), + plan_type: Some("pro".to_owned()), + }), + }), + ); + let verification = service + .verify( + &secret(&claims(get_current_timestamp() + 3_600), None), + &CodexIdentityExpectation::imported( + Some("account-signed".to_owned()), + Some("user-signed".to_owned()), + ) + .expect("import expectation"), + ) + .await + .expect("complete authenticated identity"); + let CodexIdentityVerification::Complete(profile) = verification else { + panic!("usage facts must complete the identity"); + }; + + assert_eq!(profile.oauth_subject, "subject-signed"); + assert_eq!(profile.chatgpt_account_id, "account-signed"); + assert_eq!(profile.chatgpt_user_id, "user-signed"); + assert_eq!(profile.email.as_deref(), Some("usage@example.invalid")); +} + +#[tokio::test] +async fn account_identity_rejects_imported_workspace_conflict() { + let (signed, _) = verifier(); + let service = CodexAccountIdentityService::new( + Arc::new(signed), + Arc::new(StaticAccountSource { + result: Ok(CodexAuthenticatedAccount { + chatgpt_account_id: "account-signed".to_owned(), + chatgpt_user_id: "user-signed".to_owned(), + email: None, + plan_type: None, + }), + }), + ); + let error = service + .verify( + &secret(&claims(get_current_timestamp() + 3_600), None), + &CodexIdentityExpectation::imported(Some("other-account".to_owned()), None) + .expect("import expectation"), + ) + .await + .expect_err("document workspace cannot override authenticated usage"); + + assert_eq!(error, CodexIdentityVerificationError::Rejected); +} + +#[tokio::test] +async fn account_identity_keeps_usage_unavailable_out_of_ready_profile() { + let (signed, _) = verifier(); + let service = CodexAccountIdentityService::new( + Arc::new(signed), + Arc::new(StaticAccountSource { + result: Err(CodexIdentityVerificationError::Unavailable), + }), + ); + let verification = service + .verify( + &secret(&claims(get_current_timestamp() + 3_600), None), + &CodexIdentityExpectation::default(), + ) + .await + .expect("signed identity remains explicit"); + + assert!(matches!( + &verification, + CodexIdentityVerification::SignedOnly(_) + )); + assert_eq!( + verification.into_complete().expect_err("not Ready"), + CodexIdentityVerificationError::Unavailable + ); +} + #[tokio::test] async fn verifier_should_reject_expired_wrong_issuer_audience_or_unbound_tokens() { let (verifier, _) = verifier(); @@ -173,14 +309,11 @@ async fn verifier_should_reject_expired_wrong_issuer_audience_or_unbound_tokens( ]; cases[1]["iss"] = json!("https://attacker.invalid"); cases[2]["aud"] = json!(["https://attacker.invalid"]); - cases[3]["https://api.openai.com/auth"] - .as_object_mut() - .expect("auth claims") - .remove("chatgpt_account_id"); + cases[3]["sub"] = json!(""); for claims in cases { let error = verifier - .verify(&secret(&claims, None)) + .verify_access(&secret(&claims, None)) .await .expect_err("invalid identity token must fail"); assert_eq!(error, CodexIdentityVerificationError::Rejected); @@ -193,7 +326,7 @@ async fn verifier_should_reject_user_rebinding_and_invalid_refresh_shape() { let mut rebound = claims(get_current_timestamp() + 3_600); rebound["https://api.openai.com/auth"]["user_id"] = json!("other-user"); let rebound_error = verifier - .verify(&secret(&rebound, None)) + .verify_access(&secret(&rebound, None)) .await .expect_err("signed aliases must not disagree"); let access = token(&claims(get_current_timestamp() + 3_600)); @@ -203,7 +336,7 @@ async fn verifier_should_reject_user_rebinding_and_invalid_refresh_shape() { id_token: None, }; let refresh_error = verifier - .verify(&reused) + .verify_access(&reused) .await .expect_err("refresh token must have independent shape"); @@ -215,8 +348,14 @@ async fn verifier_should_reject_user_rebinding_and_invalid_refresh_shape() { async fn verifier_should_cache_jwks_and_fail_closed_when_source_is_unavailable() { let (verifier, calls) = verifier(); let valid = secret(&claims(get_current_timestamp() + 3_600), None); - verifier.verify(&valid).await.expect("first verification"); - verifier.verify(&valid).await.expect("cached verification"); + verifier + .verify_access(&valid) + .await + .expect("first verification"); + verifier + .verify_access(&valid) + .await + .expect("cached verification"); assert_eq!(calls.load(Ordering::SeqCst), 1); let unavailable = CodexJwtIdentityVerifier::new(Box::new(StaticJwksSource { @@ -225,7 +364,7 @@ async fn verifier_should_cache_jwks_and_fail_closed_when_source_is_unavailable() error: Some(CodexIdentityVerificationError::Unavailable), })); let error = unavailable - .verify(&valid) + .verify_access(&valid) .await .expect_err("unavailable JWKS must fail closed"); assert_eq!(error, CodexIdentityVerificationError::Unavailable); @@ -305,7 +444,7 @@ async fn supplied_real_access_only_account_should_pass_official_signature_verifi ReqwestOpenAiJwksSource::new().expect("build official JWKS source"), )); let profile = verifier - .verify(&CodexOAuthSecret { + .verify_access(&CodexOAuthSecret { access_token: SecretString::new(account.token.into()), refresh_token: account .refresh_token @@ -315,7 +454,6 @@ async fn supplied_real_access_only_account_should_pass_official_signature_verifi .await .expect("verify official access-only Codex account"); - assert!(!profile.chatgpt_account_id.is_empty()); - assert!(profile.access_token_expires_at.is_some()); - assert!(profile.next_refresh_at.is_none()); + assert!(!profile.oauth_subject().is_empty()); + assert!(profile.access_token_expires_at() > chrono::Utc::now()); } diff --git a/backend/crates/providers/openai/tests/credential/oauth.rs b/backend/crates/providers/openai/tests/credential/oauth.rs index e1fd805a4..0a0dbec71 100644 --- a/backend/crates/providers/openai/tests/credential/oauth.rs +++ b/backend/crates/providers/openai/tests/credential/oauth.rs @@ -1,21 +1,30 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use gateway_admin::model::provider_credentials::{ + AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision}; +use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::routing::ProviderKind; use provider_openai::credential::token_client::{ AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant, AuthorizationTokenSet, }; use provider_openai::credential::{ - CodexAccountProfile, CodexAuthorizationTokenVerifier, CodexCredentialAdmin, - CodexIdentityVerificationError, CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService, - CodexOAuthFlowBinding, CodexOAuthPendingStore, CodexOAuthPendingStoreError, CodexOAuthSecret, - CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CreateCodexCredential, - StartCodexOAuthAuthorization, StartCodexOAuthReauthorization, StoredCodexPendingAuthorization, + CodexAccountIdentityVerifier, CodexCredentialAdmin, CodexIdentityExpectation, + CodexIdentityVerification, CodexIdentityVerificationError, CodexOAuthAdmin, + CodexOAuthAdminError, CodexOAuthAdminService, CodexOAuthPendingStore, + CodexOAuthPendingStoreError, CodexOAuthSecret, CodexPendingAuthorization, + CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, CreateCodexCredential, + StartCodexOAuthAuthorization, StoredCodexPendingAuthorization, }; use secrecy::{ExposeSecret, SecretString}; +use sha2::{Digest as _, Sha256}; use url::Url; -use crate::support::{MemoryAccountStore, instance_id, profile, secret}; +use crate::support::{MemoryAccountStore, instance_id, profile, runtime_policy, secret}; #[derive(Default)] struct PendingStore { @@ -48,6 +57,7 @@ impl CodexOAuthPendingStore for PendingStore { reauthorization_credential_revision: pending .reauthorization() .map(|target| target.credential_revision().get()), + mutation: pending.mutation().clone(), }); Ok(()) } @@ -96,18 +106,29 @@ impl AuthorizationCodeExchanger for Exchanger { struct Verifier; #[async_trait] -impl CodexAuthorizationTokenVerifier for Verifier { +impl CodexAccountIdentityVerifier for Verifier { + async fn verify( + &self, + _secret: &CodexOAuthSecret, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } + async fn verify_authorization( &self, _secret: &CodexOAuthSecret, id_token: &SecretString, expected_nonce: &SecretString, - ) -> Result { + _expectation: &CodexIdentityExpectation, + ) -> Result { if id_token.expose_secret() != "oauth-id-token" || expected_nonce.expose_secret().len() < 16 { return Err(CodexIdentityVerificationError::Rejected); } - Ok(profile("chatgpt-oauth")) + Ok(CodexIdentityVerification::Complete(profile( + "chatgpt-oauth", + ))) } } @@ -117,13 +138,14 @@ fn service() -> CodexOAuthAdminService { fn service_with_store( store: Arc, - verifier: Arc, + verifier: Arc, ) -> CodexOAuthAdminService { CodexOAuthAdminService::new( Arc::new(PendingStore::default()), Arc::new(Exchanger), verifier, store, + runtime_policy(), CodexCredentialAdmin, ) } @@ -133,14 +155,55 @@ async fn started( ) -> provider_openai::credential::CodexOAuthAuthorizationStarted { service .start_authorization(StartCodexOAuthAuthorization { - binding: CodexOAuthFlowBinding::new("owner-one", "request-one").expect("binding"), - provider_instance_id: "inst_openai_primary".to_owned(), - name: "OAuth Account".to_owned(), + mutation: create_mutation("request-one"), }) .await .expect("start OAuth") } +fn owner_context(request_id: &str) -> MutationContext { + MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "owner-one".to_owned(), + }, + request_id: request_id.to_owned(), + } +} + +fn create_mutation(request_id: &str) -> PendingAuthorizationMutation { + let context = owner_context(request_id); + PendingAuthorizationMutation::new( + Revision::new(1).expect("revision"), + ProviderKind::new("openai").expect("provider"), + AuthorizationMutationTarget::Create { + provider_instance_id: instance_id(), + name: "OAuth Account".to_owned(), + }, + AuthorizationOwnerBinding::from_context(&context), + ) +} + +fn reauthorization_mutation(request_id: &str, account_id: &str) -> PendingAuthorizationMutation { + let context = owner_context(request_id); + PendingAuthorizationMutation::new( + Revision::new(1).expect("revision"), + ProviderKind::new("openai").expect("provider"), + AuthorizationMutationTarget::Reauthorize { + provider_instance_id: instance_id(), + account_id: ProviderAccountId::new(account_id).expect("account id"), + expected_credential_revision: Revision::new(1).expect("revision"), + }, + AuthorizationOwnerBinding::from_context(&context), + ) +} + +fn owner_ref() -> String { + let mut digest = Sha256::new(); + digest.update(b"admin-session\0"); + digest.update(b"owner-one"); + URL_SAFE_NO_PAD.encode(digest.finalize()) +} + #[tokio::test] async fn start_authorization_uses_pkce_nonce_and_fixed_official_redirect() { let service = service(); @@ -176,7 +239,7 @@ async fn completion_returns_verified_core_account_without_writing_store() { .expect("state"); let prepared = service .complete_authorization(CompleteCodexOAuthAuthorization { - owner_ref: "owner-one".to_owned(), + owner_ref: owner_ref(), flow_id: started.flow_id, callback_url: SecretString::from(format!( "http://localhost:1455/auth/callback?code=authorization-code&state={state}" @@ -184,6 +247,9 @@ async fn completion_returns_verified_core_account_without_writing_store() { }) .await .expect("complete OAuth"); + let CompletedCodexOAuthCredential::Create(prepared) = prepared.credential else { + panic!("expected prepared create"); + }; assert_eq!(prepared.account.provider().as_str(), "openai"); assert!(prepared.account.id().as_str().starts_with("acct_")); assert_eq!( @@ -210,7 +276,7 @@ async fn callback_state_mismatch_fails_closed_and_consumes_flow() { let started = started(&service).await; let flow_id = started.flow_id; let command = || CompleteCodexOAuthAuthorization { - owner_ref: "owner-one".to_owned(), + owner_ref: owner_ref(), flow_id: flow_id.clone(), callback_url: SecretString::from( "http://localhost:1455/auth/callback?code=authorization-code&state=wrong-state-value", @@ -256,20 +322,24 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati name: "reauthorize".to_owned(), secret: secret("old-oauth-access"), account: profile("chatgpt-oauth"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await .expect("create account"); + let original_account = store + .account("acct_oauth_reauth") + .expect("original account"); + let original_installation_id = store + .repository() + .load_runtime_credential(&original_account) + .await + .expect("original credential") + .installation_id; let service = service_with_store(store.clone(), Arc::new(Verifier)); let started = service - .start_reauthorization(StartCodexOAuthReauthorization { - binding: CodexOAuthFlowBinding::new("owner-one", "request-reauth").expect("binding"), - account_id: gateway_core::engine::credential::ProviderAccountId::new( - "acct_oauth_reauth", - ) - .expect("account id"), - expected_credential_revision: - gateway_core::engine::credential::CredentialRevision::new(1).expect("revision"), + .start_authorization(StartCodexOAuthAuthorization { + mutation: reauthorization_mutation("request-reauth", "acct_oauth_reauth"), }) .await .expect("start reauthorization"); @@ -279,8 +349,8 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) .expect("state"); let prepared = service - .complete_reauthorization(CompleteCodexOAuthAuthorization { - owner_ref: "owner-one".to_owned(), + .complete_authorization(CompleteCodexOAuthAuthorization { + owner_ref: owner_ref(), flow_id: started.flow_id, callback_url: SecretString::from(format!( "http://localhost:1455/auth/callback?code=authorization-code&state={state}" @@ -288,6 +358,9 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati }) .await .expect("complete reauthorization"); + let CompletedCodexOAuthCredential::Reauthorize(prepared) = prepared.credential else { + panic!("expected prepared reauthorization"); + }; assert_eq!( prepared.credential.account_id().as_str(), @@ -301,6 +374,7 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati runtime.secret.access_token.expose_secret(), "oauth-access-token" ); + assert_eq!(runtime.installation_id, original_installation_id); assert_eq!( store .account("acct_oauth_reauth") @@ -314,14 +388,25 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati struct RebindingVerifier; #[async_trait] -impl CodexAuthorizationTokenVerifier for RebindingVerifier { +impl CodexAccountIdentityVerifier for RebindingVerifier { + async fn verify( + &self, + _: &CodexOAuthSecret, + _: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } + async fn verify_authorization( &self, _: &CodexOAuthSecret, _: &SecretString, _: &SecretString, - ) -> Result { - Ok(profile("chatgpt-different-owner")) + _: &CodexIdentityExpectation, + ) -> Result { + Ok(CodexIdentityVerification::Complete(profile( + "chatgpt-different-owner", + ))) } } @@ -336,20 +421,15 @@ async fn reauthorization_rejects_identity_rebinding() { name: "reauthorize".to_owned(), secret: secret("old-oauth-access"), account: profile("chatgpt-oauth"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await .expect("create account"); let service = service_with_store(store, Arc::new(RebindingVerifier)); let started = service - .start_reauthorization(StartCodexOAuthReauthorization { - binding: CodexOAuthFlowBinding::new("owner-one", "request-rebind").expect("binding"), - account_id: gateway_core::engine::credential::ProviderAccountId::new( - "acct_oauth_rebind", - ) - .expect("account id"), - expected_credential_revision: - gateway_core::engine::credential::CredentialRevision::new(1).expect("revision"), + .start_authorization(StartCodexOAuthAuthorization { + mutation: reauthorization_mutation("request-rebind", "acct_oauth_rebind"), }) .await .expect("start reauthorization"); @@ -359,8 +439,8 @@ async fn reauthorization_rejects_identity_rebinding() { .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) .expect("state"); let error = service - .complete_reauthorization(CompleteCodexOAuthAuthorization { - owner_ref: "owner-one".to_owned(), + .complete_authorization(CompleteCodexOAuthAuthorization { + owner_ref: owner_ref(), flow_id: started.flow_id, callback_url: SecretString::from(format!( "http://localhost:1455/auth/callback?code=authorization-code&state={state}" diff --git a/backend/crates/providers/openai/tests/credential/quota.rs b/backend/crates/providers/openai/tests/credential/quota.rs index b8452fb97..53b1e4d3c 100644 --- a/backend/crates/providers/openai/tests/credential/quota.rs +++ b/backend/crates/providers/openai/tests/credential/quota.rs @@ -2,21 +2,21 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use chrono::{TimeZone, Utc}; +use futures::future::join_all; use gateway_core::engine::credential::{ AccountAvailability, AccountStateChange, ProviderAccountStore, }; use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderKind}; -use provider_openai::CodexEndpointPolicy; use provider_openai::credential::{ CodexCredentialQuotaService, CodexQuotaSyncSummary, CodexQuotaWindowKind, CreateCodexCredential, parse_codex_quota_usage, }; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; use serde_json::json; -use wiremock::matchers::{header, headers, method, path}; +use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use crate::support::{MemoryAccountStore, instance_id, profile, secret}; +use crate::support::{MemoryAccountStore, instance_id, loopback_origin_policy, profile, secret}; #[test] fn parser_extracts_dynamic_windows_without_a_fixed_database_shape() { @@ -133,6 +133,42 @@ fn wire_profile() -> CodexWireProfileState { }) } +fn quota_service(store: &Arc) -> CodexCredentialQuotaService { + CodexCredentialQuotaService::new( + store.repository(), + wire_profile(), + reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"), + loopback_origin_policy(), + ) +} + +#[tokio::test] +async fn concurrent_cold_scheduling_hydration_reads_quota_once() { + let store = Arc::new(MemoryAccountStore::default()); + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_hydration".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "hydration".to_owned(), + secret: secret("hydration-token"), + account: profile("chatgpt-acct_hydration"), + next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await + .expect("create account"); + let account = store.account("acct_hydration").expect("created account"); + let service = quota_service(&store); + + join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; + + assert_eq!(store.quota_reads(), 1); +} + #[tokio::test] async fn quota_service_stores_raw_provider_json_and_projects_common_state() { let server = MockServer::start().await; @@ -164,12 +200,6 @@ async fn quota_service_stores_raw_provider_json_and_projects_common_state() { Mock::given(method("GET")) .and(path("/backend-api/wham/usage")) .and(header("authorization", "Bearer quota-token")) - .and(header("openai-beta", "codex-1")) - .and(header("oai-language", "zh-CN")) - .and(header("sec-fetch-site", "none")) - .and(header("sec-fetch-mode", "no-cors")) - .and(header("sec-fetch-dest", "empty")) - .and(headers("priority", vec!["u=4", "i"])) .respond_with(ResponseTemplate::new(200).set_body_json(&raw)) .expect(2) .mount(&server) @@ -183,16 +213,12 @@ async fn quota_service_stores_raw_provider_json_and_projects_common_state() { name: "quota".to_owned(), secret: secret("quota-token"), account: profile("chatgpt-acct_quota"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await .expect("create account"); - let service = CodexCredentialQuotaService::new_with_endpoint_policy( - store.repository(), - wire_profile(), - CodexEndpointPolicy::Loopback, - ) - .expect("quota service"); + let service = quota_service(&store); let instance = ProviderInstance::new( instance_id(), ProviderKind::new("openai").expect("provider"), @@ -229,9 +255,10 @@ async fn quota_service_stores_raw_provider_json_and_projects_common_state() { assert_eq!(refreshed.account_id(), account.id()); assert_eq!(refreshed.windows()[0].used_percent(), Some(37.0)); let observation = store - .get_quota(account.id()) + .get_quotas(std::slice::from_ref(account.id())) .await .expect("quota read") + .pop() .expect("quota observation"); assert_eq!( observation @@ -270,17 +297,13 @@ async fn quota_refresh_success_must_not_clear_newer_future_cooldown() { name: "quota cooldown".to_owned(), secret: secret("quota-cooldown-token"), account: profile("chatgpt-acct_quota_cooldown"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await .expect("create account"); let account = store.account("acct_quota_cooldown").expect("quota account"); - let service = CodexCredentialQuotaService::new_with_endpoint_policy( - store.repository(), - wire_profile(), - CodexEndpointPolicy::Loopback, - ) - .expect("quota service"); + let service = quota_service(&store); let instance = ProviderInstance::new( instance_id(), ProviderKind::new("openai").expect("provider"), @@ -335,6 +358,73 @@ async fn non_exhausted_quota_refresh_should_restore_quota_exhausted_account() { ); } +#[tokio::test] +async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision_fence() { + let store = Arc::new(MemoryAccountStore::default()); + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_passive_quota".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "passive quota".to_owned(), + secret: secret("passive-quota-token"), + account: profile("chatgpt-acct_passive_quota"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await + .expect("create account"); + let account = store.account("acct_passive_quota").expect("account"); + let service = quota_service(&store); + let reset_at = 1_900_000_000_i64; + let headers = vec![ + ("x-codex-active-limit".to_owned(), "codex".to_owned()), + ("x-codex-primary-used-percent".to_owned(), "100".to_owned()), + ( + "x-codex-primary-window-minutes".to_owned(), + "300".to_owned(), + ), + ("x-codex-primary-reset-at".to_owned(), reset_at.to_string()), + ("x-codex-limit-reached".to_owned(), "true".to_owned()), + ]; + + assert!( + service + .synchronize_passive_headers(&account, &headers) + .await + .expect("passive quota") + ); + let current = store + .account("acct_passive_quota") + .expect("current account"); + assert_eq!(current.availability(), AccountAvailability::QuotaExhausted); + let observation = store + .get_quotas(std::slice::from_ref(account.id())) + .await + .expect("quota") + .pop() + .expect("quota observation"); + let quota = serde_json::Value::Object( + observation + .quota + .expect("quota JSON") + .expose_to_provider() + .clone(), + ); + assert_eq!( + quota + .pointer("/rate_limit/primary_window/limit_window_seconds") + .and_then(serde_json::Value::as_u64), + Some(18_000) + ); + assert_eq!( + quota + .pointer("/rate_limit/primary_window/reset_at") + .and_then(serde_json::Value::as_i64), + Some(reset_at) + ); +} + #[tokio::test] async fn exhausted_quota_refresh_should_upgrade_future_cooldown() { let cooldown_until = SystemTime::now() + Duration::from_secs(3_600); @@ -417,6 +507,7 @@ async fn synchronize_account_from_state( name: account_id.to_owned(), secret: secret(&format!("token-{account_id}")), account: profile(&format!("chatgpt-{account_id}")), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await @@ -433,12 +524,7 @@ async fn synchronize_account_from_state( }) .await .expect("apply initial account state"); - let service = CodexCredentialQuotaService::new_with_endpoint_policy( - store.repository(), - wire_profile(), - CodexEndpointPolicy::Loopback, - ) - .expect("quota service"); + let service = quota_service(&store); let instance = ProviderInstance::new( instance_id(), ProviderKind::new("openai").expect("provider"), diff --git a/backend/crates/providers/openai/tests/credential/refresh.rs b/backend/crates/providers/openai/tests/credential/refresh.rs index dd4e14e3d..d86e23f06 100644 --- a/backend/crates/providers/openai/tests/credential/refresh.rs +++ b/backend/crates/providers/openai/tests/credential/refresh.rs @@ -1,18 +1,23 @@ use std::collections::VecDeque; use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::time::{Duration, SystemTime}; use async_trait::async_trait; +use futures::future::BoxFuture; use gateway_core::engine::credential::AccountAvailability; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest, + ProviderStoreError, +}; use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use provider_openai::credential::{ - CodexCredentialRefreshOutcome, CodexCredentialRefreshService, CodexRefreshLeaseAcquisition, - CodexRefreshLeaseCoordinator, CodexRefreshLeaseError, CodexRefreshLeaseRequest, - CreateCodexCredential, + CodexAccountIdentityVerifier, CodexCredentialRefreshOutcome, CodexCredentialRefreshService, + CodexIdentityExpectation, CodexIdentityVerification, CodexIdentityVerificationError, + CodexOAuthSecret, CodexSignedIdentity, CreateCodexCredential, RotateCodexCredential, }; -use secrecy::ExposeSecret; +use secrecy::{ExposeSecret, SecretString}; -use crate::support::{MemoryAccountStore, instance_id, profile, secret}; +use crate::support::{MemoryAccountStore, instance_id, profile, runtime_policy, secret}; struct Refresher { outcomes: Mutex>>, @@ -49,23 +54,163 @@ impl TokenRefresher for Refresher { struct RefreshLeases { available: bool, - requests: Mutex>, + requests: Mutex>, +} + +struct VerifiedIdentity; + +#[async_trait] +impl CodexAccountIdentityVerifier for VerifiedIdentity { + async fn verify( + &self, + _secret: &CodexOAuthSecret, + expectation: &CodexIdentityExpectation, + ) -> Result { + let account_id = expectation + .chatgpt_account_id() + .ok_or(CodexIdentityVerificationError::Rejected)?; + Ok(CodexIdentityVerification::Complete(profile(account_id))) + } + + async fn verify_authorization( + &self, + _secret: &CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } +} + +struct RejectedIdentity; + +#[async_trait] +impl CodexAccountIdentityVerifier for RejectedIdentity { + async fn verify( + &self, + _secret: &CodexOAuthSecret, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } + + async fn verify_authorization( + &self, + _secret: &CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } +} + +struct UnavailableIdentity; + +#[async_trait] +impl CodexAccountIdentityVerifier for UnavailableIdentity { + async fn verify( + &self, + _secret: &CodexOAuthSecret, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Unavailable) + } + + async fn verify_authorization( + &self, + _secret: &CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Unavailable) + } +} + +struct SignedOnlyIdentity(CodexSignedIdentity); + +struct RacingPermanentRefresher { + store: Arc, } #[async_trait] -impl CodexRefreshLeaseCoordinator for RefreshLeases { - async fn try_acquire( +impl TokenRefresher for RacingPermanentRefresher { + async fn refresh(&self, _: &str) -> Result { + self.store + .repository() + .rotate_oauth_secret(RotateCodexCredential { + account_id: "acct_refresh".to_owned(), + expected_credential_revision: 1, + secret: secret("concurrent-access"), + verified_account: profile("chatgpt-acct_refresh"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + }) + .await + .expect("concurrent credential rotation"); + Err(RefreshFailure::InvalidGrant) + } +} + +#[async_trait] +impl CodexAccountIdentityVerifier for SignedOnlyIdentity { + async fn verify( &self, - request: &CodexRefreshLeaseRequest, - ) -> Result { - self.requests - .lock() - .expect("refresh lease lock") - .push(request.clone()); - Ok(if self.available { - CodexRefreshLeaseAcquisition::Acquired(Box::new(())) - } else { - CodexRefreshLeaseAcquisition::Unavailable + _secret: &CodexOAuthSecret, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Ok(CodexIdentityVerification::SignedOnly(self.0.clone())) + } + + async fn verify_authorization( + &self, + _secret: &CodexOAuthSecret, + _id_token: &SecretString, + _expected_nonce: &SecretString, + _expectation: &CodexIdentityExpectation, + ) -> Result { + Err(CodexIdentityVerificationError::Rejected) + } +} + +impl ProviderLeasePort for RefreshLeases { + fn load_state<'a>( + &'a self, + _: &'a gateway_core::routing::ProviderInstanceId, + _: &'a [gateway_core::engine::credential::ProviderAccountId], + ) -> BoxFuture< + 'a, + Result, + > { + Box::pin(async { + Ok(gateway_core::provider_ports::ProviderSchedulingState::new( + Default::default(), + None, + 0, + )) + }) + } + + fn try_acquire( + &self, + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + match request { + ProviderLeaseRequest::RefreshCapacity(_) => {} + ProviderLeaseRequest::Refresh(request) => self + .requests + .lock() + .expect("refresh lease lock") + .push(request), + ProviderLeaseRequest::Scheduling(_) => panic!("unexpected scheduling lease"), + } + Ok(if self.available { + ProviderLeaseAcquisition::Acquired(Box::new(())) + } else { + ProviderLeaseAcquisition::Busy { retry_after: None } + }) }) } } @@ -87,6 +232,7 @@ async fn setup( name: "refresh".to_owned(), secret: secret("old-access"), account: profile("chatgpt-acct_refresh"), + next_refresh_at: Some(chrono::Utc::now() - chrono::Duration::seconds(1)), enabled: true, }) .await @@ -95,13 +241,13 @@ async fn setup( let service = CodexCredentialRefreshService::new( store.repository(), refresher.clone(), + Arc::new(VerifiedIdentity), Arc::new(RefreshLeases { available: lease_available, requests: Mutex::new(Vec::new()), }), - Duration::from_secs(60 * 60), - ) - .expect("refresh service"); + runtime_policy(), + ); (store, refresher, service) } @@ -116,7 +262,14 @@ fn success_tokens() -> TokenPair { #[tokio::test] async fn successful_refresh_uses_redis_lease_and_cas_rotates_plaintext_tokens() { let (store, refresher, service) = setup(Ok(success_tokens()), true).await; - let outcomes = service.refresh_due(10).await.expect("refresh due"); + let original_account = store.account("acct_refresh").expect("original account"); + let original_installation_id = store + .repository() + .load_runtime_credential(&original_account) + .await + .expect("original credential") + .installation_id; + let outcomes = service.refresh_due().await.expect("refresh due"); assert!(matches!( outcomes.as_slice(), [CodexCredentialRefreshOutcome::Refreshed { @@ -132,16 +285,131 @@ async fn successful_refresh_uses_redis_lease_and_cas_rotates_plaintext_tokens() .await .expect("rotated runtime credential"); assert_eq!(runtime.secret.access_token.expose_secret(), "new-access"); + assert_eq!(runtime.installation_id, original_installation_id); assert_eq!( refresher.seen.lock().expect("seen tokens lock").as_slice(), ["rt-old-access"] ); } +#[tokio::test] +async fn refreshed_identity_rejection_revision_fences_account_as_invalid() { + let (store, _, _) = setup(Ok(success_tokens()), true).await; + let service = CodexCredentialRefreshService::new( + store.repository(), + Arc::new(Refresher::new(Ok(success_tokens()))), + Arc::new(RejectedIdentity), + Arc::new(RefreshLeases { + available: true, + requests: Mutex::new(Vec::new()), + }), + runtime_policy(), + ); + + let outcomes = service.refresh_due().await.expect("refresh due"); + + assert!(matches!( + outcomes.as_slice(), + [CodexCredentialRefreshOutcome::Invalidated { .. }] + )); + assert_eq!( + store + .account("acct_refresh") + .expect("invalid account") + .availability(), + AccountAvailability::Invalid + ); +} + +#[tokio::test] +async fn unavailable_signature_verification_persists_refresh_backoff() { + let (store, _, _) = setup(Ok(success_tokens()), true).await; + let service = CodexCredentialRefreshService::new( + store.repository(), + Arc::new(Refresher::new(Ok(success_tokens()))), + Arc::new(UnavailableIdentity), + Arc::new(RefreshLeases { + available: true, + requests: Mutex::new(Vec::new()), + }), + runtime_policy(), + ); + let before = SystemTime::now(); + + let outcomes = service.refresh_due().await.expect("refresh due"); + + assert!(matches!( + outcomes.as_slice(), + [CodexCredentialRefreshOutcome::Ambiguous { .. }] + )); + let account = store.account("acct_refresh").expect("deferred account"); + assert_eq!(account.revision().get(), 1); + assert_eq!(account.availability(), AccountAvailability::Invalid); + assert!( + account + .next_refresh_at() + .is_some_and(|retry| retry < before) + ); +} + +#[tokio::test] +async fn unavailable_usage_preserves_rotated_tokens_and_persists_backoff() { + let store = Arc::new(MemoryAccountStore::default()); + let signed = super::identity::signed_identity_fixture().await; + let mut account_profile = profile("account-signed"); + account_profile.oauth_subject = signed.oauth_subject().to_owned(); + account_profile.poid = signed.poid().map(str::to_owned); + account_profile.chatgpt_user_id = "user-signed".to_owned(); + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_signed_only".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "signed only".to_owned(), + secret: secret("old-access"), + account: account_profile, + next_refresh_at: Some(chrono::Utc::now() - chrono::Duration::seconds(1)), + enabled: true, + }) + .await + .expect("create account"); + let service = CodexCredentialRefreshService::new( + store.repository(), + Arc::new(Refresher::new(Ok(success_tokens()))), + Arc::new(SignedOnlyIdentity(signed)), + Arc::new(RefreshLeases { + available: true, + requests: Mutex::new(Vec::new()), + }), + runtime_policy(), + ); + let before = SystemTime::now(); + + let outcomes = service.refresh_due().await.expect("refresh due"); + + assert!(matches!( + outcomes.as_slice(), + [CodexCredentialRefreshOutcome::Transient { .. }] + )); + let account = store.account("acct_signed_only").expect("rotated account"); + assert_eq!(account.revision().get(), 2); + assert!( + account + .next_refresh_at() + .is_some_and(|retry| retry > before) + ); + let runtime = store + .repository() + .load_runtime_credential(&account) + .await + .expect("rotated credential"); + assert_eq!(runtime.secret.access_token.expose_secret(), "new-access"); +} + #[tokio::test] async fn invalid_grant_marks_unified_account_expired() { let (store, _, service) = setup(Err(RefreshFailure::InvalidGrant), true).await; - let outcomes = service.refresh_due(10).await.expect("refresh due"); + let outcomes = service.refresh_due().await.expect("refresh due"); assert!(matches!( outcomes.as_slice(), [CodexCredentialRefreshOutcome::Invalidated { .. }] @@ -155,10 +423,35 @@ async fn invalid_grant_marks_unified_account_expired() { ); } +#[tokio::test] +async fn permanent_refresh_failure_cannot_overwrite_a_newer_credential_revision() { + let (store, _, _) = setup(Ok(success_tokens()), true).await; + let service = CodexCredentialRefreshService::new( + store.repository(), + Arc::new(RacingPermanentRefresher { + store: Arc::clone(&store), + }), + Arc::new(VerifiedIdentity), + Arc::new(RefreshLeases { + available: true, + requests: Mutex::new(Vec::new()), + }), + runtime_policy(), + ); + + assert!(matches!( + service.refresh_due().await.expect("refresh cycle").as_slice(), + [CodexCredentialRefreshOutcome::Stale { account_id }] if account_id == "acct_refresh" + )); + let account = store.account("acct_refresh").expect("account"); + assert_eq!(account.revision().get(), 2); + assert_eq!(account.availability(), AccountAvailability::Ready); +} + #[tokio::test] async fn upstream_ban_marks_unified_account_banned() { let (store, _, service) = setup(Err(RefreshFailure::Banned), true).await; - service.refresh_due(10).await.expect("refresh due"); + service.refresh_due().await.expect("refresh due"); assert_eq!( store .account("acct_refresh") @@ -171,33 +464,46 @@ async fn upstream_ban_marks_unified_account_banned() { #[tokio::test] async fn proven_pre_send_transport_failure_uses_short_cooldown() { let (store, _, service) = setup(Err(RefreshFailure::RetryableTransport), true).await; - let outcomes = service.refresh_due(10).await.expect("refresh due"); + let outcomes = service.refresh_due().await.expect("refresh due"); assert!(matches!( outcomes.as_slice(), [CodexCredentialRefreshOutcome::Transient { .. }] )); let account = store.account("acct_refresh").expect("account"); - assert_eq!(account.availability(), AccountAvailability::Cooldown); - assert!(account.cooldown_until().is_some()); + assert_eq!(account.availability(), AccountAvailability::Ready); + assert_eq!(account.revision().get(), 2); + assert!( + account + .next_refresh_at() + .is_some_and(|retry| retry > SystemTime::now()) + ); } #[tokio::test] async fn ambiguous_refresh_does_not_mutate_tokens_or_account_state() { - let (store, _, service) = setup(Err(RefreshFailure::Transport), true).await; - let outcomes = service.refresh_due(10).await.expect("refresh due"); + let (store, refresher, service) = setup(Err(RefreshFailure::Transport), true).await; + let outcomes = service.refresh_due().await.expect("refresh due"); assert!(matches!( outcomes.as_slice(), [CodexCredentialRefreshOutcome::Ambiguous { .. }] )); let account = store.account("acct_refresh").expect("account"); assert_eq!(account.revision().get(), 1); - assert_eq!(account.availability(), AccountAvailability::Ready); + assert_eq!(account.availability(), AccountAvailability::Invalid); + assert!( + service + .refresh_due() + .await + .expect("second refresh cycle") + .is_empty() + ); + assert_eq!(refresher.seen.lock().expect("seen tokens lock").len(), 1); } #[tokio::test] async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() { let (store, refresher, service) = setup(Ok(success_tokens()), false).await; - let outcomes = service.refresh_due(10).await.expect("refresh due"); + let outcomes = service.refresh_due().await.expect("refresh due"); assert!(matches!( outcomes.as_slice(), [CodexCredentialRefreshOutcome::LeaseUnavailable { .. }] @@ -216,8 +522,14 @@ async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() { #[tokio::test] async fn invalid_refresh_batch_limit_fails_before_scanning_accounts() { let (_, _, service) = setup(Ok(success_tokens()), true).await; - assert!(service.refresh_due(0).await.is_err()); - assert!(service.refresh_due(1_001).await.is_err()); + assert_eq!( + service + .refresh_due() + .await + .expect("bounded provider batch") + .len(), + 1 + ); } #[tokio::test] @@ -232,6 +544,7 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() { name: account_id.to_owned(), secret: secret(account_id), account: profile(&format!("chatgpt-{account_id}")), + next_refresh_at: Some(chrono::Utc::now() - chrono::Duration::seconds(1)), enabled: true, }) .await @@ -248,18 +561,15 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() { let service = CodexCredentialRefreshService::new( store.repository(), refresher, + Arc::new(VerifiedIdentity), Arc::new(RefreshLeases { available: true, requests: Mutex::new(Vec::new()), }), - Duration::from_secs(60 * 60), - ) - .expect("refresh service"); + runtime_policy(), + ); - let outcomes = service - .refresh_due(10) - .await - .expect("isolated refresh cycle"); + let outcomes = service.refresh_due().await.expect("isolated refresh cycle"); assert!(matches!( outcomes.as_slice(), diff --git a/backend/crates/providers/openai/tests/credential/types.rs b/backend/crates/providers/openai/tests/credential/types.rs index fe1320b68..2764a49f0 100644 --- a/backend/crates/providers/openai/tests/credential/types.rs +++ b/backend/crates/providers/openai/tests/credential/types.rs @@ -1,6 +1,7 @@ use chrono::Utc; use provider_openai::credential::{ - CodexAccountProfile, CodexCookie, CodexCredentialData, CodexOAuthSecret, + CodexAccountProfile, CodexCookie, CodexCredentialData, CodexCredentialPrincipal, + CodexOAuthSecret, }; use secrecy::SecretString; @@ -21,11 +22,12 @@ fn oauth_secret_debug_redacts_every_token() { fn account_profile_debug_redacts_identity_fields() { let profile = CodexAccountProfile { email: Some("private@example.com".to_owned()), + oauth_subject: "subject-private".to_owned(), + poid: Some("poid-private".to_owned()), chatgpt_account_id: "chatgpt-private".to_owned(), - chatgpt_user_id: Some("user-private".to_owned()), + chatgpt_user_id: "user-private".to_owned(), plan_type: Some("pro".to_owned()), access_token_expires_at: Some(Utc::now()), - next_refresh_at: None, }; let debug = format!("{profile:?}"); assert!(!debug.contains("private@example.com")); @@ -38,6 +40,11 @@ fn account_profile_debug_redacts_identity_fields() { fn plaintext_provider_schema_round_trips_dynamic_cookie_data() { let data = CodexCredentialData { schema_version: 1, + principal: CodexCredentialPrincipal { + oauth_subject: "subject-private".to_owned(), + poid: Some("poid-private".to_owned()), + }, + installation_id: "00000000-0000-4000-8000-000000000001".to_owned(), access_token: "at".to_owned(), refresh_token: Some("rt".to_owned()), id_token: None, @@ -65,9 +72,19 @@ fn plaintext_provider_schema_round_trips_dynamic_cookie_data() { fn provider_schema_rejects_unknown_public_layer_fields() { let value = serde_json::json!({ "schema_version": 1, + "principal": {"oauth_subject": "subject-private", "poid": null}, + "installation_id": "00000000-0000-4000-8000-000000000001", "access_token": "at", "cookies": [], "unknown_field": 9 }); assert!(serde_json::from_value::(value).is_err()); + assert!( + serde_json::from_value::(serde_json::json!({ + "schema_version": 1, + "access_token": "at", + "cookies": [] + })) + .is_err() + ); } diff --git a/backend/crates/providers/openai/tests/main.rs b/backend/crates/providers/openai/tests/main.rs index 6d4472d7f..c4aa9a5f2 100644 --- a/backend/crates/providers/openai/tests/main.rs +++ b/backend/crates/providers/openai/tests/main.rs @@ -1,3 +1,5 @@ +mod admin; +mod config; mod credential; mod provider; mod support; diff --git a/backend/crates/providers/openai/tests/provider/config.rs b/backend/crates/providers/openai/tests/provider/config.rs index 8028a6778..7e58a523f 100644 --- a/backend/crates/providers/openai/tests/provider/config.rs +++ b/backend/crates/providers/openai/tests/provider/config.rs @@ -1,9 +1,11 @@ use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderInstanceId, ProviderKind}; use provider_openai::{ - CodexEndpointPolicy, CodexProviderConfigError, CodexProviderInstanceConfig, - CodexProviderTransport, OFFICIAL_CODEX_BASE_PATH, + CodexProviderConfigError, CodexProviderInstanceConfig, CodexProviderTransport, + OFFICIAL_CODEX_BASE_PATH, OfficialCodexOriginPolicy, }; +use crate::support::LoopbackCodexOriginPolicy; + fn instance(provider: &str, base_url: &str) -> ProviderInstance { ProviderInstance::new( ProviderInstanceId::new("inst_openai").expect("instance id"), @@ -16,13 +18,16 @@ fn instance(provider: &str, base_url: &str) -> ProviderInstance { #[test] fn official_instance_requires_exact_https_origin_and_base_path() { - let compiled = CodexProviderInstanceConfig::from_snapshot(&instance( - "openai", - "https://chatgpt.com/backend-api/", - )) + let compiled = CodexProviderInstanceConfig::from_snapshot( + &instance("openai", "https://chatgpt.com/backend-api/"), + &OfficialCodexOriginPolicy, + ) .expect("official endpoint"); assert_eq!(compiled.base_url().path(), OFFICIAL_CODEX_BASE_PATH); - assert_eq!(compiled.transport(), CodexProviderTransport::HttpSse); + assert_eq!( + compiled.transport(), + CodexProviderTransport::PreferWebSocket + ); } #[test] @@ -37,7 +42,10 @@ fn official_policy_rejects_lookalike_private_and_metadata_origins() { "http://chatgpt.com/backend-api", ] { assert_eq!( - CodexProviderInstanceConfig::from_snapshot(&instance("openai", endpoint)), + CodexProviderInstanceConfig::from_snapshot( + &instance("openai", endpoint), + &OfficialCodexOriginPolicy, + ), Err(CodexProviderConfigError::UnsafeBaseUrl), "endpoint must fail closed: {endpoint}" ); @@ -48,20 +56,16 @@ fn official_policy_rejects_lookalike_private_and_metadata_origins() { fn loopback_policy_requires_explicit_numeric_loopback_host() { let numeric = instance("openai", "http://127.0.0.1:43123/backend-api"); assert_eq!( - CodexProviderInstanceConfig::from_snapshot(&numeric), + CodexProviderInstanceConfig::from_snapshot(&numeric, &OfficialCodexOriginPolicy), Err(CodexProviderConfigError::UnsafeBaseUrl) ); assert!( - CodexProviderInstanceConfig::from_snapshot_with_policy( - &numeric, - CodexEndpointPolicy::Loopback, - ) - .is_ok() + CodexProviderInstanceConfig::from_snapshot(&numeric, &LoopbackCodexOriginPolicy).is_ok() ); assert_eq!( - CodexProviderInstanceConfig::from_snapshot_with_policy( + CodexProviderInstanceConfig::from_snapshot( &instance("openai", "http://localhost:43123/backend-api"), - CodexEndpointPolicy::Loopback, + &LoopbackCodexOriginPolicy, ), Err(CodexProviderConfigError::UnsafeBaseUrl) ); @@ -70,10 +74,10 @@ fn loopback_policy_requires_explicit_numeric_loopback_host() { #[test] fn non_codex_instance_is_rejected_before_transport_initialization() { assert_eq!( - CodexProviderInstanceConfig::from_snapshot(&instance( - "xai", - "https://chatgpt.com/backend-api", - )), + CodexProviderInstanceConfig::from_snapshot( + &instance("xai", "https://chatgpt.com/backend-api"), + &OfficialCodexOriginPolicy, + ), Err(CodexProviderConfigError::ProviderMismatch) ); } diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index df838b440..1df095d38 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -6,34 +6,43 @@ use std::time::{Duration, SystemTime}; use chrono::Utc; use futures::{SinkExt, StreamExt}; use gateway_core::engine::continuation::{ - NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, +}; +use gateway_core::engine::credential::{ + AccountAvailability, ProviderAccountId, ProviderAccountStore, }; -use gateway_core::engine::credential::{AccountAvailability, ProviderAccountId}; use gateway_core::engine::provider::{Provider, ProviderRequest}; -use gateway_core::engine::{AttemptContext, CancellationToken, ModelRequestId}; +use gateway_core::engine::{ + AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId, + ProviderAccountStateOwner, UpstreamSendState, +}; use gateway_core::error::{ProviderError, ProviderErrorKind, SafeUpstreamValue}; +use gateway_core::event::{GatewayEvent, UpstreamHttpVersion}; use gateway_core::operation::{ ContentPart, ContinuationMode, GenerateRequest, Message, MessageRole, Operation, OperationKind, - ProviderOptions, + ProtocolPayload, ProviderOptions, }; use gateway_core::routing::{ ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId, }; +use provider_openai::CodexProvider; use provider_openai::credential::{ - CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialSelector, - CreateCodexCredential, + CodexCookie, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialQuotaService, + CodexCredentialSelector, CreateCodexCredential, RotateCodexCredential, }; +use provider_openai::transport::CodexWebSocketPool; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; -use provider_openai::{CodexEndpointPolicy, CodexProvider}; use serde_json::{Map, Value, json}; -use tokio::net::TcpListener; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; use tokio_tungstenite::tungstenite::Message as WebSocketMessage; -use wiremock::matchers::{method, path}; +use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use crate::support::{ - MemoryAccountStore, TestLeaseCoordinator, account_policy, instance_id, profile, secret, + MemoryAccountStore, TestLeaseCoordinator, account_policy, instance_id, loopback_origin_policy, + profile, secret, }; fn wire_profile() -> CodexWireProfileState { @@ -59,12 +68,51 @@ async fn create_account(store: &Arc) { name: "provider".to_owned(), secret: secret("provider-access"), account: profile("chatgpt-acct_provider"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), enabled: true, }) .await .expect("create account"); } +async fn create_other_account(store: &Arc) { + store + .repository() + .create_oauth_credential(CreateCodexCredential { + account_id: "acct_provider_other".to_owned(), + provider_instance_id: instance_id().to_string(), + name: "provider-other".to_owned(), + secret: secret("provider-other-access"), + account: profile("chatgpt-acct_provider_other"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await + .expect("create other account"); +} + +async fn attach_cookie(store: &Arc, account_id: &str, value: &str) { + let repository = store.repository(); + let account = store.account(account_id).expect("cookie account"); + let mut data = repository + .load_complete_data(&account) + .await + .expect("load credential data"); + data.cookies.push(CodexCookie { + name: "cf_clearance".to_owned(), + value: value.to_owned(), + domain: "127.0.0.1".to_owned(), + path: "/backend-api".to_owned(), + host_only: true, + secure: false, + expires_at: None, + }); + repository + .compare_and_swap_data(&account, data) + .await + .expect("persist test cookie"); +} + fn instance(base_url: &str) -> ProviderInstance { ProviderInstance::new( instance_id(), @@ -76,24 +124,118 @@ fn instance(base_url: &str) -> ProviderInstance { } fn provider(store: &Arc) -> CodexProvider { + provider_with_pool(store, Arc::new(CodexWebSocketPool::default())) +} + +fn provider_with_pool( + store: &Arc, + websocket_pool: Arc, +) -> CodexProvider { + provider_with_pool_and_cookie_policy( + store, + websocket_pool, + CodexCookiePolicy::official().expect("cookie policy"), + ) +} + +fn provider_with_pool_and_cookie_policy( + store: &Arc, + websocket_pool: Arc, + cookie_policy: CodexCookiePolicy, +) -> CodexProvider { let profile = wire_profile(); + let http = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"); + let origin = loopback_origin_policy(); + let catalog = Arc::new(CodexCredentialCatalogService::new( + store.repository(), + profile.clone(), + http.clone(), + Arc::clone(&origin), + )); + let quota = Arc::new(CodexCredentialQuotaService::new( + store.repository(), + profile.clone(), + http.clone(), + Arc::clone(&origin), + )); let selector = Arc::new(CodexCredentialSelector::new( store.repository(), Arc::new(TestLeaseCoordinator::default()), - CodexCookiePolicy::official().expect("cookie policy"), + Arc::clone(&catalog), + Arc::clone("a), + cookie_policy, )); - let catalog = Arc::new(CodexCredentialCatalogService::new_with_endpoint_policy( - store.repository(), - profile.clone(), - reqwest::Client::builder() - .no_proxy() - .build() - .expect("client"), - CodexEndpointPolicy::Loopback, - Duration::from_secs(60), - )); - CodexProvider::new_with_policy(selector, catalog, profile, CodexEndpointPolicy::Loopback) - .expect("Codex provider") + CodexProvider::new( + selector, + catalog, + quota, + http, + profile, + websocket_pool, + origin, + ) +} + +#[tokio::test] +async fn provider_uses_the_injected_account_lifecycle_pool() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("websocket connection"); + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let _ = websocket.next().await.expect("request").expect("frame"); + websocket + .send(WebSocketMessage::Text(created("shared-pool").into())) + .await + .expect("created event"); + websocket + .send(WebSocketMessage::Text(completed("shared-pool").into())) + .await + .expect("completed event"); + let close = tokio::time::timeout(Duration::from_secs(2), websocket.next()) + .await + .expect("shared lifecycle pool should close the idle socket") + .expect("close frame") + .expect("valid close frame"); + std::assert_matches!(close, WebSocketMessage::Close(_)); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let pool = Arc::new(CodexWebSocketPool::default()); + let provider = provider_with_pool(&store, Arc::clone(&pool)); + let operation = operation("shared-lifecycle-pool"); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context("req_shared_lifecycle_pool", None, CancellationToken::new()), + ) + .await + .expect("provider stream"); + let first = stream + .next() + .await + .expect("transport observation") + .expect("transport observation"); + let observation = first.response_observation().expect("response facts"); + assert_eq!(observation.transport().as_str(), "websocket"); + assert_eq!( + observation.http_version(), + Some(UpstreamHttpVersion::Http11) + ); + assert_eq!(observation.status_code(), Some(101)); + assert!(observation.timings().connect_ms.is_some()); + while let Some(event) = stream.next().await { + event.expect("websocket event"); + } + + pool.evict_account("acct_provider").await; + server.await.expect("server task"); } fn operation(conversation_id: &str) -> Operation { @@ -126,7 +268,193 @@ fn operation(conversation_id: &str) -> Operation { ) } +fn pool_identity_operation( + prompt_cache_key: &str, + session_id: &str, + thread_id: &str, + conversation_id: &str, +) -> Operation { + let payload = ProtocolPayload::json_object( + "openai", + json!({ + "model": "gpt-5.4", + "input": [{"role": "user", "content": "pool identity"}], + "prompt_cache_key": prompt_cache_key, + "session_id": session_id, + "thread_id": thread_id, + "conversation_id": conversation_id + }) + .as_object() + .cloned() + .expect("pool identity request"), + ) + .expect("OpenAI payload"); + let mut options = ProviderOptions::new(); + options + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), Value::from(1)), + ( + "transport".to_owned(), + Value::String("websocket".to_owned()), + ), + ]), + ) + .expect("provider options"); + Operation::Generate( + GenerateRequest::from_protocol_payload(Vec::new(), payload).with_provider_options(options), + ) +} + +#[tokio::test] +async fn account_scoped_pool_identity_is_stable_isolated_and_uses_documented_priority() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let mut handlers = tokio::task::JoinSet::new(); + let mut next_connection_id = 0_u64; + loop { + tokio::select! { + _ = &mut shutdown_rx => break, + accepted = listener.accept() => { + let (stream, _) = accepted.expect("websocket connection"); + next_connection_id += 1; + let connection_id = next_connection_id; + let observed_tx = observed_tx.clone(); + handlers.spawn(async move { + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let mut response_sequence = 0_u64; + while let Some(frame) = websocket.next().await { + let frame = frame.expect("valid websocket frame"); + let WebSocketMessage::Text(payload) = frame else { + continue; + }; + response_sequence += 1; + let payload: Value = serde_json::from_str(&payload) + .expect("response.create payload"); + observed_tx + .send((connection_id, payload)) + .expect("pool observation receiver"); + let response_id = + format!("resp_pool_{connection_id}_{response_sequence}"); + websocket + .send(WebSocketMessage::Text(created(&response_id).into())) + .await + .expect("created event"); + websocket + .send(WebSocketMessage::Text(completed(&response_id).into())) + .await + .expect("completed event"); + } + }); + } + joined = handlers.join_next(), if !handlers.is_empty() => { + joined.expect("pool connection task").expect("pool connection handler"); + } + } + } + handlers.abort_all(); + while handlers.join_next().await.is_some() {} + }); + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let pool = Arc::new(CodexWebSocketPool::default()); + let provider = provider_with_pool(&store, Arc::clone(&pool)); + let primary = ProviderAccountId::new("acct_provider").expect("primary account"); + let secondary = ProviderAccountId::new("acct_provider_other").expect("secondary account"); + let requests = [ + ( + pool_identity_operation("cache-a", "session-a", "thread-a", "conversation-a"), + BTreeSet::from([secondary.clone()]), + ), + ( + pool_identity_operation("cache-a", "session-a", "thread-a", "conversation-a"), + BTreeSet::from([secondary.clone()]), + ), + ( + pool_identity_operation("cache-a", "session-b", "thread-b", "conversation-b"), + BTreeSet::from([secondary.clone()]), + ), + ( + pool_identity_operation("cache-b", "session-a", "thread-a", "conversation-a"), + BTreeSet::from([secondary.clone()]), + ), + ( + pool_identity_operation("cache-a", "session-a", "thread-a", "conversation-a"), + BTreeSet::from([primary]), + ), + ]; + let mut observations = Vec::new(); + for (index, (operation, excluded)) in requests.into_iter().enumerate() { + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context_with_exclusions( + &format!("req_pool_identity_{index}"), + NonZeroU32::new(1).expect("first attempt"), + excluded, + None, + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + while let Some(event) = stream.next().await { + event.expect("provider event"); + } + observations.push( + tokio::time::timeout(Duration::from_secs(2), observed_rx.recv()) + .await + .expect("pool observation timeout") + .expect("pool observation"), + ); + } + + assert_eq!(observations[0].0, observations[1].0); + assert_eq!(observations[0].0, observations[2].0); + assert_ne!(observations[0].0, observations[3].0); + assert_ne!(observations[0].0, observations[4].0); + assert_eq!(observations[0].1["prompt_cache_key"], "cache-a"); + assert_eq!(observations[2].1["session_id"], "session-b"); + + pool.shutdown().await; + let _ = shutdown_tx.send(()); + server.await.expect("pool identity server"); +} + fn http_operation() -> Operation { + let message = Message::new( + MessageRole::User, + vec![ContentPart::Text("hello".to_owned())], + ) + .expect("message"); + let mut options = ProviderOptions::new(); + options + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), Value::from(1)), + ("transport".to_owned(), Value::String("http_sse".to_owned())), + ]), + ) + .expect("provider options"); + Operation::Generate( + GenerateRequest::new(vec![message]) + .expect("generate request") + .with_provider_options(options), + ) +} + +fn preferred_operation() -> Operation { let message = Message::new( MessageRole::User, vec![ContentPart::Text("hello".to_owned())], @@ -135,28 +463,141 @@ fn http_operation() -> Operation { Operation::Generate(GenerateRequest::new(vec![message]).expect("generate request")) } -async fn rejected_http_operation(status: u16) -> (ProviderError, [AccountAvailability; 2]) { +fn identity_tainted_operation() -> Operation { + let body = json!({ + "model": "gpt-5.4", + "input": [{"role": "user", "content": "continue"}], + "accountId": "attacker-account", + "access_token": "attacker-token", + "installationId": "attacker-installation", + "conversation_id": "different-client-conversation-after", + "previous_response_id": "attacker-response", + "turnState": "attacker-turn-state", + "turnMetadata": "{\"accountId\":\"attacker-account\",\"installationId\":\"attacker-installation\",\"request_kind\":\"review\"}", + "client_metadata": { + "accountId": "attacker-account", + "refreshToken": "attacker-refresh", + "installationId": "attacker-installation", + "turnMetadata": "{\"accountId\":\"attacker-account\",\"installationId\":\"attacker-installation\",\"subagent_kind\":\"review\"}" + }, + "future_official_field": {"preserved": true} + }); + let payload = + ProtocolPayload::json_object("openai", body.as_object().cloned().expect("request body")) + .expect("OpenAI payload"); + let mut options = ProviderOptions::new(); + options + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), Value::from(1)), + ( + "transport".to_owned(), + Value::String("websocket".to_owned()), + ), + ]), + ) + .expect("provider options"); + Operation::Generate( + GenerateRequest::from_protocol_payload(Vec::new(), payload).with_provider_options(options), + ) +} + +fn failover_identity_operation() -> Operation { + let body = json!({ + "model": "gpt-5.4", + "input": [ + { + "type": "message", + "role": "user", + "id": "user_input_id", + "encrypted_content": "user_account_state", + "content": [{ + "type": "input_text", + "id": "nested_user_content_id", + "encrypted_content": "nested_user_semantics", + "text": "same client turn" + }] + }, + { + "type": "message", + "role": "assistant", + "id": "source_message_id", + "encrypted_content": "source_message_encrypted", + "content": [{ + "type": "output_text", + "id": "nested_output_id", + "encrypted_content": "nested_output_semantics", + "text": "source answer" + }] + }, + { + "type": "reasoning", + "id": "source_reasoning_id", + "encrypted_content": "source_reasoning_encrypted", + "summary": [{"type": "summary_text", "id": "nested_summary_id", "text": "summary"}] + }, + { + "type": "function_call", + "id": "source_function_id", + "call_id": "call_source", + "name": "lookup", + "arguments": "{\"query\":\"preserved\"}" + }, + {"type": "compaction", "id": "source_compaction_id", "encrypted_content": "source_compaction"}, + {"type": "compaction_summary", "id": "source_compaction_summary_id"}, + {"type": "context_compaction", "id": "source_context_compaction_id"} + ], + "accountId": "attacker-account", + "access_token": "attacker-token", + "cookie": "attacker-cookie", + "installationId": "attacker-installation", + "conversation_id": "account-bound-client-conversation", + "session_id": "client-session-stable", + "turnState": "client-turn-state", + "turnMetadata": "{\"installationId\":\"attacker-installation\",\"accountId\":\"attacker-account\",\"opaque\":\"preserved\"}", + "client_metadata": { + "accountId": "attacker-account", + "cookie": "attacker-cookie", + "installationId": "attacker-installation", + "session_id": "client-session-stable", + "x-codex-turn-state": "client-metadata-turn-state", + "x-codex-turn-metadata": "{\"installationId\":\"attacker-installation\",\"accountId\":\"attacker-account\",\"opaque\":\"metadata-preserved\"}" + } + }); + let payload = ProtocolPayload::json_object( + "openai", + body.as_object().cloned().expect("failover request body"), + ) + .expect("OpenAI payload"); + let mut options = ProviderOptions::new(); + options + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), Value::from(1)), + ("transport".to_owned(), Value::String("http_sse".to_owned())), + ]), + ) + .expect("provider options"); + Operation::Generate( + GenerateRequest::from_protocol_payload(Vec::new(), payload).with_provider_options(options), + ) +} + +async fn rejected_http_operation( + response: ResponseTemplate, +) -> (ProviderError, [AccountAvailability; 2]) { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/backend-api/codex/responses")) - .respond_with(ResponseTemplate::new(status)) + .respond_with(response) .expect(1) .mount(&server) .await; let store = Arc::new(MemoryAccountStore::default()); create_account(&store).await; - store - .repository() - .create_oauth_credential(CreateCodexCredential { - account_id: "acct_provider_other".to_owned(), - provider_instance_id: instance_id().to_string(), - name: "provider-other".to_owned(), - secret: secret("provider-other-access"), - account: profile("chatgpt-acct_provider_other"), - enabled: true, - }) - .await - .expect("create other account"); + create_other_account(&store).await; let provider = provider(&store); let operation = http_operation(); let mut stream = provider @@ -166,11 +607,125 @@ async fn rejected_http_operation(status: u16) -> (ProviderError, [AccountAvailab ) .await .expect("provider stream"); - let error = stream - .next() + let error = next_provider_error(&mut stream).await; + let availability = [ + store + .account("acct_provider") + .expect("provider account") + .availability(), + store + .account("acct_provider_other") + .expect("other provider account") + .availability(), + ]; + (error, availability) +} + +async fn in_band_failure( + code: &str, + status: Option, + message: &str, +) -> (ProviderError, [AccountAvailability; 2]) { + let server = MockServer::start().await; + let mut event = json!({ + "type": "response.failed", + "response": { + "id": "resp_in_band_failure", + "status": "failed", + "error": {"code": code, "message": message} + } + }); + if let Some(status) = status { + event + .as_object_mut() + .expect("failure event") + .insert("status_code".to_owned(), Value::from(status)); + } + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(format!("event: response.failed\ndata: {event}\n\n")), + ) + .expect(1) + .mount(&server) + .await; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let provider = provider(&store); + let secondary = ProviderAccountId::new("acct_provider_other").expect("secondary account"); + let operation = http_operation(); + let mut stream = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context_with_exclusions( + "req_in_band_failure", + NonZeroU32::new(1).expect("first attempt"), + BTreeSet::from([secondary]), + None, + None, + CancellationToken::new(), + ), + ) .await - .expect("error event") - .expect_err("HTTP rejection"); + .expect("provider stream"); + let error = next_provider_error(&mut stream).await; + let availability = [ + store + .account("acct_provider") + .expect("provider account") + .availability(), + store + .account("acct_provider_other") + .expect("other provider account") + .availability(), + ]; + (error, availability) +} + +async fn rejected_websocket_opening( + status: u16, + body: &str, +) -> (ProviderError, [AccountAvailability; 2]) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let body = body.to_owned(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.expect("websocket opening"); + let _ = read_http_headers(&mut stream).await; + let response = format!( + "HTTP/1.1 {status} Rejected\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream + .write_all(response.as_bytes()) + .await + .expect("opening rejection"); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let provider = provider(&store); + let operation = operation("opening-rejection"); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context( + "req_websocket_opening_rejection", + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + let error = next_provider_error(&mut stream).await; + server.await.expect("server task"); let availability = [ store .account("acct_provider") @@ -210,17 +765,54 @@ fn planned_request(base_url: &str, operation: &Operation) -> ProviderRequest { ProviderRequest::new(operation.clone(), plan.candidates()[0].clone()) } +async fn next_provider_error( + stream: &mut gateway_core::engine::provider::ProviderStream, +) -> ProviderError { + loop { + match stream.next().await.expect("provider failure event") { + Ok(event) => assert!(!event.has_client_event()), + Err(error) => return error, + } + } +} + fn context( request_id: &str, - continuation: Option, + continuation: Option, + cancellation: CancellationToken, +) -> AttemptContext { + let account_state_owner = continuation + .as_ref() + .and_then(ContinuationBinding::pinned) + .map(ProviderAccountStateOwner::from_continuation); + context_with_exclusions( + request_id, + NonZeroU32::new(1).expect("attempt"), + BTreeSet::new(), + account_state_owner, + continuation, + cancellation, + ) +} + +fn context_with_exclusions( + request_id: &str, + attempt_index: NonZeroU32, + excluded_accounts: BTreeSet, + account_state_owner: Option, + continuation: Option, cancellation: CancellationToken, ) -> AttemptContext { AttemptContext::new( ModelRequestId::new(request_id).expect("request id"), - NonZeroU32::new(1).expect("attempt"), + attempt_index, SystemTime::now() + Duration::from_secs(30), account_policy(), - gateway_core::engine::AccountSelectionConstraints::default(), + gateway_core::engine::AccountAttemptContext::new( + excluded_accounts, + None, + account_state_owner, + ), continuation, cancellation, ) @@ -248,6 +840,745 @@ fn created(response_id: &str) -> String { .to_string() } +async fn serve_http_after_dropped_websocket(listener: TcpListener) { + let (websocket, _) = listener.accept().await.expect("websocket connection"); + drop(websocket); + + let (mut http, _) = listener.accept().await.expect("HTTP connection"); + let request = read_http_headers(&mut http).await; + assert!(request.starts_with(b"POST /backend-api/codex/responses HTTP/1.1\r\n")); + + let body = format!( + "data: {}\n\ndata: {}\n\n", + created("native-http-fallback"), + completed("native-http-fallback") + ); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nx-codex-active-limit: codex\r\nx-codex-primary-used-percent: 25\r\nx-codex-primary-window-minutes: 300\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + http.write_all(response.as_bytes()) + .await + .expect("HTTP response"); +} + +async fn read_http_headers(stream: &mut TcpStream) -> Vec { + let mut request = Vec::new(); + let mut buffer = [0_u8; 4_096]; + loop { + let read = stream.read(&mut buffer).await.expect("HTTP request"); + assert!(read > 0, "HTTP request must contain headers"); + request.extend_from_slice(&buffer[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + return request; + } + assert!( + request.len() <= 32 * 1_024, + "HTTP headers are unexpectedly large" + ); + } +} + +#[tokio::test] +async fn default_transport_prefers_websocket_and_accepts_same_account_http_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(serve_http_after_dropped_websocket(listener)); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let operation = preferred_operation(); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context( + "req_provider_preferred_fallback", + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + let first = stream + .next() + .await + .expect("transport observation") + .expect("transport observation"); + let observation = first.response_observation().expect("response facts"); + assert_eq!(observation.transport().as_str(), "http_sse"); + assert_eq!( + observation.http_version(), + Some(UpstreamHttpVersion::Http11) + ); + assert_eq!(observation.status_code(), Some(200)); + assert!(observation.timings().transport_decision_wait_ms.is_some()); + while let Some(event) = stream.next().await { + event.expect("canonical HTTP fallback event"); + } + let account = store.account("acct_provider").expect("account"); + let quota = store + .get_quotas(std::slice::from_ref(account.id())) + .await + .expect("passive quota") + .pop() + .and_then(|observation| observation.quota) + .expect("quota JSON"); + assert_eq!( + quota + .expose_to_provider() + .get("rate_limit") + .and_then(Value::as_object) + .and_then(|limit| limit.get("primary_window")) + .and_then(Value::as_object) + .and_then(|window| window.get("used_percent")) + .and_then(Value::as_f64), + Some(25.0) + ); + server.await.expect("server task"); +} + +#[tokio::test] +async fn external_previous_response_id_is_forwarded_without_a_local_account_pin() { + let server = MockServer::start().await; + let body = format!( + "event: response.created\ndata: {}\n\nevent: response.completed\ndata: {}\n\n", + created("resp_external_result"), + completed("resp_external_result") + ); + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .expect(1) + .mount(&server) + .await; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let provider = provider(&store); + let operation = http_operation(); + let external = PreviousResponseId::new("external-provider-response").expect("external ID"); + let mut stream = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context( + "req_external_continuation", + Some(ContinuationBinding::External(external)), + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + while let Some(event) = stream.next().await { + event.expect("external continuation response"); + } + let requests = server + .received_requests() + .await + .expect("external continuation request"); + let body: Value = serde_json::from_slice(&requests[0].body).expect("request JSON"); + assert_eq!(body["previous_response_id"], "external-provider-response"); +} + +#[tokio::test] +async fn account_retry_rebinds_every_provider_identity_and_preserves_client_session_semantics() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .and(header("authorization", "Bearer provider-access")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": {"code": "rate_limit_exceeded", "message": "retry another account"} + }))) + .expect(1) + .mount(&server) + .await; + let completed_body = format!( + "event: response.created\ndata: {}\n\nevent: response.completed\ndata: {}\n\n", + created("resp_second_account"), + completed("resp_second_account") + ); + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .and(header("authorization", "Bearer provider-other-access")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(completed_body), + ) + .expect(1) + .mount(&server) + .await; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + attach_cookie(&store, "acct_provider", "cookie-primary").await; + attach_cookie(&store, "acct_provider_other", "cookie-secondary").await; + let provider = provider_with_pool_and_cookie_policy( + &store, + Arc::new(CodexWebSocketPool::default()), + CodexCookiePolicy::new(["cf_clearance"], ["127.0.0.1"]).expect("loopback cookie policy"), + ); + let operation = failover_identity_operation(); + let primary = ProviderAccountId::new("acct_provider").expect("primary account"); + let secondary = ProviderAccountId::new("acct_provider_other").expect("secondary account"); + + let mut first = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context_with_exclusions( + "req_identity_retry_first", + NonZeroU32::new(1).expect("first attempt"), + BTreeSet::from([secondary.clone()]), + None, + None, + CancellationToken::new(), + ), + ) + .await + .expect("first provider stream"); + assert_eq!(first.metadata().provider_account_id(), Some(&primary)); + let first_error = next_provider_error(&mut first).await; + assert_eq!(first_error.kind(), ProviderErrorKind::RateLimited); + assert!(first_error.replay_is_safe()); + + let mut second = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context_with_exclusions( + "req_identity_retry_second", + NonZeroU32::new(2).expect("second attempt"), + BTreeSet::from([primary.clone()]), + Some(ProviderAccountStateOwner::new( + ProviderKind::new("openai").expect("provider"), + instance_id(), + primary.clone(), + )), + None, + CancellationToken::new(), + ), + ) + .await + .expect("second provider stream"); + assert_eq!(second.metadata().provider_account_id(), Some(&secondary)); + while let Some(event) = second.next().await { + event.expect("second account response"); + } + + let requests = server + .received_requests() + .await + .expect("received failover requests"); + assert_eq!(requests.len(), 2); + for (index, (request, account_id, upstream_account_id, access_token, cookie)) in [ + ( + &requests[0], + "acct_provider", + "chatgpt-acct_provider", + "provider-access", + "cookie-primary", + ), + ( + &requests[1], + "acct_provider_other", + "chatgpt-acct_provider_other", + "provider-other-access", + "cookie-secondary", + ), + ] + .into_iter() + .enumerate() + { + let account = store.account(account_id).expect("selected account"); + let credential = store + .repository() + .load_complete_data(&account) + .await + .expect("selected credential data"); + let header_value = |name: &str| { + request + .headers + .get(name) + .and_then(|value| value.to_str().ok()) + }; + assert_eq!( + header_value("authorization"), + Some(format!("Bearer {access_token}").as_str()) + ); + assert_eq!( + header_value("chatgpt-account-id"), + Some(upstream_account_id) + ); + assert_eq!( + header_value("x-codex-installation-id"), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + header_value("cookie"), + Some(format!("cf_clearance={cookie}").as_str()) + ); + assert_eq!(header_value("session-id"), Some("client-session-stable")); + + let body: Value = serde_json::from_slice(&request.body).expect("upstream request JSON"); + let input = body["input"].as_array().expect("input array"); + assert_eq!( + input[0].pointer("/content/0/text").and_then(Value::as_str), + Some("same client turn") + ); + assert_eq!( + input[0] + .pointer("/content/0/encrypted_content") + .and_then(Value::as_str), + Some("nested_user_semantics") + ); + if index == 0 { + assert_eq!(input.len(), 7); + assert_eq!(input[0].get("id"), Some(&json!("user_input_id"))); + assert_eq!( + input[1].get("encrypted_content"), + Some(&json!("source_message_encrypted")) + ); + assert!(input.iter().any(|item| item["type"] == "compaction")); + } else { + assert_eq!(input.len(), 4); + assert!(input.iter().all(|item| item.get("id").is_none())); + assert!( + input + .iter() + .all(|item| item.get("encrypted_content").is_none()) + ); + assert!(input.iter().all(|item| { + !matches!( + item.get("type").and_then(Value::as_str), + Some("compaction" | "compaction_summary" | "context_compaction") + ) + })); + assert_eq!( + input[1].pointer("/content/0/id").and_then(Value::as_str), + Some("nested_output_id") + ); + assert_eq!( + input[1] + .pointer("/content/0/encrypted_content") + .and_then(Value::as_str), + Some("nested_output_semantics") + ); + assert_eq!( + input[2].pointer("/summary/0/id").and_then(Value::as_str), + Some("nested_summary_id") + ); + assert_eq!( + input[3].get("call_id").and_then(Value::as_str), + Some("call_source") + ); + assert_eq!( + input[3].get("arguments").and_then(Value::as_str), + Some("{\"query\":\"preserved\"}") + ); + } + assert!(body.get("accountId").is_none()); + assert!(body.get("access_token").is_none()); + assert!(body.get("cookie").is_none()); + if index == 0 { + assert_eq!( + body.get("conversation_id").and_then(Value::as_str), + Some("account-bound-client-conversation") + ); + assert_eq!( + header_value("x-codex-turn-state"), + Some("client-turn-state") + ); + let turn_metadata: Value = serde_json::from_str( + header_value("x-codex-turn-metadata").expect("first attempt turn metadata"), + ) + .expect("valid turn metadata"); + assert!(turn_metadata.get("accountId").is_none()); + assert_eq!( + turn_metadata.get("installationId").and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + turn_metadata.get("opaque").and_then(Value::as_str), + Some("preserved") + ); + } else { + assert!(body.get("conversation_id").is_none()); + assert!(header_value("x-codex-turn-state").is_none()); + let turn_metadata: Value = serde_json::from_str( + header_value("x-codex-turn-metadata").expect("cross-account metadata"), + ) + .expect("valid cross-account turn metadata"); + assert!(turn_metadata.get("accountId").is_none()); + assert_eq!( + turn_metadata.get("installationId").and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + turn_metadata.get("opaque").and_then(Value::as_str), + Some("preserved") + ); + } + assert_eq!( + body.get("installationId").and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + body.get("session_id").and_then(Value::as_str), + Some("client-session-stable") + ); + assert_eq!( + body.pointer("/client_metadata/installationId") + .and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + body.pointer("/client_metadata/x-codex-installation-id") + .and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + body.pointer("/client_metadata/session_id") + .and_then(Value::as_str), + Some("client-session-stable") + ); + assert!(body.pointer("/client_metadata/accountId").is_none()); + assert!(body.pointer("/client_metadata/cookie").is_none()); + if index == 0 { + assert_eq!( + body.pointer("/client_metadata/x-codex-turn-state") + .and_then(Value::as_str), + Some("client-metadata-turn-state") + ); + } else { + assert!( + body.pointer("/client_metadata/x-codex-turn-state") + .is_none() + ); + let metadata_turn: Value = serde_json::from_str( + body.pointer("/client_metadata/x-codex-turn-metadata") + .and_then(Value::as_str) + .expect("cross-account client metadata turn metadata"), + ) + .expect("valid cross-account client metadata turn metadata"); + assert!(metadata_turn.get("accountId").is_none()); + assert_eq!( + metadata_turn.get("installationId").and_then(Value::as_str), + Some(credential.installation_id.as_str()) + ); + assert_eq!( + metadata_turn.get("opaque").and_then(Value::as_str), + Some("metadata-preserved") + ); + } + } + assert_ne!( + requests[0] + .headers + .get("x-codex-installation-id") + .expect("primary installation"), + requests[1] + .headers + .get("x-codex-installation-id") + .expect("secondary installation") + ); +} + +#[tokio::test] +async fn explicit_websocket_connect_failure_is_not_sent_and_does_not_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (websocket, _) = listener.accept().await.expect("websocket connection"); + drop(websocket); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let operation = operation("required-websocket"); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context( + "req_provider_required_websocket", + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + let error = next_provider_error(&mut stream).await; + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + server.await.expect("server task"); +} + +#[tokio::test] +async fn explicit_websocket_post_send_disconnect_remains_ambiguous() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("websocket connection"); + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let _ = websocket.next().await.expect("request").expect("frame"); + drop(websocket); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let operation = operation("ambiguous-websocket"); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context( + "req_provider_ambiguous_websocket", + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + let error = next_provider_error(&mut stream).await; + assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); + server.await.expect("server task"); +} + +#[tokio::test] +async fn account_invalidation_closes_its_idle_websocket_connections() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("websocket connection"); + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let _ = websocket.next().await.expect("request").expect("frame"); + websocket + .send(WebSocketMessage::Text(created("evict-me").into())) + .await + .expect("created event"); + websocket + .send(WebSocketMessage::Text(completed("evict-me").into())) + .await + .expect("completed event"); + + let (mut http, _) = listener.accept().await.expect("HTTP connection"); + let _ = read_http_headers(&mut http).await; + let body = r#"{"error":{"code":"token_expired","message":"access token expired"}}"#; + let response = format!( + "HTTP/1.1 401 Unauthorized\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + http.write_all(response.as_bytes()) + .await + .expect("HTTP rejection"); + + loop { + match tokio::time::timeout(Duration::from_secs(2), websocket.next()).await { + Ok(Some(Ok(WebSocketMessage::Close(_)))) | Ok(None) => break, + Ok(Some(Ok(_))) => {} + Ok(Some(Err(error))) => panic!("websocket close failed: {error}"), + Err(_) => panic!("account invalidation did not close pooled websocket"), + } + } + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + + let websocket_operation = operation("eviction-lifecycle"); + let mut first = provider + .execute( + planned_request(&base_url, &websocket_operation), + context( + "req_pool_before_invalidation", + None, + CancellationToken::new(), + ), + ) + .await + .expect("websocket provider stream"); + while let Some(event) = first.next().await { + event.expect("websocket event"); + } + + let http_operation = http_operation(); + let mut second = provider + .execute( + planned_request(&base_url, &http_operation), + context("req_pool_invalidation", None, CancellationToken::new()), + ) + .await + .expect("HTTP provider stream"); + let error = next_provider_error(&mut second).await; + + assert_eq!(error.kind(), ProviderErrorKind::Unauthorized); + assert_eq!( + store + .account("acct_provider") + .expect("provider account") + .availability(), + AccountAvailability::Expired + ); + server.await.expect("server task"); +} + +#[tokio::test] +async fn credential_revision_change_does_not_split_established_logical_websocket() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.expect("first websocket"); + let mut first = crate::transport::accept_codex_test_websocket(first_stream).await; + let _ = first.next().await.expect("first request").expect("frame"); + first + .send(WebSocketMessage::Text(created("before-rotation").into())) + .await + .expect("first created"); + first + .send(WebSocketMessage::Text(completed("before-rotation").into())) + .await + .expect("first completed"); + + let _ = first.next().await.expect("second request").expect("frame"); + first + .send(WebSocketMessage::Text(created("after-rotation").into())) + .await + .expect("second created"); + first + .send(WebSocketMessage::Text(completed("after-rotation").into())) + .await + .expect("second completed"); + assert!( + tokio::time::timeout(Duration::from_millis(300), listener.accept()) + .await + .is_err(), + "credential revision must not split an established logical connection" + ); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let operation = operation("credential-rotation"); + + let mut first = provider + .execute( + planned_request(&base_url, &operation), + context("req_before_rotation", None, CancellationToken::new()), + ) + .await + .expect("first provider stream"); + while let Some(event) = first.next().await { + event.expect("first websocket event"); + } + + store + .repository() + .rotate_oauth_secret(RotateCodexCredential { + account_id: "acct_provider".to_owned(), + expected_credential_revision: 1, + secret: secret("rotated-provider-access"), + verified_account: profile("chatgpt-acct_provider"), + next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)), + }) + .await + .expect("rotate credential"); + + let mut second = provider + .execute( + planned_request(&base_url, &operation), + context("req_after_rotation", None, CancellationToken::new()), + ) + .await + .expect("second provider stream"); + while let Some(event) = second.next().await { + event.expect("second websocket event"); + } + server.await.expect("server task"); +} + +#[tokio::test] +async fn provider_instance_reuses_the_same_websocket_origin_breaker() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let mut accepted = 0; + loop { + tokio::select! { + _ = &mut stop_rx => break, + accepted_connection = listener.accept() => { + let (connection, _) = accepted_connection.expect("websocket connection"); + accepted += 1; + drop(connection); + } + } + } + accepted + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + + for attempt in 0..4 { + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("gpt-public")), + ( + "input".to_owned(), + json!([{"role": "user", "content": "warm connection"}]), + ), + ("stream".to_owned(), json!(true)), + ("store".to_owned(), json!(false)), + ("generate".to_owned(), json!(false)), + ("conversation_id".to_owned(), json!("breaker-origin")), + ]), + ) + .expect("warmup payload"); + let operation = + Operation::Generate(GenerateRequest::from_protocol_payload(Vec::new(), payload)); + let mut stream = provider + .execute( + planned_request(&base_url, &operation), + context( + &format!("req_breaker_{attempt}"), + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + let error = next_provider_error(&mut stream).await; + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } + + stop_tx.send(()).expect("stop server"); + assert_eq!(server.await.expect("server task"), 3); +} + #[tokio::test] async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); @@ -264,6 +1595,44 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { }; let first: Value = serde_json::from_str(&first).expect("first request JSON"); assert!(first.get("previous_response_id").is_none()); + assert!(first.get("accountId").is_none()); + assert!(first.get("access_token").is_none()); + assert_eq!( + first.get("conversation_id").and_then(Value::as_str), + Some("different-client-conversation-after") + ); + assert_eq!( + first.get("turnState").and_then(Value::as_str), + Some("attacker-turn-state") + ); + assert_ne!(first["installationId"], "attacker-installation"); + let first_turn_metadata: Value = + serde_json::from_str(first["turnMetadata"].as_str().expect("first turn metadata")) + .expect("first turn metadata JSON"); + assert!(first_turn_metadata.get("accountId").is_none()); + assert_ne!( + first_turn_metadata["installationId"], + "attacker-installation" + ); + assert_eq!(first_turn_metadata["request_kind"], "review"); + let first_metadata = first["client_metadata"] + .as_object() + .expect("first client metadata"); + assert!(first_metadata.get("accountId").is_none()); + assert!(first_metadata.get("refreshToken").is_none()); + let first_client_turn_metadata: Value = serde_json::from_str( + first_metadata + .get("turnMetadata") + .and_then(Value::as_str) + .expect("first client turn metadata"), + ) + .expect("first client turn metadata JSON"); + assert!(first_client_turn_metadata.get("accountId").is_none()); + assert_ne!( + first_client_turn_metadata["installationId"], + "attacker-installation" + ); + assert_eq!(first_client_turn_metadata["subagent_kind"], "review"); websocket .send(WebSocketMessage::Text( created("native-upstream-handle").into(), @@ -288,6 +1657,31 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { let second: Value = serde_json::from_str(&second).expect("second request JSON"); assert_eq!(second["previous_response_id"], "native-upstream-handle"); assert_ne!(second["previous_response_id"], "resp_gateway_client_id"); + assert!(second.get("accountId").is_none()); + assert!(second.get("access_token").is_none()); + assert_ne!(second["installationId"], "attacker-installation"); + assert_eq!( + second.pointer("/future_official_field/preserved"), + Some(&json!(true)) + ); + let turn_metadata: Value = serde_json::from_str( + second["turnMetadata"] + .as_str() + .expect("sanitized turn metadata"), + ) + .expect("turn metadata JSON"); + assert!(turn_metadata.get("accountId").is_none()); + assert_ne!(turn_metadata["installationId"], "attacker-installation"); + assert_eq!(turn_metadata["request_kind"], "review"); + let client_metadata = second["client_metadata"] + .as_object() + .expect("client metadata"); + assert!(client_metadata.get("accountId").is_none()); + assert!(client_metadata.get("refreshToken").is_none()); + assert_ne!( + client_metadata.get("installationId"), + Some(&json!("attacker-installation")) + ); websocket .send(WebSocketMessage::Text(created("native-second").into())) .await @@ -301,7 +1695,7 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store).await; let provider = provider(&store); - let first_operation = operation("client-conversation-before"); + let first_operation = identity_tainted_operation(); let mut first = provider .execute( planned_request(&base_url, &first_operation), @@ -322,11 +1716,15 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { ProviderAccountId::new("acct_provider").expect("account id"), NativeContinuationReuse::Reusable, ); - let second_operation = operation("different-client-conversation-after"); + let second_operation = identity_tainted_operation(); let mut second = provider .execute( planned_request(&base_url, &second_operation), - context("req_provider_second", Some(pin), CancellationToken::new()), + context( + "req_provider_second", + Some(ContinuationBinding::Pinned(pin)), + CancellationToken::new(), + ), ) .await .expect("second provider stream"); @@ -336,6 +1734,194 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { server.await.expect("server task"); } +#[tokio::test] +async fn connection_replay_restores_full_input_without_reusing_native_handle() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept"); + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let first = websocket.next().await.expect("first frame").expect("frame"); + let WebSocketMessage::Text(first) = first else { + panic!("first request must be text"); + }; + let first: Value = serde_json::from_str(&first).expect("first request JSON"); + let first_input = first["input"].as_array().expect("first input").to_vec(); + websocket + .send(WebSocketMessage::Text( + created("native-replay-source").into(), + )) + .await + .expect("first response start"); + websocket + .send(WebSocketMessage::Text( + completed("native-replay-source").into(), + )) + .await + .expect("first response"); + + let second = websocket + .next() + .await + .expect("second frame") + .expect("frame"); + let WebSocketMessage::Text(second) = second else { + panic!("second request must be text"); + }; + let second: Value = serde_json::from_str(&second).expect("second request JSON"); + assert!(second.get("previous_response_id").is_none()); + let second_input = second["input"].as_array().expect("second input"); + assert!(second_input.starts_with(&first_input)); + assert!(second_input.len() > first_input.len()); + websocket + .send(WebSocketMessage::Text(created("native-replayed").into())) + .await + .expect("second response start"); + websocket + .send(WebSocketMessage::Text(completed("native-replayed").into())) + .await + .expect("second response"); + }); + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let first_operation = identity_tainted_operation(); + let mut first = provider + .execute( + planned_request(&base_url, &first_operation), + context("req_replay_first", None, CancellationToken::new()), + ) + .await + .expect("first provider stream"); + let mut provider_state = None; + while let Some(event) = first.next().await { + let mut event = event.expect("first canonical event"); + provider_state = event.take_session_update().or(provider_state); + } + drop(first); + + let pin = NativeContinuationPin::new( + PreviousResponseId::new("resp_gateway_replay").expect("gateway response id"), + SafeUpstreamValue::new("native-replay-source").expect("native handle"), + ProviderKind::new("openai").expect("provider"), + instance_id(), + ProviderAccountId::new("acct_provider").expect("account id"), + NativeContinuationReuse::Reusable, + ); + let second_operation = identity_tainted_operation() + .with_provider_session_state(provider_state.expect("connection replay state")); + let replay_context = context( + "req_replay_second", + Some(ContinuationBinding::Pinned(pin)), + CancellationToken::new(), + ) + .with_continuation_attempt(ContinuationAttempt::ReplayOwner); + let mut second = provider + .execute( + planned_request(&base_url, &second_operation), + replay_context, + ) + .await + .expect("second provider stream"); + while let Some(event) = second.next().await { + event.expect("second canonical event"); + } + server.await.expect("server task"); +} + +#[tokio::test] +async fn connection_replay_preserves_openai_observability_semantics() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!( + "http://{}/backend-api", + listener.local_addr().expect("address") + ); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept"); + let mut websocket = crate::transport::accept_codex_test_websocket(stream).await; + let _ = websocket.next().await.expect("request").expect("frame"); + websocket + .send(WebSocketMessage::Text(created("semantic-source").into())) + .await + .expect("response start"); + websocket + .send(WebSocketMessage::Text(completed("semantic-source").into())) + .await + .expect("response"); + }); + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let first_payload = ProtocolPayload::json_object( + "openai", + json!({ + "model": "gpt-5.4", + "reasoning": {"effort": "max"}, + "input": [ + { + "type": "message", + "role": "developer", + "content": [{ + "type": "input_text", + "text": "Proactive multi-agent delegation is active." + }] + }, + {"role": "user", "content": "remember this mode"} + ], + "conversation_id": "semantic-conversation", + "stream": true, + "store": false + }) + .as_object() + .expect("first request object") + .clone(), + ) + .expect("first payload"); + let first_operation = Operation::Generate(GenerateRequest::from_protocol_payload( + Vec::new(), + first_payload, + )); + let mut first = provider + .execute( + planned_request(&base_url, &first_operation), + context("req_semantics_first", None, CancellationToken::new()), + ) + .await + .expect("first provider stream"); + let mut provider_state = None; + while let Some(event) = first.next().await { + let mut event = event.expect("first canonical event"); + provider_state = event.take_session_update().or(provider_state); + } + server.await.expect("server task"); + + let next_payload = ProtocolPayload::json_object( + "openai", + json!({ + "model": "gpt-5.4", + "reasoning": {"effort": "max"}, + "input": [{"role": "user", "content": "continue"}], + "stream": true, + "store": false + }) + .as_object() + .expect("next request object") + .clone(), + ) + .expect("next payload"); + let next_operation = Operation::Generate( + GenerateRequest::from_protocol_payload(Vec::new(), next_payload) + .with_provider_session_state(provider_state.expect("connection state")), + ); + + let observation = provider.request_observation(&next_operation); + assert_eq!(observation.reasoning_preset.as_deref(), Some("ultra")); +} + #[tokio::test] async fn provider_metadata_uses_unified_account_resource_and_real_transport() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); @@ -401,16 +1987,490 @@ async fn cancelled_attempt_fails_before_selecting_or_sending() { #[tokio::test] async fn explicit_http_429_marks_provider_error_replay_safe() { - let (error, _) = rejected_http_operation(429).await; + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit exceeded, retry in 12 seconds" + } + }))) + .await; assert_eq!(error.kind(), ProviderErrorKind::RateLimited); assert_eq!(error.upstream_status(), Some(429)); + assert_eq!(error.send_state(), UpstreamSendState::Sent); assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_rate_limit_failure_updates_account_state_and_remains_replay_safe() { + let server = MockServer::start().await; + let body = concat!( + "event: response.failed\n", + "data: {\"type\":\"response.failed\",\"status_code\":429,\"response\":{\"id\":\"resp_rate_limited\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"sensitive upstream detail\"}}}\n\n", + ); + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .expect(1) + .mount(&server) + .await; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let provider = provider(&store); + let secondary = ProviderAccountId::new("acct_provider_other").expect("secondary account"); + let operation = http_operation(); + let mut stream = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context_with_exclusions( + "req_in_band_rate_limit", + NonZeroU32::new(1).expect("first attempt"), + BTreeSet::from([secondary]), + None, + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + let error = next_provider_error(&mut stream).await; + assert_eq!(error.kind(), ProviderErrorKind::RateLimited); + assert_eq!(error.send_state(), UpstreamSendState::Sent); + assert_eq!(error.upstream_status(), Some(429)); + assert_eq!( + error.upstream_code().map(SafeUpstreamValue::as_str), + Some("rate_limit_exceeded") + ); + assert!(error.replay_is_safe()); + assert!(error.sensitive_context_was_redacted()); + assert_eq!( + [ + store + .account("acct_provider") + .expect("primary account") + .availability(), + store + .account("acct_provider_other") + .expect("secondary account") + .availability(), + ], + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_failure_after_same_chunk_output_preserves_output_and_forbids_replay() { + let marker = "same-chunk-provider-secret"; + let server = MockServer::start().await; + let body = format!( + concat!( + "event: response.created\n", + "data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_partial\",\"model\":\"gpt-5.4\"}}}}\n\n", + "event: response.content_part.added\n", + "data: {{\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{{\"type\":\"output_text\"}}}}\n\n", + "event: response.output_text.delta\n", + "data: {{\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}}\n\n", + "event: response.failed\n", + "data: {{\"type\":\"response.failed\",\"status_code\":429,\"response\":{{\"id\":\"resp_partial\",\"status\":\"failed\",\"error\":{{\"code\":\"rate_limit_exceeded\",\"message\":\"{}\"}}}}}}\n\n" + ), + marker + ); + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .expect(1) + .mount(&server) + .await; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + create_other_account(&store).await; + let provider = provider(&store); + let secondary = ProviderAccountId::new("acct_provider_other").expect("secondary account"); + let operation = http_operation(); + let mut stream = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context_with_exclusions( + "req_in_band_output_failure", + NonZeroU32::new(1).expect("first attempt"), + BTreeSet::from([secondary]), + None, + None, + CancellationToken::new(), + ), + ) + .await + .expect("provider stream"); + + let mut saw_text = false; + let error = loop { + match stream.next().await.expect("typed failure terminal") { + Ok(event) => { + assert_ne!( + event.wire_event().and_then(|wire| wire.event_type()), + Some("response.failed"), + "raw failed event must not cross the Provider boundary" + ); + saw_text |= event.canonical_facts().iter().any( + |fact| matches!(fact, GatewayEvent::TextDelta(delta) if delta.text == "hello"), + ); + assert!(!format!("{event:?}").contains(marker)); + } + Err(error) => break error, + } + }; + + assert!( + saw_text, + "semantic output must precede the terminal failure" + ); + assert_eq!(error.kind(), ProviderErrorKind::RateLimited); + assert!(!error.replay_is_safe()); + assert_eq!( + error.upstream_code().map(SafeUpstreamValue::as_str), + Some("rate_limit_exceeded") + ); + assert!(!format!("{error:?} {error}").contains(marker)); + assert_eq!( + store + .account("acct_provider") + .expect("primary account") + .availability(), + AccountAvailability::Cooldown + ); + assert_eq!( + store + .account("acct_provider_other") + .expect("secondary account") + .availability(), + AccountAvailability::Ready + ); +} + +#[tokio::test] +async fn in_band_token_invalid_expires_only_the_selected_account() { + let (error, availability) = + in_band_failure("token_invalid", Some(401), "credential rejected").await; + + assert_eq!(error.kind(), ProviderErrorKind::Unauthorized); + assert!(error.replay_is_safe()); + assert_eq!( + error.upstream_code().map(SafeUpstreamValue::as_str), + Some("token_invalid") + ); + assert_eq!( + availability, + [AccountAvailability::Expired, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_model_not_supported_preserves_account_health() { + let (error, availability) = + in_band_failure("model_not_supported", Some(404), "model is unavailable").await; + + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_quota_exhaustion_updates_only_the_selected_account() { + let (error, availability) = + in_band_failure("insufficient_quota", Some(402), "quota exhausted").await; + + assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [ + AccountAvailability::QuotaExhausted, + AccountAvailability::Ready + ] + ); +} + +#[tokio::test] +async fn in_band_cyber_policy_is_not_replayed_or_applied_to_account_health() { + let (error, availability) = + in_band_failure("cyber_policy", Some(400), "request rejected").await; + + assert_eq!(error.kind(), ProviderErrorKind::InvalidRequest); + assert!(!error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_overload_is_not_replayed_or_applied_to_account_health() { + let (error, availability) = + in_band_failure("server_is_overloaded", None, "upstream overloaded").await; + + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert!(!error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn in_band_unknown_code_is_classified_without_persisting_opaque_text() { + let marker = "unknown-secret-account-marker"; + let (error, availability) = in_band_failure(marker, None, marker).await; + + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert!(error.upstream_code().is_none()); + assert!(!error.replay_is_safe()); + assert!(error.sensitive_context_was_redacted()); + assert!(!format!("{error:?}").contains(marker)); + assert!(!error.to_string().contains(marker)); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn rejected_response_passively_persists_rate_limit_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("x-codex-active-limit", "codex") + .insert_header("x-codex-primary-used-percent", "100") + .insert_header("x-codex-primary-window-minutes", "300") + .insert_header("x-codex-limit-reached", "true") + .set_body_json(json!({ + "error": {"code": "rate_limit_exceeded", "message": "rate limited"} + })), + ) + .expect(1) + .mount(&server) + .await; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store).await; + let provider = provider(&store); + let operation = http_operation(); + let mut stream = provider + .execute( + planned_request(&format!("{}/backend-api", server.uri()), &operation), + context("req_passive_rejection", None, CancellationToken::new()), + ) + .await + .expect("provider stream"); + + let error = next_provider_error(&mut stream).await; + assert_eq!(error.kind(), ProviderErrorKind::RateLimited); + let account = store.account("acct_provider").expect("account"); + let quota = store + .get_quotas(std::slice::from_ref(account.id())) + .await + .expect("quota") + .pop() + .and_then(|observation| observation.quota) + .expect("quota JSON"); + assert_eq!( + quota + .expose_to_provider() + .get("rate_limit") + .and_then(Value::as_object) + .and_then(|limit| limit.get("limit_reached")) + .and_then(Value::as_bool), + Some(true) + ); +} + +#[tokio::test] +async fn websocket_opening_429_is_proven_not_sent_and_replay_safe() { + let (error, availability) = rejected_websocket_opening( + 429, + r#"{"error":{"code":"rate_limit_exceeded","message":"Rate limit exceeded"}}"#, + ) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::RateLimited); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn websocket_opening_cloudflare_challenge_is_replay_safe_before_payload() { + let (error, availability) = + rejected_websocket_opening(403, "Just a moment...").await; + + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn http_cloudflare_challenge_is_replay_safe_after_explicit_rejection() { + let (error, availability) = rejected_http_operation( + ResponseTemplate::new(403).set_body_string("Just a moment..."), + ) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert_eq!(error.send_state(), UpstreamSendState::Sent); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn empty_http_404_path_block_is_replay_safe_after_explicit_rejection() { + let (error, availability) = rejected_http_operation(ResponseTemplate::new(404)).await; + + assert_eq!(error.kind(), ProviderErrorKind::Unavailable); + assert_eq!(error.send_state(), UpstreamSendState::Sent); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Cooldown, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn history_failures_keep_only_their_allowlisted_machine_code() { + for code in [ + "previous_response_not_found", + "invalid_encrypted_content", + "missing_tool_output", + "no_tool_output", + ] { + let marker = format!("{code}-sensitive-marker"); + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(400).set_body_json(json!({ + "error": {"code": code, "message": marker} + }))) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::InvalidRequest); + assert_eq!( + error.upstream_code().map(SafeUpstreamValue::as_str), + Some(code) + ); + assert!(error.replay_is_safe()); + assert_eq!( + error.continuation_failure(), + Some(gateway_core::error::ContinuationFailure::HistoryUnavailable) + ); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); + assert!(!format!("{error:?} {error}").contains(&marker)); + } +} + +#[tokio::test] +async fn deactivated_workspace_is_banned_before_payment_classification() { + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(402).set_body_json(json!({ + "detail": {"code": "deactivated_workspace", "message": "workspace deactivated"} + }))) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied); + assert!(error.replay_is_safe()); + assert_eq!( + error.upstream_code().map(SafeUpstreamValue::as_str), + Some("deactivated_workspace") + ); + assert_eq!( + availability, + [AccountAvailability::Banned, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn unauthorized_expires_only_the_rejected_account() { + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(401).set_body_json(json!({ + "error": {"code": "token_expired", "message": "access token expired"} + }))) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::Unauthorized); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Expired, AccountAvailability::Ready] + ); +} + +#[tokio::test] +async fn payment_required_exhausts_only_the_rejected_account() { + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"code": "insufficient_quota", "message": "quota exhausted"} + }))) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [ + AccountAvailability::QuotaExhausted, + AccountAvailability::Ready + ] + ); +} + +#[tokio::test] +async fn model_not_supported_is_instance_fallback_fact_without_account_poisoning() { + let (error, availability) = + rejected_http_operation(ResponseTemplate::new(404).set_body_json(json!({ + "error": {"code": "model_not_supported", "message": "model not supported"} + }))) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert!(error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); } #[tokio::test] async fn explicit_http_408_does_not_mark_provider_error_replay_safe() { - let (error, _) = rejected_http_operation(408).await; + let (error, _) = rejected_http_operation(ResponseTemplate::new(408)).await; assert_eq!(error.kind(), ProviderErrorKind::Timeout); assert!(!error.replay_is_safe()); @@ -422,7 +2482,7 @@ async fn generic_http_403_and_500_do_not_mark_provider_error_replay_safe() { (403, ProviderErrorKind::PermissionDenied), (500, ProviderErrorKind::Unavailable), ] { - let (error, availability) = rejected_http_operation(status).await; + let (error, availability) = rejected_http_operation(ResponseTemplate::new(status)).await; assert_eq!(error.kind(), expected_kind); assert!(!error.replay_is_safe()); @@ -433,6 +2493,18 @@ async fn generic_http_403_and_500_do_not_mark_provider_error_replay_safe() { } } +#[tokio::test] +async fn permission_denied_failure_keeps_every_account_ready() { + let (error, availability) = rejected_http_operation(ResponseTemplate::new(403)).await; + + assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied); + assert!(!error.replay_is_safe()); + assert_eq!( + availability, + [AccountAvailability::Ready, AccountAvailability::Ready] + ); +} + #[tokio::test] async fn capability_query_compiles_only_explicit_catalog_evidence() { let server = MockServer::start().await; diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index a6d2cfe40..5b9a70d95 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -2,10 +2,12 @@ use std::collections::BTreeMap; use std::num::NonZeroU32; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use async_trait::async_trait; +use futures::future::BoxFuture; use gateway_core::engine::credential::{ AccountAvailability, AccountRuntimeSignals, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, LoadedCredential, NewProviderAccount, ProviderAccount, @@ -13,12 +15,18 @@ use gateway_core::engine::credential::{ QuotaWriteOutcome, }; use gateway_core::error::{StoreError, StoreErrorKind}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest, ProviderSchedulingState, + ProviderStoreError, +}; use gateway_core::routing::{ProviderInstanceId, ProviderKind}; +use provider_openai::CodexOriginPolicy; use provider_openai::credential::{ - CodexAccountProfile, CodexCredentialRepository, CodexOAuthSecret, CredentialLeaseCoordinator, - CredentialLeaseCoordinatorError, CredentialLeaseRequest, LeaseAcquisition, + CodexAccountProfile, CodexCredentialRepository, CodexOAuthSecret, }; use secrecy::SecretString; +use url::{Host, Url}; #[derive(Clone)] struct StoredAccount { @@ -30,6 +38,7 @@ struct StoredAccount { #[derive(Default)] pub(crate) struct MemoryAccountStore { accounts: Mutex>, + quota_reads: AtomicUsize, } impl MemoryAccountStore { @@ -45,6 +54,10 @@ impl MemoryAccountStore { .get(&id) .map(|stored| stored.account.clone()) } + + pub(crate) fn quota_reads(&self) -> usize { + self.quota_reads.load(Ordering::SeqCst) + } } #[async_trait] @@ -156,21 +169,20 @@ impl ProviderAccountStore for MemoryAccountStore { }, ); stored.credential = credential; + stored.quota = None; Ok(CredentialCasOutcome::Updated(next)) } - async fn get_quota( + async fn get_quotas( &self, - account: &ProviderAccountId, - ) -> Result, StoreError> { - Ok(self - .accounts - .lock() - .expect("account store lock") - .get(account) - .ok_or_else(|| store_error(StoreErrorKind::InvalidData))? - .quota - .clone()) + accounts: &[ProviderAccountId], + ) -> Result, StoreError> { + self.quota_reads.fetch_add(1, Ordering::SeqCst); + let stored = self.accounts.lock().expect("account store lock"); + Ok(accounts + .iter() + .filter_map(|account| stored.get(account)?.quota.clone()) + .collect()) } async fn compare_and_swap_quota( @@ -272,6 +284,21 @@ const fn store_error(kind: StoreErrorKind) -> StoreError { StoreError::new(kind) } +#[derive(Debug)] +pub(crate) struct LoopbackCodexOriginPolicy; + +impl CodexOriginPolicy for LoopbackCodexOriginPolicy { + fn allows(&self, url: &Url) -> bool { + url.scheme() == "http" + && matches!(url.host(), Some(Host::Ipv4(host)) if host.is_loopback()) + && url.port().is_some() + } +} + +pub(crate) fn loopback_origin_policy() -> Arc { + Arc::new(LoopbackCodexOriginPolicy) +} + struct AccountRebuild { revision: CredentialRevision, enabled: bool, @@ -315,78 +342,98 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid #[derive(Default)] pub(crate) struct TestLeaseCoordinator { - pub(crate) requests: Mutex>, + pub(crate) requests: Mutex>, pub(crate) busy: Mutex, round_robin_cursor: Mutex, } -#[async_trait] -impl CredentialLeaseCoordinator for TestLeaseCoordinator { - async fn runtime_signals( - &self, - accounts: &[ProviderAccountId], - ) -> Result, CredentialLeaseCoordinatorError> - { - Ok(accounts - .iter() - .cloned() - .map(|account| { - ( - account, - AccountRuntimeSignals { - in_flight: 0, - last_started_at: None, - quota_reset_at: None, - quota_remaining_rank: None, - }, - ) - }) - .collect()) +impl ProviderLeasePort for TestLeaseCoordinator { + fn load_state<'a>( + &'a self, + _provider_instance_id: &'a ProviderInstanceId, + accounts: &'a [ProviderAccountId], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let signals = accounts + .iter() + .cloned() + .map(|account| { + ( + account, + AccountRuntimeSignals { + in_flight: 0, + last_started_at: None, + quota_reset_at: None, + quota_remaining_rank: None, + }, + ) + }) + .collect(); + let mut cursor = self + .round_robin_cursor + .lock() + .expect("round robin cursor lock"); + let current = *cursor; + *cursor = cursor.wrapping_add(1); + Ok(ProviderSchedulingState::new(signals, None, current)) + }) } - fn next_round_robin_cursor( + fn try_acquire( &self, - _provider_instance_id: &ProviderInstanceId, - ) -> Result { - let mut cursor = self - .round_robin_cursor - .lock() - .expect("round robin cursor lock"); - let current = *cursor; - *cursor = cursor.wrapping_add(1); - Ok(current) - } - - async fn try_acquire( - &self, - request: CredentialLeaseRequest, - ) -> Result { - self.requests - .lock() - .expect("lease requests lock") - .push(request); - if *self.busy.lock().expect("lease busy lock") { - Ok(LeaseAcquisition::Busy { - retry_after: Some(Duration::from_millis(25)), - }) - } else { - Ok(LeaseAcquisition::Acquired(Box::new(()))) - } + request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let ProviderLeaseRequest::Scheduling(request) = request else { + panic!("expected scheduling lease request"); + }; + self.requests + .lock() + .expect("lease requests lock") + .push(request); + if *self.busy.lock().expect("lease busy lock") { + Ok(ProviderLeaseAcquisition::Busy { + retry_after: Some(Duration::from_millis(25)), + }) + } else { + Ok(ProviderLeaseAcquisition::Acquired(Box::new(()))) + } + }) } } pub(crate) fn profile(account_id: &str) -> CodexAccountProfile { let now = chrono::Utc::now(); CodexAccountProfile { + oauth_subject: format!("subject-{account_id}"), + poid: Some(format!("poid-{account_id}")), chatgpt_account_id: account_id.to_owned(), - chatgpt_user_id: Some(format!("user-{account_id}")), + chatgpt_user_id: format!("user-{account_id}"), email: Some(format!("{account_id}@example.com")), plan_type: Some("pro".to_owned()), access_token_expires_at: Some(now + chrono::Duration::hours(1)), - next_refresh_at: Some(now + chrono::Duration::minutes(45)), } } +pub(crate) struct StaticRuntimePolicy; + +impl ProviderRuntimePolicyPort for StaticRuntimePolicy { + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + ProviderRefreshPolicy::try_new( + Duration::from_secs(60 * 60), + NonZeroU32::new(2).expect("positive concurrency"), + ) + }) + } +} + +pub(crate) fn runtime_policy() -> Arc { + Arc::new(StaticRuntimePolicy) +} + pub(crate) fn secret(access_token: &str) -> CodexOAuthSecret { CodexOAuthSecret { access_token: SecretString::from(access_token.to_owned()), diff --git a/backend/crates/providers/openai/tests/transport/canonical.rs b/backend/crates/providers/openai/tests/transport/canonical.rs index af8279b30..ecd7e6634 100644 --- a/backend/crates/providers/openai/tests/transport/canonical.rs +++ b/backend/crates/providers/openai/tests/transport/canonical.rs @@ -1,8 +1,35 @@ use gateway_core::engine::UpstreamSendState; use gateway_core::error::ProviderErrorKind; -use gateway_core::event::{ContentKind, FinishReason, GatewayEvent}; +use gateway_core::event::{ContentKind, FinishReason, GatewayEvent, ProviderEvent}; -use provider_openai::CodexCanonicalDecoder; +use provider_openai::transport::canonical::{ + CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalFailure, CodexCanonicalOutcome, +}; + +const METADATA_PREFIX_FIXTURE: &str = include_str!( + "../../../../gateway-api/tests/openai/responses/fixtures/http_sse/metadata_only_prefix.sse" +); + +#[test] +fn decoder_should_preserve_existing_metadata_fixture_as_openai_wire() { + let events = CodexCanonicalDecoder::new("fallback") + .push(METADATA_PREFIX_FIXTURE.as_bytes()) + .expect("metadata fixture should remain open-world"); + let wire_types = events + .iter() + .filter_map(ProviderEvent::wire_event) + .filter_map(|wire| wire.event_type()) + .collect::>(); + + assert_eq!( + wire_types, + vec!["response.created", "codex.rate_limits", "response.metadata"] + ); + assert!(matches!( + canonical_facts(&events).as_slice(), + [GatewayEvent::Started(_)] + )); +} #[test] fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() { @@ -16,7 +43,7 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() { .push(body.as_bytes()) .expect("canonical priced response"); - assert!(events.iter().any(|event| matches!( + assert!(canonical_facts(&events).into_iter().any(|event| matches!( event, GatewayEvent::CalculatedCost(cost) if cost.total().amount().scaled() == 3_437_500 @@ -38,17 +65,18 @@ fn decoder_should_normalize_text_usage_and_completion() { let events = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) .expect("canonical response"); + let canonical = canonical_facts(&events); - assert!(matches!(events[0], GatewayEvent::Started(_))); + assert!(matches!(canonical[0], GatewayEvent::Started(_))); assert!(matches!( - events[1], - GatewayEvent::ContentAdded(ref item) if item.kind() == ContentKind::Text + canonical[1], + GatewayEvent::ContentAdded(item) if item.kind() == ContentKind::Text )); - assert!(matches!(events[2], GatewayEvent::TextDelta(_))); - assert!(matches!(events[3], GatewayEvent::Usage(_))); + assert!(matches!(canonical[2], GatewayEvent::TextDelta(_))); + assert!(matches!(canonical[3], GatewayEvent::Usage(_))); assert!(matches!( - events[4], - GatewayEvent::Completed(ref meta) + canonical[4], + GatewayEvent::Completed(meta) if meta.finish_reason() == Some(FinishReason::Stop) )); } @@ -71,9 +99,10 @@ fn decoder_should_accept_official_codex_metadata_lifecycle_events() { let events = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) .expect("official Codex metadata lifecycle events"); + let canonical = canonical_facts(&events); assert!(matches!( - events.as_slice(), + canonical.as_slice(), [ GatewayEvent::Started(_), GatewayEvent::ContentAdded(_), @@ -100,6 +129,7 @@ fn decoder_should_accept_whole_function_call_without_argument_deltas() { .expect("canonical function call"); let tool_deltas = events .iter() + .flat_map(ProviderEvent::canonical_facts) .filter_map(|event| match event { GatewayEvent::ToolCallDelta(delta) => Some(delta), _ => None, @@ -112,21 +142,57 @@ fn decoder_should_accept_whole_function_call_without_argument_deltas() { } #[test] -fn decoder_should_fail_closed_for_unknown_events_without_echoing_body() { +fn decoder_should_preserve_unknown_events_without_exposing_wire_data_in_debug() { let body = concat!( "event: response.created\n", "data: {\"response\":{\"id\":\"resp_1\",\"model\":\"gpt-test\"}}\n\n", "event: response.secret_future_event\n", "data: {\"secret\":\"must-not-leak\"}\n\n", ); - let error = CodexCanonicalDecoder::new("fallback") + let events = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) - .expect_err("unknown event must fail closed"); + .expect("unknown OpenAI event should remain wire-only"); + let unknown = events[1] + .wire_event() + .expect("unknown event should retain wire data"); - assert_eq!(error.kind(), ProviderErrorKind::Unsupported); - assert_eq!(error.send_state(), UpstreamSendState::Sent); - assert!(!format!("{error:?}").contains("must-not-leak")); - assert!(!format!("{error}").contains("must-not-leak")); + assert_eq!(unknown.event_type(), Some("response.secret_future_event")); + assert_eq!( + unknown.data().get("secret"), + Some(&serde_json::json!("must-not-leak")) + ); + assert!(!format!("{unknown:?}").contains("must-not-leak")); +} + +#[test] +fn decoder_should_keep_media_and_hosted_tool_events_as_openai_wire() { + for event_type in [ + "response.image_generation_call.partial_image", + "response.audio.delta", + "response.web_search_call.searching", + "response.code_interpreter_call.in_progress", + "response.computer_tool_call.in_progress", + ] { + let body = format!( + "event: response.created\ndata: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_open_world\",\"model\":\"gpt-test\"}}}}\n\nevent: {event_type}\ndata: {{\"type\":\"{event_type}\",\"opaque\":{{\"future\":true}}}}\n\n" + ); + let events = CodexCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("open-world event should remain protocol wire"); + let event = &events[1]; + + assert!(event.canonical_facts().is_empty(), "{event_type}"); + assert_eq!( + event.wire_event().and_then(|wire| wire.event_type()), + Some(event_type) + ); + assert_eq!( + event + .wire_event() + .and_then(|wire| wire.data().pointer("/opaque/future")), + Some(&serde_json::json!(true)) + ); + } } #[test] @@ -144,7 +210,10 @@ fn decoder_finish_should_parse_a_final_frame_without_blank_line() { decoder.push(tail.as_bytes()).expect("buffer partial frame"); let events = decoder.finish().expect("finish partial frame"); - assert!(matches!(events.as_slice(), [GatewayEvent::Completed(_)])); + assert!(matches!( + canonical_facts(&events).as_slice(), + [GatewayEvent::Completed(_)] + )); } #[test] @@ -161,52 +230,78 @@ fn decoder_should_accept_done_only_after_terminal_event() { .push(body.as_bytes()) .expect("terminal done marker"); - assert!(matches!(events.last(), Some(GatewayEvent::Completed(_)))); + assert!(matches!( + canonical_facts(&events).last(), + Some(GatewayEvent::Completed(_)) + )); } #[test] fn decoder_should_classify_official_token_invalid_failure() { - assert_failed_event( - "token_invalid", - ProviderErrorKind::Unauthorized, - "auth-secret-marker", - ); + assert_failed_event("token_invalid", "auth-secret-marker"); } #[test] fn decoder_should_classify_official_model_not_supported_failure() { - assert_failed_event( - "model_not_supported", - ProviderErrorKind::Unsupported, - "model-secret-marker", - ); + assert_failed_event("model_not_supported", "model-secret-marker"); } #[test] fn decoder_should_classify_official_quota_failure() { - assert_failed_event( - "quota_exceeded", - ProviderErrorKind::QuotaExhausted, - "quota-secret-marker", - ); + assert_failed_event("quota_exceeded", "quota-secret-marker"); } #[test] fn decoder_should_classify_official_server_overloaded_failure() { - assert_failed_event( - "server_is_overloaded", - ProviderErrorKind::Unavailable, - "server-secret-marker", - ); + assert_failed_event("server_is_overloaded", "server-secret-marker"); } #[test] fn decoder_should_classify_official_cyber_policy_as_an_invalid_request() { - assert_failed_event( - "cyber_policy", - ProviderErrorKind::InvalidRequest, - "policy-secret-marker", + assert_failed_event("cyber_policy", "policy-secret-marker"); +} + +#[test] +fn decoder_should_preserve_same_chunk_output_before_typed_failure() { + let marker = "same-chunk-secret-marker"; + let body = format!( + concat!( + "event: response.created\n", + "data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_partial\",\"model\":\"gpt-test\"}}}}\n\n", + "event: response.content_part.added\n", + "data: {{\"type\":\"response.content_part.added\",\"output_index\":0,\"content_index\":0,\"part\":{{\"type\":\"output_text\"}}}}\n\n", + "event: response.output_text.delta\n", + "data: {{\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"hello\"}}\n\n", + "event: response.failed\n", + "data: {{\"type\":\"response.failed\",\"response\":{{\"id\":\"resp_partial\",\"status\":\"failed\",\"error\":{{\"code\":\"rate_limit_exceeded\",\"message\":\"{}\"}}}}}}\n\n" + ), + marker ); + + let failure = CodexCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect_err("response.failed must retain the preceding batch"); + let facts = canonical_facts(failure.events()); + let wire_types = failure + .events() + .iter() + .filter_map(ProviderEvent::wire_event) + .filter_map(|wire| wire.event_type()) + .collect::>(); + + assert!(failure.semantic_output_seen()); + assert!(matches!(facts[0], GatewayEvent::Started(_))); + assert!(matches!(facts[1], GatewayEvent::ContentAdded(_))); + assert!(matches!(facts[2], GatewayEvent::TextDelta(_))); + assert_eq!( + wire_types, + vec![ + "response.created", + "response.content_part.added", + "response.output_text.delta" + ] + ); + assert!(!format!("{failure:?}").contains(marker)); } #[test] @@ -233,25 +328,31 @@ fn decoder_should_reject_a_changed_upstream_response_id() { "data: {\"response\":{\"id\":\"resp_changed\",\"model\":\"gpt-test\"}}\n\n", ); - let error = CodexCanonicalDecoder::new("fallback") + let failure = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) .expect_err("response ID changes must fail closed"); + let CodexCanonicalError::Protocol(error) = failure.error() else { + panic!("response ID changes are protocol failures"); + }; assert_eq!(error.kind(), ProviderErrorKind::Protocol); assert_eq!(error.send_state(), UpstreamSendState::Sent); } -fn assert_failed_event(code: &str, expected: ProviderErrorKind, marker: &str) { +fn assert_failed_event(code: &str, marker: &str) { let body = format!( "event: response.failed\ndata: {{\"response\":{{\"id\":\"resp_failed\",\"status\":\"failed\",\"error\":{{\"code\":\"{code}\",\"message\":\"{marker}\"}}}}}}\n\n" ); - let error = CodexCanonicalDecoder::new("fallback") + let failure = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) .expect_err("failed event must become a typed error"); - assert_eq!(error.kind(), expected); - assert_eq!(error.send_state(), UpstreamSendState::Sent); - assert!(!format!("{error:?} {error}").contains(marker)); + let CodexCanonicalError::Upstream(upstream) = failure.error() else { + panic!("response.failed must preserve its typed upstream failure"); + }; + assert_eq!(upstream.upstream_code.as_deref(), Some(code)); + assert!(!failure.semantic_output_seen()); + assert!(!format!("{failure:?}").contains(marker)); } fn assert_incomplete_reason(reason: &str, expected: FinishReason) { @@ -261,10 +362,40 @@ fn assert_incomplete_reason(reason: &str, expected: FinishReason) { let events = CodexCanonicalDecoder::new("fallback") .push(body.as_bytes()) .expect("incomplete response is terminal"); - let finish_reason = events.iter().find_map(|event| match event { - GatewayEvent::Completed(meta) => meta.finish_reason(), - _ => None, - }); + let finish_reason = canonical_facts(&events) + .into_iter() + .find_map(|event| match event { + GatewayEvent::Completed(meta) => meta.finish_reason(), + _ => None, + }); assert_eq!(finish_reason, Some(expected)); } + +trait CanonicalOutcomeAssertions { + fn expect(self, message: &str) -> Vec; + fn expect_err(self, message: &str) -> CodexCanonicalFailure; +} + +impl CanonicalOutcomeAssertions for CodexCanonicalOutcome { + fn expect(self, message: &str) -> Vec { + match self { + Self::Events(events) => events, + Self::Failed(failure) => panic!("{message}: {failure:?}"), + } + } + + fn expect_err(self, message: &str) -> CodexCanonicalFailure { + match self { + Self::Events(events) => panic!("{message}: decoded {} events", events.len()), + Self::Failed(failure) => failure, + } + } +} + +fn canonical_facts(events: &[ProviderEvent]) -> Vec<&GatewayEvent> { + events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .collect() +} diff --git a/backend/crates/providers/openai/tests/transport/client.rs b/backend/crates/providers/openai/tests/transport/client.rs index 7173f7db2..174034637 100644 --- a/backend/crates/providers/openai/tests/transport/client.rs +++ b/backend/crates/providers/openai/tests/transport/client.rs @@ -1,3 +1,4 @@ +use provider_openai::transport::CodexUpstreamSendPhase; use provider_openai::transport::{CodexBackendTransport, CodexClientError, CodexRequestContext}; #[test] @@ -11,6 +12,8 @@ fn upstream_error_formatting_should_redact_body() { set_cookie_headers: vec!["session=secret".to_owned()], rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, + transport_metrics: Box::default(), + send_phase: CodexUpstreamSendPhase::AfterPayload, }; assert!(!format!("{error}").contains(secret)); diff --git a/backend/crates/providers/openai/tests/transport/headers.rs b/backend/crates/providers/openai/tests/transport/headers.rs index c6dfcb6bc..0625ac7c0 100644 --- a/backend/crates/providers/openai/tests/transport/headers.rs +++ b/backend/crates/providers/openai/tests/transport/headers.rs @@ -265,9 +265,13 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f for (name, expected) in [ ("x-client-request-id", "req_ws_security"), ("x-codex-installation-id", "install-123"), + ("x-openai-internal-codex-residency", "us"), ("x-codex-turn-state", "turn-state"), + ("x-codex-turn-metadata", "{\"thread_source\":\"subagent\"}"), ("x-codex-beta-features", "feature-a"), + ("x-responsesapi-include-timing-metrics", "true"), ("version", "26.318.11754"), + ("x-codex-window-id", "cw_derived"), ("x-codex-parent-thread-id", "parent-456"), ("x-openai-subagent", "review"), ("x-openai-memgen-request", "true"), @@ -283,6 +287,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f "content-type", "accept", "session_id", + "thread-id", "x-openai-internal-codex-responses-lite", ] { assert!(headers.iter().all(|(header, _)| header != forbidden)); @@ -380,3 +385,123 @@ async fn backend_http_should_send_codex_context_without_browser_headers() { assert!(header_names.iter().all(|name| name != forbidden)); } } + +#[tokio::test] +async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_latest_wire_profile() +{ + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind profile server"); + let address = listener.local_addr().expect("profile server address"); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.expect("first websocket"); + let mut first_user_agent = String::new(); + let mut first = accept_codex_test_websocket_with(first_stream, |request, _response| { + first_user_agent = request + .headers() + .get("user-agent") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(); + }) + .await; + let _ = first + .next() + .await + .expect("first response.create") + .expect("valid first response.create"); + first + .send(Message::Text( + completed_websocket_response("resp_profile_first", 1, 1).into(), + )) + .await + .expect("first response.completed"); + + let _ = timeout(Duration::from_secs(2), first.next()) + .await + .expect("exact continuation should keep the original websocket") + .expect("continuation response.create") + .expect("valid continuation response.create"); + first + .send(Message::Text( + completed_websocket_response("resp_profile_continued", 1, 1).into(), + )) + .await + .expect("continuation response.completed"); + + let (second_stream, _) = timeout(Duration::from_secs(2), listener.accept()) + .await + .expect("new chain should use a new wire profile connection") + .expect("second websocket"); + let mut second_user_agent = String::new(); + let mut second = accept_codex_test_websocket_with(second_stream, |request, _response| { + second_user_agent = request + .headers() + .get("user-agent") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(); + }) + .await; + let _ = second + .next() + .await + .expect("second response.create") + .expect("valid second response.create"); + second + .send(Message::Text( + completed_websocket_response("resp_profile_second", 1, 1).into(), + )) + .await + .expect("second response.completed"); + + (first_user_agent, second_user_agent) + }); + + let profile = test_wire_profile(); + let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let client = CodexBackendClient::new( + reqwest::Client::builder() + .no_proxy() + .build() + .expect("HTTP client"), + format!("http://{address}"), + profile.clone(), + ) + .with_websocket_pool(pool); + let mut request = CodexResponsesRequest::new_http_sse("gpt-test", "", Vec::new()); + request.use_websocket = true; + request.local_conversation_id = Some("profile-rotation".to_owned()); + + client + .create_response( + &request, + request_context("req_profile_first", Some("acct-profile")), + ) + .await + .expect("first response"); + profile.update_desktop_release("26.900.1", "7001"); + + let mut continuation = request.clone(); + continuation.set_previous_response_id(Some("resp_profile_first".to_owned())); + continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + client + .create_response( + &continuation, + request_context("req_profile_continued", Some("acct-profile")), + ) + .await + .expect("continued response"); + client + .create_response( + &request, + request_context("req_profile_second", Some("acct-profile")), + ) + .await + .expect("second response"); + + let (first_user_agent, second_user_agent) = server.await.expect("profile server task"); + assert!(first_user_agent.contains("1.2.3")); + assert!(second_user_agent.contains("26.900.1")); + assert_ne!(first_user_agent, second_user_agent); +} diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index 5edea277d..2ef623672 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -452,9 +452,6 @@ async fn connection_local_continuation_should_use_the_exact_socket() { .await .unwrap(); let mut second_request = first_request; - // DB 只保存 upstream response ID，不保存客户端 conversation ID；续接必须仍能 - // 按 origin/account/native handle 找回原连接。 - second_request.local_conversation_id = Some("conversation-after-restart".to_string()); second_request.set_previous_response_id(Some("resp_exact_seed".to_string())); second_request.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); let second = backend @@ -466,9 +463,9 @@ async fn connection_local_continuation_should_use_the_exact_socket() { .unwrap(); server.await.unwrap(); - assert!(first.connection_local_continuation_expires_at.is_some()); + assert!(first.connection_local_continuation); assert!(first.transport_metrics.first_event_ms.is_some()); - assert!(second.connection_local_continuation_expires_at.is_some()); + assert!(second.connection_local_continuation); assert_eq!( second.transport_metrics.decision, Some(CodexTransportDecision::ExactWebSocket) diff --git a/backend/crates/providers/openai/tests/transport/profile.rs b/backend/crates/providers/openai/tests/transport/profile.rs index 7d1d214f7..7cbcb871e 100644 --- a/backend/crates/providers/openai/tests/transport/profile.rs +++ b/backend/crates/providers/openai/tests/transport/profile.rs @@ -1,5 +1,38 @@ +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + use chrono::{TimeZone, Utc}; -use provider_openai::transport::profile::CodexWireProfile; +use futures::future::BoxFuture; +use provider_openai::transport::profile::{ + CodexDesktopRelease, CodexDesktopReleaseError, CodexDesktopReleaseService, + CodexDesktopReleaseTransport, CodexWireProfile, CodexWireProfileState, parse_desktop_release, +}; + +struct ReleaseTransport { + outcomes: Mutex>>, +} + +impl ReleaseTransport { + fn new( + outcomes: impl IntoIterator>, + ) -> Self { + Self { + outcomes: Mutex::new(outcomes.into_iter().collect()), + } + } +} + +impl CodexDesktopReleaseTransport for ReleaseTransport { + fn fetch(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { + self.outcomes + .lock() + .expect("release outcomes") + .pop_front() + .expect("release outcome") + }) + } +} #[test] fn wire_profile_should_generate_codex_core_user_agent() { @@ -23,3 +56,141 @@ fn wire_profile_should_generate_codex_core_user_agent() { "Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.707.72221)" ); } + +#[test] +fn desktop_appcast_should_publish_the_first_complete_release() { + let release = parse_desktop_release( + r#" + + + + Incomplete release without artifact identity + + + Sun, 19 Jul 2026 08:00:00 +0000 + 14.0 + arm64 + + + + + + + "#, + ) + .expect("valid appcast"); + + assert_eq!(release.version, "26.720.1"); + assert_eq!(release.build, "72001"); + assert_eq!(release.download_size, Some(123_456)); + assert!(release.signature_present); +} + +#[test] +fn desktop_appcast_should_reject_invalid_version_and_build() { + for (version, build, expected) in [ + ("latest", "72001", CodexDesktopReleaseError::InvalidVersion), + ( + "26.720.1", + "build-72001", + CodexDesktopReleaseError::InvalidBuild, + ), + ] { + let error = parse_desktop_release(&format!( + r#""#, + )) + .expect_err("invalid release identity"); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + } +} + +#[tokio::test] +async fn wire_profile_release_update_should_change_only_desktop_identity() { + let original = wire_profile(); + let state = CodexWireProfileState::new(original.clone()); + let service = CodexDesktopReleaseService::new( + state.clone(), + Arc::new(ReleaseTransport::new([Ok(release("26.720.1", "72001"))])), + ); + + service.refresh().await.expect("release refresh"); + + let updated = state.snapshot(); + assert_eq!(updated.desktop_version, "26.720.1"); + assert_eq!(updated.desktop_build, "72001"); + assert_eq!(updated.codex_version, original.codex_version); + assert_eq!(updated.originator, original.originator); + assert_eq!(updated.os_type, original.os_type); + assert_eq!(updated.os_version, original.os_version); + assert_eq!(updated.arch, original.arch); + assert_eq!(updated.terminal, original.terminal); + assert_eq!(updated.verified_at, original.verified_at); + let status = service.status().snapshot(); + assert_eq!(status.latest, Some(release("26.720.1", "72001"))); + assert!(status.checked_at.is_some()); + assert!(status.last_error.is_none()); +} + +#[tokio::test] +async fn failed_release_refresh_should_preserve_the_last_successful_profile_and_release() { + let state = CodexWireProfileState::new(wire_profile()); + let latest = release("26.720.1", "72001"); + let service = CodexDesktopReleaseService::new( + state.clone(), + Arc::new(ReleaseTransport::new([ + Ok(latest.clone()), + Err(CodexDesktopReleaseError::InvalidDocument), + ])), + ); + service.refresh().await.expect("first release refresh"); + let successful_profile = state.snapshot(); + + service + .refresh() + .await + .expect_err("second release refresh should fail"); + + assert_eq!(state.snapshot(), successful_profile); + let status = service.status().snapshot(); + assert_eq!(status.latest, Some(latest)); + assert_eq!( + status.last_error.as_deref(), + Some("Codex Desktop appcast document is invalid") + ); + assert!(status.checked_at.is_some()); +} + +fn wire_profile() -> CodexWireProfile { + CodexWireProfile { + originator: "Codex Desktop".to_owned(), + codex_version: "0.144.2".to_owned(), + desktop_version: "26.707.72221".to_owned(), + desktop_build: "72221".to_owned(), + os_type: "Mac OS".to_owned(), + os_version: "15.7.1".to_owned(), + arch: "arm64".to_owned(), + terminal: "unknown".to_owned(), + verified_at: Utc + .with_ymd_and_hms(2026, 7, 18, 0, 0, 0) + .single() + .expect("valid fixture time"), + } +} + +fn release(version: &str, build: &str) -> CodexDesktopRelease { + CodexDesktopRelease { + version: version.to_owned(), + build: build.to_owned(), + published_at: None, + minimum_system_version: Some("14.0".to_owned()), + hardware_requirements: Some("arm64".to_owned()), + download_url: Some("https://persistent.oaistatic.com/codex.dmg".to_owned()), + download_size: Some(123_456), + signature_present: true, + } +} diff --git a/backend/crates/providers/openai/tests/transport/protocol/responses.rs b/backend/crates/providers/openai/tests/transport/protocol/responses.rs index 9f270293f..fb11159ef 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/responses.rs @@ -215,15 +215,21 @@ fn transport_should_require_exact_websocket_for_connection_local_continuation() #[test] fn exact_continuation_should_remain_stricter_than_forced_websocket() { - let mut request = CodexResponsesRequest::new_http_sse("gpt-test", "", Vec::new()); - request.set_previous_response_id(Some("resp_local".to_owned())); - request.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); - request.force_websocket = true; - + let mut preferred = CodexResponsesRequest::new_http_sse("gpt-test", "be brief", Vec::new()); + preferred.use_websocket = true; assert_eq!( - transport_requirement(&request), + transport_requirement(&preferred), + TransportRequirement::NewChain + ); + assert!(!transport_requirement(&preferred).requires_websocket()); + + preferred.set_previous_response_id(Some("resp_previous".to_owned())); + preferred.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + assert_eq!( + transport_requirement(&preferred), TransportRequirement::ExactWebSocketContinuation ); + assert!(transport_requirement(&preferred).requires_websocket()); } #[test] diff --git a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs index 0b1c1ab75..c56e69962 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs @@ -90,15 +90,23 @@ fn websocket_response_create_payload_should_keep_explicit_empty_instructions() { #[test] fn required_websocket_audit_should_forbid_http_fallback() { - let mut request = CodexResponsesRequest::new_http_sse("gpt-test", "", Vec::new()); - request.use_websocket = true; - request.force_websocket = true; + let request = CodexResponsesRequest::from_body( + json!({ + "model": "gpt-test", + "input": [], + "stream": true, + "store": false, + "generate": false + }) + .as_object() + .expect("request object") + .clone(), + ); let payload = websocket_payload_audit_snapshot(&request); - let artifact = websocket_audit_artifact_from_attempt(&request, OpeningAuditSnapshot::default(), payload); - assert_eq!(artifact.transport_mode, "websocket_required"); + assert_eq!(artifact.transport_mode, "explicit_websocket_warmup"); assert!(!artifact.fallback_allowed); } diff --git a/backend/crates/providers/openai/tests/transport/request.rs b/backend/crates/providers/openai/tests/transport/request.rs index 95ca649ee..67f895ec4 100644 --- a/backend/crates/providers/openai/tests/transport/request.rs +++ b/backend/crates/providers/openai/tests/transport/request.rs @@ -1,11 +1,70 @@ use gateway_core::operation::{ - ContentPart, GenerateRequest, Message, MessageRole, ProviderOptions, ReasoningEffort, - ReasoningRequirement, ReasoningSummary, ResponsePersistence, ToolDefinition, + ContentPart, GenerateRequest, Message, MessageRole, ProtocolPayload, ProviderOptions, + ReasoningEffort, ReasoningRequirement, ReasoningSummary, ResponsePersistence, ToolDefinition, }; use serde_json::{Map, Value, json}; use provider_openai::{CodexRequestEncodeError, codex_request_semantics, encode_generate_request}; +#[test] +fn encoder_should_preserve_openai_wire_fields_without_deriving_accountless_pool_identity() { + let body = Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ( + "input".to_owned(), + json!([{"role":"user","content":"private stable prompt"}]), + ), + ("include".to_owned(), json!(["reasoning.encrypted_content"])), + ("tool_choice".to_owned(), json!("auto")), + ("service_tier".to_owned(), json!("priority")), + ("conversation_id".to_owned(), json!("private-conversation")), + ("session_id".to_owned(), json!("private-session")), + ("turnState".to_owned(), json!("private-turn-state")), + ("future_official_field".to_owned(), json!({"enabled": true})), + ]); + let payload = ProtocolPayload::json_object("openai", body).expect("OpenAI payload"); + let request = GenerateRequest::from_protocol_payload(Vec::new(), payload); + + let encoded = encode_generate_request(&request, "gpt-routed").expect("encode wire payload"); + + assert_eq!(encoded.body().get("model"), Some(&json!("gpt-routed"))); + assert!(encoded.body().get("stream").is_none()); + assert!(encoded.body().get("store").is_none()); + assert_eq!(encoded.body().get("tool_choice"), Some(&json!("auto"))); + assert_eq!( + encoded.body().get("future_official_field"), + Some(&json!({"enabled": true})) + ); + assert_eq!(encoded.turn_state.as_deref(), Some("private-turn-state")); + assert_eq!( + encoded.client_session_id.as_deref(), + Some("private-session") + ); + assert!(encoded.local_conversation_id.is_none()); + assert!(!format!("{encoded:?}").contains("private stable prompt")); +} + +#[test] +fn encoder_should_never_hash_prompt_content_into_an_accountless_pool_identity() { + let request = |input: &str| { + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("input".to_owned(), Value::String(input.to_owned())), + ]), + ) + .expect("OpenAI payload"); + GenerateRequest::from_protocol_payload(Vec::new(), payload) + }; + + for input in ["private stable prompt", "different private prompt"] { + let encoded = + encode_generate_request(&request(input), "gpt-routed").expect("encoded request"); + assert!(encoded.local_conversation_id.is_none()); + } +} + #[test] fn encoder_should_project_typed_generate_semantics() { let message = Message::new( @@ -42,7 +101,6 @@ fn encoder_should_project_typed_generate_semantics() { assert_eq!(body.pointer("/tools/0/strict"), Some(&json!(true))); assert_eq!(body.pointer("/reasoning/effort"), Some(&json!("high"))); assert!(!encoded.force_http_sse); - assert!(!encoded.force_websocket); } #[test] @@ -58,7 +116,26 @@ fn downstream_store_intent_should_not_enable_upstream_storage() { let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); - assert_eq!(encoded.body().get("store"), Some(&json!(false))); + assert_eq!(encoded.body().get("store"), Some(&json!(true))); +} + +#[test] +fn encoder_should_forward_the_common_prompt_cache_key() { + let message = Message::new( + MessageRole::User, + vec![ContentPart::Text("cache prefix".to_owned())], + ) + .expect("message"); + let request = GenerateRequest::new(vec![message]) + .expect("generate") + .with_prompt_cache_key("cache-route"); + + let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); + + assert_eq!( + encoded.body().get("prompt_cache_key"), + Some(&json!("cache-route")) + ); } #[test] @@ -85,11 +162,38 @@ fn encoder_should_project_explicit_websocket_transport_without_touching_body() { let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); assert!(encoded.use_websocket); - assert!(encoded.force_websocket); assert!(!encoded.force_http_sse); assert!(encoded.body().get("transport").is_none()); } +#[test] +fn encoder_should_project_explicit_http_transport_without_touching_body() { + let message = Message::new( + MessageRole::User, + vec![ContentPart::Text("prompt".to_owned())], + ) + .expect("message"); + let mut providers = ProviderOptions::new(); + providers + .insert( + "openai", + Map::from_iter([ + ("schema_version".to_owned(), json!(1)), + ("transport".to_owned(), json!("http_sse")), + ]), + ) + .expect("provider options"); + let request = GenerateRequest::new(vec![message]) + .expect("generate") + .with_provider_options(providers); + + let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); + + assert!(!encoded.use_websocket); + assert!(encoded.force_http_sse); + assert!(encoded.body().get("transport").is_none()); +} + #[test] fn encoder_should_reject_unknown_codex_options_without_echoing_values() { let message = Message::new( @@ -119,11 +223,21 @@ fn encoder_should_reject_unknown_codex_options_without_echoing_values() { #[test] fn encoder_should_project_lite_and_memgen_options_to_transport_state_only() { - let message = Message::new( - MessageRole::User, - vec![ContentPart::Text("prompt".to_owned())], + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("input".to_owned(), json!("prompt")), + ( + "client_metadata".to_owned(), + json!({ + "ws_request_header_x_openai_internal_codex_responses_lite": "false", + "x-openai-memgen-request": "false" + }), + ), + ]), ) - .expect("message"); + .expect("OpenAI payload"); let mut providers = ProviderOptions::new(); providers .insert( @@ -135,8 +249,7 @@ fn encoder_should_project_lite_and_memgen_options_to_transport_state_only() { ]), ) .expect("provider options"); - let request = GenerateRequest::new(vec![message]) - .expect("generate") + let request = GenerateRequest::from_protocol_payload(Vec::new(), payload) .with_provider_options(providers); let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); @@ -177,3 +290,35 @@ fn observability_semantics_should_reuse_codex_turn_metadata() { assert_eq!(semantics.subagent_kind.as_deref(), Some("review")); assert!(semantics.compact); } + +#[test] +fn observability_semantics_should_use_the_transparent_openai_payload() { + let payload = ProtocolPayload::json_object( + "openai", + json!({ + "model": "gpt-test", + "reasoning": {"effort": "max"}, + "input": [ + { + "type": "message", + "role": "developer", + "content": [{ + "type": "input_text", + "text": "Proactive multi-agent delegation is active." + }] + }, + {"type": "compaction_trigger"} + ] + }) + .as_object() + .expect("request object") + .clone(), + ) + .expect("OpenAI payload"); + let request = GenerateRequest::from_protocol_payload(Vec::new(), payload); + + let semantics = codex_request_semantics(&request); + + assert_eq!(semantics.reasoning_preset, Some("ultra")); + assert!(semantics.compact); +} diff --git a/backend/crates/providers/openai/tests/transport/usage.rs b/backend/crates/providers/openai/tests/transport/usage.rs index 526080ee5..cf5d32f73 100644 --- a/backend/crates/providers/openai/tests/transport/usage.rs +++ b/backend/crates/providers/openai/tests/transport/usage.rs @@ -128,12 +128,6 @@ async fn fetch_should_use_wham_usage_headers_only() { ("originator", "codex_cli_rs"), ("accept", "application/json"), ("cookie", "session=old"), - ("openai-beta", "codex-1"), - ("oai-language", "zh-CN"), - ("sec-fetch-site", "none"), - ("sec-fetch-mode", "no-cors"), - ("sec-fetch-dest", "empty"), - ("priority", "u=4, i"), ] { assert_eq!( headers.get(name).and_then(|value| value.to_str().ok()), @@ -154,13 +148,7 @@ async fn fetch_should_use_wham_usage_headers_only() { "authorization", "chatgpt-account-id", "cookie", - "oai-language", - "openai-beta", "originator", - "priority", - "sec-fetch-dest", - "sec-fetch-mode", - "sec-fetch-site", "user-agent", ]) ); diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index b02dc9e52..240d50c9e 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -70,6 +70,67 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn exact_continuation_reuses_owning_socket_after_connection_profile_updates() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted = Arc::new(AtomicUsize::new(0)); + let server_accepted = Arc::clone(&accepted); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + server_accepted.fetch_add(1, Ordering::SeqCst); + let mut websocket = accept_codex_test_websocket(stream).await; + for response_id in ["resp_connection_local", "resp_continued_after_profile"] { + let _request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response(response_id, 2, 1).into(), + )) + .await + .unwrap(); + } + websocket.close(None).await.unwrap(); + }); + let profile = test_wire_profile(); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + profile.clone(), + ) + .with_websocket_pool(Arc::new(CodexWebSocketPool::new(1, Duration::from_mins(1)))); + let mut first_request = pooled_websocket_request("wire-profile-continuation"); + first_request.set_store(false); + let first = backend + .create_response( + &first_request, + CodexRequestContext { + cookie_header: Some("cf_clearance=before-handshake"), + ..request_context("req_local_first", Some("chatgpt-account")) + }, + ) + .await + .expect("connection-local response"); + profile.update_desktop_release("2.0.0", "200"); + let mut continuation = first_request; + continuation.set_previous_response_id(Some("resp_connection_local".to_owned())); + continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + let second = backend + .create_response( + &continuation, + CodexRequestContext { + cookie_header: Some("cf_clearance=after-handshake"), + ..request_context("req_local_second", Some("chatgpt-account")) + }, + ) + .await + .expect("exact continuation must keep the owning socket"); + server.await.unwrap(); + + assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new"); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "reuse"); + assert_eq!(accepted.load(Ordering::SeqCst), 1); +} + /// idle 连接被上游静默关闭后，后台 pump 会实时把它标记为 closed。 /// 复用前的零成本 `is_closed` 检查应直接丢弃它并新建连接，不经过 /// “发请求 → 等首帧超时 → stale-reuse 重试” 的长尾（无需任何 maintenance sweep）。 @@ -535,6 +596,94 @@ async fn websocket_pool_should_use_http_while_exact_key_is_busy() { ); } +#[tokio::test] +async fn account_eviction_should_let_busy_stream_finish_without_returning_it_to_the_pool() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (finish_first_tx, finish_first_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = accept_codex_test_websocket(first_stream).await; + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "busy stream remains alive" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + finish_first_rx.await.unwrap(); + first_websocket + .send(Message::Text( + completed_websocket_response("resp_busy_evicted", 2, 1).into(), + )) + .await + .unwrap(); + let close = timeout(Duration::from_secs(2), first_websocket.next()) + .await + .expect("completed evicted stream should be closed instead of pooled") + .expect("close frame") + .expect("valid close frame"); + std::assert_matches!(close, Message::Close(_)); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_websocket = accept_codex_test_websocket(second_stream).await; + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_after_busy_eviction", 3, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::default()); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + test_wire_profile(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let request = pooled_websocket_request("conversation-busy-eviction"); + let mut first = backend + .create_response_stream( + &request, + request_context("req_busy_eviction", Some("chatgpt-account")), + ) + .await + .expect("busy websocket stream") + .body; + let first_chunk = first + .next() + .await + .expect("first stream event") + .expect("valid first stream event"); + assert!( + std::str::from_utf8(&first_chunk) + .unwrap() + .contains("busy stream remains alive") + ); + + pool.evict_account("chatgpt-account").await; + finish_first_tx.send(()).unwrap(); + while first.next().await.transpose().unwrap().is_some() {} + let second = backend + .create_response( + &request, + request_context("req_after_busy_eviction", Some("chatgpt-account")), + ) + .await + .expect("request after eviction should open a fresh websocket"); + server.await.unwrap(); + + assert!(second.body.contains("resp_after_busy_eviction")); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); +} + #[tokio::test] async fn websocket_pool_should_release_slot_when_client_drops_stream() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/crates/providers/xai/Cargo.toml b/backend/crates/providers/xai/Cargo.toml index 5b21af23c..f02619509 100644 --- a/backend/crates/providers/xai/Cargo.toml +++ b/backend/crates/providers/xai/Cargo.toml @@ -13,6 +13,7 @@ async-trait = "0.1.89" base64 = "0.22.1" chrono = { version = "0.4.45", features = ["clock", "serde"] } futures = "0.3.32" +gateway-admin = { path = "../../gateway-admin" } gateway-core = { path = "../../gateway-core" } gateway-protocol = { path = "../../gateway-protocol" } getrandom = "0.3.4" diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs new file mode 100644 index 000000000..616481224 --- /dev/null +++ b/backend/crates/providers/xai/src/admin.rs @@ -0,0 +1,1314 @@ +//! xAI 管理边界：Provider preparation 与 Redis OAuth pending 适配。 + +use std::collections::BTreeSet; +use std::sync::Arc; +use std::time::SystemTime; + +use async_trait::async_trait; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, TimeDelta, Utc}; +use gateway_admin::model::accounts::{ + AccountAvailability as AdminAccountAvailability, AccountRecord, +}; +use gateway_admin::model::observability::{ + CalculatedBillingBreakdown, CurrencyCost, DecimalAmount, ProviderBillingInput, +}; +use gateway_admin::model::provider_credentials::{ + AuthorizationMutationTarget, AuthorizationOwner, AuthorizationOwnerBinding, + AuthorizationStarted, CompleteAuthorization, CredentialCommitGuard, + PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedAuthorizationCredential, + PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation, + PreparedCredentialRotationFacts, ProviderDocument, ProviderExport, + ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota, + ProviderQuotaWindow, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision}; +use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind}; +use gateway_core::accounting::Money; +use gateway_core::engine::credential::{ + AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount, + OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId, + ProviderAccountStore, +}; +use gateway_core::error::StoreErrorKind; +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome, + OAuthPendingTakeOutcome, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, + ProviderStoreErrorKind, +}; +use gateway_core::routing::{ProviderInstanceId, ProviderKind, UpstreamModelId}; +use serde::Deserialize; +use serde_json::{Map, Number, Value}; +use sha2::{Digest as _, Sha256}; +use url::Url; +use uuid::Uuid; + +use crate::credential::{ + AuthorizationCallback, FailureClass, GrokAccountProfile, GrokCredentialAdmin, + GrokCredentialCatalogError, GrokCredentialCatalogService, GrokCredentialQuotaService, + GrokCredentialRefreshError, GrokCredentialRefreshService, GrokCredentialRepository, + GrokCredentialRepositoryError, GrokOAuthClient, GrokOAuthConfig, GrokOAuthImportCandidate, + GrokOAuthImportDocument, GrokOAuthImportMetadata, GrokOAuthImportTokens, GrokOAuthSecret, + GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, OAuthError, PendingAuthorization, + PreparedGrokCredentialRotation, PreparedGrokCredentialRotationGuard, RedirectUriAllowlist, + RotateManagedGrokCredential, SecretValue, VerifiedGrokAccount, VerifiedTokenSet, +}; +use crate::transport::{ + GROK_CLI_BASE_URL, GROK_CLIENT_VERSION, XAI_PROVIDER_NAME, grok_billing_breakdown, +}; + +const PENDING_SCHEMA_VERSION: u64 = 1; +const PENDING_TTL: TimeDelta = TimeDelta::minutes(10); +const MAX_PENDING_TEXT_BYTES: usize = 512; + +pub(crate) struct XaiAdminProvider { + provider_kind: ProviderKind, + accounts: Arc, + repository: GrokCredentialRepository, + oauth_config: GrokOAuthConfig, + oauth: Arc, + pending: Arc, + refresh: Arc, + quota: Arc, + catalog: Arc, + runtime_policy: Arc, +} + +pub(crate) struct XaiAdminServices { + pub(crate) repository: GrokCredentialRepository, + pub(crate) oauth_config: GrokOAuthConfig, + pub(crate) oauth: Arc, + pub(crate) pending: Arc, + pub(crate) refresh: Arc, + pub(crate) quota: Arc, + pub(crate) catalog: Arc, + pub(crate) runtime_policy: Arc, +} + +impl XaiAdminProvider { + #[must_use] + pub(crate) fn new( + provider_kind: ProviderKind, + accounts: Arc, + services: XaiAdminServices, + ) -> Self { + Self { + provider_kind, + accounts, + repository: services.repository, + oauth_config: services.oauth_config, + oauth: services.oauth, + pending: services.pending, + refresh: services.refresh, + quota: services.quota, + catalog: services.catalog, + runtime_policy: services.runtime_policy, + } + } + + async fn account( + &self, + account_id: &ProviderAccountId, + ) -> Result { + self.accounts + .get_account(account_id) + .await + .map_err(map_store_error)? + .filter(|account| account.provider() == &self.provider_kind) + .ok_or_else(|| provider_error(ProviderAdminErrorKind::NotFound)) + } + + async fn store_pending( + &self, + mutation: PendingAuthorizationMutation, + pending: PendingAuthorization, + ) -> Result { + let authorization_url = pending.authorization_url().to_string(); + let server_state = pending.into_server_state().map_err(map_oauth_error)?; + let flow_id = random_flow_id()?; + let expires_at = Utc::now() + .checked_add_signed(PENDING_TTL) + .ok_or_else(|| provider_error(ProviderAdminErrorKind::Internal))?; + let owner_ref = owner_ref(mutation.owner_binding().owner()); + let ttl = (expires_at - Utc::now()) + .to_std() + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?; + let payload = encode_pending(&StoredXaiAuthorization { + flow_id: flow_id.clone(), + owner_ref: owner_ref.clone(), + expires_at, + server_state: server_state.expose().to_owned(), + mutation, + }); + let flow = NewOAuthPendingFlow::try_new( + self.provider_kind.clone(), + binding(&flow_id)?, + binding(&owner_ref)?, + ttl, + OpaqueProviderData::new(payload), + ) + .map_err(map_provider_store_error)?; + match self + .pending + .put_if_absent(flow) + .await + .map_err(map_provider_store_error)? + { + OAuthPendingPutOutcome::Stored => Ok(AuthorizationStarted { + flow_id, + authorization_url, + expires_at, + }), + OAuthPendingPutOutcome::AlreadyExists => { + Err(provider_error(ProviderAdminErrorKind::Conflict)) + } + } + } + + async fn take_pending( + &self, + context: &MutationContext, + flow_id: &str, + ) -> Result { + let flow = binding(flow_id)?; + let owner_binding = AuthorizationOwnerBinding::from_context(context); + let owner_ref = owner_ref(owner_binding.owner()); + let owner = binding(&owner_ref)?; + match self + .pending + .take_if_owner(&self.provider_kind, &flow, &owner) + .await + .map_err(map_provider_store_error)? + { + OAuthPendingTakeOutcome::Taken(payload) => { + let stored = decode_pending(payload, &self.oauth_config)?; + if stored.flow_id != flow_id || stored.owner_ref != owner_ref { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + Ok(stored) + } + OAuthPendingTakeOutcome::NotFound | OAuthPendingTakeOutcome::OwnerMismatch => { + Err(provider_error(ProviderAdminErrorKind::NotFound)) + } + } + } +} + +#[async_trait] +impl ProviderAdmin for XaiAdminProvider { + fn provider_kind(&self) -> &ProviderKind { + &self.provider_kind + } + + async fn account_unavailable(&self, _account_id: &ProviderAccountId) {} + + fn dashboard_wire_profile( + &self, + ) -> Option { + None + } + + fn calculated_billing( + &self, + input: &ProviderBillingInput, + ) -> Result, ProviderAdminError> { + let (Some(input_tokens), Some(output_tokens)) = (input.input_tokens, input.output_tokens) + else { + return Ok(None); + }; + let Some(breakdown) = grok_billing_breakdown( + &input.upstream_model_id, + input_tokens, + output_tokens, + input.cached_tokens.unwrap_or_default(), + ) else { + return Ok(None); + }; + let total_amount = currency_cost(breakdown.total_amount())?; + if total_amount != input.total { + return Ok(None); + } + Ok(Some(CalculatedBillingBreakdown { + input_amount: currency_cost(breakdown.input_amount())?, + output_amount: currency_cost(breakdown.output_amount())?, + cache_read_amount: currency_cost(breakdown.cache_read_amount())?, + cache_write_amount: currency_cost(breakdown.cache_write_amount())?, + standard_amount: currency_cost(breakdown.standard_amount())?, + total_amount, + input_price_per_million: currency_cost(breakdown.input_price_per_million())?, + output_price_per_million: currency_cost(breakdown.output_price_per_million())?, + cache_read_price_per_million: currency_cost(breakdown.cache_read_price_per_million())?, + cache_write_price_per_million: currency_cost( + breakdown.cache_write_price_per_million(), + )?, + service_tier: breakdown.service_tier().map(str::to_owned), + multiplier_percent: breakdown.multiplier_percent(), + })) + } + + async fn prepare_import( + &self, + command: PrepareCredentialImport, + ) -> Result { + let document = serde_json::to_vec(&Value::Object( + command.document.into_provider_data().into_inner(), + )) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let document = GrokOAuthImportDocument::parse_json(&document) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(map_provider_store_error)?; + let mut subjects = BTreeSet::new(); + let mut credentials = Vec::new(); + for entry in document.into_entries() { + let name = entry.name().to_owned(); + let email = entry.email().map(str::to_owned); + let tokens = self + .oauth + .verify_imported_credential(&discovery, entry.into_candidate()) + .await + .map_err(|error| map_failure_class(error.class()))?; + if !subjects.insert(tokens.evidence().subject().to_owned()) { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + let prepared = GrokCredentialAdmin + .prepare_verified_account( + &VerifiedGrokAccount { + account_id: ProviderAccountId::new(format!( + "acct_{}", + Uuid::now_v7().simple() + )) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + provider_instance_id: command.provider_instance_id.clone(), + name, + email, + upstream_account_id: None, + plan_type: None, + tokens, + enabled: true, + }, + policy, + ) + .map_err(map_repository_error)?; + credentials.push(prepared_create(prepared, Utc::now())?); + } + Ok(PreparedCredentialImport { + provider_kind: self.provider_kind.clone(), + provider_instance_id: command.provider_instance_id, + credentials, + }) + } + + async fn start_authorization( + &self, + pending: PendingAuthorizationMutation, + ) -> Result { + if pending.provider_kind() != &self.provider_kind { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + if let AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } = pending.target() + { + let revision = CredentialRevision::new(expected_credential_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let current = self + .accounts + .load_credential(account_id, revision) + .await + .map_err(map_store_error)?; + if current.account.provider() != &self.provider_kind + || current.account.instance() != provider_instance_id + { + return Err(provider_error(ProviderAdminErrorKind::NotFound)); + } + } + let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; + let redirect = RedirectUriAllowlist::new([crate::OFFICIAL_REDIRECT_URI]) + .and_then(|allowlist| allowlist.authorize(crate::OFFICIAL_REDIRECT_URI)) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?; + let authorization = self + .oauth + .start_authorization_code(&discovery, redirect, None) + .map_err(map_oauth_error)?; + self.store_pending(pending, authorization).await + } + + async fn complete_authorization( + &self, + command: CompleteAuthorization, + ) -> Result { + let stored = self + .take_pending(&command.context, &command.flow_id) + .await?; + let authorization = PendingAuthorization::from_server_state( + &self.oauth_config, + &SecretValue::new(stored.server_state), + ) + .map_err(map_oauth_error)?; + let callback = callback(&command.callback_url)?; + let grant = authorization + .accept_callback(callback) + .map_err(map_oauth_error)?; + let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; + let tokens = self + .oauth + .exchange_authorization_code(&discovery, grant) + .await + .map_err(map_oauth_error)?; + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(map_provider_store_error)?; + let credential = match stored.mutation.target() { + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + } => { + let prepared = GrokCredentialAdmin + .prepare_verified_account( + &VerifiedGrokAccount { + account_id: ProviderAccountId::new(format!( + "acct_{}", + Uuid::now_v7().simple() + )) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + provider_instance_id: provider_instance_id.clone(), + name: name.clone(), + email: None, + upstream_account_id: None, + plan_type: None, + tokens, + enabled: true, + }, + policy, + ) + .map_err(map_repository_error)?; + PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?) + } + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => { + let revision = CredentialRevision::new(expected_credential_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?; + let current = self + .accounts + .load_credential(account_id, revision) + .await + .map_err(map_store_error)?; + if current.account.provider() != &self.provider_kind + || current.account.instance() != provider_instance_id + { + return Err(provider_error(ProviderAdminErrorKind::NotFound)); + } + let prepared = verified_rotation(current, tokens, policy)?; + PreparedAuthorizationCredential::Reauthorize(prepared_rotation( + prepared, + provider_instance_id.clone(), + self.provider_kind.clone(), + )?) + } + }; + Ok(PreparedAuthorizationCommit { + pending: stored.mutation, + credential, + }) + } + + async fn prepare_rotation( + &self, + command: PrepareCredentialRotation, + ) -> Result { + validate_account_record(&command.account, &self.provider_kind)?; + if command.account.credential_revision != command.expected_credential_revision { + return Err(provider_error(ProviderAdminErrorKind::Conflict)); + } + let account_id = ProviderAccountId::new(command.account.id.clone()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let revision = CredentialRevision::new(command.expected_credential_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let current = self + .accounts + .load_credential(&account_id, revision) + .await + .map_err(map_store_error)?; + if !account_matches_record(¤t.account, &command.account) { + return Err(provider_error(ProviderAdminErrorKind::Conflict)); + } + let rotation: RotationDocument = serde_json::from_value(Value::Object( + command.provider_material.into_provider_data().into_inner(), + )) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let tokens = match rotation.id_token { + Some(id_token) => GrokOAuthImportTokens::new( + SecretValue::new(rotation.access_token), + SecretValue::new(rotation.refresh_token), + SecretValue::new(id_token), + ), + None => GrokOAuthImportTokens::without_id_token( + SecretValue::new(rotation.access_token), + SecretValue::new(rotation.refresh_token), + ), + }; + let candidate = GrokOAuthImportCandidate::new( + tokens, + GrokOAuthImportMetadata::new( + "Bearer".to_owned(), + crate::OFFICIAL_CLIENT_ID.to_owned(), + rotation.scope, + GROK_CLI_BASE_URL.to_owned(), + Utc::now(), + rotation.expires_at, + ), + ); + let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; + let tokens = self + .oauth + .verify_imported_credential(&discovery, candidate) + .await + .map_err(|error| map_failure_class(error.class()))?; + let policy = self + .runtime_policy + .load_refresh_policy() + .await + .map_err(map_provider_store_error)?; + let prepared = verified_rotation(current, tokens, policy)?; + prepared_rotation( + prepared, + command.account.provider_instance_id, + command.account.provider_kind, + ) + } + + async fn prepare_refresh( + &self, + command: PrepareCredentialRefresh, + ) -> Result { + validate_account_record(&command.account, &self.provider_kind)?; + let account_id = ProviderAccountId::new(command.account.id.clone()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let current = self.account(&account_id).await?; + if !account_matches_record(¤t, &command.account) { + return Err(provider_error(ProviderAdminErrorKind::Conflict)); + } + let revision = CredentialRevision::new(command.account.credential_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let prepared = self + .refresh + .prepare_manual_refresh(&account_id, revision) + .await + .map_err(map_refresh_error)?; + prepared_rotation( + prepared, + command.account.provider_instance_id, + command.account.provider_kind, + ) + } + + async fn quota( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + self.account(account_id).await?; + let lifecycle = self + .repository + .read_lifecycle(account_id) + .await + .map_err(map_repository_error)?; + let snapshot = if refresh { + Some( + self.quota + .refresh_account(account_id) + .await + .map_err(map_quota_error)?, + ) + } else { + self.quota + .read_account(account_id) + .await + .map_err(map_quota_error)? + }; + Ok(project_quota( + snapshot, + lifecycle.refresh_token_expires_at().copied(), + )) + } + + async fn models( + &self, + account_id: &ProviderAccountId, + refresh: bool, + ) -> Result { + let account = self.account(account_id).await?; + let catalog = if refresh { + Some( + self.catalog + .refresh_account_catalog(account_id, GROK_CLIENT_VERSION) + .await + .map_err(map_catalog_error)?, + ) + } else { + self.catalog + .read_account_catalog(account_id, account.revision()) + .await + .map_err(map_catalog_error)? + }; + let Some(catalog) = catalog else { + return Ok(ProviderModels { + models: Vec::new(), + observed_at: None, + }); + }; + let models = catalog + .seed() + .models() + .iter() + .map(|model| { + Ok(ProviderModel { + id: UpstreamModelId::new(model.clone()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + name: model.clone(), + }) + }) + .collect::, ProviderAdminError>>()?; + Ok(ProviderModels { + models, + observed_at: Some(catalog.observed_at()), + }) + } + + async fn export_credentials( + &self, + credentials: Vec, + ) -> Result { + let mut account_ids = Vec::with_capacity(credentials.len()); + let mut loaded = Vec::with_capacity(credentials.len()); + for input in credentials { + validate_account_record(&input.account, &self.provider_kind)?; + let current = LoadedCredential { + account: account_from_record(&input.account)?, + credential: PlaintextCredential::new( + input.provider_material.into_provider_data().into_inner(), + ), + }; + account_ids.push(current.account.id().clone()); + loaded.push(current); + } + let document = GrokCredentialAdmin + .export_oauth_bundle(&loaded, Utc::now()) + .map_err(map_repository_error)? + .into_value(); + let Value::Object(document) = document else { + return Err(provider_error(ProviderAdminErrorKind::Internal)); + }; + Ok(ProviderExport { + provider_kind: self.provider_kind.clone(), + account_ids, + document: ProviderDocument::new(OpaqueProviderData::new(document)), + }) + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RotationDocument { + access_token: String, + refresh_token: String, + id_token: Option, + scope: String, + expires_at: DateTime, +} + +struct XaiCredentialCommitGuard { + _guard: PreparedGrokCredentialRotationGuard, +} + +impl CredentialCommitGuard for XaiCredentialCommitGuard { + fn finish(self: Box) {} +} + +fn verified_rotation( + current: LoadedCredential, + tokens: VerifiedTokenSet, + policy: ProviderRefreshPolicy, +) -> Result { + let expires_in = tokens + .expires_in() + .ok_or_else(|| provider_error(ProviderAdminErrorKind::Invalid))?; + let refresh_token = tokens + .refresh_token() + .cloned() + .ok_or_else(|| provider_error(ProviderAdminErrorKind::Invalid))?; + let now = SystemTime::now(); + let access_token_expires_at = now + .checked_add(expires_in) + .ok_or_else(|| provider_error(ProviderAdminErrorKind::Invalid))?; + let next_refresh_at = policy + .next_attempt_at(current.account.id(), access_token_expires_at, now) + .map_err(map_provider_store_error)?; + GrokCredentialAdmin + .prepare_rotation(&RotateManagedGrokCredential { + secret: GrokOAuthSecret { + access_token: tokens.access_token().clone(), + refresh_token, + id_token: tokens.id_token().cloned(), + scope: tokens.scope().to_owned(), + }, + verified_account: GrokAccountProfile { + subject: tokens.evidence().subject().to_owned(), + email: current.account.email().map(str::to_owned), + upstream_account_id: current.account.upstream_account_id().map(str::to_owned), + plan_type: current.account.plan_type().map(str::to_owned), + access_token_expires_at: access_token_expires_at.into(), + refresh_token_expires_at: None, + }, + current, + next_refresh_at: next_refresh_at.into(), + }) + .map_err(map_repository_error) +} + +fn prepared_create( + prepared: NewProviderAccount, + observed_at: DateTime, +) -> Result { + let NewProviderAccount { + account, + credential, + } = prepared; + Ok(PreparedCredentialCreate { + account_id: account.id().clone(), + provider_instance_id: account.instance().clone(), + provider_kind: account.provider().clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + upstream_user_id: account.upstream_user_id().to_owned(), + upstream_account_id: account.upstream_account_id().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + provider_material: ProviderDocument::new(OpaqueProviderData::new(credential.into_inner())), + has_refresh_token: account.has_refresh_token(), + access_token_expires_at: account.access_token_expires_at().into(), + next_refresh_at: account.next_refresh_at().map(Into::into), + enabled: account.enabled(), + availability: admin_availability(account.availability()), + availability_reason: None, + cooldown_until: account.cooldown_until().map(Into::into), + availability_observed_at: observed_at, + }) +} + +fn prepared_rotation( + prepared: PreparedGrokCredentialRotation, + provider_instance_id: ProviderInstanceId, + provider_kind: ProviderKind, +) -> Result { + let (profile, credential, guard) = prepared.into_parts(); + let ( + account_id, + expected_revision, + credential_profile, + credential, + has_refresh_token, + access_token_expires_at, + next_refresh_at, + ) = credential.into_parts(); + if profile != credential_profile || profile.account_id != account_id { + return Err(provider_error(ProviderAdminErrorKind::Internal)); + } + Ok(PreparedCredentialRotation::new( + PreparedCredentialRotationFacts { + account_id, + provider_instance_id, + provider_kind, + expected_credential_revision: Revision::new(expected_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + name: profile.name, + email: profile.email, + plan_type: profile.plan_type, + provider_material: ProviderDocument::new(OpaqueProviderData::new( + credential.into_inner(), + )), + has_refresh_token, + access_token_expires_at: access_token_expires_at.into(), + next_refresh_at: next_refresh_at.map(Into::into), + }, + Box::new(XaiCredentialCommitGuard { _guard: guard }), + )) +} + +fn validate_account_record( + account: &AccountRecord, + provider_kind: &ProviderKind, +) -> Result<(), ProviderAdminError> { + if &account.provider_kind != provider_kind || provider_kind.as_str() != XAI_PROVIDER_NAME { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + ProviderAccountId::new(account.id.clone()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + Ok(()) +} + +fn account_from_record(account: &AccountRecord) -> Result { + let account_id = ProviderAccountId::new(account.id.clone()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let revision = CredentialRevision::new(account.credential_revision.get()) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + Ok(ProviderAccount::new( + account_id, + account.provider_instance_id.clone(), + account.provider_kind.clone(), + account.name.clone(), + account.upstream_user_id.clone(), + revision, + account.access_token_expires_at.into(), + ) + .with_profile( + account.email.clone(), + account.upstream_account_id.clone(), + account.plan_type.clone(), + ) + .with_runtime_state( + account.enabled, + core_availability(account.availability), + account.cooldown_until.map(Into::into), + ) + .with_refresh_schedule( + account.has_refresh_token, + account.next_refresh_at.map(Into::into), + )) +} + +fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool { + account.id().as_str() == record.id + && account.instance() == &record.provider_instance_id + && account.provider() == &record.provider_kind + && account.revision().get() == record.credential_revision.get() + && account.name() == record.name + && account.email() == record.email.as_deref() + && account.upstream_user_id() == record.upstream_user_id + && account.upstream_account_id() == record.upstream_account_id.as_deref() + && account.plan_type() == record.plan_type.as_deref() + && account.enabled() == record.enabled + && admin_availability(account.availability()) == record.availability + && account.cooldown_until().map(DateTime::::from) == record.cooldown_until + && DateTime::::from(account.access_token_expires_at()) + == record.access_token_expires_at + && account.next_refresh_at().map(DateTime::::from) == record.next_refresh_at + && account.has_refresh_token() == record.has_refresh_token +} + +const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability { + match value { + AccountAvailability::Unknown => AdminAccountAvailability::Unknown, + AccountAvailability::Ready => AdminAccountAvailability::Ready, + AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown, + AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted, + AccountAvailability::Expired => AdminAccountAvailability::Expired, + AccountAvailability::Banned => AdminAccountAvailability::Banned, + AccountAvailability::Invalid => AdminAccountAvailability::Invalid, + } +} + +const fn core_availability(value: AdminAccountAvailability) -> AccountAvailability { + match value { + AdminAccountAvailability::Unknown => AccountAvailability::Unknown, + AdminAccountAvailability::Ready => AccountAvailability::Ready, + AdminAccountAvailability::Cooldown => AccountAvailability::Cooldown, + AdminAccountAvailability::QuotaExhausted => AccountAvailability::QuotaExhausted, + AdminAccountAvailability::Expired => AccountAvailability::Expired, + AdminAccountAvailability::Banned => AccountAvailability::Banned, + AdminAccountAvailability::Invalid => AccountAvailability::Invalid, + } +} + +fn project_quota( + snapshot: Option, + refresh_token_expires_at: Option>, +) -> ProviderQuota { + let Some(snapshot) = snapshot else { + return ProviderQuota { + observed_at: None, + refresh_token_expires_at, + windows: Vec::new(), + provider_data: None, + }; + }; + let billing = snapshot.billing(); + let window = if billing.has_authoritative_quota() { + let mut data = Map::new(); + data.insert( + "periodStart".to_owned(), + billing + .period_start() + .map_or(Value::Null, |value| Value::String(value.to_owned())), + ); + data.insert( + "periodEnd".to_owned(), + billing + .period_end() + .map_or(Value::Null, |value| Value::String(value.to_owned())), + ); + for (key, value) in [ + ("monthlyLimitCents", billing.monthly_limit_cents()), + ("includedUsedCents", billing.included_used_cents()), + ("onDemandCapCents", billing.on_demand_cap_cents()), + ("onDemandUsedCents", billing.on_demand_used_cents()), + ("prepaidBalanceCents", billing.prepaid_balance_cents()), + ] { + data.insert( + key.to_owned(), + value.map_or(Value::Null, |value| Value::Number(Number::from(value))), + ); + } + ProviderQuotaWindow { + key: billing.period_type().unwrap_or("billing").to_owned(), + group: quota_group(billing.period_kind()).to_owned(), + source: None, + window_seconds: quota_window_seconds(billing.period_start(), billing.period_end()), + used_percent: billing.used_percent(), + reset_at: billing.period_end().and_then(parse_utc), + provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))), + } + } else { + ProviderQuotaWindow { + key: "free-rolling-24h".to_owned(), + group: "shortTerm".to_owned(), + source: None, + window_seconds: Some(crate::GROK_FREE_ROLLING_WINDOW_SECONDS), + used_percent: None, + reset_at: None, + provider_data: None, + } + }; + ProviderQuota { + observed_at: Some(snapshot.observed_at()), + refresh_token_expires_at, + windows: vec![window], + provider_data: None, + } +} + +const fn quota_group(kind: GrokQuotaPeriodKind) -> &'static str { + match kind { + GrokQuotaPeriodKind::Weekly => "shortTerm", + GrokQuotaPeriodKind::Monthly => "monthly", + GrokQuotaPeriodKind::Other => "other", + } +} + +fn quota_window_seconds(start: Option<&str>, end: Option<&str>) -> Option { + let start = start.and_then(parse_utc)?; + let end = end.and_then(parse_utc)?; + end.signed_duration_since(start) + .num_seconds() + .try_into() + .ok() +} + +fn parse_utc(value: &str) -> Option> { + DateTime::parse_from_rfc3339(value) + .ok() + .map(|value| value.with_timezone(&Utc)) +} + +struct StoredXaiAuthorization { + flow_id: String, + owner_ref: String, + expires_at: DateTime, + server_state: String, + mutation: PendingAuthorizationMutation, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PendingDocument { + schema_version: u64, + flow_id: String, + owner_ref: String, + expires_at: DateTime, + server_state: String, + mutation: PendingMutationDocument, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct PendingMutationDocument { + expected_config_revision: u64, + provider_kind: String, + target: PendingTargetDocument, + owner: PendingOwnerDocument, + started_request_id: String, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum PendingTargetDocument { + Create { + provider_instance_id: String, + name: String, + }, + Reauthorize { + provider_instance_id: String, + account_id: String, + expected_credential_revision: u64, + }, +} + +#[derive(Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum PendingOwnerDocument { + AdminSession { admin_user_id: String }, + AdminApiKey, + System, +} + +fn encode_pending(pending: &StoredXaiAuthorization) -> Map { + let mut document = Map::new(); + document.insert( + "schema_version".to_owned(), + Value::Number(Number::from(PENDING_SCHEMA_VERSION)), + ); + document.insert("flow_id".to_owned(), Value::String(pending.flow_id.clone())); + document.insert( + "owner_ref".to_owned(), + Value::String(pending.owner_ref.clone()), + ); + document.insert( + "expires_at".to_owned(), + Value::String(pending.expires_at.to_rfc3339()), + ); + document.insert( + "server_state".to_owned(), + Value::String(pending.server_state.clone()), + ); + document.insert( + "mutation".to_owned(), + Value::Object(encode_mutation(&pending.mutation)), + ); + document +} + +fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map { + let mut document = Map::new(); + document.insert( + "expected_config_revision".to_owned(), + Value::Number(Number::from(mutation.expected_config_revision().get())), + ); + document.insert( + "provider_kind".to_owned(), + Value::String(mutation.provider_kind().as_str().to_owned()), + ); + document.insert( + "target".to_owned(), + Value::Object(encode_target(mutation.target())), + ); + document.insert( + "owner".to_owned(), + Value::Object(encode_owner(mutation.owner_binding().owner())), + ); + document.insert( + "started_request_id".to_owned(), + Value::String(mutation.owner_binding().started_request_id().to_owned()), + ); + document +} + +fn encode_target(target: &AuthorizationMutationTarget) -> Map { + let mut document = Map::new(); + match target { + AuthorizationMutationTarget::Create { + provider_instance_id, + name, + } => { + document.insert("kind".to_owned(), Value::String("create".to_owned())); + document.insert( + "provider_instance_id".to_owned(), + Value::String(provider_instance_id.to_string()), + ); + document.insert("name".to_owned(), Value::String(name.clone())); + } + AuthorizationMutationTarget::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => { + document.insert("kind".to_owned(), Value::String("reauthorize".to_owned())); + document.insert( + "provider_instance_id".to_owned(), + Value::String(provider_instance_id.to_string()), + ); + document.insert( + "account_id".to_owned(), + Value::String(account_id.to_string()), + ); + document.insert( + "expected_credential_revision".to_owned(), + Value::Number(Number::from(expected_credential_revision.get())), + ); + } + } + document +} + +fn encode_owner(owner: &AuthorizationOwner) -> Map { + let mut document = Map::new(); + match owner { + AuthorizationOwner::AdminSession { admin_user_id } => { + document.insert("kind".to_owned(), Value::String("admin_session".to_owned())); + document.insert( + "admin_user_id".to_owned(), + Value::String(admin_user_id.clone()), + ); + } + AuthorizationOwner::AdminApiKey => { + document.insert("kind".to_owned(), Value::String("admin_api_key".to_owned())); + } + AuthorizationOwner::System => { + document.insert("kind".to_owned(), Value::String("system".to_owned())); + } + } + document +} + +fn decode_pending( + payload: OpaqueProviderData, + oauth_config: &GrokOAuthConfig, +) -> Result { + let document: PendingDocument = serde_json::from_value(Value::Object(payload.into_inner())) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + if document.schema_version != PENDING_SCHEMA_VERSION + || !valid_pending_text(&document.flow_id) + || !valid_pending_text(&document.owner_ref) + || document.expires_at <= Utc::now() + { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + let mutation = decode_mutation(document.mutation)?; + if mutation.provider_kind().as_str() != XAI_PROVIDER_NAME + || owner_ref(mutation.owner_binding().owner()) != document.owner_ref + { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + PendingAuthorization::from_server_state( + oauth_config, + &SecretValue::new(document.server_state.clone()), + ) + .map_err(map_oauth_error)?; + Ok(StoredXaiAuthorization { + flow_id: document.flow_id, + owner_ref: document.owner_ref, + expires_at: document.expires_at, + server_state: document.server_state, + mutation, + }) +} + +fn decode_mutation( + document: PendingMutationDocument, +) -> Result { + let provider_kind = ProviderKind::new(document.provider_kind) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + let target = match document.target { + PendingTargetDocument::Create { + provider_instance_id, + name, + } => AuthorizationMutationTarget::Create { + provider_instance_id: ProviderInstanceId::new(provider_instance_id) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?, + name, + }, + PendingTargetDocument::Reauthorize { + provider_instance_id, + account_id, + expected_credential_revision, + } => AuthorizationMutationTarget::Reauthorize { + provider_instance_id: ProviderInstanceId::new(provider_instance_id) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?, + account_id: ProviderAccountId::new(account_id) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?, + expected_credential_revision: Revision::new(expected_credential_revision) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?, + }, + }; + let actor = match document.owner { + PendingOwnerDocument::AdminSession { admin_user_id } => { + MutationActor::AdminSession { admin_user_id } + } + PendingOwnerDocument::AdminApiKey => MutationActor::AdminApiKey, + PendingOwnerDocument::System => MutationActor::System, + }; + let context = MutationContext { + actor, + request_id: document.started_request_id, + }; + Ok(PendingAuthorizationMutation::new( + Revision::new(document.expected_config_revision) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?, + provider_kind, + target, + AuthorizationOwnerBinding::from_context(&context), + )) +} + +fn callback(value: &str) -> Result { + let mut callback = + Url::parse(value).map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; + if callback.fragment().is_some() { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + let query = callback.query().unwrap_or_default().to_owned(); + callback.set_query(None); + let expected = Url::parse(crate::OFFICIAL_REDIRECT_URI) + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?; + if callback != expected { + return Err(provider_error(ProviderAdminErrorKind::Invalid)); + } + AuthorizationCallback::parse(&query) + .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid)) +} + +fn random_flow_id() -> Result { + let mut random = [0_u8; 32]; + getrandom::fill(&mut random) + .map_err(|_| provider_error(ProviderAdminErrorKind::Unavailable))?; + Ok(URL_SAFE_NO_PAD.encode(random)) +} + +fn owner_ref(owner: &AuthorizationOwner) -> String { + let mut digest = Sha256::new(); + match owner { + AuthorizationOwner::AdminSession { admin_user_id } => { + digest.update(b"admin-session\0"); + digest.update(admin_user_id.as_bytes()); + } + AuthorizationOwner::AdminApiKey => digest.update(b"admin-api-key"), + AuthorizationOwner::System => digest.update(b"system"), + } + URL_SAFE_NO_PAD.encode(digest.finalize()) +} + +fn binding(value: &str) -> Result { + OAuthPendingBinding::try_new(value.to_owned()).map_err(map_provider_store_error) +} + +fn valid_pending_text(value: &str) -> bool { + !value.is_empty() + && value.len() <= MAX_PENDING_TEXT_BYTES + && !value.chars().any(char::is_control) +} + +fn currency_cost(money: Money) -> Result { + Ok(CurrencyCost { + currency: money.currency().as_str().to_owned(), + amount: money + .amount() + .to_string() + .parse::() + .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?, + }) +} + +fn provider_error(kind: ProviderAdminErrorKind) -> ProviderAdminError { + ProviderAdminError::new(kind) +} + +fn map_failure_class(class: FailureClass) -> ProviderAdminError { + provider_error(match class { + FailureClass::Transient => ProviderAdminErrorKind::Unavailable, + FailureClass::Ambiguous => ProviderAdminErrorKind::Conflict, + FailureClass::CredentialPermanent + | FailureClass::ConfigurationPermanent + | FailureClass::UserActionRequired + | FailureClass::Security => ProviderAdminErrorKind::Invalid, + FailureClass::Unsupported => ProviderAdminErrorKind::Unsupported, + }) +} + +fn map_oauth_error(error: OAuthError) -> ProviderAdminError { + map_failure_class(error.class()) +} + +fn map_provider_store_error(error: ProviderStoreError) -> ProviderAdminError { + provider_error(match error.kind() { + ProviderStoreErrorKind::InvalidData => ProviderAdminErrorKind::Invalid, + ProviderStoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict, + ProviderStoreErrorKind::Unavailable => ProviderAdminErrorKind::Unavailable, + }) +} + +fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError { + provider_error(match error.kind() { + StoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict, + StoreErrorKind::InvalidData | StoreErrorKind::InvalidState => { + ProviderAdminErrorKind::NotFound + } + StoreErrorKind::Unavailable => ProviderAdminErrorKind::Unavailable, + _ => ProviderAdminErrorKind::Internal, + }) +} + +fn map_repository_error(error: GrokCredentialRepositoryError) -> ProviderAdminError { + use GrokCredentialRepositoryError as Error; + provider_error(match error { + Error::InvalidInput(_) + | Error::WrongProviderKind + | Error::IdentityRebind + | Error::InvalidCredentialData => ProviderAdminErrorKind::Invalid, + Error::CredentialNotFound => ProviderAdminErrorKind::NotFound, + Error::StaleCredentialRevision | Error::Conflict | Error::RevisionOverflow => { + ProviderAdminErrorKind::Conflict + } + Error::Store => ProviderAdminErrorKind::Unavailable, + }) +} + +fn map_refresh_error(error: GrokCredentialRefreshError) -> ProviderAdminError { + use GrokCredentialRefreshError as Error; + match error { + Error::Repository(error) => map_repository_error(error), + Error::Lease(error) => map_provider_store_error(error), + Error::LeaseBusy => provider_error(ProviderAdminErrorKind::Conflict), + Error::InvalidRefreshResponse => provider_error(ProviderAdminErrorKind::Invalid), + Error::Preparation => provider_error(ProviderAdminErrorKind::Unavailable), + Error::ManualFailure(failure) => map_failure_class(match failure { + crate::GrokRefreshFailure::Transient => FailureClass::Transient, + crate::GrokRefreshFailure::Ambiguous => FailureClass::Ambiguous, + crate::GrokRefreshFailure::InvalidGrant + | crate::GrokRefreshFailure::Banned + | crate::GrokRefreshFailure::Rejected => FailureClass::CredentialPermanent, + }), + } +} + +fn map_quota_error(error: GrokQuotaError) -> ProviderAdminError { + use GrokQuotaError as Error; + provider_error(match error { + Error::AccountUnavailable => ProviderAdminErrorKind::NotFound, + Error::StaleCredentialSnapshot => ProviderAdminErrorKind::Conflict, + Error::InvalidData => ProviderAdminErrorKind::Invalid, + Error::Upstream | Error::Store => ProviderAdminErrorKind::Unavailable, + }) +} + +fn map_catalog_error(error: GrokCredentialCatalogError) -> ProviderAdminError { + use GrokCredentialCatalogError as Error; + provider_error(match error { + Error::InvalidInstance | Error::InvalidCredentialData | Error::ConflictingModelFacts => { + ProviderAdminErrorKind::Invalid + } + Error::NoEligibleCredential => ProviderAdminErrorKind::NotFound, + Error::StaleCredentialSnapshot => ProviderAdminErrorKind::Conflict, + Error::Upstream | Error::Cache | Error::Store => ProviderAdminErrorKind::Unavailable, + }) +} diff --git a/backend/crates/providers/xai/src/config.rs b/backend/crates/providers/xai/src/config.rs new file mode 100644 index 000000000..12b087715 --- /dev/null +++ b/backend/crates/providers/xai/src/config.rs @@ -0,0 +1,30 @@ +//! xAI Provider 启动配置。 + +use std::path::Path; + +use serde::Deserialize; + +use crate::{GrokOAuthConfig, transport::GROK_CLIENT_VERSION}; + +/// xAI 的协议端点、client 与版本均为经过校验的官方固定事实。 +#[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct XaiConfig {} + +impl XaiConfig { + /// 校验 Provider-owned 固定协议配置；当前没有相对路径字段。 + pub fn resolve_and_validate(&mut self, _source_dir: &Path) -> Result<(), XaiConfigError> { + self.oauth_config().map(|_| ()) + } + + pub(crate) fn oauth_config(self) -> Result { + GrokOAuthConfig::official(GROK_CLIENT_VERSION).map_err(|_| XaiConfigError::InvalidProtocol) + } +} + +/// xAI 固定协议配置无法构造。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum XaiConfigError { + #[error("xAI official protocol configuration is invalid")] + InvalidProtocol, +} diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index 2074e8dda..d23cf3bb0 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -2,18 +2,23 @@ use std::collections::BTreeMap; use std::fmt; -use std::sync::Arc; -use std::time::SystemTime; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant, SystemTime}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use futures::{StreamExt, TryStreamExt, stream}; use gateway_core::engine::credential::{ - AccountAvailability, CredentialRevision, ProviderAccountId, + AccountAvailability, AccountQuotaSignals, CredentialRevision, OpaqueProviderData, + ProviderAccount, ProviderAccountId, QuotaObservation, }; -use gateway_core::routing::{ConfigRevision, ProviderInstance}; +use gateway_core::engine::provider::ProviderCatalogGeneration; +use gateway_core::provider_ports::{ProviderCatalogCacheKey, ProviderCatalogCachePort}; +use gateway_core::routing::{ConfigRevision, ProviderInstance, ProviderKind}; +use tokio::sync::Mutex; use super::repository::{GrokCredentialRepository, LoadedGrokCredential}; +use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState}; use crate::transport::GROK_CLIENT_VERSION; use crate::transport::catalog::{MAX_CATALOG_MODELS, valid_model_slug, validate_etag}; use crate::{ @@ -23,6 +28,9 @@ use crate::{ }; const MAX_CONCURRENT_CATALOG_REQUESTS: usize = 8; +const CATALOG_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5 * 60); +const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); +const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); /// Grok Build Free 额度的滚动观察窗口。 pub const GROK_FREE_ROLLING_WINDOW_SECONDS: u64 = 86_400; @@ -184,6 +192,169 @@ pub enum GrokQuotaError { pub struct GrokCredentialQuotaService { repository: GrokCredentialRepository, client: Arc, + scheduling: GrokQuotaSchedulingProjection, +} + +#[derive(Clone, Default)] +struct GrokQuotaSchedulingProjection { + state: Arc>, + hydration: Arc>, +} + +#[derive(Default)] +struct GrokQuotaProjectionState { + next_version: u64, + entries: BTreeMap, +} + +#[derive(Debug, Clone, Copy)] +struct GrokQuotaSchedulingEntry { + version: u64, + revision: CredentialRevision, + expires_at: Instant, + signals: Option, +} + +#[derive(Clone)] +struct GrokQuotaHydrationTarget { + account: ProviderAccount, + expected_version: Option, +} + +impl GrokQuotaSchedulingProjection { + fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + accounts + .iter() + .filter_map(|account| { + let current = state.entries.get(account.id()); + current + .is_none_or(|entry| { + entry.revision != account.revision() || now >= entry.expires_at + }) + .then(|| GrokQuotaHydrationTarget { + account: account.clone(), + expected_version: current.map(|entry| entry.version), + }) + }) + .collect() + } + + fn observe(&self, snapshot: &GrokQuotaSnapshot) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + self.replace( + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn observe_if_unchanged( + &self, + target: &GrokQuotaHydrationTarget, + snapshot: &GrokQuotaSnapshot, + ) -> bool { + let Some(remaining_ttl) = quota_projection_ttl(snapshot.observed_at()) else { + return false; + }; + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return true; + } + insert_grok_projection_entry( + &mut state, + snapshot.account_id().clone(), + snapshot.credential_revision(), + remaining_ttl, + quota_scheduling_signals(snapshot.billing()), + ); + true + } + + fn mark_unknown_if_unchanged(&self, target: &GrokQuotaHydrationTarget, ttl: Duration) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .entries + .get(target.account.id()) + .map(|entry| entry.version) + != target.expected_version + { + return; + } + insert_grok_projection_entry( + &mut state, + target.account.id().clone(), + target.account.revision(), + ttl, + None, + ); + } + + fn replace( + &self, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, + ) { + let mut state = self + .state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + insert_grok_projection_entry(&mut state, account_id, revision, ttl, signals); + } + + fn signals(&self, account: &ProviderAccount) -> Option { + let state = self + .state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .entries + .get(account.id()) + .filter(|entry| { + entry.revision == account.revision() && Instant::now() < entry.expires_at + }) + .and_then(|entry| entry.signals) + } +} + +fn insert_grok_projection_entry( + state: &mut GrokQuotaProjectionState, + account_id: ProviderAccountId, + revision: CredentialRevision, + ttl: Duration, + signals: Option, +) { + state.next_version = state.next_version.saturating_add(1); + state.entries.insert( + account_id, + GrokQuotaSchedulingEntry { + version: state.next_version, + revision, + expires_at: Instant::now() + ttl, + signals, + }, + ); } impl GrokCredentialQuotaService { @@ -195,9 +366,54 @@ impl GrokCredentialQuotaService { Self { repository, client: Arc::new(GrokBillingClient::new(transport)), + scheduling: GrokQuotaSchedulingProjection::default(), } } + /// 批量预热请求级额度投影；持久层或 Provider JSON 异常只退化为未知额度。 + pub async fn prepare_scheduling(&self, accounts: &[ProviderAccount]) { + if self.scheduling.hydration_targets(accounts).is_empty() { + return; + } + let _hydration = self.scheduling.hydration.lock().await; + let pending = self.scheduling.hydration_targets(accounts); + if pending.is_empty() { + return; + } + let pending_accounts = pending + .iter() + .map(|target| target.account.clone()) + .collect::>(); + let Ok(observations) = self.repository.quota_observations(&pending_accounts).await else { + for target in &pending { + self.scheduling + .mark_unknown_if_unchanged(target, QUOTA_HYDRATION_FAILURE_TTL); + } + return; + }; + let observations = observations + .into_iter() + .map(|observation| (observation.account_id.clone(), observation)) + .collect::>(); + for target in pending { + let snapshot = observations + .get(target.account.id()) + .filter(|observation| observation.expected_revision == target.account.revision()) + .and_then(quota_snapshot_from_observation); + if !snapshot + .is_some_and(|snapshot| self.scheduling.observe_if_unchanged(&target, &snapshot)) + { + self.scheduling + .mark_unknown_if_unchanged(&target, QUOTA_SCHEDULING_TTL); + } + } + } + + #[must_use] + pub fn scheduling_signals(&self, account: &ProviderAccount) -> Option { + self.scheduling.signals(account) + } + /// 立即刷新一个账号的动态 billing document，并以 credential revision CAS 写回。 pub async fn refresh_account( &self, @@ -231,12 +447,41 @@ impl GrokCredentialQuotaService { ) .await .map_err(map_quota_repository_error)?; - Ok(GrokQuotaSnapshot { + let snapshot = GrokQuotaSnapshot { account_id: loaded.account.id().clone(), credential_revision: loaded.account.revision(), observed_at, billing: presentation, - }) + }; + self.scheduling.observe(&snapshot); + let current = self + .repository + .load_current(account_id) + .await + .map_err(map_quota_repository_error)?; + if quota_refresh_may_update_state(&loaded.account, ¤t.account) + && let Some(availability) = quota_refresh_availability( + current.account.availability(), + quota_is_exhausted(snapshot.billing()), + ) + { + self.repository + .update_state(&UpdateGrokCredentialState { + account_id: current.account.id().clone(), + expected_revision: current.account.revision(), + availability, + availability_reason: matches!( + availability, + GrokCredentialAvailability::QuotaExhausted + ) + .then_some("quota_exhausted".to_owned()), + cooldown_until: None, + observed_at, + }) + .await + .map_err(map_quota_repository_error)?; + } + Ok(snapshot) } /// 读取并重新验证 Store 中的 Provider-owned quota JSON。 @@ -256,12 +501,14 @@ impl GrokCredentialQuotaService { return Err(GrokQuotaError::InvalidData); }; let document = quota.into_inner(); - Ok(Some(GrokQuotaSnapshot { + let snapshot = GrokQuotaSnapshot { account_id: observation.account_id, credential_revision: observation.expected_revision, observed_at: observed_at.into(), billing: billing_presentation(&document)?, - })) + }; + self.scheduling.observe(&snapshot); + Ok(Some(snapshot)) } } @@ -403,12 +650,148 @@ pub trait GrokCredentialCatalogCache: Send + Sync { revision: CredentialRevision, ) -> Result, GrokCatalogCacheError>; - async fn permits( + async fn observed_model_support( &self, account_id: &ProviderAccountId, revision: CredentialRevision, model: &str, - ) -> Result; + ) -> Result, GrokCatalogCacheError>; +} + +/// xAI 负责解释 catalog 文档，Store 只保存 opaque JSON。 +pub struct GrokCatalogCache { + port: Arc, + provider_kind: ProviderKind, +} + +impl GrokCatalogCache { + pub fn new(port: Arc) -> Result { + Ok(Self { + port, + provider_kind: ProviderKind::new("xai") + .map_err(|_| GrokCatalogCacheError::InvalidData)?, + }) + } + + fn key( + &self, + account_id: &ProviderAccountId, + revision: CredentialRevision, + ) -> ProviderCatalogCacheKey { + ProviderCatalogCacheKey::new(self.provider_kind.clone(), account_id.clone(), revision) + } + + fn encode(catalog: &GrokAccountCatalog) -> OpaqueProviderData { + let mut document = serde_json::Map::new(); + document.insert("version".to_owned(), serde_json::Value::from(1)); + document.insert( + "observedAt".to_owned(), + serde_json::Value::String(catalog.observed_at().to_rfc3339()), + ); + if let Some(etag) = catalog.seed().etag() { + document.insert( + "etag".to_owned(), + serde_json::Value::String(etag.to_owned()), + ); + } + document.insert( + "models".to_owned(), + serde_json::Value::Array( + catalog + .seed() + .models() + .iter() + .cloned() + .map(serde_json::Value::String) + .collect(), + ), + ); + OpaqueProviderData::new(document) + } + + fn decode( + account_id: &ProviderAccountId, + revision: CredentialRevision, + document: OpaqueProviderData, + ) -> Result { + let mut fields = document.into_inner(); + if fields.remove("version").and_then(|value| value.as_u64()) != Some(1) { + return Err(GrokCatalogCacheError::InvalidData); + } + let observed_at = fields + .remove("observedAt") + .and_then(|value| value.as_str().map(ToOwned::to_owned)) + .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) + .map(|value| value.with_timezone(&Utc)) + .ok_or(GrokCatalogCacheError::InvalidData)?; + let etag = match fields.remove("etag") { + None => None, + Some(serde_json::Value::String(value)) => Some(value), + Some(_) => return Err(GrokCatalogCacheError::InvalidData), + }; + let models = fields + .remove("models") + .and_then(|value| value.as_array().cloned()) + .ok_or(GrokCatalogCacheError::InvalidData)? + .into_iter() + .map(|value| { + value + .as_str() + .map(ToOwned::to_owned) + .ok_or(GrokCatalogCacheError::InvalidData) + }) + .collect::, _>>()?; + if !fields.is_empty() { + return Err(GrokCatalogCacheError::InvalidData); + } + let seed = GrokCredentialCatalogSeed::new(models, etag) + .map_err(|_| GrokCatalogCacheError::InvalidData)?; + Ok(GrokAccountCatalog::new( + account_id.clone(), + revision, + observed_at, + seed, + )) + } +} + +#[async_trait] +impl GrokCredentialCatalogCache for GrokCatalogCache { + async fn replace(&self, catalog: GrokAccountCatalog) -> Result<(), GrokCatalogCacheError> { + self.port + .replace( + &self.key(catalog.account_id(), catalog.revision()), + &Self::encode(&catalog), + CATALOG_CACHE_TTL, + ) + .await + .map_err(|_| GrokCatalogCacheError::Unavailable) + } + + async fn read( + &self, + account_id: &ProviderAccountId, + revision: CredentialRevision, + ) -> Result, GrokCatalogCacheError> { + self.port + .read(&self.key(account_id, revision)) + .await + .map_err(|_| GrokCatalogCacheError::Unavailable)? + .map(|document| Self::decode(account_id, revision, document)) + .transpose() + } + + async fn observed_model_support( + &self, + account_id: &ProviderAccountId, + revision: CredentialRevision, + model: &str, + ) -> Result, GrokCatalogCacheError> { + Ok(self + .read(account_id, revision) + .await? + .map(|catalog| catalog.seed().permits(model))) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] @@ -481,6 +864,13 @@ pub struct GrokCredentialCatalogService { repository: GrokCredentialRepository, client: Arc, cache: Arc, + published: Arc>, +} + +#[derive(Default)] +struct PublishedCatalogState { + generation: u64, + models: BTreeMap>, } impl GrokCredentialCatalogService { @@ -494,9 +884,19 @@ impl GrokCredentialCatalogService { repository, client: Arc::new(GrokModelCatalogClient::new(transport)), cache, + published: Arc::new(RwLock::new(PublishedCatalogState::default())), } } + #[must_use] + pub fn catalog_generation(&self) -> ProviderCatalogGeneration { + let published = self + .published + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + ProviderCatalogGeneration::new(published.generation) + } + pub async fn fetch_seed( &self, access_token: SecretValue, @@ -646,6 +1046,7 @@ impl GrokCredentialCatalogService { .await .map_err(|_| GrokCredentialCatalogError::Cache)?; } + self.publish_models(instance_config.id().as_str(), &models); Ok(FetchedInstanceCatalog { provider_instance_id: instance_config.id().as_str().to_owned(), @@ -654,6 +1055,24 @@ impl GrokCredentialCatalogService { models, }) } + + fn publish_models(&self, instance: &str, models: &[GrokCatalogModel]) { + let mut published = self + .published + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if published + .models + .get(instance) + .is_some_and(|current| current == models) + { + return; + } + published + .models + .insert(instance.to_owned(), models.to_vec()); + published.generation = published.generation.saturating_add(1); + } } impl fmt::Debug for GrokCredentialCatalogService { @@ -692,10 +1111,8 @@ fn eligible_catalog_candidate(candidate: &LoadedGrokCredential) -> bool { && match account.availability() { AccountAvailability::Unknown | AccountAvailability::Ready + | AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted => true, - AccountAvailability::Cooldown => { - account.cooldown_until().is_some_and(|until| until <= now) - } AccountAvailability::Expired | AccountAvailability::Banned | AccountAvailability::Invalid => false, @@ -805,6 +1222,102 @@ fn billing_presentation( }) } +fn quota_snapshot_from_observation(observation: &QuotaObservation) -> Option { + let observed_at = DateTime::::from(observation.observed_at?); + let document = observation.quota.as_ref()?.expose_to_provider(); + Some(GrokQuotaSnapshot { + account_id: observation.account_id.clone(), + credential_revision: observation.expected_revision, + observed_at, + billing: billing_presentation(document).ok()?, + }) +} + +fn quota_projection_ttl(observed_at: DateTime) -> Option { + let age = SystemTime::now() + .duration_since(SystemTime::from(observed_at)) + .unwrap_or(Duration::ZERO); + QUOTA_SCHEDULING_TTL + .checked_sub(age) + .filter(|remaining| !remaining.is_zero()) +} + +fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option { + let remaining_rank = billing + .used_percent() + .filter(|used| used.is_finite() && (0.0..=100.0).contains(used)) + .map(|used| ((100.0 - used) * 100.0).round() as u64); + let now = SystemTime::now(); + let reset_at = billing + .period_end() + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| SystemTime::from(value.to_utc())) + .filter(|reset_at| *reset_at > now); + (remaining_rank.is_some() || reset_at.is_some()) + .then(|| AccountQuotaSignals::new(reset_at, remaining_rank)) +} + +fn quota_is_exhausted(billing: &GrokBillingPresentation) -> bool { + if billing + .prepaid_balance_cents() + .is_some_and(|balance| balance > 0) + || billing + .used_percent() + .is_some_and(|used| used.is_finite() && used < 100.0) + || matches!( + (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), + (Some(used), Some(cap)) if cap > 0 && used < cap + ) + || matches!( + (billing.included_used_cents(), billing.monthly_limit_cents()), + (Some(used), Some(limit)) if limit > 0 && used < limit + ) + { + return false; + } + billing + .used_percent() + .is_some_and(|used| used.is_finite() && used >= 100.0) + || matches!( + (billing.included_used_cents(), billing.monthly_limit_cents()), + (Some(used), Some(limit)) if limit > 0 && used >= limit + ) + || matches!( + (billing.on_demand_used_cents(), billing.on_demand_cap_cents()), + (Some(used), Some(cap)) if cap > 0 && used >= cap + ) +} + +fn quota_refresh_may_update_state(before: &ProviderAccount, current: &ProviderAccount) -> bool { + current.enabled() + && current.access_token_expires_at() > SystemTime::now() + && current.revision() == before.revision() + && current.availability() == before.availability() + && current.cooldown_until() == before.cooldown_until() +} + +fn quota_refresh_availability( + current: AccountAvailability, + exhausted: bool, +) -> Option { + match current { + AccountAvailability::Invalid + | AccountAvailability::Expired + | AccountAvailability::Banned => None, + AccountAvailability::QuotaExhausted => { + (!exhausted).then_some(GrokCredentialAvailability::Ready) + } + AccountAvailability::Cooldown | AccountAvailability::Unknown => Some(if exhausted { + GrokCredentialAvailability::QuotaExhausted + } else { + GrokCredentialAvailability::Ready + }), + AccountAvailability::Ready => { + exhausted.then_some(GrokCredentialAvailability::QuotaExhausted) + } + } +} + fn dynamic_string( object: Option<&serde_json::Map>, field: &str, diff --git a/backend/crates/providers/xai/src/credential/import.rs b/backend/crates/providers/xai/src/credential/import.rs index fb30f6368..776d630ca 100644 --- a/backend/crates/providers/xai/src/credential/import.rs +++ b/backend/crates/providers/xai/src/credential/import.rs @@ -156,7 +156,11 @@ impl GrokOAuthImportDocument { } let wire: OAuthAccountDocumentWire = serde_json::from_slice(document) .map_err(|_| GrokOAuthImportError::InvalidField("document"))?; - if wire.version.is_some_and(|version| version != 1) + if wire + .kind + .as_deref() + .is_some_and(|kind| kind != "external-account-bundle") + || wire.version.is_some_and(|version| version != 1) || !wire.proxies.is_empty() || wire.accounts.is_empty() || wire.accounts.len() > MAX_IMPORT_ACCOUNTS @@ -244,7 +248,7 @@ impl fmt::Debug for GrokOAuthImportDocument { #[serde(deny_unknown_fields)] struct OAuthAccountDocumentWire { #[serde(default, rename = "type")] - _kind: Option, + kind: Option, #[serde(default)] version: Option, exported_at: DateTime, diff --git a/backend/crates/providers/xai/src/credential/mod.rs b/backend/crates/providers/xai/src/credential/mod.rs index b843ab853..6289eeb47 100644 --- a/backend/crates/providers/xai/src/credential/mod.rs +++ b/backend/crates/providers/xai/src/credential/mod.rs @@ -27,9 +27,10 @@ pub use config::{ pub use catalog::{ GROK_FREE_ROLLING_WINDOW_SECONDS, GrokAccountCatalog, GrokBillingPresentation, - GrokCatalogCacheError, GrokCredentialCatalogCache, GrokCredentialCatalogError, - GrokCredentialCatalogSeed, GrokCredentialCatalogService, GrokCredentialCatalogSnapshot, - GrokCredentialQuotaService, GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, + GrokCatalogCache, GrokCatalogCacheError, GrokCredentialCatalogCache, + GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, + GrokCredentialCatalogSnapshot, GrokCredentialQuotaService, GrokQuotaError, GrokQuotaPeriodKind, + GrokQuotaSnapshot, }; pub use discovery::DiscoveryDocument; pub use error::{ @@ -49,19 +50,14 @@ pub use pkce::Pkce; pub use refresh::{ DueGrokCredential, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, GrokCredentialRefresher, GrokOAuthRefreshClient, - GrokRefreshFailure, GrokRefreshLeaseAcquisition, GrokRefreshLeaseCoordinator, - GrokRefreshLeaseError, GrokRefreshLeaseGuard, GrokRefreshLeaseRequest, GrokRefreshTokens, + GrokRefreshFailure, GrokRefreshTokens, }; pub use repository::{ GrokAccountExport, GrokCredentialAdmin, GrokCredentialLifecycle, GrokCredentialRepository, GrokCredentialRepositoryError, VerifiedGrokAccount, }; pub use secret::SecretValue; -pub use selector::{ - GrokAccountSchedulingState, GrokAccountSessionSelector, GrokCredentialLeaseAcquisition, - GrokCredentialLeaseCoordinator, GrokCredentialLeaseCoordinatorError, GrokCredentialLeaseGuard, - GrokCredentialLeaseRequest, -}; +pub use selector::GrokAccountSessionSelector; pub use token::{ OAuthPrincipal, RefreshTokenGrant, RefreshedTokenSet, VerifiedTokenSet, parse_oauth_error, parse_refresh_success, diff --git a/backend/crates/providers/xai/src/credential/oidc_verifier.rs b/backend/crates/providers/xai/src/credential/oidc_verifier.rs index 499e261ac..044402993 100644 --- a/backend/crates/providers/xai/src/credential/oidc_verifier.rs +++ b/backend/crates/providers/xai/src/credential/oidc_verifier.rs @@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet}; use std::fmt; +use std::sync::Arc; use std::time::{Duration, Instant}; use base64::Engine; @@ -15,7 +16,7 @@ use serde::Deserialize; use tokio::sync::Mutex; use url::Url; -use crate::transport::network::{BoundedBody, GrokEndpointPolicy, build_client, collect_bounded}; +use crate::transport::network::{BoundedBody, GrokEndpointPolicy, collect_bounded}; use crate::{ OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, TokenCandidate, TokenVerificationContext, TokenVerifier, VerificationEvidence, VerificationFailure, VerificationFlow, VerificationFuture, @@ -23,7 +24,6 @@ use crate::{ const ES256_JWA: &str = "ES256"; const JWT_TYPE: &str = "JWT"; -const JWKS_CACHE_TTL: Duration = Duration::from_secs(60 * 60); const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const MAX_JWKS_BODY_BYTES: usize = 256 * 1024; const MAX_USERINFO_BODY_BYTES: usize = 64 * 1024; @@ -37,7 +37,7 @@ const MAX_AUDIENCES: usize = 16; /// 使用官方 JWKS 与 user-info endpoint 的严格生产 verifier。 pub struct ReqwestOidcTokenVerifier { client: Client, - endpoint_policy: GrokEndpointPolicy, + endpoint_policy: Arc, jwks_cache: Mutex, cache_ttl: Duration, } @@ -48,32 +48,16 @@ impl ReqwestOidcTokenVerifier { /// # Errors /// /// TLS client 初始化失败时返回 transport build error。 - pub fn new() -> Result { - let endpoint_policy = GrokEndpointPolicy::official_oauth(); - Self::with_endpoint_policy(endpoint_policy) - } - - /// 构建使用显式 endpoint policy 的 verifier。 - /// - /// # Errors - /// - /// HTTP client 初始化失败时返回 transport build error。 - pub fn with_endpoint_policy( - endpoint_policy: GrokEndpointPolicy, - ) -> Result { - Self::with_endpoint_policy_and_cache_ttl(endpoint_policy, JWKS_CACHE_TTL) - } - /// 构建使用显式 endpoint policy 与 JWKS cache TTL 的 verifier。 /// /// # Errors /// /// HTTP client 初始化失败时返回 transport build error。 - pub fn with_endpoint_policy_and_cache_ttl( - endpoint_policy: GrokEndpointPolicy, + pub fn new( + endpoint_policy: Arc, cache_ttl: Duration, ) -> Result { - let client = build_client(&endpoint_policy, Some(REQUEST_TIMEOUT))?; + let client = endpoint_policy.build_oauth_client(Some(REQUEST_TIMEOUT))?; Ok(Self { client, endpoint_policy, diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index 7dd473866..4e514362f 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -7,9 +7,15 @@ use std::time::{Duration, SystemTime}; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use futures::{StreamExt as _, stream}; use gateway_core::engine::credential::{ AccountAvailability, CredentialRevision, ProviderAccountId, }; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, + ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_retry_at, +}; use gateway_core::routing::ProviderInstanceId; use super::catalog::GrokCredentialCatalogService; @@ -30,7 +36,7 @@ const MAX_REFRESH_BATCH: u32 = 100; const MAX_REFRESH_EXCLUSIONS: usize = 400; const MAX_SECRET_BYTES: usize = 64 * 1_024; const DISCOVERY_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(5); -const TRANSIENT_COOLDOWN: Duration = Duration::from_secs(30); +const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60); /// 一个到期且已按 revision 读取明文 RT 的 xAI account。 pub struct DueGrokCredential { @@ -118,9 +124,7 @@ pub enum GrokRefreshFailure { #[async_trait] pub trait GrokCredentialRefresher: Send + Sync { - async fn prepare_cycle(&self) -> Result<(), GrokRefreshFailure> { - Ok(()) - } + async fn prepare_cycle(&self) -> Result<(), GrokRefreshFailure>; async fn refresh( &self, @@ -203,36 +207,6 @@ impl GrokCredentialRefresher for GrokOAuthRefreshClient { } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct GrokRefreshLeaseRequest { - pub provider_instance_id: ProviderInstanceId, - pub account_id: ProviderAccountId, - pub credential_revision: CredentialRevision, -} - -pub trait GrokRefreshLeaseGuard: Send + Sync + 'static {} - -impl GrokRefreshLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum GrokRefreshLeaseAcquisition { - Acquired(Box), - Unavailable, -} - -#[derive(Debug, thiserror::Error, Clone, Copy, PartialEq, Eq)] -pub enum GrokRefreshLeaseError { - #[error("xAI credential refresh lease is unavailable")] - Unavailable, -} - -#[async_trait] -pub trait GrokRefreshLeaseCoordinator: Send + Sync { - async fn try_acquire( - &self, - request: &GrokRefreshLeaseRequest, - ) -> Result; -} - #[derive(Clone, Debug, PartialEq, Eq)] pub enum GrokCredentialRefreshOutcome { Refreshed { @@ -267,9 +241,9 @@ pub enum GrokCredentialRefreshError { #[error(transparent)] Repository(#[from] GrokCredentialRepositoryError), #[error(transparent)] - Lease(#[from] GrokRefreshLeaseError), - #[error("xAI credential refresh configuration is invalid")] - InvalidConfiguration, + Lease(#[from] ProviderStoreError), + #[error("xAI credential refresh lease is busy")] + LeaseBusy, #[error("xAI OAuth refresh response is invalid")] InvalidRefreshResponse, #[error("xAI OAuth refresh preparation failed")] @@ -282,8 +256,8 @@ pub struct GrokCredentialRefreshService { repository: GrokCredentialRepository, refresher: Arc, catalog: Arc, - leases: Arc, - refresh_margin: Duration, + leases: Arc, + runtime_policy: Arc, } impl GrokCredentialRefreshService { @@ -291,33 +265,29 @@ impl GrokCredentialRefreshService { repository: GrokCredentialRepository, refresher: Arc, catalog: Arc, - leases: Arc, - refresh_margin: Duration, - ) -> Result { - if refresh_margin.is_zero() || chrono::Duration::from_std(refresh_margin).is_err() { - return Err(GrokCredentialRefreshError::InvalidConfiguration); - } - Ok(Self { + leases: Arc, + runtime_policy: Arc, + ) -> Self { + Self { repository, refresher, catalog, leases, - refresh_margin, - }) + runtime_policy, + } } pub async fn refresh_due( &self, - limit: u32, ) -> Result, GrokCredentialRefreshError> { if !self.prepare_cycle_if_due().await? { return Ok(Vec::new()); } - self.refresh_due_excluding(limit, &[]).await + self.refresh_due_excluding(&[]).await } pub async fn prepare_cycle_if_due(&self) -> Result { - if !self.repository.has_due_refresh(self.refresh_margin).await? { + if !self.repository.has_due_refresh().await? { return Ok(false); } self.refresher @@ -338,6 +308,7 @@ impl GrokCredentialRefreshService { .load_managed(account_id, expected_revision) .await?; let loaded = super::repository::loaded_from_core(current.clone())?; + let policy = self.runtime_policy.load_refresh_policy().await?; if loaded .refresh_token_expires_at .is_some_and(|expires_at| expires_at <= Utc::now()) @@ -350,18 +321,28 @@ impl GrokCredentialRefreshService { .prepare_cycle() .await .map_err(GrokCredentialRefreshError::ManualFailure)?; - let lease = self + let capacity_guard = match self .leases - .try_acquire(&GrokRefreshLeaseRequest { - provider_instance_id: current.account.instance().clone(), - account_id: account_id.clone(), - credential_revision: expected_revision, - }) - .await?; - let guard = match lease { - GrokRefreshLeaseAcquisition::Acquired(guard) => guard, - GrokRefreshLeaseAcquisition::Unavailable => { - return Err(GrokRefreshLeaseError::Unavailable.into()); + .try_acquire(ProviderLeaseRequest::RefreshCapacity( + ProviderRefreshCapacityRequest::new(policy.concurrency()), + )) + .await? + { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { + return Err(GrokCredentialRefreshError::LeaseBusy); + } + }; + let account_guard = match self + .leases + .try_acquire(ProviderLeaseRequest::Refresh( + ProviderRefreshLeaseRequest::new(account_id.clone(), expected_revision), + )) + .await? + { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { + return Err(GrokCredentialRefreshError::LeaseBusy); } }; let tokens = self @@ -380,7 +361,7 @@ impl GrokCredentialRefreshService { return Err(GrokCredentialRefreshError::InvalidRefreshResponse); } let (access_token_expires_at, next_refresh_at) = - refreshed_deadlines(tokens.expires_in, self.refresh_margin) + refreshed_deadlines(account_id, tokens.expires_in, policy) .ok_or(GrokCredentialRefreshError::InvalidRefreshResponse)?; let prepared = GrokCredentialAdmin .prepare_rotation(&RotateManagedGrokCredential { @@ -402,21 +383,21 @@ impl GrokCredentialRefreshService { } else { loaded.refresh_token_expires_at }, - next_refresh_at: Some(next_refresh_at), }, + next_refresh_at, }) .map_err(GrokCredentialRefreshError::from)?; - Ok(prepared.with_refresh_guard(guard)) + Ok(prepared.with_refresh_guards(capacity_guard, account_guard)) } pub async fn refresh_due_excluding( &self, - limit: u32, excluded_account_ids: &[ProviderAccountId], ) -> Result, GrokCredentialRefreshError> { + let policy = self.runtime_policy.load_refresh_policy().await?; let batch = self .repository - .list_due_refresh_batch_excluding(self.refresh_margin, limit, excluded_account_ids) + .list_due_refresh_batch_excluding(excluded_account_ids) .await?; let mut outcomes = batch .failed_account_ids @@ -424,9 +405,18 @@ impl GrokCredentialRefreshService { .map(|account_id| GrokCredentialRefreshOutcome::Failed { account_id }) .collect::>(); outcomes.reserve(batch.credentials.len()); - for credential in batch.credentials { + let refreshed = stream::iter(batch.credentials.into_iter().map(|credential| async move { let account_id = credential.account_id.clone(); - match self.refresh_one(credential).await { + ( + account_id, + self.refresh_one_with_policy(credential, policy).await, + ) + })) + .buffer_unordered(policy.concurrency().get() as usize) + .collect::>() + .await; + for (account_id, result) in refreshed { + match result { Ok(outcome) => outcomes.push(outcome), Err(_) => outcomes.push(GrokCredentialRefreshOutcome::Failed { account_id }), } @@ -437,6 +427,15 @@ impl GrokCredentialRefreshService { pub async fn refresh_one( &self, credential: DueGrokCredential, + ) -> Result { + let policy = self.runtime_policy.load_refresh_policy().await?; + self.refresh_one_with_policy(credential, policy).await + } + + async fn refresh_one_with_policy( + &self, + credential: DueGrokCredential, + policy: ProviderRefreshPolicy, ) -> Result { let account_id = credential.account_id.clone(); if credential @@ -451,23 +450,36 @@ impl GrokCredentialRefreshService { ) .await; } + let capacity = self + .leases + .try_acquire(ProviderLeaseRequest::RefreshCapacity( + ProviderRefreshCapacityRequest::new(policy.concurrency()), + )) + .await?; + let _capacity_guard = match capacity { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { + return Ok(GrokCredentialRefreshOutcome::LeaseUnavailable { account_id }); + } + }; let lease = self .leases - .try_acquire(&GrokRefreshLeaseRequest { - provider_instance_id: credential.provider_instance_id.clone(), - account_id: account_id.clone(), - credential_revision: credential.credential_revision, - }) + .try_acquire(ProviderLeaseRequest::Refresh( + ProviderRefreshLeaseRequest::new( + account_id.clone(), + credential.credential_revision, + ), + )) .await?; let _guard = match lease { - GrokRefreshLeaseAcquisition::Acquired(guard) => guard, - GrokRefreshLeaseAcquisition::Unavailable => { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { .. } => { return Ok(GrokCredentialRefreshOutcome::LeaseUnavailable { account_id }); } }; match self.refresher.refresh(&credential.refresh_token).await { - Ok(tokens) => self.persist_success(credential, tokens).await, + Ok(tokens) => self.persist_success(credential, tokens, policy).await, Err(GrokRefreshFailure::InvalidGrant) => { self.persist_terminal_failure( credential, @@ -493,14 +505,24 @@ impl GrokCredentialRefreshService { .await .map(|_| GrokCredentialRefreshOutcome::Ambiguous { account_id }), Err(GrokRefreshFailure::Transient) => { - self.persist_cooldown(&credential, "refresh_transient") - .await?; - Ok(GrokCredentialRefreshOutcome::Transient { account_id }) + if self + .persist_backoff(&credential, "refresh-transient") + .await? + { + Ok(GrokCredentialRefreshOutcome::Transient { account_id }) + } else { + Ok(GrokCredentialRefreshOutcome::Stale { account_id }) + } } Err(GrokRefreshFailure::Rejected) => { - self.persist_cooldown(&credential, "refresh_rejected") - .await?; - Ok(GrokCredentialRefreshOutcome::Rejected { account_id }) + if self + .persist_backoff(&credential, "refresh-rejected") + .await? + { + Ok(GrokCredentialRefreshOutcome::Rejected { account_id }) + } else { + Ok(GrokCredentialRefreshOutcome::Stale { account_id }) + } } } } @@ -509,6 +531,7 @@ impl GrokCredentialRefreshService { &self, credential: DueGrokCredential, tokens: GrokRefreshTokens, + policy: ProviderRefreshPolicy, ) -> Result { if tokens.access_token.is_empty() || tokens.access_token.len() > MAX_SECRET_BYTES { return Err(GrokCredentialRefreshError::InvalidRefreshResponse); @@ -520,10 +543,10 @@ impl GrokCredentialRefreshService { if refresh_token.is_empty() || refresh_token.len() > MAX_SECRET_BYTES { return Err(GrokCredentialRefreshError::InvalidRefreshResponse); } - let (access_expires_at, next_refresh_at) = - refreshed_deadlines(tokens.expires_in, self.refresh_margin) - .ok_or(GrokCredentialRefreshError::InvalidRefreshResponse)?; let account_id = credential.account_id.clone(); + let (access_expires_at, next_refresh_at) = + refreshed_deadlines(&account_id, tokens.expires_in, policy) + .ok_or(GrokCredentialRefreshError::InvalidRefreshResponse)?; let access_token = tokens.access_token.clone(); let subject = credential.subject.clone(); let email = credential.email.clone(); @@ -549,8 +572,8 @@ impl GrokCredentialRefreshService { } else { credential.refresh_token_expires_at }, - next_refresh_at: Some(next_refresh_at), }, + next_refresh_at, }) .await { @@ -589,6 +612,17 @@ impl GrokCredentialRefreshService { reason: &str, ) -> Result { let account_id = credential.account_id.clone(); + match self + .repository + .load(&account_id, credential.credential_revision) + .await + { + Ok(_) => {} + Err(error) if stale_repository_error(&error) => { + return Ok(GrokCredentialRefreshOutcome::Stale { account_id }); + } + Err(error) => return Err(error.into()), + } match self .repository .update_state(&UpdateGrokCredentialState { @@ -609,59 +643,54 @@ impl GrokCredentialRefreshService { } } - async fn persist_cooldown( + async fn persist_backoff( &self, credential: &DueGrokCredential, - reason: &str, - ) -> Result<(), GrokCredentialRefreshError> { - let cooldown_until = chrono::Duration::from_std(TRANSIENT_COOLDOWN) - .ok() - .and_then(|duration| Utc::now().checked_add_signed(duration)); - self.repository - .update_state(&UpdateGrokCredentialState { - account_id: credential.account_id.clone(), - expected_revision: credential.credential_revision, - availability: GrokCredentialAvailability::Cooldown, - availability_reason: Some(reason.to_owned()), - cooldown_until, - observed_at: Utc::now(), - }) + reason: &'static str, + ) -> Result { + let retry_at = provider_refresh_retry_at( + &credential.account_id, + SystemTime::now(), + REFRESH_RETRY_DELAY, + reason, + )?; + match self + .repository + .defer_refresh( + &credential.account_id, + credential.credential_revision, + retry_at, + ) .await - .map_err(Into::into) + { + Ok(()) => Ok(true), + Err(error) if stale_repository_error(&error) => Ok(false), + Err(error) => Err(error.into()), + } } } impl GrokCredentialRepository { - pub async fn has_due_refresh( - &self, - refresh_before: Duration, - ) -> Result { + pub async fn has_due_refresh(&self) -> Result { Ok(self .list_all_accounts() .await? .iter() - .any(|account| account_due(account, refresh_before, SystemTime::now()))) + .any(|account| account_due(account, SystemTime::now()))) } pub async fn list_due_refresh( &self, - refresh_before: Duration, - limit: u32, ) -> Result, GrokCredentialRepositoryError> { - self.list_due_refresh_batch_excluding(refresh_before, limit, &[]) + self.list_due_refresh_batch_excluding(&[]) .await .map(|batch| batch.credentials) } async fn list_due_refresh_batch_excluding( &self, - refresh_before: Duration, - limit: u32, excluded_account_ids: &[ProviderAccountId], ) -> Result { - if limit == 0 || limit > MAX_REFRESH_BATCH { - return Err(GrokCredentialRepositoryError::InvalidInput("refresh_limit")); - } if excluded_account_ids.len() > MAX_REFRESH_EXCLUSIONS { return Err(GrokCredentialRepositoryError::InvalidInput( "refresh_exclusions", @@ -676,9 +705,7 @@ impl GrokCredentialRepository { .list_all_accounts() .await? .into_iter() - .filter(|account| { - !excluded.contains(account.id()) && account_due(account, refresh_before, now) - }) + .filter(|account| !excluded.contains(account.id()) && account_due(account, now)) .collect::>(); accounts.sort_by_key(|account| { ( @@ -687,7 +714,7 @@ impl GrokCredentialRepository { account.id().clone(), ) }); - accounts.truncate(limit as usize); + accounts.truncate(MAX_REFRESH_BATCH as usize); let mut due = Vec::with_capacity(accounts.len()); let mut failed_account_ids = Vec::new(); @@ -722,7 +749,6 @@ impl GrokCredentialRepository { fn account_due( account: &gateway_core::engine::credential::ProviderAccount, - refresh_before: Duration, now: SystemTime, ) -> bool { let availability = account.availability(); @@ -736,11 +762,7 @@ fn account_due( ) && (availability != AccountAvailability::Cooldown || account.cooldown_until().is_some_and(|until| until <= now)) - && (account.next_refresh_at().is_some_and(|next| next <= now) - || account - .access_token_expires_at() - .duration_since(now) - .is_ok_and(|remaining| remaining <= refresh_before)) + && account.next_refresh_at().is_some_and(|next| next <= now) } fn refreshed_tokens(tokens: RefreshedTokenSet) -> Result { @@ -775,20 +797,20 @@ fn classify_oauth_refresh_error(error: OAuthError) -> GrokRefreshFailure { } fn refreshed_deadlines( + account_id: &ProviderAccountId, expires_in: Duration, - refresh_margin: Duration, + policy: ProviderRefreshPolicy, ) -> Option<(DateTime, DateTime)> { if expires_in.is_zero() { return None; } - let lead = refresh_margin.min(expires_in / 2); - if lead.is_zero() || lead >= expires_in { - return None; - } let expires = chrono::Duration::from_std(expires_in).ok()?; - let lead = chrono::Duration::from_std(lead).ok()?; - let access_expires_at = Utc::now().checked_add_signed(expires)?; - let next_refresh_at = access_expires_at.checked_sub_signed(lead)?; + let observed_at = SystemTime::now(); + let access_expires_at = DateTime::::from(observed_at).checked_add_signed(expires)?; + let next_refresh_at = policy + .next_attempt_at(account_id, access_expires_at.into(), observed_at) + .ok() + .map(DateTime::::from)?; Some((access_expires_at, next_refresh_at)) } diff --git a/backend/crates/providers/xai/src/credential/repository.rs b/backend/crates/providers/xai/src/credential/repository.rs index 591a44f53..bac64d364 100644 --- a/backend/crates/providers/xai/src/credential/repository.rs +++ b/backend/crates/providers/xai/src/credential/repository.rs @@ -2,7 +2,7 @@ use std::collections::BTreeSet; use std::sync::Arc; -use std::time::{Duration, SystemTime}; +use std::time::SystemTime; use chrono::{DateTime, Utc}; use gateway_core::engine::credential::{ @@ -12,6 +12,7 @@ use gateway_core::engine::credential::{ ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome, }; use gateway_core::error::StoreErrorKind; +use gateway_core::provider_ports::ProviderRefreshPolicy; use gateway_core::routing::{ProviderInstanceId, ProviderKind}; use serde_json::Value; use thiserror::Error; @@ -42,7 +43,6 @@ pub struct VerifiedGrokAccount { pub plan_type: Option, pub tokens: VerifiedTokenSet, pub enabled: bool, - pub refresh_margin: Duration, } impl std::fmt::Debug for VerifiedGrokAccount { @@ -60,7 +60,6 @@ impl std::fmt::Debug for VerifiedGrokAccount { .field("plan_type", &self.plan_type) .field("tokens", &self.tokens) .field("enabled", &self.enabled) - .field("refresh_margin", &self.refresh_margin) .finish() } } @@ -175,7 +174,7 @@ impl GrokCredentialAdmin { availability(input.initial_availability), input.initial_cooldown_until.map(to_system_time), ) - .with_refresh_schedule(true, input.account.next_refresh_at.map(to_system_time)); + .with_refresh_schedule(true, Some(to_system_time(input.next_refresh_at))); Ok(NewProviderAccount { account, credential: encode_secret(&input.secret, &input.account)?, @@ -206,12 +205,8 @@ impl GrokCredentialAdmin { pub fn prepare_verified_account( &self, input: &VerifiedGrokAccount, + policy: ProviderRefreshPolicy, ) -> Result { - if input.refresh_margin.is_zero() { - return Err(GrokCredentialRepositoryError::InvalidInput( - "refresh_margin", - )); - } let expires_in = input .tokens .expires_in() @@ -224,9 +219,9 @@ impl GrokCredentialAdmin { let access_token_expires_at = now .checked_add(expires_in) .ok_or(GrokCredentialRepositoryError::InvalidInput("expires_in"))?; - let next_refresh_at = access_token_expires_at - .checked_sub(input.refresh_margin) - .map_or(now, |planned| planned.max(now)); + let next_refresh_at = policy + .next_attempt_at(&input.account_id, access_token_expires_at, now) + .map_err(|_| GrokCredentialRepositoryError::InvalidInput("next_refresh_at"))?; self.prepare_import(&CreateGrokCredential { account_id: input.account_id.clone(), provider_instance_id: input.provider_instance_id.clone(), @@ -244,8 +239,8 @@ impl GrokCredentialAdmin { plan_type: input.plan_type.clone(), access_token_expires_at: access_token_expires_at.into(), refresh_token_expires_at: None, - next_refresh_at: Some(next_refresh_at.into()), }, + next_refresh_at: next_refresh_at.into(), enabled: input.enabled, initial_availability: GrokCredentialAvailability::Ready, initial_availability_reason: None, @@ -335,6 +330,7 @@ impl GrokCredentialAdmin { account.revision(), &input.secret, &input.verified_account, + input.next_refresh_at, ) } } @@ -390,6 +386,7 @@ impl GrokCredentialRepository { input.expected_revision, &input.secret, &input.verified_account, + input.next_refresh_at, )?; let outcome = self .store @@ -407,6 +404,48 @@ impl GrokCredentialRepository { }) } + /// 临时失败时只推进持久刷新时刻，不改变 OAuth material 或账号可用性。 + pub(crate) async fn defer_refresh( + &self, + account_id: &ProviderAccountId, + expected_revision: CredentialRevision, + next_refresh_at: SystemTime, + ) -> Result<(), GrokCredentialRepositoryError> { + let current = self + .store + .load_credential(account_id, expected_revision) + .await + .map_err(map_rotation_load_error)?; + ensure_xai(¤t.account)?; + let profile = ProviderAccountUpdate { + account_id: account_id.clone(), + name: current.account.name().to_owned(), + email: current.account.email().map(str::to_owned), + plan_type: current.account.plan_type().map(str::to_owned), + }; + let update = CredentialCasUpdate::new( + account_id.clone(), + expected_revision, + profile, + current.credential, + true, + current.account.access_token_expires_at(), + Some(next_refresh_at), + ) + .map_err(|_| GrokCredentialRepositoryError::InvalidCredentialData)?; + match self + .store + .compare_and_swap_credential(update) + .await + .map_err(map_store_error)? + { + CredentialCasOutcome::Updated(_) => Ok(()), + CredentialCasOutcome::Conflict => { + Err(GrokCredentialRepositoryError::StaleCredentialRevision) + } + } + } + /// 用同一 credential revision fence 更新可用性。 pub async fn update_state( &self, @@ -540,7 +579,25 @@ impl GrokCredentialRepository { ) -> Result, GrokCredentialRepositoryError> { self.ensure_xai_account(account_id).await?; self.store - .get_quota(account_id) + .get_quotas(std::slice::from_ref(account_id)) + .await + .map_err(map_store_error) + .map(|mut observations| observations.pop()) + } + + pub(crate) async fn quota_observations( + &self, + accounts: &[ProviderAccount], + ) -> Result, GrokCredentialRepositoryError> { + for account in accounts { + ensure_xai(account)?; + } + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + self.store + .get_quotas(&account_ids) .await .map_err(map_store_error) } @@ -613,8 +670,10 @@ fn prepare_rotation( expected_revision: CredentialRevision, secret: &GrokOAuthSecret, verified_account: &GrokAccountProfile, + next_refresh_at: DateTime, ) -> Result { validate_profile(verified_account)?; + validate_refresh_schedule(next_refresh_at, verified_account.access_token_expires_at)?; validate_secret(secret)?; ensure_xai(current)?; if current.revision() != expected_revision { @@ -648,7 +707,7 @@ fn prepare_rotation( encode_secret(&merged_secret, verified_account)?, true, to_system_time(verified_account.access_token_expires_at), - verified_account.next_refresh_at.map(to_system_time), + Some(to_system_time(next_refresh_at)), ) .map_err(|_| GrokCredentialRepositoryError::InvalidCredentialData)?; Ok(PreparedGrokCredentialRotation::new(profile, credential)) @@ -674,6 +733,7 @@ fn decode_secret( fn validate_create(input: &CreateGrokCredential) -> Result<(), GrokCredentialRepositoryError> { validate_name(&input.name)?; validate_profile(&input.account)?; + validate_refresh_schedule(input.next_refresh_at, input.account.access_token_expires_at)?; validate_secret(&input.secret)?; validate_reason(input.initial_availability_reason.as_deref())?; validate_cooldown(input.initial_availability, input.initial_cooldown_until) @@ -692,9 +752,6 @@ fn validate_profile(profile: &GrokAccountProfile) -> Result<(), GrokCredentialRe || profile .refresh_token_expires_at .is_some_and(|expires_at| expires_at <= profile.access_token_expires_at) - || profile - .next_refresh_at - .is_some_and(|next| next >= profile.access_token_expires_at) { return Err(GrokCredentialRepositoryError::InvalidInput( "token_lifetime", @@ -703,6 +760,18 @@ fn validate_profile(profile: &GrokAccountProfile) -> Result<(), GrokCredentialRe Ok(()) } +fn validate_refresh_schedule( + next_refresh_at: DateTime, + access_token_expires_at: DateTime, +) -> Result<(), GrokCredentialRepositoryError> { + if next_refresh_at >= access_token_expires_at { + return Err(GrokCredentialRepositoryError::InvalidInput( + "next_refresh_at", + )); + } + Ok(()) +} + fn validate_secret(secret: &GrokOAuthSecret) -> Result<(), GrokCredentialRepositoryError> { for value in [&secret.access_token, &secret.refresh_token] { let exposed = value.expose(); diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index 5fe4d1c0c..f48f9d083 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -1,19 +1,19 @@ //! xAI OAuth account 选择、Redis lease 与失败反馈。 -use std::collections::BTreeMap; use std::fmt; use std::sync::Arc; use std::time::{Duration, SystemTime}; -use async_trait::async_trait; use chrono::Utc; use gateway_core::engine::credential::{ - AccountCandidate, AccountRuntimeSignals, AccountSelectionContext, AccountSelector, - CredentialRevision, ProviderAccountId, + AccountCandidate, AccountSelectionContext, AccountSelector, +}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, + ProviderSchedulingLeaseRequest, }; -use gateway_core::routing::ProviderInstanceId; -use super::catalog::GrokCredentialCatalogCache; +use super::catalog::{GrokCredentialCatalogCache, GrokCredentialQuotaService}; use super::repository::GrokCredentialRepository; use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState}; use crate::{ @@ -25,70 +25,12 @@ use crate::{ const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60); const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60); -/// Redis coordinator 接收的全局调度策略事实。 -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct GrokCredentialLeaseRequest { - pub provider_instance_id: ProviderInstanceId, - pub account_id: ProviderAccountId, - pub credential_revision: CredentialRevision, - pub max_concurrent_per_account: u32, - pub request_interval: Duration, -} - -pub trait GrokCredentialLeaseGuard: Send + Sync + 'static {} - -impl GrokCredentialLeaseGuard for T where T: Send + Sync + 'static {} - -pub enum GrokCredentialLeaseAcquisition { - Acquired(Box), - Unavailable { retry_after: Option }, -} - -impl fmt::Debug for GrokCredentialLeaseAcquisition { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Acquired(_) => formatter.write_str("Acquired([LEASE])"), - Self::Unavailable { retry_after } => formatter - .debug_struct("Unavailable") - .field("retry_after", retry_after) - .finish(), - } - } -} - -/// 一次 instance 的可失效调度状态；持久事实仍来自 PostgreSQL。 -#[derive(Debug, Clone)] -pub struct GrokAccountSchedulingState { - pub signals: BTreeMap, - pub sticky_account: Option, - pub round_robin_cursor: u64, -} - -#[async_trait] -pub trait GrokCredentialLeaseCoordinator: Send + Sync { - async fn load_scheduling_state( - &self, - provider_instance_id: &ProviderInstanceId, - accounts: &[ProviderAccountId], - ) -> Result; - - async fn try_acquire( - &self, - request: &GrokCredentialLeaseRequest, - ) -> Result; -} - -#[derive(Debug, thiserror::Error, PartialEq, Eq)] -pub enum GrokCredentialLeaseCoordinatorError { - #[error("Grok credential lease coordinator is unavailable")] - Unavailable, -} - /// 仅经 Core account port、TTL catalog cache 和 Redis lease 选择一个 OAuth session。 pub struct GrokAccountSessionSelector { repository: GrokCredentialRepository, catalog_cache: Arc, - lease_coordinator: Arc, + quota: Arc, + scheduling: Arc, } impl GrokAccountSessionSelector { @@ -96,12 +38,14 @@ impl GrokAccountSessionSelector { pub fn new( repository: GrokCredentialRepository, catalog_cache: Arc, - lease_coordinator: Arc, + quota: Arc, + scheduling: Arc, ) -> Self { Self { repository, catalog_cache, - lease_coordinator, + quota, + scheduling, } } @@ -114,18 +58,18 @@ impl GrokAccountSessionSelector { .list_accounts_for_instance(request.provider_instance_id()) .await .map_err(|_| GrokSessionSelectorError::Unavailable)?; + self.quota.prepare_scheduling(&accounts).await; let mut catalog_eligible = Vec::new(); for account in accounts { - if self + let support = self .catalog_cache - .permits( + .observed_model_support( account.id(), account.revision(), request.upstream_model().as_str(), ) - .await - .map_err(|_| GrokSessionSelectorError::Unavailable)? - { + .await; + if !matches!(support, Ok(Some(false))) { catalog_eligible.push(account); } } @@ -138,18 +82,19 @@ impl GrokAccountSessionSelector { .map(|account| account.id().clone()) .collect::>(); let scheduling = self - .lease_coordinator - .load_scheduling_state(request.provider_instance_id(), &account_ids) + .scheduling + .load_state(request.provider_instance_id(), &account_ids) .await .map_err(|_| GrokSessionSelectorError::Unavailable)?; let mut candidates = catalog_eligible .into_iter() .map(|account| { let signals = scheduling - .signals + .signals() .get(account.id()) .cloned() - .ok_or(GrokSessionSelectorError::Unavailable)?; + .ok_or(GrokSessionSelectorError::Unavailable)? + .with_provider_quota(self.quota.scheduling_signals(&account)); Ok(AccountCandidate { account, signals }) }) .collect::, GrokSessionSelectorError>>()?; @@ -164,8 +109,8 @@ impl GrokAccountSessionSelector { sticky_account: request .required_account() .cloned() - .or(scheduling.sticky_account), - round_robin_cursor: scheduling.round_robin_cursor, + .or_else(|| scheduling.sticky_account().cloned()), + round_robin_cursor: scheduling.round_robin_cursor(), }; let mut capacity_denied = false; let mut retry_after = None; @@ -173,22 +118,24 @@ impl GrokAccountSessionSelector { let selected_id = selected.account.id().clone(); let selected_revision = selected.account.revision(); let lease = self - .lease_coordinator - .try_acquire(&GrokCredentialLeaseRequest { - provider_instance_id: request.provider_instance_id().clone(), - account_id: selected_id.clone(), - credential_revision: selected_revision, - max_concurrent_per_account: request - .account_selection_policy() - .max_concurrent_per_account() - .get(), - request_interval: request.account_selection_policy().request_interval(), - }) + .scheduling + .try_acquire(ProviderLeaseRequest::Scheduling( + ProviderSchedulingLeaseRequest::new( + request.provider_instance_id().clone(), + selected_id.clone(), + selected_revision, + request + .account_selection_policy() + .max_concurrent_per_account(), + request.account_selection_policy().request_interval(), + request.deadline(), + ), + )) .await .map_err(|_| GrokSessionSelectorError::Unavailable)?; let guard = match lease { - GrokCredentialLeaseAcquisition::Acquired(guard) => guard, - GrokCredentialLeaseAcquisition::Unavailable { + ProviderLeaseAcquisition::Acquired(guard) => guard, + ProviderLeaseAcquisition::Busy { retry_after: candidate_retry, } => { capacity_denied = true; @@ -240,7 +187,8 @@ impl fmt::Debug for GrokAccountSessionSelector { .debug_struct("GrokAccountSessionSelector") .field("repository", &self.repository) .field("catalog_cache", &"[TTL_CACHE]") - .field("lease_coordinator", &"[LEASE_COORDINATOR]") + .field("quota", &self.quota) + .field("scheduling", &"[SCHEDULING_PORT]") .finish() } } diff --git a/backend/crates/providers/xai/src/credential/types.rs b/backend/crates/providers/xai/src/credential/types.rs index 8d910d4d5..60d50a710 100644 --- a/backend/crates/providers/xai/src/credential/types.rs +++ b/backend/crates/providers/xai/src/credential/types.rs @@ -7,11 +7,11 @@ use gateway_core::engine::credential::{ CredentialCasUpdate, CredentialRevision, LoadedCredential, ProviderAccountId, ProviderAccountUpdate, }; +use gateway_core::provider_ports::ProviderLeaseGuard; use gateway_core::routing::ProviderInstanceId; use serde::{Deserialize, Serialize}; use zeroize::{Zeroize, ZeroizeOnDrop}; -use super::refresh::GrokRefreshLeaseGuard; use crate::SecretValue; /// 官方 xAI OAuth token pair；不接受 API Key。 @@ -43,7 +43,6 @@ pub struct GrokAccountProfile { pub access_token_expires_at: DateTime, /// Provider 明确返回时才保存；普通账号导出通常不携带 RT 过期时间。 pub refresh_token_expires_at: Option>, - pub next_refresh_at: Option>, } impl fmt::Debug for GrokAccountProfile { @@ -59,7 +58,6 @@ impl fmt::Debug for GrokAccountProfile { .field("plan_type", &self.plan_type) .field("access_token_expires_at", &self.access_token_expires_at) .field("refresh_token_expires_at", &self.refresh_token_expires_at) - .field("next_refresh_at", &self.next_refresh_at) .finish() } } @@ -83,6 +81,7 @@ pub struct CreateGrokCredential { pub name: String, pub secret: GrokOAuthSecret, pub account: GrokAccountProfile, + pub next_refresh_at: DateTime, pub enabled: bool, pub initial_availability: GrokCredentialAvailability, pub initial_availability_reason: Option, @@ -98,6 +97,7 @@ impl fmt::Debug for CreateGrokCredential { .field("name", &self.name) .field("secret", &"[REDACTED]") .field("account", &self.account) + .field("next_refresh_at", &self.next_refresh_at) .field("enabled", &self.enabled) .field("initial_availability", &self.initial_availability) .finish_non_exhaustive() @@ -110,6 +110,7 @@ pub struct RotateGrokCredential { pub expected_revision: CredentialRevision, pub secret: GrokOAuthSecret, pub verified_account: GrokAccountProfile, + pub next_refresh_at: DateTime, } /// App 已读取的当前账号与 xAI 已验证的新 OAuth 身份材料。 @@ -117,6 +118,7 @@ pub struct RotateManagedGrokCredential { pub current: LoadedCredential, pub secret: GrokOAuthSecret, pub verified_account: GrokAccountProfile, + pub next_refresh_at: DateTime, } impl fmt::Debug for RotateManagedGrokCredential { @@ -126,6 +128,7 @@ impl fmt::Debug for RotateManagedGrokCredential { .field("current", &self.current) .field("secret", &"[REDACTED]") .field("verified_account", &self.verified_account) + .field("next_refresh_at", &self.next_refresh_at) .finish() } } @@ -134,7 +137,12 @@ impl fmt::Debug for RotateManagedGrokCredential { pub struct PreparedGrokCredentialRotation { pub profile: ProviderAccountUpdate, pub credential: CredentialCasUpdate, - refresh_guard: Option>, + refresh_guards: Option, +} + +struct ProviderRefreshGuards { + _capacity: Box, + _account: Box, } impl fmt::Debug for PreparedGrokCredentialRotation { @@ -144,8 +152,8 @@ impl fmt::Debug for PreparedGrokCredentialRotation { .field("profile", &self.profile) .field("credential", &self.credential) .field( - "refresh_guard", - &self.refresh_guard.as_ref().map(|_| ""), + "refresh_guards", + &self.refresh_guards.as_ref().map(|_| ""), ) .finish() } @@ -156,18 +164,25 @@ impl PreparedGrokCredentialRotation { Self { profile, credential, - refresh_guard: None, + refresh_guards: None, } } - pub(crate) fn with_refresh_guard(mut self, guard: Box) -> Self { - self.refresh_guard = Some(guard); + pub(crate) fn with_refresh_guards( + mut self, + capacity: Box, + account: Box, + ) -> Self { + self.refresh_guards = Some(ProviderRefreshGuards { + _capacity: capacity, + _account: account, + }); self } #[must_use] pub const fn holds_refresh_lease(&self) -> bool { - self.refresh_guard.is_some() + self.refresh_guards.is_some() } /// 将 command 与 lease 一起交给 App；返回的 guard 必须活到 CAS 提交结束。 @@ -182,13 +197,13 @@ impl PreparedGrokCredentialRotation { ( self.profile, self.credential, - PreparedGrokCredentialRotationGuard(self.refresh_guard), + PreparedGrokCredentialRotationGuard(self.refresh_guards), ) } } /// 手工刷新从 token exchange 到数据库 CAS 完成期间持有的 Redis lease。 -pub struct PreparedGrokCredentialRotationGuard(Option>); +pub struct PreparedGrokCredentialRotationGuard(Option); impl fmt::Debug for PreparedGrokCredentialRotationGuard { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -207,6 +222,7 @@ impl fmt::Debug for RotateGrokCredential { .field("expected_revision", &self.expected_revision) .field("secret", &"[REDACTED]") .field("verified_account", &self.verified_account) + .field("next_refresh_at", &self.next_refresh_at) .finish() } } diff --git a/backend/crates/providers/xai/src/lib.rs b/backend/crates/providers/xai/src/lib.rs index 6f867c552..830a676f8 100644 --- a/backend/crates/providers/xai/src/lib.rs +++ b/backend/crates/providers/xai/src/lib.rs @@ -1,30 +1,43 @@ //! Official Grok Build OAuth Provider boundary. +mod admin; +pub mod config; pub mod credential; mod provider; pub mod transport; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use gateway_admin::ports::provider::ProviderAdmin; +use gateway_core::engine::credential::ProviderAccountStore; +use gateway_core::engine::provider::Provider; +use gateway_core::provider_ports::ProviderStorePorts; +use gateway_core::routing::ProviderKind; +use gateway_core::task::WorkerContribution; + +use crate::admin::{XaiAdminProvider, XaiAdminServices}; + +pub use config::{XaiConfig, XaiConfigError}; + pub use credential::{ AllowedRedirectUri, AuthorizationCallback, AuthorizationCodeGrant, CallbackRejection, ConfigError, CreateGrokCredential, DiscoveryDocument, DueGrokCredential, FailClosedTokenVerifier, FailureClass, FormField, FormValue, GROK_FREE_ROLLING_WINDOW_SECONDS, - GrokAccountCatalog, GrokAccountExport, GrokAccountProfile, GrokAccountSchedulingState, - GrokAccountSessionSelector, GrokBillingPresentation, GrokCatalogCacheError, - GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, - GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, - GrokCredentialCatalogSnapshot, GrokCredentialLeaseAcquisition, GrokCredentialLeaseCoordinator, - GrokCredentialLeaseCoordinatorError, GrokCredentialLeaseGuard, GrokCredentialLeaseRequest, + GrokAccountCatalog, GrokAccountExport, GrokAccountProfile, GrokAccountSessionSelector, + GrokBillingPresentation, GrokCatalogCache, GrokCatalogCacheError, GrokCredentialAdmin, + GrokCredentialAvailability, GrokCredentialCatalogCache, GrokCredentialCatalogError, + GrokCredentialCatalogSeed, GrokCredentialCatalogService, GrokCredentialCatalogSnapshot, GrokCredentialLifecycle, GrokCredentialQuotaService, GrokCredentialRecord, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, GrokCredentialRefresher, GrokCredentialRepository, GrokCredentialRepositoryError, GrokOAuthClient, GrokOAuthConfig, GrokOAuthImportCandidate, GrokOAuthImportDocument, GrokOAuthImportEntry, GrokOAuthImportError, GrokOAuthImportMetadata, GrokOAuthImportTokens, GrokOAuthRefreshClient, GrokOAuthSecret, GrokQuotaError, GrokQuotaPeriodKind, - GrokQuotaSnapshot, GrokRefreshFailure, GrokRefreshLeaseAcquisition, - GrokRefreshLeaseCoordinator, GrokRefreshLeaseError, GrokRefreshLeaseGuard, - GrokRefreshLeaseRequest, GrokRefreshTokens, HttpHeader, HttpMethod, OAuthError, OAuthErrorCode, - OAuthHttpRequest, OAuthHttpResponse, OAuthHttpTransport, OAuthOperation, OAuthPrincipal, - OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, OFFICIAL_REDIRECT_URI, OFFICIAL_SCOPES, + GrokQuotaSnapshot, GrokRefreshFailure, GrokRefreshTokens, HttpHeader, HttpMethod, OAuthError, + OAuthErrorCode, OAuthHttpRequest, OAuthHttpResponse, OAuthHttpTransport, OAuthOperation, + OAuthPrincipal, OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, OFFICIAL_REDIRECT_URI, OFFICIAL_SCOPES, PendingAuthorization, Pkce, PreparedGrokCredentialRotation, PreparedGrokCredentialRotationGuard, ProtocolViolation, RedirectUriAllowlist, RefreshTokenGrant, RefreshedTokenSet, ReqwestOidcTokenVerifier, RotateGrokCredential, @@ -50,10 +63,154 @@ pub use transport::{ GrokModelCatalogSnapshot, GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokProviderConfigError, GrokProviderInstanceConfig, - GrokProviderTransport, GrokRequestEncodeError, GrokResponsesRequest, GrokSessionBinding, - GrokSessionDataError, GrokSessionLeaseGuard, GrokSessionSelection, GrokSessionSelector, - GrokSessionSelectorError, GrokSessionSelectorFuture, MAX_GROK_BILLING_BYTES, - MAX_GROK_MODEL_CATALOG_BYTES, ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, - ReqwestOAuthTransport, SelectedGrokSession, XAI_PROVIDER_NAME, build_grok_headers, - grok_billing_breakdown, parse_grok_billing, parse_grok_model_catalog, + GrokProviderTransport, GrokRequestEncodeError, GrokReqwestTransportBuildError, + GrokResponsesRequest, GrokSessionBinding, GrokSessionDataError, GrokSessionLeaseGuard, + GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, GrokSessionSelectorFuture, + MAX_GROK_BILLING_BYTES, MAX_GROK_MODEL_CATALOG_BYTES, OfficialGrokEndpointPolicy, + ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport, + SelectedGrokSession, XAI_PROVIDER_NAME, build_grok_headers, grok_billing_breakdown, + parse_grok_billing, parse_grok_model_catalog, }; + +/// xAI 初始化后交给组装根的最小能力集。 +pub struct ProviderBundle { + core_provider: Arc, + admin_provider: Arc, + worker_contributions: Vec, +} + +/// 构造 xAI 数据面、管理面准备器与 Provider-owned 后台任务。 +pub async fn initialize( + mut config: XaiConfig, + ports: ProviderStorePorts, +) -> Result { + config + .resolve_and_validate(Path::new(".")) + .map_err(XaiInitializeError::Config)?; + let oauth_config = config.oauth_config().map_err(XaiInitializeError::Config)?; + let provider_kind = + ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| XaiInitializeError::ProviderKind)?; + let accounts: Arc = ports.accounts(); + let leases = ports.leases(); + let runtime_policy = ports.runtime_policy(); + let instances = ports.instances(); + let repository = GrokCredentialRepository::new(Arc::clone(&accounts)); + let endpoint_policy: Arc = Arc::new(OfficialGrokEndpointPolicy); + + let catalog_cache = Arc::new( + GrokCatalogCache::new(ports.catalog_cache()).map_err(|_| XaiInitializeError::Catalog)?, + ); + let catalog_cache_port: Arc = catalog_cache.clone(); + let upstream_catalog = Arc::new( + ReqwestGrokModelCatalogTransport::new(Arc::clone(&endpoint_policy)) + .map_err(|_| XaiInitializeError::Transport)?, + ); + let catalog_transport: Arc = upstream_catalog.clone(); + let billing_transport: Arc = upstream_catalog; + let catalog = Arc::new(GrokCredentialCatalogService::new( + repository.clone(), + catalog_transport, + catalog_cache_port, + )); + let quota = Arc::new(GrokCredentialQuotaService::new( + repository.clone(), + billing_transport, + )); + let selector: Arc = Arc::new(GrokAccountSessionSelector::new( + repository.clone(), + catalog_cache, + Arc::clone("a), + Arc::clone(&leases), + )); + let inference: Arc = Arc::new( + ReqwestGrokInferenceTransport::new(Arc::clone(&endpoint_policy)) + .map_err(|_| XaiInitializeError::Transport)?, + ); + let core_provider: Arc = Arc::new(GrokBuildProvider::new( + selector, + inference, + Arc::clone(&catalog), + )); + + let oauth_transport: Arc = Arc::new( + ReqwestOAuthTransport::new(Arc::clone(&endpoint_policy)) + .map_err(|_| XaiInitializeError::Transport)?, + ); + let token_verifier: Arc = Arc::new( + ReqwestOidcTokenVerifier::new(endpoint_policy, Duration::from_secs(60 * 60)) + .map_err(|_| XaiInitializeError::Transport)?, + ); + let oauth = Arc::new(GrokOAuthClient::new( + oauth_config.clone(), + oauth_transport, + token_verifier, + )); + let refresher: Arc = + Arc::new(GrokOAuthRefreshClient::new(Arc::clone(&oauth))); + let refresh = Arc::new(GrokCredentialRefreshService::new( + repository.clone(), + refresher, + Arc::clone(&catalog), + Arc::clone(&leases), + Arc::clone(&runtime_policy), + )); + let admin_provider: Arc = Arc::new(XaiAdminProvider::new( + provider_kind.clone(), + Arc::clone(&accounts), + XaiAdminServices { + repository, + oauth_config, + oauth, + pending: ports.oauth_pending(), + refresh: Arc::clone(&refresh), + quota: Arc::clone("a), + catalog: Arc::clone(&catalog), + runtime_policy, + }, + )); + let worker_contributions = + provider::worker_contributions(refresh, quota, catalog, accounts, instances, provider_kind) + .map_err(|_| XaiInitializeError::Worker)?; + + Ok(ProviderBundle { + core_provider, + admin_provider, + worker_contributions, + }) +} + +impl ProviderBundle { + #[must_use] + pub fn core_provider(&self) -> Arc { + Arc::clone(&self.core_provider) + } + + #[must_use] + pub fn admin_provider(&self) -> Arc { + Arc::clone(&self.admin_provider) + } + + /// 一次性移交 Host 任务计划，防止同一 owner 被重复注册。 + pub fn take_worker_contributions(&mut self) -> Vec { + std::mem::take(&mut self.worker_contributions) + } +} + +/// xAI 初始化失败的脱敏分类。 +#[derive(Debug, thiserror::Error)] +pub enum XaiInitializeError { + #[error(transparent)] + Config(XaiConfigError), + #[error("xAI runtime policy is unavailable")] + RuntimePolicy, + #[error("xAI Provider kind is invalid")] + ProviderKind, + #[error("xAI transport could not initialize")] + Transport, + #[error("xAI model catalog could not initialize")] + Catalog, + #[error("xAI credential refresh could not initialize")] + Refresh, + #[error("xAI worker plan is invalid")] + Worker, +} diff --git a/backend/crates/providers/xai/src/provider.rs b/backend/crates/providers/xai/src/provider.rs index 52ccc413c..f5c8fecab 100644 --- a/backend/crates/providers/xai/src/provider.rs +++ b/backend/crates/providers/xai/src/provider.rs @@ -5,21 +5,34 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use async_trait::async_trait; -use futures::StreamExt; +use futures::{StreamExt, future::BoxFuture}; +use gateway_core::engine::credential::{ + AccountAvailability, ProviderAccount, ProviderAccountStore, +}; use gateway_core::engine::provider::{ - EventStream, Provider, ProviderCallMetadata, ProviderModelCapabilities, ProviderRequest, - ProviderStream, UpstreamTransport, + EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, + ProviderModelCapabilities, ProviderRequest, ProviderStream, UpstreamTransport, }; use gateway_core::engine::{AttemptContext, UpstreamSendState}; use gateway_core::error::{ProviderError, ProviderErrorKind}; -use gateway_core::event::GatewayEvent; +use gateway_core::event::{GatewayEvent, ProviderEvent, ProviderResponseObservation}; use gateway_core::operation::{Feature, Operation, OperationKind}; +use gateway_core::provider_ports::ProviderInstanceCatalogPort; use gateway_core::routing::{ - ModelCapabilities, ProviderInstance, ProviderKind, SupportLevel, UpstreamModelId, + InstanceHealth, ModelCapabilities, ProviderInstance, ProviderKind, SupportLevel, + UpstreamModelId, +}; +use gateway_core::task::{ + ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId, + WorkerKind, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, + WorkerTaskError, }; use crate::GrokCatalogCapabilityEvidence; -use crate::credential::GrokCredentialCatalogService; +use crate::credential::{ + GrokCredentialCatalogService, GrokCredentialQuotaService, GrokCredentialRefreshOutcome, + GrokCredentialRefreshService, GrokQuotaError, +}; use crate::transport::canonical::GrokCanonicalDecoder; use crate::transport::config::XAI_PROVIDER_NAME; use crate::transport::headers::build_grok_headers; @@ -78,6 +91,10 @@ impl Provider for GrokBuildProvider { XAI_PROVIDER_NAME } + fn catalog_generation(&self) -> ProviderCatalogGeneration { + self.catalog.catalog_generation() + } + async fn query_model_capabilities( &self, instance: &ProviderInstance, @@ -164,6 +181,7 @@ impl Provider for GrokBuildProvider { context.excluded_accounts().clone(), context.required_account().cloned(), context.account_selection_policy(), + context.deadline(), ); let selection_deadline = remaining(context.deadline()).ok_or_else(|| { provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent) @@ -247,7 +265,6 @@ fn cold_http_sse_stream( } let headers = build_grok_headers( &instance, - &request, &session, context.request_id(), &upstream_model, @@ -267,29 +284,49 @@ fn cold_http_sse_stream( return; }; let cancellation = context.cancellation().clone(); - let response = tokio::select! { + let boundary = tokio::select! { biased; - _ = cancellation.cancelled() => Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Ambiguous, - )), - _ = tokio::time::sleep(handshake_deadline) => Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Ambiguous, - )), - response = transport.execute(inference_request) => { - response.map_err(map_transport_error) - } + _ = cancellation.cancelled() => InferenceBoundary::Cancelled, + _ = tokio::time::sleep(handshake_deadline) => InferenceBoundary::Deadline, + response = transport.execute(inference_request) => InferenceBoundary::Response(response), }; - let response = match response { - Ok(response) => response, - Err(error) => { - let error = record_failure(selector.as_ref(), &session, error).await; + let response = match boundary { + InferenceBoundary::Cancelled => { + Err(provider_error(ProviderErrorKind::Cancelled, UpstreamSendState::Ambiguous))?; + return; + } + InferenceBoundary::Deadline => { + Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Ambiguous))?; + return; + } + InferenceBoundary::Response(Ok(response)) => response, + InferenceBoundary::Response(Err(error)) => { + let observation = xai_error_observation(&error)?; + let error = record_failure( + selector.as_ref(), + &session, + map_transport_error(error), + ) + .await; + yield ProviderEvent::observation(observation); Err(error)?; return; } }; + let mut observation = ProviderResponseObservation::new( + UpstreamTransport::new(HTTP_SSE_TRANSPORT).map_err(|_| provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + ))?, + ) + .with_http_version(response.http_version()) + .with_status_code(response.status_code()); + if let Some(request_id) = response.request_id().cloned() { + observation = observation.with_request_id(request_id); + } + yield ProviderEvent::observation(observation); + let mut body = response.into_body(); let mut decoder = GrokCanonicalDecoder::new(upstream_model.as_str()); loop { @@ -332,6 +369,7 @@ fn cold_http_sse_stream( }; let completed = events .iter() + .flat_map(ProviderEvent::canonical_facts) .any(|event| matches!(event, GatewayEvent::Completed(_))); for event in events { ensure_sent_context(&context)?; @@ -356,6 +394,31 @@ fn cold_http_sse_stream( }) } +enum InferenceBoundary { + Response(Result), + Cancelled, + Deadline, +} + +fn xai_error_observation( + error: &GrokInferenceTransportError, +) -> Result { + let mut observation = ProviderResponseObservation::new( + UpstreamTransport::new(HTTP_SSE_TRANSPORT) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, error.send_state()))?, + ); + if let Some(http_version) = error.http_version() { + observation = observation.with_http_version(http_version); + } + if let Some(status_code) = error.status() { + observation = observation.with_status_code(status_code); + } + if let Some(request_id) = error.request_id().cloned() { + observation = observation.with_request_id(request_id); + } + Ok(observation) +} + async fn record_failure( selector: &dyn GrokSessionSelector, session: &SelectedGrokSession, @@ -419,11 +482,9 @@ fn ensure_sent_context(context: &AttemptContext) -> Result<(), ProviderError> { fn map_request_error(error: GrokRequestEncodeError) -> ProviderError { let kind = match error { - GrokRequestEncodeError::InvalidProviderOptions | GrokRequestEncodeError::InvalidImage => { - ProviderErrorKind::InvalidRequest - } - GrokRequestEncodeError::UnsupportedProviderOption - | GrokRequestEncodeError::UnsupportedContent => ProviderErrorKind::Unsupported, + GrokRequestEncodeError::InvalidProtocolPayload + | GrokRequestEncodeError::InvalidProviderOptions => ProviderErrorKind::InvalidRequest, + GrokRequestEncodeError::UnsupportedProviderOption => ProviderErrorKind::Unsupported, GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol, }; provider_error(kind, UpstreamSendState::NotSent) @@ -495,6 +556,9 @@ fn map_transport_error_with_state( if let Some(retry_after) = error.retry_after() { mapped = mapped.with_retry_after(retry_after); } + if let Some(request_id) = error.request_id().cloned() { + mapped = mapped.with_upstream_request_id(request_id); + } if error.sensitive_context_was_redacted() { mapped = mapped.redact_sensitive_context("upstream transport context"); } @@ -520,3 +584,178 @@ fn remaining(deadline: SystemTime) -> Option { .ok() .filter(|remaining| !remaining.is_zero()) } + +const WORKER_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60); +const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); +const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); +const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); +const QUOTA_CATALOG_INTERVAL: Duration = Duration::from_secs(5 * 60); + +pub(crate) fn worker_contributions( + refresh: Arc, + quota: Arc, + catalog: Arc, + accounts: Arc, + instances: Arc, + provider_kind: ProviderKind, +) -> Result, WorkerDefinitionError> { + let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, XAI_PROVIDER_NAME)?; + let catalog_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, XAI_PROVIDER_NAME)?; + Ok(vec![ + WorkerContribution::Registration(scheduled_registration( + refresh_id, + OAUTH_REFRESH_INTERVAL, + Box::new(XaiOAuthRefreshTask { service: refresh }), + )?), + WorkerContribution::Registration(scheduled_registration( + catalog_id, + QUOTA_CATALOG_INTERVAL, + Box::new(XaiQuotaCatalogTask { + accounts, + instances, + quota, + catalog, + provider_kind, + }), + )?), + ]) +} + +fn scheduled_registration( + id: WorkerId, + interval: Duration, + task: Box, +) -> Result { + let schedule = WorkerSchedule::try_new( + interval, + WORKER_INITIAL_BACKOFF, + WORKER_MAXIMUM_BACKOFF, + WORKER_LEASE_TTL, + WORKER_LEASE_RENEWAL, + )?; + let lease = WorkerLeaseRequest::try_new(id.clone(), WORKER_LEASE_TTL)?; + WorkerRegistration::try_new( + id, + WorkerRunnable::Scheduled { + schedule, + lease: Some(lease), + task, + }, + ) +} + +struct XaiOAuthRefreshTask { + service: Arc, +} + +impl ScheduledTask for XaiOAuthRefreshTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + if context.cancellation().is_cancelled() { + return Ok(()); + } + let outcomes = self + .service + .refresh_due() + .await + .map_err(|_| WorkerTaskError::safe("xAI OAuth refresh failed"))?; + let failures = outcomes + .iter() + .filter(|outcome| { + matches!( + outcome, + GrokCredentialRefreshOutcome::Ambiguous { .. } + | GrokCredentialRefreshOutcome::Transient { .. } + | GrokCredentialRefreshOutcome::Failed { .. } + ) + }) + .count(); + if failures > 0 { + tracing::warn!(failures, "xAI OAuth refresh cycle contained failures"); + } + Ok(()) + }) + } +} + +struct XaiQuotaCatalogTask { + accounts: Arc, + instances: Arc, + quota: Arc, + catalog: Arc, + provider_kind: ProviderKind, +} + +impl ScheduledTask for XaiQuotaCatalogTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let instances = self + .instances + .list_instances(&self.provider_kind, false) + .await + .map_err(|_| WorkerTaskError::safe("xAI Provider instances unavailable"))?; + let mut failures = 0_u64; + for config in instances { + if context.cancellation().is_cancelled() { + return Ok(()); + } + if !config.enabled() || config.provider_kind() != &self.provider_kind { + failures = failures.saturating_add(1); + continue; + } + let instance = ProviderInstance::new( + config.id().clone(), + config.provider_kind().clone(), + config.base_url().to_owned(), + true, + InstanceHealth::Healthy, + ); + match self.accounts.list_for_instance(instance.id()).await { + Ok(accounts) => { + let now = SystemTime::now(); + for account in accounts + .into_iter() + .filter(|account| eligible_quota_worker_account(account, now)) + { + if context.cancellation().is_cancelled() { + return Ok(()); + } + match self.quota.refresh_account(account.id()).await { + Ok(_) | Err(GrokQuotaError::AccountUnavailable) => {} + Err(_) => failures = failures.saturating_add(1), + } + } + } + Err(_) => failures = failures.saturating_add(1), + } + if self.catalog.query_instance_models(&instance).await.is_err() { + failures = failures.saturating_add(1); + } + } + if failures == 0 { + Ok(()) + } else { + Err(WorkerTaskError::safe( + "xAI quota or catalog synchronization failed", + )) + } + }) + } +} + +fn eligible_quota_worker_account(account: &ProviderAccount, now: SystemTime) -> bool { + account.enabled() + && account.access_token_expires_at() > now + && match account.availability() { + AccountAvailability::Unknown + | AccountAvailability::Ready + | AccountAvailability::QuotaExhausted => true, + AccountAvailability::Cooldown => { + account.cooldown_until().is_some_and(|until| until <= now) + } + AccountAvailability::Expired + | AccountAvailability::Banned + | AccountAvailability::Invalid => false, + } +} diff --git a/backend/crates/providers/xai/src/transport/canonical.rs b/backend/crates/providers/xai/src/transport/canonical.rs index 10985671b..b36170752 100644 --- a/backend/crates/providers/xai/src/transport/canonical.rs +++ b/backend/crates/providers/xai/src/transport/canonical.rs @@ -3,14 +3,14 @@ use std::collections::{BTreeMap, BTreeSet}; use gateway_core::accounting::{ - CalculatedCost, CalculatedCostBreakdown, CurrencyCode, Decimal, Money, ProviderReportedCost, - Usage, + CalculatedCost, CalculatedCostAmounts, CalculatedCostBreakdown, CalculatedCostRates, + CurrencyCode, Decimal, Money, ProviderReportedCost, Usage, }; use gateway_core::engine::UpstreamSendState; use gateway_core::error::{ProviderError, ProviderErrorKind}; use gateway_core::event::{ - ContentItem, ContentKind, FinishReason, GatewayEvent, ReasoningDelta, ResponseMeta, TextDelta, - ToolCallDelta, + ContentItem, ContentKind, FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, + ReasoningDelta, ResponseMeta, TextDelta, ToolCallDelta, }; use gateway_protocol::openai::events::{TokenUsage, extract_usage}; use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder}; @@ -94,16 +94,20 @@ pub fn grok_billing_breakdown( .checked_add(cache_read_amount_ticks)? .checked_add(output_amount_ticks)?; Some(CalculatedCostBreakdown::new( - usd_money(input_amount_ticks)?, - usd_money(output_amount_ticks)?, - usd_money(cache_read_amount_ticks)?, - usd_money(0)?, - usd_money(standard_amount_ticks)?, - usd_money(standard_amount_ticks)?, - usd_price_per_million(rates.input_ticks)?, - usd_price_per_million(rates.output_ticks)?, - usd_price_per_million(rates.cached_input_ticks)?, - usd_money(0)?, + CalculatedCostAmounts::new( + usd_money(input_amount_ticks)?, + usd_money(output_amount_ticks)?, + usd_money(cache_read_amount_ticks)?, + usd_money(0)?, + usd_money(standard_amount_ticks)?, + usd_money(standard_amount_ticks)?, + ), + CalculatedCostRates::new( + usd_price_per_million(rates.input_ticks)?, + usd_price_per_million(rates.output_ticks)?, + usd_price_per_million(rates.cached_input_ticks)?, + usd_money(0)?, + ), Some("default".to_owned()), 100, )) @@ -122,9 +126,7 @@ fn usd_price_per_million(per_token_ticks: u128) -> Option { /// Incremental decoder for one official Grok Responses attempt. /// -/// Unknown output kinds and events fail closed because gateway-core does not -/// currently have canonical representations for Grok backend tools, images, -/// searches, or MCP activity. No raw upstream data is retained in errors. +/// 每个上游 event 同时保留 OpenAI wire，并在可识别时附加 canonical facts。 pub struct GrokCanonicalDecoder { decoder: SseEventDecoder, fallback_model: String, @@ -150,23 +152,29 @@ impl GrokCanonicalDecoder { } } - pub fn push(&mut self, chunk: &[u8]) -> Result, ProviderError> { + pub fn push(&mut self, chunk: &[u8]) -> Result, ProviderError> { let events = self.decoder.push(chunk).map_err(protocol_error)?; self.decode(events) } - pub fn finish(&mut self) -> Result, ProviderError> { + pub fn finish(&mut self) -> Result, ProviderError> { let events = self.decoder.finish().map_err(protocol_error)?; - let canonical = self.decode(events)?; + let output = self.decode(events)?; if !self.completed { return Err(protocol_error_marker()); } - Ok(canonical) + Ok(output) } - fn decode(&mut self, events: Vec) -> Result, ProviderError> { - let mut canonical = Vec::new(); + fn decode(&mut self, events: Vec) -> Result, ProviderError> { + let mut output = Vec::new(); for event in events { + if event.data.trim() == "[DONE]" { + if !self.completed { + return Err(protocol_error_marker()); + } + continue; + } let value = serde_json::from_str::(&event.data).map_err(protocol_error)?; let body_type = value.get("type").and_then(Value::as_str); let event_type = match (event.event.as_deref(), body_type) { @@ -177,9 +185,17 @@ impl GrokCanonicalDecoder { (None, Some(body)) => body, (None, None) => return Err(protocol_error_marker()), }; + let mut canonical = Vec::new(); self.decode_event(event_type, &value, &mut canonical)?; + let wire = ProtocolWireEvent::json("openai", event.event, value) + .map_err(|_| protocol_error_marker())?; + output.push(if canonical.is_empty() { + ProviderEvent::wire(wire) + } else { + ProviderEvent::canonical_with_wire(canonical, wire) + }); } - Ok(canonical) + Ok(output) } fn decode_event( @@ -206,7 +222,9 @@ impl GrokCanonicalDecoder { "response.reasoning_summary_text.delta" | "response.reasoning_text.delta" => { self.reasoning_delta(value, output) } - "response.function_call_arguments.delta" => self.tool_delta(value, output), + "response.function_call_arguments.delta" | "response.custom_tool_call_input.delta" => { + self.tool_delta(value, output) + } "response.completed" | "response.incomplete" => { self.complete(event_type, value, output) } @@ -216,12 +234,13 @@ impl GrokCanonicalDecoder { | "response.reasoning_summary_text.done" | "response.reasoning_text.done" | "response.function_call_arguments.done" + | "response.custom_tool_call_input.done" | "response.content_part.done" | "response.reasoning_summary_part.done" | "response.reasoning_part.done" | "response.rate_limits.updated" | "rate_limits.updated" => Ok(()), - _ => Err(unsupported_event_error()), + _ => Ok(()), } } @@ -266,7 +285,7 @@ impl GrokCanonicalDecoder { let item = value.get("item").ok_or_else(protocol_error_marker)?; let output_index = event_index(value, "output_index")?; match item.get("type").and_then(Value::as_str) { - Some("function_call") => { + Some("function_call" | "custom_tool_call") => { let index = content_index(output_index, 0)?; self.add_content(index, ContentKind::ToolCall, output)?; let call_id = item @@ -293,8 +312,7 @@ impl GrokCanonicalDecoder { ContentKind::Reasoning, output, ), - Some("message") => Ok(()), - Some(_) | None => Err(unsupported_event_error()), + _ => Ok(()), } } @@ -307,8 +325,8 @@ impl GrokCanonicalDecoder { let item = value.get("item").ok_or_else(protocol_error_marker)?; match item.get("type").and_then(Value::as_str) { Some("message" | "reasoning") => return Ok(()), - Some("function_call") => {} - Some(_) | None => return Err(unsupported_event_error()), + Some("function_call" | "custom_tool_call") => {} + Some(_) | None => return Ok(()), } let output_index = event_index(value, "output_index")?; let index = content_index(output_index, 0)?; @@ -317,6 +335,7 @@ impl GrokCanonicalDecoder { } let Some(arguments) = item .get("arguments") + .or_else(|| item.get("input")) .and_then(Value::as_str) .filter(|value| !value.is_empty()) else { @@ -360,7 +379,7 @@ impl GrokCanonicalDecoder { let kind = match part.get("type").and_then(Value::as_str) { Some("output_text" | "refusal") => ContentKind::Text, Some("summary_text" | "reasoning_text") => ContentKind::Reasoning, - Some(_) | None => return Err(unsupported_event_error()), + Some(_) | None => return Ok(()), }; self.ensure_content(content_index(output_index, part_index)?, kind, output) } @@ -710,8 +729,3 @@ fn protocol_error_marker() -> ProviderError { ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::Sent) .redact_sensitive_context("invalid upstream event") } - -fn unsupported_event_error() -> ProviderError { - ProviderError::new(ProviderErrorKind::Unsupported, UpstreamSendState::Sent) - .redact_sensitive_context("unsupported upstream event") -} diff --git a/backend/crates/providers/xai/src/transport/catalog.rs b/backend/crates/providers/xai/src/transport/catalog.rs index b6593b7f2..c2cb94a64 100644 --- a/backend/crates/providers/xai/src/transport/catalog.rs +++ b/backend/crates/providers/xai/src/transport/catalog.rs @@ -250,16 +250,6 @@ impl GrokBillingRequest { fn from_session(session: &GrokModelCatalogSession) -> Result { let endpoint = Url::parse(GROK_BILLING_URL).map_err(|_| GrokBillingError::InvalidRequest)?; - Self::for_endpoint(endpoint, session) - } - - /// 构造由 transport endpoint policy 再次校验的 billing 请求。 - /// - /// 生产 client 只调用固定官方 URL；此入口仅供显式注入的 loopback transport 测试。 - pub fn for_endpoint( - endpoint: Url, - session: &GrokModelCatalogSession, - ) -> Result { let mut headers = vec![ GrokHeader::sensitive( "authorization", diff --git a/backend/crates/providers/xai/src/transport/config.rs b/backend/crates/providers/xai/src/transport/config.rs index 8a2ad1380..c831da1a0 100644 --- a/backend/crates/providers/xai/src/transport/config.rs +++ b/backend/crates/providers/xai/src/transport/config.rs @@ -74,12 +74,6 @@ pub enum GrokProviderConfigError { UnsafeBaseUrl, } -pub(crate) fn valid_header_value(value: &str, max_bytes: usize) -> bool { - !value.is_empty() - && value.len() <= max_bytes - && value.bytes().all(|byte| (0x20..=0x7e).contains(&byte)) -} - fn validate_official_base_url(value: &str) -> Result { let mut url = Url::parse(value).map_err(|_| GrokProviderConfigError::InvalidBaseUrl)?; if url.scheme() != "https" diff --git a/backend/crates/providers/xai/src/transport/headers.rs b/backend/crates/providers/xai/src/transport/headers.rs index ce711a1b9..3eaf36f2c 100644 --- a/backend/crates/providers/xai/src/transport/headers.rs +++ b/backend/crates/providers/xai/src/transport/headers.rs @@ -5,7 +5,6 @@ use gateway_core::routing::UpstreamModelId; use crate::SecretValue; -use super::request::GrokResponsesRequest; use super::{GrokProviderInstanceConfig, SelectedGrokSession}; const CLIENT_IDENTIFIER: &str = "codex-proxy-rs"; @@ -83,16 +82,11 @@ impl GrokHeader { pub fn build_grok_headers( instance: &GrokProviderInstanceConfig, - request: &GrokResponsesRequest, session: &SelectedGrokSession, request_id: &ModelRequestId, model: &UpstreamModelId, ) -> Vec { - let options = request.header_options(); let request_id = request_id.as_str(); - let conversation_id = options.conversation_id().unwrap_or(request_id); - let session_id = options.session_id().unwrap_or(request_id); - let agent_id = options.agent_id().unwrap_or(CLIENT_IDENTIFIER); let mut headers = vec![ GrokHeader::sensitive( @@ -112,14 +106,14 @@ pub fn build_grok_headers( ), GrokHeader::public("content-type", "application/json"), GrokHeader::public("accept", "text/event-stream"), - GrokHeader::sensitive( - "x-grok-conv-id", - SecretValue::new(conversation_id.to_owned()), - ), + GrokHeader::sensitive("x-grok-conv-id", SecretValue::new(request_id.to_owned())), GrokHeader::sensitive("x-grok-req-id", SecretValue::new(request_id.to_owned())), GrokHeader::public("x-grok-model-override", model.as_str()), - GrokHeader::sensitive("x-grok-session-id", SecretValue::new(session_id.to_owned())), - GrokHeader::sensitive("x-grok-agent-id", SecretValue::new(agent_id.to_owned())), + GrokHeader::sensitive("x-grok-session-id", SecretValue::new(request_id.to_owned())), + GrokHeader::sensitive( + "x-grok-agent-id", + SecretValue::new(CLIENT_IDENTIFIER.to_owned()), + ), ]; if let Some(email) = session.email() { headers.push(GrokHeader::sensitive("x-email", email.clone())); diff --git a/backend/crates/providers/xai/src/transport/mod.rs b/backend/crates/providers/xai/src/transport/mod.rs index 5d9268771..b0eb34f8f 100644 --- a/backend/crates/providers/xai/src/transport/mod.rs +++ b/backend/crates/providers/xai/src/transport/mod.rs @@ -33,7 +33,7 @@ pub use config::{ pub use headers::{GrokHeader, GrokHeaderValue, build_grok_headers}; pub use network::{ GrokDnsResolutionError, GrokDnsResolutionPlan, GrokDnsResolutionPolicy, GrokEndpointPolicy, - GrokReqwestTransportBuildError, ReqwestGrokInferenceTransport, + GrokReqwestTransportBuildError, OfficialGrokEndpointPolicy, ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport, }; pub use request::{GrokRequestEncodeError, GrokResponsesRequest}; diff --git a/backend/crates/providers/xai/src/transport/network.rs b/backend/crates/providers/xai/src/transport/network.rs index 6da632e65..55dc72c9a 100644 --- a/backend/crates/providers/xai/src/transport/network.rs +++ b/backend/crates/providers/xai/src/transport/network.rs @@ -27,6 +27,8 @@ use crate::{ TransportFailure, TransportFailureKind, TransportFuture, }; use gateway_core::engine::UpstreamSendState; +use gateway_core::error::SafeUpstreamValue; +use gateway_core::event::UpstreamHttpVersion; use gateway_protocol::openai::sse::MAX_SSE_EVENT_BUFFER_BYTES; pub(crate) const OFFICIAL_OAUTH_HOST: &str = "auth.x.ai"; @@ -57,9 +59,6 @@ pub enum GrokReqwestTransportBuildError { /// Reqwest TLS/client 初始化失败。 #[error("Grok reqwest transport initialization failed")] ClientInitialization, - /// Injected origin 不是数字 loopback HTTP origin。 - #[error("Grok injected transport origin is invalid")] - InvalidEndpointOrigin, } /// 固定官方 host 的 DNS 解析路径；只有系统结果全部为公网地址时才直接使用。 @@ -150,136 +149,82 @@ impl GrokDnsResolutionPolicy { #[error("Grok official DNS resolution was rejected")] pub struct GrokDnsResolutionError; -#[derive(Debug, Clone)] -enum EndpointPolicyKind { - Official { host: &'static str }, - Loopback { origin: LoopbackOrigin }, +/// HTTP client construction and endpoint validation are one injected security boundary. +pub trait GrokEndpointPolicy: fmt::Debug + Send + Sync { + fn build_oauth_client( + &self, + timeout: Option, + ) -> Result; + fn build_inference_client( + &self, + timeout: Option, + ) -> Result; + fn validate_oauth(&self, url: &Url) -> bool; + fn validate_inference(&self, url: &Url) -> bool; + fn validate_model_catalog(&self, url: &Url) -> bool; + fn route_billing(&self, url: &Url) -> Option; + fn validate_jwks(&self, url: &Url) -> bool; + fn validate_userinfo(&self, url: &Url) -> bool; } -/// Auditable endpoint policy injected into reqwest transports and OIDC verification. -#[derive(Debug, Clone)] -pub struct GrokEndpointPolicy(EndpointPolicyKind); +#[derive(Debug, Default)] +pub struct OfficialGrokEndpointPolicy; -impl GrokEndpointPolicy { - /// Restricts requests to the official xAI OAuth origin. - #[must_use] - pub const fn official_oauth() -> Self { - Self(EndpointPolicyKind::Official { - host: OFFICIAL_OAUTH_HOST, - }) - } - - /// Restricts requests to the official Grok CLI inference origin. - #[must_use] - pub const fn official_inference() -> Self { - Self(EndpointPolicyKind::Official { - host: OFFICIAL_INFERENCE_HOST, - }) - } - - /// Creates an explicitly injected numeric-loopback HTTP origin policy. - /// - /// # Errors - /// - /// Returns an error for domains, non-loopback addresses, missing ports, or - /// origins containing credentials, query, fragment, or a non-root path. - pub fn loopback(origin: &Url) -> Result { - Ok(Self(EndpointPolicyKind::Loopback { - origin: LoopbackOrigin::parse(origin)?, - })) - } - - /// Applies the same public-address predicate used by the strict DNS resolver. +impl OfficialGrokEndpointPolicy { #[must_use] pub fn accepts_resolved_address(address: IpAddr) -> bool { is_public_ip(address) } +} + +impl GrokEndpointPolicy for OfficialGrokEndpointPolicy { + fn build_oauth_client( + &self, + timeout: Option, + ) -> Result { + build_official_client(OFFICIAL_OAUTH_HOST, timeout) + } + + fn build_inference_client( + &self, + timeout: Option, + ) -> Result { + build_official_client(OFFICIAL_INFERENCE_HOST, timeout) + } fn validate_oauth(&self, url: &Url) -> bool { - match &self.0 { - EndpointPolicyKind::Official { host } => valid_official_url(url, host, None), - EndpointPolicyKind::Loopback { origin } => origin.matches(url), - } + valid_official_url(url, OFFICIAL_OAUTH_HOST, None) } fn validate_inference(&self, url: &Url) -> bool { - match &self.0 { - EndpointPolicyKind::Official { host } => { - valid_official_url(url, host, Some(OFFICIAL_INFERENCE_PATH)) - } - EndpointPolicyKind::Loopback { origin } => origin.matches(url), - } + valid_official_url(url, OFFICIAL_INFERENCE_HOST, Some(OFFICIAL_INFERENCE_PATH)) } fn validate_model_catalog(&self, url: &Url) -> bool { - match &self.0 { - EndpointPolicyKind::Official { host } => { - valid_official_url(url, host, Some(OFFICIAL_MODEL_CATALOG_PATH)) - } - EndpointPolicyKind::Loopback { origin } => { - origin.matches(url) - && url.path() == OFFICIAL_MODEL_CATALOG_PATH - && url.query().is_none() - && url.fragment().is_none() - } - } + valid_official_url( + url, + OFFICIAL_INFERENCE_HOST, + Some(OFFICIAL_MODEL_CATALOG_PATH), + ) } - fn validate_billing(&self, url: &Url) -> bool { - let valid_query = url.query_pairs().count() == 1 - && url - .query_pairs() - .next() - .is_some_and(|(key, value)| key == "format" && value == "credits"); - if !valid_query || url.fragment().is_some() { - return false; - } - match &self.0 { - EndpointPolicyKind::Official { host } => { - url.scheme() == "https" - && url.host_str() == Some(*host) - && url.port_or_known_default() == Some(443) - && url.path() == OFFICIAL_BILLING_PATH - && url.username().is_empty() - && url.password().is_none() - } - EndpointPolicyKind::Loopback { origin } => { - origin.matches(url) && url.path() == OFFICIAL_BILLING_PATH - } - } + fn route_billing(&self, url: &Url) -> Option { + valid_billing_url(url, OFFICIAL_INFERENCE_HOST).then(|| url.clone()) } - pub(crate) fn validate_jwks(&self, url: &Url) -> bool { - self.validate_exact_oauth_path(url, OFFICIAL_JWKS_PATH) + fn validate_jwks(&self, url: &Url) -> bool { + valid_official_url(url, OFFICIAL_OAUTH_HOST, Some(OFFICIAL_JWKS_PATH)) } - pub(crate) fn validate_userinfo(&self, url: &Url) -> bool { - self.validate_exact_oauth_path(url, OFFICIAL_USERINFO_PATH) + fn validate_userinfo(&self, url: &Url) -> bool { + valid_official_url(url, OFFICIAL_OAUTH_HOST, Some(OFFICIAL_USERINFO_PATH)) } - - fn validate_exact_oauth_path(&self, url: &Url, path: &str) -> bool { - match &self.0 { - EndpointPolicyKind::Official { host } => valid_official_url(url, host, Some(path)), - EndpointPolicyKind::Loopback { origin } => { - origin.matches(url) - && url.path() == path - && url.query().is_none() - && url.fragment().is_none() - } - } - } -} - -#[derive(Debug, Clone)] -struct LoopbackOrigin { - host: IpAddr, - port: u16, } /// 官方 OAuth HTTP transport。只允许 `auth.x.ai:443`。 pub struct ReqwestOAuthTransport { client: Client, - endpoint_policy: GrokEndpointPolicy, + endpoint_policy: Arc, } impl ReqwestOAuthTransport { @@ -288,20 +233,10 @@ impl ReqwestOAuthTransport { /// # Errors /// /// TLS client 初始化失败时返回错误。 - pub fn new() -> Result { - let endpoint_policy = GrokEndpointPolicy::official_oauth(); - Self::with_endpoint_policy(endpoint_policy) - } - - /// Builds an OAuth transport with an explicit endpoint policy. - /// - /// # Errors - /// - /// Returns an error when the reqwest client cannot be initialized. - pub fn with_endpoint_policy( - endpoint_policy: GrokEndpointPolicy, + pub fn new( + endpoint_policy: Arc, ) -> Result { - let client = build_client(&endpoint_policy, Some(OAUTH_REQUEST_TIMEOUT))?; + let client = endpoint_policy.build_oauth_client(Some(OAUTH_REQUEST_TIMEOUT))?; Ok(Self { client, endpoint_policy, @@ -359,7 +294,7 @@ impl OAuthHttpTransport for ReqwestOAuthTransport { /// 官方 Grok Responses HTTP SSE transport。 pub struct ReqwestGrokInferenceTransport { client: Client, - endpoint_policy: GrokEndpointPolicy, + endpoint_policy: Arc, } impl ReqwestGrokInferenceTransport { @@ -368,20 +303,10 @@ impl ReqwestGrokInferenceTransport { /// # Errors /// /// TLS client 初始化失败时返回错误。 - pub fn new() -> Result { - let endpoint_policy = GrokEndpointPolicy::official_inference(); - Self::with_endpoint_policy(endpoint_policy) - } - - /// Builds an inference transport with an explicit endpoint policy. - /// - /// # Errors - /// - /// Returns an error when the reqwest client cannot be initialized. - pub fn with_endpoint_policy( - endpoint_policy: GrokEndpointPolicy, + pub fn new( + endpoint_policy: Arc, ) -> Result { - let client = build_client(&endpoint_policy, None)?; + let client = endpoint_policy.build_inference_client(None)?; Ok(Self { client, endpoint_policy, @@ -429,6 +354,9 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport { UpstreamSendState::Sent, )); } + let http_version = upstream_http_version(response.version()); + let status_code = response.status().as_u16(); + let request_id = upstream_request_id(&response); let body = response.bytes_stream().map(|chunk| match chunk { Ok(chunk) if chunk.len() <= MAX_SSE_EVENT_BUFFER_BYTES => Ok(chunk.to_vec()), Ok(_) => Err(GrokInferenceTransportError::new( @@ -437,7 +365,12 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport { )), Err(error) => Err(classify_inference_stream_error(&error)), }); - Ok(GrokInferenceResponse::new(Box::pin(body))) + Ok(GrokInferenceResponse::new( + Box::pin(body), + http_version, + status_code, + request_id, + )) }) } } @@ -445,7 +378,7 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport { /// 官方 Grok CLI proxy 模型目录 GET transport。 pub struct ReqwestGrokModelCatalogTransport { client: Client, - endpoint_policy: GrokEndpointPolicy, + endpoint_policy: Arc, } impl ReqwestGrokModelCatalogTransport { @@ -454,19 +387,10 @@ impl ReqwestGrokModelCatalogTransport { /// # Errors /// /// TLS client 初始化失败时返回错误。 - pub fn new() -> Result { - Self::with_endpoint_policy(GrokEndpointPolicy::official_inference()) - } - - /// 使用显式 endpoint policy 构建测试或生产 transport。 - /// - /// # Errors - /// - /// Reqwest client 无法初始化时返回错误。 - pub fn with_endpoint_policy( - endpoint_policy: GrokEndpointPolicy, + pub fn new( + endpoint_policy: Arc, ) -> Result { - let client = build_client(&endpoint_policy, Some(OAUTH_REQUEST_TIMEOUT))?; + let client = endpoint_policy.build_inference_client(Some(OAUTH_REQUEST_TIMEOUT))?; Ok(Self { client, endpoint_policy, @@ -540,12 +464,12 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport { let client = self.client.clone(); let endpoint_policy = self.endpoint_policy.clone(); Box::pin(async move { - if !endpoint_policy.validate_billing(request.endpoint()) { - return Err(GrokBillingTransportError::new( - GrokBillingTransportErrorKind::Protocol, - )); - } - let mut builder = client.get(request.endpoint().clone()); + let endpoint = endpoint_policy + .route_billing(request.endpoint()) + .ok_or_else(|| { + GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol) + })?; + let mut builder = client.get(endpoint); for header in request.headers() { builder = builder.header(header.name(), header.value().expose()); } @@ -575,8 +499,8 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport { } } -pub(crate) fn build_client( - policy: &GrokEndpointPolicy, +fn build_official_client( + host: &'static str, timeout: Option, ) -> Result { let mut builder = Client::builder() @@ -584,16 +508,9 @@ pub(crate) fn build_client( .no_proxy() .connect_timeout(CONNECT_TIMEOUT) .pool_idle_timeout(POOL_IDLE_TIMEOUT) - .tcp_nodelay(true); - match &policy.0 { - EndpointPolicyKind::Official { host } => { - let resolver = StrictDnsResolver::new(host)?; - builder = builder.https_only(true).dns_resolver(Arc::new(resolver)); - } - EndpointPolicyKind::Loopback { .. } => { - builder = builder.https_only(false); - } - } + .tcp_nodelay(true) + .https_only(true) + .dns_resolver(Arc::new(StrictDnsResolver::new(host)?)); if let Some(timeout) = timeout { builder = builder.timeout(timeout); } @@ -602,46 +519,6 @@ pub(crate) fn build_client( .map_err(|_| GrokReqwestTransportBuildError::ClientInitialization) } -impl LoopbackOrigin { - fn parse(url: &Url) -> Result { - let host = url - .host() - .and_then(|host| match host { - url::Host::Ipv4(address) => Some(IpAddr::V4(address)), - url::Host::Ipv6(address) => Some(IpAddr::V6(address)), - url::Host::Domain(_) => None, - }) - .filter(IpAddr::is_loopback) - .ok_or(GrokReqwestTransportBuildError::InvalidEndpointOrigin)?; - if url.scheme() != "http" - || url.path() != "/" - || !url.username().is_empty() - || url.password().is_some() - || url.query().is_some() - || url.fragment().is_some() - { - return Err(GrokReqwestTransportBuildError::InvalidEndpointOrigin); - } - let port = url - .port() - .ok_or(GrokReqwestTransportBuildError::InvalidEndpointOrigin)?; - Ok(Self { host, port }) - } - - fn matches(&self, url: &Url) -> bool { - url.scheme() == "http" - && url.port() == Some(self.port) - && url.host().is_some_and(|host| match host { - url::Host::Ipv4(address) => IpAddr::V4(address) == self.host, - url::Host::Ipv6(address) => IpAddr::V6(address) == self.host, - url::Host::Domain(_) => false, - }) - && url.username().is_empty() - && url.password().is_none() - && url.fragment().is_none() - } -} - fn valid_official_url(url: &Url, host: &str, path: Option<&str>) -> bool { url.scheme() == "https" && url.host_str() == Some(host) @@ -653,6 +530,21 @@ fn valid_official_url(url: &Url, host: &str, path: Option<&str>) -> bool { && url.fragment().is_none() } +fn valid_billing_url(url: &Url, host: &str) -> bool { + url.scheme() == "https" + && url.host_str() == Some(host) + && url.port_or_known_default() == Some(443) + && url.path() == OFFICIAL_BILLING_PATH + && url.username().is_empty() + && url.password().is_none() + && url.fragment().is_none() + && url.query_pairs().count() == 1 + && url + .query_pairs() + .next() + .is_some_and(|(key, value)| key == "format" && value == "credits") +} + #[derive(Debug)] struct StrictDnsResolver { policy: GrokDnsResolutionPolicy, @@ -910,6 +802,8 @@ fn classify_inference_stream_error(error: &reqwest::Error) -> GrokInferenceTrans async fn classify_inference_status(response: Response) -> GrokInferenceTransportError { let status = response.status(); let retry_after = retry_after(&response); + let http_version = upstream_http_version(response.version()); + let request_id = upstream_request_id(&response); let kind = match status { StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND @@ -929,6 +823,7 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport let _ = collect_bounded(response, MAX_ERROR_BODY_BYTES).await; let mut error = GrokInferenceTransportError::new(kind, UpstreamSendState::Sent) .with_status(status) + .with_response_facts(http_version, request_id) .redact_sensitive_context("upstream response body"); if let Some(retry_after) = retry_after { error = error.with_retry_after(retry_after); @@ -936,6 +831,25 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport error } +fn upstream_http_version(version: reqwest::Version) -> UpstreamHttpVersion { + match version { + reqwest::Version::HTTP_09 => UpstreamHttpVersion::Http09, + reqwest::Version::HTTP_10 => UpstreamHttpVersion::Http10, + reqwest::Version::HTTP_11 => UpstreamHttpVersion::Http11, + reqwest::Version::HTTP_2 => UpstreamHttpVersion::Http2, + reqwest::Version::HTTP_3 => UpstreamHttpVersion::Http3, + _ => UpstreamHttpVersion::Unknown, + } +} + +fn upstream_request_id(response: &Response) -> Option { + ["x-request-id", "request-id", "cf-ray"] + .into_iter() + .find_map(|name| response.headers().get(name)) + .and_then(|value| value.to_str().ok()) + .and_then(|value| SafeUpstreamValue::new(value.to_owned()).ok()) +} + async fn classify_model_catalog_status(response: Response) -> GrokModelCatalogTransportError { let status = response.status(); let kind = match status { diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs index ac874f023..e097102eb 100644 --- a/backend/crates/providers/xai/src/transport/request.rs +++ b/backend/crates/providers/xai/src/transport/request.rs @@ -1,35 +1,61 @@ use std::fmt; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD; -use gateway_core::operation::{ - ContentPart, GenerateRequest, ImageSource, MessageRole, OutputFormat, ReasoningEffort, - ReasoningSummary, -}; -use serde_json::{Map, Value, json}; -use url::Url; +use gateway_core::operation::GenerateRequest; +use serde_json::{Map, Value}; use super::XAI_PROVIDER_NAME; -use super::config::valid_header_value; -const GROK_OPTION_FIELDS: &[&str] = &[ +const GROK_OPTION_FIELDS: &[&str] = &["schema_version", "transport"]; +const IDENTITY_FIELDS: &[&str] = &[ + "Authorization", + "authorization", + "Cookie", + "cookie", + "accessToken", + "access_token", + "accountId", + "account_id", + "cookies", + "email", + "idToken", + "id_token", + "refreshToken", + "refresh_token", + "sessionToken", + "session_token", + "teamId", + "team_id", + "token", + "userId", + "user_id", + "x-email", + "x-grok-user-id", + "x-userid", +]; +const ACCOUNT_BOUND_FIELDS: &[&str] = &[ + "agentId", "agent_id", + "conversation", + "conversationId", "conversation_id", - "schema_version", + "previousResponseId", + "previous_response_id", + "responseId", + "response_id", + "sessionId", "session_id", - "transport", + "x-grok-agent-id", + "x-grok-conv-id", + "x-grok-session-id", ]; -const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; -const MAX_IMAGE_URL_BYTES: usize = 8 * 1024; -/// Strictly encoded official Responses request and typed Grok header options. +/// 保留客户端 OpenAI Responses object 的 xAI 上游请求。 pub struct GrokResponsesRequest { body: Map, - header_options: GrokRequestHeaderOptions, } impl GrokResponsesRequest { - /// Returns the typed JSON object sent to `/v1/responses`. + /// 返回发送到 `/v1/responses` 的 JSON object。 #[must_use] pub const fn body(&self) -> &Map { &self.body @@ -39,86 +65,24 @@ impl GrokResponsesRequest { request: &GenerateRequest, upstream_model: &str, ) -> Result { - let mut body = Map::new(); + let payload = request + .protocol_payload() + .filter(|payload| payload.protocol() == "openai") + .ok_or(GrokRequestEncodeError::InvalidProtocolPayload)?; + let mut body = payload.body().clone(); + sanitize_account_identity(&mut body); body.insert("model".to_owned(), Value::String(upstream_model.to_owned())); - body.insert("input".to_owned(), Value::Array(encode_messages(request)?)); body.insert("stream".to_owned(), Value::Bool(true)); - // xAI OAuth session 只执行本轮推理；连续对话由 Gateway portable history 重放。 body.insert("store".to_owned(), Value::Bool(false)); - - if !request.tools().is_empty() { - body.insert( - "tools".to_owned(), - Value::Array( - request - .tools() - .iter() - .map(|tool| { - let mut value = Map::new(); - value.insert("type".to_owned(), Value::String("function".to_owned())); - value.insert("name".to_owned(), Value::String(tool.name().to_owned())); - if let Some(description) = tool.description() { - value.insert( - "description".to_owned(), - Value::String(description.to_owned()), - ); - } - value.insert( - "parameters".to_owned(), - Value::Object(tool.input_schema().clone()), - ); - value.insert("strict".to_owned(), Value::Bool(tool.strict())); - Value::Object(value) - }) - .collect(), - ), - ); + if let Some(options) = request.provider_options().get(XAI_PROVIDER_NAME) { + validate_provider_options(options)?; } - if !matches!(request.output_format(), OutputFormat::Text) { - body.insert( - "text".to_owned(), - encode_output_format(request.output_format()), - ); - } - if let Some(reasoning) = request.reasoning() { - let mut value = Map::new(); - if let Some(effort) = reasoning.effort { - value.insert( - "effort".to_owned(), - Value::String(reasoning_effort(effort).to_owned()), - ); - } - if let Some(summary) = reasoning.summary { - value.insert( - "summary".to_owned(), - Value::String(reasoning_summary(summary).to_owned()), - ); - } - body.insert("reasoning".to_owned(), Value::Object(value)); - } - if let Some(tokens) = request.max_output_tokens() { - body.insert("max_output_tokens".to_owned(), Value::from(tokens)); - } - - let header_options = request - .provider_options() - .get(XAI_PROVIDER_NAME) - .map(parse_provider_options) - .transpose()? - .unwrap_or_default(); - Ok(Self { - body, - header_options, - }) + Ok(Self { body }) } pub(crate) fn to_json_bytes(&self) -> Result, GrokRequestEncodeError> { serde_json::to_vec(&self.body).map_err(|_| GrokRequestEncodeError::Serialization) } - - pub(crate) const fn header_options(&self) -> &GrokRequestHeaderOptions { - &self.header_options - } } impl fmt::Debug for GrokResponsesRequest { @@ -127,42 +91,6 @@ impl fmt::Debug for GrokResponsesRequest { .debug_struct("GrokResponsesRequest") .field("body_keys", &self.body.keys().collect::>()) .field("body", &"") - .field("header_options", &self.header_options) - .finish() - } -} - -#[derive(Clone, Default)] -pub(crate) struct GrokRequestHeaderOptions { - conversation_id: Option, - session_id: Option, - agent_id: Option, -} - -impl GrokRequestHeaderOptions { - pub(crate) fn conversation_id(&self) -> Option<&str> { - self.conversation_id.as_deref() - } - - pub(crate) fn session_id(&self) -> Option<&str> { - self.session_id.as_deref() - } - - pub(crate) fn agent_id(&self) -> Option<&str> { - self.agent_id.as_deref() - } -} - -impl fmt::Debug for GrokRequestHeaderOptions { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("GrokRequestHeaderOptions") - .field( - "conversation_id", - &self.conversation_id.as_ref().map(|_| ""), - ) - .field("session_id", &self.session_id.as_ref().map(|_| "")) - .field("agent_id", &self.agent_id.as_ref().map(|_| "")) .finish() } } @@ -171,146 +99,21 @@ impl fmt::Debug for GrokRequestHeaderOptions { /// values. #[derive(Debug, thiserror::Error, Clone, Copy, PartialEq, Eq)] pub enum GrokRequestEncodeError { + /// 数据面只接受 OpenAI adapter 保留的原始 Responses object。 + #[error("Grok Build request is missing its OpenAI protocol payload")] + InvalidProtocolPayload, /// Grok provider option schema or value is malformed. #[error("Grok Build provider option schema is invalid")] InvalidProviderOptions, /// A Grok-specific option is unknown or unsupported. #[error("Grok Build provider option is unsupported")] UnsupportedProviderOption, - /// Canonical content cannot be represented by this Responses adapter. - #[error("Grok Build request contains unsupported content")] - UnsupportedContent, - /// Image URL, MIME type, or size is unsafe. - #[error("Grok Build request contains an invalid image")] - InvalidImage, - /// Typed JSON serialization unexpectedly failed. + /// JSON serialization unexpectedly failed. #[error("Grok Build request serialization failed")] Serialization, } -fn encode_messages(request: &GenerateRequest) -> Result, GrokRequestEncodeError> { - let mut input = Vec::new(); - for message in request.messages() { - let mut content = Vec::new(); - for part in message.content() { - match part { - ContentPart::Text(text) => content.push(json!({ - "type": match message.role() { - MessageRole::Assistant => "output_text", - MessageRole::System | MessageRole::Developer | MessageRole::User => { - "input_text" - } - }, - "text": text, - })), - ContentPart::Image(source) => { - if message.role() != MessageRole::User { - return Err(GrokRequestEncodeError::UnsupportedContent); - } - content.push(json!({ - "type": "input_image", - "image_url": encode_image(source)?, - })); - } - ContentPart::ToolResult { call_id, output } => { - input.push(json!({ - "type": "function_call_output", - "call_id": call_id, - "output": output, - })); - } - ContentPart::ToolCall { - call_id, - name, - arguments, - } => { - if message.role() != MessageRole::Assistant { - return Err(GrokRequestEncodeError::UnsupportedContent); - } - input.push(json!({ - "type": "function_call", - "call_id": call_id, - "name": name, - "arguments": arguments, - })); - } - _ => return Err(GrokRequestEncodeError::UnsupportedContent), - } - } - if !content.is_empty() { - input.push(json!({ - "type": "message", - "role": message_role(message.role()), - "content": content, - })); - } - } - if input.is_empty() { - return Err(GrokRequestEncodeError::UnsupportedContent); - } - Ok(input) -} - -fn encode_image(source: &ImageSource) -> Result { - match source { - ImageSource::Url(value) => { - if value.len() > MAX_IMAGE_URL_BYTES { - return Err(GrokRequestEncodeError::InvalidImage); - } - let url = Url::parse(value).map_err(|_| GrokRequestEncodeError::InvalidImage)?; - if url.scheme() != "https" - || url.host().is_none() - || !url.username().is_empty() - || url.password().is_some() - || url.fragment().is_some() - { - return Err(GrokRequestEncodeError::InvalidImage); - } - Ok(value.clone()) - } - ImageSource::Bytes { media_type, data } => { - if data.is_empty() - || data.len() > MAX_IMAGE_BYTES - || !matches!( - media_type.as_str(), - "image/gif" | "image/jpeg" | "image/png" | "image/webp" - ) - { - return Err(GrokRequestEncodeError::InvalidImage); - } - Ok(format!( - "data:{media_type};base64,{}", - STANDARD.encode(data) - )) - } - } -} - -fn encode_output_format(format: &OutputFormat) -> Value { - let format = match format { - OutputFormat::Text => json!({"type": "text"}), - OutputFormat::JsonObject => json!({"type": "json_object"}), - OutputFormat::JsonSchema(schema) => { - let mut value = Map::new(); - value.insert("type".to_owned(), Value::String("json_schema".to_owned())); - value.insert("name".to_owned(), Value::String(schema.name().to_owned())); - if let Some(description) = schema.description() { - value.insert( - "description".to_owned(), - Value::String(description.to_owned()), - ); - } - value.insert("schema".to_owned(), Value::Object(schema.schema().clone())); - value.insert("strict".to_owned(), Value::Bool(schema.strict())); - Value::Object(value) - } - }; - json!({"format": format}) -} - -fn parse_provider_options( - options: &Map, -) -> Result { +fn validate_provider_options(options: &Map) -> Result<(), GrokRequestEncodeError> { if options .keys() .any(|field| !GROK_OPTION_FIELDS.contains(&field.as_str())) @@ -325,48 +128,17 @@ fn parse_provider_options( { return Err(GrokRequestEncodeError::InvalidProviderOptions); } - Ok(GrokRequestHeaderOptions { - conversation_id: optional_header_value(options, "conversation_id")?, - session_id: optional_header_value(options, "session_id")?, - agent_id: optional_header_value(options, "agent_id")?, - }) + Ok(()) } -fn optional_header_value( - options: &Map, - field: &str, -) -> Result, GrokRequestEncodeError> { - match options.get(field) { - None => Ok(None), - Some(Value::String(value)) if valid_header_value(value, 256) => Ok(Some(value.clone())), - Some(_) => Err(GrokRequestEncodeError::InvalidProviderOptions), +fn sanitize_account_identity(body: &mut Map) { + for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { + body.remove(*field); } -} - -const fn message_role(role: MessageRole) -> &'static str { - match role { - MessageRole::System => "system", - MessageRole::Developer => "developer", - MessageRole::User => "user", - MessageRole::Assistant => "assistant", - } -} - -const fn reasoning_effort(effort: ReasoningEffort) -> &'static str { - match effort { - ReasoningEffort::Minimal => "minimal", - ReasoningEffort::Low => "low", - ReasoningEffort::Medium => "medium", - ReasoningEffort::High => "high", - ReasoningEffort::ExtraHigh => "xhigh", - } -} - -const fn reasoning_summary(summary: ReasoningSummary) -> &'static str { - match summary { - ReasoningSummary::Auto => "auto", - ReasoningSummary::Concise => "concise", - ReasoningSummary::Detailed => "detailed", - ReasoningSummary::None => "none", + let Some(Value::Object(metadata)) = body.get_mut("metadata") else { + return; + }; + for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { + metadata.remove(*field); } } diff --git a/backend/crates/providers/xai/src/transport/session.rs b/backend/crates/providers/xai/src/transport/session.rs index 7562fbf95..e82d6c85e 100644 --- a/backend/crates/providers/xai/src/transport/session.rs +++ b/backend/crates/providers/xai/src/transport/session.rs @@ -2,7 +2,7 @@ use std::collections::BTreeSet; use std::fmt; use std::future::Future; use std::pin::Pin; -use std::time::Duration; +use std::time::{Duration, SystemTime}; use gateway_core::engine::credential::{ AccountSelectionPolicy, CredentialRevision, ProviderAccountId, @@ -167,6 +167,7 @@ pub struct GrokSessionSelection { excluded_accounts: BTreeSet, required_account: Option, account_selection_policy: AccountSelectionPolicy, + deadline: SystemTime, } impl GrokSessionSelection { @@ -178,6 +179,7 @@ impl GrokSessionSelection { excluded_accounts: BTreeSet, required_account: Option, account_selection_policy: AccountSelectionPolicy, + deadline: SystemTime, ) -> Self { Self { provider_instance_id, @@ -185,6 +187,7 @@ impl GrokSessionSelection { excluded_accounts, required_account, account_selection_policy, + deadline, } } @@ -217,6 +220,12 @@ impl GrokSessionSelection { pub const fn account_selection_policy(&self) -> AccountSelectionPolicy { self.account_selection_policy } + + /// Returns the absolute deadline that bounds the scheduling lease. + #[must_use] + pub const fn deadline(&self) -> SystemTime { + self.deadline + } } /// Future returned by a Grok session selector. diff --git a/backend/crates/providers/xai/src/transport/transport.rs b/backend/crates/providers/xai/src/transport/transport.rs index 15f49c8eb..0801619fd 100644 --- a/backend/crates/providers/xai/src/transport/transport.rs +++ b/backend/crates/providers/xai/src/transport/transport.rs @@ -5,6 +5,8 @@ use std::time::Duration; use futures::Stream; use gateway_core::engine::UpstreamSendState; +use gateway_core::error::SafeUpstreamValue; +use gateway_core::event::UpstreamHttpVersion; use url::Url; use zeroize::Zeroizing; @@ -81,13 +83,41 @@ pub type GrokInferenceChunkStream = /// [`GrokInferenceTransportError`] instead. pub struct GrokInferenceResponse { body: GrokInferenceChunkStream, + http_version: UpstreamHttpVersion, + status_code: u16, + request_id: Option, } impl GrokInferenceResponse { /// Wraps one accepted SSE response body. #[must_use] - pub fn new(body: GrokInferenceChunkStream) -> Self { - Self { body } + pub fn new( + body: GrokInferenceChunkStream, + http_version: UpstreamHttpVersion, + status_code: u16, + request_id: Option, + ) -> Self { + Self { + body, + http_version, + status_code, + request_id, + } + } + + #[must_use] + pub const fn http_version(&self) -> UpstreamHttpVersion { + self.http_version + } + + #[must_use] + pub const fn status_code(&self) -> u16 { + self.status_code + } + + #[must_use] + pub const fn request_id(&self) -> Option<&SafeUpstreamValue> { + self.request_id.as_ref() } pub fn into_body(self) -> GrokInferenceChunkStream { @@ -99,6 +129,9 @@ impl fmt::Debug for GrokInferenceResponse { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("GrokInferenceResponse") + .field("http_version", &self.http_version) + .field("status_code", &self.status_code) + .field("request_id", &self.request_id) .field("body", &"[SSE STREAM]") .finish() } @@ -138,6 +171,8 @@ pub struct GrokInferenceTransportError { send_state: UpstreamSendState, status: Option, retry_after: Option, + http_version: Option, + request_id: Option, sensitive_context_redacted: bool, } @@ -150,6 +185,8 @@ impl GrokInferenceTransportError { send_state, status: None, retry_after: None, + http_version: None, + request_id: None, sensitive_context_redacted: false, } } @@ -170,6 +207,17 @@ impl GrokInferenceTransportError { self } + #[must_use] + pub fn with_response_facts( + mut self, + http_version: UpstreamHttpVersion, + request_id: Option, + ) -> Self { + self.http_version = Some(http_version); + self.request_id = request_id; + self + } + /// Discards a possibly sensitive upstream body while retaining that fact. #[must_use] pub fn redact_sensitive_context(mut self, _value: impl AsRef) -> Self { @@ -201,6 +249,16 @@ impl GrokInferenceTransportError { self.retry_after } + #[must_use] + pub const fn http_version(&self) -> Option { + self.http_version + } + + #[must_use] + pub const fn request_id(&self) -> Option<&SafeUpstreamValue> { + self.request_id.as_ref() + } + /// Reports whether a sensitive body was discarded. #[must_use] pub const fn sensitive_context_was_redacted(&self) -> bool { @@ -216,6 +274,8 @@ impl fmt::Debug for GrokInferenceTransportError { .field("send_state", &self.send_state) .field("status", &self.status) .field("retry_after", &self.retry_after) + .field("http_version", &self.http_version) + .field("request_id", &self.request_id) .field( "sensitive_context", &self.sensitive_context_redacted.then_some("[REDACTED]"), diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs new file mode 100644 index 000000000..9dd44bab6 --- /dev/null +++ b/backend/crates/providers/xai/tests/admin.rs @@ -0,0 +1,615 @@ +use std::collections::BTreeMap; +use std::num::NonZeroU32; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime}; + +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, Utc}; +use futures::future::BoxFuture; +use gateway_admin::model::accounts::{ + AccountAvailability as AdminAccountAvailability, AccountRecord, +}; +use gateway_admin::model::observability::{CurrencyCost, ProviderBillingInput}; +use gateway_admin::model::provider_credentials::{ + CompleteAuthorization, PrepareCredentialImport, PrepareCredentialRefresh, + PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput, +}; +use gateway_admin::model::{MutationActor, MutationContext, Revision}; +use gateway_admin::ports::provider::ProviderAdminErrorKind; +use gateway_core::engine::credential::{ + AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccount, + ProviderAccountId, ProviderAccountStore, +}; +use gateway_core::provider_ports::{ + NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome, + OAuthPendingTakeOutcome, ProviderCatalogCacheKey, ProviderCatalogCachePort, + ProviderCatalogPorts, ProviderCooldown, ProviderCooldownPort, ProviderCredentialState, + ProviderCredentialStatePort, ProviderInstanceCatalogPort, ProviderInstanceConfig, + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, + ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts, +}; +use gateway_core::routing::ProviderKind; +use gateway_core::task::WorkerKind; +use provider_xai::{ + DiscoveryDocument, GrokOAuthConfig, PendingAuthorization, RedirectUriAllowlist, XaiConfig, +}; +use serde_json::{Map, Value, json}; +use sha2::{Digest as _, Sha256}; + +use crate::support::{MemoryProviderAccountStore, create_input, instance_id, seed_input}; + +#[tokio::test] +async fn xai_bundle_exposes_core_admin_and_drains_two_worker_contributions_once() { + let mut bundle = provider_xai::initialize(XaiConfig::default(), provider_ports()) + .await + .expect("xAI bundle"); + + assert_eq!(bundle.core_provider().name(), "xai"); + assert_eq!(bundle.admin_provider().provider_kind().as_str(), "xai"); + let contributions = bundle.take_worker_contributions(); + assert_eq!(contributions.len(), 2); + assert!( + contributions + .iter() + .any(|item| item.kind() == WorkerKind::OAuthRefresh) + ); + assert!( + contributions + .iter() + .any(|item| item.kind() == WorkerKind::QuotaCatalogHealth) + ); + assert!(bundle.take_worker_contributions().is_empty()); +} + +#[tokio::test] +async fn xai_admin_provider_validates_known_billing_breakdown() { + let bundle = provider_xai::initialize(XaiConfig::default(), provider_ports()) + .await + .expect("xAI bundle"); + let billing = bundle + .admin_provider() + .calculated_billing(&ProviderBillingInput { + upstream_model_id: "grok-4.5".to_owned(), + input_tokens: Some(100), + output_tokens: Some(10), + cached_tokens: Some(25), + cache_write_tokens: Some(0), + total: CurrencyCost { + currency: "USD".to_owned(), + amount: "0.0002175".parse().expect("amount"), + }, + }) + .expect("billing") + .expect("known pricing"); + + assert_eq!(billing.total_amount.amount.as_str(), "0.0002175"); + assert_eq!(billing.input_price_per_million.amount.as_str(), "2"); + assert_eq!(billing.output_price_per_million.amount.as_str(), "6"); +} + +#[tokio::test] +async fn xai_admin_provider_restores_full_pending_envelope_and_binds_owner() { + let pending = Arc::new(TestOAuthPending::default()); + let bundle = provider_xai::initialize( + XaiConfig::default(), + provider_ports_with( + Arc::new(MemoryProviderAccountStore::default()), + Arc::clone(&pending), + ), + ) + .await + .expect("xAI bundle"); + let owner = MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "admin-owner".to_owned(), + }, + request_id: "request-start".to_owned(), + }; + let flow_id = "pending-envelope"; + let owner_ref = admin_session_owner_ref("admin-owner"); + let payload = pending_payload(flow_id, &owner_ref); + pending.insert(flow_id, owner_ref, payload.clone()); + + let mutation = payload + .expose_to_provider() + .get("mutation") + .and_then(Value::as_object) + .expect("pending mutation"); + assert_eq!( + mutation + .get("expected_config_revision") + .and_then(Value::as_u64), + Some(7) + ); + assert_eq!( + mutation.get("started_request_id").and_then(Value::as_str), + Some("request-start") + ); + + let wrong_owner = bundle + .admin_provider() + .complete_authorization(CompleteAuthorization { + context: MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "different-owner".to_owned(), + }, + request_id: "request-complete".to_owned(), + }, + flow_id: flow_id.to_owned(), + callback_url: format!( + "{}?code=unused&state=wrong-owner", + provider_xai::OFFICIAL_REDIRECT_URI + ), + }) + .await + .expect_err("wrong owner"); + assert_eq!(wrong_owner.kind(), ProviderAdminErrorKind::NotFound); + assert_eq!(pending.len(), 1); + + let invalid_callback = bundle + .admin_provider() + .complete_authorization(CompleteAuthorization { + context: owner, + flow_id: flow_id.to_owned(), + callback_url: format!( + "{}?code=unused&state=invalid-state", + provider_xai::OFFICIAL_REDIRECT_URI + ), + }) + .await + .expect_err("restored pending must validate callback state"); + assert_eq!(invalid_callback.kind(), ProviderAdminErrorKind::Invalid); + assert_eq!(pending.len(), 0); + + let mismatched_flow = "pending-envelope-mismatch"; + let owner_ref = admin_session_owner_ref("admin-owner"); + pending.insert( + mismatched_flow, + owner_ref.clone(), + pending_payload("different-envelope", &owner_ref), + ); + let mismatched_envelope = bundle + .admin_provider() + .complete_authorization(CompleteAuthorization { + context: MutationContext { + actor: MutationActor::AdminSession { + admin_user_id: "admin-owner".to_owned(), + }, + request_id: "request-complete-mismatch".to_owned(), + }, + flow_id: mismatched_flow.to_owned(), + callback_url: format!( + "{}?code=unused&state=invalid-state", + provider_xai::OFFICIAL_REDIRECT_URI + ), + }) + .await + .expect_err("flow binding mismatch"); + assert_eq!(mismatched_envelope.kind(), ProviderAdminErrorKind::Invalid); + assert_eq!(pending.len(), 0); +} + +#[tokio::test] +async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() { + let store = Arc::new(MemoryProviderAccountStore::default()); + let input = create_input("admin_projection", "subject-admin-projection"); + seed_input(&store, &input).await.expect("create account"); + let account = store.account(&input.account_id).expect("stored account"); + let record = account_record(&account); + let bundle = provider_xai::initialize( + XaiConfig::default(), + provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), + ) + .await + .expect("xAI bundle"); + let admin = bundle.admin_provider(); + + let quota = admin + .quota(account.id(), false) + .await + .expect("cached quota"); + assert!(quota.windows.is_empty()); + let models = admin + .models(account.id(), false) + .await + .expect("cached models"); + assert!(models.models.is_empty()); + let loaded = store + .load_credential(account.id(), account.revision()) + .await + .expect("loaded credential"); + let exported = admin + .export_credentials(vec![ProviderExportCredentialInput { + account: record, + provider_material: ProviderDocument::new(OpaqueProviderData::new( + loaded.credential.into_inner(), + )), + }]) + .await + .expect("canonical export"); + assert_eq!(exported.account_ids, vec![input.account_id]); + let document = exported.document.expose_to_provider().expose_to_provider(); + assert_eq!(document.get("version").and_then(Value::as_u64), Some(1)); + assert_eq!( + document.get("type").and_then(Value::as_str), + Some("oauth-account-bundle") + ); +} + +#[tokio::test] +async fn xai_admin_provider_rejects_unprepared_mutations_before_store_commit() { + let store = Arc::new(MemoryProviderAccountStore::default()); + let input = create_input("admin_invalid", "subject-admin-invalid"); + seed_input(&store, &input).await.expect("create account"); + let account = store.account(&input.account_id).expect("stored account"); + let record = account_record(&account); + let bundle = provider_xai::initialize( + XaiConfig::default(), + provider_ports_with(store, Arc::new(TestOAuthPending::default())), + ) + .await + .expect("xAI bundle"); + let admin = bundle.admin_provider(); + + let import_error = admin + .prepare_import(PrepareCredentialImport { + provider_instance_id: instance_id(), + document: ProviderDocument::new(OpaqueProviderData::new(Map::new())), + }) + .await + .expect_err("invalid import"); + assert_eq!(import_error.kind(), ProviderAdminErrorKind::Invalid); + let rotation_error = admin + .prepare_rotation(PrepareCredentialRotation { + account: record.clone(), + expected_credential_revision: record.credential_revision, + provider_material: ProviderDocument::new(OpaqueProviderData::new(Map::new())), + }) + .await + .expect_err("invalid rotation"); + assert_eq!(rotation_error.kind(), ProviderAdminErrorKind::Invalid); + let mut missing = record; + missing.id = "acct_admin_missing".to_owned(); + let refresh_error = admin + .prepare_refresh(PrepareCredentialRefresh { account: missing }) + .await + .expect_err("missing refresh target"); + assert_eq!(refresh_error.kind(), ProviderAdminErrorKind::NotFound); +} + +fn pending_payload(flow_id: &str, owner_ref: &str) -> OpaqueProviderData { + let config = GrokOAuthConfig::official(provider_xai::transport::GROK_CLIENT_VERSION) + .expect("official config"); + let discovery = DiscoveryDocument::parse( + &config, + include_bytes!("credential/fixtures/discovery.json"), + ) + .expect("discovery fixture"); + let redirect = RedirectUriAllowlist::new([provider_xai::OFFICIAL_REDIRECT_URI]) + .expect("redirect allowlist") + .authorize(provider_xai::OFFICIAL_REDIRECT_URI) + .expect("official redirect"); + let pending = PendingAuthorization::start(&config, &discovery, redirect, None) + .expect("pending authorization"); + let server_state = pending + .into_server_state() + .expect("server state") + .expose() + .to_owned(); + let value = json!({ + "schema_version": 1, + "flow_id": flow_id, + "owner_ref": owner_ref, + "expires_at": (Utc::now() + chrono::Duration::minutes(10)).to_rfc3339(), + "server_state": server_state, + "mutation": { + "expected_config_revision": 7, + "provider_kind": "xai", + "target": { + "kind": "create", + "provider_instance_id": instance_id().to_string(), + "name": "OAuth account" + }, + "owner": { + "kind": "admin_session", + "admin_user_id": "admin-owner" + }, + "started_request_id": "request-start" + } + }); + OpaqueProviderData::new(value.as_object().expect("pending object").clone()) +} + +fn admin_session_owner_ref(admin_user_id: &str) -> String { + let mut digest = Sha256::new(); + digest.update(b"admin-session\0"); + digest.update(admin_user_id.as_bytes()); + URL_SAFE_NO_PAD.encode(digest.finalize()) +} + +fn provider_ports() -> ProviderStorePorts { + provider_ports_with( + Arc::new(MemoryProviderAccountStore::default()), + Arc::new(TestOAuthPending::default()), + ) +} + +fn provider_ports_with( + accounts: Arc, + pending: Arc, +) -> ProviderStorePorts { + ProviderStorePorts::new( + accounts, + Arc::new(TestLeases), + ProviderCatalogPorts::new( + Arc::new(TestInstanceCatalog), + Arc::new(TestCatalogCache::default()), + ), + Arc::new(TestCredentialState), + Arc::new(TestCooldown), + Arc::new(TestRuntimePolicy), + pending, + ) +} + +fn account_record(account: &ProviderAccount) -> AccountRecord { + let now = Utc::now(); + AccountRecord { + id: account.id().to_string(), + provider_instance_id: account.instance().clone(), + provider_kind: account.provider().clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + upstream_user_id: account.upstream_user_id().to_owned(), + upstream_account_id: account.upstream_account_id().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + credential_revision: Revision::new(account.revision().get()).expect("revision"), + has_refresh_token: account.has_refresh_token(), + access_token_expires_at: DateTime::::from(account.access_token_expires_at()), + next_refresh_at: account.next_refresh_at().map(DateTime::::from), + enabled: account.enabled(), + availability: admin_availability(account.availability()), + availability_reason: None, + cooldown_until: account.cooldown_until().map(DateTime::::from), + availability_observed_at: now, + quota_observed_at: None, + created_at: now, + updated_at: now, + } +} + +const fn admin_availability(value: AccountAvailability) -> AdminAccountAvailability { + match value { + AccountAvailability::Unknown => AdminAccountAvailability::Unknown, + AccountAvailability::Ready => AdminAccountAvailability::Ready, + AccountAvailability::Cooldown => AdminAccountAvailability::Cooldown, + AccountAvailability::QuotaExhausted => AdminAccountAvailability::QuotaExhausted, + AccountAvailability::Expired => AdminAccountAvailability::Expired, + AccountAvailability::Banned => AdminAccountAvailability::Banned, + AccountAvailability::Invalid => AdminAccountAvailability::Invalid, + } +} + +struct TestLeases; + +impl ProviderLeasePort for TestLeases { + fn load_state<'a>( + &'a self, + _: &'a gateway_core::routing::ProviderInstanceId, + _: &'a [ProviderAccountId], + ) -> BoxFuture< + 'a, + Result, + > { + Box::pin(async { + Ok(gateway_core::provider_ports::ProviderSchedulingState::new( + Default::default(), + None, + 0, + )) + }) + } + + fn try_acquire( + &self, + _request: ProviderLeaseRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(ProviderLeaseAcquisition::Acquired(Box::new(()))) }) + } +} + +struct TestInstanceCatalog; + +impl ProviderInstanceCatalogPort for TestInstanceCatalog { + fn list_instances<'a>( + &'a self, + provider_kind: &'a ProviderKind, + _include_disabled: bool, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(vec![ProviderInstanceConfig::new( + instance_id(), + provider_kind.clone(), + provider_xai::GROK_CLI_BASE_URL.to_owned(), + true, + )]) + }) + } +} + +#[derive(Default)] +struct TestCatalogCache { + values: Mutex>, +} + +impl ProviderCatalogCachePort for TestCatalogCache { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + _ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>> { + Box::pin(async move { + self.values + .lock() + .expect("catalog cache") + .insert(key.account_id().to_string(), catalog.clone()); + Ok(()) + }) + } + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(self + .values + .lock() + .expect("catalog cache") + .get(key.account_id().as_str()) + .cloned()) + }) + } +} + +struct TestCredentialState; + +impl ProviderCredentialStatePort for TestCredentialState { + fn replace( + &self, + _state: ProviderCredentialState, + ) -> BoxFuture<'_, Result<(), ProviderStoreError>> { + Box::pin(async { Ok(()) }) + } + + fn read<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async { Ok(None) }) + } + + fn clear<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(false) }) + } +} + +struct TestCooldown; + +impl ProviderCooldownPort for TestCooldown { + fn put_if_later( + &self, + _cooldown: ProviderCooldown, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(false) }) + } + + fn read<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async { Ok(None) }) + } + + fn clear<'a>( + &'a self, + _account_id: &'a ProviderAccountId, + _through_revision: CredentialRevision, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(false) }) + } +} + +struct TestRuntimePolicy; + +impl ProviderRuntimePolicyPort for TestRuntimePolicy { + fn load_refresh_policy( + &self, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + ProviderRefreshPolicy::try_new( + Duration::from_secs(300), + NonZeroU32::new(4).expect("nonzero concurrency"), + ) + }) + } +} + +#[derive(Default)] +struct TestOAuthPending { + values: Mutex>, +} + +type PendingKey = (String, String); +type PendingValue = (String, OpaqueProviderData, SystemTime); + +impl TestOAuthPending { + fn insert(&self, flow_id: &str, owner: String, payload: OpaqueProviderData) { + self.values.lock().expect("OAuth pending").insert( + ("xai".to_owned(), flow_id.to_owned()), + (owner, payload, SystemTime::now() + Duration::from_secs(600)), + ); + } + + fn len(&self) -> usize { + self.values.lock().expect("OAuth pending").len() + } +} + +impl OAuthPendingFlowPort for TestOAuthPending { + fn put_if_absent( + &self, + flow: NewOAuthPendingFlow, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let key = ( + flow.provider_kind().as_str().to_owned(), + flow.flow().expose_to_store().to_owned(), + ); + let mut values = self.values.lock().expect("OAuth pending"); + if values.contains_key(&key) { + return Ok(OAuthPendingPutOutcome::AlreadyExists); + } + values.insert( + key, + ( + flow.owner().expose_to_store().to_owned(), + flow.payload().clone(), + SystemTime::now() + flow.ttl(), + ), + ); + Ok(OAuthPendingPutOutcome::Stored) + }) + } + + fn take_if_owner<'a>( + &'a self, + provider_kind: &'a ProviderKind, + flow: &'a OAuthPendingBinding, + owner: &'a OAuthPendingBinding, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let key = ( + provider_kind.as_str().to_owned(), + flow.expose_to_store().to_owned(), + ); + let mut values = self.values.lock().expect("OAuth pending"); + let Some((stored_owner, _, expires_at)) = values.get(&key) else { + return Ok(OAuthPendingTakeOutcome::NotFound); + }; + if expires_at <= &SystemTime::now() { + values.remove(&key); + return Ok(OAuthPendingTakeOutcome::NotFound); + } + if stored_owner != owner.expose_to_store() { + return Ok(OAuthPendingTakeOutcome::OwnerMismatch); + } + let (_, payload, _) = values.remove(&key).expect("checked pending"); + Ok(OAuthPendingTakeOutcome::Taken(payload)) + }) + } +} diff --git a/backend/crates/providers/xai/tests/config.rs b/backend/crates/providers/xai/tests/config.rs new file mode 100644 index 000000000..33b048cea --- /dev/null +++ b/backend/crates/providers/xai/tests/config.rs @@ -0,0 +1,13 @@ +use std::path::Path; + +use provider_xai::XaiConfig; + +#[test] +fn xai_config_accepts_only_the_fixed_official_protocol_shape() { + let mut config: XaiConfig = serde_json::from_str("{}").expect("empty xAI config"); + config + .resolve_and_validate(Path::new("/srv/gateway")) + .expect("fixed official protocol"); + + assert!(serde_json::from_str::(r#"{"client_version":"override"}"#).is_err()); +} diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 9c5b10952..7c7e3d734 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -1,11 +1,13 @@ use std::collections::VecDeque; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; +use futures::future::join_all; use gateway_core::engine::credential::{ - AccountAvailability, AccountStateChange, CredentialRevision, OpaqueProviderData, - ProviderAccountStore, QuotaObservation, + AccountAvailability, AccountStateChange, CredentialCasUpdate, CredentialRevision, + OpaqueProviderData, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, + QuotaObservation, }; use gateway_core::routing::{ConfigRevision, InstanceHealth, ProviderInstance, ProviderKind}; use provider_xai::{ @@ -102,6 +104,40 @@ impl GrokBillingTransport for QueueBillingTransport { } } +enum BillingMutation { + State(AccountStateChange), + Credential(CredentialCasUpdate), +} + +struct MutatingBillingTransport { + store: Arc, + mutation: Mutex>, + body: Vec, +} + +impl GrokBillingTransport for MutatingBillingTransport { + fn execute(&self, _: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + let store = Arc::clone(&self.store); + let mutation = self.mutation.lock().expect("mutation").take(); + let body = self.body.clone(); + Box::pin(async move { + match mutation.expect("one mutation per request") { + BillingMutation::State(change) => store + .apply_state_change(change) + .await + .expect("apply concurrent state"), + BillingMutation::Credential(update) => { + store + .compare_and_swap_credential(update) + .await + .expect("apply concurrent credential rotation"); + } + } + Ok(GrokBillingTransportResponse::new(body)) + }) + } +} + fn instance() -> ProviderInstance { ProviderInstance::new( instance_id(), @@ -126,15 +162,53 @@ async fn repository_with_accounts( (store, repository) } +async fn set_account_state( + store: &MemoryProviderAccountStore, + id: &ProviderAccountId, + availability: AccountAvailability, + cooldown_until: Option, +) { + store + .apply_state_change(AccountStateChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + availability, + reason: Some("test isolation".to_owned()), + cooldown_until, + observed_at: SystemTime::now(), + }) + .await + .expect("set account state"); +} + +#[tokio::test] +async fn concurrent_cold_scheduling_hydration_reads_quota_once() { + let (store, repository) = + repository_with_accounts(&[("quota-hydration", "subject-hydration")]).await; + let account = store + .account(&account_id("quota-hydration")) + .expect("created account"); + let service = GrokCredentialQuotaService::new(repository, QueueBillingTransport::failure()); + + join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await; + + assert_eq!(store.quota_reads(), 1); +} + #[tokio::test] async fn synchronization_caches_each_account_and_returns_strict_union() { let (_, repository) = repository_with_accounts(&[("catalog-a", "subject-a"), ("catalog-b", "subject-b")]).await; let cache = MemoryGrokCatalogCache::shared(); let cache_port: Arc = cache.clone(); - let transport = - QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec(), OFFICIAL_FIXTURE.to_vec()]); + let transport = QueueCatalogTransport::from_bodies([ + OFFICIAL_FIXTURE.to_vec(), + OFFICIAL_FIXTURE.to_vec(), + OFFICIAL_FIXTURE.to_vec(), + OFFICIAL_FIXTURE.to_vec(), + ]); let service = GrokCredentialCatalogService::new(repository, transport, cache_port); + assert_eq!(service.catalog_generation().get(), 0); let snapshot = service .synchronize_instance( &instance(), @@ -142,20 +216,27 @@ async fn synchronization_caches_each_account_and_returns_strict_union() { ) .await .expect("catalog sync"); + assert_eq!(service.catalog_generation().get(), 1); + service + .query_instance_models(&instance()) + .await + .expect("same catalog sync"); + assert_eq!(service.catalog_generation().get(), 1); assert_eq!(snapshot.accounts().len(), 2); assert_eq!(snapshot.models().len(), 1); assert_eq!(snapshot.models()[0].request_model().as_str(), "grok-4.5"); for id in [account_id("catalog-a"), account_id("catalog-b")] { - assert!( + assert_eq!( cache - .permits( + .observed_model_support( &id, gateway_core::engine::credential::CredentialRevision::new(1).expect("revision"), "grok-4.5", ) .await - .expect("cache lookup") + .expect("cache lookup"), + Some(true) ); } } @@ -342,9 +423,10 @@ async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fie .await .expect("refresh quota"); let persisted = store - .get_quota(&account_id("quota")) + .get_quotas(&[account_id("quota")]) .await .expect("read persisted quota") + .pop() .expect("quota exists"); let document = persisted.quota.expect("provider quota"); @@ -366,6 +448,144 @@ async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fie ); } +#[tokio::test] +async fn recovered_quota_refresh_releases_existing_quota_and_cooldown_isolation() { + for (suffix, availability, cooldown_until) in [ + ("recover-quota", AccountAvailability::QuotaExhausted, None), + ( + "recover-cooldown", + AccountAvailability::Cooldown, + SystemTime::now().checked_add(Duration::from_secs(300)), + ), + ] { + let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; + let id = account_id(suffix); + set_account_state(&store, &id, availability, cooldown_until).await; + GrokCredentialQuotaService::new( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh recovered quota"); + + let account = store.account(&id).expect("account"); + assert_eq!(account.availability(), AccountAvailability::Ready); + assert_eq!(account.cooldown_until(), None); + } +} + +#[tokio::test] +async fn authoritative_exhaustion_preserves_quota_exhausted_state() { + let (store, repository) = + repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await; + let id = account_id("still-exhausted"); + set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await; + + GrokCredentialQuotaService::new( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":100}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh exhausted quota"); + + assert_eq!( + store.account(&id).expect("account").availability(), + AccountAvailability::QuotaExhausted + ); +} + +#[tokio::test] +async fn recovered_quota_does_not_clear_terminal_account_states() { + for (suffix, availability) in [ + ("keep-banned", AccountAvailability::Banned), + ("keep-expired", AccountAvailability::Expired), + ("keep-invalid", AccountAvailability::Invalid), + ] { + let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await; + let id = account_id(suffix); + set_account_state(&store, &id, availability, None).await; + GrokCredentialQuotaService::new( + repository, + QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#), + ) + .refresh_account(&id) + .await + .expect("refresh terminal account quota"); + + assert_eq!( + store.account(&id).expect("account").availability(), + availability + ); + } +} + +#[tokio::test] +async fn quota_refresh_does_not_overwrite_a_newer_cooldown() { + let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await; + let id = account_id("new-cooldown"); + let new_cooldown = SystemTime::now() + .checked_add(Duration::from_secs(600)) + .expect("cooldown time"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + availability: AccountAvailability::Cooldown, + reason: Some("new rate limit".to_owned()), + cooldown_until: Some(new_cooldown), + observed_at: SystemTime::now(), + }))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + GrokCredentialQuotaService::new(repository, transport) + .refresh_account(&id) + .await + .expect("refresh around concurrent cooldown"); + + let account = store.account(&id).expect("account"); + assert_eq!(account.availability(), AccountAvailability::Cooldown); + assert_eq!(account.cooldown_until(), Some(new_cooldown)); +} + +#[tokio::test] +async fn quota_refresh_rejects_a_concurrent_credential_revision() { + let (store, repository) = repository_with_accounts(&[("new-revision", "new-revision")]).await; + let id = account_id("new-revision"); + let account = store.account(&id).expect("account"); + let update = CredentialCasUpdate::new( + id.clone(), + account.revision(), + ProviderAccountUpdate { + account_id: id.clone(), + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + plan_type: account.plan_type().map(str::to_owned), + }, + store.credential(&id).expect("credential"), + account.has_refresh_token(), + account.access_token_expires_at(), + account.next_refresh_at(), + ) + .expect("credential update"); + let transport = Arc::new(MutatingBillingTransport { + store: Arc::clone(&store), + mutation: Mutex::new(Some(BillingMutation::Credential(update))), + body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(), + }); + + assert!(matches!( + GrokCredentialQuotaService::new(repository, transport) + .refresh_account(&id) + .await, + Err(GrokQuotaError::StaleCredentialSnapshot) + )); + assert_eq!(store.account(&id).expect("account").revision().get(), 2); +} + #[tokio::test] async fn quota_projection_falls_back_to_legacy_monthly_usage() { let (_, repository) = @@ -518,9 +738,10 @@ async fn failed_quota_refresh_does_not_replace_last_good_observation() { Err(GrokQuotaError::Upstream) )); let persisted = store - .get_quota(&account_id("stable")) + .get_quotas(&[account_id("stable")]) .await .expect("read quota") + .pop() .expect("quota remains") .quota .expect("provider document"); diff --git a/backend/crates/providers/xai/tests/credential/client.rs b/backend/crates/providers/xai/tests/credential/client.rs index c3914e715..917642660 100644 --- a/backend/crates/providers/xai/tests/credential/client.rs +++ b/backend/crates/providers/xai/tests/credential/client.rs @@ -3,7 +3,6 @@ use std::collections::VecDeque; use std::future::ready; use std::sync::{Arc, Mutex}; -use std::time::Duration; use gateway_core::engine::credential::AccountAvailability; use provider_xai::{ @@ -16,7 +15,7 @@ use provider_xai::{ VerifiedGrokAccount, }; -use crate::support::{account_id, instance_id}; +use crate::support::{account_id, instance_id, refresh_policy}; const DISCOVERY: &[u8] = include_bytes!("fixtures/discovery.json"); const TOKEN_SUCCESS: &[u8] = include_bytes!("fixtures/token_success.json"); @@ -114,17 +113,19 @@ async fn credential_import_should_require_exact_metadata_and_official_userinfo() assert_eq!(verified.evidence().method(), VerificationMethod::UserInfo); let prepared = GrokCredentialAdmin - .prepare_verified_account(&VerifiedGrokAccount { - account_id: account_id("verified-import"), - provider_instance_id: instance_id(), - name: "verified import".to_owned(), - email: Some("verified@example.com".to_owned()), - upstream_account_id: None, - plan_type: None, - tokens: verified, - enabled: true, - refresh_margin: Duration::from_secs(300), - }) + .prepare_verified_account( + &VerifiedGrokAccount { + account_id: account_id("verified-import"), + provider_instance_id: instance_id(), + name: "verified import".to_owned(), + email: Some("verified@example.com".to_owned()), + upstream_account_id: None, + plan_type: None, + tokens: verified, + enabled: true, + }, + refresh_policy(), + ) .expect("Provider projects verified token lifetime"); assert_eq!(prepared.account.availability(), AccountAvailability::Ready); assert!(prepared.account.next_refresh_at().is_some()); diff --git a/backend/crates/providers/xai/tests/credential/contract.rs b/backend/crates/providers/xai/tests/credential/contract.rs index b7fd50184..aba87fa89 100644 --- a/backend/crates/providers/xai/tests/credential/contract.rs +++ b/backend/crates/providers/xai/tests/credential/contract.rs @@ -99,6 +99,7 @@ async fn rotate_uses_revision_cas_and_replaces_plaintext_tokens() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject-rotate"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }) .expect("rotate"); assert!(matches!( @@ -145,6 +146,7 @@ async fn stale_rotate_does_not_modify_tokens() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject-stale"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }) .expect("winning rotation"); let stale = GrokCredentialAdmin @@ -157,6 +159,7 @@ async fn stale_rotate_does_not_modify_tokens() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject-stale"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }) .expect("stale command"); assert!(matches!( @@ -203,6 +206,7 @@ async fn rotate_rejects_verified_identity_rebind() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject-b"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }); assert!(matches!( result, @@ -303,6 +307,7 @@ async fn repository_rejects_account_owned_by_another_provider() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }), Err(GrokCredentialRepositoryError::WrongProviderKind) )); diff --git a/backend/crates/providers/xai/tests/credential/import.rs b/backend/crates/providers/xai/tests/credential/import.rs index 7a1f488e3..6f01f7d1b 100644 --- a/backend/crates/providers/xai/tests/credential/import.rs +++ b/backend/crates/providers/xai/tests/credential/import.rs @@ -1,19 +1,19 @@ use std::future::ready; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::{Duration as StdDuration, SystemTime}; +use std::time::SystemTime; use chrono::{Duration, Utc}; use provider_xai::{ FailClosedTokenVerifier, FailureClass, GrokCredentialAdmin, GrokOAuthClient, GrokOAuthConfig, GrokOAuthImportCandidate, GrokOAuthImportDocument, GrokOAuthImportError, GrokOAuthImportMetadata, GrokOAuthImportTokens, OAuthHttpRequest, OAuthHttpResponse, - OAuthHttpTransport, ReqwestOAuthTransport, ReqwestOidcTokenVerifier, SecretValue, - TokenCandidate, TokenVerificationContext, TokenVerifier, TransportFuture, VerificationEvidence, - VerificationFuture, VerifiedGrokAccount, + OAuthHttpTransport, OfficialGrokEndpointPolicy, ReqwestOAuthTransport, + ReqwestOidcTokenVerifier, SecretValue, TokenCandidate, TokenVerificationContext, TokenVerifier, + TransportFuture, VerificationEvidence, VerificationFuture, VerifiedGrokAccount, }; -use crate::support::{account_id, instance_id}; +use crate::support::{account_id, instance_id, refresh_policy}; const DISCOVERY: &[u8] = include_bytes!("fixtures/discovery.json"); const INVALID_GRANT: &[u8] = include_bytes!("fixtures/invalid_grant.json"); @@ -139,17 +139,23 @@ async fn verified_import_inside_refresh_margin_should_schedule_immediate_refresh .expect("still-valid imported credential"); let before = SystemTime::now(); let prepared = GrokCredentialAdmin - .prepare_verified_account(&VerifiedGrokAccount { - account_id: account_id("refresh-window-import"), - provider_instance_id: instance_id(), - name: "refresh-window-import".to_owned(), - email: None, - upstream_account_id: None, - plan_type: None, - tokens, - enabled: true, - refresh_margin: StdDuration::from_secs(60 * 60), - }) + .prepare_verified_account( + &VerifiedGrokAccount { + account_id: account_id("refresh-window-import"), + provider_instance_id: instance_id(), + name: "refresh-window-import".to_owned(), + email: None, + upstream_account_id: None, + plan_type: None, + tokens, + enabled: true, + }, + gateway_core::provider_ports::ProviderRefreshPolicy::try_new( + std::time::Duration::from_secs(60 * 60), + std::num::NonZeroU32::new(2).expect("positive concurrency"), + ) + .expect("valid refresh policy"), + ) .expect("valid credential inside refresh window must be imported"); let after = SystemTime::now(); let next_refresh_at = prepared @@ -310,8 +316,13 @@ fn real_oauth_account_should_match_provider_import_contract() { } #[tokio::test] -#[ignore = "requires XAI_REAL_ACCOUNT_FIXTURE and official xAI network access"] +#[ignore = "requires a disposable XAI_REAL_ACCOUNT_FIXTURE and explicit destructive refresh approval"] async fn real_oauth_accounts_should_cross_the_official_verification_boundary() { + assert_eq!( + std::env::var("XAI_ALLOW_DESTRUCTIVE_FIXTURE_REFRESH").as_deref(), + Ok("1"), + "official verification may rotate a refresh token; use a disposable fixture and set XAI_ALLOW_DESTRUCTIVE_FIXTURE_REFRESH=1", + ); let path = std::env::var_os("XAI_REAL_ACCOUNT_FIXTURE") .expect("XAI_REAL_ACCOUNT_FIXTURE must point to a local secret fixture"); let document = std::fs::read(path).expect("read local secret fixture"); @@ -319,10 +330,17 @@ async fn real_oauth_accounts_should_cross_the_official_verification_boundary() { .expect("real document must match the import contract") .into_entries(); let expected_accounts = entries.len(); + let endpoint_policy = Arc::new(OfficialGrokEndpointPolicy); let client = GrokOAuthClient::new( GrokOAuthConfig::official("0.2.101").expect("official config"), - Arc::new(ReqwestOAuthTransport::new().expect("production OAuth transport")), - Arc::new(ReqwestOidcTokenVerifier::new().expect("production token verifier")), + Arc::new( + ReqwestOAuthTransport::new(endpoint_policy.clone()) + .expect("production OAuth transport"), + ), + Arc::new( + ReqwestOidcTokenVerifier::new(endpoint_policy, std::time::Duration::from_secs(60 * 60)) + .expect("production token verifier"), + ), ); let discovery = client.discover().await.expect("official discovery"); let mut prepared_accounts = 0_usize; @@ -333,17 +351,19 @@ async fn real_oauth_accounts_should_cross_the_official_verification_boundary() { .expect("real OAuth credential must pass official verification"); let suffix = format!("real-import-{index}"); GrokCredentialAdmin - .prepare_verified_account(&VerifiedGrokAccount { - account_id: account_id(&suffix), - provider_instance_id: instance_id(), - name: suffix, - email: None, - upstream_account_id: None, - plan_type: None, - tokens, - enabled: true, - refresh_margin: StdDuration::from_secs(60 * 60), - }) + .prepare_verified_account( + &VerifiedGrokAccount { + account_id: account_id(&suffix), + provider_instance_id: instance_id(), + name: suffix, + email: None, + upstream_account_id: None, + plan_type: None, + tokens, + enabled: true, + }, + refresh_policy(), + ) .expect("verified real credential must cross account preparation"); prepared_accounts += 1; } diff --git a/backend/crates/providers/xai/tests/credential/oidc_verifier.rs b/backend/crates/providers/xai/tests/credential/oidc_verifier.rs index 599734126..ce5a51226 100644 --- a/backend/crates/providers/xai/tests/credential/oidc_verifier.rs +++ b/backend/crates/providers/xai/tests/credential/oidc_verifier.rs @@ -14,11 +14,13 @@ use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use provider_xai::{ - GrokEndpointPolicy, OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, ReqwestOidcTokenVerifier, SecretValue, - TokenCandidate, TokenVerificationContext, TokenVerifier, VerificationFailure, VerificationFlow, + OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, ReqwestOidcTokenVerifier, SecretValue, TokenCandidate, + TokenVerificationContext, TokenVerifier, VerificationFailure, VerificationFlow, VerificationMethod, }; +use crate::support::loopback_endpoint_policy; + const CACHE_TTL: Duration = Duration::from_secs(60 * 60); #[tokio::test] @@ -289,9 +291,8 @@ fn endpoints(server: &MockServer) -> TestEndpoints { } fn verifier(origin: &Url, cache_ttl: Duration) -> ReqwestOidcTokenVerifier { - let endpoint_policy = GrokEndpointPolicy::loopback(origin).expect("loopback policy"); - ReqwestOidcTokenVerifier::with_endpoint_policy_and_cache_ttl(endpoint_policy, cache_ttl) - .expect("injected loopback verifier") + ReqwestOidcTokenVerifier::new(loopback_endpoint_policy(origin), cache_ttl) + .expect("loopback verifier") } fn auth_context<'a>( diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index b7b34d867..a060c0518 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -8,19 +8,20 @@ use chrono::Utc; use gateway_core::engine::credential::{ AccountAvailability, CredentialCasOutcome, CredentialRevision, ProviderAccountStore, }; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError, +}; use provider_xai::{ GrokCredentialAdmin, GrokCredentialCatalogCache, GrokCredentialCatalogService, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, GrokCredentialRefresher, GrokCredentialRepository, GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, - GrokRefreshFailure, GrokRefreshLeaseAcquisition, GrokRefreshLeaseCoordinator, - GrokRefreshLeaseError, GrokRefreshLeaseRequest, GrokRefreshTokens, RotateManagedGrokCredential, - SecretValue, + GrokRefreshFailure, GrokRefreshTokens, RotateManagedGrokCredential, SecretValue, }; use crate::support::{ MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input, credential_object, profile, - seed_input, + runtime_policy, seed_input, }; const OFFICIAL_FIXTURE: &[u8] = @@ -76,17 +77,39 @@ struct TestRefreshLeases { calls: AtomicUsize, } -#[async_trait] -impl GrokRefreshLeaseCoordinator for TestRefreshLeases { - async fn try_acquire( +impl ProviderLeasePort for TestRefreshLeases { + fn load_state<'a>( + &'a self, + _: &'a gateway_core::routing::ProviderInstanceId, + _: &'a [gateway_core::engine::credential::ProviderAccountId], + ) -> futures::future::BoxFuture< + 'a, + Result, + > { + Box::pin(async { + Ok(gateway_core::provider_ports::ProviderSchedulingState::new( + Default::default(), + None, + 0, + )) + }) + } + + fn try_acquire( &self, - _: &GrokRefreshLeaseRequest, - ) -> Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Ok(if self.available { - GrokRefreshLeaseAcquisition::Acquired(Box::new(())) - } else { - GrokRefreshLeaseAcquisition::Unavailable + request: ProviderLeaseRequest, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + assert!(matches!( + request, + ProviderLeaseRequest::RefreshCapacity(_) | ProviderLeaseRequest::Refresh(_) + )); + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(if self.available { + ProviderLeaseAcquisition::Acquired(Box::new(())) + } else { + ProviderLeaseAcquisition::Busy { retry_after: None } + }) }) } } @@ -102,7 +125,7 @@ fn success_tokens(rotated: Option<&str>) -> GrokRefreshTokens { fn due_input(suffix: &str) -> provider_xai::CreateGrokCredential { let mut input = create_input(suffix, &format!("subject-{suffix}")); input.account.access_token_expires_at = Utc::now() + chrono::Duration::minutes(2); - input.account.next_refresh_at = Some(Utc::now() + chrono::Duration::minutes(1)); + input.next_refresh_at = Utc::now() - chrono::Duration::seconds(1); input } @@ -152,9 +175,8 @@ async fn fixture_many( refresher_port, catalog, leases, - Duration::from_secs(5 * 60), - ) - .expect("refresh service"); + runtime_policy(), + ); (store, repository, refresher, service) } @@ -164,7 +186,7 @@ async fn successful_refresh_rotates_plaintext_tokens_once() { let id = input.account_id.clone(); let (store, _, refresher, service) = fixture(input, [Ok(success_tokens(Some("new-refresh")))], true).await; - let outcomes = service.refresh_due(10).await.expect("refresh cycle"); + let outcomes = service.refresh_due().await.expect("refresh cycle"); assert!(matches!( outcomes.as_slice(), @@ -253,7 +275,7 @@ async fn omitted_rotated_refresh_token_preserves_existing_rt() { let input = due_input("preserve"); let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Ok(success_tokens(None))], true).await; - service.refresh_due(10).await.expect("refresh"); + service.refresh_due().await.expect("refresh"); let credential = store.credential(&id).expect("credential"); assert_eq!( credential_object(&credential) @@ -271,7 +293,7 @@ async fn unknown_refresh_token_expiry_does_not_block_refresh() { input.account.refresh_token_expires_at = None; let (store, _, _, service) = fixture(input, [Ok(success_tokens(None))], true).await; - let outcomes = service.refresh_due(10).await.expect("refresh cycle"); + let outcomes = service.refresh_due().await.expect("refresh cycle"); assert!(matches!( outcomes.as_slice(), [GrokCredentialRefreshOutcome::Refreshed { account_id, .. }] if account_id == &id @@ -283,13 +305,7 @@ async fn unknown_refresh_token_expiry_does_not_block_refresh() { async fn empty_due_set_does_not_resolve_discovery_or_call_upstream() { let input = create_input("not-due", "subject-not-due"); let (_, _, refresher, service) = fixture(input, [], true).await; - assert!( - service - .refresh_due(10) - .await - .expect("empty cycle") - .is_empty() - ); + assert!(service.refresh_due().await.expect("empty cycle").is_empty()); assert_eq!(refresher.prepare_calls.load(Ordering::SeqCst), 0); } @@ -298,7 +314,7 @@ async fn lease_unavailable_never_calls_refresh_exchange() { let input = due_input("lease"); let id = input.account_id.clone(); let (_, _, refresher, service) = fixture(input, [], false).await; - let outcomes = service.refresh_due(10).await.expect("refresh cycle"); + let outcomes = service.refresh_due().await.expect("refresh cycle"); assert!(matches!( outcomes.as_slice(), [GrokCredentialRefreshOutcome::LeaseUnavailable { account_id }] if account_id == &id @@ -312,7 +328,7 @@ async fn invalid_grant_marks_account_expired() { let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::InvalidGrant)], true).await; - service.refresh_due(10).await.expect("refresh"); + service.refresh_due().await.expect("refresh"); assert_eq!( store.account(&id).expect("account").availability(), AccountAvailability::Expired @@ -324,7 +340,7 @@ async fn banned_failure_marks_account_banned() { let input = due_input("banned"); let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::Banned)], true).await; - service.refresh_due(10).await.expect("refresh"); + service.refresh_due().await.expect("refresh"); assert_eq!( store.account(&id).expect("account").availability(), AccountAvailability::Banned @@ -336,7 +352,7 @@ async fn ambiguous_refresh_fails_closed_and_is_not_retried() { let input = due_input("ambiguous"); let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::Ambiguous)], true).await; - let outcomes = service.refresh_due(10).await.expect("refresh"); + let outcomes = service.refresh_due().await.expect("refresh"); assert!(matches!( outcomes.as_slice(), [GrokCredentialRefreshOutcome::Ambiguous { account_id }] if account_id == &id @@ -347,7 +363,7 @@ async fn ambiguous_refresh_fails_closed_and_is_not_retried() { ); assert!( service - .refresh_due(10) + .refresh_due() .await .expect("second cycle") .is_empty() @@ -359,10 +375,15 @@ async fn pre_send_transient_failure_applies_bounded_cooldown() { let input = due_input("transient"); let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::Transient)], true).await; - service.refresh_due(10).await.expect("refresh"); + service.refresh_due().await.expect("refresh"); let account = store.account(&id).expect("account"); - assert_eq!(account.availability(), AccountAvailability::Cooldown); - assert!(account.cooldown_until().is_some()); + assert_eq!(account.availability(), AccountAvailability::Unknown); + assert_eq!(account.revision().get(), 2); + assert!( + account + .next_refresh_at() + .is_some_and(|retry| retry > std::time::SystemTime::now()) + ); } #[tokio::test] @@ -370,9 +391,9 @@ async fn stale_due_snapshot_cannot_overwrite_newer_revision() { let input = due_input("stale"); let id = input.account_id.clone(); let (store, repository, _, service) = - fixture(input, [Ok(success_tokens(Some("worker-rt")))], true).await; + fixture(input, [Err(GrokRefreshFailure::InvalidGrant)], true).await; let due = repository - .list_due_refresh(Duration::from_secs(300), 10) + .list_due_refresh() .await .expect("due list") .pop() @@ -391,6 +412,7 @@ async fn stale_due_snapshot_cannot_overwrite_newer_revision() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: profile("subject-stale"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }) .expect("admin rotation"); assert!(matches!( @@ -421,7 +443,7 @@ async fn invalid_refresh_lifetime_is_rejected_without_cas_write() { ) .await; assert!(matches!( - service.refresh_due(10).await.expect("isolated refresh cycle").as_slice(), + service.refresh_due().await.expect("isolated refresh cycle").as_slice(), [GrokCredentialRefreshOutcome::Failed { account_id }] if account_id == &id )); assert_eq!(store.account(&id).expect("account").revision().get(), 1); @@ -429,14 +451,14 @@ async fn invalid_refresh_lifetime_is_rejected_without_cas_write() { #[tokio::test] async fn malformed_account_refresh_does_not_stop_later_accounts() { - let due_at = Utc::now() + chrono::Duration::minutes(1); + let due_at = Utc::now() - chrono::Duration::seconds(1); let access_expires_at = Utc::now() + chrono::Duration::minutes(2); let mut bad = due_input("bad"); - bad.account.next_refresh_at = Some(due_at); + bad.next_refresh_at = due_at; bad.account.access_token_expires_at = access_expires_at; let bad_id = bad.account_id.clone(); let mut good = due_input("good"); - good.account.next_refresh_at = Some(due_at); + good.next_refresh_at = due_at; good.account.access_token_expires_at = access_expires_at; let good_id = good.account_id.clone(); let (store, _, _, service) = fixture_many( @@ -453,10 +475,7 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() { ) .await; - let outcomes = service - .refresh_due(10) - .await - .expect("isolated refresh cycle"); + let outcomes = service.refresh_due().await.expect("isolated refresh cycle"); assert!(matches!( outcomes.as_slice(), diff --git a/backend/crates/providers/xai/tests/main.rs b/backend/crates/providers/xai/tests/main.rs index 6d4472d7f..c4aa9a5f2 100644 --- a/backend/crates/providers/xai/tests/main.rs +++ b/backend/crates/providers/xai/tests/main.rs @@ -1,3 +1,5 @@ +mod admin; +mod config; mod credential; mod provider; mod support; diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index a2f90a04a..0f72d3280 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -6,21 +6,18 @@ use std::time::{Duration, SystemTime}; use futures::{StreamExt, stream}; use gateway_core::engine::continuation::{ - NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, }; use gateway_core::engine::credential::{ AccountSelectionPolicy, CredentialRevision, ProviderAccountStore, RotationStrategy, }; use gateway_core::engine::provider::{Provider, ProviderRequest}; use gateway_core::engine::{ - AccountSelectionConstraints, AttemptContext, CancellationToken, ModelRequestId, - UpstreamSendState, + AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, UpstreamSendState, }; use gateway_core::error::{ProviderError, ProviderErrorKind, SafeUpstreamValue}; -use gateway_core::event::GatewayEvent; -use gateway_core::operation::{ - ContentPart, GenerateRequest, Message, MessageRole, Operation, OperationKind, -}; +use gateway_core::event::{GatewayEvent, UpstreamHttpVersion}; +use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProtocolPayload}; use gateway_core::routing::{ ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId, @@ -35,6 +32,7 @@ use provider_xai::{ GrokSessionSelector, GrokSessionSelectorError, GrokSessionSelectorFuture, SecretValue, SelectedGrokSession, }; +use serde_json::{Map, json}; use crate::support::{ MemoryGrokCatalogCache, MemoryProviderAccountStore, account_id, create_input, instance_id, @@ -178,15 +176,19 @@ impl GrokInferenceTransport for StubInferenceTransport { .expect("one transport mode"); Box::pin(async move { match mode { - InferenceMode::Success => { - Ok(GrokInferenceResponse::new(Box::pin(stream::iter([Ok( - SUCCESS_SSE.to_vec(), - )])))) - } + InferenceMode::Success => Ok(GrokInferenceResponse::new( + Box::pin(stream::iter([Ok(SUCCESS_SSE.to_vec())])), + UpstreamHttpVersion::Http2, + 200, + None, + )), InferenceMode::Error(error) => Err(error), - InferenceMode::StreamError(error) => Ok(GrokInferenceResponse::new(Box::pin( - stream::iter([Err(error)]), - ))), + InferenceMode::StreamError(error) => Ok(GrokInferenceResponse::new( + Box::pin(stream::iter([Err(error)])), + UpstreamHttpVersion::Http2, + 200, + None, + )), } }) } @@ -245,24 +247,30 @@ async fn mapped_transport_error( ) .await .expect("stream"); - stream - .next() - .await - .expect("error event") - .expect_err("transport rejection") + next_provider_error(&mut stream).await +} + +async fn next_provider_error( + stream: &mut gateway_core::engine::provider::ProviderStream, +) -> ProviderError { + loop { + match stream.next().await.expect("error event") { + Ok(event) => assert!(!event.has_client_event()), + Err(error) => return error, + } + } } fn operation() -> Operation { - Operation::Generate( - GenerateRequest::new(vec![ - Message::new( - MessageRole::User, - vec![ContentPart::Text("hello".to_owned())], - ) - .expect("message"), - ]) - .expect("generate request"), + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("input".to_owned(), json!("hello")), + ]), ) + .expect("OpenAI payload"); + Operation::Generate(GenerateRequest::from_protocol_payload(Vec::new(), payload)) } fn instance(provider: &str) -> ProviderInstance { @@ -314,22 +322,26 @@ fn selection_policy() -> AccountSelectionPolicy { fn context( cancellation: CancellationToken, - continuation: Option, + continuation: Option, ) -> AttemptContext { context_with_required(cancellation, continuation, None) } fn context_with_required( cancellation: CancellationToken, - continuation: Option, + continuation: Option, required_account: Option, ) -> AttemptContext { + let account_state_owner = continuation + .as_ref() + .and_then(ContinuationBinding::pinned) + .map(gateway_core::engine::ProviderAccountStateOwner::from_continuation); AttemptContext::new( ModelRequestId::new("req_xai").expect("request ID"), NonZeroU32::new(1).expect("attempt"), SystemTime::now() + Duration::from_secs(30), selection_policy(), - AccountSelectionConstraints::new(BTreeSet::new(), required_account), + AccountAttemptContext::new(BTreeSet::new(), required_account, account_state_owner), continuation, cancellation, ) @@ -380,10 +392,20 @@ async fn execute_returns_cold_stream_and_records_selected_account() { let events = stream.by_ref().collect::>().await; assert_eq!(transport.calls.load(Ordering::SeqCst), 1); assert!(events.iter().all(Result::is_ok)); + let observation = events[0] + .as_ref() + .expect("response observation") + .response_observation() + .expect("transport facts"); + assert_eq!(observation.transport().as_str(), "http_sse"); + assert_eq!(observation.http_version(), Some(UpstreamHttpVersion::Http2)); + assert_eq!(observation.status_code(), Some(200)); assert!( events .iter() - .any(|event| matches!(event, Ok(GatewayEvent::Completed(_)))) + .filter_map(|event| event.as_ref().ok()) + .flat_map(|event| event.canonical_facts()) + .any(|event| matches!(event, GatewayEvent::Completed(_))) ); } @@ -435,11 +457,7 @@ async fn unauthorized_transport_feedback_is_bound_to_selected_account() { ) .await .expect("stream"); - let error = stream - .next() - .await - .expect("error event") - .expect_err("transport error"); + let error = next_provider_error(&mut stream).await; assert_eq!(error.kind(), ProviderErrorKind::Unauthorized); assert_eq!( selector.feedback.lock().expect("feedback").as_slice(), @@ -497,11 +515,7 @@ async fn generic_http_403_does_not_mutate_or_switch_credential() { ) .await .expect("stream"); - let error = stream - .next() - .await - .expect("error event") - .expect_err("transport rejection"); + let error = next_provider_error(&mut stream).await; assert_eq!(error.kind(), ProviderErrorKind::PermissionDenied); assert!(!error.replay_is_safe()); @@ -579,7 +593,10 @@ async fn native_previous_response_is_rejected_before_selection() { let error = match provider .execute( provider_request("xai"), - context(CancellationToken::new(), Some(pin)), + context( + CancellationToken::new(), + Some(ContinuationBinding::Pinned(pin)), + ), ) .await { @@ -590,6 +607,27 @@ async fn native_previous_response_is_rejected_before_selection() { assert_eq!(selector.calls.load(Ordering::SeqCst), 0); } +#[tokio::test] +async fn external_previous_response_is_rejected_before_selection() { + let selector = StubSelector::success(); + let provider = provider(selector.clone(), StubInferenceTransport::success()).await; + let binding = ContinuationBinding::External( + PreviousResponseId::new("external-provider-response").expect("response ID"), + ); + let error = match provider + .execute( + provider_request("xai"), + context(CancellationToken::new(), Some(binding)), + ) + .await + { + Ok(_) => panic!("xAI external continuation must fail"), + Err(error) => error, + }; + assert_eq!(error.kind(), ProviderErrorKind::Unsupported); + assert_eq!(selector.calls.load(Ordering::SeqCst), 0); +} + #[tokio::test] async fn cancellation_before_poll_never_calls_transport() { let transport = StubInferenceTransport::success(); diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index 529293dda..4adb0655f 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -1,6 +1,9 @@ use std::collections::BTreeMap; +use std::net::IpAddr; +use std::num::NonZeroU32; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; use async_trait::async_trait; use chrono::Utc; @@ -11,12 +14,142 @@ use gateway_core::engine::credential::{ QuotaWriteOutcome, }; use gateway_core::error::{StoreError, StoreErrorKind}; +use gateway_core::provider_ports::{ + ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, +}; use gateway_core::routing::ProviderInstanceId; use provider_xai::{ - CreateGrokCredential, GrokAccountCatalog, GrokAccountProfile, GrokCatalogCacheError, - GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, - GrokCredentialRepositoryError, GrokOAuthSecret, SecretValue, + CreateGrokCredential, GROK_BILLING_URL, GrokAccountCatalog, GrokAccountProfile, + GrokCatalogCacheError, GrokCredentialAdmin, GrokCredentialAvailability, + GrokCredentialCatalogCache, GrokCredentialRepositoryError, GrokEndpointPolicy, GrokOAuthSecret, + GrokReqwestTransportBuildError, SecretValue, }; +use reqwest::Client; +use reqwest::redirect::Policy; +use url::{Host, Url}; + +#[derive(Debug, Clone, Copy)] +pub struct LoopbackGrokEndpointPolicy { + address: IpAddr, + port: u16, +} + +impl LoopbackGrokEndpointPolicy { + pub fn for_origin(origin: &Url) -> Self { + assert_eq!(origin.scheme(), "http", "test origin must use HTTP"); + assert!( + origin.username().is_empty(), + "test origin must not contain credentials" + ); + assert!( + origin.password().is_none(), + "test origin must not contain credentials" + ); + assert!( + origin.query().is_none(), + "test origin must not contain a query" + ); + assert!( + origin.fragment().is_none(), + "test origin must not contain a fragment" + ); + let address = match origin.host() { + Some(Host::Ipv4(address)) if address.is_loopback() => IpAddr::V4(address), + Some(Host::Ipv6(address)) if address.is_loopback() => IpAddr::V6(address), + _ => panic!("test origin must use a numeric loopback address"), + }; + Self { + address, + port: origin.port().expect("test origin must contain a port"), + } + } + + fn accepts_origin(self, url: &Url) -> bool { + url.scheme() == "http" + && url.username().is_empty() + && url.password().is_none() + && url.host().and_then(|host| match host { + Host::Ipv4(address) => Some(IpAddr::V4(address)), + Host::Ipv6(address) => Some(IpAddr::V6(address)), + Host::Domain(_) => None, + }) == Some(self.address) + && url.port() == Some(self.port) + && url.fragment().is_none() + } + + fn build_client(timeout: Option) -> Result { + let mut builder = Client::builder() + .redirect(Policy::none()) + .no_proxy() + .https_only(false) + .tcp_nodelay(true); + if let Some(timeout) = timeout { + builder = builder.timeout(timeout); + } + builder + .build() + .map_err(|_| GrokReqwestTransportBuildError::ClientInitialization) + } +} + +impl GrokEndpointPolicy for LoopbackGrokEndpointPolicy { + fn build_oauth_client( + &self, + timeout: Option, + ) -> Result { + Self::build_client(timeout) + } + + fn build_inference_client( + &self, + timeout: Option, + ) -> Result { + Self::build_client(timeout) + } + + fn validate_oauth(&self, url: &Url) -> bool { + self.accepts_origin(url) + } + + fn validate_inference(&self, url: &Url) -> bool { + self.accepts_origin(url) && url.path() == "/v1/responses" && url.query().is_none() + } + + fn validate_model_catalog(&self, url: &Url) -> bool { + self.accepts_origin(url) && url.path() == "/v1/models" && url.query().is_none() + } + + fn route_billing(&self, url: &Url) -> Option { + (url.as_str() == GROK_BILLING_URL).then(|| { + let mut endpoint = self.origin(); + endpoint.set_path("/v1/billing"); + endpoint.set_query(Some("format=credits")); + endpoint + }) + } + + fn validate_jwks(&self, url: &Url) -> bool { + self.accepts_origin(url) && url.path() == "/.well-known/jwks.json" && url.query().is_none() + } + + fn validate_userinfo(&self, url: &Url) -> bool { + self.accepts_origin(url) && url.path() == "/oauth2/userinfo" && url.query().is_none() + } +} + +impl LoopbackGrokEndpointPolicy { + fn origin(self) -> Url { + let origin = match self.address { + IpAddr::V4(address) => format!("http://{address}:{}", self.port), + IpAddr::V6(address) => format!("http://[{address}]:{}", self.port), + }; + Url::parse(&origin).expect("validated loopback origin") + } +} + +pub fn loopback_endpoint_policy(origin: &Url) -> Arc { + Arc::new(LoopbackGrokEndpointPolicy::for_origin(origin)) +} #[derive(Clone)] struct StoredAccount { @@ -28,6 +161,7 @@ struct StoredAccount { #[derive(Default)] pub struct MemoryProviderAccountStore { accounts: Mutex>, + quota_reads: AtomicUsize, } impl MemoryProviderAccountStore { @@ -50,6 +184,10 @@ impl MemoryProviderAccountStore { pub fn len(&self) -> usize { lock(&self.accounts).len() } + + pub fn quota_reads(&self) -> usize { + self.quota_reads.load(Ordering::SeqCst) + } } #[async_trait] @@ -146,16 +284,20 @@ impl ProviderAccountStore for MemoryProviderAccountStore { }, ); stored.credential = credential; + stored.quota = None; Ok(CredentialCasOutcome::Updated(next)) } - async fn get_quota( + async fn get_quotas( &self, - account: &ProviderAccountId, - ) -> Result, StoreError> { - Ok(lock(&self.accounts) - .get(account) - .and_then(|stored| stored.quota.clone())) + accounts: &[ProviderAccountId], + ) -> Result, StoreError> { + self.quota_reads.fetch_add(1, Ordering::SeqCst); + let stored = lock(&self.accounts); + Ok(accounts + .iter() + .filter_map(|account| stored.get(account)?.quota.clone()) + .collect()) } async fn compare_and_swap_quota( @@ -313,15 +455,15 @@ impl GrokCredentialCatalogCache for MemoryGrokCatalogCache { .cloned()) } - async fn permits( + async fn observed_model_support( &self, account_id: &ProviderAccountId, revision: CredentialRevision, model: &str, - ) -> Result { + ) -> Result, GrokCatalogCacheError> { Ok(lock(&self.entries) .get(&(account_id.clone(), revision)) - .is_some_and(|catalog| catalog.seed().permits(model))) + .map(|catalog| catalog.seed().permits(model))) } } @@ -342,11 +484,11 @@ pub fn profile(subject: &str) -> GrokAccountProfile { plan_type: Some("standard".to_owned()), access_token_expires_at: now + chrono::Duration::hours(1), refresh_token_expires_at: Some(now + chrono::Duration::days(30)), - next_refresh_at: Some(now + chrono::Duration::minutes(30)), } } pub fn create_input(suffix: &str, subject: &str) -> CreateGrokCredential { + let now = Utc::now(); CreateGrokCredential { account_id: account_id(suffix), provider_instance_id: instance_id(), @@ -358,6 +500,7 @@ pub fn create_input(suffix: &str, subject: &str) -> CreateGrokCredential { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, account: profile(subject), + next_refresh_at: now + chrono::Duration::minutes(30), enabled: true, initial_availability: GrokCredentialAvailability::Unknown, initial_availability_reason: None, @@ -365,6 +508,28 @@ pub fn create_input(suffix: &str, subject: &str) -> CreateGrokCredential { } } +pub struct StaticRuntimePolicy; + +impl ProviderRuntimePolicyPort for StaticRuntimePolicy { + fn load_refresh_policy( + &self, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async { Ok(refresh_policy()) }) + } +} + +pub fn refresh_policy() -> ProviderRefreshPolicy { + ProviderRefreshPolicy::try_new( + Duration::from_secs(5 * 60), + NonZeroU32::new(2).expect("positive concurrency"), + ) + .expect("valid refresh policy") +} + +pub fn runtime_policy() -> Arc { + Arc::new(StaticRuntimePolicy) +} + pub fn prepare_input( input: &CreateGrokCredential, ) -> Result { diff --git a/backend/crates/providers/xai/tests/transport/canonical.rs b/backend/crates/providers/xai/tests/transport/canonical.rs index bd091f93f..78840899d 100644 --- a/backend/crates/providers/xai/tests/transport/canonical.rs +++ b/backend/crates/providers/xai/tests/transport/canonical.rs @@ -1,4 +1,5 @@ -use gateway_core::event::{ContentKind, FinishReason, GatewayEvent}; +use gateway_core::error::ProviderError; +use gateway_core::event::{ContentKind, FinishReason, GatewayEvent, ProviderEvent}; use provider_xai::{GrokCanonicalDecoder, grok_billing_breakdown}; @@ -25,9 +26,25 @@ fn terminal_cost_events( cached_tokens = cached_tokens, provider_cost = provider_cost, ); + decode_canonical(body.as_bytes()).expect("canonical cost response") +} + +fn decode_canonical(body: &[u8]) -> Result, ProviderError> { GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical cost response") + .push(body) + .map(|events| { + events + .into_iter() + .flat_map(|event| event.into_parts().0) + .collect() + }) +} + +fn wire_events(events: &[ProviderEvent]) -> Vec<&gateway_core::event::ProtocolWireEvent> { + events + .iter() + .filter_map(ProviderEvent::wire_event) + .collect() } fn calculated_cost_ticks(events: &[GatewayEvent]) -> Option { @@ -58,9 +75,7 @@ fn decoder_should_normalize_text_usage_and_completion() { "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"model\":\"grok-code-test\",\"status\":\"completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":2,\"total_tokens\":5,\"cost_in_usd_ticks\":37756000}}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical response"); + let events = decode_canonical(body.as_bytes()).expect("canonical response"); assert!(matches!(events[0], GatewayEvent::Started(_))); assert!(matches!( @@ -89,9 +104,7 @@ fn decoder_should_leave_cost_unavailable_when_upstream_omits_it() { "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_no_cost\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical response"); + let events = decode_canonical(body.as_bytes()).expect("canonical response"); assert!(!events.iter().any(|event| matches!( event, @@ -184,9 +197,7 @@ fn decoder_should_leave_incomplete_usage_pricing_unavailable() { "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_partial\",\"model\":\"grok-4.5\",\"status\":\"completed\",\"usage\":{\"total_tokens\":3}}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical partial usage response"); + let events = decode_canonical(body.as_bytes()).expect("canonical partial usage response"); assert_eq!(calculated_cost_ticks(&events), None); } @@ -232,9 +243,7 @@ fn decoder_should_normalize_function_call_and_tool_finish() { "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_tool\",\"model\":\"grok-code-test\",\"status\":\"completed\"}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical function call"); + let events = decode_canonical(body.as_bytes()).expect("canonical function call"); assert!(events.iter().any(|event| matches!( event, @@ -257,9 +266,7 @@ fn decoder_should_coalesce_reasoning_item_part_and_summary_index() { "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_reason\",\"status\":\"completed\"}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("canonical reasoning"); + let events = decode_canonical(body.as_bytes()).expect("canonical reasoning"); assert_eq!( events @@ -272,27 +279,74 @@ fn decoder_should_coalesce_reasoning_item_part_and_summary_index() { } #[test] -fn decoder_should_fail_closed_for_unknown_or_mismatched_events() { - for body in [ - concat!( - "event: response.created\n", - "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n", - "event: response.backend_tool_call.started\n", - "data: {\"type\":\"response.backend_tool_call.started\",\"secret\":\"must-not-leak\"}\n\n", - ), - concat!( - "event: response.created\n", - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\"}}\n\n", - ), - ] { - let error = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect_err("unsupported event must fail"); - let debug = format!("{error:?}"); +fn decoder_should_preserve_unknown_events_as_openai_wire() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n", + "event: response.backend_tool_call.started\n", + "data: {\"type\":\"response.backend_tool_call.started\",\"future\":{\"nested\":[1,true]}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"status\":\"completed\"}}\n\n", + ); + let events = GrokCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("unknown wire event"); + let wire = wire_events(&events); - assert!(!debug.contains("must-not-leak")); - assert!(error.sensitive_context_was_redacted()); - } + assert_eq!(wire.len(), 3); + assert_eq!(wire[1].protocol(), "openai"); + assert_eq!( + wire[1].event_type(), + Some("response.backend_tool_call.started") + ); + assert_eq!( + wire[1].data().pointer("/future/nested"), + Some(&serde_json::json!([1, true])) + ); + assert!(events[1].canonical_facts().is_empty()); +} + +#[test] +fn decoder_should_preserve_image_and_hosted_tool_items_without_inventing_facts() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_items\"}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"id\":\"image_1\",\"future\":true}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"type\":\"web_search_call\",\"id\":\"search_1\",\"status\":\"searching\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_items\",\"status\":\"completed\"}}\n\n", + ); + let events = GrokCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("opaque output items"); + + assert_eq!(wire_events(&events).len(), 4); + assert!(events[1].canonical_facts().is_empty()); + assert!(events[2].canonical_facts().is_empty()); + assert_eq!( + events[1] + .wire_event() + .expect("image wire") + .data() + .pointer("/item/future"), + Some(&serde_json::json!(true)) + ); +} + +#[test] +fn decoder_should_reject_mismatched_event_identity_without_retaining_body() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\"},\"secret\":\"must-not-leak\"}\n\n", + ); + let error = GrokCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect_err("mismatched event must fail"); + + assert!(!format!("{error:?}").contains("must-not-leak")); + assert!(error.sensitive_context_was_redacted()); } #[test] @@ -320,9 +374,7 @@ fn decoder_should_preserve_incomplete_length_reason() { "event: response.incomplete\n", "data: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_short\",\"status\":\"incomplete\",\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n", ); - let events = GrokCanonicalDecoder::new("fallback") - .push(body.as_bytes()) - .expect("incomplete response"); + let events = decode_canonical(body.as_bytes()).expect("incomplete response"); assert!(matches!( events.last(), diff --git a/backend/crates/providers/xai/tests/transport/failure_feedback.rs b/backend/crates/providers/xai/tests/transport/failure_feedback.rs index 5216f710d..d6bc4916b 100644 --- a/backend/crates/providers/xai/tests/transport/failure_feedback.rs +++ b/backend/crates/providers/xai/tests/transport/failure_feedback.rs @@ -2,20 +2,23 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; -use async_trait::async_trait; use chrono::Utc; use gateway_core::engine::credential::{ AccountAvailability, AccountRuntimeSignals, AccountSelectionPolicy, CredentialCasOutcome, - CredentialRevision, ProviderAccountId, ProviderAccountStore, RotationStrategy, + CredentialRevision, OpaqueProviderData, ProviderAccountId, ProviderAccountStore, + QuotaObservation, QuotaWriteOutcome, RotationStrategy, +}; +use gateway_core::provider_ports::{ + ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingState, + ProviderStoreError, }; use provider_xai::{ - GrokAccountCatalog, GrokAccountSchedulingState, GrokAccountSessionSelector, + GrokAccountCatalog, GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport, + GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, - GrokCredentialCatalogSeed, GrokCredentialFailure, GrokCredentialLeaseAcquisition, - GrokCredentialLeaseCoordinator, GrokCredentialLeaseCoordinatorError, - GrokCredentialLeaseRequest, GrokCredentialRepository, GrokSessionSelection, - GrokSessionSelector, GrokSessionSelectorError, RotateManagedGrokCredential, - UpdateGrokCredentialState, + GrokCredentialCatalogSeed, GrokCredentialFailure, GrokCredentialQuotaService, + GrokCredentialRepository, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, + RotateManagedGrokCredential, UpdateGrokCredentialState, }; use crate::support::{ @@ -28,43 +31,62 @@ struct SchedulingCoordinator { denied: Mutex>, } -#[async_trait] -impl GrokCredentialLeaseCoordinator for SchedulingCoordinator { - async fn load_scheduling_state( - &self, - _: &gateway_core::routing::ProviderInstanceId, - _: &[ProviderAccountId], - ) -> Result { - Ok(GrokAccountSchedulingState { - signals: self.signals.lock().expect("signals").clone(), - sticky_account: None, - round_robin_cursor: 0, +struct UnavailableBillingTransport; + +impl GrokBillingTransport for UnavailableBillingTransport { + fn execute(&self, _: GrokBillingRequest) -> GrokBillingTransportFuture<'_> { + Box::pin(async { + Err(GrokBillingTransportError::new( + GrokBillingTransportErrorKind::Unavailable, + )) + }) + } +} + +impl ProviderLeasePort for SchedulingCoordinator { + fn load_state<'a>( + &'a self, + _: &'a gateway_core::routing::ProviderInstanceId, + _: &'a [ProviderAccountId], + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + Ok(ProviderSchedulingState::new( + self.signals.lock().expect("signals").clone(), + None, + 0, + )) }) } - async fn try_acquire( + fn try_acquire( &self, - request: &GrokCredentialLeaseRequest, - ) -> Result { - Ok( - if self - .denied - .lock() - .expect("denied") - .contains(&request.account_id) - { - GrokCredentialLeaseAcquisition::Unavailable { - retry_after: Some(Duration::from_millis(25)), - } - } else { - GrokCredentialLeaseAcquisition::Acquired(Box::new(())) - }, - ) + request: ProviderLeaseRequest, + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { + let ProviderLeaseRequest::Scheduling(request) = request else { + panic!("expected scheduling lease request"); + }; + Ok( + if self + .denied + .lock() + .expect("denied") + .contains(request.account_id()) + { + ProviderLeaseAcquisition::Busy { + retry_after: Some(Duration::from_millis(25)), + } + } else { + ProviderLeaseAcquisition::Acquired(Box::new(())) + }, + ) + }) } } struct SelectorFixture { store: Arc, + cache: Arc, selector: GrokAccountSessionSelector, coordinator: Arc, } @@ -105,7 +127,7 @@ impl SelectorFixture { in_flight: 0, last_started_at: None, quota_reset_at: None, - quota_remaining_rank: Some(100), + quota_remaining_rank: None, }, ); } @@ -113,12 +135,17 @@ impl SelectorFixture { signals: Mutex::new(signals), denied: Mutex::new(BTreeSet::new()), }); - let catalog_cache: Arc = cache; - let lease_port: Arc = coordinator.clone(); + let catalog_cache: Arc = cache.clone(); + let lease_port: Arc = coordinator.clone(); + let quota = Arc::new(GrokCredentialQuotaService::new( + repository.clone(), + Arc::new(UnavailableBillingTransport), + )); let selector = - GrokAccountSessionSelector::new(repository.clone(), catalog_cache, lease_port); + GrokAccountSessionSelector::new(repository.clone(), catalog_cache, quota, lease_port); Self { store, + cache, selector, coordinator, } @@ -132,6 +159,15 @@ impl SelectorFixture { &self, excluded: BTreeSet, required_account: Option, + ) -> GrokSessionSelection { + self.request_with_policy(excluded, required_account, RotationStrategy::Smart) + } + + fn request_with_policy( + &self, + excluded: BTreeSet, + required_account: Option, + strategy: RotationStrategy, ) -> GrokSessionSelection { GrokSessionSelection::new( instance_id(), @@ -139,12 +175,42 @@ impl SelectorFixture { excluded, required_account, AccountSelectionPolicy::new( - RotationStrategy::Smart, + strategy, std::num::NonZeroU32::new(2).expect("limit"), Duration::ZERO, ), + SystemTime::now() + Duration::from_secs(30), ) } + + async fn seed_quota(&self, id: &ProviderAccountId, used_percent: f64, reset_after: Duration) { + let reset_at = (Utc::now() + + chrono::Duration::from_std(reset_after).expect("valid reset duration")) + .to_rfc3339(); + let document = serde_json::json!({ + "config": { + "creditUsagePercent": used_percent, + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": Utc::now().to_rfc3339(), + "end": reset_at + } + } + }); + let outcome = self + .store + .compare_and_swap_quota(QuotaObservation { + account_id: id.clone(), + expected_revision: CredentialRevision::new(1).expect("revision"), + quota: Some(OpaqueProviderData::new( + document.as_object().expect("quota object").clone(), + )), + observed_at: Some(SystemTime::now()), + }) + .await + .expect("persist quota"); + assert_eq!(outcome, QuotaWriteOutcome::Updated); + } } #[tokio::test] @@ -313,7 +379,7 @@ async fn smart_strategy_prefers_lower_in_flight_account() { } #[tokio::test] -async fn stale_catalog_revision_fails_closed() { +async fn stale_catalog_revision_does_not_block_transparent_request() { let fixture = SelectorFixture::new(&["catalog-stale"]).await; let id = account_id("catalog-stale"); let current = fixture @@ -331,6 +397,7 @@ async fn stale_catalog_revision_fails_closed() { scope: provider_xai::OFFICIAL_SCOPES.join(" "), }, verified_account: crate::support::profile("subject-catalog-stale"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), }) .expect("rotate"); assert!(matches!( @@ -341,11 +408,141 @@ async fn stale_catalog_revision_fails_closed() { .expect("persist rotation"), CredentialCasOutcome::Updated(revision) if revision.get() == 2 )); - assert!(matches!( - fixture - .selector - .select(fixture.request(BTreeSet::new())) - .await, - Err(GrokSessionSelectorError::NoEligibleSession) - )); + let session = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("stale auxiliary catalog must not block selection"); + assert_eq!(session.account_id(), &id); +} + +#[tokio::test] +async fn explicit_catalog_non_membership_excludes_only_unsupported_account() { + let fixture = SelectorFixture::new(&["aaa-unsupported", "zzz-supported"]).await; + fixture + .cache + .replace(GrokAccountCatalog::new( + account_id("aaa-unsupported"), + CredentialRevision::new(1).expect("revision"), + Utc::now(), + GrokCredentialCatalogSeed::new(["grok-other"], None).expect("catalog"), + )) + .await + .expect("replace catalog"); + + let session = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("supported account"); + + assert_eq!(session.account_id(), &account_id("zzz-supported")); +} + +#[tokio::test] +async fn smart_strategy_uses_fresh_provider_quota_after_load_ties() { + let fixture = SelectorFixture::new(&["aaa-low-quota", "zzz-high-quota"]).await; + fixture + .seed_quota(&account_id("aaa-low-quota"), 90.0, Duration::from_secs(600)) + .await; + fixture + .seed_quota( + &account_id("zzz-high-quota"), + 10.0, + Duration::from_secs(600), + ) + .await; + + let session = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("quota-ranked account"); + + assert_eq!(session.account_id(), &account_id("zzz-high-quota")); +} + +#[tokio::test] +async fn smart_strategy_never_reuses_quota_projection_after_credential_rotation() { + let fixture = SelectorFixture::new(&["aaa-stale-high", "zzz-current-low"]).await; + let stale = account_id("aaa-stale-high"); + fixture + .seed_quota(&stale, 5.0, Duration::from_secs(600)) + .await; + fixture + .seed_quota( + &account_id("zzz-current-low"), + 95.0, + Duration::from_secs(600), + ) + .await; + let first = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("initial quota-ranked account"); + assert_eq!(first.account_id(), &stale); + drop(first); + + let current = fixture + .store + .load_credential(&stale, CredentialRevision::new(1).expect("revision")) + .await + .expect("current credential"); + let prepared = GrokCredentialAdmin + .prepare_rotation(&RotateManagedGrokCredential { + current, + secret: provider_xai::GrokOAuthSecret { + access_token: provider_xai::SecretValue::new("rotated-access"), + refresh_token: provider_xai::SecretValue::new("rotated-refresh"), + id_token: None, + scope: provider_xai::OFFICIAL_SCOPES.join(" "), + }, + verified_account: crate::support::profile("subject-aaa-stale-high"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), + }) + .expect("rotate"); + let outcome = fixture + .store + .compare_and_swap_credential(prepared.credential) + .await + .expect("persist rotation"); + assert!(matches!(outcome, CredentialCasOutcome::Updated(revision) if revision.get() == 2)); + + let selected = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("current quota-ranked account"); + + assert_eq!(selected.account_id(), &account_id("zzz-current-low")); +} + +#[tokio::test] +async fn quota_reset_strategy_uses_provider_reported_earliest_reset() { + let fixture = SelectorFixture::new(&["aaa-later-reset", "zzz-earlier-reset"]).await; + fixture + .seed_quota( + &account_id("aaa-later-reset"), + 10.0, + Duration::from_secs(1_200), + ) + .await; + fixture + .seed_quota( + &account_id("zzz-earlier-reset"), + 90.0, + Duration::from_secs(600), + ) + .await; + let request = + fixture.request_with_policy(BTreeSet::new(), None, RotationStrategy::QuotaResetPriority); + + let session = fixture + .selector + .select(request) + .await + .expect("reset-ranked account"); + + assert_eq!(session.account_id(), &account_id("zzz-earlier-reset")); } diff --git a/backend/crates/providers/xai/tests/transport/headers.rs b/backend/crates/providers/xai/tests/transport/headers.rs index 258bbc3e6..f15a64e96 100644 --- a/backend/crates/providers/xai/tests/transport/headers.rs +++ b/backend/crates/providers/xai/tests/transport/headers.rs @@ -35,23 +35,8 @@ fn instance() -> GrokProviderInstanceConfig { #[test] fn header_debug_should_redact_oauth_and_identity_values() { - let request = provider_xai::GrokResponsesRequest::encode( - &gateway_core::operation::GenerateRequest::new(vec![ - gateway_core::operation::Message::new( - gateway_core::operation::MessageRole::User, - vec![gateway_core::operation::ContentPart::Text( - "hello".to_owned(), - )], - ) - .expect("message"), - ]) - .expect("request"), - "grok-code-test", - ) - .expect("encoded request"); let headers = build_grok_headers( &instance(), - &request, &selected_session(), &ModelRequestId::new("req_grok_test").expect("request ID"), &UpstreamModelId::new("grok-code-test").expect("model"), @@ -68,3 +53,26 @@ fn header_debug_should_redact_oauth_and_identity_values() { assert!(!debug.contains(secret), "debug output was {debug}"); } } + +#[test] +fn headers_should_bind_identity_to_the_selected_oauth_account() { + let headers = build_grok_headers( + &instance(), + &selected_session(), + &ModelRequestId::new("req_grok_identity").expect("request ID"), + &UpstreamModelId::new("grok-code-test").expect("model"), + ); + let value = |name: &str| { + headers + .iter() + .find(|header| header.name().eq_ignore_ascii_case(name)) + .map(|header| header.value().expose()) + }; + + assert_eq!(value("x-userid"), Some("fixture-user-id")); + assert_eq!(value("x-grok-user-id"), Some("fixture-user-id")); + assert_eq!(value("x-email"), Some("fixture@example.test")); + assert_eq!(value("x-grok-conv-id"), Some("req_grok_identity")); + assert_eq!(value("x-grok-session-id"), Some("req_grok_identity")); + assert_eq!(value("x-grok-agent-id"), Some("codex-proxy-rs")); +} diff --git a/backend/crates/providers/xai/tests/transport/network.rs b/backend/crates/providers/xai/tests/transport/network.rs index 2ace6b6a2..b29fd227e 100644 --- a/backend/crates/providers/xai/tests/transport/network.rs +++ b/backend/crates/providers/xai/tests/transport/network.rs @@ -12,14 +12,16 @@ use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; use provider_xai::{ - FailClosedTokenVerifier, FormField, GrokBillingRequest, GrokBillingTransport, - GrokDnsResolutionPlan, GrokDnsResolutionPolicy, GrokEndpointPolicy, GrokInferenceRequest, - GrokInferenceTransport, GrokInferenceTransportErrorKind, GrokModelCatalogSession, - GrokOAuthClient, GrokOAuthConfig, GrokSessionBinding, HttpMethod, OAuthHttpRequest, - OAuthHttpTransport, ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, + FailClosedTokenVerifier, FormField, GrokBillingClient, GrokDnsResolutionPlan, + GrokDnsResolutionPolicy, GrokInferenceRequest, GrokInferenceTransport, + GrokInferenceTransportErrorKind, GrokModelCatalogSession, GrokOAuthClient, GrokOAuthConfig, + GrokSessionBinding, HttpMethod, OAuthHttpRequest, OAuthHttpTransport, + OfficialGrokEndpointPolicy, ReqwestGrokInferenceTransport, ReqwestGrokModelCatalogTransport, ReqwestOAuthTransport, SecretValue, }; +use crate::support::loopback_endpoint_policy; + #[tokio::test] async fn oauth_transport_should_post_form_once_without_redirect() { let server = MockServer::start().await; @@ -74,6 +76,11 @@ async fn inference_transport_should_stream_one_official_shape_response() { ); let response = transport.execute(request).await.expect("SSE response"); + assert_eq!( + response.http_version(), + gateway_core::event::UpstreamHttpVersion::Http11 + ); + assert_eq!(response.status_code(), 200); let chunks = response.into_body().collect::>().await; assert_eq!(chunks.len(), 1); @@ -114,8 +121,15 @@ async fn inference_transport_should_classify_http_failures_without_retaining_bod error.status(), error.send_state(), error.sensitive_context_was_redacted(), + error.http_version(), ), - (expected_kind, Some(status), UpstreamSendState::Sent, true,) + ( + expected_kind, + Some(status), + UpstreamSendState::Sent, + true, + Some(gateway_core::event::UpstreamHttpVersion::Http11), + ) ); assert!(!rendered.contains(&secret)); } @@ -187,9 +201,9 @@ async fn billing_transport_should_get_exact_credits_resource_without_redirect() .mount(&server) .await; let origin = Url::parse(&server.uri()).expect("wiremock origin"); - let policy = GrokEndpointPolicy::loopback(&origin).expect("loopback policy"); let transport = Arc::new( - ReqwestGrokModelCatalogTransport::with_endpoint_policy(policy).expect("billing transport"), + ReqwestGrokModelCatalogTransport::new(loopback_endpoint_policy(&origin)) + .expect("billing transport"), ); let session = GrokModelCatalogSession::new( SecretValue::new("access-token".to_owned()), @@ -198,15 +212,10 @@ async fn billing_transport_should_get_exact_credits_resource_without_redirect() "0.2.101", ) .expect("billing session"); - let request = GrokBillingRequest::for_endpoint( - origin - .join("v1/billing?format=credits") - .expect("billing URL"), - &session, - ) - .expect("billing request"); - - transport.execute(request).await.expect("billing response"); + GrokBillingClient::new(transport) + .fetch(&session) + .await + .expect("billing response"); } #[test] @@ -220,7 +229,7 @@ fn endpoint_policy_should_reject_private_and_documentation_addresses() { "::1", ] { assert!( - !GrokEndpointPolicy::accepts_resolved_address( + !OfficialGrokEndpointPolicy::accepts_resolved_address( address.parse().expect("fixture address") ), "{address} must be rejected" @@ -297,7 +306,10 @@ async fn official_oauth_transport_should_resolve_through_the_production_policy_w } let client = GrokOAuthClient::new( GrokOAuthConfig::official("0.2.101").expect("official config"), - Arc::new(ReqwestOAuthTransport::new().expect("production OAuth transport")), + Arc::new( + ReqwestOAuthTransport::new(Arc::new(OfficialGrokEndpointPolicy)) + .expect("production OAuth transport"), + ), Arc::new(FailClosedTokenVerifier), ); @@ -320,15 +332,12 @@ fn oauth_request_method_should_remain_typed() { } fn oauth_transport(origin: &Url) -> ReqwestOAuthTransport { - let endpoint_policy = GrokEndpointPolicy::loopback(origin).expect("loopback policy"); - ReqwestOAuthTransport::with_endpoint_policy(endpoint_policy) - .expect("injected loopback transport") + ReqwestOAuthTransport::new(loopback_endpoint_policy(origin)).expect("loopback transport") } fn inference_transport(origin: &Url) -> ReqwestGrokInferenceTransport { - let endpoint_policy = GrokEndpointPolicy::loopback(origin).expect("loopback policy"); - ReqwestGrokInferenceTransport::with_endpoint_policy(endpoint_policy) - .expect("injected loopback transport") + ReqwestGrokInferenceTransport::new(loopback_endpoint_policy(origin)) + .expect("loopback transport") } fn inference_request(origin: &Url) -> GrokInferenceRequest { diff --git a/backend/crates/providers/xai/tests/transport/request.rs b/backend/crates/providers/xai/tests/transport/request.rs index 5a28bc986..59f091b28 100644 --- a/backend/crates/providers/xai/tests/transport/request.rs +++ b/backend/crates/providers/xai/tests/transport/request.rs @@ -1,92 +1,155 @@ use gateway_core::operation::{ - ContentPart, GenerateRequest, Message, MessageRole, ProviderOptions, ReasoningEffort, - ReasoningRequirement, ReasoningSummary, ResponsePersistence, ToolDefinition, + ContentPart, GenerateRequest, Message, MessageRole, ProtocolPayload, ProviderOptions, }; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use provider_xai::{GrokRequestEncodeError, GrokResponsesRequest}; -#[test] -fn encoder_should_project_typed_generate_semantics() { - let message = Message::new( - MessageRole::User, - vec![ContentPart::Text("private prompt".to_owned())], +fn raw_request(body: Value) -> GenerateRequest { + let Value::Object(body) = body else { + panic!("request fixture must be an object"); + }; + GenerateRequest::from_protocol_payload( + Vec::new(), + ProtocolPayload::json_object("openai", body).expect("OpenAI payload"), ) - .expect("message is valid"); - let tool = ToolDefinition::new( - "lookup", - Some("lookup description".to_owned()), - json!({"type":"object"}) - .as_object() - .cloned() - .expect("schema is object"), - ) - .expect("tool is valid") - .with_strict(true); - let request = GenerateRequest::new(vec![message]) - .expect("generate is valid") - .with_tools(vec![tool]) - .with_max_output_tokens(256) - .with_response_persistence(ResponsePersistence::DoNotStore) - .with_reasoning(ReasoningRequirement { - effort: Some(ReasoningEffort::High), - summary: Some(ReasoningSummary::Concise), - }); +} - let encoded = - GrokResponsesRequest::encode(&request, "grok-code-test").expect("typed request encodes"); +#[test] +fn encoder_should_preserve_raw_images_hosted_tools_and_unknown_fields() { + let request = raw_request(json!({ + "model": "client-model", + "input": [{ + "role": "user", + "content": [ + {"type": "input_text", "text": "describe"}, + { + "type": "input_image", + "image_url": "data:image/png;base64,AQID", + "detail": "original", + "future_image_field": {"keep": true} + } + ] + }], + "tools": [ + {"type": "web_search_preview", "search_context_size": "high"}, + {"type": "code_interpreter", "container": {"type": "auto"}}, + {"type": "future_hosted_tool", "future": [1, 2, 3]} + ], + "tool_choice": "auto", + "future_official_field": {"nested": [true, 7]}, + "stream": false, + "store": true + })); + + let encoded = GrokResponsesRequest::encode(&request, "grok-routed").expect("raw request"); let body = Value::Object(encoded.body().clone()); + assert_eq!(body.pointer("/model"), Some(&json!("grok-routed"))); + assert_eq!(body.pointer("/stream"), Some(&json!(true))); + assert_eq!(body.pointer("/store"), Some(&json!(false))); assert_eq!( - ( - body.pointer("/store"), - body.pointer("/tools/0/type"), - body.pointer("/reasoning/effort") - ), - ( - Some(&json!(false)), - Some(&json!("function")), - Some(&json!("high")) - ) + body.pointer("/input/0/content/1/future_image_field"), + Some(&json!({"keep": true})) + ); + assert_eq!( + body.pointer("/tools/2"), + Some(&json!({"type": "future_hosted_tool", "future": [1, 2, 3]})) + ); + assert_eq!( + body.pointer("/future_official_field"), + Some(&json!({"nested": [true, 7]})) ); } #[test] -fn downstream_store_intent_should_not_enable_upstream_storage() { - let message = Message::new( - MessageRole::User, - vec![ContentPart::Text("persist inside gateway".to_owned())], - ) - .expect("message is valid"); - let request = GenerateRequest::new(vec![message]) - .expect("generate is valid") - .with_response_persistence(ResponsePersistence::Store); +fn account_identity_should_be_removed_without_touching_prompt_content() { + let request = raw_request(json!({ + "model": "client-model", + "input": [{ + "role": "user", + "content": [{ + "type": "input_text", + "text": "account_id and x-userid are ordinary prompt text" + }] + }], + "authorization": "Bearer attacker", + "account_id": "attacker-account", + "user_id": "attacker-user", + "team_id": "attacker-team", + "conversation_id": "attacker-conversation", + "previous_response_id": "attacker-response", + "metadata": { + "accountId": "nested-attacker-account", + "session_id": "nested-attacker-session", + "application_tag": "preserve-me" + } + })); - let encoded = - GrokResponsesRequest::encode(&request, "grok-code-test").expect("typed request encodes"); + let encoded = GrokResponsesRequest::encode(&request, "grok-routed").expect("sanitized request"); + let body = Value::Object(encoded.body().clone()); - assert_eq!(encoded.body().get("store"), Some(&json!(false))); + for pointer in [ + "/authorization", + "/account_id", + "/user_id", + "/team_id", + "/conversation_id", + "/previous_response_id", + "/metadata/accountId", + "/metadata/session_id", + ] { + assert_eq!( + body.pointer(pointer), + None, + "identity survived at {pointer}" + ); + } + assert_eq!( + body.pointer("/input/0/content/0/text"), + Some(&json!("account_id and x-userid are ordinary prompt text")) + ); + assert_eq!( + body.pointer("/metadata/application_tag"), + Some(&json!("preserve-me")) + ); } #[test] -fn request_debug_should_not_expose_prompt() { +fn typed_projection_should_not_exist_as_a_second_request_path() { let message = Message::new( MessageRole::User, vec![ContentPart::Text("private prompt".to_owned())], ) - .expect("message is valid"); - let request = GenerateRequest::new(vec![message]).expect("generate is valid"); - let encoded = - GrokResponsesRequest::encode(&request, "grok-code-test").expect("typed request encodes"); + .expect("message"); + let request = GenerateRequest::new(vec![message]).expect("typed operation"); - assert!(!format!("{encoded:?}").contains("private prompt")); + assert_eq!( + GrokResponsesRequest::encode(&request, "grok-routed") + .expect_err("missing raw payload must fail"), + GrokRequestEncodeError::InvalidProtocolPayload + ); } #[test] -fn encoder_should_reject_unknown_or_non_ascii_provider_options() { +fn request_debug_should_not_expose_prompt_or_unknown_values() { + let request = raw_request(json!({ + "model": "client-model", + "input": "private prompt", + "future_secret_shaped_value": "must-not-leak" + })); + let encoded = GrokResponsesRequest::encode(&request, "grok-routed").expect("raw request"); + let debug = format!("{encoded:?}"); + + assert!(!debug.contains("private prompt")); + assert!(!debug.contains("must-not-leak")); +} + +#[test] +fn provider_options_should_only_select_the_supported_transport() { for raw_options in [ - json!({"schema_version": 1, "raw_body": {"store": true}}), - json!({"schema_version": 1, "conversation_id": "会话"}), + json!({"schema_version": 1, "conversation_id": "attacker"}), + json!({"schema_version": 1, "transport": "websocket"}), ] { let mut provider_options = ProviderOptions::new(); provider_options @@ -98,23 +161,28 @@ fn encoder_should_reject_unknown_or_non_ascii_provider_options() { .expect("provider options object"), ) .expect("provider options"); - let request = GenerateRequest::new(vec![ - Message::new( - MessageRole::User, - vec![ContentPart::Text("private prompt".to_owned())], - ) - .expect("message"), - ]) - .expect("request") - .with_provider_options(provider_options); - - let error = GrokResponsesRequest::encode(&request, "grok-code-test") - .expect_err("unsafe provider option must fail"); + let request = raw_request(json!({"model": "client", "input": "hello"})) + .with_provider_options(provider_options); assert!(matches!( - error, - GrokRequestEncodeError::UnsupportedProviderOption - | GrokRequestEncodeError::InvalidProviderOptions + GrokResponsesRequest::encode(&request, "grok-routed"), + Err(GrokRequestEncodeError::UnsupportedProviderOption + | GrokRequestEncodeError::InvalidProviderOptions) )); } + + let mut provider_options = ProviderOptions::new(); + provider_options + .insert( + "xai", + Map::from_iter([ + ("schema_version".to_owned(), json!(1)), + ("transport".to_owned(), json!("http_sse")), + ]), + ) + .expect("provider options"); + let request = raw_request(json!({"model": "client", "input": "hello"})) + .with_provider_options(provider_options); + + assert!(GrokResponsesRequest::encode(&request, "grok-routed").is_ok()); } diff --git a/deploy/compose.yaml b/deploy/compose.yaml index cc559c1fe..35555a589 100644 --- a/deploy/compose.yaml +++ b/deploy/compose.yaml @@ -8,9 +8,6 @@ x-json-logging: &json-logging services: postgres: - extends: - file: config.yaml - service: postgres-credentials image: postgres:18-bookworm@sha256:d9c83446333daec3f0588cc709adb80c26090b7f9f0f7ec8d43c243385d79818 restart: unless-stopped command: @@ -23,6 +20,7 @@ services: TZ: Asia/Shanghai POSTGRES_DB: codex_proxy POSTGRES_USER: codex_proxy + POSTGRES_PASSWORD: ${CPR_DATABASE_PASSWORD:?CPR_DATABASE_PASSWORD is required} POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 volumes: - type: bind @@ -42,11 +40,14 @@ services: logging: *json-logging redis: - extends: - file: config.yaml - service: redis-credentials image: redis:8-alpine@sha256:9d317178eceac8454a2284a9e6df2466b93c745529947f0cd42a0fa9609d7005 restart: unless-stopped + command: + - redis-server + - --appendonly + - 'yes' + - --requirepass + - ${CPR_REDIS_PASSWORD:?CPR_REDIS_PASSWORD is required} volumes: - type: bind source: ../.runtime/redis @@ -56,6 +57,13 @@ services: ports: - '127.0.0.1:6379:6379' healthcheck: + test: + - CMD + - redis-cli + - --no-auth-warning + - -a + - ${CPR_REDIS_PASSWORD:?CPR_REDIS_PASSWORD is required} + - ping interval: 5s timeout: 5s retries: 10 @@ -120,6 +128,8 @@ services: CPR_SERVER_PORT: '8080' CPR_DATABASE_URL: postgres://codex_proxy@postgres:5432/codex_proxy CPR_REDIS_URL: redis://redis:6379/ + CPR_DATABASE_PASSWORD: ${CPR_DATABASE_PASSWORD:?CPR_DATABASE_PASSWORD is required} + CPR_REDIS_PASSWORD: ${CPR_REDIS_PASSWORD:?CPR_REDIS_PASSWORD is required} CPR_DEPLOYMENT_MODE: docker CPR_UPDATE_REPOSITORY: zyycn/codex-proxy-rs CPR_UPDATE_CHANNEL: stable diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 9fa8221d8..aedb30219 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -7,48 +7,13 @@ # # 三个密码建议分别执行 `openssl rand -hex 24` 生成。 -x-cpr: - schema_version: 1 +schema_version: 1 - server: - # 本地 cargo run 使用；Docker 内由 compose.yaml 固定覆盖为 0.0.0.0:8080。 +host: + listen: + # 本地 cargo run 使用；Docker 内由 CPR_SERVER_HOST/PORT 覆盖。 host: '127.0.0.1' port: 8080 - - database: - # 本地连接地址；Docker 内由 compose.yaml 固定改用 postgres 服务名。 - # 密码不要写进 URL。 - url: 'postgres://codex_proxy@127.0.0.1:5432/codex_proxy' - # 必填：48 位十六进制字符。 - password: &postgres_password '' - - redis: - # 本地连接地址；Docker 内由 compose.yaml 固定改用 redis 服务名。 - # 密码不要写进 URL。 - url: 'redis://127.0.0.1:6379/' - # 必填：48 位十六进制字符。 - password: &redis_password '' - - wire_profile: - # 该画像来自 Desktop 制品与 bundled Codex Core 的联合审计。 - # appcast 成功检查后会自动同步 Desktop 版本和 Build；其余字段保持此处配置。 - originator: 'Codex Desktop' - codex_version: '0.144.2' - desktop_version: '26.707.72221' - desktop_build: '5307' - # 固定参考主机画像；必须在重新抓包核验后再修改。 - os_type: 'Mac OS' - os_version: '15.7.1' - arch: 'arm64' - terminal: 'unknown' - verified_at: '2026-07-15T02:25:50+08:00' - - admin: - session_ttl_minutes: 1440 - default_username: 'admin@cpr.local' - # 仅首次创建管理员时使用；至少 12 个字符且不能包含 `$`。 - default_password: '' - logging: level: 'info' stdout: true @@ -59,25 +24,43 @@ x-cpr: retention_days: 14 max_file_size_mb: 20 max_files: 20 + drain_timeout_seconds: 30 + worker_shutdown_timeout_seconds: 30 -# Compose 凭据桥接区。通常不需要修改。 -services: - postgres-credentials: - environment: - POSTGRES_PASSWORD: *postgres_password +store: + database: + # Docker 内由 CPR_DATABASE_URL 覆盖；密码不要写进 URL。 + url: 'postgres://codex_proxy@127.0.0.1:5432/codex_proxy' + password: '' + redis: + # Docker 内由 CPR_REDIS_URL 覆盖；密码不要写进 URL。 + url: 'redis://127.0.0.1:6379/' + password: '' - redis-credentials: - command: - - redis-server - - --appendonly - - 'yes' - - --requirepass - - *redis_password - healthcheck: - test: - - CMD - - redis-cli - - --no-auth-warning - - -a - - *redis_password - - ping +admin: + session_ttl_minutes: 1440 + default_username: 'admin@cpr.local' + # 仅首次创建管理员时使用；至少 12 个字符且不能包含 `$`。 + default_password: '' + +api: + # Docker 内由 CPR_WEB_DIST_DIR 覆盖。 + asset_directory: '../frontend/dist' + cors_allowed_origins: [] + request_timeout_seconds: null + request_id_header: 'x-request-id' + +openai: + wire_profile: + # 该画像来自 Desktop 制品与 bundled Codex Core 的联合审计。 + originator: 'Codex Desktop' + codex_version: '0.144.2' + desktop_version: '26.707.72221' + desktop_build: '5307' + os_type: 'Mac OS' + os_version: '15.7.1' + arch: 'arm64' + terminal: 'unknown' + verified_at: '2026-07-15T02:25:50+08:00' + +xai: {}