From 2a0a50adc5851d2798f2ddc2522cbbc2e916f5d9 Mon Sep 17 00:00:00 2001 From: Zyy Date: Tue, 1 Sep 2026 17:17:57 +0800 Subject: [PATCH] refactor(backend): complete module boundary cleanup Extract runtime snapshot lifecycle into its semantic owner, split oversized provider and admin modules, and enforce mirrored src/tests module layouts. Co-authored-by: Codex --- backend/apps/gateway/tests/architecture.rs | 88 +- backend/crates/gateway-admin/src/lib.rs | 3 +- .../src/use_case/account_groups.rs | 3 +- .../gateway-admin/src/use_case/accounts.rs | 3 +- .../gateway-admin/src/use_case/backup.rs | 2 +- .../gateway-admin/src/use_case/client_keys.rs | 2 +- .../crates/gateway-admin/src/use_case/mod.rs | 3 +- .../gateway-admin/src/use_case/openai.rs | 2 +- .../gateway-admin/src/use_case/settings.rs | 2 +- .../crates/gateway-admin/src/use_case/xai.rs | 2 +- .../tests/model/{backup.rs => backup/mod.rs} | 2 +- .../gateway-admin/tests/use_case/mod.rs | 3 +- .../crates/gateway-api/src/admin/accounts.rs | 2346 ---------- .../src/admin/accounts/credentials.rs | 488 ++ .../src/admin/accounts/handlers.rs | 573 +++ .../gateway-api/src/admin/accounts/mod.rs | 80 + .../src/admin/accounts/presenter.rs | 413 ++ .../gateway-api/src/admin/accounts/wire.rs | 838 ++++ .../admin/{accounts.rs => accounts/mod.rs} | 0 backend/crates/gateway-api/tests/admin/mod.rs | 5 +- .../crates/gateway-api/tests/architecture.rs | 10 +- .../tests/openai/{images.rs => images/mod.rs} | 0 .../crates/gateway-api/tests/openai/mod.rs | 2 +- .../gateway-core/src/engine/execution.rs | 2 +- backend/crates/gateway-core/src/lib.rs | 7 +- .../gateway-core/src/routing/snapshot.rs | 353 +- .../crates/gateway-core/src/runtime/mod.rs | 355 ++ .../tests/{account.rs => account/mod.rs} | 0 .../gateway-core/tests/engine/execution.rs | 2 +- backend/crates/gateway-core/tests/main.rs | 1 + .../gateway-core/tests/routing/snapshot.rs | 166 +- .../crates/gateway-core/tests/runtime/mod.rs | 235 + .../gateway-store/src/redis/runtime_change.rs | 6 +- .../crates/gateway-store/tests/backup/mod.rs | 1 + .../gateway-store/tests/{ => backup}/s3.rs | 0 backend/crates/gateway-store/tests/main.rs | 2 +- .../mod.rs} | 0 .../mod.rs} | 0 .../crates/providers/openai/src/provider.rs | 3391 -------------- .../openai/src/provider/execution.rs | 916 ++++ .../providers/openai/src/provider/failure.rs | 964 ++++ .../providers/openai/src/provider/mod.rs | 494 ++ .../openai/src/provider/observation.rs | 796 ++++ .../providers/openai/src/provider/workers.rs | 260 ++ .../{websocket.rs => websocket/mod.rs} | 0 backend/crates/providers/xai/src/provider.rs | 2307 ---------- .../xai/src/provider/continuation.rs | 333 ++ .../providers/xai/src/provider/failure.rs | 557 +++ .../crates/providers/xai/src/provider/mod.rs | 627 +++ .../providers/xai/src/provider/stream.rs | 596 +++ .../providers/xai/src/provider/workers.rs | 228 + .../providers/xai/src/transport/request.rs | 4043 ----------------- .../xai/src/transport/request/history.rs | 578 +++ .../xai/src/transport/request/mod.rs | 923 ++++ .../xai/src/transport/request/response.rs | 839 ++++ .../xai/src/transport/request/tools.rs | 1752 +++++++ .../providers/xai/tests/provider/contract.rs | 2 +- .../transport/{request.rs => request/mod.rs} | 0 docs/architecture.md | 21 +- 59 files changed, 12999 insertions(+), 12628 deletions(-) rename backend/crates/gateway-admin/tests/model/{backup.rs => backup/mod.rs} (98%) delete mode 100644 backend/crates/gateway-api/src/admin/accounts.rs create mode 100644 backend/crates/gateway-api/src/admin/accounts/credentials.rs create mode 100644 backend/crates/gateway-api/src/admin/accounts/handlers.rs create mode 100644 backend/crates/gateway-api/src/admin/accounts/mod.rs create mode 100644 backend/crates/gateway-api/src/admin/accounts/presenter.rs create mode 100644 backend/crates/gateway-api/src/admin/accounts/wire.rs rename backend/crates/gateway-api/tests/admin/{accounts.rs => accounts/mod.rs} (100%) rename backend/crates/gateway-api/tests/openai/{images.rs => images/mod.rs} (100%) create mode 100644 backend/crates/gateway-core/src/runtime/mod.rs rename backend/crates/gateway-core/tests/{account.rs => account/mod.rs} (100%) create mode 100644 backend/crates/gateway-core/tests/runtime/mod.rs create mode 100644 backend/crates/gateway-store/tests/backup/mod.rs rename backend/crates/gateway-store/tests/{ => backup}/s3.rs (100%) rename backend/crates/gateway-store/tests/postgres/{observability.rs => observability/mod.rs} (100%) rename backend/crates/gateway-store/tests/postgres/{provider_accounts.rs => provider_accounts/mod.rs} (100%) delete mode 100644 backend/crates/providers/openai/src/provider.rs create mode 100644 backend/crates/providers/openai/src/provider/execution.rs create mode 100644 backend/crates/providers/openai/src/provider/failure.rs create mode 100644 backend/crates/providers/openai/src/provider/mod.rs create mode 100644 backend/crates/providers/openai/src/provider/observation.rs create mode 100644 backend/crates/providers/openai/src/provider/workers.rs rename backend/crates/providers/openai/tests/transport/{websocket.rs => websocket/mod.rs} (100%) delete mode 100644 backend/crates/providers/xai/src/provider.rs create mode 100644 backend/crates/providers/xai/src/provider/continuation.rs create mode 100644 backend/crates/providers/xai/src/provider/failure.rs create mode 100644 backend/crates/providers/xai/src/provider/mod.rs create mode 100644 backend/crates/providers/xai/src/provider/stream.rs create mode 100644 backend/crates/providers/xai/src/provider/workers.rs delete mode 100644 backend/crates/providers/xai/src/transport/request.rs create mode 100644 backend/crates/providers/xai/src/transport/request/history.rs create mode 100644 backend/crates/providers/xai/src/transport/request/mod.rs create mode 100644 backend/crates/providers/xai/src/transport/request/response.rs create mode 100644 backend/crates/providers/xai/src/transport/request/tools.rs rename backend/crates/providers/xai/tests/transport/{request.rs => request/mod.rs} (100%) diff --git a/backend/apps/gateway/tests/architecture.rs b/backend/apps/gateway/tests/architecture.rs index 0f51124a9..0994f93f1 100644 --- a/backend/apps/gateway/tests/architecture.rs +++ b/backend/apps/gateway/tests/architecture.rs @@ -1,7 +1,7 @@ //! docs/architecture.md 依赖 DAG 与生产源码纪律的 workspace 级机器校验。 use std::{ - collections::BTreeSet, + collections::{BTreeMap, BTreeSet}, fs, path::{Path, PathBuf}, }; @@ -98,7 +98,7 @@ const PACKAGE_TO_MEMBER: &[(&str, &str)] = &[ ("provider-xai", "crates/providers/xai"), ]; -/// Adapter/provider 根门面暂时允许公开的模块；收窄公开面时必须同步缩小本表。 +/// Adapter/provider 根门面的稳定合同模块；任何增减都必须同步完成边界审计。 const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[ ("crates/gateway-api", &["admin", "openai"]), ( @@ -122,6 +122,12 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[ ), ]; +/// 不对应单一生产模块、而是校验 crate/workspace 整体契约的根级测试场景。 +const ROOT_TEST_SCENARIOS: &[(&str, &[&str])] = &[ + ("apps/gateway", &["architecture"]), + ("crates/gateway-api", &["architecture"]), +]; + /// 冻结的 workspace 内部运行时依赖边；新增/删除任何边都必须同步本表。 const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[ ("codex-proxy-rs", "gateway-admin"), @@ -254,6 +260,84 @@ fn workspace_modules_follow_conventional_file_layout() { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ModuleLayout { + File, + Directory, +} + +#[test] +fn integration_tests_mirror_production_module_tree() { + for member in WORKSPACE_MEMBERS { + let member_root = backend_root().join(member); + let tests_root = member_root.join("tests"); + if !tests_root.is_dir() { + continue; + } + + let production = module_layouts(&member_root.join("src"), &["lib.rs", "main.rs"]); + let tests = module_layouts(&tests_root, &["main.rs"]); + for (module, test_layout) in tests { + if module == Path::new("support") { + continue; + } + + if let Some(production_layout) = production.get(&module) { + assert_eq!( + test_layout, + *production_layout, + "{} must mirror the production module layout for {}", + tests_root.join(module.with_extension("rs")).display(), + member_root.join("src").join(&module).display(), + ); + continue; + } + + let has_production_owner = module + .parent() + .into_iter() + .flat_map(Path::ancestors) + .take_while(|ancestor| !ancestor.as_os_str().is_empty()) + .any(|ancestor| production.contains_key(ancestor)); + assert!( + has_production_owner || root_test_scenario_allowed(member, &module), + "{} has no mirrored production module owner", + tests_root.join(module.with_extension("rs")).display(), + ); + } + } +} + +fn module_layouts(root: &Path, crate_roots: &[&str]) -> BTreeMap { + super::rust_files(root) + .into_iter() + .filter(|relative| { + !crate_roots + .iter() + .any(|candidate| relative == Path::new(candidate)) + }) + .map(|relative| { + if relative.file_name().and_then(|value| value.to_str()) == Some("mod.rs") { + ( + relative.parent().expect("mod.rs parent").to_path_buf(), + ModuleLayout::Directory, + ) + } else { + (relative.with_extension(""), ModuleLayout::File) + } + }) + .collect() +} + +fn root_test_scenario_allowed(member: &str, module: &Path) -> bool { + let Some(module) = module.to_str() else { + return false; + }; + ROOT_TEST_SCENARIOS + .iter() + .any(|(candidate, scenarios)| *candidate == member && scenarios.contains(&module)) +} + fn assert_module_tree(root: &Path, crate_roots: &[&str]) { let files = super::rust_files(root); for relative in &files { diff --git a/backend/crates/gateway-admin/src/lib.rs b/backend/crates/gateway-admin/src/lib.rs index 14d515427..e5041a9b7 100644 --- a/backend/crates/gateway-admin/src/lib.rs +++ b/backend/crates/gateway-admin/src/lib.rs @@ -6,7 +6,8 @@ use std::{fmt, path::Path, sync::Arc, time::Duration}; use gateway_core::{ engine::probe::AccountProbe, - routing::{ProviderKind, snapshot::SnapshotControl}, + routing::ProviderKind, + runtime::SnapshotControl, task::{ DaemonRestartPolicy, WorkerContribution, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, diff --git a/backend/crates/gateway-admin/src/use_case/account_groups.rs b/backend/crates/gateway-admin/src/use_case/account_groups.rs index d97e22c82..f6321877a 100644 --- a/backend/crates/gateway-admin/src/use_case/account_groups.rs +++ b/backend/crates/gateway-admin/src/use_case/account_groups.rs @@ -9,7 +9,8 @@ use std::{ use async_trait::async_trait; use gateway_core::{ account::{AccountStatus, resolve_account_status}, - routing::{AccountGroupId, snapshot::SnapshotControl}, + routing::AccountGroupId, + runtime::SnapshotControl, }; use uuid::Uuid; diff --git a/backend/crates/gateway-admin/src/use_case/accounts.rs b/backend/crates/gateway-admin/src/use_case/accounts.rs index 0a7f11350..2d0620369 100644 --- a/backend/crates/gateway-admin/src/use_case/accounts.rs +++ b/backend/crates/gateway-admin/src/use_case/accounts.rs @@ -8,7 +8,8 @@ use futures::StreamExt as _; use gateway_core::{ account::ProviderAccountId, engine::probe::{AccountProbe, AccountProbeRequest}, - routing::{ProviderKind, UpstreamModelId, snapshot::SnapshotControl}, + routing::{ProviderKind, UpstreamModelId}, + runtime::SnapshotControl, }; use crate::{ diff --git a/backend/crates/gateway-admin/src/use_case/backup.rs b/backend/crates/gateway-admin/src/use_case/backup.rs index 87334b13c..9ad0ee90f 100644 --- a/backend/crates/gateway-admin/src/use_case/backup.rs +++ b/backend/crates/gateway-admin/src/use_case/backup.rs @@ -5,7 +5,7 @@ use std::time::Duration; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use gateway_core::routing::snapshot::SnapshotControl; +use gateway_core::runtime::SnapshotControl; use secrecy::ExposeSecret as _; use uuid::Uuid; diff --git a/backend/crates/gateway-admin/src/use_case/client_keys.rs b/backend/crates/gateway-admin/src/use_case/client_keys.rs index ef1f38bdf..7517ef717 100644 --- a/backend/crates/gateway-admin/src/use_case/client_keys.rs +++ b/backend/crates/gateway-admin/src/use_case/client_keys.rs @@ -5,7 +5,7 @@ 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 gateway_core::runtime::SnapshotControl; use rand_core::{OsRng, RngCore as _}; use uuid::Uuid; diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index f36935b8a..4223529d2 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -31,7 +31,8 @@ use crate::{ }; use gateway_core::{ account::ProviderAccountId, - routing::{ConfigRevision, ProviderKind, snapshot::SnapshotControl}, + routing::{ConfigRevision, ProviderKind}, + runtime::SnapshotControl, }; fn map_store_error(error: AdminStoreError, resource: &'static str) -> AdminError { diff --git a/backend/crates/gateway-admin/src/use_case/openai.rs b/backend/crates/gateway-admin/src/use_case/openai.rs index 9b04532de..85f47638f 100644 --- a/backend/crates/gateway-admin/src/use_case/openai.rs +++ b/backend/crates/gateway-admin/src/use_case/openai.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use async_trait::async_trait; -use gateway_core::{account::ProviderAccountId, routing::snapshot::SnapshotControl}; +use gateway_core::{account::ProviderAccountId, runtime::SnapshotControl}; use crate::{ model::{ diff --git a/backend/crates/gateway-admin/src/use_case/settings.rs b/backend/crates/gateway-admin/src/use_case/settings.rs index 5e62eaddd..7040234d2 100644 --- a/backend/crates/gateway-admin/src/use_case/settings.rs +++ b/backend/crates/gateway-admin/src/use_case/settings.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use async_trait::async_trait; use gateway_core::policy::CodexClientVersion; -use gateway_core::routing::snapshot::SnapshotControl; +use gateway_core::runtime::SnapshotControl; use rand_core::{OsRng, RngCore as _}; use crate::{ diff --git a/backend/crates/gateway-admin/src/use_case/xai.rs b/backend/crates/gateway-admin/src/use_case/xai.rs index b32e044aa..82e452977 100644 --- a/backend/crates/gateway-admin/src/use_case/xai.rs +++ b/backend/crates/gateway-admin/src/use_case/xai.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use async_trait::async_trait; -use gateway_core::routing::snapshot::SnapshotControl; +use gateway_core::runtime::SnapshotControl; use crate::{ model::{ diff --git a/backend/crates/gateway-admin/tests/model/backup.rs b/backend/crates/gateway-admin/tests/model/backup/mod.rs similarity index 98% rename from backend/crates/gateway-admin/tests/model/backup.rs rename to backend/crates/gateway-admin/tests/model/backup/mod.rs index 23ed38c54..c3fe2e772 100644 --- a/backend/crates/gateway-admin/tests/model/backup.rs +++ b/backend/crates/gateway-admin/tests/model/backup/mod.rs @@ -1,4 +1,4 @@ -//! 备份领域模型（model/backup.rs）的纯逻辑测试。 +//! 备份领域模型（model/backup）的纯逻辑测试。 use chrono::{TimeZone as _, Utc}; use secrecy::SecretString; diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs index 22781c17b..e2d0f2e3f 100644 --- a/backend/crates/gateway-admin/tests/use_case/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -74,7 +74,8 @@ use gateway_core::{ engine::probe::{AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult}, error::{GatewayError, GatewayErrorKind}, policy::ClientApiKeyId, - routing::{ConfigRevision, ProviderKind, snapshot::SnapshotControl}, + routing::{ConfigRevision, ProviderKind}, + runtime::SnapshotControl, }; pub(super) struct AdminHarness { diff --git a/backend/crates/gateway-api/src/admin/accounts.rs b/backend/crates/gateway-api/src/admin/accounts.rs deleted file mode 100644 index 89aab2f2e..000000000 --- a/backend/crates/gateway-api/src/admin/accounts.rs +++ /dev/null @@ -1,2346 +0,0 @@ -//! 管理端账号目录、查询校验与安全响应 wire。 - -use std::{collections::BTreeSet, convert::Infallible, fmt}; - -use axum::{ - Router, - body::Body, - extract::State, - http::{HeaderName, HeaderValue, StatusCode, header}, - response::{ - IntoResponse, Response, - sse::{Event, KeepAlive, Sse}, - }, - routing::{get, post}, -}; -use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; -use chrono::{DateTime, FixedOffset, Utc}; -use futures::{Stream, StreamExt as _}; -use gateway_admin::model::{ - AdminError as AdminServiceError, PageSize, - accounts::{ - AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent, - AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort, - AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage, - AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount, - }, - provider_credentials::{ - AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult, - AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit, - CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation, - CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels, - ProviderProfileActivityInsights, ProviderProfileAvatar, ProviderProfileDailyUsage, - ProviderProfileInvocation, ProviderProfileStatistics, ProviderProfileStatisticsSummary, - ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult, - ProviderResetCredits, RotateCredential, StartAuthorization, - }, -}; -use gateway_core::{ - account::{OpaqueProviderData, ProviderAccountId}, - routing::{AccountGroupId, ProviderKind, UpstreamModelId}, -}; -use serde::{Deserialize, Deserializer, Serialize}; -use serde_json::{Map, Value}; -use uuid::{Uuid, Version}; - -use super::presenter::{format_compact_number, format_decimal_currency, format_number}; -use super::{ - AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse, AdminSessionState, - PageMeta, WireValidationError, -}; - -const DEFAULT_PAGE_SIZE: u32 = 50; -const MAX_PAGE_SIZE: u32 = 200; -const MAX_SEARCH_BYTES: usize = 256; -const MAX_ID_BYTES: usize = 256; -const MAX_NAME_BYTES: usize = 512; -const MAX_IMPORT_DATA_BYTES: usize = 64 * 1024 * 1024; -const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024; -const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024; -const MAX_ID_TOKEN_BYTES: usize = 16 * 1024; -const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024; -const MAX_ACCOUNT_DELETE_BATCH: usize = 200; -const MAX_ACCOUNT_GROUP_BATCH: usize = 1000; -const MAX_AVATAR_VERSION_BYTES: usize = 32; - -/// 账号列表查询参数。 -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct ListQuery { - pub page: Option, - pub page_size: Option, - pub provider: Option, - pub group_id: Option, - pub search: Option, - pub status: Option, - pub sort_by: Option, - pub sort_direction: Option, -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct BatchUpdateAccountsRequest { - pub account_ids: Vec, - pub enabled: bool, - #[serde(deserialize_with = "deserialize_required_nullable")] - pub concurrency_limit: Option, - pub weight: u64, - pub group_ids: Vec, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct BatchUpdatedAccountsData { - account_ids: Vec, - config_revision: u64, -} - -impl From for BatchUpdatedAccountsData { - fn from(result: AccountsUpdateResult) -> Self { - Self { - account_ids: result - .account_ids - .into_iter() - .map(|id| id.to_string()) - .collect(), - config_revision: result.config_revision.get(), - } - } -} - -impl BatchUpdateAccountsRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - if self.account_ids.is_empty() - || self.account_ids.len() > MAX_ACCOUNT_GROUP_BATCH - || self - .account_ids - .iter() - .any(|id| require_account_id(id, "accountIds").is_err()) - || self.account_ids.iter().collect::>().len() != self.account_ids.len() - { - return Err(WireValidationError::new("accountIds")); - } - validate_wire_group_ids(&self.group_ids)?; - parse_concurrency_limit(self.concurrency_limit)?; - parse_account_weight(self.weight)?; - Ok(()) - } - - fn into_command(self) -> Result { - self.validate()?; - Ok(BatchUpdateAccounts { - account_ids: self.account_ids, - enabled: self.enabled, - concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, - weight: parse_account_weight(self.weight)?, - group_ids: validate_wire_group_ids(&self.group_ids)?, - }) - } -} - -impl ListQuery { - /// 解析并校验全部 wire 字段，生成 Admin 查询命令。 - pub fn validate(self) -> Result { - let page = self.page.unwrap_or(1); - if page == 0 { - return Err(WireValidationError::new("page")); - } - let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE); - if page_size == 0 || page_size > MAX_PAGE_SIZE { - return Err(WireValidationError::new("pageSize")); - } - let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?; - let group_filter = match self.group_id.as_deref().map(str::trim) { - None | Some("") => None, - Some("ungrouped") => Some(AccountGroupFilter::Ungrouped), - Some(value) => Some(AccountGroupFilter::Group( - AccountGroupId::new(value.to_owned()) - .map_err(|_| WireValidationError::new("groupId"))?, - )), - }; - let search = self - .search - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()); - if search.as_deref().is_some_and(|value| { - value.len() > MAX_SEARCH_BYTES || value.chars().any(char::is_control) - }) { - return Err(WireValidationError::new("search")); - } - let status = match self.status.as_deref().map(str::trim) { - None | Some("") => None, - Some(value) => Some( - DomainAccountStatus::parse(&value.to_ascii_lowercase()) - .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: 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(AccountListQuery { - page, - page_size: PageSize::new( - u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?, - ) - .map_err(|_| WireValidationError::new("pageSize"))?, - provider_kind, - group_filter, - search, - status, - sort, - }) - } -} - -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")) -} - -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, - } -} - -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, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountPageData { - pub items: Vec, - pub page: PageMeta, - pub summary: AccountSummaryView, -} - -/// 账号概览计数。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountSummaryView { - pub total: u64, - pub normal: u64, - pub quota_exhausted: u64, - pub rate_limited: u64, - pub disabled: u64, - pub error: u64, -} - -/// 一条安全账号视图。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountView { - pub id: String, - pub name: String, - pub provider: String, - pub groups: Vec, - pub resource_ref: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub authentication_kind: String, - pub has_refresh_token: bool, - pub status: String, - /// `status == "error"` 时的具体原因；其余状态为 `null`。 - pub error_reason: Option, - /// 最近一次失败的上游错误描述；仅错误状态存在。 - pub error_message: Option, - pub enabled: bool, - pub concurrency_limit: Option, - pub weight: u16, - pub access_token_expires_at: Option, - pub access_token_expires_at_display: Option, - pub refresh_token_expires_at: Option, - pub next_refresh_at: Option, - pub added_at: String, - pub added_at_display: String, - pub updated_at: String, - pub updated_at_display: String, - pub quota: AccountQuotaView, - pub usage: AccountUsageView, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountGroupRefView { - pub id: String, - pub name: String, - pub color: String, - pub enabled: bool, -} - -/// Provider quota 安全视图。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountQuotaView { - pub refreshed_at_display: String, - pub limit_reached: bool, - /// 429 临时限流（Redis 冷却）到期时间展示；非限流中为 `null`。 - pub rate_limited_until: Option, - pub windows: Vec, -} - -/// 一个 quota 时间窗口。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountQuotaWindowView { - pub key: String, - pub group: String, - pub limit_id: Option, - pub limit_name: Option, - pub role: Option, - pub window_seconds: Option, - pub label_display: String, - pub window_label_display: String, - pub used_percent: Option, - pub used_percent_display: String, - pub limit_reached: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub local_usage: Option, - pub reset_at_display: String, -} - -/// 账号观测用量。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountUsageView { - pub request_count: Option, - pub request_count_display: String, - pub input_tokens: Option, - pub input_tokens_display: String, - pub output_tokens: Option, - pub output_tokens_display: String, - pub cached_tokens: Option, - pub cached_tokens_display: String, - pub image_input_tokens: Option, - pub image_input_tokens_display: String, - pub image_output_tokens: Option, - pub image_output_tokens_display: String, - pub image_request_count: Option, - pub image_request_count_display: String, - pub image_request_failed_count: Option, - pub image_request_failed_count_display: String, - pub total_tokens: Option, - pub total_tokens_display: String, - pub created_tokens: Option, - pub created_tokens_display: String, - pub read_tokens: Option, - pub read_tokens_display: String, - pub last_used_at: Option, - pub last_used_at_display: String, - pub cost_estimate_status: String, - pub known_cost_count: Option, - pub partial_cost_count: Option, - pub unknown_cost_count: Option, - pub costs: Vec, - pub models: Vec, -} - -/// 凭据在单个上游模型上的观测用量。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ModelUsageView { - pub model: String, - pub request_count: u64, - pub request_count_display: String, - pub success_rate: Option, - pub success_rate_display: String, - pub input_tokens: Option, - pub input_tokens_display: String, - pub output_tokens: Option, - pub output_tokens_display: String, - pub cached_tokens: Option, - pub cached_tokens_display: String, - pub image_input_tokens: Option, - pub image_input_tokens_display: String, - pub image_output_tokens: Option, - pub image_output_tokens_display: String, - pub image_request_count: u64, - pub image_request_count_display: String, - pub image_request_failed_count: u64, - pub image_request_failed_count_display: String, - pub total_tokens: Option, - pub total_tokens_display: String, - pub billing_amount_usd: Option, - pub billing_amount_usd_display: String, - pub cost_estimate_status: String, - pub known_cost_count: u64, - pub partial_cost_count: u64, - pub unknown_cost_count: u64, - pub costs: Vec, - pub last_used_at: String, - pub last_used_at_display: String, -} - -/// 单一货币的可查询成本。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct CurrencyCostView { - pub currency: String, - pub estimated_amount: String, - pub estimated_amount_display: String, -} - -/// 账号详情类 GET 的固定 ID query。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountIdQuery { - pub account_id: String, -} - -impl AccountIdQuery { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId") - } - - fn into_id(self) -> Result { - self.validate()?; - ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) - } -} - -/// 头像 GET 的固定 query；`version` 只参与浏览器缓存键，不进入上游请求。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountProfileAvatarQuery { - pub account_id: String, - pub version: Option, -} - -impl AccountProfileAvatarQuery { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId")?; - if self.version.as_deref().is_some_and(|version| { - version.is_empty() - || version.len() > MAX_AVATAR_VERSION_BYTES - || !version.bytes().all(|byte| byte.is_ascii_alphanumeric()) - }) { - return Err(WireValidationError::new("version")); - } - Ok(()) - } - - fn into_id(self) -> Result { - self.validate()?; - ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) - } -} - -/// 敏感导出的固定 query；IDs 使用逗号分隔，禁止隐式导出全部账号。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountExportQuery { - pub account_ids: String, - pub confirm: String, -} - -impl AccountExportQuery { - pub fn into_ids(self) -> Result, WireValidationError> { - if self.confirm != "export_sensitive_accounts" { - return Err(WireValidationError::new("confirm")); - } - let ids = self - .account_ids - .split(',') - .map(str::trim) - .filter(|id| !id.is_empty()) - .map(ToOwned::to_owned) - .collect::>(); - if ids.is_empty() - || ids.len() > 200 - || ids - .iter() - .any(|id| require_account_id(id, "accountIds").is_err()) - { - return Err(WireValidationError::new("accountIds")); - } - let unique = ids.iter().collect::>(); - if unique.len() != ids.len() { - return Err(WireValidationError::new("accountIds")); - } - ids.into_iter() - .map(|id| { - ProviderAccountId::new(id).map_err(|_| WireValidationError::new("accountIds")) - }) - .collect() - } -} - -/// 账号运行期动作。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountActionRequest { - pub account_id: String, -} - -impl AccountActionRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId") - } - - fn into_id(self) -> Result { - self.validate()?; - ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) - } -} - -/// 主动额度重置卡消费请求。幂等键由 UI 生成并在不确定重试时复用，与官方一致。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountResetCreditConsumeRequest { - pub account_id: String, - #[serde(default)] - pub credit_id: Option, - pub redeem_request_id: String, -} - -impl AccountResetCreditConsumeRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId")?; - if self.credit_id.as_deref().is_some_and(|value| { - let value = value.trim(); - value.is_empty() || value.len() > 512 || value.chars().any(char::is_control) - }) { - return Err(WireValidationError::new("creditId")); - } - let value = self.redeem_request_id.trim(); - let redeem_request_id = - Uuid::parse_str(value).map_err(|_| WireValidationError::new("redeemRequestId"))?; - if redeem_request_id.get_version() != Some(Version::Random) - || redeem_request_id.hyphenated().to_string() != value - { - return Err(WireValidationError::new("redeemRequestId")); - } - Ok(()) - } - - fn into_command(self) -> Result { - self.validate()?; - Ok(ConsumeProviderResetCredit { - account_id: ProviderAccountId::new(self.account_id) - .map_err(|_| WireValidationError::new("accountId"))?, - credit_id: self.credit_id.map(|value| value.trim().to_owned()), - redeem_request_id: Uuid::parse_str(&self.redeem_request_id) - .map_err(|_| WireValidationError::new("redeemRequestId"))?, - }) - } -} - -/// 手工 OAuth 刷新会变更持久 credential。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountRefreshRequest { - pub account_id: String, -} - -impl AccountRefreshRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId") - } - - fn into_command(self) -> Result { - self.validate()?; - ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) - } -} - -/// 连接测试 query；测试仍经唯一 Core/Provider 模型请求路径执行。 -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountTestQuery { - pub account_id: String, - pub model_id: String, -} - -impl AccountTestQuery { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId")?; - if self.model_id.trim().is_empty() || self.model_id.chars().any(char::is_control) { - return Err(WireValidationError::new("modelId")); - } - Ok(()) - } - - fn into_command(self) -> Result<(ProviderAccountId, UpstreamModelId), WireValidationError> { - self.validate()?; - Ok(( - ProviderAccountId::new(self.account_id) - .map_err(|_| WireValidationError::new("accountId"))?, - UpstreamModelId::new(self.model_id).map_err(|_| WireValidationError::new("modelId"))?, - )) - } -} - -#[derive(Debug, Clone, Serialize)] -pub struct AccountModelView { - pub id: String, - pub label: String, -} - -#[derive(Debug, Clone, Serialize)] -pub struct AccountModelsData { - pub models: Vec, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountRefreshData { - pub account: AccountView, -} - -#[derive(Debug, Clone, Serialize)] -pub struct AccountQuotaData { - pub account: AccountView, -} - -/// Provider 官方个人资料统计响应。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountProfileStatisticsData { - pub display_name: Option, - pub username: Option, - pub image_url: Option, - pub has_stats_error: bool, - pub summary: AccountProfileStatisticsSummaryView, - pub daily_usage: Option>, - pub activity_insights: AccountProfileActivityInsightsView, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountProfileStatisticsSummaryView { - pub total_text_tokens: Option, - pub peak_tokens: Option, - pub longest_task_duration_ms: Option, - pub current_streak_days: Option, - pub longest_streak_days: Option, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountProfileDailyUsageView { - pub date: String, - pub tokens: u64, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountProfileActivityInsightsView { - pub fast_mode_percent: Option, - pub reasoning_effort: Option, - pub reasoning_effort_percent: Option, - pub skills_explored: Option, - pub total_skills_used: Option, - pub total_threads: Option, - pub invocations: Option>, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountProfileInvocationView { - #[serde(rename = "type")] - pub invocation_type: String, - pub plugin_id: Option, - pub plugin_name: Option, - pub skill_id: Option, - pub skill_name: Option, - pub usage_count: Option, -} - -impl From for AccountProfileStatisticsData { - fn from(statistics: ProviderProfileStatistics) -> Self { - Self { - display_name: statistics.display_name, - username: statistics.username, - image_url: statistics.image_url, - has_stats_error: statistics.has_stats_error, - summary: profile_statistics_summary_view(statistics.summary), - daily_usage: statistics - .daily_usage - .map(|daily| daily.into_iter().map(profile_daily_usage_view).collect()), - activity_insights: profile_activity_insights_view(statistics.activity_insights), - } - } -} - -const fn profile_statistics_summary_view( - summary: ProviderProfileStatisticsSummary, -) -> AccountProfileStatisticsSummaryView { - AccountProfileStatisticsSummaryView { - total_text_tokens: summary.total_text_tokens, - peak_tokens: summary.peak_tokens, - longest_task_duration_ms: summary.longest_task_duration_ms, - current_streak_days: summary.current_streak_days, - longest_streak_days: summary.longest_streak_days, - } -} - -fn profile_daily_usage_view(daily: ProviderProfileDailyUsage) -> AccountProfileDailyUsageView { - AccountProfileDailyUsageView { - date: daily.date.to_string(), - tokens: daily.tokens, - } -} - -fn profile_activity_insights_view( - insights: ProviderProfileActivityInsights, -) -> AccountProfileActivityInsightsView { - AccountProfileActivityInsightsView { - fast_mode_percent: insights.fast_mode_percent, - reasoning_effort: insights.reasoning_effort, - reasoning_effort_percent: insights.reasoning_effort_percent, - skills_explored: insights.skills_explored, - total_skills_used: insights.total_skills_used, - total_threads: insights.total_threads, - invocations: insights.invocations.map(|invocations| { - invocations - .into_iter() - .map(profile_invocation_view) - .collect() - }), - } -} - -fn profile_invocation_view(invocation: ProviderProfileInvocation) -> AccountProfileInvocationView { - AccountProfileInvocationView { - invocation_type: invocation.invocation_type, - plugin_id: invocation.plugin_id, - plugin_name: invocation.plugin_name, - skill_id: invocation.skill_id, - skill_name: invocation.skill_name, - usage_count: invocation.usage_count, - } -} - -/// 主动额度重置卡列表响应。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountResetCreditsData { - pub available_count: u64, - pub credits: Vec, -} - -/// 一张安全主动额度重置卡视图。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountResetCreditView { - pub id: String, - pub status: Option, - pub title: Option, - pub expires_at: Option, - pub reset_type: Option, -} - -/// 主动额度重置卡消费响应。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountResetCreditResultData { - pub code: String, - pub credit: Option, -} - -impl From for AccountResetCreditView { - fn from(credit: ProviderResetCredit) -> Self { - Self { - id: credit.id, - status: credit.status, - title: credit.title, - expires_at: credit.expires_at.map(|value| value.to_rfc3339()), - reset_type: credit.reset_type, - } - } -} - -impl From for AccountResetCreditsData { - fn from(credits: ProviderResetCredits) -> Self { - Self { - available_count: credits.available_count, - credits: credits - .credits - .into_iter() - .map(AccountResetCreditView::from) - .collect(), - } - } -} - -impl From for AccountResetCreditResultData { - fn from(result: ProviderResetCreditResult) -> Self { - Self { - code: result.code, - credit: result.credit.map(AccountResetCreditView::from), - } - } -} - -/// Provider-owned 明文导出文档；Debug 永远不输出内部 JSON。 -#[derive(Serialize)] -#[serde(transparent)] -pub struct AccountExportData(Value); - -impl AccountExportData { - #[must_use] - 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 { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("AccountExportData()") - } -} - -pub struct AccountConnectionTestEvent { - pub data: Value, -} - -impl From for AccountConnectionTestEvent { - fn from(event: DomainConnectionTestEvent) -> Self { - let data = match event { - DomainConnectionTestEvent::Started { model } => serde_json::json!({ - "type": "test_start", - "model": model, - "text": "正在连接上游 Responses" - }), - 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 => serde_json::json!({ - "type": "test_complete", - "success": true - }), - DomainConnectionTestEvent::Failed { - source, - gateway_error_code, - send_state, - message, - provider_error_code, - provider_error_type, - upstream_status, - upstream_content_type, - upstream_body, - } => serde_json::json!({ - "type": "error", - "source": source.as_str(), - "gatewayErrorCode": gateway_error_code.as_str(), - "sendState": send_state.map(gateway_core::upstream::UpstreamSendState::as_str), - "error": message, - "providerErrorCode": provider_error_code, - "providerErrorType": provider_error_type, - "upstreamStatus": upstream_status, - "upstreamContentType": upstream_content_type, - "upstreamBody": upstream_body - }), - }; - Self { data } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AccountProvider { - OpenAi, - Xai, -} - -impl AccountProvider { - fn parse(value: &str) -> Result { - match value.trim() { - "openai" => Ok(Self::OpenAi), - "xai" => Ok(Self::Xai), - _ => Err(WireValidationError::new("provider")), - } - } -} - -/// Provider-owned 账号导入请求；公共 API 不解释 `data` 内部字段。 -#[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountImportRequest { - pub provider: String, - pub data: Value, -} - -impl AccountImportRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - AccountProvider::parse(&self.provider)?; - if !self.data.is_object() - || serde_json::to_vec(&self.data) - .map_or(true, |encoded| encoded.len() > MAX_IMPORT_DATA_BYTES) - { - return Err(WireValidationError::new("data")); - } - Ok(()) - } - - fn into_command( - self, - context: gateway_admin::model::MutationContext, - ) -> Result<(AccountProvider, ImportCredentials), WireValidationError> { - self.validate()?; - let provider = AccountProvider::parse(&self.provider)?; - Ok(( - provider, - ImportCredentials { - context, - document: provider_document(self.data, "data")?, - }, - )) - } -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct StartAccountAuthorizationRequest { - pub provider: String, - pub name: String, - pub account_id: Option, -} - -impl StartAccountAuthorizationRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - AccountProvider::parse(&self.provider)?; - require_text(&self.name, MAX_NAME_BYTES, "name")?; - if let Some(account_id) = self.account_id.as_deref() { - require_account_id(account_id, "accountId")?; - } - Ok(()) - } - - fn into_command( - self, - context: gateway_admin::model::MutationContext, - ) -> Result<(AccountProvider, StartAuthorization), WireValidationError> { - self.validate()?; - let provider = AccountProvider::parse(&self.provider)?; - let reauthorization = self - .account_id - .map(ProviderAccountId::new) - .transpose() - .map_err(|_| WireValidationError::new("accountId"))?; - Ok(( - provider, - StartAuthorization { - context, - name: self.name, - reauthorization, - }, - )) - } -} - -#[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct CompleteAccountAuthorizationRequest { - pub provider: String, - pub flow_id: String, - pub callback_url: String, -} - -impl CompleteAccountAuthorizationRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - let provider = AccountProvider::parse(&self.provider)?; - match provider { - AccountProvider::OpenAi => { - if !URL_SAFE_NO_PAD - .decode(&self.flow_id) - .is_ok_and(|decoded| decoded.len() == 32) - { - return Err(WireValidationError::new("flowId")); - } - } - AccountProvider::Xai => require_wire_id(&self.flow_id, "flowId")?, - } - require_text(&self.callback_url, MAX_CALLBACK_URL_BYTES, "callbackUrl") - } - - fn into_command( - self, - context: gateway_admin::model::MutationContext, - ) -> Result<(AccountProvider, CompleteAuthorization), WireValidationError> { - self.validate()?; - let provider = AccountProvider::parse(&self.provider)?; - Ok(( - provider, - CompleteAuthorization { - context, - flow_id: self.flow_id, - callback_url: self.callback_url, - }, - )) - } -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct UpdateAccountRequest { - pub account_id: String, - pub enabled: bool, - #[serde(deserialize_with = "deserialize_required_nullable")] - pub concurrency_limit: Option, - pub weight: u64, - pub group_ids: Vec, -} - -impl UpdateAccountRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - require_account_id(&self.account_id, "accountId")?; - parse_concurrency_limit(self.concurrency_limit)?; - parse_account_weight(self.weight)?; - validate_wire_group_ids(&self.group_ids)?; - Ok(()) - } - - fn into_command(self) -> Result { - self.validate()?; - Ok(UpdateAccount { - account_id: self.account_id, - enabled: self.enabled, - concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, - weight: parse_account_weight(self.weight)?, - group_ids: validate_wire_group_ids(&self.group_ids)?, - }) - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdatedAccountData { - pub account_id: String, - pub config_revision: u64, -} - -impl From for UpdatedAccountData { - fn from(result: AccountUpdateResult) -> Self { - Self { - account_id: result.account_id.to_string(), - config_revision: result.config_revision.get(), - } - } -} - -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct AccountDeletionRequest { - pub provider: String, - pub account_ids: Vec, -} - -impl AccountDeletionRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - AccountProvider::parse(&self.provider)?; - if self.account_ids.is_empty() || self.account_ids.len() > MAX_ACCOUNT_DELETE_BATCH { - return Err(WireValidationError::new("accountIds")); - } - let mut unique = BTreeSet::new(); - for account_id in &self.account_ids { - require_account_id(account_id, "accountIds")?; - if !unique.insert(account_id.as_str()) { - return Err(WireValidationError::new("accountIds")); - } - } - Ok(()) - } - - fn into_command( - self, - context: gateway_admin::model::MutationContext, - ) -> Result<(AccountProvider, CredentialDeletion), WireValidationError> { - self.validate()?; - let provider = AccountProvider::parse(&self.provider)?; - Ok(( - provider, - CredentialDeletion { - context, - account_ids: self - .account_ids - .into_iter() - .map(ProviderAccountId::new) - .collect::, _>>() - .map_err(|_| WireValidationError::new("accountIds"))?, - }, - )) - } -} - -#[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct RotateAccountRequest { - pub provider: String, - pub account_id: String, - pub access_token: String, - pub refresh_token: Option, - pub id_token: Option, -} - -impl RotateAccountRequest { - pub fn validate(&self) -> Result<(), WireValidationError> { - if AccountProvider::parse(&self.provider)? != AccountProvider::OpenAi { - return Err(WireValidationError::new("provider")); - } - require_account_id(&self.account_id, "accountId")?; - validate_oauth_material( - &self.access_token, - self.refresh_token.as_deref(), - self.id_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), - ); - material.insert( - "id_token".to_owned(), - self.id_token.map_or(Value::Null, Value::String), - ); - Ok(RotateCredential { - mutation: CredentialMutation { - context, - account_id: ProviderAccountId::new(self.account_id) - .map_err(|_| WireValidationError::new("accountId"))?, - }, - provider_material: ProviderDocument::new(OpaqueProviderData::new(material)), - }) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountImportData { - pub imported_count: usize, - pub account_ids: Vec, -} - -impl AccountImportData { - pub fn from_result(result: CredentialImportResult) -> Self { - let account_ids = result - .credential_ids - .into_iter() - .map(|account_id| account_id.to_string()) - .collect::>(); - Self { - imported_count: account_ids.len(), - account_ids, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountAuthorizationData { - pub flow_id: String, - pub authorization_url: String, - pub expires_at: DateTime, -} - -impl From for AccountAuthorizationData { - fn from(started: AuthorizationStarted) -> Self { - Self { - flow_id: started.flow_id, - authorization_url: started.authorization_url, - expires_at: started.expires_at, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountMutationData { - pub account_id: String, -} - -impl From for AccountMutationData { - fn from(result: CredentialMutationResult) -> Self { - Self { - account_id: result.account_id.to_string(), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountDeletionData { - pub deleted_count: usize, - pub account_ids: Vec, -} - -impl From for AccountDeletionData { - fn from(result: CredentialDeletionResult) -> Self { - let account_ids = result - .account_ids - .into_iter() - .map(|account_id| account_id.to_string()) - .collect::>(); - Self { - deleted_count: account_ids.len(), - account_ids, - } - } -} - -fn require_account_id(value: &str, field: &'static str) -> Result<(), WireValidationError> { - if value.trim().is_empty() - || value.len() > 128 - || value.chars().any(char::is_control) - || !value.starts_with("acct_") - { - return Err(WireValidationError::new(field)); - } - Ok(()) -} - -fn parse_concurrency_limit( - value: Option, -) -> Result, WireValidationError> { - value - .map(|value| { - u32::try_from(value) - .ok() - .and_then(AccountConcurrencyLimit::new) - .ok_or_else(|| WireValidationError::new("concurrencyLimit")) - }) - .transpose() -} - -fn deserialize_required_nullable<'de, D>(deserializer: D) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - Option::::deserialize(deserializer) -} - -fn parse_account_weight(value: u64) -> Result { - u16::try_from(value) - .ok() - .and_then(AccountWeight::new) - .ok_or_else(|| WireValidationError::new("weight")) -} - -fn validate_wire_group_ids(values: &[String]) -> Result, WireValidationError> { - if values.len() > MAX_ACCOUNT_GROUP_BATCH - || values.iter().collect::>().len() != values.len() - { - return Err(WireValidationError::new("groupIds")); - } - values - .iter() - .cloned() - .map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds"))) - .collect() -} - -/// 构造统一账号管理路由。 -pub fn router() -> Router -where - S: AdminSessionState + Clone + Send + Sync + 'static, -{ - Router::new() - .route("/api/admin/accounts", get(list_accounts::)) - .route("/api/admin/accounts/detail", get(account_detail::)) - .route("/api/admin/accounts/export", get(export_accounts::)) - .route("/api/admin/accounts/import", post(import_accounts::)) - .route("/api/admin/accounts/refresh", post(refresh_account::)) - .route("/api/admin/accounts/recover", post(recover_account::)) - .route("/api/admin/accounts/rotate", post(rotate_account::)) - .route("/api/admin/accounts/update", post(update_account::)) - .route("/api/admin/accounts/delete", post(delete_accounts::)) - .route( - "/api/admin/accounts/batch-update", - post(batch_update_accounts::), - ) - .route("/api/admin/accounts/quota", get(account_quota::)) - .route( - "/api/admin/accounts/profile-statistics", - get(account_profile_statistics::), - ) - .route( - "/api/admin/accounts/profile-avatar", - get(account_profile_avatar::), - ) - .route( - "/api/admin/accounts/reset-credits", - get(account_reset_credits::).post(consume_account_reset_credit::), - ) - .route( - "/api/admin/accounts/quota/refresh", - post(refresh_account_quota::), - ) - .route("/api/admin/accounts/models", get(account_models::)) - .route( - "/api/admin/accounts/models/refresh", - post(refresh_account_models::), - ) - .route( - "/api/admin/accounts/connection-test", - get(test_account_connection::), - ) - .route( - "/api/admin/accounts/oauth/start", - post(start_account_authorization::), - ) - .route( - "/api/admin/accounts/oauth/complete", - post(complete_account_authorization::), - ) -} - -async fn batch_update_accounts( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let command = request.into_command().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .batch_update(&auth.context().mutation_context(), command) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(BatchUpdatedAccountsData::from(result)), - )) -} - -async fn list_accounts( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn account_detail( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn export_accounts( - auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let ids = query.into_ids().map_err(map_wire_error)?; - 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() - .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store")); - Ok(response) -} - -async fn import_accounts( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let (provider, command) = request - .into_command(auth.context().mutation_context()) - .map_err(map_wire_error)?; - let result = match provider { - AccountProvider::OpenAi => { - state - .admin_services() - .openai() - .import_document(command) - .await - } - AccountProvider::Xai => state.admin_services().xai().import_document(command).await, - } - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::CREATED, - AdminEnvelope::ok(AccountImportData::from_result(result)), - )) -} - -async fn start_account_authorization( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let (provider, command) = request - .into_command(auth.context().mutation_context()) - .map_err(map_wire_error)?; - let result = match provider { - AccountProvider::OpenAi => { - state - .admin_services() - .openai() - .start_authorization(command) - .await - } - AccountProvider::Xai => { - state - .admin_services() - .xai() - .start_authorization(command) - .await - } - } - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::CREATED, - AdminEnvelope::ok(AccountAuthorizationData::from(result)), - )) -} - -async fn complete_account_authorization( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let (provider, command) = request - .into_command(auth.context().mutation_context()) - .map_err(map_wire_error)?; - let result = match provider { - AccountProvider::OpenAi => { - state - .admin_services() - .openai() - .complete_authorization(command) - .await - } - AccountProvider::Xai => { - state - .admin_services() - .xai() - .complete_authorization(command) - .await - } - } - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::CREATED, - AdminEnvelope::ok(AccountMutationData::from(result)), - )) -} - -async fn rotate_account( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> 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(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountMutationData::from(result)), - )) -} - -async fn update_account( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let command = request.into_command().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .update(&auth.context().mutation_context(), command) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(UpdatedAccountData::from(result)), - )) -} - -async fn delete_accounts( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let (provider, command) = request - .into_command(auth.context().mutation_context()) - .map_err(map_wire_error)?; - let result = match provider { - AccountProvider::OpenAi => state.admin_services().openai().delete(command).await, - AccountProvider::Xai => state.admin_services().xai().delete(command).await, - } - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountDeletionData::from(result)), - )) -} - -async fn refresh_account( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let account_id = request.into_command().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .refresh(&auth.context().mutation_context(), account_id) - .await - .map_err(map_service_error)?; - let data = account_refresh_data(result, Utc::now()); - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) -} - -async fn recover_account( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let account_id = request.into_id().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .recover(&auth.context().mutation_context(), account_id) - .await - .map_err(map_service_error)?; - let data = account_refresh_data(result, Utc::now()); - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) -} - -async fn account_quota( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn account_profile_statistics( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let account_id = query.into_id().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .profile_statistics(&account_id) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountProfileStatisticsData::from(result)), - )) -} - -async fn account_profile_avatar( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let account_id = query.into_id().map_err(map_wire_error)?; - let avatar = state - .admin_services() - .accounts() - .profile_avatar(&account_id) - .await - .map_err(map_service_error)?; - Ok(profile_avatar_response(avatar)) -} - -/// 将 Provider 头像流投影为受保护的同源 HTTP 响应。 -#[must_use] -pub fn profile_avatar_response(avatar: ProviderProfileAvatar) -> Response { - let ProviderProfileAvatar { - content_type, - content_length, - etag, - body, - } = avatar; - let mut response = Response::new(Body::from_stream(body)); - *response.status_mut() = StatusCode::OK; - let headers = response.headers_mut(); - headers.insert( - header::CONTENT_TYPE, - content_type - .and_then(|value| HeaderValue::from_str(&value).ok()) - .unwrap_or_else(|| HeaderValue::from_static("application/octet-stream")), - ); - if let Some(value) = - content_length.and_then(|length| HeaderValue::from_str(&length.to_string()).ok()) - { - headers.insert(header::CONTENT_LENGTH, value); - } - if let Some(value) = etag.and_then(|value| HeaderValue::from_str(&value).ok()) { - headers.insert(header::ETAG, value); - } - headers.insert( - header::CACHE_CONTROL, - HeaderValue::from_static("private, max-age=3600"), - ); - headers.insert( - HeaderName::from_static("x-content-type-options"), - HeaderValue::from_static("nosniff"), - ); - headers.insert( - HeaderName::from_static("cross-origin-resource-policy"), - HeaderValue::from_static("same-origin"), - ); - headers.insert( - HeaderName::from_static("content-security-policy"), - HeaderValue::from_static("sandbox; default-src 'none'"), - ); - response -} - -async fn refresh_account_quota( - _auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn account_reset_credits( - auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let account_id = query.into_id().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .reset_credits(&auth.context().mutation_context(), account_id) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountResetCreditsData::from(result)), - )) -} - -async fn consume_account_reset_credit( - auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let command = request.into_command().map_err(map_wire_error)?; - let result = state - .admin_services() - .accounts() - .consume_reset_credit(&auth.context().mutation_context(), command) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountResetCreditResultData::from(result)), - )) -} - -async fn account_models( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn refresh_account_models( - _auth: AdminAuth, - State(state): State, - AdminJson(request): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - 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))) -} - -async fn test_account_connection( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result>>, AdminError> -where - S: AdminSessionState + Send + Sync, -{ - let (account_id, upstream_model) = query.into_command().map_err(map_wire_error)?; - let stream = state - .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 { - 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, - normal: result.summary.normal, - quota_exhausted: result.summary.quota_exhausted, - rate_limited: result.summary.rate_limited, - disabled: result.summary.disabled, - error: result.summary.error, - }, - } -} - -fn account_refresh_data(result: AccountRefreshResult, now: DateTime) -> AccountRefreshData { - AccountRefreshData { - 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, - projection, - usage, - quota, - } = item; - let status = projection.status.as_str().to_owned(); - let rate_limited_until = projection.rate_limited_until.map(DateTime::::from); - let expires_at = account.access_token_expires_at.as_ref().map(china_rfc3339); - 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(quota, rate_limited_until, now); - AccountView { - id: account.id.clone(), - name: account.name, - provider: account.provider_kind.to_string(), - groups: account - .groups - .into_iter() - .map(|group| AccountGroupRefView { - id: group.id.to_string(), - name: group.name, - color: group.color.as_str().to_owned(), - enabled: group.enabled, - }) - .collect(), - resource_ref: account.id, - email: account.email, - account_id: account.upstream_account_id, - user_id: account.upstream_user_id, - label: None, - plan_type: account.plan_type, - authentication_kind: account.authentication_kind, - has_refresh_token: account.has_refresh_token, - status, - error_reason: projection - .error_reason - .map(|reason| reason.as_str().to_owned()), - error_message: projection.error_message, - enabled: account.enabled, - concurrency_limit: account.concurrency_limit.map(|limit| limit.get()), - weight: account.weight.get(), - access_token_expires_at: expires_at, - access_token_expires_at_display: account - .access_token_expires_at - .as_ref() - .map(china_datetime), - 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( - mut quota: ProviderQuota, - rate_limited_until: Option>, - now: DateTime, -) -> (AccountQuotaView, Option) { - quota.apply_limit_reached_display(); - 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(quota_window_view).collect(); - let rate_limited_until = rate_limited_until.map(|until| china_datetime(&until)); - ( - AccountQuotaView { - refreshed_at_display, - limit_reached: quota.limit_reached, - rate_limited_until, - windows, - }, - refresh_token_expires_at, - ) -} - -fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView { - let ProviderQuotaWindow { - key, - group, - label, - limit_id, - limit_name, - role, - local_usage_attribution: _, - window_seconds, - used_percent, - reset_at, - limit_reached, - local_usage, - provider_data: _, - } = window; - let label_display = limit_name - .as_ref() - .map_or_else(|| label.clone(), |name| format!("{name} · {label}")); - AccountQuotaWindowView { - label_display, - window_label_display: label, - key, - group, - limit_id, - limit_name, - role: role.map(|value| value.as_str().to_owned()), - window_seconds, - used_percent, - used_percent_display: used_percent - .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), - limit_reached, - local_usage: local_usage.as_ref().map(quota_local_usage), - reset_at_display: reset_at.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), - } -} - -fn 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), - "imageInputTokens": usage.image_input_tokens.unwrap_or_default(), - "imageInputTokensDisplay": display_optional_tokens(usage.image_input_tokens), - "imageOutputTokens": usage.image_output_tokens.unwrap_or_default(), - "imageOutputTokensDisplay": display_optional_tokens(usage.image_output_tokens), - "imageRequestCount": usage.image_request_count, - "imageRequestFailedCount": usage.image_request_failed_count, - "totalTokens": total_tokens, - "totalTokensDisplay": format_compact_number(total_tokens), - "requestBuckets": usage.request_buckets.iter().map(|bucket| serde_json::json!({ - "bucketStart": bucket.bucket_start, - "requestCount": bucket.request_count, - })).collect::>(), - }) -} - -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), - image_input_tokens: usage.image_input_tokens, - image_input_tokens_display: display_optional_tokens(usage.image_input_tokens), - image_output_tokens: usage.image_output_tokens, - image_output_tokens_display: display_optional_tokens(usage.image_output_tokens), - image_request_count: Some(usage.image_request_count), - image_request_count_display: format_number(usage.image_request_count), - image_request_failed_count: Some(usage.image_request_failed_count), - image_request_failed_count_display: format_number(usage.image_request_failed_count), - 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), - image_input_tokens: usage.image_input_tokens, - image_input_tokens_display: display_optional_tokens(usage.image_input_tokens), - image_output_tokens: usage.image_output_tokens, - image_output_tokens_display: display_optional_tokens(usage.image_output_tokens), - image_request_count: usage.image_request_count, - image_request_count_display: format_number(usage.image_request_count), - image_request_failed_count: usage.image_request_failed_count, - image_request_failed_count_display: format_number(usage.image_request_failed_count), - 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_decimal_currency(cost.amount.as_str(), "USD"), - ), - 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_decimal_currency(cost.amount.as_str(), &cost.currency), - } -} - -fn display_optional_tokens(value: Option) -> String { - value.map_or_else(|| "—".to_owned(), format_compact_number) -} - -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(), - image_input_tokens: None, - image_input_tokens_display: "—".to_owned(), - image_output_tokens: None, - image_output_tokens_display: "—".to_owned(), - image_request_count: None, - image_request_count_display: "—".to_owned(), - image_request_failed_count: None, - image_request_failed_count_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: Value, - field: &'static str, -) -> Result { - match value { - Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))), - _ => Err(WireValidationError::new(field)), - } -} - -fn require_wire_id(value: &str, field: &'static str) -> Result<(), WireValidationError> { - require_text(value, MAX_ID_BYTES, field)?; - if value.starts_with("__") - || !value - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.' | b':')) - { - return Err(WireValidationError::new(field)); - } - Ok(()) -} - -fn require_text( - value: &str, - max_bytes: usize, - field: &'static str, -) -> Result<(), WireValidationError> { - if value.trim().is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) { - return Err(WireValidationError::new(field)); - } - Ok(()) -} - -fn validate_oauth_material( - access_token: &str, - refresh_token: Option<&str>, - id_token: Option<&str>, -) -> Result<(), WireValidationError> { - if access_token.len() > MAX_ACCESS_TOKEN_BYTES - || !valid_visible_ascii(access_token) - || !valid_compact_jwt_shape(access_token) - { - return Err(WireValidationError::new("accessToken")); - } - if refresh_token.is_some_and(|token| { - token.len() > MAX_REFRESH_TOKEN_BYTES - || !valid_visible_ascii(token) - || token == access_token - }) { - return Err(WireValidationError::new("refreshToken")); - } - if id_token.is_some_and(|token| { - token.len() > MAX_ID_TOKEN_BYTES - || !valid_visible_ascii(token) - || !valid_compact_jwt_shape(token) - }) { - return Err(WireValidationError::new("idToken")); - } - Ok(()) -} - -fn valid_compact_jwt_shape(value: &str) -> bool { - let mut segments = value.split('.'); - matches!( - (segments.next(), segments.next(), segments.next(), segments.next()), - (Some(header), Some(payload), Some(signature), None) - if !header.is_empty() && !payload.is_empty() && !signature.is_empty() - ) -} - -fn valid_visible_ascii(value: &str) -> bool { - !value.is_empty() && value.bytes().all(|byte| (0x21..=0x7e).contains(&byte)) -} - -fn provider_document_value(document: ProviderDocument) -> Value { - Value::Object(document.into_provider_data().into_inner()) -} - -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!("{} 字段不合法", error.field())) -} - -fn map_service_error(error: AdminServiceError) -> AdminError { - super::wire::map_admin_service_error(error) -} diff --git a/backend/crates/gateway-api/src/admin/accounts/credentials.rs b/backend/crates/gateway-api/src/admin/accounts/credentials.rs new file mode 100644 index 000000000..0895f6279 --- /dev/null +++ b/backend/crates/gateway-api/src/admin/accounts/credentials.rs @@ -0,0 +1,488 @@ +//! Provider-owned 凭据请求、命令转换与敏感材料校验。 + +use super::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum AccountProvider { + OpenAi, + Xai, +} + +impl AccountProvider { + fn parse(value: &str) -> Result { + match value.trim() { + "openai" => Ok(Self::OpenAi), + "xai" => Ok(Self::Xai), + _ => Err(WireValidationError::new("provider")), + } + } +} + +/// Provider-owned 账号导入请求；公共 API 不解释 `data` 内部字段。 +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountImportRequest { + pub provider: String, + pub data: Value, +} + +impl AccountImportRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + AccountProvider::parse(&self.provider)?; + if !self.data.is_object() + || serde_json::to_vec(&self.data) + .map_or(true, |encoded| encoded.len() > MAX_IMPORT_DATA_BYTES) + { + return Err(WireValidationError::new("data")); + } + Ok(()) + } + + pub(super) fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result<(AccountProvider, ImportCredentials), WireValidationError> { + self.validate()?; + let provider = AccountProvider::parse(&self.provider)?; + Ok(( + provider, + ImportCredentials { + context, + document: provider_document(self.data, "data")?, + }, + )) + } +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct StartAccountAuthorizationRequest { + pub provider: String, + pub name: String, + pub account_id: Option, +} + +impl StartAccountAuthorizationRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + AccountProvider::parse(&self.provider)?; + require_text(&self.name, MAX_NAME_BYTES, "name")?; + if let Some(account_id) = self.account_id.as_deref() { + require_account_id(account_id, "accountId")?; + } + Ok(()) + } + + pub(super) fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result<(AccountProvider, StartAuthorization), WireValidationError> { + self.validate()?; + let provider = AccountProvider::parse(&self.provider)?; + let reauthorization = self + .account_id + .map(ProviderAccountId::new) + .transpose() + .map_err(|_| WireValidationError::new("accountId"))?; + Ok(( + provider, + StartAuthorization { + context, + name: self.name, + reauthorization, + }, + )) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct CompleteAccountAuthorizationRequest { + pub provider: String, + pub flow_id: String, + pub callback_url: String, +} + +impl CompleteAccountAuthorizationRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + let provider = AccountProvider::parse(&self.provider)?; + match provider { + AccountProvider::OpenAi => { + if !URL_SAFE_NO_PAD + .decode(&self.flow_id) + .is_ok_and(|decoded| decoded.len() == 32) + { + return Err(WireValidationError::new("flowId")); + } + } + AccountProvider::Xai => require_wire_id(&self.flow_id, "flowId")?, + } + require_text(&self.callback_url, MAX_CALLBACK_URL_BYTES, "callbackUrl") + } + + pub(super) fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result<(AccountProvider, CompleteAuthorization), WireValidationError> { + self.validate()?; + let provider = AccountProvider::parse(&self.provider)?; + Ok(( + provider, + CompleteAuthorization { + context, + flow_id: self.flow_id, + callback_url: self.callback_url, + }, + )) + } +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct UpdateAccountRequest { + pub account_id: String, + pub enabled: bool, + #[serde(deserialize_with = "deserialize_required_nullable")] + pub concurrency_limit: Option, + pub weight: u64, + pub group_ids: Vec, +} + +impl UpdateAccountRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId")?; + parse_concurrency_limit(self.concurrency_limit)?; + parse_account_weight(self.weight)?; + validate_wire_group_ids(&self.group_ids)?; + Ok(()) + } + + pub(super) fn into_command(self) -> Result { + self.validate()?; + Ok(UpdateAccount { + account_id: self.account_id, + enabled: self.enabled, + concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, + weight: parse_account_weight(self.weight)?, + group_ids: validate_wire_group_ids(&self.group_ids)?, + }) + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdatedAccountData { + pub account_id: String, + pub config_revision: u64, +} + +impl From for UpdatedAccountData { + fn from(result: AccountUpdateResult) -> Self { + Self { + account_id: result.account_id.to_string(), + config_revision: result.config_revision.get(), + } + } +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountDeletionRequest { + pub provider: String, + pub account_ids: Vec, +} + +impl AccountDeletionRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + AccountProvider::parse(&self.provider)?; + if self.account_ids.is_empty() || self.account_ids.len() > MAX_ACCOUNT_DELETE_BATCH { + return Err(WireValidationError::new("accountIds")); + } + let mut unique = BTreeSet::new(); + for account_id in &self.account_ids { + require_account_id(account_id, "accountIds")?; + if !unique.insert(account_id.as_str()) { + return Err(WireValidationError::new("accountIds")); + } + } + Ok(()) + } + + pub(super) fn into_command( + self, + context: gateway_admin::model::MutationContext, + ) -> Result<(AccountProvider, CredentialDeletion), WireValidationError> { + self.validate()?; + let provider = AccountProvider::parse(&self.provider)?; + Ok(( + provider, + CredentialDeletion { + context, + account_ids: self + .account_ids + .into_iter() + .map(ProviderAccountId::new) + .collect::, _>>() + .map_err(|_| WireValidationError::new("accountIds"))?, + }, + )) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct RotateAccountRequest { + pub provider: String, + pub account_id: String, + pub access_token: String, + pub refresh_token: Option, + pub id_token: Option, +} + +impl RotateAccountRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + if AccountProvider::parse(&self.provider)? != AccountProvider::OpenAi { + return Err(WireValidationError::new("provider")); + } + require_account_id(&self.account_id, "accountId")?; + validate_oauth_material( + &self.access_token, + self.refresh_token.as_deref(), + self.id_token.as_deref(), + ) + } + + pub(super) 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), + ); + material.insert( + "id_token".to_owned(), + self.id_token.map_or(Value::Null, Value::String), + ); + Ok(RotateCredential { + mutation: CredentialMutation { + context, + account_id: ProviderAccountId::new(self.account_id) + .map_err(|_| WireValidationError::new("accountId"))?, + }, + provider_material: ProviderDocument::new(OpaqueProviderData::new(material)), + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountImportData { + pub imported_count: usize, + pub account_ids: Vec, +} + +impl AccountImportData { + pub fn from_result(result: CredentialImportResult) -> Self { + let account_ids = result + .credential_ids + .into_iter() + .map(|account_id| account_id.to_string()) + .collect::>(); + Self { + imported_count: account_ids.len(), + account_ids, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountAuthorizationData { + pub flow_id: String, + pub authorization_url: String, + pub expires_at: DateTime, +} + +impl From for AccountAuthorizationData { + fn from(started: AuthorizationStarted) -> Self { + Self { + flow_id: started.flow_id, + authorization_url: started.authorization_url, + expires_at: started.expires_at, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountMutationData { + pub account_id: String, +} + +impl From for AccountMutationData { + fn from(result: CredentialMutationResult) -> Self { + Self { + account_id: result.account_id.to_string(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountDeletionData { + pub deleted_count: usize, + pub account_ids: Vec, +} + +impl From for AccountDeletionData { + fn from(result: CredentialDeletionResult) -> Self { + let account_ids = result + .account_ids + .into_iter() + .map(|account_id| account_id.to_string()) + .collect::>(); + Self { + deleted_count: account_ids.len(), + account_ids, + } + } +} + +pub(super) fn require_account_id( + value: &str, + field: &'static str, +) -> Result<(), WireValidationError> { + if value.trim().is_empty() + || value.len() > 128 + || value.chars().any(char::is_control) + || !value.starts_with("acct_") + { + return Err(WireValidationError::new(field)); + } + Ok(()) +} + +pub(super) fn parse_concurrency_limit( + value: Option, +) -> Result, WireValidationError> { + value + .map(|value| { + u32::try_from(value) + .ok() + .and_then(AccountConcurrencyLimit::new) + .ok_or_else(|| WireValidationError::new("concurrencyLimit")) + }) + .transpose() +} + +pub(super) fn deserialize_required_nullable<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + Option::::deserialize(deserializer) +} + +pub(super) fn parse_account_weight(value: u64) -> Result { + u16::try_from(value) + .ok() + .and_then(AccountWeight::new) + .ok_or_else(|| WireValidationError::new("weight")) +} + +pub(super) fn validate_wire_group_ids( + values: &[String], +) -> Result, WireValidationError> { + if values.len() > MAX_ACCOUNT_GROUP_BATCH + || values.iter().collect::>().len() != values.len() + { + return Err(WireValidationError::new("groupIds")); + } + values + .iter() + .cloned() + .map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds"))) + .collect() +} + +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 require_wire_id(value: &str, field: &'static str) -> Result<(), WireValidationError> { + require_text(value, MAX_ID_BYTES, field)?; + if value.starts_with("__") + || !value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.' | b':')) + { + return Err(WireValidationError::new(field)); + } + Ok(()) +} + +fn require_text( + value: &str, + max_bytes: usize, + field: &'static str, +) -> Result<(), WireValidationError> { + if value.trim().is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) { + return Err(WireValidationError::new(field)); + } + Ok(()) +} + +fn validate_oauth_material( + access_token: &str, + refresh_token: Option<&str>, + id_token: Option<&str>, +) -> Result<(), WireValidationError> { + if access_token.len() > MAX_ACCESS_TOKEN_BYTES + || !valid_visible_ascii(access_token) + || !valid_compact_jwt_shape(access_token) + { + return Err(WireValidationError::new("accessToken")); + } + if refresh_token.is_some_and(|token| { + token.len() > MAX_REFRESH_TOKEN_BYTES + || !valid_visible_ascii(token) + || token == access_token + }) { + return Err(WireValidationError::new("refreshToken")); + } + if id_token.is_some_and(|token| { + token.len() > MAX_ID_TOKEN_BYTES + || !valid_visible_ascii(token) + || !valid_compact_jwt_shape(token) + }) { + return Err(WireValidationError::new("idToken")); + } + Ok(()) +} + +fn valid_compact_jwt_shape(value: &str) -> bool { + let mut segments = value.split('.'); + matches!( + (segments.next(), segments.next(), segments.next(), segments.next()), + (Some(header), Some(payload), Some(signature), None) + if !header.is_empty() && !payload.is_empty() && !signature.is_empty() + ) +} + +fn valid_visible_ascii(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|byte| (0x21..=0x7e).contains(&byte)) +} + +pub(super) fn provider_document_value(document: ProviderDocument) -> Value { + Value::Object(document.into_provider_data().into_inner()) +} diff --git a/backend/crates/gateway-api/src/admin/accounts/handlers.rs b/backend/crates/gateway-api/src/admin/accounts/handlers.rs new file mode 100644 index 000000000..118459ffb --- /dev/null +++ b/backend/crates/gateway-api/src/admin/accounts/handlers.rs @@ -0,0 +1,573 @@ +//! 账号管理路由与 HTTP handler 编排。 + +use super::*; + +/// 构造统一账号管理路由。 +pub fn router() -> Router +where + S: AdminSessionState + Clone + Send + Sync + 'static, +{ + Router::new() + .route("/api/admin/accounts", get(list_accounts::)) + .route("/api/admin/accounts/detail", get(account_detail::)) + .route("/api/admin/accounts/export", get(export_accounts::)) + .route("/api/admin/accounts/import", post(import_accounts::)) + .route("/api/admin/accounts/refresh", post(refresh_account::)) + .route("/api/admin/accounts/recover", post(recover_account::)) + .route("/api/admin/accounts/rotate", post(rotate_account::)) + .route("/api/admin/accounts/update", post(update_account::)) + .route("/api/admin/accounts/delete", post(delete_accounts::)) + .route( + "/api/admin/accounts/batch-update", + post(batch_update_accounts::), + ) + .route("/api/admin/accounts/quota", get(account_quota::)) + .route( + "/api/admin/accounts/profile-statistics", + get(account_profile_statistics::), + ) + .route( + "/api/admin/accounts/profile-avatar", + get(account_profile_avatar::), + ) + .route( + "/api/admin/accounts/reset-credits", + get(account_reset_credits::).post(consume_account_reset_credit::), + ) + .route( + "/api/admin/accounts/quota/refresh", + post(refresh_account_quota::), + ) + .route("/api/admin/accounts/models", get(account_models::)) + .route( + "/api/admin/accounts/models/refresh", + post(refresh_account_models::), + ) + .route( + "/api/admin/accounts/connection-test", + get(test_account_connection::), + ) + .route( + "/api/admin/accounts/oauth/start", + post(start_account_authorization::), + ) + .route( + "/api/admin/accounts/oauth/complete", + post(complete_account_authorization::), + ) +} + +async fn batch_update_accounts( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .batch_update(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(BatchUpdatedAccountsData::from(result)), + )) +} + +async fn list_accounts( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn account_detail( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn export_accounts( + auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let ids = query.into_ids().map_err(map_wire_error)?; + 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() + .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store")); + Ok(response) +} + +async fn import_accounts( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (provider, command) = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = match provider { + AccountProvider::OpenAi => { + state + .admin_services() + .openai() + .import_document(command) + .await + } + AccountProvider::Xai => state.admin_services().xai().import_document(command).await, + } + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::CREATED, + AdminEnvelope::ok(AccountImportData::from_result(result)), + )) +} + +async fn start_account_authorization( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (provider, command) = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = match provider { + AccountProvider::OpenAi => { + state + .admin_services() + .openai() + .start_authorization(command) + .await + } + AccountProvider::Xai => { + state + .admin_services() + .xai() + .start_authorization(command) + .await + } + } + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::CREATED, + AdminEnvelope::ok(AccountAuthorizationData::from(result)), + )) +} + +async fn complete_account_authorization( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (provider, command) = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = match provider { + AccountProvider::OpenAi => { + state + .admin_services() + .openai() + .complete_authorization(command) + .await + } + AccountProvider::Xai => { + state + .admin_services() + .xai() + .complete_authorization(command) + .await + } + } + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::CREATED, + AdminEnvelope::ok(AccountMutationData::from(result)), + )) +} + +async fn rotate_account( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> 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(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountMutationData::from(result)), + )) +} + +async fn update_account( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .update(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(UpdatedAccountData::from(result)), + )) +} + +async fn delete_accounts( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let (provider, command) = request + .into_command(auth.context().mutation_context()) + .map_err(map_wire_error)?; + let result = match provider { + AccountProvider::OpenAi => state.admin_services().openai().delete(command).await, + AccountProvider::Xai => state.admin_services().xai().delete(command).await, + } + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountDeletionData::from(result)), + )) +} + +async fn refresh_account( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let account_id = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .refresh(&auth.context().mutation_context(), account_id) + .await + .map_err(map_service_error)?; + let data = account_refresh_data(result, Utc::now()); + Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) +} + +async fn recover_account( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let account_id = request.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .recover(&auth.context().mutation_context(), account_id) + .await + .map_err(map_service_error)?; + let data = account_refresh_data(result, Utc::now()); + Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))) +} + +async fn account_quota( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn account_profile_statistics( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let account_id = query.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .profile_statistics(&account_id) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountProfileStatisticsData::from(result)), + )) +} + +async fn account_profile_avatar( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let account_id = query.into_id().map_err(map_wire_error)?; + let avatar = state + .admin_services() + .accounts() + .profile_avatar(&account_id) + .await + .map_err(map_service_error)?; + Ok(profile_avatar_response(avatar)) +} + +/// 将 Provider 头像流投影为受保护的同源 HTTP 响应。 +#[must_use] +pub fn profile_avatar_response(avatar: ProviderProfileAvatar) -> Response { + let ProviderProfileAvatar { + content_type, + content_length, + etag, + body, + } = avatar; + let mut response = Response::new(Body::from_stream(body)); + *response.status_mut() = StatusCode::OK; + let headers = response.headers_mut(); + headers.insert( + header::CONTENT_TYPE, + content_type + .and_then(|value| HeaderValue::from_str(&value).ok()) + .unwrap_or_else(|| HeaderValue::from_static("application/octet-stream")), + ); + if let Some(value) = + content_length.and_then(|length| HeaderValue::from_str(&length.to_string()).ok()) + { + headers.insert(header::CONTENT_LENGTH, value); + } + if let Some(value) = etag.and_then(|value| HeaderValue::from_str(&value).ok()) { + headers.insert(header::ETAG, value); + } + headers.insert( + header::CACHE_CONTROL, + HeaderValue::from_static("private, max-age=3600"), + ); + headers.insert( + HeaderName::from_static("x-content-type-options"), + HeaderValue::from_static("nosniff"), + ); + headers.insert( + HeaderName::from_static("cross-origin-resource-policy"), + HeaderValue::from_static("same-origin"), + ); + headers.insert( + HeaderName::from_static("content-security-policy"), + HeaderValue::from_static("sandbox; default-src 'none'"), + ); + response +} + +async fn refresh_account_quota( + _auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn account_reset_credits( + auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let account_id = query.into_id().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .reset_credits(&auth.context().mutation_context(), account_id) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountResetCreditsData::from(result)), + )) +} + +async fn consume_account_reset_credit( + auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let command = request.into_command().map_err(map_wire_error)?; + let result = state + .admin_services() + .accounts() + .consume_reset_credit(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountResetCreditResultData::from(result)), + )) +} + +async fn account_models( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn refresh_account_models( + _auth: AdminAuth, + State(state): State, + AdminJson(request): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + 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))) +} + +async fn test_account_connection( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result>>, AdminError> +where + S: AdminSessionState + Send + Sync, +{ + let (account_id, upstream_model) = query.into_command().map_err(map_wire_error)?; + let stream = state + .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())) +} diff --git a/backend/crates/gateway-api/src/admin/accounts/mod.rs b/backend/crates/gateway-api/src/admin/accounts/mod.rs new file mode 100644 index 000000000..b2c24f31a --- /dev/null +++ b/backend/crates/gateway-api/src/admin/accounts/mod.rs @@ -0,0 +1,80 @@ +//! 管理端账号目录的 HTTP 合同、凭据动作、路由处理与安全响应投影。 + +use std::{collections::BTreeSet, convert::Infallible, fmt}; + +use axum::{ + Router, + body::Body, + extract::State, + http::{HeaderName, HeaderValue, StatusCode, header}, + response::{ + IntoResponse, Response, + sse::{Event, KeepAlive, Sse}, + }, + routing::{get, post}, +}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, FixedOffset, Utc}; +use futures::{Stream, StreamExt as _}; +use gateway_admin::model::{ + AdminError as AdminServiceError, PageSize, + accounts::{ + AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent, + AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort, + AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage, + AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount, + }, + provider_credentials::{ + AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult, + AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit, + CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation, + CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels, + ProviderProfileActivityInsights, ProviderProfileAvatar, ProviderProfileDailyUsage, + ProviderProfileInvocation, ProviderProfileStatistics, ProviderProfileStatisticsSummary, + ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult, + ProviderResetCredits, RotateCredential, StartAuthorization, + }, +}; +use gateway_core::{ + account::{OpaqueProviderData, ProviderAccountId}, + routing::{AccountGroupId, ProviderKind, UpstreamModelId}, +}; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value}; +use uuid::{Uuid, Version}; + +use super::presenter::{format_compact_number, format_decimal_currency, format_number}; +use super::{ + AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse, AdminSessionState, + PageMeta, WireValidationError, +}; + +const DEFAULT_PAGE_SIZE: u32 = 50; +const MAX_PAGE_SIZE: u32 = 200; +const MAX_SEARCH_BYTES: usize = 256; +const MAX_ID_BYTES: usize = 256; +const MAX_NAME_BYTES: usize = 512; +const MAX_IMPORT_DATA_BYTES: usize = 64 * 1024 * 1024; +const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024; +const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024; +const MAX_ID_TOKEN_BYTES: usize = 16 * 1024; +const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024; +const MAX_ACCOUNT_DELETE_BATCH: usize = 200; +const MAX_ACCOUNT_GROUP_BATCH: usize = 1000; +const MAX_AVATAR_VERSION_BYTES: usize = 32; + +mod credentials; +mod handlers; +mod presenter; +mod wire; + +pub use credentials::*; +pub use handlers::{profile_avatar_response, router}; +pub use wire::*; + +use credentials::{ + AccountProvider, deserialize_required_nullable, parse_account_weight, parse_concurrency_limit, + provider_document_value, require_account_id, validate_wire_group_ids, +}; +use presenter::*; +use wire::BatchUpdatedAccountsData; diff --git a/backend/crates/gateway-api/src/admin/accounts/presenter.rs b/backend/crates/gateway-api/src/admin/accounts/presenter.rs new file mode 100644 index 000000000..72fc12fa1 --- /dev/null +++ b/backend/crates/gateway-api/src/admin/accounts/presenter.rs @@ -0,0 +1,413 @@ +//! Admin 领域结果到安全 HTTP wire 的展示投影。 + +use super::*; + +pub(super) 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 { + 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, + normal: result.summary.normal, + quota_exhausted: result.summary.quota_exhausted, + rate_limited: result.summary.rate_limited, + disabled: result.summary.disabled, + error: result.summary.error, + }, + } +} + +pub(super) fn account_refresh_data( + result: AccountRefreshResult, + now: DateTime, +) -> AccountRefreshData { + AccountRefreshData { + account: account_view(result.account, now), + } +} + +pub(super) 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(), + } +} + +pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime) -> AccountView { + let AccountDirectoryItem { + account, + projection, + usage, + quota, + } = item; + let status = projection.status.as_str().to_owned(); + let rate_limited_until = projection.rate_limited_until.map(DateTime::::from); + let expires_at = account.access_token_expires_at.as_ref().map(china_rfc3339); + 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(quota, rate_limited_until, now); + AccountView { + id: account.id.clone(), + name: account.name, + provider: account.provider_kind.to_string(), + groups: account + .groups + .into_iter() + .map(|group| AccountGroupRefView { + id: group.id.to_string(), + name: group.name, + color: group.color.as_str().to_owned(), + enabled: group.enabled, + }) + .collect(), + resource_ref: account.id, + email: account.email, + account_id: account.upstream_account_id, + user_id: account.upstream_user_id, + label: None, + plan_type: account.plan_type, + authentication_kind: account.authentication_kind, + has_refresh_token: account.has_refresh_token, + status, + error_reason: projection + .error_reason + .map(|reason| reason.as_str().to_owned()), + error_message: projection.error_message, + enabled: account.enabled, + concurrency_limit: account.concurrency_limit.map(|limit| limit.get()), + weight: account.weight.get(), + access_token_expires_at: expires_at, + access_token_expires_at_display: account + .access_token_expires_at + .as_ref() + .map(china_datetime), + 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), + } +} + +pub(super) fn account_quota_view( + mut quota: ProviderQuota, + rate_limited_until: Option>, + now: DateTime, +) -> (AccountQuotaView, Option) { + quota.apply_limit_reached_display(); + 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(quota_window_view).collect(); + let rate_limited_until = rate_limited_until.map(|until| china_datetime(&until)); + ( + AccountQuotaView { + refreshed_at_display, + limit_reached: quota.limit_reached, + rate_limited_until, + windows, + }, + refresh_token_expires_at, + ) +} + +pub(super) fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView { + let ProviderQuotaWindow { + key, + group, + label, + limit_id, + limit_name, + role, + local_usage_attribution: _, + window_seconds, + used_percent, + reset_at, + limit_reached, + local_usage, + provider_data: _, + } = window; + let label_display = limit_name + .as_ref() + .map_or_else(|| label.clone(), |name| format!("{name} · {label}")); + AccountQuotaWindowView { + label_display, + window_label_display: label, + key, + group, + limit_id, + limit_name, + role: role.map(|value| value.as_str().to_owned()), + window_seconds, + used_percent, + used_percent_display: used_percent + .map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")), + limit_reached, + local_usage: local_usage.as_ref().map(quota_local_usage), + reset_at_display: reset_at.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)), + } +} + +pub(super) fn 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), + "imageInputTokens": usage.image_input_tokens.unwrap_or_default(), + "imageInputTokensDisplay": display_optional_tokens(usage.image_input_tokens), + "imageOutputTokens": usage.image_output_tokens.unwrap_or_default(), + "imageOutputTokensDisplay": display_optional_tokens(usage.image_output_tokens), + "imageRequestCount": usage.image_request_count, + "imageRequestFailedCount": usage.image_request_failed_count, + "totalTokens": total_tokens, + "totalTokensDisplay": format_compact_number(total_tokens), + "requestBuckets": usage.request_buckets.iter().map(|bucket| serde_json::json!({ + "bucketStart": bucket.bucket_start, + "requestCount": bucket.request_count, + })).collect::>(), + }) +} + +pub(super) 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), + image_input_tokens: usage.image_input_tokens, + image_input_tokens_display: display_optional_tokens(usage.image_input_tokens), + image_output_tokens: usage.image_output_tokens, + image_output_tokens_display: display_optional_tokens(usage.image_output_tokens), + image_request_count: Some(usage.image_request_count), + image_request_count_display: format_number(usage.image_request_count), + image_request_failed_count: Some(usage.image_request_failed_count), + image_request_failed_count_display: format_number(usage.image_request_failed_count), + 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(), + } +} + +pub(super) 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), + image_input_tokens: usage.image_input_tokens, + image_input_tokens_display: display_optional_tokens(usage.image_input_tokens), + image_output_tokens: usage.image_output_tokens, + image_output_tokens_display: display_optional_tokens(usage.image_output_tokens), + image_request_count: usage.image_request_count, + image_request_count_display: format_number(usage.image_request_count), + image_request_failed_count: usage.image_request_failed_count, + image_request_failed_count_display: format_number(usage.image_request_failed_count), + 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_decimal_currency(cost.amount.as_str(), "USD"), + ), + 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), + } +} + +pub(super) 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_decimal_currency(cost.amount.as_str(), &cost.currency), + } +} + +pub(super) fn display_optional_tokens(value: Option) -> String { + value.map_or_else(|| "—".to_owned(), format_compact_number) +} + +pub(super) 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(), + image_input_tokens: None, + image_input_tokens_display: "—".to_owned(), + image_output_tokens: None, + image_output_tokens_display: "—".to_owned(), + image_request_count: None, + image_request_count_display: "—".to_owned(), + image_request_failed_count: None, + image_request_failed_count_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(), + } +} + +pub(super) 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()) +} + +pub(super) fn china_offset() -> FixedOffset { + FixedOffset::east_opt(8 * 60 * 60).expect("UTC+8 is a valid fixed offset") +} + +pub(super) fn china_rfc3339(value: &DateTime) -> String { + value.with_timezone(&china_offset()).to_rfc3339() +} + +pub(super) fn china_datetime(value: &DateTime) -> String { + value + .with_timezone(&china_offset()) + .format("%Y-%m-%d %H:%M:%S") + .to_string() +} + +pub(super) fn map_wire_error(error: WireValidationError) -> AdminError { + AdminError::bad_request(format!("{} 字段不合法", error.field())) +} + +pub(super) fn map_service_error(error: AdminServiceError) -> AdminError { + super::super::wire::map_admin_service_error(error) +} diff --git a/backend/crates/gateway-api/src/admin/accounts/wire.rs b/backend/crates/gateway-api/src/admin/accounts/wire.rs new file mode 100644 index 000000000..c8227147c --- /dev/null +++ b/backend/crates/gateway-api/src/admin/accounts/wire.rs @@ -0,0 +1,838 @@ +//! 账号管理请求、响应与查询 wire contract。 + +use super::*; + +/// 账号列表查询参数。 +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ListQuery { + pub page: Option, + pub page_size: Option, + pub provider: Option, + pub group_id: Option, + pub search: Option, + pub status: Option, + pub sort_by: Option, + pub sort_direction: Option, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct BatchUpdateAccountsRequest { + pub account_ids: Vec, + pub enabled: bool, + #[serde(deserialize_with = "deserialize_required_nullable")] + pub concurrency_limit: Option, + pub weight: u64, + pub group_ids: Vec, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub(super) struct BatchUpdatedAccountsData { + account_ids: Vec, + config_revision: u64, +} + +impl From for BatchUpdatedAccountsData { + fn from(result: AccountsUpdateResult) -> Self { + Self { + account_ids: result + .account_ids + .into_iter() + .map(|id| id.to_string()) + .collect(), + config_revision: result.config_revision.get(), + } + } +} + +impl BatchUpdateAccountsRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + if self.account_ids.is_empty() + || self.account_ids.len() > MAX_ACCOUNT_GROUP_BATCH + || self + .account_ids + .iter() + .any(|id| require_account_id(id, "accountIds").is_err()) + || self.account_ids.iter().collect::>().len() != self.account_ids.len() + { + return Err(WireValidationError::new("accountIds")); + } + validate_wire_group_ids(&self.group_ids)?; + parse_concurrency_limit(self.concurrency_limit)?; + parse_account_weight(self.weight)?; + Ok(()) + } + + pub(super) fn into_command(self) -> Result { + self.validate()?; + Ok(BatchUpdateAccounts { + account_ids: self.account_ids, + enabled: self.enabled, + concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, + weight: parse_account_weight(self.weight)?, + group_ids: validate_wire_group_ids(&self.group_ids)?, + }) + } +} + +impl ListQuery { + /// 解析并校验全部 wire 字段，生成 Admin 查询命令。 + pub fn validate(self) -> Result { + let page = self.page.unwrap_or(1); + if page == 0 { + return Err(WireValidationError::new("page")); + } + let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE); + if page_size == 0 || page_size > MAX_PAGE_SIZE { + return Err(WireValidationError::new("pageSize")); + } + let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?; + let group_filter = match self.group_id.as_deref().map(str::trim) { + None | Some("") => None, + Some("ungrouped") => Some(AccountGroupFilter::Ungrouped), + Some(value) => Some(AccountGroupFilter::Group( + AccountGroupId::new(value.to_owned()) + .map_err(|_| WireValidationError::new("groupId"))?, + )), + }; + let search = self + .search + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()); + if search.as_deref().is_some_and(|value| { + value.len() > MAX_SEARCH_BYTES || value.chars().any(char::is_control) + }) { + return Err(WireValidationError::new("search")); + } + let status = match self.status.as_deref().map(str::trim) { + None | Some("") => None, + Some(value) => Some( + DomainAccountStatus::parse(&value.to_ascii_lowercase()) + .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: 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(AccountListQuery { + page, + page_size: PageSize::new( + u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?, + ) + .map_err(|_| WireValidationError::new("pageSize"))?, + provider_kind, + group_filter, + search, + status, + sort, + }) + } +} + +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")) +} + +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, + } +} + +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, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountPageData { + pub items: Vec, + pub page: PageMeta, + pub summary: AccountSummaryView, +} + +/// 账号概览计数。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountSummaryView { + pub total: u64, + pub normal: u64, + pub quota_exhausted: u64, + pub rate_limited: u64, + pub disabled: u64, + pub error: u64, +} + +/// 一条安全账号视图。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountView { + pub id: String, + pub name: String, + pub provider: String, + pub groups: Vec, + pub resource_ref: String, + pub email: Option, + pub account_id: Option, + pub user_id: Option, + pub label: Option, + pub plan_type: Option, + pub authentication_kind: String, + pub has_refresh_token: bool, + pub status: String, + /// `status == "error"` 时的具体原因；其余状态为 `null`。 + pub error_reason: Option, + /// 最近一次失败的上游错误描述；仅错误状态存在。 + pub error_message: Option, + pub enabled: bool, + pub concurrency_limit: Option, + pub weight: u16, + pub access_token_expires_at: Option, + pub access_token_expires_at_display: Option, + pub refresh_token_expires_at: Option, + pub next_refresh_at: Option, + pub added_at: String, + pub added_at_display: String, + pub updated_at: String, + pub updated_at_display: String, + pub quota: AccountQuotaView, + pub usage: AccountUsageView, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountGroupRefView { + pub id: String, + pub name: String, + pub color: String, + pub enabled: bool, +} + +/// Provider quota 安全视图。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountQuotaView { + pub refreshed_at_display: String, + pub limit_reached: bool, + /// 429 临时限流（Redis 冷却）到期时间展示；非限流中为 `null`。 + pub rate_limited_until: Option, + pub windows: Vec, +} + +/// 一个 quota 时间窗口。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountQuotaWindowView { + pub key: String, + pub group: String, + pub limit_id: Option, + pub limit_name: Option, + pub role: Option, + pub window_seconds: Option, + pub label_display: String, + pub window_label_display: String, + pub used_percent: Option, + pub used_percent_display: String, + pub limit_reached: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub local_usage: Option, + pub reset_at_display: String, +} + +/// 账号观测用量。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountUsageView { + pub request_count: Option, + pub request_count_display: String, + pub input_tokens: Option, + pub input_tokens_display: String, + pub output_tokens: Option, + pub output_tokens_display: String, + pub cached_tokens: Option, + pub cached_tokens_display: String, + pub image_input_tokens: Option, + pub image_input_tokens_display: String, + pub image_output_tokens: Option, + pub image_output_tokens_display: String, + pub image_request_count: Option, + pub image_request_count_display: String, + pub image_request_failed_count: Option, + pub image_request_failed_count_display: String, + pub total_tokens: Option, + pub total_tokens_display: String, + pub created_tokens: Option, + pub created_tokens_display: String, + pub read_tokens: Option, + pub read_tokens_display: String, + pub last_used_at: Option, + pub last_used_at_display: String, + pub cost_estimate_status: String, + pub known_cost_count: Option, + pub partial_cost_count: Option, + pub unknown_cost_count: Option, + pub costs: Vec, + pub models: Vec, +} + +/// 凭据在单个上游模型上的观测用量。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ModelUsageView { + pub model: String, + pub request_count: u64, + pub request_count_display: String, + pub success_rate: Option, + pub success_rate_display: String, + pub input_tokens: Option, + pub input_tokens_display: String, + pub output_tokens: Option, + pub output_tokens_display: String, + pub cached_tokens: Option, + pub cached_tokens_display: String, + pub image_input_tokens: Option, + pub image_input_tokens_display: String, + pub image_output_tokens: Option, + pub image_output_tokens_display: String, + pub image_request_count: u64, + pub image_request_count_display: String, + pub image_request_failed_count: u64, + pub image_request_failed_count_display: String, + pub total_tokens: Option, + pub total_tokens_display: String, + pub billing_amount_usd: Option, + pub billing_amount_usd_display: String, + pub cost_estimate_status: String, + pub known_cost_count: u64, + pub partial_cost_count: u64, + pub unknown_cost_count: u64, + pub costs: Vec, + pub last_used_at: String, + pub last_used_at_display: String, +} + +/// 单一货币的可查询成本。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CurrencyCostView { + pub currency: String, + pub estimated_amount: String, + pub estimated_amount_display: String, +} + +/// 账号详情类 GET 的固定 ID query。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountIdQuery { + pub account_id: String, +} + +impl AccountIdQuery { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId") + } + + pub(super) fn into_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) + } +} + +/// 头像 GET 的固定 query；`version` 只参与浏览器缓存键，不进入上游请求。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountProfileAvatarQuery { + pub account_id: String, + pub version: Option, +} + +impl AccountProfileAvatarQuery { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId")?; + if self.version.as_deref().is_some_and(|version| { + version.is_empty() + || version.len() > MAX_AVATAR_VERSION_BYTES + || !version.bytes().all(|byte| byte.is_ascii_alphanumeric()) + }) { + return Err(WireValidationError::new("version")); + } + Ok(()) + } + + pub(super) fn into_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) + } +} + +/// 敏感导出的固定 query；IDs 使用逗号分隔，禁止隐式导出全部账号。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountExportQuery { + pub account_ids: String, + pub confirm: String, +} + +impl AccountExportQuery { + pub fn into_ids(self) -> Result, WireValidationError> { + if self.confirm != "export_sensitive_accounts" { + return Err(WireValidationError::new("confirm")); + } + let ids = self + .account_ids + .split(',') + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(ToOwned::to_owned) + .collect::>(); + if ids.is_empty() + || ids.len() > 200 + || ids + .iter() + .any(|id| require_account_id(id, "accountIds").is_err()) + { + return Err(WireValidationError::new("accountIds")); + } + let unique = ids.iter().collect::>(); + if unique.len() != ids.len() { + return Err(WireValidationError::new("accountIds")); + } + ids.into_iter() + .map(|id| { + ProviderAccountId::new(id).map_err(|_| WireValidationError::new("accountIds")) + }) + .collect() + } +} + +/// 账号运行期动作。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountActionRequest { + pub account_id: String, +} + +impl AccountActionRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId") + } + + pub(super) fn into_id(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) + } +} + +/// 主动额度重置卡消费请求。幂等键由 UI 生成并在不确定重试时复用，与官方一致。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountResetCreditConsumeRequest { + pub account_id: String, + #[serde(default)] + pub credit_id: Option, + pub redeem_request_id: String, +} + +impl AccountResetCreditConsumeRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId")?; + if self.credit_id.as_deref().is_some_and(|value| { + let value = value.trim(); + value.is_empty() || value.len() > 512 || value.chars().any(char::is_control) + }) { + return Err(WireValidationError::new("creditId")); + } + let value = self.redeem_request_id.trim(); + let redeem_request_id = + Uuid::parse_str(value).map_err(|_| WireValidationError::new("redeemRequestId"))?; + if redeem_request_id.get_version() != Some(Version::Random) + || redeem_request_id.hyphenated().to_string() != value + { + return Err(WireValidationError::new("redeemRequestId")); + } + Ok(()) + } + + pub(super) fn into_command(self) -> Result { + self.validate()?; + Ok(ConsumeProviderResetCredit { + account_id: ProviderAccountId::new(self.account_id) + .map_err(|_| WireValidationError::new("accountId"))?, + credit_id: self.credit_id.map(|value| value.trim().to_owned()), + redeem_request_id: Uuid::parse_str(&self.redeem_request_id) + .map_err(|_| WireValidationError::new("redeemRequestId"))?, + }) + } +} + +/// 手工 OAuth 刷新会变更持久 credential。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountRefreshRequest { + pub account_id: String, +} + +impl AccountRefreshRequest { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId") + } + + pub(super) fn into_command(self) -> Result { + self.validate()?; + ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId")) + } +} + +/// 连接测试 query；测试仍经唯一 Core/Provider 模型请求路径执行。 +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountTestQuery { + pub account_id: String, + pub model_id: String, +} + +impl AccountTestQuery { + pub fn validate(&self) -> Result<(), WireValidationError> { + require_account_id(&self.account_id, "accountId")?; + if self.model_id.trim().is_empty() || self.model_id.chars().any(char::is_control) { + return Err(WireValidationError::new("modelId")); + } + Ok(()) + } + + pub(super) fn into_command( + self, + ) -> Result<(ProviderAccountId, UpstreamModelId), WireValidationError> { + self.validate()?; + Ok(( + ProviderAccountId::new(self.account_id) + .map_err(|_| WireValidationError::new("accountId"))?, + UpstreamModelId::new(self.model_id).map_err(|_| WireValidationError::new("modelId"))?, + )) + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct AccountModelView { + pub id: String, + pub label: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct AccountModelsData { + pub models: Vec, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountRefreshData { + pub account: AccountView, +} + +#[derive(Debug, Clone, Serialize)] +pub struct AccountQuotaData { + pub account: AccountView, +} + +/// Provider 官方个人资料统计响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountProfileStatisticsData { + pub display_name: Option, + pub username: Option, + pub image_url: Option, + pub has_stats_error: bool, + pub summary: AccountProfileStatisticsSummaryView, + pub daily_usage: Option>, + pub activity_insights: AccountProfileActivityInsightsView, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountProfileStatisticsSummaryView { + pub total_text_tokens: Option, + pub peak_tokens: Option, + pub longest_task_duration_ms: Option, + pub current_streak_days: Option, + pub longest_streak_days: Option, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountProfileDailyUsageView { + pub date: String, + pub tokens: u64, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountProfileActivityInsightsView { + pub fast_mode_percent: Option, + pub reasoning_effort: Option, + pub reasoning_effort_percent: Option, + pub skills_explored: Option, + pub total_skills_used: Option, + pub total_threads: Option, + pub invocations: Option>, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountProfileInvocationView { + #[serde(rename = "type")] + pub invocation_type: String, + pub plugin_id: Option, + pub plugin_name: Option, + pub skill_id: Option, + pub skill_name: Option, + pub usage_count: Option, +} + +impl From for AccountProfileStatisticsData { + fn from(statistics: ProviderProfileStatistics) -> Self { + Self { + display_name: statistics.display_name, + username: statistics.username, + image_url: statistics.image_url, + has_stats_error: statistics.has_stats_error, + summary: profile_statistics_summary_view(statistics.summary), + daily_usage: statistics + .daily_usage + .map(|daily| daily.into_iter().map(profile_daily_usage_view).collect()), + activity_insights: profile_activity_insights_view(statistics.activity_insights), + } + } +} + +const fn profile_statistics_summary_view( + summary: ProviderProfileStatisticsSummary, +) -> AccountProfileStatisticsSummaryView { + AccountProfileStatisticsSummaryView { + total_text_tokens: summary.total_text_tokens, + peak_tokens: summary.peak_tokens, + longest_task_duration_ms: summary.longest_task_duration_ms, + current_streak_days: summary.current_streak_days, + longest_streak_days: summary.longest_streak_days, + } +} + +fn profile_daily_usage_view(daily: ProviderProfileDailyUsage) -> AccountProfileDailyUsageView { + AccountProfileDailyUsageView { + date: daily.date.to_string(), + tokens: daily.tokens, + } +} + +fn profile_activity_insights_view( + insights: ProviderProfileActivityInsights, +) -> AccountProfileActivityInsightsView { + AccountProfileActivityInsightsView { + fast_mode_percent: insights.fast_mode_percent, + reasoning_effort: insights.reasoning_effort, + reasoning_effort_percent: insights.reasoning_effort_percent, + skills_explored: insights.skills_explored, + total_skills_used: insights.total_skills_used, + total_threads: insights.total_threads, + invocations: insights.invocations.map(|invocations| { + invocations + .into_iter() + .map(profile_invocation_view) + .collect() + }), + } +} + +fn profile_invocation_view(invocation: ProviderProfileInvocation) -> AccountProfileInvocationView { + AccountProfileInvocationView { + invocation_type: invocation.invocation_type, + plugin_id: invocation.plugin_id, + plugin_name: invocation.plugin_name, + skill_id: invocation.skill_id, + skill_name: invocation.skill_name, + usage_count: invocation.usage_count, + } +} + +/// 主动额度重置卡列表响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountResetCreditsData { + pub available_count: u64, + pub credits: Vec, +} + +/// 一张安全主动额度重置卡视图。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountResetCreditView { + pub id: String, + pub status: Option, + pub title: Option, + pub expires_at: Option, + pub reset_type: Option, +} + +/// 主动额度重置卡消费响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountResetCreditResultData { + pub code: String, + pub credit: Option, +} + +impl From for AccountResetCreditView { + fn from(credit: ProviderResetCredit) -> Self { + Self { + id: credit.id, + status: credit.status, + title: credit.title, + expires_at: credit.expires_at.map(|value| value.to_rfc3339()), + reset_type: credit.reset_type, + } + } +} + +impl From for AccountResetCreditsData { + fn from(credits: ProviderResetCredits) -> Self { + Self { + available_count: credits.available_count, + credits: credits + .credits + .into_iter() + .map(AccountResetCreditView::from) + .collect(), + } + } +} + +impl From for AccountResetCreditResultData { + fn from(result: ProviderResetCreditResult) -> Self { + Self { + code: result.code, + credit: result.credit.map(AccountResetCreditView::from), + } + } +} + +/// Provider-owned 明文导出文档；Debug 永远不输出内部 JSON。 +#[derive(Serialize)] +#[serde(transparent)] +pub struct AccountExportData(Value); + +impl AccountExportData { + #[must_use] + pub const fn new(value: Value) -> Self { + Self(value) + } + + pub(super) 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 { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("AccountExportData()") + } +} + +pub struct AccountConnectionTestEvent { + pub data: Value, +} + +impl From for AccountConnectionTestEvent { + fn from(event: DomainConnectionTestEvent) -> Self { + let data = match event { + DomainConnectionTestEvent::Started { model } => serde_json::json!({ + "type": "test_start", + "model": model, + "text": "正在连接上游 Responses" + }), + 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 => serde_json::json!({ + "type": "test_complete", + "success": true + }), + DomainConnectionTestEvent::Failed { + source, + gateway_error_code, + send_state, + message, + provider_error_code, + provider_error_type, + upstream_status, + upstream_content_type, + upstream_body, + } => serde_json::json!({ + "type": "error", + "source": source.as_str(), + "gatewayErrorCode": gateway_error_code.as_str(), + "sendState": send_state.map(gateway_core::upstream::UpstreamSendState::as_str), + "error": message, + "providerErrorCode": provider_error_code, + "providerErrorType": provider_error_type, + "upstreamStatus": upstream_status, + "upstreamContentType": upstream_content_type, + "upstreamBody": upstream_body + }), + }; + Self { data } + } +} diff --git a/backend/crates/gateway-api/tests/admin/accounts.rs b/backend/crates/gateway-api/tests/admin/accounts/mod.rs similarity index 100% rename from backend/crates/gateway-api/tests/admin/accounts.rs rename to backend/crates/gateway-api/tests/admin/accounts/mod.rs diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index a6a339a5b..4cbced33b 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -72,9 +72,8 @@ use gateway_core::{ account::{AccountStatusFacts, CredentialState, ProviderAccountId, QuotaState}, engine::probe::{AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult}, policy::{ClientApiKeyId, RateLimits}, - routing::{ - ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId, snapshot::SnapshotControl, - }, + routing::{ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId}, + runtime::SnapshotControl, }; mod account_groups; diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index 9091d0ab1..0365dfce1 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -34,7 +34,11 @@ fn source_tree_should_match_frozen_machine_manifest() { actual.sort(); let mut expected = vec![ "src/admin/account_groups.rs", - "src/admin/accounts.rs", + "src/admin/accounts/credentials.rs", + "src/admin/accounts/handlers.rs", + "src/admin/accounts/mod.rs", + "src/admin/accounts/presenter.rs", + "src/admin/accounts/wire.rs", "src/admin/auth.rs", "src/admin/backups.rs", "src/admin/client_keys.rs", @@ -81,7 +85,7 @@ fn test_tree_should_match_frozen_rust_mirror() { actual.sort(); let mut expected = vec![ "tests/admin/account_groups.rs", - "tests/admin/accounts.rs", + "tests/admin/accounts/mod.rs", "tests/admin/auth.rs", "tests/admin/client_keys.rs", "tests/admin/errors.rs", @@ -97,7 +101,7 @@ fn test_tree_should_match_frozen_rust_mirror() { "tests/main.rs", "tests/openai/auth.rs", "tests/openai/error.rs", - "tests/openai/images.rs", + "tests/openai/images/mod.rs", "tests/openai/mod.rs", "tests/openai/models.rs", "tests/openai/responses/http.rs", diff --git a/backend/crates/gateway-api/tests/openai/images.rs b/backend/crates/gateway-api/tests/openai/images/mod.rs similarity index 100% rename from backend/crates/gateway-api/tests/openai/images.rs rename to backend/crates/gateway-api/tests/openai/images/mod.rs diff --git a/backend/crates/gateway-api/tests/openai/mod.rs b/backend/crates/gateway-api/tests/openai/mod.rs index 4124bd056..eea8f48d5 100644 --- a/backend/crates/gateway-api/tests/openai/mod.rs +++ b/backend/crates/gateway-api/tests/openai/mod.rs @@ -36,11 +36,11 @@ use gateway_core::policy::{ ClientApiKeyId, ClientPolicy, CodexClientMinVersions, CodexClientVersion, PlaintextClientApiKey, RateLimits, }; -use gateway_core::routing::snapshot::RuntimeSnapshotHandle; use gateway_core::routing::{ ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId, }; +use gateway_core::runtime::RuntimeSnapshotHandle; use gateway_core::upstream::UpstreamSendState; pub(super) async fn api_router(execution: Arc) -> axum::Router { diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index db7d3a3dc..78d8138f0 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -34,11 +34,11 @@ use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader}; use crate::lifecycle::CancellationToken; use crate::operation::{Operation, ProviderSessionState}; use crate::policy::{ClientApiKeyId, ClientPolicy}; -use crate::routing::snapshot::RuntimeSnapshotHandle; use crate::routing::{ ProviderKind, PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot, UpstreamModelId, }; +use crate::runtime::RuntimeSnapshotHandle; const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60); const COORDINATION_TIMEOUT: Duration = Duration::from_millis(100); diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index f6f353567..238ae8364 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -14,6 +14,7 @@ pub mod operation; pub mod policy; pub mod provider_ports; pub mod routing; +pub mod runtime; pub mod task; pub mod upstream; @@ -31,9 +32,9 @@ use engine::execution::{ use engine::probe::AccountProbe; use engine::provider::ProviderRegistry; use health::HealthProbe; -use routing::snapshot::{ - RuntimeSnapshotCompiler, RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, - SnapshotStorePort, SnapshotSubscriptionPort, +use routing::snapshot::{RuntimeSnapshotCompiler, SnapshotStorePort}; +use runtime::{ + RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, SnapshotSubscriptionPort, }; use task::{WorkerContribution, WorkerDefinitionError}; diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index 0b5b6bff9..3e10f047d 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -2,29 +2,19 @@ use std::collections::{BTreeMap, BTreeSet}; use std::num::NonZeroU32; -use std::pin::Pin; -use std::sync::{Arc, RwLock}; +use std::sync::Arc; 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::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy}; use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry}; use crate::error::RoutingError; -use crate::health::{HealthProbe, HealthState}; -use crate::lifecycle::CancellationToken; use crate::operation::Operation; use crate::policy::{ ClientApiKeyId, ClientPolicy, CodexClientMinVersions, CodexClientVersion, PlaintextClientApiKey, RateLimits, }; -use crate::task::{ - DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, - WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, - WorkerSchedule, WorkerTaskError, -}; use super::{ AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, @@ -32,11 +22,6 @@ use super::{ RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, 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 在一个一致性读取中提供的调度设置事实。 @@ -203,22 +188,6 @@ impl SnapshotStoreError { } } -/// 不泄漏订阅基础设施细节的通知错误。 -#[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>; @@ -226,18 +195,6 @@ pub trait SnapshotStorePort: Send + Sync { 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 { @@ -264,6 +221,16 @@ impl RuntimeSnapshotCompiler { Self { store, providers } } + pub(crate) fn store(&self) -> Arc { + Arc::clone(&self.store) + } + + pub(crate) fn provider_catalog_generations( + &self, + ) -> BTreeMap { + self.providers.catalog_generations() + } + /// 读取一个 revision，并为已注册 Provider 查询实时模型目录。 pub async fn compile(&self) -> Result { for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS { @@ -868,301 +835,3 @@ impl RuntimeSnapshot { }) } } - -/// 请求级冻结失败；此状态必须 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/runtime/mod.rs b/backend/crates/gateway-core/src/runtime/mod.rs new file mode 100644 index 000000000..97a353892 --- /dev/null +++ b/backend/crates/gateway-core/src/runtime/mod.rs @@ -0,0 +1,355 @@ +//! 运行时快照的原子发布、健康状态与跨进程版本收敛。 + +use std::collections::BTreeMap; +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::provider::ProviderCatalogGeneration; +use crate::health::{HealthProbe, HealthState}; +use crate::lifecycle::CancellationToken; +use crate::routing::snapshot::{ + RuntimeSnapshot, RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotStorePort, +}; +use crate::routing::{ConfigRevision, ProviderKind}; +use crate::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerRegistration, WorkerRunnable, + WorkerSchedule, WorkerTaskError, +}; + +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); + +/// 不泄漏订阅基础设施细节的通知错误。 +#[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>>; + +/// 跨进程 revision 通知的基础设施中立端口。 +pub trait SnapshotSubscriptionPort: Send + Sync { + fn publish_snapshot_revision( + &self, + revision: ConfigRevision, + ) -> BoxFuture<'_, Result<(), SnapshotSubscriptionError>>; + + fn subscribe_snapshot_revisions( + &self, + ) -> BoxFuture<'_, Result>; +} + +/// 请求级冻结失败；此状态必须 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.provider_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: 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/tests/account.rs b/backend/crates/gateway-core/tests/account/mod.rs similarity index 100% rename from backend/crates/gateway-core/tests/account.rs rename to backend/crates/gateway-core/tests/account/mod.rs diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index 4fbf92dab..23f4b81bf 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -41,12 +41,12 @@ use gateway_core::operation::{ RawJsonPayload, }; use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits}; -use gateway_core::routing::snapshot::RuntimeSnapshotHandle; use gateway_core::routing::{ ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId, }; +use gateway_core::runtime::RuntimeSnapshotHandle; use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; use serde_json::json; diff --git a/backend/crates/gateway-core/tests/main.rs b/backend/crates/gateway-core/tests/main.rs index bd3aadb87..cbb9914ad 100644 --- a/backend/crates/gateway-core/tests/main.rs +++ b/backend/crates/gateway-core/tests/main.rs @@ -9,5 +9,6 @@ mod operation; mod policy; mod provider_ports; mod routing; +mod runtime; mod task; mod upstream; diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index 12260d3b5..3d5aa4e94 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -1,14 +1,11 @@ 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::account::{AccountSelectionPolicy, RotationStrategy}; use gateway_core::engine::AttemptContext; use gateway_core::engine::provider::{ Provider, ProviderCatalogGeneration, ProviderModelCapabilities, ProviderRegistry, @@ -18,17 +15,14 @@ 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, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, - SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, SnapshotProviderAccountFacts, - SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort, - SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh, + RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, + SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts, SnapshotFacts, + SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort, }; use gateway_core::routing::{ ConfigRevision, ModelCapabilities, ModelPresentation, ProviderKind, PublicModelId, - RuntimeSnapshot, UpstreamModelId, + UpstreamModelId, }; -use gateway_core::task::WorkerKind; use gateway_core::upstream::UpstreamSendState; #[derive(Clone)] @@ -57,32 +51,6 @@ impl SnapshotStorePort for TestSnapshotStore { } } -#[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, @@ -164,14 +132,6 @@ impl Provider for PublishingCatalogProvider { } } -#[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); @@ -234,109 +194,6 @@ fn compiler_retries_when_provider_publishes_catalog_during_compilation() { ); } -#[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, - ], - ); -} - #[test] fn compiler_should_freeze_valid_client_min_versions() { let store = Arc::new(TestSnapshotStore::new(Ok(facts_with_min_versions( @@ -422,21 +279,6 @@ 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/runtime/mod.rs b/backend/crates/gateway-core/tests/runtime/mod.rs new file mode 100644 index 000000000..34b5cc1dd --- /dev/null +++ b/backend/crates/gateway-core/tests/runtime/mod.rs @@ -0,0 +1,235 @@ +use std::collections::BTreeMap; +use std::num::NonZeroU32; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures::executor::block_on; +use futures::future::BoxFuture; + +use gateway_core::account::{AccountSelectionPolicy, RotationStrategy}; +use gateway_core::engine::provider::ProviderRegistry; +use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}; +use gateway_core::routing::snapshot::{ + RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, + SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts, + SnapshotStoreError, SnapshotStorePort, +}; +use gateway_core::routing::{ConfigRevision, RuntimeSnapshot}; +use gateway_core::runtime::{ + RuntimeSnapshotHandle, RuntimeSnapshotPublisher, SnapshotControl, SnapshotRevisionStream, + SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh, +}; +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(SnapshotFacts::config_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) }) + } +} + +#[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 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([("public-model".to_owned(), "upstream-model".to_owned())]), + None, + None, + ), + vec![SnapshotClientPolicyFacts::new( + ClientApiKeyId::new("key_one").expect("key ID"), + PlaintextClientApiKey::new("sk_test").expect("plaintext key"), + Vec::new(), + RateLimits::unlimited(), + )], + Vec::::new(), + Vec::::new(), + Vec::::new(), + ) +} + +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-store/src/redis/runtime_change.rs b/backend/crates/gateway-store/src/redis/runtime_change.rs index 99633b979..be6ae3cba 100644 --- a/backend/crates/gateway-store/src/redis/runtime_change.rs +++ b/backend/crates/gateway-store/src/redis/runtime_change.rs @@ -4,9 +4,9 @@ use std::pin::Pin; use async_trait::async_trait; use futures::{Stream, StreamExt}; -use gateway_core::routing::{ - ConfigRevision, - snapshot::{SnapshotRevisionStream, SnapshotSubscriptionError, SnapshotSubscriptionPort}, +use gateway_core::{ + routing::ConfigRevision, + runtime::{SnapshotRevisionStream, SnapshotSubscriptionError, SnapshotSubscriptionPort}, }; use redis::AsyncCommands; use serde::{Deserialize, Serialize}; diff --git a/backend/crates/gateway-store/tests/backup/mod.rs b/backend/crates/gateway-store/tests/backup/mod.rs new file mode 100644 index 000000000..67039598d --- /dev/null +++ b/backend/crates/gateway-store/tests/backup/mod.rs @@ -0,0 +1 @@ +mod s3; diff --git a/backend/crates/gateway-store/tests/s3.rs b/backend/crates/gateway-store/tests/backup/s3.rs similarity index 100% rename from backend/crates/gateway-store/tests/s3.rs rename to backend/crates/gateway-store/tests/backup/s3.rs diff --git a/backend/crates/gateway-store/tests/main.rs b/backend/crates/gateway-store/tests/main.rs index 0570001a2..f1acfc4ba 100644 --- a/backend/crates/gateway-store/tests/main.rs +++ b/backend/crates/gateway-store/tests/main.rs @@ -1,5 +1,5 @@ +mod backup; mod config; mod postgres; mod redis; -mod s3; mod support; diff --git a/backend/crates/gateway-store/tests/postgres/observability.rs b/backend/crates/gateway-store/tests/postgres/observability/mod.rs similarity index 100% rename from backend/crates/gateway-store/tests/postgres/observability.rs rename to backend/crates/gateway-store/tests/postgres/observability/mod.rs diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs similarity index 100% rename from backend/crates/gateway-store/tests/postgres/provider_accounts.rs rename to backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs deleted file mode 100644 index df15e01cf..000000000 --- a/backend/crates/providers/openai/src/provider.rs +++ /dev/null @@ -1,3391 +0,0 @@ -//! Codex 的 `gateway-core` Provider adapter。 - -use std::collections::BTreeSet; -use std::fmt; -use std::num::NonZeroU32; -use std::sync::Arc; -use std::time::{Duration, Instant, SystemTime}; - -use async_trait::async_trait; -use bytes::Bytes; -use futures::{StreamExt, future::BoxFuture}; -use gateway_core::account::{AccountFeedbackStats, ProviderAccount}; -use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope}; -use gateway_core::engine::provider::{ - EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, - ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, - ProviderSelectionObservation, ProviderStream, -}; -use gateway_core::engine::{AttemptContext, AttemptTransport, ContinuationAttempt}; -use gateway_core::error::{ - ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure, - OpaqueUpstreamValue, ProviderError, ProviderErrorKind, -}; -use gateway_core::event::{ - FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, - ProviderResponseMetadata, ProviderResponseObservation, ProviderResponseTimings, ResponseMeta, - UpstreamHttpVersion, WebSocketPoolKind, -}; -use gateway_core::lifecycle::CancellationToken; -use gateway_core::operation::{ - GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProviderSessionState, -}; -use gateway_core::provider_ports::ProviderSessionAffinityKey; -use gateway_core::routing::{ - ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId, -}; -use gateway_core::task::{ - DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, - WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration, - WorkerRunnable, WorkerSchedule, WorkerTaskError, -}; -use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; -use gateway_protocol::openai::events::{ - ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs, -}; -use reqwest::Client; -use secrecy::{ExposeSecret, SecretString}; -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value, json}; -use url::Url; -use uuid::Uuid; - -use crate::credential::{ - CodexAccountFailure, CodexCredentialCatalogError, CodexCredentialCatalogService, - CodexCredentialLease, CodexCredentialQuotaService, CodexCredentialRefreshOutcome, - CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, - CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, - SelectCodexCredential, SelectCodexProviderEndpointCredential, - derive_codex_cyber_policy_session_key, derive_codex_session_affinity, -}; -use crate::session_transport::CodexSessionTransportFallbacks; -use crate::transport::canonical::{ - CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, -}; -use crate::transport::catalog::{ - CodexCatalogCapabilityEvidence, CodexCatalogModel, CodexCatalogVisibility, -}; -use crate::transport::diagnostics::{ - CodexFailureCategory, CodexUpstreamFailure, CodexUpstreamSendPhase, -}; -use crate::transport::profile::{ - APPCAST_POLL_INTERVAL, CodexDesktopReleaseService, CodexWireProfileState, -}; -use crate::transport::protocol::responses::{ - CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, - PreviousResponseScope, ResponseEventSignals, TransportRequirement, transport_requirement, -}; -use crate::transport::protocol::websocket::WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE; -use crate::transport::request::{ - CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account, -}; -use crate::transport::session::CodexSessionIdentity; -use crate::transport::usage::normalize_service_tier; -use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason}; -use crate::transport::{ - CODEX_IMAGE_EDITS_PATH, CODEX_IMAGE_GENERATIONS_PATH, CODEX_RESPONSES_PATH, - CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse, - CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates, - CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics, - CodexWebSocketPool, endpoint_url, -}; - -const PROVIDER_NAME: &str = "openai"; -const HTTP_SSE_TRANSPORT: &str = "http_sse"; -const HTTP_JSON_TRANSPORT: &str = "http_json"; -const WEBSOCKET_TRANSPORT: &str = "websocket"; -const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024; -/// 提交边界前最多保留 64 KiB 原始上游 chunk；达到阈值后结束无感换号窗口， -/// 但不会把上游数据改写成协议失败。 -const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; -/// 短暂保留 response.created 等结构事件，让随后到达的明确拒绝可以无感换号； -/// 到期即放行，避免模型长时间思考时让客户端一直收不到首事件。 -const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(1_200); -// 额度拒绝后先给上游额度结算留出时间，再以受限时长同步 usage 快照。 -const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2); -const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); -pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api"; -pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api"; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CodexProviderTransport { - HttpOnly, - PreferWebSocket, -} - -#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -pub enum CodexProviderConfigError { - #[error("Codex provider URL is invalid")] - InvalidBaseUrl, -} - -pub struct CodexProvider { - selector: Arc, - catalog: Arc, - quota: Arc, - account_feedback: Arc, - client: CodexBackendClient, - responses_url: Url, - image_generations_url: Url, - image_edits_url: Url, - session_identity: Option, - session_transport_fallbacks: CodexSessionTransportFallbacks, - stream_max_retries: u32, -} - -impl CodexProvider { - // Provider 构造集中装配独立领域服务和透明传输依赖，拆分参数会模糊所有权。 - #[expect(clippy::too_many_arguments)] - pub fn new( - selector: Arc, - catalog: Arc, - quota: Arc, - account_feedback: Arc, - http: Client, - profile: CodexWireProfileState, - base_url: String, - websocket_pool: Arc, - stream_max_retries: u32, - ) -> Result { - let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH)) - .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; - let image_generations_url = - Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_GENERATIONS_PATH)) - .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; - let image_edits_url = Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_EDITS_PATH)) - .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; - let client = - CodexBackendClient::new(http, base_url, profile).with_websocket_pool(websocket_pool); - Ok(Self { - selector, - catalog, - quota, - account_feedback, - client, - responses_url, - image_generations_url, - image_edits_url, - session_identity: None, - session_transport_fallbacks: CodexSessionTransportFallbacks::default(), - stream_max_retries, - }) - } - - pub(crate) fn with_session_identity(mut self, identity: CodexSessionIdentity) -> Self { - self.session_identity = Some(identity); - self - } -} - -impl fmt::Debug for CodexProvider { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("CodexProvider") - .field("selector", &"") - .field("catalog", &"") - .finish() - } -} - -#[async_trait] -impl Provider for CodexProvider { - fn name(&self) -> &'static str { - 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, reasoning_effort) = encode_generate_request(request, "observability") - .map(|encoded| { - let semantics = encoded.semantics(); - let reasoning_effort = semantics.reasoning_effort.clone(); - (semantics, reasoning_effort) - }) - .unwrap_or_default(); - ProviderRequestObservation { - reasoning_effort, - 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, - ) -> Result, ProviderError> { - let snapshot = self.catalog.synchronize().await.map_err(|_| { - provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) - })?; - Ok(snapshot - .models() - .iter() - .map(compile_model_capabilities) - .collect()) - } - - async fn execute( - &self, - request: ProviderRequest, - context: AttemptContext, - ) -> Result { - if request.candidate().provider().as_str() != PROVIDER_NAME { - return Err(provider_error( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - )); - } - let candidate = request.candidate(); - if context.cancellation().is_cancelled() { - return Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::NotSent, - )); - } - if remaining(context.deadline()).is_none() { - return Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::NotSent, - )); - } - if let Operation::GenerateImage(image) = request.operation() { - return self.execute_image(image, candidate, context).await; - } - let Operation::Generate(generate) = request.operation() else { - return Err(provider_error( - ProviderErrorKind::Unsupported, - UpstreamSendState::NotSent, - )); - }; - let Some(upstream_model) = candidate.upstream_model() else { - return Err(provider_error( - ProviderErrorKind::Protocol, - UpstreamSendState::NotSent, - )); - }; - let previous_session = decode_openai_session_state(generate); - let continuation_requested = generate.native_continuation_requested(); - let mut upstream_request = encode_generate_request(generate, upstream_model.as_str()) - .map_err(map_request_error)?; - 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 let Some(identity) = &self.session_identity { - identity.prepare_local_conversation(&mut upstream_request); - } - if previous_session.as_ref().is_some_and(|state| { - same_client_turn( - state.client_turn_id.as_deref(), - upstream_request.client_turn_id.as_deref(), - ) - }) { - upstream_request.turn_state = previous_session - .as_ref() - .and_then(|state| state.turn_state.clone()); - } - let session_affinity = - derive_codex_session_affinity(&upstream_request, context.client_api_key_ref()); - let cyber_policy_session_key = - derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref()); - - let selection_started_at = Instant::now(); - let lease = self - .selector - .select_with_cyber_policy( - &SelectCodexCredential { - upstream_model: upstream_model.as_str(), - request_url: &self.responses_url, - attempt: &context, - session_affinity_key: session_affinity.as_ref().map(|affinity| affinity.key()), - }, - cyber_policy_session_key.as_ref(), - session_affinity.as_ref(), - ) - .await - .map_err(map_selection_error)?; - let account_selection_wait_ms = - u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); - let lease = Arc::new(lease); - // 首字计时的起点：账号选择完成之后、上游建立之前。 - let output_started_at = Instant::now(); - let provider_kind = ProviderKind::new(PROVIDER_NAME) - .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - let state_owner_cross_account = context - .account_state_owner() - .is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id())) - || previous_session - .as_ref() - .is_some_and(|state| state.account_id != lease.account_id().as_str()); - let has_explicit_state_owner = - context.account_state_owner().is_some() || previous_session.is_some(); - let account_scope = - if state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()) { - RequestAccountScope::Different - } else if has_explicit_state_owner || lease.affinity_hit() { - RequestAccountScope::Same - } else { - RequestAccountScope::Unknown - }; - if context.continuation_attempt() != ContinuationAttempt::None - && let Some(continuation) = context.continuation() - { - match continuation { - ContinuationBinding::Pinned(continuation) => { - let previous_response_scope = match context.continuation_attempt() { - ContinuationAttempt::Native => match previous_session - .as_ref() - .map(|state| state.continuation_scope) - { - Some(OpenAiContinuationScope::Persisted) => { - PreviousResponseScope::Persisted - } - Some(OpenAiContinuationScope::ConnectionLocal) => { - PreviousResponseScope::ConnectionLocal - } - Some(OpenAiContinuationScope::ReplayRequired) => { - PreviousResponseScope::ExternalUnknown - } - None => match continuation.scope() { - NativeContinuationScope::Persisted => { - PreviousResponseScope::Persisted - } - NativeContinuationScope::ConnectionLocal => { - PreviousResponseScope::ConnectionLocal - } - }, - }, - ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => { - PreviousResponseScope::ExternalUnknown - } - ContinuationAttempt::None => PreviousResponseScope::ExternalUnknown, - }; - 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); - } - } - } - if upstream_request.previous_response_id().is_some() - && !account_scope.can_reuse_account_state() - { - return Err(continuation_replay_required_error()); - } - scope_request_to_account( - &mut upstream_request, - lease.installation_id(), - account_scope, - ); - let requirement = transport_requirement(&upstream_request); - let requested_transport = selected_transport(&upstream_request); - let sticky_http_fallback = context.transport() == AttemptTransport::Default - && requested_transport == CodexProviderTransport::PreferWebSocket - && !requirement.requires_websocket() - && session_affinity.as_ref().is_some_and(|affinity| { - self.session_transport_fallbacks - .is_http_only(affinity.key()) - }); - let transport = if requirement.requires_websocket() { - CodexProviderTransport::PreferWebSocket - } else { - match context.transport() { - AttemptTransport::Default if sticky_http_fallback => { - CodexProviderTransport::HttpOnly - } - AttemptTransport::Default => requested_transport, - AttemptTransport::Retry(_) => requested_transport, - AttemptTransport::Fallback => CodexProviderTransport::HttpOnly, - } - }; - if sticky_http_fallback { - tracing::info!( - request_id = %context.request_id(), - attempt_index = context.attempt_index().get(), - session_affinity_key_hash = session_affinity - .as_ref() - .map_or("", CodexSessionAffinity::key_hash), - transport_requirement = requirement.as_str(), - "Reusing session-sticky OpenAI upstream HTTP fallback" - ); - } - apply_transport(&mut upstream_request, transport); - let metadata = ProviderCallMetadata::new( - provider_kind, - upstream_model.clone(), - lease.account_id().clone(), - UpstreamTransport::new(transport_name(transport)).map_err(|_| { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) - })?, - ) - .with_selection_observation(ProviderSelectionObservation::new( - account_selection_wait_ms, - lease.capacity_snapshot(), - )); - let response_store = upstream_request.store(); - let session_capture = - (!continuation_requested || previous_session.is_some()).then(|| OpenAiSessionCapture { - account_id: lease.account_id().as_str().to_owned(), - conversation_id: upstream_request.local_conversation_id.clone(), - turn_state: upstream_request.turn_state.clone(), - client_turn_id: upstream_request.client_turn_id.clone(), - response_store, - continuation_scope: None, - }); - let allows_account_state_mutation = lease.allows_account_state_mutation(); - let session_affinity_key_hash = session_affinity - .as_ref() - .map(|affinity| affinity.key_hash().to_owned()); - let session_affinity_key = session_affinity.map(CodexSessionAffinity::into_key); - let websocket_retry_count = match context.transport() { - AttemptTransport::Retry(retry_index) => retry_index.get(), - AttemptTransport::Default | AttemptTransport::Fallback => 0, - }; - let events = cold_response_stream(ColdResponse { - client: self.client.clone(), - response_origin: self.responses_url.clone(), - request: upstream_request, - upstream_model: upstream_model.clone(), - transport_policy: transport, - context, - selector: Arc::clone(&self.selector), - quota: Arc::clone(&self.quota), - catalog: Arc::clone(&self.catalog), - lease: Arc::clone(&lease), - output_started_at, - session_affinity_key, - session_affinity_key_hash, - session_transport_fallbacks: self.session_transport_fallbacks.clone(), - websocket_retry_count, - stream_max_retries: self.stream_max_retries, - session_capture, - }); - let stream = ProviderStream::new(metadata, events, lease); - Ok(if allows_account_state_mutation { - stream.with_account_feedback(Arc::clone(&self.account_feedback)) - } else { - stream - }) - } -} - -impl CodexProvider { - async fn execute_image( - &self, - image: &ImageRequest, - candidate: &ProviderCandidate, - context: AttemptContext, - ) -> Result { - if image.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() { - return Err(provider_error( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - )); - } - let (endpoint_path, response_origin) = match image.kind() { - ImageRequestKind::Generation => ( - CODEX_IMAGE_GENERATIONS_PATH, - self.image_generations_url.clone(), - ), - ImageRequestKind::Edit => (CODEX_IMAGE_EDITS_PATH, self.image_edits_url.clone()), - }; - let selection_started_at = Instant::now(); - let lease = self - .selector - .select_for_provider_endpoint(&SelectCodexProviderEndpointCredential { - request_url: &response_origin, - attempt: &context, - }) - .await - .map_err(map_selection_error)?; - let account_selection_wait_ms = - u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); - let lease = Arc::new(lease); - let allows_account_state_mutation = lease.allows_account_state_mutation(); - let provider_kind = ProviderKind::new(PROVIDER_NAME) - .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - let metadata = ProviderCallMetadata::for_provider_endpoint( - provider_kind, - lease.account_id().clone(), - UpstreamTransport::new(HTTP_JSON_TRANSPORT).map_err(|_| { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) - })?, - ) - .with_selection_observation(ProviderSelectionObservation::new( - account_selection_wait_ms, - lease.capacity_snapshot(), - )); - let image_turn_id = image - .payload() - .context() - .get("image_turn_id") - .and_then(Value::as_str) - .map(str::to_owned); - let events = cold_json_response_stream(ColdJsonResponse { - client: self.client.clone(), - response_origin, - endpoint_path, - body: image.payload().body().clone(), - image_turn_id, - context, - selector: Arc::clone(&self.selector), - quota: Arc::clone(&self.quota), - lease: Arc::clone(&lease), - output_started_at: Instant::now(), - }); - let stream = ProviderStream::new(metadata, events, lease); - Ok(if allows_account_state_mutation { - stream.with_account_feedback(Arc::clone(&self.account_feedback)) - } else { - stream - }) - } -} - -struct ColdResponse { - client: CodexBackendClient, - response_origin: Url, - request: CodexResponsesRequest, - upstream_model: UpstreamModelId, - transport_policy: CodexProviderTransport, - context: AttemptContext, - selector: Arc, - quota: Arc, - catalog: Arc, - lease: Arc, - output_started_at: Instant, - session_affinity_key: Option, - session_affinity_key_hash: Option, - session_transport_fallbacks: CodexSessionTransportFallbacks, - websocket_retry_count: u32, - stream_max_retries: u32, - session_capture: Option, -} - -struct ColdJsonResponse { - client: CodexBackendClient, - response_origin: Url, - endpoint_path: &'static str, - body: Bytes, - image_turn_id: Option, - context: AttemptContext, - selector: Arc, - quota: Arc, - lease: Arc, - output_started_at: Instant, -} - -#[derive(Clone, Serialize, Deserialize)] -struct OpenAiSessionState { - account_id: String, - conversation_id: Option, - #[serde(default)] - turn_state: Option, - #[serde(default)] - client_turn_id: Option, - continuation_scope: OpenAiContinuationScope, -} - -#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -enum OpenAiContinuationScope { - Persisted, - ConnectionLocal, - ReplayRequired, -} - -struct OpenAiSessionCapture { - account_id: String, - conversation_id: Option, - turn_state: Option, - client_turn_id: Option, - response_store: bool, - continuation_scope: Option, -} - -fn same_client_turn(previous: Option<&str>, current: Option<&str>) -> bool { - previous - .zip(current) - .is_some_and(|(previous, current)| !previous.is_empty() && previous == current) -} - -fn decode_openai_session_state(request: &GenerateRequest) -> Option { - request - .provider_session_state(PROVIDER_NAME) - .and_then(|state| serde_json::from_value(Value::Object(state.payload().clone())).ok()) -} - -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, -) { - let Some(terminal_index) = events - .iter() - .position(|event| terminal_response_output(event).is_some()) - else { - return; - }; - let Some(capture) = capture.take() else { - return; - }; - let Some(continuation_scope) = capture.continuation_scope else { - return; - }; - let Ok(update) = encode_openai_session_state(OpenAiSessionState { - account_id: capture.account_id, - conversation_id: capture.conversation_id, - turn_state: capture.turn_state, - client_turn_id: capture.client_turn_id, - continuation_scope, - }) else { - return; - }; - events[terminal_index].attach_session_update(update); -} - -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() -} - -enum CodexHandshakeAttemptError { - Client(CodexClientError), - Cancelled, - Timeout, -} - -async fn create_response_attempt( - client: &CodexBackendClient, - request: &CodexResponsesRequest, - request_context: CodexRequestContext<'_>, - account_id: &str, - deadline: SystemTime, - cancellation: &CancellationToken, -) -> Result { - let Some(handshake_deadline) = remaining(deadline) else { - return Err(CodexHandshakeAttemptError::Timeout); - }; - tokio::select! { - biased; - _ = cancellation.cancelled() => Err(CodexHandshakeAttemptError::Cancelled), - _ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout), - response = client.create_response_stream_with_deferred_websocket_recovery( - request, - request_context, - Some(account_id), - ) => response.map_err(CodexHandshakeAttemptError::Client), - } -} - -fn map_handshake_attempt_error(error: CodexHandshakeAttemptError) -> MappedProviderFailure { - match error { - CodexHandshakeAttemptError::Client(error) => map_handshake_error(error), - CodexHandshakeAttemptError::Cancelled => MappedProviderFailure::plain(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Ambiguous, - )), - CodexHandshakeAttemptError::Timeout => MappedProviderFailure::plain(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Ambiguous, - )), - } -} - -async fn create_json_attempt( - request: &ColdJsonResponse, - account: &ProviderAccount, - installation_id: &str, - authorization: &SecretString, - cookie_header: Option<&SecretString>, - account_selection: CodexAccountSelectionTelemetry<'_>, -) -> Result { - let Some(handshake_deadline) = remaining(request.context.deadline()) else { - return Err(CodexHandshakeAttemptError::Timeout); - }; - let request_id = request.context.request_id().as_str(); - let mut request_context = CodexRequestContext::auxiliary( - authorization.expose_secret(), - account.upstream_account_id(), - request_id, - Some(installation_id), - ); - request_context.cookie_header = cookie_header.map(ExposeSecret::expose_secret); - request_context.account_selection = account_selection; - tokio::select! { - biased; - _ = request.context.cancellation().cancelled() => Err(CodexHandshakeAttemptError::Cancelled), - _ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout), - response = request.client.post_raw_json( - request.endpoint_path, - request.body.clone(), - request.image_turn_id.as_deref(), - request_context, - ) => response.map_err(CodexHandshakeAttemptError::Client), - } -} - -fn cold_json_response_stream(request: ColdJsonResponse) -> EventStream { - Box::pin(async_stream::try_stream! { - let allows_account_state_mutation = request.lease.allows_account_state_mutation(); - let failure_context = OpenAiFailureContext { - client: &request.client, - selector: &request.selector, - quota: &request.quota, - response_origin: &request.response_origin, - cyber_policy_scope: None, - allows_account_state_mutation, - }; - let active_account = request.lease.account().clone(); - let cookie_header = build_cookie_header(request.lease.cookies())?; - let authorization = request - .lease - .authentication() - .authorization_header() - .map_err(|_| { - provider_error( - ProviderErrorKind::Unauthorized, - UpstreamSendState::NotSent, - ) - })?; - let account_selection = CodexAccountSelectionTelemetry::new( - request.lease.affinity_hit(), - request.lease.escape_reason(), - request.lease.account_switch(), - ); - let response = create_json_attempt( - &request, - &active_account, - request.lease.installation_id(), - &authorization, - cookie_header.as_ref(), - account_selection, - ) - .await; - if let Err(CodexHandshakeAttemptError::Client(error)) = &response { - log_client_upstream_error( - UpstreamErrorLogContext::new(&request.context, &active_account, None), - error, - ); - } - let response = match response.map_err(map_handshake_attempt_error) { - Ok(response) => response, - Err(mut failure) => { - if let Some(observation) = failure.observation.take() { - yield ProviderEvent::observation(observation); - } - apply_failure(&failure_context, &active_account, &failure).await; - Err(failure.error)?; - return; - } - }; - - let mut metrics = response.transport_metrics.clone(); - metrics.first_event_ms = Some( - i64::try_from(request.output_started_at.elapsed().as_millis()).unwrap_or(i64::MAX), - ); - if let Some(observation) = codex_response_observation( - CodexBackendTransport::HttpJson, - &response.diagnostics, - &response.response_metadata, - &metrics, - None, - openai_response_timings(&metrics, &response.response_metadata), - ) { - yield ProviderEvent::observation(observation); - } - if allows_account_state_mutation { - synchronize_passive_quota_headers( - &request.quota, - &active_account, - &response.rate_limit_headers, - ) - .await; - if !response.set_cookie_headers.is_empty() - && let Err(error) = request - .selector - .capture_response_cookies( - &active_account, - &request.response_origin, - &response.set_cookie_headers, - ) - .await - { - tracing::warn!( - account_id = %active_account.id(), - error = %error, - "Failed to persist OpenAI image response cookies" - ); - } - } - - let response_meta = - ResponseMeta::for_provider_endpoint(request.context.request_id().as_str()); - yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone())); - let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent) - })?; - yield ProviderEvent::wire(wire); - yield ProviderEvent::canonical(GatewayEvent::Completed( - response_meta.with_finish_reason(FinishReason::Stop), - )); - }) -} - -fn cold_response_stream(response: ColdResponse) -> EventStream { - let ColdResponse { - client, - response_origin, - request, - upstream_model, - transport_policy, - context, - selector, - quota, - catalog, - lease, - output_started_at, - session_affinity_key, - session_affinity_key_hash, - session_transport_fallbacks, - websocket_retry_count, - stream_max_retries, - mut session_capture, - } = response; - Box::pin(async_stream::try_stream! { - let cyber_policy_scope = lease.cyber_policy_scope().cloned(); - let allows_account_state_mutation = lease.allows_account_state_mutation(); - let failure_context = OpenAiFailureContext { - client: &client, - selector: &selector, - quota: "a, - response_origin: &response_origin, - cyber_policy_scope: cyber_policy_scope.as_ref(), - allows_account_state_mutation, - }; - let mut active_account = lease.account().clone(); - let cookie_header = build_cookie_header(lease.cookies())?; - let authorization = lease - .authentication() - .authorization_header() - .map_err(|_| { - provider_error( - ProviderErrorKind::Unauthorized, - UpstreamSendState::NotSent, - ) - })?; - let request_id = context.request_id().as_str().to_owned(); - let cancellation = context.cancellation().clone(); - let account_selection = CodexAccountSelectionTelemetry::new( - lease.affinity_hit(), - lease.escape_reason(), - lease.account_switch(), - ); - let request_transport_requirement = transport_requirement(&request); - let response = create_response_attempt( - &client, - &request, - codex_request_context( - &request, - &request_id, - &active_account, - lease.installation_id(), - &authorization, - cookie_header.as_ref(), - account_selection, - ), - active_account.id().as_str(), - context.deadline(), - &cancellation, - ) - .await; - let websocket_failure_policy = match &response { - Err(CodexHandshakeAttemptError::Client(error)) - if transport_policy == CodexProviderTransport::PreferWebSocket => - { - websocket_client_failure_policy(error) - } - _ => None, - }; - if let Err(CodexHandshakeAttemptError::Client(error)) = &response { - log_client_upstream_error( - UpstreamErrorLogContext::new(&context, &active_account, None), - error, - ); - } - let response = response.map_err(map_handshake_attempt_error); - let response = match response { - Ok(response) => response, - Err(mut failure) => { - if let Some(policy) = websocket_failure_policy { - apply_websocket_recovery_policy( - &mut failure, - WebSocketRecoveryContext { - policy, - requirement: request_transport_requirement, - retry_count: websocket_retry_count, - max_retries: stream_max_retries, - request_id: context.request_id().as_str(), - attempt_index: context.attempt_index().get(), - account_id: active_account.id().as_str(), - session_affinity_key: session_affinity_key.as_ref(), - session_affinity_key_hash: session_affinity_key_hash.as_deref(), - session_transport_fallbacks: &session_transport_fallbacks, - }, - ); - } - if let Some(observation) = failure.observation.take() { - yield ProviderEvent::observation(observation); - } - apply_failure(&failure_context, &active_account, &failure) - .await; - Err(failure.error)?; - return; - } - }; - if !accepts_backend_transport(transport_policy, response.transport) { - let failure = MappedProviderFailure::plain(provider_error( - ProviderErrorKind::Protocol, - UpstreamSendState::Sent, - )); - apply_failure(&failure_context, &active_account, &failure) - .await; - Err(failure.error)?; - return; - } - 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 - }); - capture.turn_state = response.turn_state.clone().or(capture.turn_state.clone()); - } - let mut observation_state = OpenAiResponseObservationState::from_backend_response( - &response, - &request, - active_account.id().as_str(), - context.attempt_index().get(), - ); - if let Some(observation) = observation_state.observation() { - yield ProviderEvent::observation(observation); - } - if let Some(etag) = response.response_metadata.models_etag.as_deref() - && let Err(error) = catalog.observe_response_etag(etag) - { - tracing::warn!( - error = %error, - "OpenAI model ETag observation was rejected" - ); - } - if allows_account_state_mutation - && !response.set_cookie_headers.is_empty() - && let Ok(outcome) = selector - .capture_response_cookies( - &active_account, - &response_origin, - &response.set_cookie_headers, - ) - .await - && let Some(revision) = outcome.credential_revision - && let Ok(current) = selector.current_account(active_account.id()).await - && current.revision().get() == revision - { - active_account = current; - } - let response_transport = response.transport; - let websocket_connection_id = response.websocket_connection_id; - 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 mut passive_quota_observation = - OpenAiPassiveQuotaObservation::new(response.rate_limit_headers); - let rate_limit_updates = response.rate_limit_updates; - let turn_state_updates = response.turn_state_update; - // OpenAI 线路为透明代理：HTTP SSE 与 WebSocket 两条上游均启用 raw 透传， - // 下游按字节转发上游原文，避免 serde 往返改写数值/精度（大整数→f64、logprobs 等）。 - // WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON - // （transport/protocol/websocket.rs），push_frames 抽出的 data 即上游原文。 - let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str()) - .with_requested_service_tier(request.service_tier()) - .with_request_tool_pricing(upstream_model.as_str(), request.tools()) - .with_raw_sse_passthrough(); - let mut pre_commit_events = PreCommitClientEvents::new(); - loop { - let Some(stream_deadline) = remaining(context.deadline()) else { - if allows_account_state_mutation { - synchronize_passive_quota( - "a, - &active_account, - passive_quota_observation.rate_limits(), - ) - .await; - } - Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?; - return; - }; - let replay_grace_deadline = pre_commit_events.replay_grace_deadline(); - let next = tokio::select! { - biased; - _ = cancellation.cancelled() => Err(MappedProviderFailure::plain(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Sent, - ))), - _ = tokio::time::sleep(stream_deadline) => Err(MappedProviderFailure::plain(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - ))), - _ = wait_for_replay_grace(replay_grace_deadline) => Ok(PreCommitPoll::GraceElapsed), - chunk = body.next() => match chunk { - Some(Ok(chunk)) => Ok(PreCommitPoll::Upstream(Some(chunk))), - Some(Err(error)) => { - log_client_upstream_error( - UpstreamErrorLogContext::new( - &context, - &active_account, - websocket_connection_id, - ), - &error, - ); - Err(map_stream_error(error)) - } - None => Ok(PreCommitPoll::Upstream(None)), - }, - }; - let next = match next { - Ok(PreCommitPoll::Upstream(next)) => next, - Ok(PreCommitPoll::GraceElapsed) => { - for event in pre_commit_events.commit_pending() { - yield event; - } - continue; - } - Err(mut failure) => { - if failure.websocket_transport_retryable - && response_transport == CodexBackendTransport::WebSocket - && !pre_commit_events.is_committed() - { - apply_websocket_recovery_policy( - &mut failure, - WebSocketRecoveryContext { - policy: WebSocketFailurePolicy::Budgeted, - requirement: request_transport_requirement, - retry_count: websocket_retry_count, - max_retries: stream_max_retries, - request_id: context.request_id().as_str(), - attempt_index: context.attempt_index().get(), - account_id: active_account.id().as_str(), - session_affinity_key: session_affinity_key.as_ref(), - session_affinity_key_hash: session_affinity_key_hash.as_deref(), - session_transport_fallbacks: &session_transport_fallbacks, - }, - ); - } - let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; - if !updates.is_empty() { - passive_quota_observation.observe(&updates); - let update_headers = rate_limit_update_headers(&updates); - if observation_state.merge_rate_limit_headers(&update_headers) - && let Some(observation) = observation_state.observation() - { - yield ProviderEvent::observation(observation); - } - } - if allows_account_state_mutation { - synchronize_passive_quota( - "a, - &active_account, - passive_quota_observation.rate_limits(), - ) - .await; - } - apply_failure(&failure_context, &active_account, &failure) - .await; - Err(failure.error)?; - return; - } - }; - let Some(chunk) = next else { break; }; - let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; - let rate_limits_changed = if updates.is_empty() { - false - } else { - passive_quota_observation.observe(&updates); - observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates)) - }; - let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref() - && let Some(turn_state) = updates.lock().await.take() - { - if let Some(capture) = session_capture.as_mut() { - capture.turn_state = Some(turn_state.clone()); - } - observation_state.merge_client_header("x-codex-turn-state", &turn_state) - } else { - false - }; - let first_event_changed = - observation_state.observe_stream_chunk(&chunk, output_started_at); - let chunk_len = chunk.len(); - let (mut events, canonical_failure) = match decoder.push(&chunk) { - CodexCanonicalOutcome::Events(events) => (events, None), - CodexCanonicalOutcome::Failed(failure) => { - let (events, error, semantic_output_seen) = failure.into_parts(); - (events, Some((error, semantic_output_seen))) - } - }; - pre_commit_events.observe_chunk(chunk_len); - let service_tier_changed = observation_state - .observe_upstream_service_tier(decoder.response_service_tier()); - let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| { - log_canonical_upstream_error( - UpstreamErrorLogContext::new( - &context, - &active_account, - websocket_connection_id, - ), - response_transport, - &error, - ); - let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_)); - ( - map_canonical_error( - error, - &failure_diagnostics, - &failure_set_cookie_headers, - &failure_rate_limit_headers, - ReplayBoundary::from_semantic_output( - semantic_output_seen || pre_commit_events.is_committed(), - ), - ), - atomic_upstream_failure, - ) - }); - let timing_signals = decoder.take_timing_signals(); - let timing_changed = first_event_changed - || observation_state - .observe_timing_signals(timing_signals, output_started_at); - let completed = events - .iter() - .flat_map(ProviderEvent::canonical_facts) - .any(|event| matches!(event, GatewayEvent::Completed(_))); - let terminal_changed = completed - && observation_state.mark_completed(terminal_response_is_incomplete(&events)); - if allows_account_state_mutation && (completed || terminal_failure.is_some()) { - synchronize_passive_quota( - "a, - &active_account, - passive_quota_observation.rate_limits(), - ) - .await; - } - if let Some((failure, _)) = terminal_failure.as_ref() { - apply_failure(&failure_context, &active_account, failure) - .await; - } - attach_openai_session_update(&mut events, &mut session_capture); - if allows_account_state_mutation && completed && terminal_failure.is_none() { - // 完成事件一旦交给下游，Core 可以立刻停止轮询 Provider stream； - // 在此之前持久化亲和关系，保证成功请求不会因流被提前 drop 而丢失绑定。 - selector - .record_success( - &active_account, - session_affinity_key.as_ref(), - lease.affinity_expected_account_id(), - ) - .await; - selector - .observe_cyber_policy_success(cyber_policy_scope.as_ref()) - .await; - } - if (rate_limits_changed - || service_tier_changed - || timing_changed - || turn_state_changed - || terminal_changed) - && let Some(observation) = observation_state.observation() - { - yield ProviderEvent::observation(observation); - } - if let Some((mut failure, atomic_upstream_failure)) = terminal_failure { - let failure_after_commit = - timing_signals.semantic_output || pre_commit_events.is_committed(); - if failure_after_commit { - for event in pre_commit_events.commit(events) { - yield event; - } - } else if atomic_upstream_failure { - failure.error = failure - .error - .with_atomic_client_events(pre_commit_events.take_for_failure(events)); - } - Err(failure.error)?; - return; - } - let events = pre_commit_events.stage(events, timing_signals, completed); - for event in events { - yield event; - } - if completed { - return; - } - } - let (mut events, canonical_failure) = match decoder.finish() { - CodexCanonicalOutcome::Events(events) => (events, None), - CodexCanonicalOutcome::Failed(failure) => { - let (events, error, semantic_output_seen) = failure.into_parts(); - (events, Some((error, semantic_output_seen))) - } - }; - let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| { - log_canonical_upstream_error( - UpstreamErrorLogContext::new( - &context, - &active_account, - websocket_connection_id, - ), - response_transport, - &error, - ); - let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_)); - ( - map_canonical_error( - error, - &failure_diagnostics, - &failure_set_cookie_headers, - &failure_rate_limit_headers, - ReplayBoundary::from_semantic_output( - semantic_output_seen || pre_commit_events.is_committed(), - ), - ), - atomic_upstream_failure, - ) - }); - let timing_signals = decoder.take_timing_signals(); - let service_tier_changed = observation_state - .observe_upstream_service_tier(decoder.response_service_tier()); - let timing_changed = observation_state - .observe_timing_signals(timing_signals, output_started_at); - let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; - let rate_limits_changed = if updates.is_empty() { - false - } else { - passive_quota_observation.observe(&updates); - observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates)) - }; - if allows_account_state_mutation { - synchronize_passive_quota( - "a, - &active_account, - passive_quota_observation.rate_limits(), - ) - .await; - } - if let Some((failure, _)) = terminal_failure.as_ref() { - apply_failure(&failure_context, &active_account, failure) - .await; - } - let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref() - && let Some(turn_state) = updates.lock().await.take() - { - if let Some(capture) = session_capture.as_mut() { - capture.turn_state = Some(turn_state.clone()); - } - observation_state.merge_client_header("x-codex-turn-state", &turn_state) - } else { - false - }; - attach_openai_session_update(&mut events, &mut session_capture); - let completed = events - .iter() - .flat_map(ProviderEvent::canonical_facts) - .any(|event| matches!(event, GatewayEvent::Completed(_))); - let terminal_changed = completed - && observation_state.mark_completed(terminal_response_is_incomplete(&events)); - if allows_account_state_mutation && completed && terminal_failure.is_none() { - // 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。 - selector - .record_success( - &active_account, - session_affinity_key.as_ref(), - lease.affinity_expected_account_id(), - ) - .await; - selector - .observe_cyber_policy_success(cyber_policy_scope.as_ref()) - .await; - } - if (service_tier_changed - || timing_changed - || rate_limits_changed - || turn_state_changed - || terminal_changed) - && let Some(observation) = observation_state.observation() - { - yield ProviderEvent::observation(observation); - } - if let Some((mut failure, atomic_upstream_failure)) = terminal_failure { - let failure_after_commit = - timing_signals.semantic_output || pre_commit_events.is_committed(); - if failure_after_commit { - for event in pre_commit_events.commit(events) { - yield event; - } - } else if atomic_upstream_failure { - failure.error = failure - .error - .with_atomic_client_events(pre_commit_events.take_for_failure(events)); - } - Err(failure.error)?; - return; - } - let events = pre_commit_events.finish(events, timing_signals, completed); - for event in events { - yield event; - } - }) -} - -/// OpenAI Responses 的观测状态完全归 Provider 所有。 -/// -/// 每次原始 SSE/WS 事件推进状态后重新生成不可变 observation，Core 只负责携带、 -/// 持久化和展示这份安全快照，不解释任何 OpenAI 协议字段。 -struct OpenAiResponseObservationState { - transport: CodexBackendTransport, - diagnostics: CodexUpstreamDiagnostics, - response_metadata: CodexResponseMetadata, - metrics: CodexTransportMetrics, - websocket_pool_decision: Option, - request_summary: Value, - requested_model: String, - stream: bool, - compact: bool, - requested_service_tier: Option, - upstream_service_tier: Option, - rate_limit_headers: Vec<(String, String)>, - account_id: String, - attempt_index: u32, - timings: ProviderResponseTimings, - terminal: Option, -} - -struct OpenAiPassiveQuotaObservation { - rate_limits: Vec, -} - -impl OpenAiPassiveQuotaObservation { - fn new(headers: Vec<(String, String)>) -> Self { - Self { - rate_limits: parse_rate_limit_headers(&headers).into_iter().collect(), - } - } - - fn observe(&mut self, updates: &[ParsedRateLimits]) { - self.rate_limits.extend_from_slice(updates); - } - - fn rate_limits(&self) -> &[ParsedRateLimits] { - &self.rate_limits - } -} - -#[derive(Clone, Copy, PartialEq, Eq)] -enum OpenAiResponseTerminal { - Completed, - Incomplete, -} - -impl OpenAiResponseObservationState { - fn from_backend_response( - response: &CodexBackendStreamingResponse, - request: &CodexResponsesRequest, - account_id: &str, - attempt_index: u32, - ) -> Self { - let semantics = request.semantics(); - Self { - transport: response.transport, - diagnostics: response.diagnostics.clone(), - response_metadata: response.response_metadata.clone(), - metrics: response.transport_metrics.clone(), - websocket_pool_decision: response.websocket_pool_decision, - request_summary: openai_response_request_summary(request, response.transport), - requested_model: request.model().to_owned(), - stream: request.stream(), - compact: semantics.compact, - requested_service_tier: normalize_service_tier(request.service_tier()), - upstream_service_tier: None, - rate_limit_headers: selected_observation_headers(&response.rate_limit_headers), - account_id: account_id.to_owned(), - attempt_index, - timings: openai_response_timings( - &response.transport_metrics, - &response.response_metadata, - ), - terminal: None, - } - } - - fn observation(&self) -> Option { - let mut observation = codex_response_observation( - self.transport, - &self.diagnostics, - &self.response_metadata, - &self.metrics, - self.websocket_pool_decision, - self.timings, - )?; - if let Some(metadata) = self.provider_metadata() { - observation = observation.with_provider_metadata(metadata); - } - if let Some(service_tier) = self.effective_service_tier() { - observation = observation.with_service_tier_if_valid(service_tier.to_owned()); - } - Some(observation) - } - - fn observe_stream_chunk(&mut self, chunk: &[u8], started_at: Instant) -> bool { - if chunk.is_empty() { - return false; - } - insert_first_timing(&mut self.timings.first_event_ms, started_at) - } - - fn observe_timing_signals( - &mut self, - signals: ResponseEventSignals, - started_at: Instant, - ) -> bool { - let mut changed = false; - // 首个非前导输出事件（结构帧也算）开启首字计时； - // 真实语义首字由 first_reasoning_ms / first_text_ms 单独观测。 - if signals.output_start { - changed |= insert_first_timing(&mut self.timings.first_token_ms, started_at); - } - if signals.reasoning_output { - changed |= insert_first_timing(&mut self.timings.first_reasoning_ms, started_at); - } - if signals.text_output { - changed |= insert_first_timing(&mut self.timings.first_text_ms, started_at); - } - changed - } - - fn observe_upstream_service_tier(&mut self, service_tier: Option<&str>) -> bool { - let Some(service_tier) = normalize_service_tier(service_tier) else { - return false; - }; - if self.upstream_service_tier.as_deref() == Some(service_tier.as_str()) { - return false; - } - self.upstream_service_tier = Some(service_tier); - true - } - - fn effective_service_tier(&self) -> Option<&str> { - self.requested_service_tier - .as_deref() - .or(self.upstream_service_tier.as_deref()) - } - - fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool { - let mut changed = false; - for (name, value) in selected_observation_headers(updates) { - if let Some(existing) = self - .rate_limit_headers - .iter_mut() - .find(|(existing_name, _)| existing_name == &name) - { - if existing.1 != value { - existing.1 = value; - changed = true; - } - } else { - self.rate_limit_headers.push((name, value)); - changed = true; - } - } - changed - } - - fn merge_client_header(&mut self, name: &str, value: &str) -> bool { - let value = Bytes::copy_from_slice(value.as_bytes()); - if let Some((_, existing)) = self - .response_metadata - .client_headers - .iter_mut() - .find(|(existing_name, _)| existing_name.eq_ignore_ascii_case(name)) - { - if *existing == value { - return false; - } - *existing = value; - return true; - } - self.response_metadata - .client_headers - .push((name.to_owned(), value)); - true - } - - fn mark_completed(&mut self, incomplete: bool) -> bool { - let terminal = if incomplete { - OpenAiResponseTerminal::Incomplete - } else { - OpenAiResponseTerminal::Completed - }; - if self.terminal == Some(terminal) { - return false; - } - self.terminal = Some(terminal); - true - } - - fn provider_metadata(&self) -> Option { - let mut metadata = Map::new(); - // 观测对象版本信封：读取端按 schemaVersion 解码，缺失视为 v0。 - metadata.insert("schemaVersion".to_owned(), json!(1)); - let effective_model = self - .response_metadata - .effective_model - .as_deref() - .filter(|model| !model.is_empty()) - .unwrap_or(&self.requested_model); - if !effective_model.is_empty() { - metadata.insert( - "effectiveModel".to_owned(), - Value::String(effective_model.to_owned()), - ); - } - metadata.insert( - "modelsEtag".to_owned(), - self.response_metadata - .models_etag - .as_ref() - .map_or(Value::Null, |etag| Value::String(etag.clone())), - ); - metadata.insert( - "reasoningIncluded".to_owned(), - Value::Bool(self.response_metadata.reasoning_included), - ); - metadata.insert("stream".to_owned(), Value::Bool(self.stream)); - metadata.insert("compact".to_owned(), Value::Bool(self.compact)); - metadata.insert("attemptIndex".to_owned(), json!(self.attempt_index)); - metadata.insert( - "attemptAccountId".to_owned(), - Value::String(self.account_id.clone()), - ); - metadata.insert( - "rateLimitHeaders".to_owned(), - json!(self.rate_limit_headers), - ); - metadata.insert( - "upstreamTraceHeaders".to_owned(), - json!(selected_observation_headers( - &self.diagnostics.trace_headers - )), - ); - metadata.insert("requestSummary".to_owned(), self.request_summary.clone()); - if let Some(service_tier) = self.effective_service_tier() { - metadata.insert( - "serviceTier".to_owned(), - Value::String(service_tier.to_owned()), - ); - } - if let Some(service_tier) = &self.requested_service_tier { - metadata.insert( - "requestedServiceTier".to_owned(), - Value::String(service_tier.clone()), - ); - } - if let Some(service_tier) = &self.upstream_service_tier { - metadata.insert( - "upstreamServiceTier".to_owned(), - Value::String(service_tier.clone()), - ); - } - if let Some(status) = self.diagnostics.status_code { - metadata.insert("upstreamStatus".to_owned(), json!(status)); - } - if let Some(decision) = self.metrics.decision { - metadata.insert( - "transportDecision".to_owned(), - Value::String(decision.as_str().to_owned()), - ); - } - if let Some(decision) = self.websocket_pool_decision { - metadata.insert( - "websocketPool".to_owned(), - json!({ "kind": decision.kind() }), - ); - } - if let Some(version) = self.metrics.http_version.as_deref() { - metadata.insert("httpVersion".to_owned(), Value::String(version.to_owned())); - } - if let Some((_, cf_ray)) = self - .diagnostics - .trace_headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("cf-ray")) - && let Some((_, cf_ray)) = selected_observation_header("cf-ray", cf_ray) - { - metadata.insert("cfRay".to_owned(), Value::String(cf_ray)); - } - insert_optional_metadata_millis( - &mut metadata, - "transportDecisionWaitMs", - self.timings.transport_decision_wait_ms, - ); - insert_optional_metadata_millis(&mut metadata, "wsConnectMs", self.timings.connect_ms); - insert_optional_metadata_millis( - &mut metadata, - "upstreamHeadersMs", - self.timings.headers_ms, - ); - insert_optional_metadata_millis(&mut metadata, "firstEventMs", self.timings.first_event_ms); - insert_optional_metadata_millis( - &mut metadata, - "firstReasoningMs", - self.timings.first_reasoning_ms, - ); - insert_optional_metadata_millis(&mut metadata, "firstTextMs", self.timings.first_text_ms); - insert_optional_metadata_millis(&mut metadata, "firstTokenMs", self.timings.first_token_ms); - insert_optional_metadata_millis( - &mut metadata, - "openaiProcessingMs", - self.timings.provider_processing_ms, - ); - if let Some(terminal) = self.terminal { - let incomplete = terminal == OpenAiResponseTerminal::Incomplete; - metadata.insert("completed".to_owned(), Value::Bool(!incomplete)); - metadata.insert("incomplete".to_owned(), Value::Bool(incomplete)); - metadata.insert("eventStatusCode".to_owned(), json!(200_u16)); - } - ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?) - } -} - -fn codex_response_observation( - transport: CodexBackendTransport, - diagnostics: &CodexUpstreamDiagnostics, - response_metadata: &CodexResponseMetadata, - metrics: &CodexTransportMetrics, - websocket_pool_decision: Option, - timings: ProviderResponseTimings, -) -> Option { - let mut observation = ProviderResponseObservation::new( - UpstreamTransport::new(actual_transport_name(transport)).ok()?, - ) - .with_timings(timings); - if let Some(version) = metrics - .http_version - .as_deref() - .and_then(UpstreamHttpVersion::parse) - { - observation = observation.with_http_version(version); - } - if let Some(decision) = websocket_pool_decision { - observation = observation.with_websocket_pool(if decision.is_reuse() { - WebSocketPoolKind::Reuse - } else { - WebSocketPoolKind::New - }); - } - // WebSocket opening 的 101 是成功升级事实，不是业务请求的失败 HTTP 状态。 - // opening 明确拒绝仍由 `codex_error_observation` 保存真实上游状态。 - if transport != CodexBackendTransport::WebSocket - && let Some(status_code) = diagnostics.status_code - { - observation = observation.with_status_code(status_code); - } - if let Some(request_id) = diagnostics.request_id.as_deref() { - observation = observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); - } - let client_headers = response_metadata - .client_headers - .iter() - .map(|(name, value)| ProviderResponseHeader::new(name.to_owned(), value.clone())) - .collect(); - observation = observation.with_client_headers(client_headers); - Some(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, - None, - openai_response_timings(transport_metrics, &CodexResponseMetadata::default()), - )? - .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() { - observation = - observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); - } - } - _ => {} - } - Some(observation) -} - -fn openai_response_timings( - metrics: &CodexTransportMetrics, - response_metadata: &CodexResponseMetadata, -) -> ProviderResponseTimings { - 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), - provider_processing_ms: openai_processing_ms(response_metadata), - ..ProviderResponseTimings::default() - } -} - -fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option { - response_metadata - .client_headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms")) - .and_then(|(_, value)| std::str::from_utf8(value).ok()) - .and_then(|value| value.trim().parse::().ok()) -} - -fn openai_response_request_summary( - request: &CodexResponsesRequest, - transport: CodexBackendTransport, -) -> Value { - let body = request.body(); - let input = body.get("input"); - let tools = body.get("tools"); - let semantics = request.semantics(); - json!({ - "model": request.model(), - "stream": request.stream(), - "store": request.store(), - "compact": semantics.compact, - "requestKind": semantics.request_kind, - "subagentKind": semantics.subagent_kind, - "reasoningPreset": semantics.reasoning_preset, - "transport": actual_transport_name(transport), - "inputType": json_value_kind(input), - "inputItemsCount": input.and_then(Value::as_array).map(Vec::len), - "toolsType": json_value_kind(tools), - "toolsCount": tools.and_then(Value::as_array).map(Vec::len), - "topLevelFields": body.keys().cloned().collect::>(), - "previousResponseIdPresent": request.previous_response_id().is_some(), - "serviceTier": request.service_tier(), - "localTransport": { - "useWebsocket": request.use_websocket, - "forceHttpSse": request.force_http_sse, - }, - }) -} - -const fn json_value_kind(value: Option<&Value>) -> &'static str { - match value { - None => "missing", - Some(Value::Null) => "null", - Some(Value::Bool(_)) => "boolean", - Some(Value::Number(_)) => "number", - Some(Value::String(_)) => "string", - Some(Value::Array(_)) => "array", - Some(Value::Object(_)) => "object", - } -} - -fn selected_observation_headers(headers: &[(String, String)]) -> Vec<(String, String)> { - headers - .iter() - .filter_map(|(name, value)| selected_observation_header(name, value)) - .collect() -} - -fn selected_observation_header(name: &str, value: &str) -> Option<(String, String)> { - let name = name.trim().to_ascii_lowercase(); - (!name.is_empty()).then(|| (name, value.to_owned())) -} - -fn insert_first_timing(target: &mut Option, started_at: Instant) -> bool { - if target.is_some() { - return false; - } - let elapsed = u64::try_from(started_at.elapsed().as_millis()) - .unwrap_or(u64::MAX) - .max(1); - *target = Some(elapsed); - true -} - -fn insert_optional_metadata_millis( - metadata: &mut Map, - name: &str, - value: Option, -) { - if let Some(value) = value { - metadata.insert(name.to_owned(), json!(value)); - } -} - -async fn take_rate_limit_updates(updates: Option<&CodexRateLimitUpdates>) -> Vec { - let Some(updates) = updates else { - return Vec::new(); - }; - std::mem::take(&mut *updates.lock().await) -} - -fn rate_limit_update_headers(updates: &[ParsedRateLimits]) -> Vec<(String, String)> { - updates - .iter() - .flat_map(rate_limits_to_header_pairs) - .collect() -} - -fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool { - events - .iter() - .find_map(|event| { - let wire = event.wire_event()?; - let event_type = wire - .event_type() - .or_else(|| wire.data().get("type").and_then(Value::as_str)); - match event_type { - Some("response.incomplete") => Some(true), - Some("response.completed") => Some( - wire.data() - .pointer("/response/status") - .and_then(Value::as_str) - == Some("incomplete"), - ), - _ => None, - } - }) - .unwrap_or(false) -} - -async fn synchronize_passive_quota( - quota: &CodexCredentialQuotaService, - account: &ProviderAccount, - rate_limits: &[ParsedRateLimits], -) { - if rate_limits.is_empty() { - return; - } - if let Err(error) = quota - .synchronize_passive_rate_limits(account, rate_limits) - .await - { - tracing::warn!( - account_id = %account.id(), - error = %error, - "OpenAI passive quota synchronization failed" - ); - } -} - -async fn synchronize_passive_quota_headers( - quota: &CodexCredentialQuotaService, - account: &ProviderAccount, - headers: &[(String, String)], -) { - let Some(rate_limits) = parse_rate_limit_headers(headers) else { - return; - }; - synchronize_passive_quota(quota, account, std::slice::from_ref(&rate_limits)).await; -} - -const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str { - match transport { - CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT, - CodexBackendTransport::HttpJson => HTTP_JSON_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 { - // Catalog membership is enough to publish Generate. `supported_in_api` is advisory; - // the upstream response is authoritative for normal and diagnostic requests. - let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) - .with_upstream_feature_validation(); - ProviderModelCapabilities::new(model.request_model().clone(), capabilities) - .with_presentation(codex_model_presentation(model)) -} - -fn codex_model_presentation(model: &CodexCatalogModel) -> ModelPresentation { - let capabilities = model.capabilities(); - let reasoning_efforts = capabilities.reasoning_efforts().to_vec(); - // Codex 目录不声明默认 effort；有 medium 时对齐官方 picker 缺省，否则取首项。 - let default_reasoning = reasoning_efforts - .iter() - .find(|effort| effort.as_str() == "medium") - .or_else(|| reasoning_efforts.first()) - .cloned(); - let hidden = matches!( - model.metadata().visibility(), - Some(CodexCatalogVisibility::Hide | CodexCatalogVisibility::None) - ); - - ModelPresentation::new( - Some(model.display_name().to_owned()), - model.metadata().description().map(str::to_owned), - ) - .with_reasoning(default_reasoning, reasoning_efforts) - .with_context_window_tokens( - model - .limits() - .context_window_tokens() - .map(std::num::NonZeroU64::get), - ) - .with_image_input(capabilities.image_input() == CodexCatalogCapabilityEvidence::DeclaredNative) - // Codex Responses 工具协议随 API 支持一并可用；只有明确不支持才关闭。 - .with_agent_tools( - capabilities.responses_api() != CodexCatalogCapabilityEvidence::DeclaredUnsupported, - capabilities.parallel_tool_calls() == CodexCatalogCapabilityEvidence::DeclaredNative, - ) - .with_search_tool(capabilities.web_search() == CodexCatalogCapabilityEvidence::DeclaredNative) - .with_image_detail_original( - capabilities.image_detail_original() == CodexCatalogCapabilityEvidence::DeclaredNative, - ) - .with_verbosity(capabilities.verbosity() == CodexCatalogCapabilityEvidence::DeclaredNative) - .with_hidden(hidden) -} - -fn selected_transport(request: &CodexResponsesRequest) -> CodexProviderTransport { - if request.force_http_sse { - CodexProviderTransport::HttpOnly - } else { - CodexProviderTransport::PreferWebSocket - } -} - -fn apply_transport(request: &mut CodexResponsesRequest, transport: CodexProviderTransport) { - 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::HttpOnly => HTTP_SSE_TRANSPORT, - CodexProviderTransport::PreferWebSocket => WEBSOCKET_TRANSPORT, - } -} - -const fn accepts_backend_transport( - transport: CodexProviderTransport, - actual: CodexBackendTransport, -) -> bool { - match transport { - CodexProviderTransport::HttpOnly => matches!(actual, CodexBackendTransport::HttpSse), - CodexProviderTransport::PreferWebSocket => true, - } -} - -fn codex_request_context<'a>( - request: &'a CodexResponsesRequest, - request_id: &'a str, - account: &'a ProviderAccount, - installation_id: &'a str, - authorization: &'a SecretString, - cookie_header: Option<&'a SecretString>, - account_selection: CodexAccountSelectionTelemetry<'a>, -) -> CodexRequestContext<'a> { - CodexRequestContext { - authorization: authorization.expose_secret(), - account_id: account.upstream_account_id(), - request_id, - turn_state: request.turn_state.as_deref(), - turn_metadata: request.turn_metadata.as_deref(), - beta_features: request.beta_features.as_deref(), - include_timing_metrics: request.include_timing_metrics.as_deref(), - version: request.version.as_deref(), - codex_window_id: request.codex_window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.map(ExposeSecret::expose_secret), - installation_id: Some(installation_id), - session_id: request.client_session_id.as_deref(), - thread_id: request.client_thread_id.as_deref(), - client_request_id: request.client_request_id.as_deref(), - turn_id: request.client_turn_id.as_deref(), - account_selection, - } -} - -fn build_cookie_header( - cookies: &[RuntimeCodexCookie], -) -> Result, ProviderError> { - if cookies.is_empty() { - return Ok(None); - } - let mut header = String::new(); - for cookie in cookies { - let value = cookie.value.expose_secret(); - if !valid_cookie_name(&cookie.name) - || value.is_empty() - || value.chars().any(char::is_control) - || value.contains(';') - { - return Err(provider_error( - ProviderErrorKind::Protocol, - UpstreamSendState::NotSent, - )); - } - if !header.is_empty() { - header.push_str("; "); - } - header.push_str(&cookie.name); - header.push('='); - header.push_str(value); - if header.len() > MAX_COOKIE_HEADER_BYTES { - return Err(provider_error( - ProviderErrorKind::Protocol, - UpstreamSendState::NotSent, - )); - } - } - Ok(Some(SecretString::from(header))) -} - -fn valid_cookie_name(name: &str) -> bool { - !name.is_empty() - && name.len() <= 256 - && name.bytes().all(|byte| { - byte.is_ascii_alphanumeric() - || matches!( - byte, - b'!' | b'#' - | b'$' - | b'%' - | b'&' - | b'\'' - | b'*' - | b'+' - | b'-' - | b'.' - | b'^' - | b'_' - | b'`' - | b'|' - | b'~' - ) - }) -} - -fn map_request_error(error: CodexRequestEncodeError) -> ProviderError { - let kind = match error { - CodexRequestEncodeError::InvalidProtocolPayload => ProviderErrorKind::InvalidRequest, - }; - provider_error(kind, UpstreamSendState::NotSent) -} - -fn map_selection_error(error: CredentialSelectionError) -> ProviderError { - match error { - CredentialSelectionError::CapacityUnavailable { retry_after } => { - let error = provider_error( - ProviderErrorKind::AccountCapacityUnavailable, - UpstreamSendState::NotSent, - ); - match retry_after { - Some(retry) => error.with_retry_after(retry), - None => error, - } - } - CredentialSelectionError::NoEligibleCredential => provider_error( - ProviderErrorKind::NoEligibleAccount, - UpstreamSendState::NotSent, - ), - CredentialSelectionError::InvalidCredential - | CredentialSelectionError::Store - | CredentialSelectionError::Coordinator - | CredentialSelectionError::CookiePolicy => provider_error( - ProviderErrorKind::ProviderInfrastructureUnavailable, - UpstreamSendState::NotSent, - ), - } -} - -struct MappedProviderFailure { - error: ProviderError, - websocket_transport_retryable: bool, - account_failure: Option, - /// 原始上游错误描述，仅在凭据错误状态下持久化。 - error_message: Option, - cyber_policy_failure: bool, - set_cookie_headers: Vec, - rate_limit_headers: Vec<(String, String)>, - observation: Option, - capture_response_cookies: bool, -} - -impl MappedProviderFailure { - fn plain(error: ProviderError) -> Self { - Self { - error, - websocket_transport_retryable: false, - account_failure: None, - error_message: None, - cyber_policy_failure: false, - set_cookie_headers: Vec::new(), - rate_limit_headers: Vec::new(), - observation: None, - capture_response_cookies: false, - } - } -} - -enum PreCommitPoll { - Upstream(T), - GraceElapsed, -} - -async fn wait_for_replay_grace(deadline: Option) { - if let Some(deadline) = deadline { - tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; - } else { - std::future::pending::<()>().await; - } -} - -/// 提交边界前的上游事件预取。 -/// -/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号 -/// 窗口；任一边界到达都会提交已缓存 wire，不能因网关私有资源规则伪造上游协议 -/// 失败。一旦提交，后续事件不再具备无痕重放资格。 -struct PreCommitClientEvents { - pending: Vec, - prefetched_bytes: usize, - replay_grace_deadline: Option, - committed: bool, -} - -impl PreCommitClientEvents { - const fn new() -> Self { - Self { - pending: Vec::new(), - prefetched_bytes: 0, - replay_grace_deadline: None, - committed: false, - } - } - - fn observe_chunk(&mut self, bytes: usize) { - if !self.committed { - self.prefetched_bytes = self.prefetched_bytes.saturating_add(bytes); - } - } - - fn stage( - &mut self, - incoming: Vec, - timing_signals: ResponseEventSignals, - completed: bool, - ) -> Vec { - if self.committed { - return incoming; - } - let starts_replay_grace = incoming.iter().any(ProviderEvent::has_client_event); - self.pending.extend(incoming); - if timing_signals.semantic_output - || completed - || self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES - { - return self.commit_pending(); - } - if starts_replay_grace && self.replay_grace_deadline.is_none() { - self.replay_grace_deadline = Instant::now().checked_add(STREAM_REPLAY_GRACE); - } - Vec::new() - } - - fn finish( - &mut self, - incoming: Vec, - timing_signals: ResponseEventSignals, - completed: bool, - ) -> Vec { - let events = self.stage(incoming, timing_signals, completed); - if self.committed { - events - } else { - self.commit_pending() - } - } - - fn commit(&mut self, incoming: Vec) -> Vec { - if self.committed { - return incoming; - } - self.pending.extend(incoming); - self.commit_pending() - } - - fn take_for_failure(&mut self, incoming: Vec) -> Vec { - self.pending.extend(incoming); - self.prefetched_bytes = 0; - self.replay_grace_deadline = None; - std::mem::take(&mut self.pending) - } - - const fn replay_grace_deadline(&self) -> Option { - self.replay_grace_deadline - } - - const fn is_committed(&self) -> bool { - self.committed - } - - fn commit_pending(&mut self) -> Vec { - self.committed = true; - self.prefetched_bytes = 0; - self.replay_grace_deadline = None; - std::mem::take(&mut self.pending) - } -} - -#[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) - } -} - -struct OpenAiFailureContext<'a> { - client: &'a CodexBackendClient, - selector: &'a CodexCredentialSelector, - quota: &'a Arc, - response_origin: &'a Url, - cyber_policy_scope: Option<&'a CodexCyberPolicyScope>, - allows_account_state_mutation: bool, -} - -#[derive(Clone, Copy)] -struct UpstreamErrorLogContext<'a> { - request_id: &'a str, - account_id: &'a str, - attempt_index: u32, - websocket_connection_id: Option, -} - -impl<'a> UpstreamErrorLogContext<'a> { - fn new( - context: &'a AttemptContext, - account: &'a ProviderAccount, - websocket_connection_id: Option, - ) -> Self { - Self { - request_id: context.request_id().as_str(), - account_id: account.id().as_str(), - attempt_index: context.attempt_index().get(), - websocket_connection_id, - } - } - - fn with_websocket_connection_id(mut self, connection_id: Option) -> Self { - self.websocket_connection_id = self.websocket_connection_id.or(connection_id); - self - } -} - -fn log_client_upstream_error(context: UpstreamErrorLogContext<'_>, error: &CodexClientError) { - match error { - CodexClientError::Upstream { - status, - body, - transport, - .. - } => log_raw_upstream_body( - context, - *transport, - "http_error_response", - Some(status.as_u16()), - None, - None, - body, - ), - CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => { - log_raw_upstream_body( - context, - CodexBackendTransport::WebSocket, - "websocket_opening_response", - Some(upstream.status_code), - None, - None, - &upstream.body, - ); - } - CodexClientError::WebSocket(error) => { - let Some(close) = error.close_before_terminal() else { - return; - }; - let context = context.with_websocket_connection_id(close.connection_id()); - let websocket_connection_id = context - .websocket_connection_id - .map(|connection_id| connection_id.to_string()) - .unwrap_or_default(); - let close_reason_bytes = close.reason().map_or(0, str::len); - tracing::warn!( - request_id = %context.request_id, - account_id = %context.account_id, - attempt_index = context.attempt_index, - websocket_connection_id, - upstream_transport = WEBSOCKET_TRANSPORT, - upstream_error_kind = "websocket_close", - upstream_close_code = close.code().unwrap_or_default(), - upstream_close_code_present = close.code().is_some(), - upstream_last_event_type = close.last_event_type().unwrap_or_default(), - upstream_last_event_type_present = close.last_event_type().is_some(), - upstream_terminal_seen = false, - upstream_close_reason_bytes = close_reason_bytes, - upstream_close_reason_present = close.reason().is_some(), - upstream_error_raw_present = close.reason().is_some(), - "OpenAI upstream WebSocket closed before terminal event" - ); - } - _ => {} - } -} - -fn log_canonical_upstream_error( - context: UpstreamErrorLogContext<'_>, - transport: CodexBackendTransport, - error: &CodexCanonicalError, -) { - let CodexCanonicalError::Upstream(failure) = error else { - return; - }; - log_raw_upstream_body( - context, - transport, - "responses_error_event", - failure.explicit_status_code, - failure.upstream_code.as_deref(), - failure.upstream_type.as_deref(), - failure.raw_body(), - ); -} - -fn log_raw_upstream_body( - context: UpstreamErrorLogContext<'_>, - transport: CodexBackendTransport, - error_kind: &'static str, - status_code: Option, - upstream_code: Option<&str>, - upstream_type: Option<&str>, - upstream_error_raw: &str, -) { - let websocket_connection_id = context - .websocket_connection_id - .map(|connection_id| connection_id.to_string()) - .unwrap_or_default(); - tracing::warn!( - request_id = %context.request_id, - account_id = %context.account_id, - attempt_index = context.attempt_index, - websocket_connection_id, - upstream_transport = actual_transport_name(transport), - upstream_error_kind = error_kind, - upstream_status_code = status_code.unwrap_or_default(), - upstream_status_code_present = status_code.is_some(), - upstream_error_code = upstream_code.unwrap_or_default(), - upstream_error_code_present = upstream_code.is_some(), - upstream_error_type = upstream_type.unwrap_or_default(), - upstream_error_type_present = upstream_type.is_some(), - upstream_error_raw, - upstream_error_raw_bytes = upstream_error_raw.len(), - "OpenAI upstream returned an error payload" - ); -} - -async fn apply_failure( - context: &OpenAiFailureContext<'_>, - account: &ProviderAccount, - failure: &MappedProviderFailure, -) { - if !context.allows_account_state_mutation { - return; - } - synchronize_passive_quota_headers(context.quota, account, &failure.rate_limit_headers).await; - let needs_authoritative_quota_refresh = matches!( - failure.account_failure, - Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. }) - ); - if failure.cyber_policy_failure { - context - .selector - .record_cyber_policy_failure(context.cyber_policy_scope, account) - .await; - } - if let Some(account_failure) = failure.account_failure { - context - .client - .evict_websocket_account(account.id().as_str()) - .await; - if let Err(error) = context - .selector - .record_failure(account, account_failure, failure.error_message.clone()) - .await - { - tracing::warn!( - account_id = %account.id(), - error = %error, - "Failed to persist OpenAI account failure state" - ); - } - } - if needs_authoritative_quota_refresh { - schedule_authoritative_quota_refresh_after_failure(context.quota, account); - } - if failure.capture_response_cookies - && !failure.set_cookie_headers.is_empty() - && let Err(error) = context - .selector - .capture_response_cookies( - account, - context.response_origin, - &failure.set_cookie_headers, - ) - .await - { - tracing::warn!( - account_id = %account.id(), - error = %error, - "Failed to persist OpenAI response cookies" - ); - } -} - -fn schedule_authoritative_quota_refresh_after_failure( - quota: &Arc, - account: &ProviderAccount, -) { - // 限额错误可以确认账号状态，但 Responses 事件中的 used_percent 可能仍停在上一结算点。 - // usage 快照在后台补齐展示基线，不得撤销同一轮真实失败，也不能阻塞原始响应。 - // 先等待 2 秒让上游结算，再查询；5 秒仅限制实际查询自身。 - let quota = Arc::clone(quota); - let account_id = account.id().clone(); - drop(tokio::spawn(async move { - tokio::time::sleep(QUOTA_FAILURE_REFRESH_DELAY).await; - match tokio::time::timeout( - QUOTA_FAILURE_REFRESH_TIMEOUT, - quota.refresh_account_after_failure(&account_id), - ) - .await - { - Ok(Ok(_)) => {} - Ok(Err(error)) => { - tracing::warn!( - account_id = %account_id, - error = %error, - "OpenAI authoritative quota synchronization after quota rejection failed" - ); - } - Err(_) => { - tracing::warn!( - account_id = %account_id, - timeout_ms = QUOTA_FAILURE_REFRESH_TIMEOUT.as_millis(), - "OpenAI authoritative quota synchronization after quota rejection timed out" - ); - } - } - })); -} - -fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { - map_client_error(error, UpstreamSendState::Ambiguous, true) -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum WebSocketFailurePolicy { - Budgeted, - ImmediateFallback, -} - -struct WebSocketRecoveryContext<'a> { - policy: WebSocketFailurePolicy, - requirement: TransportRequirement, - retry_count: u32, - max_retries: u32, - request_id: &'a str, - attempt_index: u32, - account_id: &'a str, - session_affinity_key: Option<&'a ProviderSessionAffinityKey>, - session_affinity_key_hash: Option<&'a str>, - session_transport_fallbacks: &'a CodexSessionTransportFallbacks, -} - -fn websocket_client_failure_policy(error: &CodexClientError) -> Option { - match error { - CodexClientError::Upstream { - status, - transport: CodexBackendTransport::WebSocket, - .. - } if *status == reqwest::StatusCode::UPGRADE_REQUIRED => { - Some(WebSocketFailurePolicy::ImmediateFallback) - } - CodexClientError::Upstream { - status, - transport: CodexBackendTransport::WebSocket, - .. - } if status.is_server_error() - || matches!( - *status, - reqwest::StatusCode::REQUEST_TIMEOUT - | reqwest::StatusCode::CONFLICT - | reqwest::StatusCode::TOO_EARLY - | reqwest::StatusCode::TOO_MANY_REQUESTS - ) => - { - Some(WebSocketFailurePolicy::Budgeted) - } - CodexClientError::WebSocket(error) - if !matches!(error, CodexWebSocketExchangeError::InvalidRequest(_)) => - { - Some(WebSocketFailurePolicy::Budgeted) - } - _ => None, - } -} - -fn apply_websocket_recovery_policy( - failure: &mut MappedProviderFailure, - context: WebSocketRecoveryContext<'_>, -) { - let recovery_is_safe = match failure.error.send_state() { - UpstreamSendState::NotSent => context.requirement.allows_pre_send_http_fallback(), - UpstreamSendState::Sent | UpstreamSendState::Ambiguous => { - context.requirement.allows_pre_delivery_http_fallback() - } - }; - if !recovery_is_safe { - return; - } - - // 传输恢复必须保持原账号可调度；最终 HTTP attempt 若仍失败，再按真实 HTTP - // 结果更新账号健康度，避免中间 WS 错误把同账号钉选提前冷却掉。 - failure.account_failure = None; - let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback - || context.retry_count >= context.max_retries; - if !fallback_now { - let Some(retry_index) = context.retry_count.checked_add(1).and_then(NonZeroU32::new) else { - return; - }; - let delay = failure - .error - .retry_after() - .unwrap_or_else(|| websocket_retry_backoff(retry_index)); - failure - .error - .set_pre_delivery_transport_retry(retry_index, delay); - tracing::warn!( - request_id = context.request_id, - attempt_index = context.attempt_index, - account_id = context.account_id, - websocket_failure_kind = failure.error.kind().as_str(), - websocket_failure_code = failure - .error - .upstream_code() - .map_or("", OpaqueUpstreamValue::as_str), - upstream_status_code = failure.error.upstream_status().unwrap_or_default(), - upstream_status_code_present = failure.error.upstream_status().is_some(), - websocket_retry_count = retry_index.get(), - websocket_max_retries = context.max_retries, - websocket_retry_delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX), - transport_requirement = context.requirement.as_str(), - "OpenAI upstream WebSocket failed; retrying the same account" - ); - return; - } - - failure.error.set_pre_delivery_transport_fallback(); - let sticky_state_transition = context - .session_affinity_key - .map(|key| context.session_transport_fallbacks.disable_websocket(key)); - tracing::warn!( - request_id = context.request_id, - attempt_index = context.attempt_index, - account_id = context.account_id, - websocket_failure_kind = failure.error.kind().as_str(), - websocket_failure_code = failure - .error - .upstream_code() - .map_or("", OpaqueUpstreamValue::as_str), - upstream_status_code = failure.error.upstream_status().unwrap_or_default(), - upstream_status_code_present = failure.error.upstream_status().is_some(), - websocket_retry_count = context.retry_count, - websocket_max_retries = context.max_retries, - websocket_fallback_reason = match context.policy { - WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted", - WebSocketFailurePolicy::ImmediateFallback => "upgrade_required", - }, - session_affinity_present = context.session_affinity_key.is_some(), - session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), - sticky_http_enabled = sticky_state_transition.is_some(), - sticky_state_transition = sticky_state_transition.unwrap_or(false), - "OpenAI upstream WebSocket disabled for the current session" - ); -} - -fn websocket_retry_backoff(retry_index: NonZeroU32) -> Duration { - const INITIAL_DELAY_MS: u64 = 200; - - let shift = retry_index.get().saturating_sub(1).min(63); - let base_ms = INITIAL_DELAY_MS.saturating_mul(1_u64.checked_shl(shift).unwrap_or(u64::MAX)); - let mut random = [0_u8; 2]; - let jitter_per_mille = if getrandom::fill(&mut random).is_ok() { - 900_u64 + u64::from(u16::from_le_bytes(random) % 200) - } else { - 1_000 - }; - Duration::from_millis(base_ms.saturating_mul(jitter_per_mille) / 1_000) -} - -fn continuation_replay_required_error() -> ProviderError { - ProviderError::new( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - ) - .with_continuation_failure(ContinuationFailure::HistoryUnavailable) - .with_client_visible_upstream_error(ClientVisibleUpstreamError::new( - PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, - Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()), - Some("invalid_request_error".to_owned()), - )) -} - -fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { - let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error); - let websocket_failure = error.transport() == Some(CodexBackendTransport::WebSocket); - let mut failure = map_client_error(error, UpstreamSendState::Sent, false); - failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure; - if allows_pre_delivery_retry && !websocket_failure { - failure.error = failure.error.with_pre_delivery_retry(); - } - failure -} - -fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool { - match error { - CodexClientError::Http(_) - | CodexClientError::HttpJson(_) - | CodexClientError::StreamIdleTimeout { .. } => true, - CodexClientError::WebSocket(error) => !matches!( - error, - CodexWebSocketExchangeError::InvalidRequest(_) - | CodexWebSocketExchangeError::Upstream(_) - | CodexWebSocketExchangeError::ContinuationUnavailable { .. } - ), - _ => 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::RequestBodyEncode(_) - | CodexClientError::RequestCompression(_) - | CodexClientError::ModelCatalog(_) - | 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) | CodexClientError::HttpJson(error) => { - let send_state = if error.is_connect() { - UpstreamSendState::NotSent - } else { - uncertain_state - }; - MappedProviderFailure::plain(provider_error( - if error.is_timeout() { - ProviderErrorKind::Timeout - } else { - ProviderErrorKind::Transport - }, - send_state, - )) - } - CodexClientError::WebSocket(error) => { - let close_code = error.close_before_terminal().and_then(|close| close.code()); - let client_visible_error = websocket_client_visible_error(&error); - let mut failure = MappedProviderFailure::plain(provider_error( - websocket_error_kind(&error), - websocket_send_state(&error), - )); - if let Some(close_code) = close_code { - failure.error = - failure - .error - .with_upstream_code(OpaqueUpstreamValue::new(format!( - "websocket_close_{close_code}" - ))); - } - if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { - failure.error = failure.error.with_upstream_code(OpaqueUpstreamValue::new( - WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(), - )); - } - if let Some(client_visible_error) = client_visible_error { - failure.error = failure - .error - .with_client_visible_upstream_error(client_visible_error); - } - failure - } - }; - if let Some(continuation_failure) = continuation_failure { - failure.error = failure - .error - .with_continuation_failure(continuation_failure); - } - failure.observation = observation; - failure -} - -fn websocket_client_visible_error( - error: &CodexWebSocketExchangeError, -) -> Option { - if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { - return Some(ClientVisibleUpstreamError::new( - "websocket connection limit reached", - Some(WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned()), - Some("websocket_error".to_owned()), - )); - } - let close = error.close_before_terminal()?; - let message = close - .reason() - .filter(|reason| !reason.is_empty()) - .map_or_else(|| close.to_string(), str::to_owned); - Some(ClientVisibleUpstreamError::new( - message, - close.code().map(|code| code.to_string()), - Some("websocket_close_error".to_owned()), - )) -} - -fn map_upstream_failure( - mut failure: CodexUpstreamFailure, - observation: Option, - replay_boundary: ReplayBoundary, -) -> MappedProviderFailure { - let category = failure.category(); - let cyber_policy_failure = failure - .status - .is_some_and(|status| status.is_client_error()) - && is_cyber_policy_code(failure.code.as_deref()); - 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); - if let Some(message) = failure.client_message.as_ref() { - error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( - message.clone(), - failure.client_code.clone(), - failure.client_error_type.clone(), - )); - } - if let Some(response) = failure.client_response.take() { - let response = (*response).into_parts(); - error = error.with_client_visible_upstream_response( - ClientVisibleUpstreamResponse::new( - response.status, - response.content_type, - response.body, - ) - .with_headers( - response - .client_headers - .into_iter() - .map(|(name, value)| ProviderResponseHeader::new(name, value)) - .collect(), - ), - ); - } - if let Some(status) = failure.status { - error = error.with_status(status.as_u16()); - } - if replay_boundary.permits_provider_proof() - && (failure.replay_is_safe() || cyber_policy_failure) - { - 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() { - error = error.with_upstream_code(OpaqueUpstreamValue::new(code.to_owned())); - } - if let Some(request_id) = failure.request_id.as_deref() { - error = error.with_upstream_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); - } - MappedProviderFailure { - error, - websocket_transport_retryable: false, - account_failure: account_failure( - category, - failure.retry_after_seconds, - failure.usage_limit_resets_at, - ), - error_message: failure.client_message, - cyber_policy_failure, - set_cookie_headers: failure.set_cookie_headers, - rate_limit_headers: failure.rate_limit_headers, - observation, - capture_response_cookies: !matches!( - category, - CodexFailureCategory::CloudflareChallenge | CodexFailureCategory::CloudflarePathBlocked - ), - } -} - -fn is_cyber_policy_code(code: Option<&str>) -> bool { - code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) -} - -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::UsageLimitExhausted => ProviderErrorKind::QuotaExhausted, - 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, - usage_limit_resets_at: Option, -) -> Option { - match category { - CodexFailureCategory::CredentialExpired => Some(CodexAccountFailure::CredentialExpired), - CodexFailureCategory::IdentityVerificationRequired => { - Some(CodexAccountFailure::IdentityVerificationRequired) - } - CodexFailureCategory::Banned => Some(CodexAccountFailure::Banned), - CodexFailureCategory::UsageLimitExhausted => { - Some(CodexAccountFailure::UsageLimitExhausted { - reset_at: usage_limit_resets_at - .and_then(|seconds| u64::try_from(seconds).ok()) - .and_then(|seconds| { - SystemTime::UNIX_EPOCH.checked_add(Duration::from_secs(seconds)) - }), - }) - } - 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::ConnectionLimitReached - | CodexWebSocketExchangeError::InvalidSse(_) - | 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::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::ConnectionLimitReached => ProviderErrorKind::RateLimited, - CodexWebSocketExchangeError::Transport(_) - | CodexWebSocketExchangeError::Connect(_) - | CodexWebSocketExchangeError::PostSendAmbiguous { .. } - | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) - | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { - ProviderErrorKind::Transport - } - } -} - -fn provider_error(kind: ProviderErrorKind, send_state: UpstreamSendState) -> ProviderError { - ProviderError::new(kind, send_state) -} - -fn remaining(deadline: SystemTime) -> Option { - deadline - .duration_since(SystemTime::now()) - .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 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, - quota_refresh_policy: CodexQuotaRefreshPolicy, - oauth_refresh_enabled: bool, - desktop_release: Arc, -) -> 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)?; - let mut contributions = Vec::new(); - if oauth_refresh_enabled { - contributions.push(WorkerContribution::Registration(scheduled_registration( - refresh_id, - OAUTH_REFRESH_INTERVAL, - Box::new(OpenAiOAuthRefreshTask { service: refresh }), - )?)); - } - contributions.extend([ - WorkerContribution::Registration(scheduled_registration( - quota_id, - quota_refresh_policy.interval(), - Box::new(OpenAiQuotaTask { - quota, - catalog: Arc::clone(&catalog), - }), - )?), - WorkerContribution::Registration(WorkerRegistration::try_new( - etag_id, - WorkerRunnable::Daemon { - restart: DaemonRestartPolicy::try_new( - WORKER_INITIAL_BACKOFF, - WORKER_MAXIMUM_BACKOFF, - )?, - task: Box::new(OpenAiCatalogEtagTask { catalog }), - }, - )?), - WorkerContribution::Registration(scheduled_registration( - desktop_release_id, - APPCAST_POLL_INTERVAL, - Box::new(OpenAiDesktopReleaseTask { - service: desktop_release, - }), - )?), - ]); - Ok(contributions) -} - -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(|error| { - tracing::error!(error = %error, "OpenAI OAuth refresh cycle failed"); - WorkerTaskError::safe("OpenAI OAuth refresh failed") - })?; - let mut refreshed = 0_u64; - let mut invalidated = 0_u64; - let mut banned = 0_u64; - let mut transient = 0_u64; - let mut lease_unavailable = 0_u64; - let mut stale = 0_u64; - let mut failed = 0_u64; - let mut transient_accounts = Vec::new(); - let mut failed_accounts = Vec::new(); - for outcome in &outcomes { - match outcome { - CodexCredentialRefreshOutcome::Refreshed { .. } => refreshed += 1, - CodexCredentialRefreshOutcome::Invalidated { .. } => invalidated += 1, - CodexCredentialRefreshOutcome::Banned { .. } => banned += 1, - CodexCredentialRefreshOutcome::Transient { account_id } => { - transient += 1; - transient_accounts.push(account_id); - } - CodexCredentialRefreshOutcome::LeaseUnavailable { .. } => { - lease_unavailable += 1; - } - CodexCredentialRefreshOutcome::Stale { .. } => stale += 1, - CodexCredentialRefreshOutcome::Failed { account_id } => { - failed += 1; - failed_accounts.push(account_id); - } - } - } - if !outcomes.is_empty() { - tracing::info!( - refreshed, - invalidated, - banned, - transient, - lease_unavailable, - stale, - failed, - "OpenAI OAuth refresh cycle completed" - ); - } - if transient > 0 || failed > 0 { - tracing::warn!( - refreshed, - invalidated, - banned, - transient, - lease_unavailable, - stale, - failed, - transient_accounts = ?transient_accounts, - failed_accounts = ?failed_accounts, - "OpenAI OAuth refresh cycle contained operational failures" - ); - } - Ok(()) - }) - } -} - -struct OpenAiQuotaTask { - quota: Arc, - catalog: Arc, -} - -struct OpenAiCatalogEtagTask { - catalog: Arc, -} - -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 { - if context.cancellation().is_cancelled() { - return Ok(()); - } - let mut failures = false; - match self.quota.synchronize().await { - Ok(summary) if summary.has_operational_failures() => { - tracing::warn!( - updated = summary.updated, - exhausted = summary.exhausted, - banned = summary.banned, - transient = summary.transient, - stale = summary.stale, - "OpenAI quota cycle contained operational failures" - ); - } - Ok(_) => {} - Err(error) => { - failures = true; - tracing::warn!( - error = %error, - "OpenAI quota synchronization failed" - ); - } - } - if context.cancellation().is_cancelled() { - return Ok(()); - } - match self.catalog.refresh_catalogs().await { - Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {} - Err(error) => { - failures = true; - tracing::warn!(error = %error, "OpenAI model catalog refresh failed"); - } - } - if failures { - Err(WorkerTaskError::safe( - "OpenAI quota or catalog synchronization failed", - )) - } else { - Ok(()) - } - }) - } -} - -impl DaemonTask for OpenAiCatalogEtagTask { - fn run( - &self, - cancellation: gateway_core::lifecycle::CancellationToken, - ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - loop { - tokio::select! { - () = cancellation.cancelled() => return Ok(()), - () = self.catalog.wait_for_etag_refresh() => {}, - }; - if let Err(error) = self.catalog.refresh().await { - tracing::warn!( - error = %error, - "OpenAI model catalog ETag refresh failed" - ); - } - } - }) - } -} diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs new file mode 100644 index 000000000..40266094f --- /dev/null +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -0,0 +1,916 @@ +//! OpenAI attempt 的选择、发送与响应流执行。 + +use super::*; + +impl CodexProvider { + pub(super) async fn execute_image( + &self, + image: &ImageRequest, + candidate: &ProviderCandidate, + context: AttemptContext, + ) -> Result { + if image.payload().protocol() != PROVIDER_NAME || candidate.upstream_model().is_some() { + return Err(provider_error( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + )); + } + let (endpoint_path, response_origin) = match image.kind() { + ImageRequestKind::Generation => ( + CODEX_IMAGE_GENERATIONS_PATH, + self.image_generations_url.clone(), + ), + ImageRequestKind::Edit => (CODEX_IMAGE_EDITS_PATH, self.image_edits_url.clone()), + }; + let selection_started_at = Instant::now(); + let lease = self + .selector + .select_for_provider_endpoint(&SelectCodexProviderEndpointCredential { + request_url: &response_origin, + attempt: &context, + }) + .await + .map_err(map_selection_error)?; + let account_selection_wait_ms = + u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + let lease = Arc::new(lease); + let allows_account_state_mutation = lease.allows_account_state_mutation(); + let provider_kind = ProviderKind::new(PROVIDER_NAME) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; + let metadata = ProviderCallMetadata::for_provider_endpoint( + provider_kind, + lease.account_id().clone(), + UpstreamTransport::new(HTTP_JSON_TRANSPORT).map_err(|_| { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) + })?, + ) + .with_selection_observation(ProviderSelectionObservation::new( + account_selection_wait_ms, + lease.capacity_snapshot(), + )); + let image_turn_id = image + .payload() + .context() + .get("image_turn_id") + .and_then(Value::as_str) + .map(str::to_owned); + let events = cold_json_response_stream(ColdJsonResponse { + client: self.client.clone(), + response_origin, + endpoint_path, + body: image.payload().body().clone(), + image_turn_id, + context, + selector: Arc::clone(&self.selector), + quota: Arc::clone(&self.quota), + lease: Arc::clone(&lease), + output_started_at: Instant::now(), + }); + let stream = ProviderStream::new(metadata, events, lease); + Ok(if allows_account_state_mutation { + stream.with_account_feedback(Arc::clone(&self.account_feedback)) + } else { + stream + }) + } +} + +pub(super) struct ColdResponse { + pub(super) client: CodexBackendClient, + pub(super) response_origin: Url, + pub(super) request: CodexResponsesRequest, + pub(super) upstream_model: UpstreamModelId, + pub(super) transport_policy: CodexProviderTransport, + pub(super) context: AttemptContext, + pub(super) selector: Arc, + pub(super) quota: Arc, + pub(super) catalog: Arc, + pub(super) lease: Arc, + pub(super) output_started_at: Instant, + pub(super) session_affinity_key: Option, + pub(super) session_affinity_key_hash: Option, + pub(super) session_transport_fallbacks: CodexSessionTransportFallbacks, + pub(super) websocket_retry_count: u32, + pub(super) stream_max_retries: u32, + pub(super) session_capture: Option, +} + +pub(super) struct ColdJsonResponse { + pub(super) client: CodexBackendClient, + pub(super) response_origin: Url, + pub(super) endpoint_path: &'static str, + pub(super) body: Bytes, + pub(super) image_turn_id: Option, + pub(super) context: AttemptContext, + pub(super) selector: Arc, + pub(super) quota: Arc, + pub(super) lease: Arc, + pub(super) output_started_at: Instant, +} + +#[derive(Clone, Serialize, Deserialize)] +pub(super) struct OpenAiSessionState { + pub(super) account_id: String, + pub(super) conversation_id: Option, + #[serde(default)] + pub(super) turn_state: Option, + #[serde(default)] + pub(super) client_turn_id: Option, + pub(super) continuation_scope: OpenAiContinuationScope, +} + +#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum OpenAiContinuationScope { + Persisted, + ConnectionLocal, + ReplayRequired, +} + +pub(super) struct OpenAiSessionCapture { + pub(super) account_id: String, + pub(super) conversation_id: Option, + pub(super) turn_state: Option, + pub(super) client_turn_id: Option, + pub(super) response_store: bool, + pub(super) continuation_scope: Option, +} + +pub(super) fn same_client_turn(previous: Option<&str>, current: Option<&str>) -> bool { + previous + .zip(current) + .is_some_and(|(previous, current)| !previous.is_empty() && previous == current) +} + +pub(super) fn decode_openai_session_state(request: &GenerateRequest) -> Option { + request + .provider_session_state(PROVIDER_NAME) + .and_then(|state| serde_json::from_value(Value::Object(state.payload().clone())).ok()) +} + +pub(super) 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)) +} + +pub(super) fn attach_openai_session_update( + events: &mut [ProviderEvent], + capture: &mut Option, +) { + let Some(terminal_index) = events + .iter() + .position(|event| terminal_response_output(event).is_some()) + else { + return; + }; + let Some(capture) = capture.take() else { + return; + }; + let Some(continuation_scope) = capture.continuation_scope else { + return; + }; + let Ok(update) = encode_openai_session_state(OpenAiSessionState { + account_id: capture.account_id, + conversation_id: capture.conversation_id, + turn_state: capture.turn_state, + client_turn_id: capture.client_turn_id, + continuation_scope, + }) else { + return; + }; + events[terminal_index].attach_session_update(update); +} + +pub(super) 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() +} + +pub(super) enum CodexHandshakeAttemptError { + Client(CodexClientError), + Cancelled, + Timeout, +} + +pub(super) async fn create_response_attempt( + client: &CodexBackendClient, + request: &CodexResponsesRequest, + request_context: CodexRequestContext<'_>, + account_id: &str, + deadline: SystemTime, + cancellation: &CancellationToken, +) -> Result { + let Some(handshake_deadline) = remaining(deadline) else { + return Err(CodexHandshakeAttemptError::Timeout); + }; + tokio::select! { + biased; + _ = cancellation.cancelled() => Err(CodexHandshakeAttemptError::Cancelled), + _ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout), + response = client.create_response_stream_with_deferred_websocket_recovery( + request, + request_context, + Some(account_id), + ) => response.map_err(CodexHandshakeAttemptError::Client), + } +} + +pub(super) fn map_handshake_attempt_error( + error: CodexHandshakeAttemptError, +) -> MappedProviderFailure { + match error { + CodexHandshakeAttemptError::Client(error) => map_handshake_error(error), + CodexHandshakeAttemptError::Cancelled => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Ambiguous, + )), + CodexHandshakeAttemptError::Timeout => MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Ambiguous, + )), + } +} + +pub(super) async fn create_json_attempt( + request: &ColdJsonResponse, + account: &ProviderAccount, + installation_id: &str, + authorization: &SecretString, + cookie_header: Option<&SecretString>, + account_selection: CodexAccountSelectionTelemetry<'_>, +) -> Result { + let Some(handshake_deadline) = remaining(request.context.deadline()) else { + return Err(CodexHandshakeAttemptError::Timeout); + }; + let request_id = request.context.request_id().as_str(); + let mut request_context = CodexRequestContext::auxiliary( + authorization.expose_secret(), + account.upstream_account_id(), + request_id, + Some(installation_id), + ); + request_context.cookie_header = cookie_header.map(ExposeSecret::expose_secret); + request_context.account_selection = account_selection; + tokio::select! { + biased; + _ = request.context.cancellation().cancelled() => Err(CodexHandshakeAttemptError::Cancelled), + _ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout), + response = request.client.post_raw_json( + request.endpoint_path, + request.body.clone(), + request.image_turn_id.as_deref(), + request_context, + ) => response.map_err(CodexHandshakeAttemptError::Client), + } +} + +pub(super) fn cold_json_response_stream(request: ColdJsonResponse) -> EventStream { + Box::pin(async_stream::try_stream! { + let allows_account_state_mutation = request.lease.allows_account_state_mutation(); + let failure_context = OpenAiFailureContext { + client: &request.client, + selector: &request.selector, + quota: &request.quota, + response_origin: &request.response_origin, + cyber_policy_scope: None, + allows_account_state_mutation, + }; + let active_account = request.lease.account().clone(); + let cookie_header = build_cookie_header(request.lease.cookies())?; + let authorization = request + .lease + .authentication() + .authorization_header() + .map_err(|_| { + provider_error( + ProviderErrorKind::Unauthorized, + UpstreamSendState::NotSent, + ) + })?; + let account_selection = CodexAccountSelectionTelemetry::new( + request.lease.affinity_hit(), + request.lease.escape_reason(), + request.lease.account_switch(), + ); + let response = create_json_attempt( + &request, + &active_account, + request.lease.installation_id(), + &authorization, + cookie_header.as_ref(), + account_selection, + ) + .await; + if let Err(CodexHandshakeAttemptError::Client(error)) = &response { + log_client_upstream_error( + UpstreamErrorLogContext::new(&request.context, &active_account, None), + error, + ); + } + let response = match response.map_err(map_handshake_attempt_error) { + Ok(response) => response, + Err(mut failure) => { + if let Some(observation) = failure.observation.take() { + yield ProviderEvent::observation(observation); + } + apply_failure(&failure_context, &active_account, &failure).await; + Err(failure.error)?; + return; + } + }; + + let mut metrics = response.transport_metrics.clone(); + metrics.first_event_ms = Some( + i64::try_from(request.output_started_at.elapsed().as_millis()).unwrap_or(i64::MAX), + ); + if let Some(observation) = codex_response_observation( + CodexBackendTransport::HttpJson, + &response.diagnostics, + &response.response_metadata, + &metrics, + None, + openai_response_timings(&metrics, &response.response_metadata), + ) { + yield ProviderEvent::observation(observation); + } + if allows_account_state_mutation { + synchronize_passive_quota_headers( + &request.quota, + &active_account, + &response.rate_limit_headers, + ) + .await; + if !response.set_cookie_headers.is_empty() + && let Err(error) = request + .selector + .capture_response_cookies( + &active_account, + &request.response_origin, + &response.set_cookie_headers, + ) + .await + { + tracing::warn!( + account_id = %active_account.id(), + error = %error, + "Failed to persist OpenAI image response cookies" + ); + } + } + + let response_meta = + ResponseMeta::for_provider_endpoint(request.context.request_id().as_str()); + yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone())); + let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent) + })?; + yield ProviderEvent::wire(wire); + yield ProviderEvent::canonical(GatewayEvent::Completed( + response_meta.with_finish_reason(FinishReason::Stop), + )); + }) +} + +pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { + let ColdResponse { + client, + response_origin, + request, + upstream_model, + transport_policy, + context, + selector, + quota, + catalog, + lease, + output_started_at, + session_affinity_key, + session_affinity_key_hash, + session_transport_fallbacks, + websocket_retry_count, + stream_max_retries, + mut session_capture, + } = response; + Box::pin(async_stream::try_stream! { + let cyber_policy_scope = lease.cyber_policy_scope().cloned(); + let allows_account_state_mutation = lease.allows_account_state_mutation(); + let failure_context = OpenAiFailureContext { + client: &client, + selector: &selector, + quota: "a, + response_origin: &response_origin, + cyber_policy_scope: cyber_policy_scope.as_ref(), + allows_account_state_mutation, + }; + let mut active_account = lease.account().clone(); + let cookie_header = build_cookie_header(lease.cookies())?; + let authorization = lease + .authentication() + .authorization_header() + .map_err(|_| { + provider_error( + ProviderErrorKind::Unauthorized, + UpstreamSendState::NotSent, + ) + })?; + let request_id = context.request_id().as_str().to_owned(); + let cancellation = context.cancellation().clone(); + let account_selection = CodexAccountSelectionTelemetry::new( + lease.affinity_hit(), + lease.escape_reason(), + lease.account_switch(), + ); + let request_transport_requirement = transport_requirement(&request); + let response = create_response_attempt( + &client, + &request, + codex_request_context( + &request, + &request_id, + &active_account, + lease.installation_id(), + &authorization, + cookie_header.as_ref(), + account_selection, + ), + active_account.id().as_str(), + context.deadline(), + &cancellation, + ) + .await; + let websocket_failure_policy = match &response { + Err(CodexHandshakeAttemptError::Client(error)) + if transport_policy == CodexProviderTransport::PreferWebSocket => + { + websocket_client_failure_policy(error) + } + _ => None, + }; + if let Err(CodexHandshakeAttemptError::Client(error)) = &response { + log_client_upstream_error( + UpstreamErrorLogContext::new(&context, &active_account, None), + error, + ); + } + let response = response.map_err(map_handshake_attempt_error); + let response = match response { + Ok(response) => response, + Err(mut failure) => { + if let Some(policy) = websocket_failure_policy { + apply_websocket_recovery_policy( + &mut failure, + WebSocketRecoveryContext { + policy, + requirement: request_transport_requirement, + retry_count: websocket_retry_count, + max_retries: stream_max_retries, + request_id: context.request_id().as_str(), + attempt_index: context.attempt_index().get(), + account_id: active_account.id().as_str(), + session_affinity_key: session_affinity_key.as_ref(), + session_affinity_key_hash: session_affinity_key_hash.as_deref(), + session_transport_fallbacks: &session_transport_fallbacks, + }, + ); + } + if let Some(observation) = failure.observation.take() { + yield ProviderEvent::observation(observation); + } + apply_failure(&failure_context, &active_account, &failure) + .await; + Err(failure.error)?; + return; + } + }; + if !accepts_backend_transport(transport_policy, response.transport) { + let failure = MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )); + apply_failure(&failure_context, &active_account, &failure) + .await; + Err(failure.error)?; + return; + } + 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 + }); + capture.turn_state = response.turn_state.clone().or(capture.turn_state.clone()); + } + let mut observation_state = OpenAiResponseObservationState::from_backend_response( + &response, + &request, + active_account.id().as_str(), + context.attempt_index().get(), + ); + if let Some(observation) = observation_state.observation() { + yield ProviderEvent::observation(observation); + } + if let Some(etag) = response.response_metadata.models_etag.as_deref() + && let Err(error) = catalog.observe_response_etag(etag) + { + tracing::warn!( + error = %error, + "OpenAI model ETag observation was rejected" + ); + } + if allows_account_state_mutation + && !response.set_cookie_headers.is_empty() + && let Ok(outcome) = selector + .capture_response_cookies( + &active_account, + &response_origin, + &response.set_cookie_headers, + ) + .await + && let Some(revision) = outcome.credential_revision + && let Ok(current) = selector.current_account(active_account.id()).await + && current.revision().get() == revision + { + active_account = current; + } + let response_transport = response.transport; + let websocket_connection_id = response.websocket_connection_id; + 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 mut passive_quota_observation = + OpenAiPassiveQuotaObservation::new(response.rate_limit_headers); + let rate_limit_updates = response.rate_limit_updates; + let turn_state_updates = response.turn_state_update; + // OpenAI 线路为透明代理：HTTP SSE 与 WebSocket 两条上游均启用 raw 透传， + // 下游按字节转发上游原文，避免 serde 往返改写数值/精度（大整数→f64、logprobs 等）。 + // WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON + // （transport/protocol/websocket.rs），push_frames 抽出的 data 即上游原文。 + let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str()) + .with_requested_service_tier(request.service_tier()) + .with_request_tool_pricing(upstream_model.as_str(), request.tools()) + .with_raw_sse_passthrough(); + let mut pre_commit_events = PreCommitClientEvents::new(); + loop { + let Some(stream_deadline) = remaining(context.deadline()) else { + if allows_account_state_mutation { + synchronize_passive_quota( + "a, + &active_account, + passive_quota_observation.rate_limits(), + ) + .await; + } + Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?; + return; + }; + let replay_grace_deadline = pre_commit_events.replay_grace_deadline(); + let next = tokio::select! { + biased; + _ = cancellation.cancelled() => Err(MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Sent, + ))), + _ = tokio::time::sleep(stream_deadline) => Err(MappedProviderFailure::plain(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + ))), + _ = wait_for_replay_grace(replay_grace_deadline) => Ok(PreCommitPoll::GraceElapsed), + chunk = body.next() => match chunk { + Some(Ok(chunk)) => Ok(PreCommitPoll::Upstream(Some(chunk))), + Some(Err(error)) => { + log_client_upstream_error( + UpstreamErrorLogContext::new( + &context, + &active_account, + websocket_connection_id, + ), + &error, + ); + Err(map_stream_error(error)) + } + None => Ok(PreCommitPoll::Upstream(None)), + }, + }; + let next = match next { + Ok(PreCommitPoll::Upstream(next)) => next, + Ok(PreCommitPoll::GraceElapsed) => { + for event in pre_commit_events.commit_pending() { + yield event; + } + continue; + } + Err(mut failure) => { + if failure.websocket_transport_retryable + && response_transport == CodexBackendTransport::WebSocket + && !pre_commit_events.is_committed() + { + apply_websocket_recovery_policy( + &mut failure, + WebSocketRecoveryContext { + policy: WebSocketFailurePolicy::Budgeted, + requirement: request_transport_requirement, + retry_count: websocket_retry_count, + max_retries: stream_max_retries, + request_id: context.request_id().as_str(), + attempt_index: context.attempt_index().get(), + account_id: active_account.id().as_str(), + session_affinity_key: session_affinity_key.as_ref(), + session_affinity_key_hash: session_affinity_key_hash.as_deref(), + session_transport_fallbacks: &session_transport_fallbacks, + }, + ); + } + let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; + if !updates.is_empty() { + passive_quota_observation.observe(&updates); + let update_headers = rate_limit_update_headers(&updates); + if observation_state.merge_rate_limit_headers(&update_headers) + && let Some(observation) = observation_state.observation() + { + yield ProviderEvent::observation(observation); + } + } + if allows_account_state_mutation { + synchronize_passive_quota( + "a, + &active_account, + passive_quota_observation.rate_limits(), + ) + .await; + } + apply_failure(&failure_context, &active_account, &failure) + .await; + Err(failure.error)?; + return; + } + }; + let Some(chunk) = next else { break; }; + let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; + let rate_limits_changed = if updates.is_empty() { + false + } else { + passive_quota_observation.observe(&updates); + observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates)) + }; + let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref() + && let Some(turn_state) = updates.lock().await.take() + { + if let Some(capture) = session_capture.as_mut() { + capture.turn_state = Some(turn_state.clone()); + } + observation_state.merge_client_header("x-codex-turn-state", &turn_state) + } else { + false + }; + let first_event_changed = + observation_state.observe_stream_chunk(&chunk, output_started_at); + let chunk_len = chunk.len(); + let (mut events, canonical_failure) = match decoder.push(&chunk) { + CodexCanonicalOutcome::Events(events) => (events, None), + CodexCanonicalOutcome::Failed(failure) => { + let (events, error, semantic_output_seen) = failure.into_parts(); + (events, Some((error, semantic_output_seen))) + } + }; + pre_commit_events.observe_chunk(chunk_len); + let service_tier_changed = observation_state + .observe_upstream_service_tier(decoder.response_service_tier()); + let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| { + log_canonical_upstream_error( + UpstreamErrorLogContext::new( + &context, + &active_account, + websocket_connection_id, + ), + response_transport, + &error, + ); + let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_)); + ( + map_canonical_error( + error, + &failure_diagnostics, + &failure_set_cookie_headers, + &failure_rate_limit_headers, + ReplayBoundary::from_semantic_output( + semantic_output_seen || pre_commit_events.is_committed(), + ), + ), + atomic_upstream_failure, + ) + }); + let timing_signals = decoder.take_timing_signals(); + let timing_changed = first_event_changed + || observation_state + .observe_timing_signals(timing_signals, output_started_at); + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .any(|event| matches!(event, GatewayEvent::Completed(_))); + let terminal_changed = completed + && observation_state.mark_completed(terminal_response_is_incomplete(&events)); + if allows_account_state_mutation && (completed || terminal_failure.is_some()) { + synchronize_passive_quota( + "a, + &active_account, + passive_quota_observation.rate_limits(), + ) + .await; + } + if let Some((failure, _)) = terminal_failure.as_ref() { + apply_failure(&failure_context, &active_account, failure) + .await; + } + attach_openai_session_update(&mut events, &mut session_capture); + if allows_account_state_mutation && completed && terminal_failure.is_none() { + // 完成事件一旦交给下游，Core 可以立刻停止轮询 Provider stream； + // 在此之前持久化亲和关系，保证成功请求不会因流被提前 drop 而丢失绑定。 + selector + .record_success( + &active_account, + session_affinity_key.as_ref(), + lease.affinity_expected_account_id(), + ) + .await; + selector + .observe_cyber_policy_success(cyber_policy_scope.as_ref()) + .await; + } + if (rate_limits_changed + || service_tier_changed + || timing_changed + || turn_state_changed + || terminal_changed) + && let Some(observation) = observation_state.observation() + { + yield ProviderEvent::observation(observation); + } + if let Some((mut failure, atomic_upstream_failure)) = terminal_failure { + let failure_after_commit = + timing_signals.semantic_output || pre_commit_events.is_committed(); + if failure_after_commit { + for event in pre_commit_events.commit(events) { + yield event; + } + } else if atomic_upstream_failure { + failure.error = failure + .error + .with_atomic_client_events(pre_commit_events.take_for_failure(events)); + } + Err(failure.error)?; + return; + } + let events = pre_commit_events.stage(events, timing_signals, completed); + for event in events { + yield event; + } + if completed { + return; + } + } + let (mut events, canonical_failure) = match decoder.finish() { + CodexCanonicalOutcome::Events(events) => (events, None), + CodexCanonicalOutcome::Failed(failure) => { + let (events, error, semantic_output_seen) = failure.into_parts(); + (events, Some((error, semantic_output_seen))) + } + }; + let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| { + log_canonical_upstream_error( + UpstreamErrorLogContext::new( + &context, + &active_account, + websocket_connection_id, + ), + response_transport, + &error, + ); + let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_)); + ( + map_canonical_error( + error, + &failure_diagnostics, + &failure_set_cookie_headers, + &failure_rate_limit_headers, + ReplayBoundary::from_semantic_output( + semantic_output_seen || pre_commit_events.is_committed(), + ), + ), + atomic_upstream_failure, + ) + }); + let timing_signals = decoder.take_timing_signals(); + let service_tier_changed = observation_state + .observe_upstream_service_tier(decoder.response_service_tier()); + let timing_changed = observation_state + .observe_timing_signals(timing_signals, output_started_at); + let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; + let rate_limits_changed = if updates.is_empty() { + false + } else { + passive_quota_observation.observe(&updates); + observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates)) + }; + if allows_account_state_mutation { + synchronize_passive_quota( + "a, + &active_account, + passive_quota_observation.rate_limits(), + ) + .await; + } + if let Some((failure, _)) = terminal_failure.as_ref() { + apply_failure(&failure_context, &active_account, failure) + .await; + } + let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref() + && let Some(turn_state) = updates.lock().await.take() + { + if let Some(capture) = session_capture.as_mut() { + capture.turn_state = Some(turn_state.clone()); + } + observation_state.merge_client_header("x-codex-turn-state", &turn_state) + } else { + false + }; + attach_openai_session_update(&mut events, &mut session_capture); + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .any(|event| matches!(event, GatewayEvent::Completed(_))); + let terminal_changed = completed + && observation_state.mark_completed(terminal_response_is_incomplete(&events)); + if allows_account_state_mutation && completed && terminal_failure.is_none() { + // 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。 + selector + .record_success( + &active_account, + session_affinity_key.as_ref(), + lease.affinity_expected_account_id(), + ) + .await; + selector + .observe_cyber_policy_success(cyber_policy_scope.as_ref()) + .await; + } + if (service_tier_changed + || timing_changed + || rate_limits_changed + || turn_state_changed + || terminal_changed) + && let Some(observation) = observation_state.observation() + { + yield ProviderEvent::observation(observation); + } + if let Some((mut failure, atomic_upstream_failure)) = terminal_failure { + let failure_after_commit = + timing_signals.semantic_output || pre_commit_events.is_committed(); + if failure_after_commit { + for event in pre_commit_events.commit(events) { + yield event; + } + } else if atomic_upstream_failure { + failure.error = failure + .error + .with_atomic_client_events(pre_commit_events.take_for_failure(events)); + } + Err(failure.error)?; + return; + } + let events = pre_commit_events.finish(events, timing_signals, completed); + for event in events { + yield event; + } + }) +} diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs new file mode 100644 index 000000000..a7fbb92a2 --- /dev/null +++ b/backend/crates/providers/openai/src/provider/failure.rs @@ -0,0 +1,964 @@ +//! OpenAI 上游失败分类、恢复决策与稳定错误投影。 + +use super::*; + +pub(super) struct MappedProviderFailure { + pub(super) error: ProviderError, + pub(super) websocket_transport_retryable: bool, + pub(super) account_failure: Option, + /// 原始上游错误描述，仅在凭据错误状态下持久化。 + pub(super) error_message: Option, + pub(super) cyber_policy_failure: bool, + pub(super) set_cookie_headers: Vec, + pub(super) rate_limit_headers: Vec<(String, String)>, + pub(super) observation: Option, + pub(super) capture_response_cookies: bool, +} + +impl MappedProviderFailure { + pub(super) fn plain(error: ProviderError) -> Self { + Self { + error, + websocket_transport_retryable: false, + account_failure: None, + error_message: None, + cyber_policy_failure: false, + set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), + observation: None, + capture_response_cookies: false, + } + } +} + +pub(super) enum PreCommitPoll { + Upstream(T), + GraceElapsed, +} + +pub(super) async fn wait_for_replay_grace(deadline: Option) { + if let Some(deadline) = deadline { + tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; + } else { + std::future::pending::<()>().await; + } +} + +/// 提交边界前的上游事件预取。 +/// +/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号 +/// 窗口；任一边界到达都会提交已缓存 wire，不能因网关私有资源规则伪造上游协议 +/// 失败。一旦提交，后续事件不再具备无痕重放资格。 +pub(super) struct PreCommitClientEvents { + pending: Vec, + prefetched_bytes: usize, + replay_grace_deadline: Option, + committed: bool, +} + +impl PreCommitClientEvents { + pub(super) const fn new() -> Self { + Self { + pending: Vec::new(), + prefetched_bytes: 0, + replay_grace_deadline: None, + committed: false, + } + } + + pub(super) fn observe_chunk(&mut self, bytes: usize) { + if !self.committed { + self.prefetched_bytes = self.prefetched_bytes.saturating_add(bytes); + } + } + + pub(super) fn stage( + &mut self, + incoming: Vec, + timing_signals: ResponseEventSignals, + completed: bool, + ) -> Vec { + if self.committed { + return incoming; + } + let starts_replay_grace = incoming.iter().any(ProviderEvent::has_client_event); + self.pending.extend(incoming); + if timing_signals.semantic_output + || completed + || self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES + { + return self.commit_pending(); + } + if starts_replay_grace && self.replay_grace_deadline.is_none() { + self.replay_grace_deadline = Instant::now().checked_add(STREAM_REPLAY_GRACE); + } + Vec::new() + } + + pub(super) fn finish( + &mut self, + incoming: Vec, + timing_signals: ResponseEventSignals, + completed: bool, + ) -> Vec { + let events = self.stage(incoming, timing_signals, completed); + if self.committed { + events + } else { + self.commit_pending() + } + } + + pub(super) fn commit(&mut self, incoming: Vec) -> Vec { + if self.committed { + return incoming; + } + self.pending.extend(incoming); + self.commit_pending() + } + + pub(super) fn take_for_failure(&mut self, incoming: Vec) -> Vec { + self.pending.extend(incoming); + self.prefetched_bytes = 0; + self.replay_grace_deadline = None; + std::mem::take(&mut self.pending) + } + + pub(super) const fn replay_grace_deadline(&self) -> Option { + self.replay_grace_deadline + } + + pub(super) const fn is_committed(&self) -> bool { + self.committed + } + + pub(super) fn commit_pending(&mut self) -> Vec { + self.committed = true; + self.prefetched_bytes = 0; + self.replay_grace_deadline = None; + std::mem::take(&mut self.pending) + } +} + +#[derive(Clone, Copy)] +pub(super) enum ReplayBoundary { + BeforeSemanticOutput, + AfterSemanticOutput, +} + +impl ReplayBoundary { + pub(super) const fn from_semantic_output(semantic_output_seen: bool) -> Self { + if semantic_output_seen { + Self::AfterSemanticOutput + } else { + Self::BeforeSemanticOutput + } + } + + pub(super) const fn permits_provider_proof(self) -> bool { + matches!(self, Self::BeforeSemanticOutput) + } +} + +pub(super) struct OpenAiFailureContext<'a> { + pub(super) client: &'a CodexBackendClient, + pub(super) selector: &'a CodexCredentialSelector, + pub(super) quota: &'a Arc, + pub(super) response_origin: &'a Url, + pub(super) cyber_policy_scope: Option<&'a CodexCyberPolicyScope>, + pub(super) allows_account_state_mutation: bool, +} + +#[derive(Clone, Copy)] +pub(super) struct UpstreamErrorLogContext<'a> { + request_id: &'a str, + account_id: &'a str, + attempt_index: u32, + websocket_connection_id: Option, +} + +impl<'a> UpstreamErrorLogContext<'a> { + pub(super) fn new( + context: &'a AttemptContext, + account: &'a ProviderAccount, + websocket_connection_id: Option, + ) -> Self { + Self { + request_id: context.request_id().as_str(), + account_id: account.id().as_str(), + attempt_index: context.attempt_index().get(), + websocket_connection_id, + } + } + + fn with_websocket_connection_id(mut self, connection_id: Option) -> Self { + self.websocket_connection_id = self.websocket_connection_id.or(connection_id); + self + } +} + +pub(super) fn log_client_upstream_error( + context: UpstreamErrorLogContext<'_>, + error: &CodexClientError, +) { + match error { + CodexClientError::Upstream { + status, + body, + transport, + .. + } => log_raw_upstream_body( + context, + *transport, + "http_error_response", + Some(status.as_u16()), + None, + None, + body, + ), + CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => { + log_raw_upstream_body( + context, + CodexBackendTransport::WebSocket, + "websocket_opening_response", + Some(upstream.status_code), + None, + None, + &upstream.body, + ); + } + CodexClientError::WebSocket(error) => { + let Some(close) = error.close_before_terminal() else { + return; + }; + let context = context.with_websocket_connection_id(close.connection_id()); + let websocket_connection_id = context + .websocket_connection_id + .map(|connection_id| connection_id.to_string()) + .unwrap_or_default(); + let close_reason_bytes = close.reason().map_or(0, str::len); + tracing::warn!( + request_id = %context.request_id, + account_id = %context.account_id, + attempt_index = context.attempt_index, + websocket_connection_id, + upstream_transport = WEBSOCKET_TRANSPORT, + upstream_error_kind = "websocket_close", + upstream_close_code = close.code().unwrap_or_default(), + upstream_close_code_present = close.code().is_some(), + upstream_last_event_type = close.last_event_type().unwrap_or_default(), + upstream_last_event_type_present = close.last_event_type().is_some(), + upstream_terminal_seen = false, + upstream_close_reason_bytes = close_reason_bytes, + upstream_close_reason_present = close.reason().is_some(), + upstream_error_raw_present = close.reason().is_some(), + "OpenAI upstream WebSocket closed before terminal event" + ); + } + _ => {} + } +} + +pub(super) fn log_canonical_upstream_error( + context: UpstreamErrorLogContext<'_>, + transport: CodexBackendTransport, + error: &CodexCanonicalError, +) { + let CodexCanonicalError::Upstream(failure) = error else { + return; + }; + log_raw_upstream_body( + context, + transport, + "responses_error_event", + failure.explicit_status_code, + failure.upstream_code.as_deref(), + failure.upstream_type.as_deref(), + failure.raw_body(), + ); +} + +pub(super) fn log_raw_upstream_body( + context: UpstreamErrorLogContext<'_>, + transport: CodexBackendTransport, + error_kind: &'static str, + status_code: Option, + upstream_code: Option<&str>, + upstream_type: Option<&str>, + upstream_error_raw: &str, +) { + let websocket_connection_id = context + .websocket_connection_id + .map(|connection_id| connection_id.to_string()) + .unwrap_or_default(); + tracing::warn!( + request_id = %context.request_id, + account_id = %context.account_id, + attempt_index = context.attempt_index, + websocket_connection_id, + upstream_transport = actual_transport_name(transport), + upstream_error_kind = error_kind, + upstream_status_code = status_code.unwrap_or_default(), + upstream_status_code_present = status_code.is_some(), + upstream_error_code = upstream_code.unwrap_or_default(), + upstream_error_code_present = upstream_code.is_some(), + upstream_error_type = upstream_type.unwrap_or_default(), + upstream_error_type_present = upstream_type.is_some(), + upstream_error_raw, + upstream_error_raw_bytes = upstream_error_raw.len(), + "OpenAI upstream returned an error payload" + ); +} + +pub(super) async fn apply_failure( + context: &OpenAiFailureContext<'_>, + account: &ProviderAccount, + failure: &MappedProviderFailure, +) { + if !context.allows_account_state_mutation { + return; + } + synchronize_passive_quota_headers(context.quota, account, &failure.rate_limit_headers).await; + let needs_authoritative_quota_refresh = matches!( + failure.account_failure, + Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. }) + ); + if failure.cyber_policy_failure { + context + .selector + .record_cyber_policy_failure(context.cyber_policy_scope, account) + .await; + } + if let Some(account_failure) = failure.account_failure { + context + .client + .evict_websocket_account(account.id().as_str()) + .await; + if let Err(error) = context + .selector + .record_failure(account, account_failure, failure.error_message.clone()) + .await + { + tracing::warn!( + account_id = %account.id(), + error = %error, + "Failed to persist OpenAI account failure state" + ); + } + } + if needs_authoritative_quota_refresh { + schedule_authoritative_quota_refresh_after_failure(context.quota, account); + } + if failure.capture_response_cookies + && !failure.set_cookie_headers.is_empty() + && let Err(error) = context + .selector + .capture_response_cookies( + account, + context.response_origin, + &failure.set_cookie_headers, + ) + .await + { + tracing::warn!( + account_id = %account.id(), + error = %error, + "Failed to persist OpenAI response cookies" + ); + } +} + +pub(super) fn schedule_authoritative_quota_refresh_after_failure( + quota: &Arc, + account: &ProviderAccount, +) { + // 限额错误可以确认账号状态，但 Responses 事件中的 used_percent 可能仍停在上一结算点。 + // usage 快照在后台补齐展示基线，不得撤销同一轮真实失败，也不能阻塞原始响应。 + // 先等待 2 秒让上游结算，再查询；5 秒仅限制实际查询自身。 + let quota = Arc::clone(quota); + let account_id = account.id().clone(); + drop(tokio::spawn(async move { + tokio::time::sleep(QUOTA_FAILURE_REFRESH_DELAY).await; + match tokio::time::timeout( + QUOTA_FAILURE_REFRESH_TIMEOUT, + quota.refresh_account_after_failure(&account_id), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::warn!( + account_id = %account_id, + error = %error, + "OpenAI authoritative quota synchronization after quota rejection failed" + ); + } + Err(_) => { + tracing::warn!( + account_id = %account_id, + timeout_ms = QUOTA_FAILURE_REFRESH_TIMEOUT.as_millis(), + "OpenAI authoritative quota synchronization after quota rejection timed out" + ); + } + } + })); +} + +pub(super) fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { + map_client_error(error, UpstreamSendState::Ambiguous, true) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum WebSocketFailurePolicy { + Budgeted, + ImmediateFallback, +} + +pub(super) struct WebSocketRecoveryContext<'a> { + pub(super) policy: WebSocketFailurePolicy, + pub(super) requirement: TransportRequirement, + pub(super) retry_count: u32, + pub(super) max_retries: u32, + pub(super) request_id: &'a str, + pub(super) attempt_index: u32, + pub(super) account_id: &'a str, + pub(super) session_affinity_key: Option<&'a ProviderSessionAffinityKey>, + pub(super) session_affinity_key_hash: Option<&'a str>, + pub(super) session_transport_fallbacks: &'a CodexSessionTransportFallbacks, +} + +pub(super) fn websocket_client_failure_policy( + error: &CodexClientError, +) -> Option { + match error { + CodexClientError::Upstream { + status, + transport: CodexBackendTransport::WebSocket, + .. + } if *status == reqwest::StatusCode::UPGRADE_REQUIRED => { + Some(WebSocketFailurePolicy::ImmediateFallback) + } + CodexClientError::Upstream { + status, + transport: CodexBackendTransport::WebSocket, + .. + } if status.is_server_error() + || matches!( + *status, + reqwest::StatusCode::REQUEST_TIMEOUT + | reqwest::StatusCode::CONFLICT + | reqwest::StatusCode::TOO_EARLY + | reqwest::StatusCode::TOO_MANY_REQUESTS + ) => + { + Some(WebSocketFailurePolicy::Budgeted) + } + CodexClientError::WebSocket(error) + if !matches!(error, CodexWebSocketExchangeError::InvalidRequest(_)) => + { + Some(WebSocketFailurePolicy::Budgeted) + } + _ => None, + } +} + +pub(super) fn apply_websocket_recovery_policy( + failure: &mut MappedProviderFailure, + context: WebSocketRecoveryContext<'_>, +) { + let recovery_is_safe = match failure.error.send_state() { + UpstreamSendState::NotSent => context.requirement.allows_pre_send_http_fallback(), + UpstreamSendState::Sent | UpstreamSendState::Ambiguous => { + context.requirement.allows_pre_delivery_http_fallback() + } + }; + if !recovery_is_safe { + return; + } + + // 传输恢复必须保持原账号可调度；最终 HTTP attempt 若仍失败，再按真实 HTTP + // 结果更新账号健康度，避免中间 WS 错误把同账号钉选提前冷却掉。 + failure.account_failure = None; + let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback + || context.retry_count >= context.max_retries; + if !fallback_now { + let Some(retry_index) = context.retry_count.checked_add(1).and_then(NonZeroU32::new) else { + return; + }; + let delay = failure + .error + .retry_after() + .unwrap_or_else(|| websocket_retry_backoff(retry_index)); + failure + .error + .set_pre_delivery_transport_retry(retry_index, delay); + tracing::warn!( + request_id = context.request_id, + attempt_index = context.attempt_index, + account_id = context.account_id, + websocket_failure_kind = failure.error.kind().as_str(), + websocket_failure_code = failure + .error + .upstream_code() + .map_or("", OpaqueUpstreamValue::as_str), + upstream_status_code = failure.error.upstream_status().unwrap_or_default(), + upstream_status_code_present = failure.error.upstream_status().is_some(), + websocket_retry_count = retry_index.get(), + websocket_max_retries = context.max_retries, + websocket_retry_delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX), + transport_requirement = context.requirement.as_str(), + "OpenAI upstream WebSocket failed; retrying the same account" + ); + return; + } + + failure.error.set_pre_delivery_transport_fallback(); + let sticky_state_transition = context + .session_affinity_key + .map(|key| context.session_transport_fallbacks.disable_websocket(key)); + tracing::warn!( + request_id = context.request_id, + attempt_index = context.attempt_index, + account_id = context.account_id, + websocket_failure_kind = failure.error.kind().as_str(), + websocket_failure_code = failure + .error + .upstream_code() + .map_or("", OpaqueUpstreamValue::as_str), + upstream_status_code = failure.error.upstream_status().unwrap_or_default(), + upstream_status_code_present = failure.error.upstream_status().is_some(), + websocket_retry_count = context.retry_count, + websocket_max_retries = context.max_retries, + websocket_fallback_reason = match context.policy { + WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted", + WebSocketFailurePolicy::ImmediateFallback => "upgrade_required", + }, + session_affinity_present = context.session_affinity_key.is_some(), + session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), + sticky_http_enabled = sticky_state_transition.is_some(), + sticky_state_transition = sticky_state_transition.unwrap_or(false), + "OpenAI upstream WebSocket disabled for the current session" + ); +} + +pub(super) fn websocket_retry_backoff(retry_index: NonZeroU32) -> Duration { + const INITIAL_DELAY_MS: u64 = 200; + + let shift = retry_index.get().saturating_sub(1).min(63); + let base_ms = INITIAL_DELAY_MS.saturating_mul(1_u64.checked_shl(shift).unwrap_or(u64::MAX)); + let mut random = [0_u8; 2]; + let jitter_per_mille = if getrandom::fill(&mut random).is_ok() { + 900_u64 + u64::from(u16::from_le_bytes(random) % 200) + } else { + 1_000 + }; + Duration::from_millis(base_ms.saturating_mul(jitter_per_mille) / 1_000) +} + +pub(super) fn continuation_replay_required_error() -> ProviderError { + ProviderError::new( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + ) + .with_continuation_failure(ContinuationFailure::HistoryUnavailable) + .with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, + Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()), + Some("invalid_request_error".to_owned()), + )) +} + +pub(super) fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { + let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error); + let websocket_failure = error.transport() == Some(CodexBackendTransport::WebSocket); + let mut failure = map_client_error(error, UpstreamSendState::Sent, false); + failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure; + if allows_pre_delivery_retry && !websocket_failure { + failure.error = failure.error.with_pre_delivery_retry(); + } + failure +} + +pub(super) fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool { + match error { + CodexClientError::Http(_) + | CodexClientError::HttpJson(_) + | CodexClientError::StreamIdleTimeout { .. } => true, + CodexClientError::WebSocket(error) => !matches!( + error, + CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::Upstream(_) + | CodexWebSocketExchangeError::ContinuationUnavailable { .. } + ), + _ => false, + } +} + +pub(super) 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, + ), + } +} + +pub(super) 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::RequestBodyEncode(_) + | CodexClientError::RequestCompression(_) + | CodexClientError::ModelCatalog(_) + | 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) | CodexClientError::HttpJson(error) => { + let send_state = if error.is_connect() { + UpstreamSendState::NotSent + } else { + uncertain_state + }; + MappedProviderFailure::plain(provider_error( + if error.is_timeout() { + ProviderErrorKind::Timeout + } else { + ProviderErrorKind::Transport + }, + send_state, + )) + } + CodexClientError::WebSocket(error) => { + let close_code = error.close_before_terminal().and_then(|close| close.code()); + let client_visible_error = websocket_client_visible_error(&error); + let mut failure = MappedProviderFailure::plain(provider_error( + websocket_error_kind(&error), + websocket_send_state(&error), + )); + if let Some(close_code) = close_code { + failure.error = + failure + .error + .with_upstream_code(OpaqueUpstreamValue::new(format!( + "websocket_close_{close_code}" + ))); + } + if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { + failure.error = failure.error.with_upstream_code(OpaqueUpstreamValue::new( + WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(), + )); + } + if let Some(client_visible_error) = client_visible_error { + failure.error = failure + .error + .with_client_visible_upstream_error(client_visible_error); + } + failure + } + }; + if let Some(continuation_failure) = continuation_failure { + failure.error = failure + .error + .with_continuation_failure(continuation_failure); + } + failure.observation = observation; + failure +} + +pub(super) fn websocket_client_visible_error( + error: &CodexWebSocketExchangeError, +) -> Option { + if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { + return Some(ClientVisibleUpstreamError::new( + "websocket connection limit reached", + Some(WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned()), + Some("websocket_error".to_owned()), + )); + } + let close = error.close_before_terminal()?; + let message = close + .reason() + .filter(|reason| !reason.is_empty()) + .map_or_else(|| close.to_string(), str::to_owned); + Some(ClientVisibleUpstreamError::new( + message, + close.code().map(|code| code.to_string()), + Some("websocket_close_error".to_owned()), + )) +} + +pub(super) fn map_upstream_failure( + mut failure: CodexUpstreamFailure, + observation: Option, + replay_boundary: ReplayBoundary, +) -> MappedProviderFailure { + let category = failure.category(); + let cyber_policy_failure = failure + .status + .is_some_and(|status| status.is_client_error()) + && is_cyber_policy_code(failure.code.as_deref()); + 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); + if let Some(message) = failure.client_message.as_ref() { + error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + message.clone(), + failure.client_code.clone(), + failure.client_error_type.clone(), + )); + } + if let Some(response) = failure.client_response.take() { + let response = (*response).into_parts(); + error = error.with_client_visible_upstream_response( + ClientVisibleUpstreamResponse::new( + response.status, + response.content_type, + response.body, + ) + .with_headers( + response + .client_headers + .into_iter() + .map(|(name, value)| ProviderResponseHeader::new(name, value)) + .collect(), + ), + ); + } + if let Some(status) = failure.status { + error = error.with_status(status.as_u16()); + } + if replay_boundary.permits_provider_proof() + && (failure.replay_is_safe() || cyber_policy_failure) + { + 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() { + error = error.with_upstream_code(OpaqueUpstreamValue::new(code.to_owned())); + } + if let Some(request_id) = failure.request_id.as_deref() { + error = error.with_upstream_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); + } + MappedProviderFailure { + error, + websocket_transport_retryable: false, + account_failure: account_failure( + category, + failure.retry_after_seconds, + failure.usage_limit_resets_at, + ), + error_message: failure.client_message, + cyber_policy_failure, + set_cookie_headers: failure.set_cookie_headers, + rate_limit_headers: failure.rate_limit_headers, + observation, + capture_response_cookies: !matches!( + category, + CodexFailureCategory::CloudflareChallenge | CodexFailureCategory::CloudflarePathBlocked + ), + } +} + +pub(super) fn is_cyber_policy_code(code: Option<&str>) -> bool { + code.is_some_and(|code| code.trim().eq_ignore_ascii_case("cyber_policy")) +} + +pub(super) fn is_history_failure_code(code: &str) -> bool { + matches!( + code, + "previous_response_not_found" + | "invalid_encrypted_content" + | "missing_tool_output" + | "no_tool_output" + ) +} + +pub(super) 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::UsageLimitExhausted => ProviderErrorKind::QuotaExhausted, + 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, + } +} + +pub(super) const fn upstream_send_state(phase: CodexUpstreamSendPhase) -> UpstreamSendState { + match phase { + CodexUpstreamSendPhase::BeforePayload => UpstreamSendState::NotSent, + CodexUpstreamSendPhase::AfterPayload => UpstreamSendState::Sent, + CodexUpstreamSendPhase::Ambiguous => UpstreamSendState::Ambiguous, + } +} + +pub(super) fn account_failure( + category: CodexFailureCategory, + retry_after_seconds: Option, + usage_limit_resets_at: Option, +) -> Option { + match category { + CodexFailureCategory::CredentialExpired => Some(CodexAccountFailure::CredentialExpired), + CodexFailureCategory::IdentityVerificationRequired => { + Some(CodexAccountFailure::IdentityVerificationRequired) + } + CodexFailureCategory::Banned => Some(CodexAccountFailure::Banned), + CodexFailureCategory::UsageLimitExhausted => { + Some(CodexAccountFailure::UsageLimitExhausted { + reset_at: usage_limit_resets_at + .and_then(|seconds| u64::try_from(seconds).ok()) + .and_then(|seconds| { + SystemTime::UNIX_EPOCH.checked_add(Duration::from_secs(seconds)) + }), + }) + } + 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, + } +} + +pub(super) 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::ConnectionLimitReached + | CodexWebSocketExchangeError::InvalidSse(_) + | CodexWebSocketExchangeError::UnexpectedBinaryEvent => UpstreamSendState::Sent, + CodexWebSocketExchangeError::Transport(_) + | CodexWebSocketExchangeError::PostSendAmbiguous { .. } + | CodexWebSocketExchangeError::SendTimeout { .. } + | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } + | CodexWebSocketExchangeError::InitialEventTimeout { .. } => UpstreamSendState::Ambiguous, + } +} + +pub(super) const fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> ProviderErrorKind { + match error { + CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::InvalidSse(_) + | 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::ConnectionLimitReached => ProviderErrorKind::RateLimited, + CodexWebSocketExchangeError::Transport(_) + | CodexWebSocketExchangeError::Connect(_) + | CodexWebSocketExchangeError::PostSendAmbiguous { .. } + | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { + ProviderErrorKind::Transport + } + } +} + +pub(super) fn provider_error( + kind: ProviderErrorKind, + send_state: UpstreamSendState, +) -> ProviderError { + ProviderError::new(kind, send_state) +} + +pub(super) fn remaining(deadline: SystemTime) -> Option { + deadline + .duration_since(SystemTime::now()) + .ok() + .filter(|remaining| !remaining.is_zero()) +} diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs new file mode 100644 index 000000000..c746e9996 --- /dev/null +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -0,0 +1,494 @@ +//! Codex 的 `gateway-core` Provider adapter。 + +use std::collections::BTreeSet; +use std::fmt; +use std::num::NonZeroU32; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::{StreamExt, future::BoxFuture}; +use gateway_core::account::{AccountFeedbackStats, ProviderAccount}; +use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope}; +use gateway_core::engine::provider::{ + EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, + ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, + ProviderSelectionObservation, ProviderStream, +}; +use gateway_core::engine::{AttemptContext, AttemptTransport, ContinuationAttempt}; +use gateway_core::error::{ + ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure, + OpaqueUpstreamValue, ProviderError, ProviderErrorKind, +}; +use gateway_core::event::{ + FinishReason, GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, + ProviderResponseMetadata, ProviderResponseObservation, ProviderResponseTimings, ResponseMeta, + UpstreamHttpVersion, WebSocketPoolKind, +}; +use gateway_core::lifecycle::CancellationToken; +use gateway_core::operation::{ + GenerateRequest, ImageRequest, ImageRequestKind, Operation, OperationKind, ProviderSessionState, +}; +use gateway_core::provider_ports::ProviderSessionAffinityKey; +use gateway_core::routing::{ + ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId, +}; +use gateway_core::task::{ + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, + WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration, + WorkerRunnable, WorkerSchedule, WorkerTaskError, +}; +use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; +use gateway_protocol::openai::events::{ + ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs, +}; +use reqwest::Client; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use url::Url; +use uuid::Uuid; + +use crate::credential::{ + CodexAccountFailure, CodexCredentialCatalogError, CodexCredentialCatalogService, + CodexCredentialLease, CodexCredentialQuotaService, CodexCredentialRefreshOutcome, + CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, + CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, + SelectCodexCredential, SelectCodexProviderEndpointCredential, + derive_codex_cyber_policy_session_key, derive_codex_session_affinity, +}; +use crate::session_transport::CodexSessionTransportFallbacks; +use crate::transport::canonical::{ + CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, +}; +use crate::transport::catalog::{ + CodexCatalogCapabilityEvidence, CodexCatalogModel, CodexCatalogVisibility, +}; +use crate::transport::diagnostics::{ + CodexFailureCategory, CodexUpstreamFailure, CodexUpstreamSendPhase, +}; +use crate::transport::profile::{ + APPCAST_POLL_INTERVAL, CodexDesktopReleaseService, CodexWireProfileState, +}; +use crate::transport::protocol::responses::{ + CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, + PreviousResponseScope, ResponseEventSignals, TransportRequirement, transport_requirement, +}; +use crate::transport::protocol::websocket::WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE; +use crate::transport::request::{ + CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account, +}; +use crate::transport::session::CodexSessionIdentity; +use crate::transport::usage::normalize_service_tier; +use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason}; +use crate::transport::{ + CODEX_IMAGE_EDITS_PATH, CODEX_IMAGE_GENERATIONS_PATH, CODEX_RESPONSES_PATH, + CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse, + CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates, + CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics, + CodexWebSocketPool, endpoint_url, +}; + +mod execution; +mod failure; +mod observation; +mod workers; + +use execution::*; +use failure::*; +use observation::*; +pub(crate) use workers::worker_contributions; + +const PROVIDER_NAME: &str = "openai"; +const HTTP_SSE_TRANSPORT: &str = "http_sse"; +const HTTP_JSON_TRANSPORT: &str = "http_json"; +const WEBSOCKET_TRANSPORT: &str = "websocket"; +const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024; +/// 提交边界前最多保留 64 KiB 原始上游 chunk；达到阈值后结束无感换号窗口， +/// 但不会把上游数据改写成协议失败。 +const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; +/// 短暂保留 response.created 等结构事件，让随后到达的明确拒绝可以无感换号； +/// 到期即放行，避免模型长时间思考时让客户端一直收不到首事件。 +const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(1_200); +// 额度拒绝后先给上游额度结算留出时间，再以受限时长同步 usage 快照。 +const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2); +const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); +pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api"; +pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexProviderTransport { + HttpOnly, + PreferWebSocket, +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum CodexProviderConfigError { + #[error("Codex provider URL is invalid")] + InvalidBaseUrl, +} + +pub struct CodexProvider { + selector: Arc, + catalog: Arc, + quota: Arc, + account_feedback: Arc, + client: CodexBackendClient, + responses_url: Url, + image_generations_url: Url, + image_edits_url: Url, + session_identity: Option, + session_transport_fallbacks: CodexSessionTransportFallbacks, + stream_max_retries: u32, +} + +impl CodexProvider { + // Provider 构造集中装配独立领域服务和透明传输依赖，拆分参数会模糊所有权。 + #[expect(clippy::too_many_arguments)] + pub fn new( + selector: Arc, + catalog: Arc, + quota: Arc, + account_feedback: Arc, + http: Client, + profile: CodexWireProfileState, + base_url: String, + websocket_pool: Arc, + stream_max_retries: u32, + ) -> Result { + let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH)) + .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; + let image_generations_url = + Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_GENERATIONS_PATH)) + .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; + let image_edits_url = Url::parse(&endpoint_url(&base_url, CODEX_IMAGE_EDITS_PATH)) + .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; + let client = + CodexBackendClient::new(http, base_url, profile).with_websocket_pool(websocket_pool); + Ok(Self { + selector, + catalog, + quota, + account_feedback, + client, + responses_url, + image_generations_url, + image_edits_url, + session_identity: None, + session_transport_fallbacks: CodexSessionTransportFallbacks::default(), + stream_max_retries, + }) + } + + pub(crate) fn with_session_identity(mut self, identity: CodexSessionIdentity) -> Self { + self.session_identity = Some(identity); + self + } +} + +impl fmt::Debug for CodexProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CodexProvider") + .field("selector", &"") + .field("catalog", &"") + .finish() + } +} + +#[async_trait] +impl Provider for CodexProvider { + fn name(&self) -> &'static str { + 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, reasoning_effort) = encode_generate_request(request, "observability") + .map(|encoded| { + let semantics = encoded.semantics(); + let reasoning_effort = semantics.reasoning_effort.clone(); + (semantics, reasoning_effort) + }) + .unwrap_or_default(); + ProviderRequestObservation { + reasoning_effort, + 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, + ) -> Result, ProviderError> { + let snapshot = self.catalog.synchronize().await.map_err(|_| { + provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) + })?; + Ok(snapshot + .models() + .iter() + .map(compile_model_capabilities) + .collect()) + } + + async fn execute( + &self, + request: ProviderRequest, + context: AttemptContext, + ) -> Result { + if request.candidate().provider().as_str() != PROVIDER_NAME { + return Err(provider_error( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + )); + } + let candidate = request.candidate(); + if context.cancellation().is_cancelled() { + return Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )); + } + if remaining(context.deadline()).is_none() { + return Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::NotSent, + )); + } + if let Operation::GenerateImage(image) = request.operation() { + return self.execute_image(image, candidate, context).await; + } + let Operation::Generate(generate) = request.operation() else { + return Err(provider_error( + ProviderErrorKind::Unsupported, + UpstreamSendState::NotSent, + )); + }; + let Some(upstream_model) = candidate.upstream_model() else { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )); + }; + let previous_session = decode_openai_session_state(generate); + let continuation_requested = generate.native_continuation_requested(); + let mut upstream_request = encode_generate_request(generate, upstream_model.as_str()) + .map_err(map_request_error)?; + 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 let Some(identity) = &self.session_identity { + identity.prepare_local_conversation(&mut upstream_request); + } + if previous_session.as_ref().is_some_and(|state| { + same_client_turn( + state.client_turn_id.as_deref(), + upstream_request.client_turn_id.as_deref(), + ) + }) { + upstream_request.turn_state = previous_session + .as_ref() + .and_then(|state| state.turn_state.clone()); + } + let session_affinity = + derive_codex_session_affinity(&upstream_request, context.client_api_key_ref()); + let cyber_policy_session_key = + derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref()); + + let selection_started_at = Instant::now(); + let lease = self + .selector + .select_with_cyber_policy( + &SelectCodexCredential { + upstream_model: upstream_model.as_str(), + request_url: &self.responses_url, + attempt: &context, + session_affinity_key: session_affinity.as_ref().map(|affinity| affinity.key()), + }, + cyber_policy_session_key.as_ref(), + session_affinity.as_ref(), + ) + .await + .map_err(map_selection_error)?; + let account_selection_wait_ms = + u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + let lease = Arc::new(lease); + // 首字计时的起点：账号选择完成之后、上游建立之前。 + let output_started_at = Instant::now(); + let provider_kind = ProviderKind::new(PROVIDER_NAME) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; + let state_owner_cross_account = context + .account_state_owner() + .is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id())) + || previous_session + .as_ref() + .is_some_and(|state| state.account_id != lease.account_id().as_str()); + let has_explicit_state_owner = + context.account_state_owner().is_some() || previous_session.is_some(); + let account_scope = + if state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()) { + RequestAccountScope::Different + } else if has_explicit_state_owner || lease.affinity_hit() { + RequestAccountScope::Same + } else { + RequestAccountScope::Unknown + }; + if context.continuation_attempt() != ContinuationAttempt::None + && let Some(continuation) = context.continuation() + { + match continuation { + ContinuationBinding::Pinned(continuation) => { + let previous_response_scope = match context.continuation_attempt() { + ContinuationAttempt::Native => match previous_session + .as_ref() + .map(|state| state.continuation_scope) + { + Some(OpenAiContinuationScope::Persisted) => { + PreviousResponseScope::Persisted + } + Some(OpenAiContinuationScope::ConnectionLocal) => { + PreviousResponseScope::ConnectionLocal + } + Some(OpenAiContinuationScope::ReplayRequired) => { + PreviousResponseScope::ExternalUnknown + } + None => match continuation.scope() { + NativeContinuationScope::Persisted => { + PreviousResponseScope::Persisted + } + NativeContinuationScope::ConnectionLocal => { + PreviousResponseScope::ConnectionLocal + } + }, + }, + ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => { + PreviousResponseScope::ExternalUnknown + } + ContinuationAttempt::None => PreviousResponseScope::ExternalUnknown, + }; + 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); + } + } + } + if upstream_request.previous_response_id().is_some() + && !account_scope.can_reuse_account_state() + { + return Err(continuation_replay_required_error()); + } + scope_request_to_account( + &mut upstream_request, + lease.installation_id(), + account_scope, + ); + let requirement = transport_requirement(&upstream_request); + let requested_transport = selected_transport(&upstream_request); + let sticky_http_fallback = context.transport() == AttemptTransport::Default + && requested_transport == CodexProviderTransport::PreferWebSocket + && !requirement.requires_websocket() + && session_affinity.as_ref().is_some_and(|affinity| { + self.session_transport_fallbacks + .is_http_only(affinity.key()) + }); + let transport = if requirement.requires_websocket() { + CodexProviderTransport::PreferWebSocket + } else { + match context.transport() { + AttemptTransport::Default if sticky_http_fallback => { + CodexProviderTransport::HttpOnly + } + AttemptTransport::Default => requested_transport, + AttemptTransport::Retry(_) => requested_transport, + AttemptTransport::Fallback => CodexProviderTransport::HttpOnly, + } + }; + if sticky_http_fallback { + tracing::info!( + request_id = %context.request_id(), + attempt_index = context.attempt_index().get(), + session_affinity_key_hash = session_affinity + .as_ref() + .map_or("", CodexSessionAffinity::key_hash), + transport_requirement = requirement.as_str(), + "Reusing session-sticky OpenAI upstream HTTP fallback" + ); + } + apply_transport(&mut upstream_request, transport); + let metadata = ProviderCallMetadata::new( + provider_kind, + upstream_model.clone(), + lease.account_id().clone(), + UpstreamTransport::new(transport_name(transport)).map_err(|_| { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) + })?, + ) + .with_selection_observation(ProviderSelectionObservation::new( + account_selection_wait_ms, + lease.capacity_snapshot(), + )); + let response_store = upstream_request.store(); + let session_capture = + (!continuation_requested || previous_session.is_some()).then(|| OpenAiSessionCapture { + account_id: lease.account_id().as_str().to_owned(), + conversation_id: upstream_request.local_conversation_id.clone(), + turn_state: upstream_request.turn_state.clone(), + client_turn_id: upstream_request.client_turn_id.clone(), + response_store, + continuation_scope: None, + }); + let allows_account_state_mutation = lease.allows_account_state_mutation(); + let session_affinity_key_hash = session_affinity + .as_ref() + .map(|affinity| affinity.key_hash().to_owned()); + let session_affinity_key = session_affinity.map(CodexSessionAffinity::into_key); + let websocket_retry_count = match context.transport() { + AttemptTransport::Retry(retry_index) => retry_index.get(), + AttemptTransport::Default | AttemptTransport::Fallback => 0, + }; + let events = cold_response_stream(ColdResponse { + client: self.client.clone(), + response_origin: self.responses_url.clone(), + request: upstream_request, + upstream_model: upstream_model.clone(), + transport_policy: transport, + context, + selector: Arc::clone(&self.selector), + quota: Arc::clone(&self.quota), + catalog: Arc::clone(&self.catalog), + lease: Arc::clone(&lease), + output_started_at, + session_affinity_key, + session_affinity_key_hash, + session_transport_fallbacks: self.session_transport_fallbacks.clone(), + websocket_retry_count, + stream_max_retries: self.stream_max_retries, + session_capture, + }); + let stream = ProviderStream::new(metadata, events, lease); + Ok(if allows_account_state_mutation { + stream.with_account_feedback(Arc::clone(&self.account_feedback)) + } else { + stream + }) + } +} diff --git a/backend/crates/providers/openai/src/provider/observation.rs b/backend/crates/providers/openai/src/provider/observation.rs new file mode 100644 index 000000000..08e576cdc --- /dev/null +++ b/backend/crates/providers/openai/src/provider/observation.rs @@ -0,0 +1,796 @@ +//! OpenAI 请求、响应与 transport 观测事实归一化。 + +use super::*; + +/// OpenAI Responses 的观测状态完全归 Provider 所有。 +/// +/// 每次原始 SSE/WS 事件推进状态后重新生成不可变 observation，Core 只负责携带、 +/// 持久化和展示这份安全快照，不解释任何 OpenAI 协议字段。 +pub(super) struct OpenAiResponseObservationState { + transport: CodexBackendTransport, + diagnostics: CodexUpstreamDiagnostics, + response_metadata: CodexResponseMetadata, + metrics: CodexTransportMetrics, + websocket_pool_decision: Option, + request_summary: Value, + requested_model: String, + stream: bool, + compact: bool, + requested_service_tier: Option, + upstream_service_tier: Option, + rate_limit_headers: Vec<(String, String)>, + account_id: String, + attempt_index: u32, + timings: ProviderResponseTimings, + terminal: Option, +} + +pub(super) struct OpenAiPassiveQuotaObservation { + rate_limits: Vec, +} + +impl OpenAiPassiveQuotaObservation { + pub(super) fn new(headers: Vec<(String, String)>) -> Self { + Self { + rate_limits: parse_rate_limit_headers(&headers).into_iter().collect(), + } + } + + pub(super) fn observe(&mut self, updates: &[ParsedRateLimits]) { + self.rate_limits.extend_from_slice(updates); + } + + pub(super) fn rate_limits(&self) -> &[ParsedRateLimits] { + &self.rate_limits + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum OpenAiResponseTerminal { + Completed, + Incomplete, +} + +impl OpenAiResponseObservationState { + pub(super) fn from_backend_response( + response: &CodexBackendStreamingResponse, + request: &CodexResponsesRequest, + account_id: &str, + attempt_index: u32, + ) -> Self { + let semantics = request.semantics(); + Self { + transport: response.transport, + diagnostics: response.diagnostics.clone(), + response_metadata: response.response_metadata.clone(), + metrics: response.transport_metrics.clone(), + websocket_pool_decision: response.websocket_pool_decision, + request_summary: openai_response_request_summary(request, response.transport), + requested_model: request.model().to_owned(), + stream: request.stream(), + compact: semantics.compact, + requested_service_tier: normalize_service_tier(request.service_tier()), + upstream_service_tier: None, + rate_limit_headers: selected_observation_headers(&response.rate_limit_headers), + account_id: account_id.to_owned(), + attempt_index, + timings: openai_response_timings( + &response.transport_metrics, + &response.response_metadata, + ), + terminal: None, + } + } + + pub(super) fn observation(&self) -> Option { + let mut observation = codex_response_observation( + self.transport, + &self.diagnostics, + &self.response_metadata, + &self.metrics, + self.websocket_pool_decision, + self.timings, + )?; + if let Some(metadata) = self.provider_metadata() { + observation = observation.with_provider_metadata(metadata); + } + if let Some(service_tier) = self.effective_service_tier() { + observation = observation.with_service_tier_if_valid(service_tier.to_owned()); + } + Some(observation) + } + + pub(super) fn observe_stream_chunk(&mut self, chunk: &[u8], started_at: Instant) -> bool { + if chunk.is_empty() { + return false; + } + insert_first_timing(&mut self.timings.first_event_ms, started_at) + } + + pub(super) fn observe_timing_signals( + &mut self, + signals: ResponseEventSignals, + started_at: Instant, + ) -> bool { + let mut changed = false; + // 首个非前导输出事件（结构帧也算）开启首字计时； + // 真实语义首字由 first_reasoning_ms / first_text_ms 单独观测。 + if signals.output_start { + changed |= insert_first_timing(&mut self.timings.first_token_ms, started_at); + } + if signals.reasoning_output { + changed |= insert_first_timing(&mut self.timings.first_reasoning_ms, started_at); + } + if signals.text_output { + changed |= insert_first_timing(&mut self.timings.first_text_ms, started_at); + } + changed + } + + pub(super) fn observe_upstream_service_tier(&mut self, service_tier: Option<&str>) -> bool { + let Some(service_tier) = normalize_service_tier(service_tier) else { + return false; + }; + if self.upstream_service_tier.as_deref() == Some(service_tier.as_str()) { + return false; + } + self.upstream_service_tier = Some(service_tier); + true + } + + pub(super) fn effective_service_tier(&self) -> Option<&str> { + self.requested_service_tier + .as_deref() + .or(self.upstream_service_tier.as_deref()) + } + + pub(super) fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool { + let mut changed = false; + for (name, value) in selected_observation_headers(updates) { + if let Some(existing) = self + .rate_limit_headers + .iter_mut() + .find(|(existing_name, _)| existing_name == &name) + { + if existing.1 != value { + existing.1 = value; + changed = true; + } + } else { + self.rate_limit_headers.push((name, value)); + changed = true; + } + } + changed + } + + pub(super) fn merge_client_header(&mut self, name: &str, value: &str) -> bool { + let value = Bytes::copy_from_slice(value.as_bytes()); + if let Some((_, existing)) = self + .response_metadata + .client_headers + .iter_mut() + .find(|(existing_name, _)| existing_name.eq_ignore_ascii_case(name)) + { + if *existing == value { + return false; + } + *existing = value; + return true; + } + self.response_metadata + .client_headers + .push((name.to_owned(), value)); + true + } + + pub(super) fn mark_completed(&mut self, incomplete: bool) -> bool { + let terminal = if incomplete { + OpenAiResponseTerminal::Incomplete + } else { + OpenAiResponseTerminal::Completed + }; + if self.terminal == Some(terminal) { + return false; + } + self.terminal = Some(terminal); + true + } + + pub(super) fn provider_metadata(&self) -> Option { + let mut metadata = Map::new(); + // 观测对象版本信封：读取端按 schemaVersion 解码，缺失视为 v0。 + metadata.insert("schemaVersion".to_owned(), json!(1)); + let effective_model = self + .response_metadata + .effective_model + .as_deref() + .filter(|model| !model.is_empty()) + .unwrap_or(&self.requested_model); + if !effective_model.is_empty() { + metadata.insert( + "effectiveModel".to_owned(), + Value::String(effective_model.to_owned()), + ); + } + metadata.insert( + "modelsEtag".to_owned(), + self.response_metadata + .models_etag + .as_ref() + .map_or(Value::Null, |etag| Value::String(etag.clone())), + ); + metadata.insert( + "reasoningIncluded".to_owned(), + Value::Bool(self.response_metadata.reasoning_included), + ); + metadata.insert("stream".to_owned(), Value::Bool(self.stream)); + metadata.insert("compact".to_owned(), Value::Bool(self.compact)); + metadata.insert("attemptIndex".to_owned(), json!(self.attempt_index)); + metadata.insert( + "attemptAccountId".to_owned(), + Value::String(self.account_id.clone()), + ); + metadata.insert( + "rateLimitHeaders".to_owned(), + json!(self.rate_limit_headers), + ); + metadata.insert( + "upstreamTraceHeaders".to_owned(), + json!(selected_observation_headers( + &self.diagnostics.trace_headers + )), + ); + metadata.insert("requestSummary".to_owned(), self.request_summary.clone()); + if let Some(service_tier) = self.effective_service_tier() { + metadata.insert( + "serviceTier".to_owned(), + Value::String(service_tier.to_owned()), + ); + } + if let Some(service_tier) = &self.requested_service_tier { + metadata.insert( + "requestedServiceTier".to_owned(), + Value::String(service_tier.clone()), + ); + } + if let Some(service_tier) = &self.upstream_service_tier { + metadata.insert( + "upstreamServiceTier".to_owned(), + Value::String(service_tier.clone()), + ); + } + if let Some(status) = self.diagnostics.status_code { + metadata.insert("upstreamStatus".to_owned(), json!(status)); + } + if let Some(decision) = self.metrics.decision { + metadata.insert( + "transportDecision".to_owned(), + Value::String(decision.as_str().to_owned()), + ); + } + if let Some(decision) = self.websocket_pool_decision { + metadata.insert( + "websocketPool".to_owned(), + json!({ "kind": decision.kind() }), + ); + } + if let Some(version) = self.metrics.http_version.as_deref() { + metadata.insert("httpVersion".to_owned(), Value::String(version.to_owned())); + } + if let Some((_, cf_ray)) = self + .diagnostics + .trace_headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("cf-ray")) + && let Some((_, cf_ray)) = selected_observation_header("cf-ray", cf_ray) + { + metadata.insert("cfRay".to_owned(), Value::String(cf_ray)); + } + insert_optional_metadata_millis( + &mut metadata, + "transportDecisionWaitMs", + self.timings.transport_decision_wait_ms, + ); + insert_optional_metadata_millis(&mut metadata, "wsConnectMs", self.timings.connect_ms); + insert_optional_metadata_millis( + &mut metadata, + "upstreamHeadersMs", + self.timings.headers_ms, + ); + insert_optional_metadata_millis(&mut metadata, "firstEventMs", self.timings.first_event_ms); + insert_optional_metadata_millis( + &mut metadata, + "firstReasoningMs", + self.timings.first_reasoning_ms, + ); + insert_optional_metadata_millis(&mut metadata, "firstTextMs", self.timings.first_text_ms); + insert_optional_metadata_millis(&mut metadata, "firstTokenMs", self.timings.first_token_ms); + insert_optional_metadata_millis( + &mut metadata, + "openaiProcessingMs", + self.timings.provider_processing_ms, + ); + if let Some(terminal) = self.terminal { + let incomplete = terminal == OpenAiResponseTerminal::Incomplete; + metadata.insert("completed".to_owned(), Value::Bool(!incomplete)); + metadata.insert("incomplete".to_owned(), Value::Bool(incomplete)); + metadata.insert("eventStatusCode".to_owned(), json!(200_u16)); + } + ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?) + } +} + +pub(super) fn codex_response_observation( + transport: CodexBackendTransport, + diagnostics: &CodexUpstreamDiagnostics, + response_metadata: &CodexResponseMetadata, + metrics: &CodexTransportMetrics, + websocket_pool_decision: Option, + timings: ProviderResponseTimings, +) -> Option { + let mut observation = ProviderResponseObservation::new( + UpstreamTransport::new(actual_transport_name(transport)).ok()?, + ) + .with_timings(timings); + if let Some(version) = metrics + .http_version + .as_deref() + .and_then(UpstreamHttpVersion::parse) + { + observation = observation.with_http_version(version); + } + if let Some(decision) = websocket_pool_decision { + observation = observation.with_websocket_pool(if decision.is_reuse() { + WebSocketPoolKind::Reuse + } else { + WebSocketPoolKind::New + }); + } + // WebSocket opening 的 101 是成功升级事实，不是业务请求的失败 HTTP 状态。 + // opening 明确拒绝仍由 `codex_error_observation` 保存真实上游状态。 + if transport != CodexBackendTransport::WebSocket + && let Some(status_code) = diagnostics.status_code + { + observation = observation.with_status_code(status_code); + } + if let Some(request_id) = diagnostics.request_id.as_deref() { + observation = observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); + } + let client_headers = response_metadata + .client_headers + .iter() + .map(|(name, value)| ProviderResponseHeader::new(name.to_owned(), value.clone())) + .collect(); + observation = observation.with_client_headers(client_headers); + Some(observation) +} + +pub(super) 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, + None, + openai_response_timings(transport_metrics, &CodexResponseMetadata::default()), + )? + .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() { + observation = + observation.with_request_id(OpaqueUpstreamValue::new(request_id.to_owned())); + } + } + _ => {} + } + Some(observation) +} + +pub(super) fn openai_response_timings( + metrics: &CodexTransportMetrics, + response_metadata: &CodexResponseMetadata, +) -> ProviderResponseTimings { + 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), + provider_processing_ms: openai_processing_ms(response_metadata), + ..ProviderResponseTimings::default() + } +} + +pub(super) fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option { + response_metadata + .client_headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms")) + .and_then(|(_, value)| std::str::from_utf8(value).ok()) + .and_then(|value| value.trim().parse::().ok()) +} + +pub(super) fn openai_response_request_summary( + request: &CodexResponsesRequest, + transport: CodexBackendTransport, +) -> Value { + let body = request.body(); + let input = body.get("input"); + let tools = body.get("tools"); + let semantics = request.semantics(); + json!({ + "model": request.model(), + "stream": request.stream(), + "store": request.store(), + "compact": semantics.compact, + "requestKind": semantics.request_kind, + "subagentKind": semantics.subagent_kind, + "reasoningPreset": semantics.reasoning_preset, + "transport": actual_transport_name(transport), + "inputType": json_value_kind(input), + "inputItemsCount": input.and_then(Value::as_array).map(Vec::len), + "toolsType": json_value_kind(tools), + "toolsCount": tools.and_then(Value::as_array).map(Vec::len), + "topLevelFields": body.keys().cloned().collect::>(), + "previousResponseIdPresent": request.previous_response_id().is_some(), + "serviceTier": request.service_tier(), + "localTransport": { + "useWebsocket": request.use_websocket, + "forceHttpSse": request.force_http_sse, + }, + }) +} + +pub(super) const fn json_value_kind(value: Option<&Value>) -> &'static str { + match value { + None => "missing", + Some(Value::Null) => "null", + Some(Value::Bool(_)) => "boolean", + Some(Value::Number(_)) => "number", + Some(Value::String(_)) => "string", + Some(Value::Array(_)) => "array", + Some(Value::Object(_)) => "object", + } +} + +pub(super) fn selected_observation_headers(headers: &[(String, String)]) -> Vec<(String, String)> { + headers + .iter() + .filter_map(|(name, value)| selected_observation_header(name, value)) + .collect() +} + +pub(super) fn selected_observation_header(name: &str, value: &str) -> Option<(String, String)> { + let name = name.trim().to_ascii_lowercase(); + (!name.is_empty()).then(|| (name, value.to_owned())) +} + +pub(super) fn insert_first_timing(target: &mut Option, started_at: Instant) -> bool { + if target.is_some() { + return false; + } + let elapsed = u64::try_from(started_at.elapsed().as_millis()) + .unwrap_or(u64::MAX) + .max(1); + *target = Some(elapsed); + true +} + +pub(super) fn insert_optional_metadata_millis( + metadata: &mut Map, + name: &str, + value: Option, +) { + if let Some(value) = value { + metadata.insert(name.to_owned(), json!(value)); + } +} + +pub(super) async fn take_rate_limit_updates( + updates: Option<&CodexRateLimitUpdates>, +) -> Vec { + let Some(updates) = updates else { + return Vec::new(); + }; + std::mem::take(&mut *updates.lock().await) +} + +pub(super) fn rate_limit_update_headers(updates: &[ParsedRateLimits]) -> Vec<(String, String)> { + updates + .iter() + .flat_map(rate_limits_to_header_pairs) + .collect() +} + +pub(super) fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool { + events + .iter() + .find_map(|event| { + let wire = event.wire_event()?; + let event_type = wire + .event_type() + .or_else(|| wire.data().get("type").and_then(Value::as_str)); + match event_type { + Some("response.incomplete") => Some(true), + Some("response.completed") => Some( + wire.data() + .pointer("/response/status") + .and_then(Value::as_str) + == Some("incomplete"), + ), + _ => None, + } + }) + .unwrap_or(false) +} + +pub(super) async fn synchronize_passive_quota( + quota: &CodexCredentialQuotaService, + account: &ProviderAccount, + rate_limits: &[ParsedRateLimits], +) { + if rate_limits.is_empty() { + return; + } + if let Err(error) = quota + .synchronize_passive_rate_limits(account, rate_limits) + .await + { + tracing::warn!( + account_id = %account.id(), + error = %error, + "OpenAI passive quota synchronization failed" + ); + } +} + +pub(super) async fn synchronize_passive_quota_headers( + quota: &CodexCredentialQuotaService, + account: &ProviderAccount, + headers: &[(String, String)], +) { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return; + }; + synchronize_passive_quota(quota, account, std::slice::from_ref(&rate_limits)).await; +} + +pub(super) const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str { + match transport { + CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT, + CodexBackendTransport::HttpJson => HTTP_JSON_TRANSPORT, + CodexBackendTransport::WebSocket => WEBSOCKET_TRANSPORT, + } +} + +pub(super) fn nonnegative_millis(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + +pub(super) fn compile_model_capabilities(model: &CodexCatalogModel) -> ProviderModelCapabilities { + // Catalog membership is enough to publish Generate. `supported_in_api` is advisory; + // the upstream response is authoritative for normal and diagnostic requests. + let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) + .with_upstream_feature_validation(); + ProviderModelCapabilities::new(model.request_model().clone(), capabilities) + .with_presentation(codex_model_presentation(model)) +} + +pub(super) fn codex_model_presentation(model: &CodexCatalogModel) -> ModelPresentation { + let capabilities = model.capabilities(); + let reasoning_efforts = capabilities.reasoning_efforts().to_vec(); + // Codex 目录不声明默认 effort；有 medium 时对齐官方 picker 缺省，否则取首项。 + let default_reasoning = reasoning_efforts + .iter() + .find(|effort| effort.as_str() == "medium") + .or_else(|| reasoning_efforts.first()) + .cloned(); + let hidden = matches!( + model.metadata().visibility(), + Some(CodexCatalogVisibility::Hide | CodexCatalogVisibility::None) + ); + + ModelPresentation::new( + Some(model.display_name().to_owned()), + model.metadata().description().map(str::to_owned), + ) + .with_reasoning(default_reasoning, reasoning_efforts) + .with_context_window_tokens( + model + .limits() + .context_window_tokens() + .map(std::num::NonZeroU64::get), + ) + .with_image_input(capabilities.image_input() == CodexCatalogCapabilityEvidence::DeclaredNative) + // Codex Responses 工具协议随 API 支持一并可用；只有明确不支持才关闭。 + .with_agent_tools( + capabilities.responses_api() != CodexCatalogCapabilityEvidence::DeclaredUnsupported, + capabilities.parallel_tool_calls() == CodexCatalogCapabilityEvidence::DeclaredNative, + ) + .with_search_tool(capabilities.web_search() == CodexCatalogCapabilityEvidence::DeclaredNative) + .with_image_detail_original( + capabilities.image_detail_original() == CodexCatalogCapabilityEvidence::DeclaredNative, + ) + .with_verbosity(capabilities.verbosity() == CodexCatalogCapabilityEvidence::DeclaredNative) + .with_hidden(hidden) +} + +pub(super) fn selected_transport(request: &CodexResponsesRequest) -> CodexProviderTransport { + if request.force_http_sse { + CodexProviderTransport::HttpOnly + } else { + CodexProviderTransport::PreferWebSocket + } +} + +pub(super) fn apply_transport( + request: &mut CodexResponsesRequest, + transport: CodexProviderTransport, +) { + match transport { + CodexProviderTransport::HttpOnly => { + request.force_http_sse = true; + request.use_websocket = false; + } + CodexProviderTransport::PreferWebSocket => { + request.force_http_sse = false; + request.use_websocket = true; + } + } +} + +pub(super) const fn transport_name(transport: CodexProviderTransport) -> &'static str { + match transport { + CodexProviderTransport::HttpOnly => HTTP_SSE_TRANSPORT, + CodexProviderTransport::PreferWebSocket => WEBSOCKET_TRANSPORT, + } +} + +pub(super) const fn accepts_backend_transport( + transport: CodexProviderTransport, + actual: CodexBackendTransport, +) -> bool { + match transport { + CodexProviderTransport::HttpOnly => matches!(actual, CodexBackendTransport::HttpSse), + CodexProviderTransport::PreferWebSocket => true, + } +} + +pub(super) fn codex_request_context<'a>( + request: &'a CodexResponsesRequest, + request_id: &'a str, + account: &'a ProviderAccount, + installation_id: &'a str, + authorization: &'a SecretString, + cookie_header: Option<&'a SecretString>, + account_selection: CodexAccountSelectionTelemetry<'a>, +) -> CodexRequestContext<'a> { + CodexRequestContext { + authorization: authorization.expose_secret(), + account_id: account.upstream_account_id(), + request_id, + turn_state: request.turn_state.as_deref(), + turn_metadata: request.turn_metadata.as_deref(), + beta_features: request.beta_features.as_deref(), + include_timing_metrics: request.include_timing_metrics.as_deref(), + version: request.version.as_deref(), + codex_window_id: request.codex_window_id.as_deref(), + parent_thread_id: request.parent_thread_id.as_deref(), + cookie_header: cookie_header.map(ExposeSecret::expose_secret), + installation_id: Some(installation_id), + session_id: request.client_session_id.as_deref(), + thread_id: request.client_thread_id.as_deref(), + client_request_id: request.client_request_id.as_deref(), + turn_id: request.client_turn_id.as_deref(), + account_selection, + } +} + +pub(super) fn build_cookie_header( + cookies: &[RuntimeCodexCookie], +) -> Result, ProviderError> { + if cookies.is_empty() { + return Ok(None); + } + let mut header = String::new(); + for cookie in cookies { + let value = cookie.value.expose_secret(); + if !valid_cookie_name(&cookie.name) + || value.is_empty() + || value.chars().any(char::is_control) + || value.contains(';') + { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )); + } + if !header.is_empty() { + header.push_str("; "); + } + header.push_str(&cookie.name); + header.push('='); + header.push_str(value); + if header.len() > MAX_COOKIE_HEADER_BYTES { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )); + } + } + Ok(Some(SecretString::from(header))) +} + +pub(super) fn valid_cookie_name(name: &str) -> bool { + !name.is_empty() + && name.len() <= 256 + && name.bytes().all(|byte| { + byte.is_ascii_alphanumeric() + || matches!( + byte, + b'!' | b'#' + | b'$' + | b'%' + | b'&' + | b'\'' + | b'*' + | b'+' + | b'-' + | b'.' + | b'^' + | b'_' + | b'`' + | b'|' + | b'~' + ) + }) +} + +pub(super) fn map_request_error(error: CodexRequestEncodeError) -> ProviderError { + let kind = match error { + CodexRequestEncodeError::InvalidProtocolPayload => ProviderErrorKind::InvalidRequest, + }; + provider_error(kind, UpstreamSendState::NotSent) +} + +pub(super) fn map_selection_error(error: CredentialSelectionError) -> ProviderError { + match error { + CredentialSelectionError::CapacityUnavailable { retry_after } => { + let error = provider_error( + ProviderErrorKind::AccountCapacityUnavailable, + UpstreamSendState::NotSent, + ); + match retry_after { + Some(retry) => error.with_retry_after(retry), + None => error, + } + } + CredentialSelectionError::NoEligibleCredential => provider_error( + ProviderErrorKind::NoEligibleAccount, + UpstreamSendState::NotSent, + ), + CredentialSelectionError::InvalidCredential + | CredentialSelectionError::Store + | CredentialSelectionError::Coordinator + | CredentialSelectionError::CookiePolicy => provider_error( + ProviderErrorKind::ProviderInfrastructureUnavailable, + UpstreamSendState::NotSent, + ), + } +} diff --git a/backend/crates/providers/openai/src/provider/workers.rs b/backend/crates/providers/openai/src/provider/workers.rs new file mode 100644 index 000000000..84a6a03fd --- /dev/null +++ b/backend/crates/providers/openai/src/provider/workers.rs @@ -0,0 +1,260 @@ +//! OpenAI Provider 向 Host 贡献的后台 worker。 + +use super::*; + +pub(super) const WORKER_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +pub(super) const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60); +pub(super) const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); +pub(super) const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); +pub(super) const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); +pub(super) const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release"; +pub(super) const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag"; + +pub(crate) fn worker_contributions( + refresh: Arc, + quota: Arc, + catalog: Arc, + quota_refresh_policy: CodexQuotaRefreshPolicy, + oauth_refresh_enabled: bool, + desktop_release: Arc, +) -> 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)?; + let mut contributions = Vec::new(); + if oauth_refresh_enabled { + contributions.push(WorkerContribution::Registration(scheduled_registration( + refresh_id, + OAUTH_REFRESH_INTERVAL, + Box::new(OpenAiOAuthRefreshTask { service: refresh }), + )?)); + } + contributions.extend([ + WorkerContribution::Registration(scheduled_registration( + quota_id, + quota_refresh_policy.interval(), + Box::new(OpenAiQuotaTask { + quota, + catalog: Arc::clone(&catalog), + }), + )?), + WorkerContribution::Registration(WorkerRegistration::try_new( + etag_id, + WorkerRunnable::Daemon { + restart: DaemonRestartPolicy::try_new( + WORKER_INITIAL_BACKOFF, + WORKER_MAXIMUM_BACKOFF, + )?, + task: Box::new(OpenAiCatalogEtagTask { catalog }), + }, + )?), + WorkerContribution::Registration(scheduled_registration( + desktop_release_id, + APPCAST_POLL_INTERVAL, + Box::new(OpenAiDesktopReleaseTask { + service: desktop_release, + }), + )?), + ]); + Ok(contributions) +} + +pub(super) 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, + }, + ) +} + +pub(super) 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(|error| { + tracing::error!(error = %error, "OpenAI OAuth refresh cycle failed"); + WorkerTaskError::safe("OpenAI OAuth refresh failed") + })?; + let mut refreshed = 0_u64; + let mut invalidated = 0_u64; + let mut banned = 0_u64; + let mut transient = 0_u64; + let mut lease_unavailable = 0_u64; + let mut stale = 0_u64; + let mut failed = 0_u64; + let mut transient_accounts = Vec::new(); + let mut failed_accounts = Vec::new(); + for outcome in &outcomes { + match outcome { + CodexCredentialRefreshOutcome::Refreshed { .. } => refreshed += 1, + CodexCredentialRefreshOutcome::Invalidated { .. } => invalidated += 1, + CodexCredentialRefreshOutcome::Banned { .. } => banned += 1, + CodexCredentialRefreshOutcome::Transient { account_id } => { + transient += 1; + transient_accounts.push(account_id); + } + CodexCredentialRefreshOutcome::LeaseUnavailable { .. } => { + lease_unavailable += 1; + } + CodexCredentialRefreshOutcome::Stale { .. } => stale += 1, + CodexCredentialRefreshOutcome::Failed { account_id } => { + failed += 1; + failed_accounts.push(account_id); + } + } + } + if !outcomes.is_empty() { + tracing::info!( + refreshed, + invalidated, + banned, + transient, + lease_unavailable, + stale, + failed, + "OpenAI OAuth refresh cycle completed" + ); + } + if transient > 0 || failed > 0 { + tracing::warn!( + refreshed, + invalidated, + banned, + transient, + lease_unavailable, + stale, + failed, + transient_accounts = ?transient_accounts, + failed_accounts = ?failed_accounts, + "OpenAI OAuth refresh cycle contained operational failures" + ); + } + Ok(()) + }) + } +} + +pub(super) struct OpenAiQuotaTask { + quota: Arc, + catalog: Arc, +} + +pub(super) struct OpenAiCatalogEtagTask { + catalog: Arc, +} + +pub(super) 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 { + if context.cancellation().is_cancelled() { + return Ok(()); + } + let mut failures = false; + match self.quota.synchronize().await { + Ok(summary) if summary.has_operational_failures() => { + tracing::warn!( + updated = summary.updated, + exhausted = summary.exhausted, + banned = summary.banned, + transient = summary.transient, + stale = summary.stale, + "OpenAI quota cycle contained operational failures" + ); + } + Ok(_) => {} + Err(error) => { + failures = true; + tracing::warn!( + error = %error, + "OpenAI quota synchronization failed" + ); + } + } + if context.cancellation().is_cancelled() { + return Ok(()); + } + match self.catalog.refresh_catalogs().await { + Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {} + Err(error) => { + failures = true; + tracing::warn!(error = %error, "OpenAI model catalog refresh failed"); + } + } + if failures { + Err(WorkerTaskError::safe( + "OpenAI quota or catalog synchronization failed", + )) + } else { + Ok(()) + } + }) + } +} + +impl DaemonTask for OpenAiCatalogEtagTask { + fn run( + &self, + cancellation: gateway_core::lifecycle::CancellationToken, + ) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + loop { + tokio::select! { + () = cancellation.cancelled() => return Ok(()), + () = self.catalog.wait_for_etag_refresh() => {}, + }; + if let Err(error) = self.catalog.refresh().await { + tracing::warn!( + error = %error, + "OpenAI model catalog ETag refresh failed" + ); + } + } + }) + } +} diff --git a/backend/crates/providers/openai/tests/transport/websocket.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs similarity index 100% rename from backend/crates/providers/openai/tests/transport/websocket.rs rename to backend/crates/providers/openai/tests/transport/websocket/mod.rs diff --git a/backend/crates/providers/xai/src/provider.rs b/backend/crates/providers/xai/src/provider.rs deleted file mode 100644 index 0efc6ad64..000000000 --- a/backend/crates/providers/xai/src/provider.rs +++ /dev/null @@ -1,2307 +0,0 @@ -//! 官方 Grok Build 会话的 `gateway-core` Provider adapter。 - -use std::collections::{BTreeMap, BTreeSet}; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant, SystemTime}; - -use async_trait::async_trait; -use futures::{StreamExt, future::BoxFuture}; -use gateway_core::account::{ - AccountEligibilityPolicy, AccountFeedbackStats, ProviderAccount, ProviderAccountId, - ProviderAccountStore, -}; -use gateway_core::engine::continuation::ContinuationBinding; -use gateway_core::engine::provider::{ - EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, - ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, - ProviderSelectionObservation, ProviderStream, -}; -use gateway_core::engine::{AttemptContext, ContinuationAttempt}; -use gateway_core::error::{ - ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind, -}; -use gateway_core::event::{ - GatewayEvent, ProviderEvent, ProviderResponseMetadata, ProviderResponseObservation, - ProviderResponseTimings, ResponseMeta, -}; -use gateway_core::operation::{ - Feature, GenerateRequest, Operation, OperationKind, ProviderSessionState, -}; -use gateway_core::routing::{ - ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, SupportLevel, - UpstreamModelId, -}; -use gateway_core::task::{ - ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId, - WorkerKind, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, - WorkerTaskError, -}; -use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; -use gateway_protocol::openai::codex_responses_request_semantics; -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; -use url::Url; - -use crate::XaiWireProfileState; -use crate::credential::{ - GrokCredentialCatalogError, GrokCredentialCatalogService, GrokCredentialQuotaService, - GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshOutcome, - GrokCredentialRefreshService, GrokQuotaError, -}; -use crate::reasoning_replay::{ - GrokReasoningReplay, GrokReasoningReplayCapture, GrokReasoningReplayKey, - valid_reasoning_ciphertext, -}; -use crate::transport::canonical::GrokCanonicalDecoder; -use crate::transport::config::XAI_PROVIDER_NAME; -use crate::transport::headers::{GrokClientIdentity, build_grok_headers}; -use crate::transport::profile::{GROK_CLI_RELEASE_POLL_INTERVAL, GrokCliReleaseService}; -use crate::transport::{ - GROK_RESPONSES_URL, GrokCompactionDecodeError, GrokCompactionRequest, - GrokCompactionSummaryDecoder, GrokCredentialFailure, GrokInferenceChunkStream, - GrokInferenceRequest, GrokInferenceResponse, GrokInferenceTransport, - GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportMetrics, - GrokProviderConfigError, GrokQuotaFailureKind, GrokRequestEncodeError, GrokResponsesRequest, - GrokSessionAffinityKey, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, - SelectedGrokSession, classify_grok_quota_failure, -}; -use crate::{GrokCatalogCapabilityEvidence, GrokCatalogModel}; - -const HTTP_SSE_TRANSPORT: &str = "http_sse"; -const DEFAULT_GROK_MODEL: &str = "grok-4.5"; -const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024; -const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096; -const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed"; -const RESPONSE_NOT_FOUND_CODE: &str = "not_found"; - -/// 官方 Grok Build Provider；会话选择与 HTTP SSE transport 均由外部注入。 -/// -/// 每次调用只选择一个 OAuth 会话。仅当 xAI 明确拒绝历史 reasoning 密文时， -/// 允许在同账号、同凭据、同会话绑定上剥离密文后有界重试一次。凭据轮换、 -/// endpoint fallback 以及公开 xAI API key 推理都不在该 adapter 内。 -pub struct GrokBuildProvider { - selector: Arc, - transport: Arc, - catalog: Arc, - credential_recovery: Arc, - account_feedback: Arc, - client_identity: GrokClientIdentity, - reasoning_replay: GrokReasoningReplay, - wire_profile: XaiWireProfileState, - responses_url: Url, -} - -impl GrokBuildProvider { - /// 在显式的会话与 transport 边界上创建 Provider。 - pub fn new( - selector: Arc, - transport: Arc, - catalog: Arc, - credential_recovery: Arc, - account_feedback: Arc, - wire_profile: XaiWireProfileState, - ) -> Result { - let responses_url = Url::parse(GROK_RESPONSES_URL) - .map_err(|_| GrokProviderConfigError::InvalidResponsesUrl)?; - Ok(Self { - selector, - transport, - catalog, - credential_recovery, - account_feedback, - client_identity: GrokClientIdentity::new(), - reasoning_replay: GrokReasoningReplay::new(), - wire_profile, - responses_url, - }) - } -} - -#[async_trait] -impl Provider for GrokBuildProvider { - fn name(&self) -> &'static str { - XAI_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 payload = request.protocol_payload(); - if payload.protocol() != "openai" { - return ProviderRequestObservation::default(); - } - let semantics = codex_responses_request_semantics(payload.body(), payload.context()); - ProviderRequestObservation { - reasoning_effort: semantics.reasoning_effort, - 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, - ) -> Result, ProviderError> { - let Ok(models) = self.catalog.query_models().await else { - return Ok(vec![default_grok_model_capabilities()?]); - }; - if models.is_empty() { - return Ok(vec![default_grok_model_capabilities()?]); - } - Ok(models - .into_iter() - .map(compile_grok_model_capabilities) - .collect()) - } - - async fn execute( - &self, - request: ProviderRequest, - context: AttemptContext, - ) -> Result { - let candidate = request.candidate(); - if candidate.provider().as_str() != XAI_PROVIDER_NAME { - return Err(provider_error( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - )); - } - preflight_context(&context)?; - - match request.operation() { - Operation::Generate(generate) => { - self.execute_generate(generate, candidate, context).await - } - _ => Err(provider_error( - ProviderErrorKind::Unsupported, - UpstreamSendState::NotSent, - )), - } - } -} - -impl GrokBuildProvider { - async fn execute_generate( - &self, - generate: &GenerateRequest, - candidate: &ProviderCandidate, - context: AttemptContext, - ) -> Result { - if crate::transport::compaction::has_terminal_compaction_trigger(generate) { - return self.execute_compaction(generate, candidate, context).await; - } - let upstream_model = candidate_upstream_model(candidate)?; - let previous_session = decode_xai_session_state(generate)?; - let continuation_account = continuation_account(&context, previous_session.as_ref())?; - let mut upstream_request = GrokResponsesRequest::encode( - generate, - upstream_model.as_str(), - context.client_api_key_ref(), - ) - .map_err(map_request_error)?; - let wire_upstream_model = UpstreamModelId::new( - upstream_request - .upstream_model() - .ok_or_else(protocol_not_sent)? - .to_owned(), - ) - .map_err(|_| protocol_not_sent())?; - let request_input = upstream_request.input_items(); - if let Some(previous) = previous_session.as_ref() { - upstream_request.inherit_session(previous.session_id.as_deref()); - } - let selection_started_at = Instant::now(); - let selected = select_grok_session( - self.selector.as_ref(), - candidate, - &wire_upstream_model, - &context, - continuation_account, - upstream_request.affinity().cloned(), - ) - .await?; - let account_selection_wait_ms = - u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); - // 首字计时的起点：账号选择完成之后、上游建立之前。 - let output_started_at = Instant::now(); - apply_continuation( - &mut upstream_request, - previous_session.as_ref(), - &context, - selected.account_id(), - request_input.as_slice(), - )?; - let inherited_replay_session_id = context - .continuation() - .is_none() - .then(|| { - previous_session - .as_ref() - .and_then(|state| state.session_id.as_deref()) - }) - .flatten(); - let reasoning_replay_key = upstream_request - .reasoning_replay_session_id() - .or(inherited_replay_session_id) - .and_then(|session_id| { - self.reasoning_replay.key( - wire_upstream_model.as_str(), - session_id, - selected.account_id().as_str(), - ) - }); - if !upstream_request.has_previous_response_id() - && let (Some(key), Some(input)) = ( - reasoning_replay_key.as_ref(), - upstream_request.replay_input_items(), - ) - && let Some(input) = self.reasoning_replay.apply(key, &input) - { - upstream_request - .set_replay_input(input) - .map_err(map_request_error)?; - } - let reasoning_replay_capture = - reasoning_replay_key.map(|key| self.reasoning_replay.capture(key)); - let session_capture = (!matches!( - context.continuation(), - Some(ContinuationBinding::Pinned(_) | ContinuationBinding::External(_)) - ) || previous_session.is_some()) - .then(|| GrokSessionCapture { - previous: previous_session, - request_input, - account_id: selected.account_id().as_str().to_owned(), - session_id: upstream_request.session_id().map(str::to_owned), - output_items: BTreeMap::new(), - }); - let selected = Arc::new(selected); - let allows_account_state_mutation = selected.allows_account_state_mutation(); - let metadata = provider_call_metadata(candidate, &selected, account_selection_wait_ms)?; - let events = cold_http_sse_stream( - Arc::clone(&self.selector), - Arc::clone(&self.transport), - GrokStreamAttempt { - client_identity: self.client_identity.clone(), - wire_profile: self.wire_profile.clone(), - credential_recovery: Arc::clone(&self.credential_recovery), - responses_url: self.responses_url.clone(), - request: upstream_request, - upstream_model: wire_upstream_model, - context, - session: Arc::clone(&selected), - output_started_at, - session_capture, - reasoning_replay_capture, - }, - ); - let stream = ProviderStream::new(metadata, events, selected); - Ok(if allows_account_state_mutation { - stream.with_account_feedback(Arc::clone(&self.account_feedback)) - } else { - stream - }) - } - - async fn execute_compaction( - &self, - generate: &GenerateRequest, - candidate: &ProviderCandidate, - context: AttemptContext, - ) -> Result { - if context.continuation().is_some() - || context.continuation_attempt() != ContinuationAttempt::None - { - return Err(provider_error( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - )); - } - let previous_session = decode_xai_session_state(generate)?; - let operation_account = previous_session - .as_ref() - .map(|previous| ProviderAccountId::new(previous.account_id.clone())) - .transpose() - .map_err(|_| protocol_not_sent())?; - let inherited_session_id = previous_session - .as_ref() - .and_then(|previous| previous.session_id.clone()); - let upstream_model = candidate_upstream_model(candidate)?; - let upstream_request = GrokCompactionRequest::encode( - generate, - upstream_model.as_str(), - context.client_api_key_ref(), - ) - .map_err(map_request_error)?; - let wire_upstream_model = UpstreamModelId::new( - upstream_request - .upstream_model() - .ok_or_else(protocol_not_sent)? - .to_owned(), - ) - .map_err(|_| protocol_not_sent())?; - let explicit_replay_session_id = upstream_request - .reasoning_replay_session_id() - .map(str::to_owned); - let upstream_session_id = inherited_session_id - .clone() - .or_else(|| explicit_replay_session_id.clone()); - let selection_started_at = Instant::now(); - let selected = Arc::new( - select_grok_session( - self.selector.as_ref(), - candidate, - &wire_upstream_model, - &context, - operation_account, - upstream_request.affinity().cloned(), - ) - .await?, - ); - let account_selection_wait_ms = - u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); - let reasoning_replay_key = explicit_replay_session_id - .as_deref() - .or(inherited_session_id.as_deref()) - .and_then(|session_id| { - self.reasoning_replay.key( - wire_upstream_model.as_str(), - session_id, - selected.account_id().as_str(), - ) - }); - let allows_account_state_mutation = selected.allows_account_state_mutation(); - let metadata = provider_call_metadata(candidate, &selected, account_selection_wait_ms)?; - let events = cold_compaction_http_sse_stream( - Arc::clone(&self.selector), - Arc::clone(&self.transport), - GrokCompactionStreamAttempt { - client_identity: self.client_identity.clone(), - wire_profile: self.wire_profile.clone(), - credential_recovery: Arc::clone(&self.credential_recovery), - responses_url: self.responses_url.clone(), - request: upstream_request, - upstream_model: wire_upstream_model, - upstream_session_id, - context, - session: Arc::clone(&selected), - reasoning_replay: self.reasoning_replay.clone(), - reasoning_replay_key, - }, - ); - let stream = ProviderStream::new(metadata, events, selected); - Ok(if allows_account_state_mutation { - stream.with_account_feedback(Arc::clone(&self.account_feedback)) - } else { - stream - }) - } -} - -async fn select_grok_session( - selector: &dyn GrokSessionSelector, - candidate: &ProviderCandidate, - wire_upstream_model: &UpstreamModelId, - context: &AttemptContext, - operation_account: Option, - affinity: Option, -) -> Result { - let required_account = context.required_account().cloned().or(operation_account); - let selection = GrokSessionSelection::new( - wire_upstream_model.clone(), - context.excluded_accounts().clone(), - required_account.clone(), - context.account_selection_policy(), - context.deadline(), - Arc::clone(candidate.account_scope()), - context.client_api_key_ref().clone(), - ) - .with_eligibility_policy(if context.is_diagnostic_required_account() { - AccountEligibilityPolicy::BypassForDiagnostic - } else { - AccountEligibilityPolicy::Enforce - }) - .with_affinity(affinity); - let selection_deadline = remaining(context.deadline()) - .ok_or_else(|| provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent))?; - let cancellation = context.cancellation().clone(); - let selected = tokio::select! { - biased; - _ = cancellation.cancelled() => Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::NotSent, - )), - _ = tokio::time::sleep(selection_deadline) => Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::NotSent, - )), - selected = selector.select(selection) => selected.map_err(map_selection_error), - }?; - if context.excluded_accounts().contains(selected.account_id()) - || required_account - .as_ref() - .is_some_and(|required| required != selected.account_id()) - { - return Err(provider_error( - ProviderErrorKind::Protocol, - UpstreamSendState::NotSent, - )); - } - Ok(selected) -} - -fn provider_call_metadata( - candidate: &ProviderCandidate, - selected: &SelectedGrokSession, - account_selection_wait_ms: u64, -) -> Result { - Ok(ProviderCallMetadata::new( - ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?, - candidate_upstream_model(candidate)?.clone(), - selected.account_id().clone(), - UpstreamTransport::new(HTTP_SSE_TRANSPORT).map_err(|_| protocol_not_sent())?, - ) - .with_selection_observation(ProviderSelectionObservation::new( - account_selection_wait_ms, - selected.capacity_snapshot(), - ))) -} - -fn candidate_upstream_model( - candidate: &ProviderCandidate, -) -> Result<&UpstreamModelId, ProviderError> { - candidate.upstream_model().ok_or_else(protocol_not_sent) -} - -fn support(evidence: GrokCatalogCapabilityEvidence) -> SupportLevel { - match evidence { - GrokCatalogCapabilityEvidence::DeclaredNative => SupportLevel::Native, - GrokCatalogCapabilityEvidence::DeclaredUnsupported => SupportLevel::Unsupported, - GrokCatalogCapabilityEvidence::Unknown => SupportLevel::Unknown, - } -} - -fn tool_support(evidence: GrokCatalogCapabilityEvidence) -> SupportLevel { - match evidence { - GrokCatalogCapabilityEvidence::DeclaredNative => SupportLevel::Native, - GrokCatalogCapabilityEvidence::DeclaredUnsupported => SupportLevel::Unsupported, - // catalog 省略该可选字段时，Grok Build 的 Responses 工具协议仍可用； - // 请求 adapter 会在发送前规范化仅客户端侧的工具结构。 - GrokCatalogCapabilityEvidence::Unknown => SupportLevel::Emulated, - } -} - -fn compile_grok_model_capabilities(model: GrokCatalogModel) -> ProviderModelCapabilities { - let mut operations = BTreeSet::new(); - if model.capabilities().responses_api() == GrokCatalogCapabilityEvidence::DeclaredNative { - operations.insert(OperationKind::Generate); - } - let capabilities = ModelCapabilities::new( - operations, - model - .limits() - .max_output_tokens() - .map(std::num::NonZeroU64::get), - ) - .with_upstream_feature_validation() - .with_feature( - Feature::Reasoning, - support(model.capabilities().reasoning_effort()), - ) - .with_feature( - Feature::Tools, - tool_support(model.capabilities().streaming_tool_calls()), - ) - .with_feature(Feature::Vision, SupportLevel::Unknown) - .with_feature(Feature::JsonSchema, SupportLevel::Unknown) - .with_feature(Feature::NativeContinuation, SupportLevel::Native); - ProviderModelCapabilities::new(model.request_model().clone(), capabilities) - .with_presentation(grok_model_presentation(&model)) -} - -fn default_grok_model_capabilities() -> Result { - let model = UpstreamModelId::new(DEFAULT_GROK_MODEL.to_owned()) - .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; - let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) - .with_upstream_feature_validation(); - Ok(ProviderModelCapabilities::new(model, capabilities) - .with_presentation(default_grok_model_presentation())) -} - -fn default_grok_model_presentation() -> ModelPresentation { - ModelPresentation::new( - Some("Grok 4.5".to_owned()), - Some("xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()), - ) - .with_context_window_tokens(Some(500_000)) - .with_image_input(true) - .with_agent_tools(true, true) -} - -fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation { - let slug = model.request_model().as_str(); - let known_grok_4_5 = is_known_grok_4_5_model(slug); - let reasoning_evidence = model.capabilities().reasoning_effort(); - let catalog_reasoning_efforts = model - .capabilities() - .reasoning_efforts() - .iter() - .map(|effort| effort.as_str().to_owned()) - .collect::>(); - let catalog_default_reasoning = model - .capabilities() - .default_reasoning_effort() - .map(|effort| effort.as_str().to_owned()); - let (default_reasoning, reasoning_efforts) = match reasoning_evidence { - GrokCatalogCapabilityEvidence::DeclaredNative if !catalog_reasoning_efforts.is_empty() => { - let default_reasoning = catalog_default_reasoning - .filter(|default| { - catalog_reasoning_efforts - .iter() - .any(|effort| effort == default) - }) - .or_else(|| catalog_reasoning_efforts.first().cloned()); - (default_reasoning, catalog_reasoning_efforts) - } - GrokCatalogCapabilityEvidence::DeclaredNative => (None, Vec::new()), - GrokCatalogCapabilityEvidence::DeclaredUnsupported => { - (Some("none".to_owned()), vec!["none".to_owned()]) - } - GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()), - }; - let context_window_tokens = model - .limits() - .context_window_tokens() - .map(std::num::NonZeroU64::get) - .or(known_grok_4_5.then_some(500_000)); - let tool_evidence = model.capabilities().streaming_tool_calls(); - - ModelPresentation::new( - model - .display_name() - .map(str::to_owned) - .or_else(|| known_grok_4_5.then(|| "Grok 4.5".to_owned())), - model - .metadata() - .description() - .map(str::to_owned) - .or_else(|| { - known_grok_4_5 - .then(|| "xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()) - }), - ) - .with_reasoning(default_reasoning, reasoning_efforts) - .with_context_window_tokens(context_window_tokens) - .with_image_input(known_grok_4_5) - .with_agent_tools( - tool_evidence != GrokCatalogCapabilityEvidence::DeclaredUnsupported, - tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative - || (known_grok_4_5 && tool_evidence == GrokCatalogCapabilityEvidence::Unknown), - ) - .with_search_tool( - model.capabilities().backend_search() == GrokCatalogCapabilityEvidence::DeclaredNative, - ) - .with_hidden(model.metadata().hidden().unwrap_or(false)) -} - -fn is_known_grok_4_5_model(slug: &str) -> bool { - matches!( - slug, - DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest" - ) -} - -#[derive(Clone, Serialize, Deserialize)] -struct XaiSessionState { - account_id: String, - session_id: Option, - transcript: Vec, -} - -#[derive(Clone, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -enum XaiReplayItem { - ClientInput(Value), - SanitizedOutput(Value), - AccountOutput { account_id: String, item: Value }, -} - -struct GrokSessionCapture { - previous: Option, - request_input: Vec, - account_id: String, - session_id: Option, - output_items: BTreeMap, -} - -fn decode_xai_session_state( - request: &GenerateRequest, -) -> Result, ProviderError> { - request - .provider_session_state(XAI_PROVIDER_NAME) - .map(|state| { - let payload = Value::Object(state.payload().clone()); - if serde_json::to_vec(&payload) - .map_err(|_| protocol_not_sent())? - .len() - > XAI_SESSION_STATE_MAX_BYTES - { - return Err(protocol_not_sent()); - } - serde_json::from_value(payload).map_err(|_| protocol_not_sent()) - }) - .transpose() -} - -fn encode_xai_session_state( - state: XaiSessionState, -) -> Result, ProviderError> { - let value = serde_json::to_value(state).map_err(|_| protocol_sent())?; - if serde_json::to_vec(&value) - .map_err(|_| protocol_sent())? - .len() - > XAI_SESSION_STATE_MAX_BYTES - { - return Ok(None); - } - let Value::Object(payload) = value else { - return Err(protocol_sent()); - }; - ProviderSessionState::new(XAI_PROVIDER_NAME, payload) - .map(Some) - .map_err(|_| protocol_sent()) -} - -fn continuation_account( - context: &AttemptContext, - previous_session: Option<&XaiSessionState>, -) -> Result, ProviderError> { - let Some(continuation) = context.continuation() else { - return Ok(None); - }; - match (context.continuation_attempt(), continuation) { - (ContinuationAttempt::Native, ContinuationBinding::Pinned(pin)) => { - if pin.provider().as_str() != XAI_PROVIDER_NAME { - return Err(invalid_continuation()); - } - Ok(Some(pin.account().clone())) - } - (ContinuationAttempt::Native, ContinuationBinding::External(_)) => { - Err(invalid_continuation()) - } - (ContinuationAttempt::ReplayOwner, _) => { - let previous = previous_session.ok_or_else(invalid_continuation)?; - gateway_core::account::ProviderAccountId::new(previous.account_id.clone()) - .map(Some) - .map_err(|_| invalid_continuation()) - } - (ContinuationAttempt::ReplayAny, _) => Ok(None), - (ContinuationAttempt::None, _) => Err(invalid_continuation()), - } -} - -fn apply_continuation( - request: &mut GrokResponsesRequest, - previous_session: Option<&XaiSessionState>, - context: &AttemptContext, - account: &gateway_core::account::ProviderAccountId, - current_input: &[Value], -) -> Result<(), ProviderError> { - let Some(continuation) = context.continuation() else { - return Ok(()); - }; - match context.continuation_attempt() { - ContinuationAttempt::Native => { - let ContinuationBinding::Pinned(pin) = continuation else { - return Err(invalid_continuation()); - }; - let provider = ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?; - if !pin.matches(&provider, account) { - return Err(invalid_continuation()); - } - request.set_previous_response_id(Some(pin.upstream_response_id().as_str().to_owned())); - Ok(()) - } - ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => { - let previous = previous_session.ok_or_else(invalid_continuation)?; - if context.continuation_attempt() == ContinuationAttempt::ReplayOwner - && previous.account_id != account.as_str() - { - return Err(invalid_continuation()); - } - let mut input = replay_input_for_account(previous, account.as_str(), true); - input.reserve(current_input.len()); - input.extend(current_input.iter().cloned()); - request.set_replay_input(input).map_err(map_request_error)?; - request.set_previous_response_id(None); - request.inherit_session(None); - Ok(()) - } - ContinuationAttempt::None => Err(invalid_continuation()), - } -} - -fn replay_input_for_account( - state: &XaiSessionState, - account_id: &str, - force_portable: bool, -) -> Vec { - state - .transcript - .iter() - .filter_map(|item| match item { - XaiReplayItem::ClientInput(value) | XaiReplayItem::SanitizedOutput(value) => { - Some(value.clone()) - } - XaiReplayItem::AccountOutput { - account_id: owner, - item, - } if owner == account_id && !force_portable => { - portable_output_item(item.clone(), false) - } - XaiReplayItem::AccountOutput { item, .. } => portable_output_item(item.clone(), true), - }) - .collect() -} - -fn project_transcript_to_account(transcript: &mut Vec, account_id: &str) { - *transcript = transcript - .drain(..) - .filter_map(|item| match item { - XaiReplayItem::AccountOutput { - account_id: owner, - item, - } if owner != account_id => { - portable_output_item(item, true).map(XaiReplayItem::SanitizedOutput) - } - item => Some(item), - }) - .collect(); -} - -fn portable_output_item(mut item: Value, strip_opaque: bool) -> Option { - let Value::Object(object) = &mut item else { - return None; - }; - let is_reasoning = object.get("type").and_then(Value::as_str) == Some("reasoning"); - if !matches!( - object.get("type").and_then(Value::as_str), - Some("reasoning" | "message" | "function_call" | "custom_tool_call") - ) { - return None; - } - object.remove("id"); - object.remove("status"); - if is_reasoning { - if strip_opaque - || object - .get("encrypted_content") - .and_then(Value::as_str) - .is_some_and(|value| !valid_reasoning_ciphertext(value)) - { - object.remove("encrypted_content"); - } - if object.get("encrypted_content").is_none() && !has_readable_reasoning(object) { - return None; - } - } - Some(item) -} - -fn has_readable_reasoning(item: &Map) -> bool { - ["summary", "content"].into_iter().any(|field| { - item.get(field) - .and_then(Value::as_array) - .is_some_and(|parts| { - parts.iter().any(|part| { - part.get("text") - .and_then(Value::as_str) - .is_some_and(|text| !text.trim().is_empty()) - }) - }) - }) -} - -fn attach_xai_session_update( - events: &mut [ProviderEvent], - capture: &mut Option, -) -> Result<(), ProviderError> { - if capture.is_none() { - return Ok(()); - } - let mut terminal_index = None; - for (index, event) in events.iter().enumerate() { - if let Some(capture) = capture.as_mut() { - capture_output_item(event, capture); - } - if event - .canonical_facts() - .iter() - .any(|fact| matches!(fact, GatewayEvent::Completed(_))) - { - terminal_index = Some(index); - } - } - let Some(terminal_index) = terminal_index else { - return Ok(()); - }; - let Some(mut capture) = capture.take() else { - return Ok(()); - }; - let mut transcript = capture - .previous - .take() - .map(|state| state.transcript) - .unwrap_or_default(); - project_transcript_to_account(&mut transcript, &capture.account_id); - transcript.extend( - capture - .request_input - .into_iter() - .map(XaiReplayItem::ClientInput), - ); - transcript.extend(capture.output_items.into_values().map(|item| { - XaiReplayItem::AccountOutput { - account_id: capture.account_id.clone(), - item, - } - })); - let state = XaiSessionState { - account_id: capture.account_id, - session_id: capture.session_id, - transcript, - }; - if let Some(update) = encode_xai_session_state(state)? { - events[terminal_index].attach_session_update(update); - } - Ok(()) -} - -fn capture_output_item(event: &ProviderEvent, capture: &mut GrokSessionCapture) { - let Some(wire) = event.wire_event() else { - return; - }; - let event_type = wire - .event_type() - .or_else(|| wire.data().get("type").and_then(Value::as_str)); - if event_type == Some("response.output_item.done") - && capture.output_items.len() < XAI_SESSION_OUTPUT_LIMIT - && let Some(item) = wire - .data() - .get("item") - .cloned() - .and_then(|item| portable_output_item(item, false)) - { - let index = wire - .data() - .get("output_index") - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) - .unwrap_or_else(|| u32::try_from(capture.output_items.len()).unwrap_or(u32::MAX)); - capture.output_items.insert(index, item); - } - if matches!( - event_type, - Some("response.completed" | "response.incomplete") - ) && let Some(output) = wire - .data() - .get("response") - .and_then(Value::as_object) - .and_then(|response| response.get("output")) - .and_then(Value::as_array) - .filter(|output| !output.is_empty()) - { - for (index, item) in output.iter().take(XAI_SESSION_OUTPUT_LIMIT).enumerate() { - let Some(index) = u32::try_from(index).ok() else { - break; - }; - if capture.output_items.contains_key(&index) { - continue; - } - if let Some(item) = portable_output_item(item.clone(), false) { - capture.output_items.insert(index, item); - } - } - } -} - -fn invalid_continuation() -> ProviderError { - ProviderError::new( - ProviderErrorKind::InvalidRequest, - UpstreamSendState::NotSent, - ) - .with_continuation_failure(ContinuationFailure::HistoryUnavailable) -} - -fn protocol_not_sent() -> ProviderError { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) -} - -fn protocol_sent() -> ProviderError { - provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent) -} - -struct GrokStreamAttempt { - client_identity: GrokClientIdentity, - wire_profile: XaiWireProfileState, - credential_recovery: Arc, - responses_url: Url, - request: GrokResponsesRequest, - upstream_model: UpstreamModelId, - context: AttemptContext, - session: Arc, - output_started_at: Instant, - session_capture: Option, - reasoning_replay_capture: Option, -} - -struct GrokCompactionStreamAttempt { - client_identity: GrokClientIdentity, - wire_profile: XaiWireProfileState, - credential_recovery: Arc, - responses_url: Url, - request: GrokCompactionRequest, - upstream_model: UpstreamModelId, - upstream_session_id: Option, - context: AttemptContext, - session: Arc, - reasoning_replay: GrokReasoningReplay, - reasoning_replay_key: Option, -} - -struct AcceptedGrokInference { - response: GrokInferenceResponse, - observation: ProviderResponseObservation, -} - -async fn next_grok_chunk( - body: &mut GrokInferenceChunkStream, - selector: &dyn GrokSessionSelector, - session: &SelectedGrokSession, - upstream_model: &UpstreamModelId, - context: &AttemptContext, -) -> Result, ProviderError> { - let Some(stream_deadline) = remaining(context.deadline()) else { - return Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - )); - }; - let cancellation = context.cancellation().clone(); - tokio::select! { - biased; - _ = cancellation.cancelled() => Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Sent, - )), - _ = tokio::time::sleep(stream_deadline) => Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - )), - chunk = body.next() => match chunk { - Some(Ok(chunk)) => Ok(Some(chunk)), - Some(Err(error)) => Err(map_and_record_stream_transport_failure( - selector, - session, - error, - upstream_model, - ).await), - None => Ok(None), - }, - } -} - -fn cold_compaction_http_sse_stream( - selector: Arc, - transport: Arc, - attempt: GrokCompactionStreamAttempt, -) -> EventStream { - let GrokCompactionStreamAttempt { - client_identity, - wire_profile, - credential_recovery, - responses_url, - mut request, - upstream_model, - upstream_session_id, - context, - session, - reasoning_replay, - reasoning_replay_key, - } = attempt; - Box::pin(async_stream::try_stream! { - let headers = build_grok_headers( - &wire_profile, - &session, - &client_identity, - context.request_id(), - upstream_session_id.as_deref(), - None, - &upstream_model, - ); - let mut invalid_encrypted_content_retried = false; - let accepted = loop { - if context.cancellation().is_cancelled() { - Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::NotSent, - ))?; - return; - } - let body = request.to_json_bytes().map_err(map_request_error)?; - let inference_request = GrokInferenceRequest::new( - responses_url.clone(), - headers.clone(), - body, - session.binding().clone(), - ); - let Some(handshake_deadline) = remaining(context.deadline()) else { - Err(mark_transient_compaction_failure(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::NotSent, - )))?; - return; - }; - let cancellation = context.cancellation().clone(); - let boundary = tokio::select! { - biased; - _ = cancellation.cancelled() => InferenceBoundary::Cancelled, - _ = tokio::time::sleep(handshake_deadline) => InferenceBoundary::Deadline, - response = transport.execute(inference_request) => InferenceBoundary::Response(response), - }; - match boundary { - InferenceBoundary::Cancelled => { - Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Ambiguous, - ))?; - return; - } - InferenceBoundary::Deadline => { - Err(mark_transient_compaction_failure(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Ambiguous, - )))?; - return; - } - InferenceBoundary::Response(Ok(response)) => { - let observation = xai_response_observation(&response) - .map_err(mark_transient_compaction_failure)?; - break AcceptedGrokInference { - response, - observation, - }; - } - InferenceBoundary::Response(Err(error)) => { - if !invalid_encrypted_content_retried - && is_invalid_encrypted_content_failure(&error) - && request.strip_invalid_encrypted_reasoning() - { - invalid_encrypted_content_retried = true; - continue; - } - let observation = xai_error_observation(&error).ok(); - let credential_failure = - transport_credential_failure(&error, &upstream_model); - let error = map_continuation_failure( - &context, - map_transport_error_for_context(error, &context), - ); - let error = recover_or_record_failure( - selector.as_ref(), - credential_recovery.as_ref(), - &session, - error, - credential_failure, - context.credential_recovery_attempted(), - ) - .await; - if let Some(observation) = observation { - yield ProviderEvent::observation(observation); - } - Err(mark_transient_compaction_failure(error))?; - return; - } - } - }; - yield ProviderEvent::observation(accepted.observation); - - let mut body = accepted.response.into_body(); - let mut canonical = GrokCanonicalDecoder::new(upstream_model.as_str()); - let mut summary = GrokCompactionSummaryDecoder::new(); - let mut facts = CompactionFacts::default(); - - 'stream: while let Some(chunk) = next_grok_chunk( - &mut body, - selector.as_ref(), - &session, - &upstream_model, - &context, - ) - .await - .map_err(mark_transient_compaction_failure)? - { - let events = match canonical.push(&chunk) { - Ok(events) => events, - Err(error) => { - let error = map_continuation_failure(&context, error); - let error = record_stream_failure( - selector.as_ref(), - &session, - error, - &upstream_model, - ) - .await; - Err(mark_transient_compaction_failure(error))?; - return; - } - }; - for event in events { - summary.observe(&event).map_err(map_compaction_decode_error)?; - facts.observe(&event); - if facts.completed.is_some() { - break 'stream; - } - } - } - - if facts.completed.is_none() { - let events = match canonical.finish_without_terminal() { - Ok(events) => events, - Err(error) => { - let error = map_continuation_failure(&context, error); - let error = record_stream_failure( - selector.as_ref(), - &session, - error, - &upstream_model, - ) - .await; - Err(mark_transient_compaction_failure(error))?; - return; - } - }; - for event in events { - summary.observe(&event).map_err(map_compaction_decode_error)?; - facts.observe(&event); - if facts.completed.is_some() { - break; - } - } - } - - let started = facts - .started - .ok_or_else(|| mark_transient_compaction_failure(protocol_sent()))?; - let completed = facts - .completed - .ok_or_else(|| mark_transient_compaction_failure(protocol_sent()))?; - let (summary, encrypted_content) = summary - .finish_with_encrypted_content() - .map_err(map_compaction_decode_error)?; - let (created, output_done, terminal) = crate::transport::compaction::compaction_wire_events( - &started, - &completed, - summary.as_deref(), - &encrypted_content, - facts.created_response.as_ref(), - facts.terminal_response.as_ref(), - ) - .map_err(|_| mark_transient_compaction_failure(protocol_sent()))? - .into_parts(); - ensure_sent_context(&context)?; - if session.allows_account_state_mutation() { - selector.record_success(&session).await; - } - if let Some(key) = reasoning_replay_key.as_ref() { - reasoning_replay.clear(key); - } - yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created); - yield ProviderEvent::wire(output_done); - let mut terminal_facts = facts.metering; - terminal_facts.push(GatewayEvent::Completed(completed)); - yield ProviderEvent::canonical_with_wire(terminal_facts, terminal); - }) -} - -#[derive(Default)] -struct CompactionFacts { - started: Option, - completed: Option, - metering: Vec, - created_response: Option, - terminal_response: Option, -} - -impl CompactionFacts { - fn observe(&mut self, event: &ProviderEvent) { - self.capture_wire_response(event); - for fact in event.canonical_facts() { - match fact { - GatewayEvent::Started(meta) if self.started.is_none() => { - self.started = Some(meta.clone()); - } - GatewayEvent::Completed(meta) if self.completed.is_none() => { - self.completed = Some(meta.clone()); - } - GatewayEvent::Usage(_) - | GatewayEvent::CalculatedCost(_) - | GatewayEvent::ProviderCost(_) => self.metering.push(fact.clone()), - _ => {} - } - } - } - - fn capture_wire_response(&mut self, event: &ProviderEvent) { - let Some(wire) = event.wire_event().filter(|wire| wire.has_json_data()) else { - return; - }; - let event_type = wire - .event_type() - .or_else(|| wire.data().get("type").and_then(Value::as_str)); - let response = wire.data().get("response").cloned(); - match event_type { - Some("response.created" | "response.in_progress") - if self.created_response.is_none() => - { - self.created_response = response; - } - Some("response.completed" | "response.incomplete") - if self.terminal_response.is_none() => - { - self.terminal_response = response; - } - _ => {} - } - } -} - -fn map_compaction_decode_error(error: GrokCompactionDecodeError) -> ProviderError { - match error { - GrokCompactionDecodeError::MissingEncryptedContent => { - mark_transient_compaction_failure(protocol_sent()) - } - } -} - -fn mark_transient_compaction_failure(error: ProviderError) -> ProviderError { - if matches!( - error.kind(), - ProviderErrorKind::RateLimited - | ProviderErrorKind::Timeout - | ProviderErrorKind::Transport - | ProviderErrorKind::Protocol - | ProviderErrorKind::Unavailable - ) { - error.with_pre_delivery_retry() - } else { - error - } -} - -fn cold_http_sse_stream( - selector: Arc, - transport: Arc, - attempt: GrokStreamAttempt, -) -> EventStream { - let GrokStreamAttempt { - client_identity, - wire_profile, - credential_recovery, - responses_url, - mut request, - upstream_model, - context, - session, - output_started_at, - mut session_capture, - mut reasoning_replay_capture, - } = attempt; - Box::pin(async_stream::try_stream! { - if context.cancellation().is_cancelled() { - Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::NotSent, - ))?; - } - let headers = build_grok_headers( - &wire_profile, - &session, - &client_identity, - context.request_id(), - request.session_id(), - None, - &upstream_model, - ); - let cancellation = context.cancellation().clone(); - let mut invalid_encrypted_content_retried = false; - let response = loop { - let body = request.to_json_bytes().map_err(map_request_error)?; - let inference_request = GrokInferenceRequest::new( - responses_url.clone(), - headers.clone(), - body, - session.binding().clone(), - ); - let Some(handshake_deadline) = remaining(context.deadline()) else { - Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::NotSent, - ))?; - return; - }; - let boundary = tokio::select! { - biased; - _ = cancellation.cancelled() => InferenceBoundary::Cancelled, - _ = tokio::time::sleep(handshake_deadline) => InferenceBoundary::Deadline, - response = transport.execute(inference_request) => InferenceBoundary::Response(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)) => break response, - InferenceBoundary::Response(Err(error)) => { - if !invalid_encrypted_content_retried - && is_invalid_encrypted_content_failure(&error) - && request.strip_invalid_encrypted_reasoning() - { - invalid_encrypted_content_retried = true; - if let Some(capture) = session_capture.as_mut() { - capture.request_input = request.input_items(); - } - continue; - } - let observation = xai_error_observation(&error)?; - let credential_failure = transport_credential_failure(&error, &upstream_model); - let error = map_continuation_failure( - &context, - map_transport_error_for_context(error, &context), - ); - let error = recover_or_record_failure( - selector.as_ref(), - credential_recovery.as_ref(), - &session, - error, - credential_failure, - context.credential_recovery_attempted(), - ) - .await; - yield ProviderEvent::observation(observation); - Err(error)?; - return; - } - } - }; - - let observation = xai_response_observation(&response)?; - let base_timings = observation.timings(); - let mut first_token_ms: Option = None; - yield ProviderEvent::observation(observation.clone()); - - let mut body = response.into_body(); - let mut decoder = GrokCanonicalDecoder::for_request(upstream_model.as_str(), &request); - loop { - let Some(stream_deadline) = remaining(context.deadline()) else { - Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - ))?; - return; - }; - let next = tokio::select! { - biased; - _ = cancellation.cancelled() => Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Sent, - )), - _ = tokio::time::sleep(stream_deadline) => Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - )), - chunk = body.next() => match chunk { - Some(Ok(chunk)) => Ok(Some(chunk)), - Some(Err(error)) => Err(map_and_record_stream_transport_failure( - selector.as_ref(), - &session, - error, - &upstream_model, - ).await), - None => Ok(None), - }, - }?; - let Some(chunk) = next else { - break; - }; - let mut events = match decoder.push(&chunk) { - Ok(events) => events, - Err(error) => { - let error = map_continuation_failure(&context, error); - let error = record_stream_failure( - selector.as_ref(), - &session, - error, - &upstream_model, - ) - .await; - Err(error)?; - return; - } - }; - // 首个非前导输出事件（结构帧也算）即上报携带 first_token_ms 的观测， - // 供 Core 覆盖会话级兜底值。 - if decoder.take_output_start() && first_token_ms.is_none() { - first_token_ms = Some( - u64::try_from(output_started_at.elapsed().as_millis()).unwrap_or(u64::MAX), - ); - yield ProviderEvent::observation( - observation - .clone() - .with_timings(ProviderResponseTimings { - first_token_ms, - ..base_timings - }), - ); - } - let completed = events - .iter() - .flat_map(ProviderEvent::canonical_facts) - .any(|event| matches!(event, GatewayEvent::Completed(_))); - attach_xai_session_update(&mut events, &mut session_capture)?; - if let Some(capture) = reasoning_replay_capture.as_mut() { - capture.observe(&events); - } - if completed && session.allows_account_state_mutation() { - selector.record_success(&session).await; - } - for event in events { - ensure_sent_context(&context)?; - yield event; - } - if completed { - return; - } - } - let mut final_events = match decoder.finish() { - Ok(events) => events, - Err(error) => { - let error = map_continuation_failure(&context, error); - let error = record_stream_failure( - selector.as_ref(), - &session, - error, - &upstream_model, - ) - .await; - Err(error)?; - return; - } - }; - let completed = final_events - .iter() - .flat_map(ProviderEvent::canonical_facts) - .any(|event| matches!(event, GatewayEvent::Completed(_))); - // 尾部 finish 补全缓冲中的首个输出帧时，同样上报首字。 - if decoder.take_output_start() && first_token_ms.is_none() { - first_token_ms = Some( - u64::try_from(output_started_at.elapsed().as_millis()).unwrap_or(u64::MAX), - ); - yield ProviderEvent::observation( - observation - .clone() - .with_timings(ProviderResponseTimings { - first_token_ms, - ..base_timings - }), - ); - } - attach_xai_session_update(&mut final_events, &mut session_capture)?; - if let Some(capture) = reasoning_replay_capture.as_mut() { - capture.observe(&final_events); - } - if completed && session.allows_account_state_mutation() { - selector.record_success(&session).await; - } - for event in final_events { - ensure_sent_context(&context)?; - yield event; - } - }) -} - -fn is_invalid_encrypted_content_failure(error: &GrokInferenceTransportError) -> bool { - error.kind() == GrokInferenceTransportErrorKind::InvalidRequest - && error.status() == Some(400) - && error.upstream_code().is_some_and(|code| { - matches!( - code.as_str(), - REASONING_DECODE_FAILED_CODE | "invalid_encrypted_content" - ) - }) -} - -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(with_xai_transport_metrics( - observation, - error.transport_metrics(), - )) -} - -fn xai_response_observation( - response: &GrokInferenceResponse, -) -> Result { - 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); - } - Ok(with_xai_transport_metrics( - observation, - response.transport_metrics(), - )) -} - -fn with_xai_transport_metrics( - mut observation: ProviderResponseObservation, - metrics: GrokInferenceTransportMetrics, -) -> ProviderResponseObservation { - observation = observation.with_timings(ProviderResponseTimings { - headers_ms: metrics.headers_ms(), - ..ProviderResponseTimings::default() - }); - if let Some(metadata) = xai_transport_metadata(metrics) { - observation = observation.with_provider_metadata(metadata); - } - observation -} - -fn xai_transport_metadata( - metrics: GrokInferenceTransportMetrics, -) -> Option { - let mut metadata = Map::new(); - if let Some(status) = metrics.client_cache_status() { - metadata.insert( - "clientCache".to_owned(), - Value::String(status.as_str().to_owned()), - ); - } - if let Some(dns) = metrics.dns() { - metadata.insert( - "dnsSource".to_owned(), - Value::String(dns.source().as_str().to_owned()), - ); - metadata.insert("dnsMs".to_owned(), Value::from(dns.duration_ms())); - } - if let Some(headers_ms) = metrics.headers_ms() { - metadata.insert("upstreamHeadersMs".to_owned(), Value::from(headers_ms)); - } - if metadata.is_empty() { - return None; - } - // 观测对象版本信封：读取端按 schemaVersion 解码，缺失视为 v0。 - metadata.insert("schemaVersion".to_owned(), Value::from(1)); - ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?) -} - -async fn record_credential_failure( - selector: &dyn GrokSessionSelector, - session: &SelectedGrokSession, - error: ProviderError, - failure: GrokCredentialFailure, -) -> ProviderError { - if session.allows_account_state_mutation() { - selector.record_failure(session, failure).await; - } - error -} - -async fn record_stream_failure( - selector: &dyn GrokSessionSelector, - session: &SelectedGrokSession, - error: ProviderError, - upstream_model: &UpstreamModelId, -) -> ProviderError { - if !session.allows_account_state_mutation() { - return error; - } - if let Some(failure) = stream_credential_failure(&error, upstream_model) { - return record_credential_failure(selector, session, error, failure).await; - } - if error.kind() != ProviderErrorKind::Cancelled { - selector - .record_failure(session, GrokCredentialFailure::StreamInterrupted) - .await; - } - error -} - -fn stream_credential_failure( - error: &ProviderError, - upstream_model: &UpstreamModelId, -) -> Option { - match error.kind() { - ProviderErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized), - ProviderErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited { - retry_after: error.retry_after(), - }), - ProviderErrorKind::QuotaExhausted => { - let visible = error.client_visible_upstream_error(); - let quota_kind = classify_grok_quota_failure( - visible - .and_then(ClientVisibleUpstreamError::code) - .or_else(|| error.upstream_code().map(|code| code.as_str())), - visible.and_then(ClientVisibleUpstreamError::error_type), - visible.map(ClientVisibleUpstreamError::message), - ); - Some(match quota_kind { - Some(GrokQuotaFailureKind::FreeAccount) => { - GrokCredentialFailure::FreeQuotaExhausted - } - Some(GrokQuotaFailureKind::FreeModelUsage) => { - GrokCredentialFailure::ModelQuotaExhausted { - upstream_model: upstream_model.clone(), - retry_after: error.retry_after(), - } - } - Some(GrokQuotaFailureKind::Account) | None => GrokCredentialFailure::QuotaExhausted, - }) - } - _ => None, - } -} - -async fn recover_or_record_failure( - selector: &dyn GrokSessionSelector, - recovery: &dyn GrokCredentialRecovery, - session: &SelectedGrokSession, - error: ProviderError, - credential_failure: Option, - recovery_attempted: bool, -) -> ProviderError { - if !session.allows_account_state_mutation() { - return error; - } - if error.requires_credential_recovery() && !recovery_attempted { - return match recovery - .recover_unauthorized(session.account_id(), session.credential_revision()) - .await - { - GrokCredentialRecoveryOutcome::Recovered => error.with_same_account_retry(), - GrokCredentialRecoveryOutcome::Rejected => error, - GrokCredentialRecoveryOutcome::Unavailable => match credential_failure { - Some(failure) => record_credential_failure(selector, session, error, failure).await, - None => error, - }, - }; - } - match credential_failure { - Some(failure) => record_credential_failure(selector, session, error, failure).await, - None => error, - } -} - -fn transport_credential_failure( - error: &GrokInferenceTransportError, - upstream_model: &UpstreamModelId, -) -> Option { - match error.kind() { - GrokInferenceTransportErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized), - GrokInferenceTransportErrorKind::PermissionDenied => { - Some(GrokCredentialFailure::AccessDenied) - } - GrokInferenceTransportErrorKind::RateLimited if model_capacity_failure(error) => { - Some(GrokCredentialFailure::ModelCapacity { - upstream_model: upstream_model.clone(), - }) - } - GrokInferenceTransportErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited { - retry_after: error.retry_after(), - }), - GrokInferenceTransportErrorKind::QuotaExhausted => { - Some(GrokCredentialFailure::QuotaExhausted) - } - GrokInferenceTransportErrorKind::FreeQuotaExhausted => { - Some(GrokCredentialFailure::FreeQuotaExhausted) - } - GrokInferenceTransportErrorKind::PaymentRequired => { - Some(GrokCredentialFailure::PaymentRequired { - retry_after: error.retry_after(), - }) - } - GrokInferenceTransportErrorKind::ModelQuotaExhausted => { - Some(GrokCredentialFailure::ModelQuotaExhausted { - upstream_model: upstream_model.clone(), - retry_after: error.retry_after(), - }) - } - GrokInferenceTransportErrorKind::ModelAccessDenied => { - Some(GrokCredentialFailure::ModelAccessDenied { - upstream_model: upstream_model.clone(), - retry_after: error.retry_after(), - }) - } - GrokInferenceTransportErrorKind::Unavailable if empty_upstream_failure(error) => { - Some(GrokCredentialFailure::EmptyUpstream) - } - GrokInferenceTransportErrorKind::Unavailable => { - Some(GrokCredentialFailure::UpstreamUnavailable) - } - _ => None, - } -} - -fn model_capacity_failure(error: &GrokInferenceTransportError) -> bool { - transport_error_contains_any( - error, - &[ - "model_capacity", - "capacity", - "overloaded", - "server_busy", - "too many concurrent", - "engine_overloaded", - ], - ) -} - -fn empty_upstream_failure(error: &GrokInferenceTransportError) -> bool { - transport_error_contains_any( - error, - &[ - "empty_upstream", - "empty model output", - "no content/tool_calls", - "no client-visible content", - "empty upstream", - ], - ) -} - -fn transport_error_contains_any(error: &GrokInferenceTransportError, signals: &[&str]) -> bool { - let code = error - .upstream_code() - .map(|code| code.as_str()) - .unwrap_or_default(); - let message = error - .client_visible_upstream_error() - .map(ClientVisibleUpstreamError::message) - .unwrap_or_default() - .to_ascii_lowercase(); - signals - .iter() - .any(|signal| code.contains(signal) || message.contains(signal)) -} - -async fn map_and_record_stream_transport_failure( - selector: &dyn GrokSessionSelector, - session: &SelectedGrokSession, - error: GrokInferenceTransportError, - upstream_model: &UpstreamModelId, -) -> ProviderError { - let request_scoped = error.kind() == GrokInferenceTransportErrorKind::SafetyRejected; - let credential_failure = transport_credential_failure(&error, upstream_model); - let error = map_stream_error(error); - match credential_failure { - Some(failure) => record_credential_failure(selector, session, error, failure).await, - None if request_scoped => error, - None => record_stream_failure(selector, session, error, upstream_model).await, - } -} - -fn map_continuation_failure(context: &AttemptContext, error: ProviderError) -> ProviderError { - let is_reasoning_decode_failure = context.continuation_attempt() == ContinuationAttempt::Native - && error.kind() == ProviderErrorKind::InvalidRequest - && error.upstream_status() == Some(400) - && error - .upstream_code() - .is_some_and(|code| code.as_str() == REASONING_DECODE_FAILED_CODE); - let is_missing_native_response = context.continuation_attempt() == ContinuationAttempt::Native - && error.kind() == ProviderErrorKind::InvalidRequest - && error.upstream_status() == Some(404) - && error - .upstream_code() - .is_some_and(|code| code.as_str() == RESPONSE_NOT_FOUND_CODE); - if is_reasoning_decode_failure || is_missing_native_response { - error - .with_continuation_failure(ContinuationFailure::HistoryUnavailable) - .with_replay_safe() - } else { - error - } -} - -fn preflight_context(context: &AttemptContext) -> Result<(), ProviderError> { - if context.cancellation().is_cancelled() { - return Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::NotSent, - )); - } - if remaining(context.deadline()).is_none() { - return Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::NotSent, - )); - } - Ok(()) -} - -fn ensure_sent_context(context: &AttemptContext) -> Result<(), ProviderError> { - if context.cancellation().is_cancelled() { - return Err(provider_error( - ProviderErrorKind::Cancelled, - UpstreamSendState::Sent, - )); - } - if remaining(context.deadline()).is_none() { - return Err(provider_error( - ProviderErrorKind::Timeout, - UpstreamSendState::Sent, - )); - } - Ok(()) -} - -fn map_request_error(error: GrokRequestEncodeError) -> ProviderError { - let kind = match error { - GrokRequestEncodeError::InvalidProtocolPayload - | GrokRequestEncodeError::InvalidRequestNormalization - | GrokRequestEncodeError::InvalidRequestField { .. } => ProviderErrorKind::InvalidRequest, - GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol, - }; - let provider_error = provider_error(kind, UpstreamSendState::NotSent); - if kind != ProviderErrorKind::InvalidRequest { - return provider_error; - } - provider_error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( - error.to_string(), - Some("invalid_request_normalization".to_owned()), - Some("invalid_request_error".to_owned()), - )) -} - -/// 将选择阶段失败映射为带结构化 code 与 retry_after 的 Provider 错误。 -fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError { - let (retry_after, message, code) = match error { - GrokSessionSelectorError::AccountCoolingDown { retry_after } => ( - retry_after, - cooling_down_message(retry_after), - "account_cooling_down", - ), - GrokSessionSelectorError::ModelCoolingDown { retry_after } => ( - retry_after, - model_cooling_down_message(retry_after), - "model_cooling_down", - ), - GrokSessionSelectorError::CapacityUnavailable { retry_after } => ( - retry_after, - "account is at its concurrency or request-interval limit".to_owned(), - "account_capacity_busy", - ), - GrokSessionSelectorError::NoEligibleSession => ( - None, - "no account is eligible for the requested model".to_owned(), - "no_eligible_account", - ), - GrokSessionSelectorError::Unavailable => ( - None, - "account scheduling state is temporarily unreadable".to_owned(), - "account_selector_unavailable", - ), - GrokSessionSelectorError::InvalidSession => { - return provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent); - } - }; - let error = provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent); - let error = match retry_after { - Some(retry_after) => error.with_retry_after(retry_after), - None => error, - }; - error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( - message, - Some(code.to_owned()), - Some("account_unavailable_error".to_owned()), - )) -} - -fn cooling_down_message(retry_after: Option) -> String { - match retry_after { - Some(retry_after) => format!( - "account is cooling down after an upstream failure; retry in {}s", - retry_after.as_secs().saturating_add(1) - ), - None => "account is cooling down after an upstream failure".to_owned(), - } -} - -fn model_cooling_down_message(retry_after: Option) -> String { - match retry_after { - Some(retry_after) => format!( - "all eligible accounts are cooling down for this model; retry in {}s", - retry_after.as_secs().saturating_add(1) - ), - None => "all eligible accounts are cooling down for this model".to_owned(), - } -} - -fn map_transport_error_for_context( - error: GrokInferenceTransportError, - context: &AttemptContext, -) -> ProviderError { - let allow_explicit_replay = context.continuation().is_none() - || error.kind() == GrokInferenceTransportErrorKind::Unauthorized; - map_transport_error_with_state(error, None, allow_explicit_replay) -} - -fn map_stream_error(error: GrokInferenceTransportError) -> ProviderError { - map_transport_error_with_state(error, Some(UpstreamSendState::Sent), false) -} - -fn map_transport_error_with_state( - error: GrokInferenceTransportError, - forced_send_state: Option, - allow_explicit_replay: bool, -) -> ProviderError { - let transport_kind = error.kind(); - let kind = match transport_kind { - GrokInferenceTransportErrorKind::InvalidRequest => ProviderErrorKind::InvalidRequest, - GrokInferenceTransportErrorKind::Unsupported => ProviderErrorKind::Unsupported, - GrokInferenceTransportErrorKind::Unauthorized => ProviderErrorKind::Unauthorized, - GrokInferenceTransportErrorKind::PermissionDenied - | GrokInferenceTransportErrorKind::ModelAccessDenied - | GrokInferenceTransportErrorKind::PaymentRequired - | GrokInferenceTransportErrorKind::SafetyRejected => ProviderErrorKind::PermissionDenied, - GrokInferenceTransportErrorKind::RateLimited => ProviderErrorKind::RateLimited, - GrokInferenceTransportErrorKind::QuotaExhausted - | GrokInferenceTransportErrorKind::FreeQuotaExhausted - | GrokInferenceTransportErrorKind::ModelQuotaExhausted => ProviderErrorKind::QuotaExhausted, - GrokInferenceTransportErrorKind::Timeout => ProviderErrorKind::Timeout, - GrokInferenceTransportErrorKind::Transport => ProviderErrorKind::Transport, - GrokInferenceTransportErrorKind::Protocol => ProviderErrorKind::Protocol, - GrokInferenceTransportErrorKind::Unavailable => ProviderErrorKind::Unavailable, - GrokInferenceTransportErrorKind::Cancelled => ProviderErrorKind::Cancelled, - }; - let mut mapped = provider_error( - kind, - forced_send_state.unwrap_or_else(|| error.send_state()), - ); - if let Some(status) = error.status() { - mapped = mapped.with_status(status); - if allow_explicit_replay - && forced_send_state.is_none() - && explicit_rejection_is_replay_safe(&error, status) - { - mapped = mapped.with_replay_safe(); - } - } - 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 let Some(code) = error.upstream_code().cloned() { - mapped = mapped.with_upstream_code(code); - } - if let Some(detail) = error.client_visible_upstream_error().cloned() { - mapped = mapped.with_client_visible_upstream_error(detail); - } - if error.requires_credential_recovery() { - mapped = mapped.with_credential_recovery().with_replay_safe(); - } - if error.sensitive_context_was_redacted() { - mapped = mapped.redact_sensitive_context("upstream transport context"); - } - mapped -} - -fn explicit_rejection_is_replay_safe(error: &GrokInferenceTransportError, status: u16) -> bool { - let kind = error.kind(); - if kind == GrokInferenceTransportErrorKind::SafetyRejected { - return false; - } - if matches!(status, 401 | 402 | 403 | 405 | 429 | 529) || (500..=599).contains(&status) { - return true; - } - status == 400 - && (matches!( - kind, - GrokInferenceTransportErrorKind::QuotaExhausted - | GrokInferenceTransportErrorKind::FreeQuotaExhausted - | GrokInferenceTransportErrorKind::ModelQuotaExhausted - | GrokInferenceTransportErrorKind::PaymentRequired - | GrokInferenceTransportErrorKind::Unavailable - ) || (kind == GrokInferenceTransportErrorKind::RateLimited - && model_capacity_failure(error))) -} - -fn provider_error(kind: ProviderErrorKind, send_state: UpstreamSendState) -> ProviderError { - ProviderError::new(kind, send_state) -} - -fn remaining(deadline: SystemTime) -> Option { - deadline - .duration_since(SystemTime::now()) - .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); -// rolling 24h 描述的是上游用量窗口，不代表从本次观测起封禁 24 小时； -// 缺少可信 reset 时间时按短周期探测策略恢复检查。 -const EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL: Duration = Duration::from_secs(10 * 60); -const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = QUOTA_CATALOG_INTERVAL; -const CLI_RELEASE_WORKER_OWNER: &str = "xai-cli-release"; - -pub(crate) fn worker_contributions( - refresh: Arc, - quota: Arc, - catalog: Arc, - accounts: Arc, - provider_kind: ProviderKind, - cli_release: Arc, -) -> Result, WorkerDefinitionError> { - let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, XAI_PROVIDER_NAME)?; - let catalog_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, XAI_PROVIDER_NAME)?; - let release_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, CLI_RELEASE_WORKER_OWNER)?; - 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, - quota, - catalog, - provider_kind, - last_periodic_refresh_at: Mutex::new(BTreeMap::new()), - }), - )?), - WorkerContribution::Registration(scheduled_registration( - release_id, - GROK_CLI_RELEASE_POLL_INTERVAL, - Box::new(XaiCliReleaseTask { - service: cli_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 XaiOAuthRefreshTask { - service: Arc, -} - -struct XaiCliReleaseTask { - service: Arc, -} - -impl ScheduledTask for XaiCliReleaseTask { - 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 { - tracing::warn!(error = %error, "xAI CLI release check failed"); - } - Ok(()) - }) - } -} - -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(|error| { - tracing::error!(error = %error, "xAI OAuth refresh cycle failed"); - WorkerTaskError::safe("xAI OAuth refresh failed") - })?; - let failures = outcomes - .iter() - .filter(|outcome| { - matches!( - outcome, - GrokCredentialRefreshOutcome::Ambiguous { .. } - | GrokCredentialRefreshOutcome::Transient { .. } - | GrokCredentialRefreshOutcome::Failed { .. } - ) - }) - .count(); - if !outcomes.is_empty() { - tracing::info!( - accounts = outcomes.len(), - failures, - "xAI OAuth refresh cycle completed" - ); - } - if failures > 0 { - tracing::warn!(failures, "xAI OAuth refresh cycle contained failures"); - } - Ok(()) - }) - } -} - -struct XaiQuotaCatalogTask { - accounts: Arc, - quota: Arc, - catalog: Arc, - provider_kind: ProviderKind, - last_periodic_refresh_at: Mutex>, -} - -impl ScheduledTask for XaiQuotaCatalogTask { - fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { - Box::pin(async move { - let accounts = self - .accounts - .list_for_provider(&self.provider_kind) - .await - .map_err(|_| WorkerTaskError::safe("xAI Provider accounts unavailable"))?; - let mut failures = 0_u64; - let now = SystemTime::now(); - let accounts = self.reserve_periodic_refreshes(accounts, now); - for account in accounts { - 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), - } - } - match self.catalog.query_models().await { - Ok(_) | Err(GrokCredentialCatalogError::NoEligibleCredential) => {} - Err(_) => failures = failures.saturating_add(1), - } - if failures == 0 { - Ok(()) - } else { - Err(WorkerTaskError::safe( - "xAI quota or catalog synchronization failed", - )) - } - }) - } -} - -impl XaiQuotaCatalogTask { - fn reserve_periodic_refreshes( - &self, - accounts: Vec, - now: SystemTime, - ) -> Vec { - let candidates = accounts - .into_iter() - .filter(|account| eligible_quota_worker_account(account, now)) - .collect::>(); - let candidate_ids = candidates - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let now = Instant::now(); - let mut last_periodic_refresh_at = self - .last_periodic_refresh_at - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - last_periodic_refresh_at.retain(|account_id, _| candidate_ids.contains(account_id)); - candidates - .into_iter() - .filter(|account| { - let due = last_periodic_refresh_at - .get(account.id()) - .is_none_or(|last| { - now.saturating_duration_since(*last) - >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL - }); - if due { - last_periodic_refresh_at.insert(account.id().clone(), now); - } - due - }) - .collect() - } -} - -fn eligible_quota_worker_account(account: &ProviderAccount, now: SystemTime) -> bool { - account.enabled() - && account - .access_token_expires_at() - .is_some_and(|expires_at| expires_at > now) - && account - .quota() - .exhaustion_refresh_due(now, EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL) -} diff --git a/backend/crates/providers/xai/src/provider/continuation.rs b/backend/crates/providers/xai/src/provider/continuation.rs new file mode 100644 index 000000000..2fc603f11 --- /dev/null +++ b/backend/crates/providers/xai/src/provider/continuation.rs @@ -0,0 +1,333 @@ +//! xAI continuation 与 reasoning replay 状态处理。 + +use super::*; + +#[derive(Clone, Serialize, Deserialize)] +pub(super) struct XaiSessionState { + pub(super) account_id: String, + pub(super) session_id: Option, + pub(super) transcript: Vec, +} + +#[derive(Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum XaiReplayItem { + ClientInput(Value), + SanitizedOutput(Value), + AccountOutput { account_id: String, item: Value }, +} + +pub(super) struct GrokSessionCapture { + pub(super) previous: Option, + pub(super) request_input: Vec, + pub(super) account_id: String, + pub(super) session_id: Option, + pub(super) output_items: BTreeMap, +} + +pub(super) fn decode_xai_session_state( + request: &GenerateRequest, +) -> Result, ProviderError> { + request + .provider_session_state(XAI_PROVIDER_NAME) + .map(|state| { + let payload = Value::Object(state.payload().clone()); + if serde_json::to_vec(&payload) + .map_err(|_| protocol_not_sent())? + .len() + > XAI_SESSION_STATE_MAX_BYTES + { + return Err(protocol_not_sent()); + } + serde_json::from_value(payload).map_err(|_| protocol_not_sent()) + }) + .transpose() +} + +pub(super) fn encode_xai_session_state( + state: XaiSessionState, +) -> Result, ProviderError> { + let value = serde_json::to_value(state).map_err(|_| protocol_sent())?; + if serde_json::to_vec(&value) + .map_err(|_| protocol_sent())? + .len() + > XAI_SESSION_STATE_MAX_BYTES + { + return Ok(None); + } + let Value::Object(payload) = value else { + return Err(protocol_sent()); + }; + ProviderSessionState::new(XAI_PROVIDER_NAME, payload) + .map(Some) + .map_err(|_| protocol_sent()) +} + +pub(super) fn continuation_account( + context: &AttemptContext, + previous_session: Option<&XaiSessionState>, +) -> Result, ProviderError> { + let Some(continuation) = context.continuation() else { + return Ok(None); + }; + match (context.continuation_attempt(), continuation) { + (ContinuationAttempt::Native, ContinuationBinding::Pinned(pin)) => { + if pin.provider().as_str() != XAI_PROVIDER_NAME { + return Err(invalid_continuation()); + } + Ok(Some(pin.account().clone())) + } + (ContinuationAttempt::Native, ContinuationBinding::External(_)) => { + Err(invalid_continuation()) + } + (ContinuationAttempt::ReplayOwner, _) => { + let previous = previous_session.ok_or_else(invalid_continuation)?; + gateway_core::account::ProviderAccountId::new(previous.account_id.clone()) + .map(Some) + .map_err(|_| invalid_continuation()) + } + (ContinuationAttempt::ReplayAny, _) => Ok(None), + (ContinuationAttempt::None, _) => Err(invalid_continuation()), + } +} + +pub(super) fn apply_continuation( + request: &mut GrokResponsesRequest, + previous_session: Option<&XaiSessionState>, + context: &AttemptContext, + account: &gateway_core::account::ProviderAccountId, + current_input: &[Value], +) -> Result<(), ProviderError> { + let Some(continuation) = context.continuation() else { + return Ok(()); + }; + match context.continuation_attempt() { + ContinuationAttempt::Native => { + let ContinuationBinding::Pinned(pin) = continuation else { + return Err(invalid_continuation()); + }; + let provider = ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?; + if !pin.matches(&provider, account) { + return Err(invalid_continuation()); + } + request.set_previous_response_id(Some(pin.upstream_response_id().as_str().to_owned())); + Ok(()) + } + ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => { + let previous = previous_session.ok_or_else(invalid_continuation)?; + if context.continuation_attempt() == ContinuationAttempt::ReplayOwner + && previous.account_id != account.as_str() + { + return Err(invalid_continuation()); + } + let mut input = replay_input_for_account(previous, account.as_str(), true); + input.reserve(current_input.len()); + input.extend(current_input.iter().cloned()); + request.set_replay_input(input).map_err(map_request_error)?; + request.set_previous_response_id(None); + request.inherit_session(None); + Ok(()) + } + ContinuationAttempt::None => Err(invalid_continuation()), + } +} + +pub(super) fn replay_input_for_account( + state: &XaiSessionState, + account_id: &str, + force_portable: bool, +) -> Vec { + state + .transcript + .iter() + .filter_map(|item| match item { + XaiReplayItem::ClientInput(value) | XaiReplayItem::SanitizedOutput(value) => { + Some(value.clone()) + } + XaiReplayItem::AccountOutput { + account_id: owner, + item, + } if owner == account_id && !force_portable => { + portable_output_item(item.clone(), false) + } + XaiReplayItem::AccountOutput { item, .. } => portable_output_item(item.clone(), true), + }) + .collect() +} + +pub(super) fn project_transcript_to_account(transcript: &mut Vec, account_id: &str) { + *transcript = transcript + .drain(..) + .filter_map(|item| match item { + XaiReplayItem::AccountOutput { + account_id: owner, + item, + } if owner != account_id => { + portable_output_item(item, true).map(XaiReplayItem::SanitizedOutput) + } + item => Some(item), + }) + .collect(); +} + +pub(super) fn portable_output_item(mut item: Value, strip_opaque: bool) -> Option { + let Value::Object(object) = &mut item else { + return None; + }; + let is_reasoning = object.get("type").and_then(Value::as_str) == Some("reasoning"); + if !matches!( + object.get("type").and_then(Value::as_str), + Some("reasoning" | "message" | "function_call" | "custom_tool_call") + ) { + return None; + } + object.remove("id"); + object.remove("status"); + if is_reasoning { + if strip_opaque + || object + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|value| !valid_reasoning_ciphertext(value)) + { + object.remove("encrypted_content"); + } + if object.get("encrypted_content").is_none() && !has_readable_reasoning(object) { + return None; + } + } + Some(item) +} + +pub(super) fn has_readable_reasoning(item: &Map) -> bool { + ["summary", "content"].into_iter().any(|field| { + item.get(field) + .and_then(Value::as_array) + .is_some_and(|parts| { + parts.iter().any(|part| { + part.get("text") + .and_then(Value::as_str) + .is_some_and(|text| !text.trim().is_empty()) + }) + }) + }) +} + +pub(super) fn attach_xai_session_update( + events: &mut [ProviderEvent], + capture: &mut Option, +) -> Result<(), ProviderError> { + if capture.is_none() { + return Ok(()); + } + let mut terminal_index = None; + for (index, event) in events.iter().enumerate() { + if let Some(capture) = capture.as_mut() { + capture_output_item(event, capture); + } + if event + .canonical_facts() + .iter() + .any(|fact| matches!(fact, GatewayEvent::Completed(_))) + { + terminal_index = Some(index); + } + } + let Some(terminal_index) = terminal_index else { + return Ok(()); + }; + let Some(mut capture) = capture.take() else { + return Ok(()); + }; + let mut transcript = capture + .previous + .take() + .map(|state| state.transcript) + .unwrap_or_default(); + project_transcript_to_account(&mut transcript, &capture.account_id); + transcript.extend( + capture + .request_input + .into_iter() + .map(XaiReplayItem::ClientInput), + ); + transcript.extend(capture.output_items.into_values().map(|item| { + XaiReplayItem::AccountOutput { + account_id: capture.account_id.clone(), + item, + } + })); + let state = XaiSessionState { + account_id: capture.account_id, + session_id: capture.session_id, + transcript, + }; + if let Some(update) = encode_xai_session_state(state)? { + events[terminal_index].attach_session_update(update); + } + Ok(()) +} + +pub(super) fn capture_output_item(event: &ProviderEvent, capture: &mut GrokSessionCapture) { + let Some(wire) = event.wire_event() else { + return; + }; + let event_type = wire + .event_type() + .or_else(|| wire.data().get("type").and_then(Value::as_str)); + if event_type == Some("response.output_item.done") + && capture.output_items.len() < XAI_SESSION_OUTPUT_LIMIT + && let Some(item) = wire + .data() + .get("item") + .cloned() + .and_then(|item| portable_output_item(item, false)) + { + let index = wire + .data() + .get("output_index") + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .unwrap_or_else(|| u32::try_from(capture.output_items.len()).unwrap_or(u32::MAX)); + capture.output_items.insert(index, item); + } + if matches!( + event_type, + Some("response.completed" | "response.incomplete") + ) && let Some(output) = wire + .data() + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("output")) + .and_then(Value::as_array) + .filter(|output| !output.is_empty()) + { + for (index, item) in output.iter().take(XAI_SESSION_OUTPUT_LIMIT).enumerate() { + let Some(index) = u32::try_from(index).ok() else { + break; + }; + if capture.output_items.contains_key(&index) { + continue; + } + if let Some(item) = portable_output_item(item.clone(), false) { + capture.output_items.insert(index, item); + } + } + } +} + +pub(super) fn invalid_continuation() -> ProviderError { + ProviderError::new( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + ) + .with_continuation_failure(ContinuationFailure::HistoryUnavailable) +} + +pub(super) fn protocol_not_sent() -> ProviderError { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) +} + +pub(super) fn protocol_sent() -> ProviderError { + provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent) +} diff --git a/backend/crates/providers/xai/src/provider/failure.rs b/backend/crates/providers/xai/src/provider/failure.rs new file mode 100644 index 000000000..e97345c6d --- /dev/null +++ b/backend/crates/providers/xai/src/provider/failure.rs @@ -0,0 +1,557 @@ +//! xAI 上游失败分类、账号反馈与恢复决策。 + +use super::*; + +pub(super) fn is_invalid_encrypted_content_failure(error: &GrokInferenceTransportError) -> bool { + error.kind() == GrokInferenceTransportErrorKind::InvalidRequest + && error.status() == Some(400) + && error.upstream_code().is_some_and(|code| { + matches!( + code.as_str(), + REASONING_DECODE_FAILED_CODE | "invalid_encrypted_content" + ) + }) +} + +pub(super) enum InferenceBoundary { + Response(Result), + Cancelled, + Deadline, +} + +pub(super) 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(with_xai_transport_metrics( + observation, + error.transport_metrics(), + )) +} + +pub(super) fn xai_response_observation( + response: &GrokInferenceResponse, +) -> Result { + 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); + } + Ok(with_xai_transport_metrics( + observation, + response.transport_metrics(), + )) +} + +pub(super) fn with_xai_transport_metrics( + mut observation: ProviderResponseObservation, + metrics: GrokInferenceTransportMetrics, +) -> ProviderResponseObservation { + observation = observation.with_timings(ProviderResponseTimings { + headers_ms: metrics.headers_ms(), + ..ProviderResponseTimings::default() + }); + if let Some(metadata) = xai_transport_metadata(metrics) { + observation = observation.with_provider_metadata(metadata); + } + observation +} + +pub(super) fn xai_transport_metadata( + metrics: GrokInferenceTransportMetrics, +) -> Option { + let mut metadata = Map::new(); + if let Some(status) = metrics.client_cache_status() { + metadata.insert( + "clientCache".to_owned(), + Value::String(status.as_str().to_owned()), + ); + } + if let Some(dns) = metrics.dns() { + metadata.insert( + "dnsSource".to_owned(), + Value::String(dns.source().as_str().to_owned()), + ); + metadata.insert("dnsMs".to_owned(), Value::from(dns.duration_ms())); + } + if let Some(headers_ms) = metrics.headers_ms() { + metadata.insert("upstreamHeadersMs".to_owned(), Value::from(headers_ms)); + } + if metadata.is_empty() { + return None; + } + // 观测对象版本信封：读取端按 schemaVersion 解码，缺失视为 v0。 + metadata.insert("schemaVersion".to_owned(), Value::from(1)); + ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?) +} + +pub(super) async fn record_credential_failure( + selector: &dyn GrokSessionSelector, + session: &SelectedGrokSession, + error: ProviderError, + failure: GrokCredentialFailure, +) -> ProviderError { + if session.allows_account_state_mutation() { + selector.record_failure(session, failure).await; + } + error +} + +pub(super) async fn record_stream_failure( + selector: &dyn GrokSessionSelector, + session: &SelectedGrokSession, + error: ProviderError, + upstream_model: &UpstreamModelId, +) -> ProviderError { + if !session.allows_account_state_mutation() { + return error; + } + if let Some(failure) = stream_credential_failure(&error, upstream_model) { + return record_credential_failure(selector, session, error, failure).await; + } + if error.kind() != ProviderErrorKind::Cancelled { + selector + .record_failure(session, GrokCredentialFailure::StreamInterrupted) + .await; + } + error +} + +pub(super) fn stream_credential_failure( + error: &ProviderError, + upstream_model: &UpstreamModelId, +) -> Option { + match error.kind() { + ProviderErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized), + ProviderErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited { + retry_after: error.retry_after(), + }), + ProviderErrorKind::QuotaExhausted => { + let visible = error.client_visible_upstream_error(); + let quota_kind = classify_grok_quota_failure( + visible + .and_then(ClientVisibleUpstreamError::code) + .or_else(|| error.upstream_code().map(|code| code.as_str())), + visible.and_then(ClientVisibleUpstreamError::error_type), + visible.map(ClientVisibleUpstreamError::message), + ); + Some(match quota_kind { + Some(GrokQuotaFailureKind::FreeAccount) => { + GrokCredentialFailure::FreeQuotaExhausted + } + Some(GrokQuotaFailureKind::FreeModelUsage) => { + GrokCredentialFailure::ModelQuotaExhausted { + upstream_model: upstream_model.clone(), + retry_after: error.retry_after(), + } + } + Some(GrokQuotaFailureKind::Account) | None => GrokCredentialFailure::QuotaExhausted, + }) + } + _ => None, + } +} + +pub(super) async fn recover_or_record_failure( + selector: &dyn GrokSessionSelector, + recovery: &dyn GrokCredentialRecovery, + session: &SelectedGrokSession, + error: ProviderError, + credential_failure: Option, + recovery_attempted: bool, +) -> ProviderError { + if !session.allows_account_state_mutation() { + return error; + } + if error.requires_credential_recovery() && !recovery_attempted { + return match recovery + .recover_unauthorized(session.account_id(), session.credential_revision()) + .await + { + GrokCredentialRecoveryOutcome::Recovered => error.with_same_account_retry(), + GrokCredentialRecoveryOutcome::Rejected => error, + GrokCredentialRecoveryOutcome::Unavailable => match credential_failure { + Some(failure) => record_credential_failure(selector, session, error, failure).await, + None => error, + }, + }; + } + match credential_failure { + Some(failure) => record_credential_failure(selector, session, error, failure).await, + None => error, + } +} + +pub(super) fn transport_credential_failure( + error: &GrokInferenceTransportError, + upstream_model: &UpstreamModelId, +) -> Option { + match error.kind() { + GrokInferenceTransportErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized), + GrokInferenceTransportErrorKind::PermissionDenied => { + Some(GrokCredentialFailure::AccessDenied) + } + GrokInferenceTransportErrorKind::RateLimited if model_capacity_failure(error) => { + Some(GrokCredentialFailure::ModelCapacity { + upstream_model: upstream_model.clone(), + }) + } + GrokInferenceTransportErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited { + retry_after: error.retry_after(), + }), + GrokInferenceTransportErrorKind::QuotaExhausted => { + Some(GrokCredentialFailure::QuotaExhausted) + } + GrokInferenceTransportErrorKind::FreeQuotaExhausted => { + Some(GrokCredentialFailure::FreeQuotaExhausted) + } + GrokInferenceTransportErrorKind::PaymentRequired => { + Some(GrokCredentialFailure::PaymentRequired { + retry_after: error.retry_after(), + }) + } + GrokInferenceTransportErrorKind::ModelQuotaExhausted => { + Some(GrokCredentialFailure::ModelQuotaExhausted { + upstream_model: upstream_model.clone(), + retry_after: error.retry_after(), + }) + } + GrokInferenceTransportErrorKind::ModelAccessDenied => { + Some(GrokCredentialFailure::ModelAccessDenied { + upstream_model: upstream_model.clone(), + retry_after: error.retry_after(), + }) + } + GrokInferenceTransportErrorKind::Unavailable if empty_upstream_failure(error) => { + Some(GrokCredentialFailure::EmptyUpstream) + } + GrokInferenceTransportErrorKind::Unavailable => { + Some(GrokCredentialFailure::UpstreamUnavailable) + } + _ => None, + } +} + +pub(super) fn model_capacity_failure(error: &GrokInferenceTransportError) -> bool { + transport_error_contains_any( + error, + &[ + "model_capacity", + "capacity", + "overloaded", + "server_busy", + "too many concurrent", + "engine_overloaded", + ], + ) +} + +pub(super) fn empty_upstream_failure(error: &GrokInferenceTransportError) -> bool { + transport_error_contains_any( + error, + &[ + "empty_upstream", + "empty model output", + "no content/tool_calls", + "no client-visible content", + "empty upstream", + ], + ) +} + +pub(super) fn transport_error_contains_any( + error: &GrokInferenceTransportError, + signals: &[&str], +) -> bool { + let code = error + .upstream_code() + .map(|code| code.as_str()) + .unwrap_or_default(); + let message = error + .client_visible_upstream_error() + .map(ClientVisibleUpstreamError::message) + .unwrap_or_default() + .to_ascii_lowercase(); + signals + .iter() + .any(|signal| code.contains(signal) || message.contains(signal)) +} + +pub(super) async fn map_and_record_stream_transport_failure( + selector: &dyn GrokSessionSelector, + session: &SelectedGrokSession, + error: GrokInferenceTransportError, + upstream_model: &UpstreamModelId, +) -> ProviderError { + let request_scoped = error.kind() == GrokInferenceTransportErrorKind::SafetyRejected; + let credential_failure = transport_credential_failure(&error, upstream_model); + let error = map_stream_error(error); + match credential_failure { + Some(failure) => record_credential_failure(selector, session, error, failure).await, + None if request_scoped => error, + None => record_stream_failure(selector, session, error, upstream_model).await, + } +} + +pub(super) fn map_continuation_failure( + context: &AttemptContext, + error: ProviderError, +) -> ProviderError { + let is_reasoning_decode_failure = context.continuation_attempt() == ContinuationAttempt::Native + && error.kind() == ProviderErrorKind::InvalidRequest + && error.upstream_status() == Some(400) + && error + .upstream_code() + .is_some_and(|code| code.as_str() == REASONING_DECODE_FAILED_CODE); + let is_missing_native_response = context.continuation_attempt() == ContinuationAttempt::Native + && error.kind() == ProviderErrorKind::InvalidRequest + && error.upstream_status() == Some(404) + && error + .upstream_code() + .is_some_and(|code| code.as_str() == RESPONSE_NOT_FOUND_CODE); + if is_reasoning_decode_failure || is_missing_native_response { + error + .with_continuation_failure(ContinuationFailure::HistoryUnavailable) + .with_replay_safe() + } else { + error + } +} + +pub(super) fn preflight_context(context: &AttemptContext) -> Result<(), ProviderError> { + if context.cancellation().is_cancelled() { + return Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )); + } + if remaining(context.deadline()).is_none() { + return Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::NotSent, + )); + } + Ok(()) +} + +pub(super) fn ensure_sent_context(context: &AttemptContext) -> Result<(), ProviderError> { + if context.cancellation().is_cancelled() { + return Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Sent, + )); + } + if remaining(context.deadline()).is_none() { + return Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + )); + } + Ok(()) +} + +pub(super) fn map_request_error(error: GrokRequestEncodeError) -> ProviderError { + let kind = match error { + GrokRequestEncodeError::InvalidProtocolPayload + | GrokRequestEncodeError::InvalidRequestNormalization + | GrokRequestEncodeError::InvalidRequestField { .. } => ProviderErrorKind::InvalidRequest, + GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol, + }; + let provider_error = provider_error(kind, UpstreamSendState::NotSent); + if kind != ProviderErrorKind::InvalidRequest { + return provider_error; + } + provider_error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + error.to_string(), + Some("invalid_request_normalization".to_owned()), + Some("invalid_request_error".to_owned()), + )) +} + +/// 将选择阶段失败映射为带结构化 code 与 retry_after 的 Provider 错误。 +pub(super) fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError { + let (retry_after, message, code) = match error { + GrokSessionSelectorError::AccountCoolingDown { retry_after } => ( + retry_after, + cooling_down_message(retry_after), + "account_cooling_down", + ), + GrokSessionSelectorError::ModelCoolingDown { retry_after } => ( + retry_after, + model_cooling_down_message(retry_after), + "model_cooling_down", + ), + GrokSessionSelectorError::CapacityUnavailable { retry_after } => ( + retry_after, + "account is at its concurrency or request-interval limit".to_owned(), + "account_capacity_busy", + ), + GrokSessionSelectorError::NoEligibleSession => ( + None, + "no account is eligible for the requested model".to_owned(), + "no_eligible_account", + ), + GrokSessionSelectorError::Unavailable => ( + None, + "account scheduling state is temporarily unreadable".to_owned(), + "account_selector_unavailable", + ), + GrokSessionSelectorError::InvalidSession => { + return provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent); + } + }; + let error = provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent); + let error = match retry_after { + Some(retry_after) => error.with_retry_after(retry_after), + None => error, + }; + error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + message, + Some(code.to_owned()), + Some("account_unavailable_error".to_owned()), + )) +} + +pub(super) fn cooling_down_message(retry_after: Option) -> String { + match retry_after { + Some(retry_after) => format!( + "account is cooling down after an upstream failure; retry in {}s", + retry_after.as_secs().saturating_add(1) + ), + None => "account is cooling down after an upstream failure".to_owned(), + } +} + +pub(super) fn model_cooling_down_message(retry_after: Option) -> String { + match retry_after { + Some(retry_after) => format!( + "all eligible accounts are cooling down for this model; retry in {}s", + retry_after.as_secs().saturating_add(1) + ), + None => "all eligible accounts are cooling down for this model".to_owned(), + } +} + +pub(super) fn map_transport_error_for_context( + error: GrokInferenceTransportError, + context: &AttemptContext, +) -> ProviderError { + let allow_explicit_replay = context.continuation().is_none() + || error.kind() == GrokInferenceTransportErrorKind::Unauthorized; + map_transport_error_with_state(error, None, allow_explicit_replay) +} + +pub(super) fn map_stream_error(error: GrokInferenceTransportError) -> ProviderError { + map_transport_error_with_state(error, Some(UpstreamSendState::Sent), false) +} + +pub(super) fn map_transport_error_with_state( + error: GrokInferenceTransportError, + forced_send_state: Option, + allow_explicit_replay: bool, +) -> ProviderError { + let transport_kind = error.kind(); + let kind = match transport_kind { + GrokInferenceTransportErrorKind::InvalidRequest => ProviderErrorKind::InvalidRequest, + GrokInferenceTransportErrorKind::Unsupported => ProviderErrorKind::Unsupported, + GrokInferenceTransportErrorKind::Unauthorized => ProviderErrorKind::Unauthorized, + GrokInferenceTransportErrorKind::PermissionDenied + | GrokInferenceTransportErrorKind::ModelAccessDenied + | GrokInferenceTransportErrorKind::PaymentRequired + | GrokInferenceTransportErrorKind::SafetyRejected => ProviderErrorKind::PermissionDenied, + GrokInferenceTransportErrorKind::RateLimited => ProviderErrorKind::RateLimited, + GrokInferenceTransportErrorKind::QuotaExhausted + | GrokInferenceTransportErrorKind::FreeQuotaExhausted + | GrokInferenceTransportErrorKind::ModelQuotaExhausted => ProviderErrorKind::QuotaExhausted, + GrokInferenceTransportErrorKind::Timeout => ProviderErrorKind::Timeout, + GrokInferenceTransportErrorKind::Transport => ProviderErrorKind::Transport, + GrokInferenceTransportErrorKind::Protocol => ProviderErrorKind::Protocol, + GrokInferenceTransportErrorKind::Unavailable => ProviderErrorKind::Unavailable, + GrokInferenceTransportErrorKind::Cancelled => ProviderErrorKind::Cancelled, + }; + let mut mapped = provider_error( + kind, + forced_send_state.unwrap_or_else(|| error.send_state()), + ); + if let Some(status) = error.status() { + mapped = mapped.with_status(status); + if allow_explicit_replay + && forced_send_state.is_none() + && explicit_rejection_is_replay_safe(&error, status) + { + mapped = mapped.with_replay_safe(); + } + } + 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 let Some(code) = error.upstream_code().cloned() { + mapped = mapped.with_upstream_code(code); + } + if let Some(detail) = error.client_visible_upstream_error().cloned() { + mapped = mapped.with_client_visible_upstream_error(detail); + } + if error.requires_credential_recovery() { + mapped = mapped.with_credential_recovery().with_replay_safe(); + } + if error.sensitive_context_was_redacted() { + mapped = mapped.redact_sensitive_context("upstream transport context"); + } + mapped +} + +pub(super) fn explicit_rejection_is_replay_safe( + error: &GrokInferenceTransportError, + status: u16, +) -> bool { + let kind = error.kind(); + if kind == GrokInferenceTransportErrorKind::SafetyRejected { + return false; + } + if matches!(status, 401 | 402 | 403 | 405 | 429 | 529) || (500..=599).contains(&status) { + return true; + } + status == 400 + && (matches!( + kind, + GrokInferenceTransportErrorKind::QuotaExhausted + | GrokInferenceTransportErrorKind::FreeQuotaExhausted + | GrokInferenceTransportErrorKind::ModelQuotaExhausted + | GrokInferenceTransportErrorKind::PaymentRequired + | GrokInferenceTransportErrorKind::Unavailable + ) || (kind == GrokInferenceTransportErrorKind::RateLimited + && model_capacity_failure(error))) +} + +pub(super) fn provider_error( + kind: ProviderErrorKind, + send_state: UpstreamSendState, +) -> ProviderError { + ProviderError::new(kind, send_state) +} + +pub(super) fn remaining(deadline: SystemTime) -> Option { + deadline + .duration_since(SystemTime::now()) + .ok() + .filter(|remaining| !remaining.is_zero()) +} diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs new file mode 100644 index 000000000..761ded783 --- /dev/null +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -0,0 +1,627 @@ +//! 官方 Grok Build 会话的 `gateway-core` Provider adapter。 + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant, SystemTime}; + +use async_trait::async_trait; +use futures::{StreamExt, future::BoxFuture}; +use gateway_core::account::{ + AccountEligibilityPolicy, AccountFeedbackStats, ProviderAccount, ProviderAccountId, + ProviderAccountStore, +}; +use gateway_core::engine::continuation::ContinuationBinding; +use gateway_core::engine::provider::{ + EventStream, Provider, ProviderCallMetadata, ProviderCatalogGeneration, + ProviderModelCapabilities, ProviderRequest, ProviderRequestObservation, + ProviderSelectionObservation, ProviderStream, +}; +use gateway_core::engine::{AttemptContext, ContinuationAttempt}; +use gateway_core::error::{ + ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind, +}; +use gateway_core::event::{ + GatewayEvent, ProviderEvent, ProviderResponseMetadata, ProviderResponseObservation, + ProviderResponseTimings, ResponseMeta, +}; +use gateway_core::operation::{ + Feature, GenerateRequest, Operation, OperationKind, ProviderSessionState, +}; +use gateway_core::routing::{ + ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, SupportLevel, + UpstreamModelId, +}; +use gateway_core::task::{ + ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId, + WorkerKind, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, + WorkerTaskError, +}; +use gateway_core::upstream::{UpstreamSendState, UpstreamTransport}; +use gateway_protocol::openai::codex_responses_request_semantics; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use url::Url; + +use crate::XaiWireProfileState; +use crate::credential::{ + GrokCredentialCatalogError, GrokCredentialCatalogService, GrokCredentialQuotaService, + GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshOutcome, + GrokCredentialRefreshService, GrokQuotaError, +}; +use crate::reasoning_replay::{ + GrokReasoningReplay, GrokReasoningReplayCapture, GrokReasoningReplayKey, + valid_reasoning_ciphertext, +}; +use crate::transport::canonical::GrokCanonicalDecoder; +use crate::transport::config::XAI_PROVIDER_NAME; +use crate::transport::headers::{GrokClientIdentity, build_grok_headers}; +use crate::transport::profile::{GROK_CLI_RELEASE_POLL_INTERVAL, GrokCliReleaseService}; +use crate::transport::{ + GROK_RESPONSES_URL, GrokCompactionDecodeError, GrokCompactionRequest, + GrokCompactionSummaryDecoder, GrokCredentialFailure, GrokInferenceChunkStream, + GrokInferenceRequest, GrokInferenceResponse, GrokInferenceTransport, + GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportMetrics, + GrokProviderConfigError, GrokQuotaFailureKind, GrokRequestEncodeError, GrokResponsesRequest, + GrokSessionAffinityKey, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, + SelectedGrokSession, classify_grok_quota_failure, +}; +use crate::{GrokCatalogCapabilityEvidence, GrokCatalogModel}; + +mod continuation; +mod failure; +mod stream; +mod workers; + +use continuation::*; +use failure::*; +use stream::*; +pub(crate) use workers::worker_contributions; + +const HTTP_SSE_TRANSPORT: &str = "http_sse"; +const DEFAULT_GROK_MODEL: &str = "grok-4.5"; +const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024; +const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096; +const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed"; +const RESPONSE_NOT_FOUND_CODE: &str = "not_found"; + +/// 官方 Grok Build Provider；会话选择与 HTTP SSE transport 均由外部注入。 +/// +/// 每次调用只选择一个 OAuth 会话。仅当 xAI 明确拒绝历史 reasoning 密文时， +/// 允许在同账号、同凭据、同会话绑定上剥离密文后有界重试一次。凭据轮换、 +/// endpoint fallback 以及公开 xAI API key 推理都不在该 adapter 内。 +pub struct GrokBuildProvider { + selector: Arc, + transport: Arc, + catalog: Arc, + credential_recovery: Arc, + account_feedback: Arc, + client_identity: GrokClientIdentity, + reasoning_replay: GrokReasoningReplay, + wire_profile: XaiWireProfileState, + responses_url: Url, +} + +impl GrokBuildProvider { + /// 在显式的会话与 transport 边界上创建 Provider。 + pub fn new( + selector: Arc, + transport: Arc, + catalog: Arc, + credential_recovery: Arc, + account_feedback: Arc, + wire_profile: XaiWireProfileState, + ) -> Result { + let responses_url = Url::parse(GROK_RESPONSES_URL) + .map_err(|_| GrokProviderConfigError::InvalidResponsesUrl)?; + Ok(Self { + selector, + transport, + catalog, + credential_recovery, + account_feedback, + client_identity: GrokClientIdentity::new(), + reasoning_replay: GrokReasoningReplay::new(), + wire_profile, + responses_url, + }) + } +} + +#[async_trait] +impl Provider for GrokBuildProvider { + fn name(&self) -> &'static str { + XAI_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 payload = request.protocol_payload(); + if payload.protocol() != "openai" { + return ProviderRequestObservation::default(); + } + let semantics = codex_responses_request_semantics(payload.body(), payload.context()); + ProviderRequestObservation { + reasoning_effort: semantics.reasoning_effort, + 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, + ) -> Result, ProviderError> { + let Ok(models) = self.catalog.query_models().await else { + return Ok(vec![default_grok_model_capabilities()?]); + }; + if models.is_empty() { + return Ok(vec![default_grok_model_capabilities()?]); + } + Ok(models + .into_iter() + .map(compile_grok_model_capabilities) + .collect()) + } + + async fn execute( + &self, + request: ProviderRequest, + context: AttemptContext, + ) -> Result { + let candidate = request.candidate(); + if candidate.provider().as_str() != XAI_PROVIDER_NAME { + return Err(provider_error( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + )); + } + preflight_context(&context)?; + + match request.operation() { + Operation::Generate(generate) => { + self.execute_generate(generate, candidate, context).await + } + _ => Err(provider_error( + ProviderErrorKind::Unsupported, + UpstreamSendState::NotSent, + )), + } + } +} + +impl GrokBuildProvider { + async fn execute_generate( + &self, + generate: &GenerateRequest, + candidate: &ProviderCandidate, + context: AttemptContext, + ) -> Result { + if crate::transport::compaction::has_terminal_compaction_trigger(generate) { + return self.execute_compaction(generate, candidate, context).await; + } + let upstream_model = candidate_upstream_model(candidate)?; + let previous_session = decode_xai_session_state(generate)?; + let continuation_account = continuation_account(&context, previous_session.as_ref())?; + let mut upstream_request = GrokResponsesRequest::encode( + generate, + upstream_model.as_str(), + context.client_api_key_ref(), + ) + .map_err(map_request_error)?; + let wire_upstream_model = UpstreamModelId::new( + upstream_request + .upstream_model() + .ok_or_else(protocol_not_sent)? + .to_owned(), + ) + .map_err(|_| protocol_not_sent())?; + let request_input = upstream_request.input_items(); + if let Some(previous) = previous_session.as_ref() { + upstream_request.inherit_session(previous.session_id.as_deref()); + } + let selection_started_at = Instant::now(); + let selected = select_grok_session( + self.selector.as_ref(), + candidate, + &wire_upstream_model, + &context, + continuation_account, + upstream_request.affinity().cloned(), + ) + .await?; + let account_selection_wait_ms = + u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + // 首字计时的起点：账号选择完成之后、上游建立之前。 + let output_started_at = Instant::now(); + apply_continuation( + &mut upstream_request, + previous_session.as_ref(), + &context, + selected.account_id(), + request_input.as_slice(), + )?; + let inherited_replay_session_id = context + .continuation() + .is_none() + .then(|| { + previous_session + .as_ref() + .and_then(|state| state.session_id.as_deref()) + }) + .flatten(); + let reasoning_replay_key = upstream_request + .reasoning_replay_session_id() + .or(inherited_replay_session_id) + .and_then(|session_id| { + self.reasoning_replay.key( + wire_upstream_model.as_str(), + session_id, + selected.account_id().as_str(), + ) + }); + if !upstream_request.has_previous_response_id() + && let (Some(key), Some(input)) = ( + reasoning_replay_key.as_ref(), + upstream_request.replay_input_items(), + ) + && let Some(input) = self.reasoning_replay.apply(key, &input) + { + upstream_request + .set_replay_input(input) + .map_err(map_request_error)?; + } + let reasoning_replay_capture = + reasoning_replay_key.map(|key| self.reasoning_replay.capture(key)); + let session_capture = (!matches!( + context.continuation(), + Some(ContinuationBinding::Pinned(_) | ContinuationBinding::External(_)) + ) || previous_session.is_some()) + .then(|| GrokSessionCapture { + previous: previous_session, + request_input, + account_id: selected.account_id().as_str().to_owned(), + session_id: upstream_request.session_id().map(str::to_owned), + output_items: BTreeMap::new(), + }); + let selected = Arc::new(selected); + let allows_account_state_mutation = selected.allows_account_state_mutation(); + let metadata = provider_call_metadata(candidate, &selected, account_selection_wait_ms)?; + let events = cold_http_sse_stream( + Arc::clone(&self.selector), + Arc::clone(&self.transport), + GrokStreamAttempt { + client_identity: self.client_identity.clone(), + wire_profile: self.wire_profile.clone(), + credential_recovery: Arc::clone(&self.credential_recovery), + responses_url: self.responses_url.clone(), + request: upstream_request, + upstream_model: wire_upstream_model, + context, + session: Arc::clone(&selected), + output_started_at, + session_capture, + reasoning_replay_capture, + }, + ); + let stream = ProviderStream::new(metadata, events, selected); + Ok(if allows_account_state_mutation { + stream.with_account_feedback(Arc::clone(&self.account_feedback)) + } else { + stream + }) + } + + async fn execute_compaction( + &self, + generate: &GenerateRequest, + candidate: &ProviderCandidate, + context: AttemptContext, + ) -> Result { + if context.continuation().is_some() + || context.continuation_attempt() != ContinuationAttempt::None + { + return Err(provider_error( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + )); + } + let previous_session = decode_xai_session_state(generate)?; + let operation_account = previous_session + .as_ref() + .map(|previous| ProviderAccountId::new(previous.account_id.clone())) + .transpose() + .map_err(|_| protocol_not_sent())?; + let inherited_session_id = previous_session + .as_ref() + .and_then(|previous| previous.session_id.clone()); + let upstream_model = candidate_upstream_model(candidate)?; + let upstream_request = GrokCompactionRequest::encode( + generate, + upstream_model.as_str(), + context.client_api_key_ref(), + ) + .map_err(map_request_error)?; + let wire_upstream_model = UpstreamModelId::new( + upstream_request + .upstream_model() + .ok_or_else(protocol_not_sent)? + .to_owned(), + ) + .map_err(|_| protocol_not_sent())?; + let explicit_replay_session_id = upstream_request + .reasoning_replay_session_id() + .map(str::to_owned); + let upstream_session_id = inherited_session_id + .clone() + .or_else(|| explicit_replay_session_id.clone()); + let selection_started_at = Instant::now(); + let selected = Arc::new( + select_grok_session( + self.selector.as_ref(), + candidate, + &wire_upstream_model, + &context, + operation_account, + upstream_request.affinity().cloned(), + ) + .await?, + ); + let account_selection_wait_ms = + u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + let reasoning_replay_key = explicit_replay_session_id + .as_deref() + .or(inherited_session_id.as_deref()) + .and_then(|session_id| { + self.reasoning_replay.key( + wire_upstream_model.as_str(), + session_id, + selected.account_id().as_str(), + ) + }); + let allows_account_state_mutation = selected.allows_account_state_mutation(); + let metadata = provider_call_metadata(candidate, &selected, account_selection_wait_ms)?; + let events = cold_compaction_http_sse_stream( + Arc::clone(&self.selector), + Arc::clone(&self.transport), + GrokCompactionStreamAttempt { + client_identity: self.client_identity.clone(), + wire_profile: self.wire_profile.clone(), + credential_recovery: Arc::clone(&self.credential_recovery), + responses_url: self.responses_url.clone(), + request: upstream_request, + upstream_model: wire_upstream_model, + upstream_session_id, + context, + session: Arc::clone(&selected), + reasoning_replay: self.reasoning_replay.clone(), + reasoning_replay_key, + }, + ); + let stream = ProviderStream::new(metadata, events, selected); + Ok(if allows_account_state_mutation { + stream.with_account_feedback(Arc::clone(&self.account_feedback)) + } else { + stream + }) + } +} + +async fn select_grok_session( + selector: &dyn GrokSessionSelector, + candidate: &ProviderCandidate, + wire_upstream_model: &UpstreamModelId, + context: &AttemptContext, + operation_account: Option, + affinity: Option, +) -> Result { + let required_account = context.required_account().cloned().or(operation_account); + let selection = GrokSessionSelection::new( + wire_upstream_model.clone(), + context.excluded_accounts().clone(), + required_account.clone(), + context.account_selection_policy(), + context.deadline(), + Arc::clone(candidate.account_scope()), + context.client_api_key_ref().clone(), + ) + .with_eligibility_policy(if context.is_diagnostic_required_account() { + AccountEligibilityPolicy::BypassForDiagnostic + } else { + AccountEligibilityPolicy::Enforce + }) + .with_affinity(affinity); + let selection_deadline = remaining(context.deadline()) + .ok_or_else(|| provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent))?; + let cancellation = context.cancellation().clone(); + let selected = tokio::select! { + biased; + _ = cancellation.cancelled() => Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )), + _ = tokio::time::sleep(selection_deadline) => Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::NotSent, + )), + selected = selector.select(selection) => selected.map_err(map_selection_error), + }?; + if context.excluded_accounts().contains(selected.account_id()) + || required_account + .as_ref() + .is_some_and(|required| required != selected.account_id()) + { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::NotSent, + )); + } + Ok(selected) +} + +fn provider_call_metadata( + candidate: &ProviderCandidate, + selected: &SelectedGrokSession, + account_selection_wait_ms: u64, +) -> Result { + Ok(ProviderCallMetadata::new( + ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?, + candidate_upstream_model(candidate)?.clone(), + selected.account_id().clone(), + UpstreamTransport::new(HTTP_SSE_TRANSPORT).map_err(|_| protocol_not_sent())?, + ) + .with_selection_observation(ProviderSelectionObservation::new( + account_selection_wait_ms, + selected.capacity_snapshot(), + ))) +} + +fn candidate_upstream_model( + candidate: &ProviderCandidate, +) -> Result<&UpstreamModelId, ProviderError> { + candidate.upstream_model().ok_or_else(protocol_not_sent) +} + +fn support(evidence: GrokCatalogCapabilityEvidence) -> SupportLevel { + match evidence { + GrokCatalogCapabilityEvidence::DeclaredNative => SupportLevel::Native, + GrokCatalogCapabilityEvidence::DeclaredUnsupported => SupportLevel::Unsupported, + GrokCatalogCapabilityEvidence::Unknown => SupportLevel::Unknown, + } +} + +fn tool_support(evidence: GrokCatalogCapabilityEvidence) -> SupportLevel { + match evidence { + GrokCatalogCapabilityEvidence::DeclaredNative => SupportLevel::Native, + GrokCatalogCapabilityEvidence::DeclaredUnsupported => SupportLevel::Unsupported, + // catalog 省略该可选字段时，Grok Build 的 Responses 工具协议仍可用； + // 请求 adapter 会在发送前规范化仅客户端侧的工具结构。 + GrokCatalogCapabilityEvidence::Unknown => SupportLevel::Emulated, + } +} + +fn compile_grok_model_capabilities(model: GrokCatalogModel) -> ProviderModelCapabilities { + let mut operations = BTreeSet::new(); + if model.capabilities().responses_api() == GrokCatalogCapabilityEvidence::DeclaredNative { + operations.insert(OperationKind::Generate); + } + let capabilities = ModelCapabilities::new( + operations, + model + .limits() + .max_output_tokens() + .map(std::num::NonZeroU64::get), + ) + .with_upstream_feature_validation() + .with_feature( + Feature::Reasoning, + support(model.capabilities().reasoning_effort()), + ) + .with_feature( + Feature::Tools, + tool_support(model.capabilities().streaming_tool_calls()), + ) + .with_feature(Feature::Vision, SupportLevel::Unknown) + .with_feature(Feature::JsonSchema, SupportLevel::Unknown) + .with_feature(Feature::NativeContinuation, SupportLevel::Native); + ProviderModelCapabilities::new(model.request_model().clone(), capabilities) + .with_presentation(grok_model_presentation(&model)) +} + +fn default_grok_model_capabilities() -> Result { + let model = UpstreamModelId::new(DEFAULT_GROK_MODEL.to_owned()) + .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; + let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None) + .with_upstream_feature_validation(); + Ok(ProviderModelCapabilities::new(model, capabilities) + .with_presentation(default_grok_model_presentation())) +} + +fn default_grok_model_presentation() -> ModelPresentation { + ModelPresentation::new( + Some("Grok 4.5".to_owned()), + Some("xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()), + ) + .with_context_window_tokens(Some(500_000)) + .with_image_input(true) + .with_agent_tools(true, true) +} + +fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation { + let slug = model.request_model().as_str(); + let known_grok_4_5 = is_known_grok_4_5_model(slug); + let reasoning_evidence = model.capabilities().reasoning_effort(); + let catalog_reasoning_efforts = model + .capabilities() + .reasoning_efforts() + .iter() + .map(|effort| effort.as_str().to_owned()) + .collect::>(); + let catalog_default_reasoning = model + .capabilities() + .default_reasoning_effort() + .map(|effort| effort.as_str().to_owned()); + let (default_reasoning, reasoning_efforts) = match reasoning_evidence { + GrokCatalogCapabilityEvidence::DeclaredNative if !catalog_reasoning_efforts.is_empty() => { + let default_reasoning = catalog_default_reasoning + .filter(|default| { + catalog_reasoning_efforts + .iter() + .any(|effort| effort == default) + }) + .or_else(|| catalog_reasoning_efforts.first().cloned()); + (default_reasoning, catalog_reasoning_efforts) + } + GrokCatalogCapabilityEvidence::DeclaredNative => (None, Vec::new()), + GrokCatalogCapabilityEvidence::DeclaredUnsupported => { + (Some("none".to_owned()), vec!["none".to_owned()]) + } + GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()), + }; + let context_window_tokens = model + .limits() + .context_window_tokens() + .map(std::num::NonZeroU64::get) + .or(known_grok_4_5.then_some(500_000)); + let tool_evidence = model.capabilities().streaming_tool_calls(); + + ModelPresentation::new( + model + .display_name() + .map(str::to_owned) + .or_else(|| known_grok_4_5.then(|| "Grok 4.5".to_owned())), + model + .metadata() + .description() + .map(str::to_owned) + .or_else(|| { + known_grok_4_5 + .then(|| "xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()) + }), + ) + .with_reasoning(default_reasoning, reasoning_efforts) + .with_context_window_tokens(context_window_tokens) + .with_image_input(known_grok_4_5) + .with_agent_tools( + tool_evidence != GrokCatalogCapabilityEvidence::DeclaredUnsupported, + tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative + || (known_grok_4_5 && tool_evidence == GrokCatalogCapabilityEvidence::Unknown), + ) + .with_search_tool( + model.capabilities().backend_search() == GrokCatalogCapabilityEvidence::DeclaredNative, + ) + .with_hidden(model.metadata().hidden().unwrap_or(false)) +} + +fn is_known_grok_4_5_model(slug: &str) -> bool { + matches!( + slug, + DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest" + ) +} diff --git a/backend/crates/providers/xai/src/provider/stream.rs b/backend/crates/providers/xai/src/provider/stream.rs new file mode 100644 index 000000000..308dae6ab --- /dev/null +++ b/backend/crates/providers/xai/src/provider/stream.rs @@ -0,0 +1,596 @@ +//! xAI inference/compaction 流执行与终态校验。 + +use super::*; + +pub(super) struct GrokStreamAttempt { + pub(super) client_identity: GrokClientIdentity, + pub(super) wire_profile: XaiWireProfileState, + pub(super) credential_recovery: Arc, + pub(super) responses_url: Url, + pub(super) request: GrokResponsesRequest, + pub(super) upstream_model: UpstreamModelId, + pub(super) context: AttemptContext, + pub(super) session: Arc, + pub(super) output_started_at: Instant, + pub(super) session_capture: Option, + pub(super) reasoning_replay_capture: Option, +} + +pub(super) struct GrokCompactionStreamAttempt { + pub(super) client_identity: GrokClientIdentity, + pub(super) wire_profile: XaiWireProfileState, + pub(super) credential_recovery: Arc, + pub(super) responses_url: Url, + pub(super) request: GrokCompactionRequest, + pub(super) upstream_model: UpstreamModelId, + pub(super) upstream_session_id: Option, + pub(super) context: AttemptContext, + pub(super) session: Arc, + pub(super) reasoning_replay: GrokReasoningReplay, + pub(super) reasoning_replay_key: Option, +} + +pub(super) struct AcceptedGrokInference { + response: GrokInferenceResponse, + observation: ProviderResponseObservation, +} + +pub(super) async fn next_grok_chunk( + body: &mut GrokInferenceChunkStream, + selector: &dyn GrokSessionSelector, + session: &SelectedGrokSession, + upstream_model: &UpstreamModelId, + context: &AttemptContext, +) -> Result, ProviderError> { + let Some(stream_deadline) = remaining(context.deadline()) else { + return Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + )); + }; + let cancellation = context.cancellation().clone(); + tokio::select! { + biased; + _ = cancellation.cancelled() => Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Sent, + )), + _ = tokio::time::sleep(stream_deadline) => Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + )), + chunk = body.next() => match chunk { + Some(Ok(chunk)) => Ok(Some(chunk)), + Some(Err(error)) => Err(map_and_record_stream_transport_failure( + selector, + session, + error, + upstream_model, + ).await), + None => Ok(None), + }, + } +} + +pub(super) fn cold_compaction_http_sse_stream( + selector: Arc, + transport: Arc, + attempt: GrokCompactionStreamAttempt, +) -> EventStream { + let GrokCompactionStreamAttempt { + client_identity, + wire_profile, + credential_recovery, + responses_url, + mut request, + upstream_model, + upstream_session_id, + context, + session, + reasoning_replay, + reasoning_replay_key, + } = attempt; + Box::pin(async_stream::try_stream! { + let headers = build_grok_headers( + &wire_profile, + &session, + &client_identity, + context.request_id(), + upstream_session_id.as_deref(), + None, + &upstream_model, + ); + let mut invalid_encrypted_content_retried = false; + let accepted = loop { + if context.cancellation().is_cancelled() { + Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + ))?; + return; + } + let body = request.to_json_bytes().map_err(map_request_error)?; + let inference_request = GrokInferenceRequest::new( + responses_url.clone(), + headers.clone(), + body, + session.binding().clone(), + ); + let Some(handshake_deadline) = remaining(context.deadline()) else { + Err(mark_transient_compaction_failure(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::NotSent, + )))?; + return; + }; + let cancellation = context.cancellation().clone(); + let boundary = tokio::select! { + biased; + _ = cancellation.cancelled() => InferenceBoundary::Cancelled, + _ = tokio::time::sleep(handshake_deadline) => InferenceBoundary::Deadline, + response = transport.execute(inference_request) => InferenceBoundary::Response(response), + }; + match boundary { + InferenceBoundary::Cancelled => { + Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Ambiguous, + ))?; + return; + } + InferenceBoundary::Deadline => { + Err(mark_transient_compaction_failure(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Ambiguous, + )))?; + return; + } + InferenceBoundary::Response(Ok(response)) => { + let observation = xai_response_observation(&response) + .map_err(mark_transient_compaction_failure)?; + break AcceptedGrokInference { + response, + observation, + }; + } + InferenceBoundary::Response(Err(error)) => { + if !invalid_encrypted_content_retried + && is_invalid_encrypted_content_failure(&error) + && request.strip_invalid_encrypted_reasoning() + { + invalid_encrypted_content_retried = true; + continue; + } + let observation = xai_error_observation(&error).ok(); + let credential_failure = + transport_credential_failure(&error, &upstream_model); + let error = map_continuation_failure( + &context, + map_transport_error_for_context(error, &context), + ); + let error = recover_or_record_failure( + selector.as_ref(), + credential_recovery.as_ref(), + &session, + error, + credential_failure, + context.credential_recovery_attempted(), + ) + .await; + if let Some(observation) = observation { + yield ProviderEvent::observation(observation); + } + Err(mark_transient_compaction_failure(error))?; + return; + } + } + }; + yield ProviderEvent::observation(accepted.observation); + + let mut body = accepted.response.into_body(); + let mut canonical = GrokCanonicalDecoder::new(upstream_model.as_str()); + let mut summary = GrokCompactionSummaryDecoder::new(); + let mut facts = CompactionFacts::default(); + + 'stream: while let Some(chunk) = next_grok_chunk( + &mut body, + selector.as_ref(), + &session, + &upstream_model, + &context, + ) + .await + .map_err(mark_transient_compaction_failure)? + { + let events = match canonical.push(&chunk) { + Ok(events) => events, + Err(error) => { + let error = map_continuation_failure(&context, error); + let error = record_stream_failure( + selector.as_ref(), + &session, + error, + &upstream_model, + ) + .await; + Err(mark_transient_compaction_failure(error))?; + return; + } + }; + for event in events { + summary.observe(&event).map_err(map_compaction_decode_error)?; + facts.observe(&event); + if facts.completed.is_some() { + break 'stream; + } + } + } + + if facts.completed.is_none() { + let events = match canonical.finish_without_terminal() { + Ok(events) => events, + Err(error) => { + let error = map_continuation_failure(&context, error); + let error = record_stream_failure( + selector.as_ref(), + &session, + error, + &upstream_model, + ) + .await; + Err(mark_transient_compaction_failure(error))?; + return; + } + }; + for event in events { + summary.observe(&event).map_err(map_compaction_decode_error)?; + facts.observe(&event); + if facts.completed.is_some() { + break; + } + } + } + + let started = facts + .started + .ok_or_else(|| mark_transient_compaction_failure(protocol_sent()))?; + let completed = facts + .completed + .ok_or_else(|| mark_transient_compaction_failure(protocol_sent()))?; + let (summary, encrypted_content) = summary + .finish_with_encrypted_content() + .map_err(map_compaction_decode_error)?; + let (created, output_done, terminal) = crate::transport::compaction::compaction_wire_events( + &started, + &completed, + summary.as_deref(), + &encrypted_content, + facts.created_response.as_ref(), + facts.terminal_response.as_ref(), + ) + .map_err(|_| mark_transient_compaction_failure(protocol_sent()))? + .into_parts(); + ensure_sent_context(&context)?; + if session.allows_account_state_mutation() { + selector.record_success(&session).await; + } + if let Some(key) = reasoning_replay_key.as_ref() { + reasoning_replay.clear(key); + } + yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created); + yield ProviderEvent::wire(output_done); + let mut terminal_facts = facts.metering; + terminal_facts.push(GatewayEvent::Completed(completed)); + yield ProviderEvent::canonical_with_wire(terminal_facts, terminal); + }) +} + +#[derive(Default)] +pub(super) struct CompactionFacts { + started: Option, + completed: Option, + metering: Vec, + created_response: Option, + terminal_response: Option, +} + +impl CompactionFacts { + fn observe(&mut self, event: &ProviderEvent) { + self.capture_wire_response(event); + for fact in event.canonical_facts() { + match fact { + GatewayEvent::Started(meta) if self.started.is_none() => { + self.started = Some(meta.clone()); + } + GatewayEvent::Completed(meta) if self.completed.is_none() => { + self.completed = Some(meta.clone()); + } + GatewayEvent::Usage(_) + | GatewayEvent::CalculatedCost(_) + | GatewayEvent::ProviderCost(_) => self.metering.push(fact.clone()), + _ => {} + } + } + } + + fn capture_wire_response(&mut self, event: &ProviderEvent) { + let Some(wire) = event.wire_event().filter(|wire| wire.has_json_data()) else { + return; + }; + let event_type = wire + .event_type() + .or_else(|| wire.data().get("type").and_then(Value::as_str)); + let response = wire.data().get("response").cloned(); + match event_type { + Some("response.created" | "response.in_progress") + if self.created_response.is_none() => + { + self.created_response = response; + } + Some("response.completed" | "response.incomplete") + if self.terminal_response.is_none() => + { + self.terminal_response = response; + } + _ => {} + } + } +} + +pub(super) fn map_compaction_decode_error(error: GrokCompactionDecodeError) -> ProviderError { + match error { + GrokCompactionDecodeError::MissingEncryptedContent => { + mark_transient_compaction_failure(protocol_sent()) + } + } +} + +pub(super) fn mark_transient_compaction_failure(error: ProviderError) -> ProviderError { + if matches!( + error.kind(), + ProviderErrorKind::RateLimited + | ProviderErrorKind::Timeout + | ProviderErrorKind::Transport + | ProviderErrorKind::Protocol + | ProviderErrorKind::Unavailable + ) { + error.with_pre_delivery_retry() + } else { + error + } +} + +pub(super) fn cold_http_sse_stream( + selector: Arc, + transport: Arc, + attempt: GrokStreamAttempt, +) -> EventStream { + let GrokStreamAttempt { + client_identity, + wire_profile, + credential_recovery, + responses_url, + mut request, + upstream_model, + context, + session, + output_started_at, + mut session_capture, + mut reasoning_replay_capture, + } = attempt; + Box::pin(async_stream::try_stream! { + if context.cancellation().is_cancelled() { + Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + ))?; + } + let headers = build_grok_headers( + &wire_profile, + &session, + &client_identity, + context.request_id(), + request.session_id(), + None, + &upstream_model, + ); + let cancellation = context.cancellation().clone(); + let mut invalid_encrypted_content_retried = false; + let response = loop { + let body = request.to_json_bytes().map_err(map_request_error)?; + let inference_request = GrokInferenceRequest::new( + responses_url.clone(), + headers.clone(), + body, + session.binding().clone(), + ); + let Some(handshake_deadline) = remaining(context.deadline()) else { + Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::NotSent, + ))?; + return; + }; + let boundary = tokio::select! { + biased; + _ = cancellation.cancelled() => InferenceBoundary::Cancelled, + _ = tokio::time::sleep(handshake_deadline) => InferenceBoundary::Deadline, + response = transport.execute(inference_request) => InferenceBoundary::Response(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)) => break response, + InferenceBoundary::Response(Err(error)) => { + if !invalid_encrypted_content_retried + && is_invalid_encrypted_content_failure(&error) + && request.strip_invalid_encrypted_reasoning() + { + invalid_encrypted_content_retried = true; + if let Some(capture) = session_capture.as_mut() { + capture.request_input = request.input_items(); + } + continue; + } + let observation = xai_error_observation(&error)?; + let credential_failure = transport_credential_failure(&error, &upstream_model); + let error = map_continuation_failure( + &context, + map_transport_error_for_context(error, &context), + ); + let error = recover_or_record_failure( + selector.as_ref(), + credential_recovery.as_ref(), + &session, + error, + credential_failure, + context.credential_recovery_attempted(), + ) + .await; + yield ProviderEvent::observation(observation); + Err(error)?; + return; + } + } + }; + + let observation = xai_response_observation(&response)?; + let base_timings = observation.timings(); + let mut first_token_ms: Option = None; + yield ProviderEvent::observation(observation.clone()); + + let mut body = response.into_body(); + let mut decoder = GrokCanonicalDecoder::for_request(upstream_model.as_str(), &request); + loop { + let Some(stream_deadline) = remaining(context.deadline()) else { + Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + ))?; + return; + }; + let next = tokio::select! { + biased; + _ = cancellation.cancelled() => Err(provider_error( + ProviderErrorKind::Cancelled, + UpstreamSendState::Sent, + )), + _ = tokio::time::sleep(stream_deadline) => Err(provider_error( + ProviderErrorKind::Timeout, + UpstreamSendState::Sent, + )), + chunk = body.next() => match chunk { + Some(Ok(chunk)) => Ok(Some(chunk)), + Some(Err(error)) => Err(map_and_record_stream_transport_failure( + selector.as_ref(), + &session, + error, + &upstream_model, + ).await), + None => Ok(None), + }, + }?; + let Some(chunk) = next else { + break; + }; + let mut events = match decoder.push(&chunk) { + Ok(events) => events, + Err(error) => { + let error = map_continuation_failure(&context, error); + let error = record_stream_failure( + selector.as_ref(), + &session, + error, + &upstream_model, + ) + .await; + Err(error)?; + return; + } + }; + // 首个非前导输出事件（结构帧也算）即上报携带 first_token_ms 的观测， + // 供 Core 覆盖会话级兜底值。 + if decoder.take_output_start() && first_token_ms.is_none() { + first_token_ms = Some( + u64::try_from(output_started_at.elapsed().as_millis()).unwrap_or(u64::MAX), + ); + yield ProviderEvent::observation( + observation + .clone() + .with_timings(ProviderResponseTimings { + first_token_ms, + ..base_timings + }), + ); + } + let completed = events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .any(|event| matches!(event, GatewayEvent::Completed(_))); + attach_xai_session_update(&mut events, &mut session_capture)?; + if let Some(capture) = reasoning_replay_capture.as_mut() { + capture.observe(&events); + } + if completed && session.allows_account_state_mutation() { + selector.record_success(&session).await; + } + for event in events { + ensure_sent_context(&context)?; + yield event; + } + if completed { + return; + } + } + let mut final_events = match decoder.finish() { + Ok(events) => events, + Err(error) => { + let error = map_continuation_failure(&context, error); + let error = record_stream_failure( + selector.as_ref(), + &session, + error, + &upstream_model, + ) + .await; + Err(error)?; + return; + } + }; + let completed = final_events + .iter() + .flat_map(ProviderEvent::canonical_facts) + .any(|event| matches!(event, GatewayEvent::Completed(_))); + // 尾部 finish 补全缓冲中的首个输出帧时，同样上报首字。 + if decoder.take_output_start() && first_token_ms.is_none() { + first_token_ms = Some( + u64::try_from(output_started_at.elapsed().as_millis()).unwrap_or(u64::MAX), + ); + yield ProviderEvent::observation( + observation + .clone() + .with_timings(ProviderResponseTimings { + first_token_ms, + ..base_timings + }), + ); + } + attach_xai_session_update(&mut final_events, &mut session_capture)?; + if let Some(capture) = reasoning_replay_capture.as_mut() { + capture.observe(&final_events); + } + if completed && session.allows_account_state_mutation() { + selector.record_success(&session).await; + } + for event in final_events { + ensure_sent_context(&context)?; + yield event; + } + }) +} diff --git a/backend/crates/providers/xai/src/provider/workers.rs b/backend/crates/providers/xai/src/provider/workers.rs new file mode 100644 index 000000000..66696be4a --- /dev/null +++ b/backend/crates/providers/xai/src/provider/workers.rs @@ -0,0 +1,228 @@ +//! xAI Provider 向 Host 贡献的后台 worker。 + +use super::*; + +pub(super) const WORKER_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +pub(super) const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60); +pub(super) const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); +pub(super) const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); +pub(super) const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); +pub(super) const QUOTA_CATALOG_INTERVAL: Duration = Duration::from_secs(5 * 60); +// rolling 24h 描述的是上游用量窗口，不代表从本次观测起封禁 24 小时； +// 缺少可信 reset 时间时按短周期探测策略恢复检查。 +pub(super) const EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL: Duration = Duration::from_secs(10 * 60); +pub(super) const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = QUOTA_CATALOG_INTERVAL; +pub(super) const CLI_RELEASE_WORKER_OWNER: &str = "xai-cli-release"; + +pub(crate) fn worker_contributions( + refresh: Arc, + quota: Arc, + catalog: Arc, + accounts: Arc, + provider_kind: ProviderKind, + cli_release: Arc, +) -> Result, WorkerDefinitionError> { + let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, XAI_PROVIDER_NAME)?; + let catalog_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, XAI_PROVIDER_NAME)?; + let release_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, CLI_RELEASE_WORKER_OWNER)?; + 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, + quota, + catalog, + provider_kind, + last_periodic_refresh_at: Mutex::new(BTreeMap::new()), + }), + )?), + WorkerContribution::Registration(scheduled_registration( + release_id, + GROK_CLI_RELEASE_POLL_INTERVAL, + Box::new(XaiCliReleaseTask { + service: cli_release, + }), + )?), + ]) +} + +pub(super) 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, + }, + ) +} + +pub(super) struct XaiOAuthRefreshTask { + service: Arc, +} + +pub(super) struct XaiCliReleaseTask { + service: Arc, +} + +impl ScheduledTask for XaiCliReleaseTask { + 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 { + tracing::warn!(error = %error, "xAI CLI release check failed"); + } + Ok(()) + }) + } +} + +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(|error| { + tracing::error!(error = %error, "xAI OAuth refresh cycle failed"); + WorkerTaskError::safe("xAI OAuth refresh failed") + })?; + let failures = outcomes + .iter() + .filter(|outcome| { + matches!( + outcome, + GrokCredentialRefreshOutcome::Ambiguous { .. } + | GrokCredentialRefreshOutcome::Transient { .. } + | GrokCredentialRefreshOutcome::Failed { .. } + ) + }) + .count(); + if !outcomes.is_empty() { + tracing::info!( + accounts = outcomes.len(), + failures, + "xAI OAuth refresh cycle completed" + ); + } + if failures > 0 { + tracing::warn!(failures, "xAI OAuth refresh cycle contained failures"); + } + Ok(()) + }) + } +} + +pub(super) struct XaiQuotaCatalogTask { + accounts: Arc, + quota: Arc, + catalog: Arc, + provider_kind: ProviderKind, + last_periodic_refresh_at: Mutex>, +} + +impl ScheduledTask for XaiQuotaCatalogTask { + fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> { + Box::pin(async move { + let accounts = self + .accounts + .list_for_provider(&self.provider_kind) + .await + .map_err(|_| WorkerTaskError::safe("xAI Provider accounts unavailable"))?; + let mut failures = 0_u64; + let now = SystemTime::now(); + let accounts = self.reserve_periodic_refreshes(accounts, now); + for account in accounts { + 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), + } + } + match self.catalog.query_models().await { + Ok(_) | Err(GrokCredentialCatalogError::NoEligibleCredential) => {} + Err(_) => failures = failures.saturating_add(1), + } + if failures == 0 { + Ok(()) + } else { + Err(WorkerTaskError::safe( + "xAI quota or catalog synchronization failed", + )) + } + }) + } +} + +impl XaiQuotaCatalogTask { + fn reserve_periodic_refreshes( + &self, + accounts: Vec, + now: SystemTime, + ) -> Vec { + let candidates = accounts + .into_iter() + .filter(|account| eligible_quota_worker_account(account, now)) + .collect::>(); + let candidate_ids = candidates + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let now = Instant::now(); + let mut last_periodic_refresh_at = self + .last_periodic_refresh_at + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + last_periodic_refresh_at.retain(|account_id, _| candidate_ids.contains(account_id)); + candidates + .into_iter() + .filter(|account| { + let due = last_periodic_refresh_at + .get(account.id()) + .is_none_or(|last| { + now.saturating_duration_since(*last) + >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL + }); + if due { + last_periodic_refresh_at.insert(account.id().clone(), now); + } + due + }) + .collect() + } +} + +pub(super) fn eligible_quota_worker_account(account: &ProviderAccount, now: SystemTime) -> bool { + account.enabled() + && account + .access_token_expires_at() + .is_some_and(|expires_at| expires_at > now) + && account + .quota() + .exhaustion_refresh_due(now, EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL) +} diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs deleted file mode 100644 index 64e86ecd5..000000000 --- a/backend/crates/providers/xai/src/transport/request.rs +++ /dev/null @@ -1,4043 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; -use std::fmt; - -use gateway_core::operation::GenerateRequest; -use gateway_core::policy::ClientApiKeyId; -use serde_json::{Map, Value, json}; -use sha2::{Digest as _, Sha256}; - -use super::{GrokSessionAffinityKey, XAI_PROVIDER_NAME}; -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", - "previousResponseId", - "previous_response_id", - "responseId", - "response_id", - "sessionId", - "session_id", - "x-grok-agent-id", - "x-grok-conv-id", - "x-grok-session-id", -]; -const SESSION_FIELDS: &[&str] = &[ - "prompt_cache_key", - "session_id", - "sessionId", - "conversation_id", - "conversationId", -]; -const PROTOCOL_CONTEXT_SESSION_FIELDS: &[&str] = &["conversation_id", "session_id", "thread_id"]; -const MAX_SESSION_SEED_BYTES: usize = 1_024; -const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"]; -/// Grok CLI 在历史条目上注入、Grok Build 无法反序列化的内部键；只在已知 -/// 注入位置剥离，避免误伤工具 schema 或输出中恰好同名的键。 -const GROK_INTERNAL_HISTORY_KEYS: &[&str] = - &["phase", "internal_chat_message_metadata_passthrough"]; - -/// 保留客户端 OpenAI Responses object 的 xAI 上游请求。 -pub struct GrokResponsesRequest { - body: Map, - session_id: Option, - reasoning_replay_session_id: Option, - affinity: Option, - response_transform: GrokResponseTransform, -} - -impl GrokResponsesRequest { - /// 返回发送到 `/v1/responses` 的 JSON object。 - #[must_use] - pub const fn body(&self) -> &Map { - &self.body - } - - /// 返回按下游租户隔离后的稳定 Grok 会话 UUID。 - #[must_use] - pub fn session_id(&self) -> Option<&str> { - self.session_id.as_deref() - } - - /// 返回只有显式会话身份才启用的 reasoning replay scope。 - pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> { - self.reasoning_replay_session_id.as_deref() - } - - /// 返回归一化后的 xAI wire 模型。 - pub(crate) fn upstream_model(&self) -> Option<&str> { - self.body.get("model").and_then(Value::as_str) - } - - pub(crate) fn has_previous_response_id(&self) -> bool { - self.body - .get("previous_response_id") - .and_then(Value::as_str) - .map(str::trim) - .is_some_and(|value| !value.is_empty()) - } - - pub(crate) fn replay_input_items(&self) -> Option> { - self.body.get("input").and_then(Value::as_array).cloned() - } - - /// 返回与会话一致、额外绑定模型的账号亲和键。 - #[must_use] - pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> { - self.affinity.as_ref() - } - - pub(crate) fn response_transform(&self) -> GrokResponseTransform { - self.response_transform.clone() - } - - pub(crate) fn input_items(&self) -> Vec { - match self.body.get("input") { - Some(Value::Array(items)) => items.clone(), - Some(Value::String(input)) => vec![json_object([ - ("type", Value::String("message".to_owned())), - ("role", Value::String("user".to_owned())), - ("content", Value::String(input.clone())), - ])], - _ => Vec::new(), - } - } - - pub(crate) fn set_replay_input( - &mut self, - input: Vec, - ) -> Result<(), GrokRequestEncodeError> { - let mut normalizer = ToolNormalizer::for_replay(self.response_transform.clone()); - let mut normalized = Vec::with_capacity(input.len()); - for item in input { - match item { - Value::Object(item) if string_field(&item, "type") == "custom_tool_call" => { - normalized.push(Value::Object( - normalizer.normalize_custom_tool_call_input(&item)?, - )); - } - item => normalized.push(item), - } - } - self.response_transform = normalizer.response; - self.body - .insert("input".to_owned(), Value::Array(normalized)); - Ok(()) - } - - /// 为同账号的一次 xAI `invalid_encrypted_content` 恢复请求移除被拒绝的密文。 - /// - /// 只在至少一个 reasoning item 含非空密文时改写；可读 summary/content、ID 与 - /// 其他历史保持不变，剥离后只剩 `type` 的空壳一并删除。 - pub(crate) fn strip_invalid_encrypted_reasoning(&mut self) -> bool { - strip_invalid_encrypted_reasoning_from_body(&mut self.body) - } - - pub(crate) fn set_previous_response_id(&mut self, response_id: Option) { - match response_id { - Some(response_id) => { - self.body.insert( - "previous_response_id".to_owned(), - Value::String(response_id), - ); - } - None => { - self.body.remove("previous_response_id"); - } - } - } - - pub(crate) fn inherit_session(&mut self, session_id: Option<&str>) { - let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else { - self.clear_session(); - return; - }; - self.session_id = Some(session_id.to_owned()); - if self.reasoning_replay_session_id.is_some() { - self.reasoning_replay_session_id = Some(session_id.to_owned()); - } - self.affinity = None; - self.body.insert( - "prompt_cache_key".to_owned(), - Value::String(session_id.to_owned()), - ); - } - - pub(crate) fn clear_session(&mut self) { - self.body.remove("prompt_cache_key"); - self.session_id = None; - self.reasoning_replay_session_id = None; - self.affinity = None; - } - - pub fn encode( - request: &GenerateRequest, - upstream_model: &str, - client_api_key_ref: &ClientApiKeyId, - ) -> Result { - Self::encode_inner(request, upstream_model, client_api_key_ref, true, false) - } - - pub(super) fn encode_compaction_source( - request: &GenerateRequest, - upstream_model: &str, - client_api_key_ref: &ClientApiKeyId, - ) -> Result { - Self::encode_inner(request, upstream_model, client_api_key_ref, false, true) - } - - fn encode_inner( - request: &GenerateRequest, - upstream_model: &str, - client_api_key_ref: &ClientApiKeyId, - allow_cache_route: bool, - should_consume_terminal_compaction_trigger: bool, - ) -> Result { - let payload = request.protocol_payload(); - if payload.protocol() != "openai" { - return Err(GrokRequestEncodeError::InvalidProtocolPayload); - } - let mut body = payload.body().clone(); - if should_consume_terminal_compaction_trigger { - consume_terminal_compaction_trigger(&mut body)?; - } - let upstream_model = resolve_grok_text_responses_model_id(upstream_model); - // 这些字段属于 Codex/OpenAI 侧请求控制，不是 xAI 上游协议字段。 - // OpenAI 透明路径会保留未知字段；这里只在 xAI adapter 内做最小剥离。 - body.remove("provider_options"); - body.remove("service_tier"); - let session_seed = explicit_session_seed(request, &body); - let enable_cache_route = allow_cache_route && session_seed.is_some(); - let identity = resolve_session_identity( - client_api_key_ref.as_str(), - &upstream_model, - session_seed.as_deref(), - &body, - ); - sanitize_account_identity(&mut body); - sanitize_client_metadata(&mut body); - normalize_build_request(&mut body, &upstream_model)?; - let mut response_transform = normalize_responses_request(&mut body)?; - response_transform.observe_client_cache_tools(); - if enable_cache_route { - enable_grok_prompt_cache_route(&mut body, &upstream_model, &mut response_transform); - } - let (session_id, affinity) = identity.map_or((None, None), |(session_id, affinity)| { - (Some(session_id), Some(affinity)) - }); - let reasoning_replay_session_id = - session_seed.is_some().then(|| session_id.clone()).flatten(); - match session_id.as_ref() { - Some(session_id) => { - body.insert( - "prompt_cache_key".to_owned(), - Value::String(session_id.clone()), - ); - } - None => { - body.remove("prompt_cache_key"); - } - } - body.insert("model".to_owned(), Value::String(upstream_model)); - body.insert("stream".to_owned(), Value::Bool(true)); - Ok(Self { - body, - session_id, - reasoning_replay_session_id, - affinity, - response_transform, - }) - } - - pub(crate) fn to_json_bytes(&self) -> Result, GrokRequestEncodeError> { - serde_json::to_vec(&self.body).map_err(|_| GrokRequestEncodeError::Serialization) - } -} - -fn resolve_grok_text_responses_model_id(model: &str) -> String { - let stripped = strip_grok_provider_prefix(model); - let normalized = stripped.to_ascii_lowercase(); - match normalized.as_str() { - "grok" | "grok-latest" | "grok-4.5" | "grok-4.5-latest" | "grok-build-latest" => { - "grok-4.5".to_owned() - } - "grok-4.6" | "grok-4.6-latest" => "grok-4.6".to_owned(), - "grok-4.3" | "grok-4.3-latest" => "grok-4.3".to_owned(), - "grok-3-mini" - | "grok-3-mini-fast" - | "grok-build-0.1" - | "grok-composer-2.5-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-0309-non-reasoning" - | "grok-4.20-multi-agent-0309" => normalized, - "grok-build" => "grok-build-0.1".to_owned(), - "grok-composer" | "composer-2.5" => "grok-composer-2.5-fast".to_owned(), - "grok-4.20-reasoning" => "grok-4.20-0309-reasoning".to_owned(), - "grok-4.20-non-reasoning" => "grok-4.20-0309-non-reasoning".to_owned(), - "grok-4.20-multi-agent" | "grok-4.20-multi-agent-latest" => { - "grok-4.20-multi-agent-0309".to_owned() - } - _ => stripped.to_owned(), - } -} - -fn consume_terminal_compaction_trigger( - body: &mut Map, -) -> Result<(), GrokRequestEncodeError> { - let Some(Value::Array(input)) = body.get_mut("input") else { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - }; - let Some(Value::Object(last)) = input.last() else { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - }; - if last.get("type").and_then(Value::as_str) != Some("compaction_trigger") { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - input.pop(); - Ok(()) -} - -impl fmt::Debug for GrokResponsesRequest { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("GrokResponsesRequest") - .field("body_keys", &self.body.keys().collect::>()) - .field("has_session", &self.session_id.is_some()) - .field( - "has_response_transform", - &!self.response_transform.is_empty(), - ) - .field("body", &"") - .finish() - } -} - -/// Generate 到 Responses 的编码错误，不保留 option 与 prompt 值。 -#[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, - /// JSON 序列化意外失败。 - #[error("Grok Build request serialization failed")] - Serialization, - /// Responses 兼容字段无法安全归一化。 - #[error("Grok Build request normalization failed")] - InvalidRequestNormalization, - /// 请求中的具体字段无法安全转换为 Grok Build 接受的形态。 - #[error("Grok Build request field `{field}` could not be normalized safely")] - InvalidRequestField { field: &'static str }, -} - -impl GrokRequestEncodeError { - fn at_field(self, field: &'static str) -> Self { - match self { - Self::InvalidRequestNormalization => Self::InvalidRequestField { field }, - error => error, - } - } -} - -fn sanitize_account_identity(body: &mut Map) { - for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { - body.remove(*field); - } - let Some(Value::Object(metadata)) = body.get_mut("metadata") else { - return; - }; - for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { - metadata.remove(*field); - } -} - -fn sanitize_client_metadata(body: &mut Map) { - // Codex 的 client_metadata 是本地 transport envelope，可能包含工作目录、仓库地址、 - // installation/session 标识。会话亲和信息已在调用本函数前提取，整个 envelope 都不能 - // 越过 Grok Build 边界。 - body.remove("client_metadata"); -} - -fn enable_grok_prompt_cache_route( - body: &mut Map, - upstream_model: &str, - response: &mut GrokResponseTransform, -) { - let tools = body - .entry("tools".to_owned()) - .or_insert_with(|| Value::Array(Vec::new())); - let Value::Array(tools) = tools else { - return; - }; - response.observe_upstream_cache_tools(tools); - if is_cache_media_model(upstream_model) || has_tool_type(tools, "image_generation") { - return; - } - if tools.is_empty() { - for tool in GROK_CACHE_ROUTE_TOOLS { - tools.push(json_object([("type", Value::String((*tool).to_owned()))])); - response.mark_injected_cache_tool(tool); - } - // 无客户端工具时以 none 选中缓存路由，同时不授予搜索能力。 - body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); - } else if !has_tool_type(tools, "x_search") { - tools.push(json_object([( - "type", - Value::String("x_search".to_owned()), - )])); - response.mark_injected_cache_tool("x_search"); - } -} - -fn is_cache_media_model(model: &str) -> bool { - let model = model.trim().to_ascii_lowercase(); - ["image", "imagine", "video"] - .into_iter() - .any(|marker| model.contains(marker)) -} - -fn explicit_session_seed(request: &GenerateRequest, body: &Map) -> Option { - request - .prompt_cache_key() - .and_then(valid_session_seed) - .map(ToOwned::to_owned) - .or_else(|| first_session_value(body)) - .or_else(|| { - body.get("metadata") - .and_then(Value::as_object) - .and_then(first_session_value) - }) - .or_else(|| { - PROTOCOL_CONTEXT_SESSION_FIELDS.iter().find_map(|field| { - request - .protocol_payload() - .context() - .get(*field) - .and_then(Value::as_str) - .and_then(valid_session_seed) - .map(ToOwned::to_owned) - }) - }) -} - -fn first_session_value(body: &Map) -> Option { - let prompt_cache_key = body - .get("prompt_cache_key") - .and_then(Value::as_str) - .and_then(valid_session_seed) - .map(ToOwned::to_owned); - prompt_cache_key - .or_else(|| { - let metadata = body.get("metadata")?.as_object()?; - ["session_id", "sessionId"] - .into_iter() - .find_map(|field| metadata.get(field).and_then(Value::as_str)) - .and_then(valid_session_seed) - .map(ToOwned::to_owned) - .or_else(|| { - metadata - .get("user_id") - .and_then(Value::as_str) - .and_then(session_seed_from_user_id) - }) - }) - .or_else(|| { - SESSION_FIELDS[1..] - .iter() - .find_map(|field| body.get(*field).and_then(Value::as_str)) - .and_then(valid_session_seed) - .map(ToOwned::to_owned) - }) -} - -fn valid_session_seed(value: &str) -> Option<&str> { - let value = value.trim(); - (!value.is_empty() - && value.len() <= MAX_SESSION_SEED_BYTES - && !value.chars().any(char::is_control)) - .then_some(value) -} - -fn session_seed_from_user_id(value: &str) -> Option { - let value = value.trim(); - if let Ok(Value::Object(embedded)) = serde_json::from_str::(value) - && let Some(seed) = ["session_id", "sessionId"] - .into_iter() - .find_map(|field| embedded.get(field).and_then(Value::as_str)) - .and_then(valid_session_seed) - { - return Some(seed.to_owned()); - } - value - .rfind("_session_") - .and_then(|index| valid_session_seed(&value[index + "_session_".len()..])) - .map(ToOwned::to_owned) -} - -fn resolve_session_identity( - client_api_key_ref: &str, - upstream_model: &str, - explicit_seed: Option<&str>, - body: &Map, -) -> Option<(String, GrokSessionAffinityKey)> { - let model = upstream_model.trim().to_ascii_lowercase(); - if client_api_key_ref.is_empty() || model.is_empty() { - return None; - } - if let Some(seed) = explicit_seed { - let upstream_source = - format!("xai:build-session:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{seed}"); - let affinity_source = format!( - "xai:build-affinity:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{model}:{seed}" - ); - return Some(( - digest_uuid(&upstream_source), - GrokSessionAffinityKey::from_digest(Sha256::digest(affinity_source).into()), - )); - } - let (system, first_user) = message_anchors(body); - let first_user = truncate_anchor(&first_user, 200); - if first_user.is_empty() { - return None; - } - let system = truncate_anchor(&system, 100); - let upstream_source = format!( - "xai:build-soft-session:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{system}:{first_user}" - ); - let affinity_source = format!( - "xai:build-soft-affinity:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{model}:{system}:{first_user}" - ); - Some(( - digest_uuid(&upstream_source), - GrokSessionAffinityKey::from_digest(Sha256::digest(affinity_source).into()), - )) -} - -fn digest_uuid(source: &str) -> String { - let digest = Sha256::digest(source); - let value = digest[..16] - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::(); - format!( - "{}-{}-{}-{}-{}", - &value[0..8], - &value[8..12], - &value[12..16], - &value[16..20], - &value[20..32] - ) -} - -fn truncate_anchor(value: &str, max_chars: usize) -> String { - value.trim().chars().take(max_chars).collect() -} - -fn message_anchors(body: &Map) -> (String, String) { - let mut system = body - .get("instructions") - .map(flatten_message_content) - .filter(|value| !value.is_empty()) - .or_else(|| body.get("system").map(flatten_message_content)) - .unwrap_or_default(); - let mut first_user = String::new(); - if let Some(messages) = body.get("messages").and_then(Value::as_array) { - collect_role_anchors(messages, &mut system, &mut first_user); - } - if first_user.is_empty() { - match body.get("input") { - Some(Value::String(value)) => first_user = value.trim().to_owned(), - Some(Value::Array(items)) => { - collect_role_anchors(items, &mut system, &mut first_user); - } - _ => {} - } - } - (system, first_user) -} - -fn collect_role_anchors(items: &[Value], system: &mut String, first_user: &mut String) { - for item in items { - let Some(item) = item.as_object() else { - continue; - }; - let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default(); - if !item_type.is_empty() && item_type != "message" { - continue; - } - let role = item - .get("role") - .and_then(Value::as_str) - .unwrap_or_default() - .trim() - .to_ascii_lowercase(); - let content = item - .get("content") - .map(flatten_message_content) - .filter(|value| !value.is_empty()) - .or_else(|| { - item.get("text") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - }) - .unwrap_or_default(); - if content.is_empty() { - continue; - } - match role.as_str() { - "system" | "developer" if system.is_empty() => *system = content, - "user" if first_user.is_empty() => *first_user = content, - "" if first_user.is_empty() => *first_user = content, - _ => {} - } - if !first_user.is_empty() && !system.is_empty() { - break; - } - } -} - -fn flatten_message_content(value: &Value) -> String { - match value { - Value::String(value) => value.trim().to_owned(), - Value::Array(parts) => parts - .iter() - .filter_map(Value::as_object) - .filter(|part| { - matches!( - part.get("type").and_then(Value::as_str).unwrap_or_default(), - "" | "text" | "input_text" | "output_text" - ) - }) - .filter_map(|part| part.get("text").and_then(Value::as_str).map(str::trim)) - .filter(|text| !text.is_empty()) - .collect::>() - .join("\n"), - _ => String::new(), - } -} - -fn normalize_responses_request( - body: &mut Map, -) -> Result { - if let Some(response_format) = body.remove("response_format") { - let text = body - .entry("text".to_owned()) - .or_insert_with(|| Value::Object(Map::new())); - if text.is_null() { - *text = Value::Object(Map::new()); - } - let text = text - .as_object_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "text" })?; - if text.get("format").is_none_or(Value::is_null) { - text.insert( - "format".to_owned(), - normalize_response_format(response_format) - .map_err(|error| error.at_field("response_format"))?, - ); - } - } - patch_reasoning_text_types(body); - ToolNormalizer::new().normalize(body) -} - -fn normalize_build_request( - body: &mut Map, - upstream_model: &str, -) -> Result<(), GrokRequestEncodeError> { - apply_build_response_defaults(body)?; - normalize_build_reasoning_effort(body, upstream_model); - sanitize_build_model_fields(body, upstream_model); - Ok(()) -} - -fn apply_build_response_defaults( - body: &mut Map, -) -> Result<(), GrokRequestEncodeError> { - if body.get("store").is_none_or(Value::is_null) { - body.insert("store".to_owned(), Value::Bool(false)); - } - - let include = body - .entry("include".to_owned()) - .or_insert_with(|| Value::Array(Vec::new())); - if include.is_null() { - *include = Value::Array(Vec::new()); - } - let include = include - .as_array_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; - if include.iter().any(|value| !value.is_string()) { - return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); - } - if !include - .iter() - .any(|value| value.as_str() == Some("reasoning.encrypted_content")) - { - include.push(Value::String("reasoning.encrypted_content".to_owned())); - } - Ok(()) -} - -fn normalize_build_reasoning_effort(body: &mut Map, upstream_model: &str) { - let supports_effort = grok_model_supports_reasoning_effort(upstream_model); - if let Some(reasoning) = body.get_mut("reasoning").and_then(Value::as_object_mut) { - if let Some(effort) = reasoning.remove("effort") - && supports_effort - && let Some(effort) = normalize_grok_reasoning_effort_value(&effort) - { - reasoning.insert("effort".to_owned(), Value::String(effort.to_owned())); - } - if reasoning.is_empty() { - body.remove("reasoning"); - } - } - - if let Some(effort) = body.remove("reasoning_effort") - && supports_effort - && let Some(effort) = normalize_grok_reasoning_effort_value(&effort) - { - body.insert( - "reasoning_effort".to_owned(), - Value::String(effort.to_owned()), - ); - } - - let camel_effort = body.remove("reasoningEffort"); - if !body.contains_key("reasoning_effort") - && supports_effort - && let Some(effort) = camel_effort - .as_ref() - .and_then(normalize_grok_reasoning_effort_value) - { - body.insert( - "reasoning_effort".to_owned(), - Value::String(effort.to_owned()), - ); - } - - if is_grok_composer_model(upstream_model) { - body.remove("reasoning"); - body.remove("reasoning_effort"); - body.remove("reasoningEffort"); - } -} - -fn normalize_grok_reasoning_effort_value(value: &Value) -> Option<&'static str> { - let normalized = value - .as_str()? - .trim() - .to_ascii_lowercase() - .chars() - .filter(|character| !matches!(character, '-' | '_' | ' ')) - .collect::(); - match normalized.as_str() { - "none" => Some("none"), - "minimal" | "low" => Some("low"), - "medium" => Some("medium"), - "high" | "xhigh" | "extrahigh" | "max" | "ultra" => Some("high"), - _ => None, - } -} - -fn strip_grok_provider_prefix(model: &str) -> &str { - let model = model.trim(); - let lower = model.to_ascii_lowercase(); - ["xai/", "x-ai/", "grok/"] - .into_iter() - .find_map(|prefix| { - lower - .starts_with(prefix) - .then(|| model[prefix.len()..].trim()) - }) - .unwrap_or(model) -} - -fn grok_model_last_slug(model: &str) -> String { - model - .trim() - .rsplit_once('/') - .map_or_else(|| model.trim(), |(_, slug)| slug.trim()) - .to_ascii_lowercase() -} - -fn is_grok_composer_model(model: &str) -> bool { - matches!( - grok_model_last_slug(model).as_str(), - "grok-composer" | "grok-composer-2.5-fast" | "composer-2.5" - ) -} - -fn grok_model_supports_reasoning_effort(model: &str) -> bool { - matches!( - strip_grok_provider_prefix(model) - .to_ascii_lowercase() - .as_str(), - "grok-4.5" - | "grok-4.5-latest" - | "grok-4.6" - | "grok-4.6-latest" - | "grok-4.3" - | "grok-4.3-latest" - | "grok-3-mini" - | "grok-3-mini-fast" - | "grok-4.20-0309-reasoning" - | "grok-4.20-reasoning" - | "grok-4.20-multi-agent-0309" - ) -} - -fn sanitize_build_model_fields(body: &mut Map, upstream_model: &str) { - for field in ["prompt_cache_retention", "safety_identifier"] { - body.remove(field); - } - if upstream_model.trim().eq_ignore_ascii_case("grok-4.5") { - for field in [ - "presence_penalty", - "presencePenalty", - "frequency_penalty", - "frequencyPenalty", - "stop", - ] { - body.remove(field); - } - } - if grok_model_last_slug(upstream_model).starts_with("grok-4.20") { - body.remove("logprobs"); - body.remove("top_logprobs"); - } - delete_json_field_recursive(body, "external_web_access"); -} - -fn delete_json_field_recursive(object: &mut Map, field: &str) { - object.remove(field); - for value in object.values_mut() { - delete_json_field_from_value(value, field); - } -} - -fn delete_json_field_from_value(value: &mut Value, field: &str) { - match value { - Value::Object(child) => delete_json_field_recursive(child, field), - Value::Array(children) => { - for child in children { - delete_json_field_from_value(child, field); - } - } - _ => {} - } -} - -fn normalize_response_format(value: Value) -> Result { - let Some(format) = value.as_object() else { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - }; - if format.get("type").and_then(Value::as_str) != Some("json_schema") { - return Ok(value); - } - let Some(schema) = format.get("json_schema").and_then(Value::as_object) else { - return Ok(value); - }; - let mut normalized = Map::new(); - normalized.insert("type".to_owned(), Value::String("json_schema".to_owned())); - normalized.extend( - schema - .iter() - .filter(|(key, _)| key.as_str() != "type") - .map(|(key, value)| (key.clone(), value.clone())), - ); - Ok(Value::Object(normalized)) -} - -fn patch_reasoning_text_types(body: &mut Map) { - let Some(items) = body.get_mut("input").and_then(Value::as_array_mut) else { - return; - }; - for item in items { - let Some(item) = item.as_object_mut() else { - continue; - }; - if item.get("type").and_then(Value::as_str) != Some("reasoning") { - continue; - } - let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else { - continue; - }; - for part in content { - let Some(part) = part.as_object_mut() else { - continue; - }; - part.entry("type".to_owned()) - .or_insert_with(|| Value::String("reasoning_text".to_owned())); - } - } -} - -const MAX_BUILD_TOOL_ALIAS_LENGTH: usize = 128; -const MAX_TOOL_SEARCH_DESCRIPTION_BYTES: usize = 16 << 10; -const MAX_BUFFERED_FUNCTION_ARGUMENTS_BYTES: usize = 1 << 20; -const MAX_TOTAL_BUFFERED_FUNCTION_ARGUMENTS_BYTES: usize = 4 << 20; -const MAX_JSON_NUMBER_TEXT_BYTES: usize = 256; -const MAX_EXACT_JSON_INTEGER_TEXT: &str = "9007199254740991"; -const MAX_EXACT_JSON_INTEGER: u64 = 9_007_199_254_740_991; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -enum ToolKind { - Function, - Custom, - ToolSearch, - ApplyPatch, -} - -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] -struct ToolIdentity { - kind: ToolKind, - namespace: String, - name: String, -} - -impl ToolIdentity { - fn new(kind: ToolKind, namespace: &str, name: &str) -> Self { - Self { - kind, - namespace: namespace.to_owned(), - name: name.to_owned(), - } - } - - fn is_root_custom_apply_patch(&self) -> bool { - self.kind == ToolKind::Custom && self.namespace.is_empty() && self.name == "apply_patch" - } - - fn custom_argument_field(&self) -> &'static str { - if self.is_root_custom_apply_patch() { - "patch" - } else { - "input" - } - } -} - -#[derive(Debug, Clone)] -struct StreamCallState { - identity: ToolIdentity, - schema: Option, - arguments: String, - passthrough: bool, - last_delta: Option>, - added_payload: Option>, -} - -#[derive(Clone, Default)] -pub(crate) struct GrokResponseTransform { - aliases: BTreeMap, - function_schemas: BTreeMap, - visible_tools: Vec, - legacy_local_shell: bool, - filter_x_search: bool, - injected_tool_types: BTreeSet, - client_declared_tools: BTreeSet, - dropped_output_indexes: BTreeSet, - dropped_item_ids: BTreeSet, - stream_calls: BTreeMap, - stream_keys: BTreeMap, - stream_argument_bytes: usize, - stream_sequence_next: Option, -} - -pub(crate) struct GrokTransformedWireEvent { - event_type: String, - value: Value, -} - -impl GrokTransformedWireEvent { - #[must_use] - pub(crate) fn event_type(&self) -> &str { - &self.event_type - } - - #[must_use] - pub(crate) fn into_value(self) -> Value { - self.value - } -} - -impl GrokResponseTransform { - #[must_use] - pub(crate) fn is_empty(&self) -> bool { - self.aliases.is_empty() - && self.function_schemas.is_empty() - && self.visible_tools.is_empty() - && !self.legacy_local_shell - && !self.filter_x_search - && self.injected_tool_types.is_empty() - } - - fn observe_client_cache_tools(&mut self) { - for tool in &self.visible_tools { - let kind = tool.get("type").and_then(Value::as_str).unwrap_or_default(); - if kind == "x_search" { - self.filter_x_search = true; - } - if matches!(kind, "function" | "custom") - && let Some(name) = tool - .get("name") - .and_then(Value::as_str) - .map(str::trim) - .filter(|name| !name.is_empty()) - { - self.client_declared_tools.insert(name.to_owned()); - } - } - } - - fn observe_upstream_cache_tools(&mut self, tools: &[Value]) { - if has_tool_type(tools, "x_search") { - self.filter_x_search = true; - } - } - - fn mark_injected_cache_tool(&mut self, kind: &str) { - self.injected_tool_types.insert(kind.to_owned()); - if kind == "x_search" { - self.filter_x_search = true; - } - } - - pub(crate) fn resequence_stream_value(&mut self, value: &mut Value) { - let Some(payload) = value.as_object_mut() else { - return; - }; - let Some(raw_sequence) = payload.get("sequence_number") else { - return; - }; - if self.stream_sequence_next.is_none() { - let Some(sequence) = exact_nonnegative_sequence(raw_sequence) else { - return; - }; - self.stream_sequence_next = Some(sequence); - } - let Some(next) = self.stream_sequence_next.as_mut() else { - return; - }; - payload.insert( - "sequence_number".to_owned(), - Value::Number(serde_json::Number::from(*next)), - ); - *next = next.saturating_add(1); - } - - pub(crate) fn rewrite_stream_event( - &mut self, - event_type: &str, - mut value: Value, - ) -> Result, GrokRequestEncodeError> { - if self.cache_filter_enabled() { - if value - .get("item") - .is_some_and(|item| self.is_internal_cache_call(item)) - { - self.record_dropped_cache_item(&value); - return Ok(Vec::new()); - } - self.filter_cache_envelope(&mut value)?; - if self.references_dropped_cache_item(&value) { - return Ok(Vec::new()); - } - self.compact_cache_output_index(&mut value); - } - if event_type == "response.output_item.added" - && let Some(item) = value.get("item").and_then(Value::as_object) - && let Some(primary) = self.remember_stream_call(item) - && self - .stream_calls - .get(&primary) - .is_some_and(|state| state.identity.kind == ToolKind::ApplyPatch) - { - if let Some(state) = self.stream_calls.get_mut(&primary) { - state.added_payload = value.as_object().cloned(); - } - return Ok(Vec::new()); - } - - if event_type == "response.function_call_arguments.delta" - && let Some(primary) = self.stream_call_key(&value) - { - let delta = value - .get("delta") - .and_then(Value::as_str) - .unwrap_or_default(); - let state = self - .stream_calls - .get_mut(&primary) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - match state.identity.kind { - ToolKind::Custom => { - state.arguments.push_str(delta); - state.last_delta = value.as_object().cloned().map(|mut payload| { - payload.remove("delta"); - payload - }); - return Ok(Vec::new()); - } - ToolKind::ToolSearch | ToolKind::ApplyPatch => return Ok(Vec::new()), - ToolKind::Function if state.schema.is_some() => { - let call_limit_exceeded = state.arguments.len().saturating_add(delta.len()) - > MAX_BUFFERED_FUNCTION_ARGUMENTS_BYTES; - let total_limit_exceeded = - self.stream_argument_bytes.saturating_add(delta.len()) - > MAX_TOTAL_BUFFERED_FUNCTION_ARGUMENTS_BYTES; - if !state.passthrough && (call_limit_exceeded || total_limit_exceeded) { - let buffered = std::mem::take(&mut state.arguments); - self.stream_argument_bytes = - self.stream_argument_bytes.saturating_sub(buffered.len()); - state.passthrough = true; - state.last_delta = None; - let mut output = Vec::with_capacity(2); - if !buffered.is_empty() { - let mut flushed = value.clone(); - flushed - .as_object_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? - .insert("delta".to_owned(), Value::String(buffered)); - output.push(GrokTransformedWireEvent { - event_type: event_type.to_owned(), - value: flushed, - }); - } - output.push(GrokTransformedWireEvent { - event_type: event_type.to_owned(), - value, - }); - return Ok(output); - } - if !state.passthrough { - state.arguments.push_str(delta); - self.stream_argument_bytes = - self.stream_argument_bytes.saturating_add(delta.len()); - state.last_delta = value.as_object().cloned().map(|mut payload| { - payload.remove("delta"); - payload - }); - return Ok(Vec::new()); - } - } - ToolKind::Function => {} - } - } - - if event_type == "response.function_call_arguments.done" - && let Some(primary) = self.stream_call_key(&value) - { - let state = self - .stream_calls - .get(&primary) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - match state.identity.kind { - ToolKind::ToolSearch | ToolKind::ApplyPatch => return Ok(Vec::new()), - ToolKind::Custom => { - let arguments = value - .get("arguments") - .and_then(Value::as_str) - .filter(|arguments| !arguments.is_empty()) - .unwrap_or(&state.arguments); - let input = decode_custom_tool_input(&state.identity, arguments); - let mut output = Vec::with_capacity(2); - if let Some(delta) = state.last_delta.as_ref() { - output.push(GrokTransformedWireEvent { - event_type: "response.custom_tool_call_input.delta".to_owned(), - value: custom_tool_stream_payload( - delta, - "response.custom_tool_call_input.delta", - "delta", - &input, - ), - }); - } - output.push(GrokTransformedWireEvent { - event_type: "response.custom_tool_call_input.done".to_owned(), - value: custom_tool_stream_payload( - value - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?, - "response.custom_tool_call_input.done", - "input", - &input, - ), - }); - if let Some(state) = self.stream_calls.get_mut(&primary) { - state.arguments.clear(); - state.last_delta = None; - } - return Ok(output); - } - ToolKind::Function if state.schema.is_some() => { - let schema = state - .schema - .clone() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let passthrough = state.passthrough; - let buffered = state.arguments.clone(); - let last_delta = state.last_delta.clone(); - let arguments = value - .get("arguments") - .and_then(Value::as_str) - .filter(|arguments| !arguments.is_empty()) - .unwrap_or(&buffered); - let normalized = normalize_function_arguments(arguments, &schema) - .unwrap_or_else(|| arguments.to_owned()); - if passthrough { - value - .as_object_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? - .insert("arguments".to_owned(), Value::String(normalized)); - return Ok(vec![GrokTransformedWireEvent { - event_type: event_type.to_owned(), - value, - }]); - } - - let mut output = Vec::with_capacity(2); - if let Some(mut delta) = last_delta { - delta.insert("delta".to_owned(), Value::String(normalized.clone())); - output.push(GrokTransformedWireEvent { - event_type: "response.function_call_arguments.delta".to_owned(), - value: Value::Object(delta), - }); - } - value - .as_object_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? - .insert("arguments".to_owned(), Value::String(normalized)); - output.push(GrokTransformedWireEvent { - event_type: event_type.to_owned(), - value, - }); - if let Some(state) = self.stream_calls.get_mut(&primary) { - self.stream_argument_bytes = self - .stream_argument_bytes - .saturating_sub(state.arguments.len()); - state.arguments.clear(); - state.last_delta = None; - } - return Ok(output); - } - ToolKind::Function => {} - } - } - - if event_type == "response.output_item.done" - && let Some(item) = value.get("item").and_then(Value::as_object) - && self - .aliases - .get(string_field(item, "name")) - .is_some_and(|identity| identity.kind == ToolKind::ApplyPatch) - { - return self.rewrite_apply_patch_done_event(value); - } - - let completed_call = (event_type == "response.output_item.done") - .then(|| value.get("item").and_then(Value::as_object)) - .flatten() - .and_then(|item| self.primary_stream_call(item)); - self.rewrite_response_value(&mut value)?; - if let Some(response) = value.get_mut("response").and_then(Value::as_object_mut) { - self.restore_visible_tools(response); - } - if let Some(primary) = completed_call { - self.take_stream_call(&primary); - } - Ok(vec![GrokTransformedWireEvent { - event_type: event_type.to_owned(), - value, - }]) - } - - fn remember_stream_call(&mut self, item: &Map) -> Option { - if string_field(item, "type") != "function_call" { - return None; - } - let identity = self.aliases.get(string_field(item, "name"))?.clone(); - let keys = ["id", "call_id"] - .into_iter() - .filter_map(|field| item.get(field).and_then(Value::as_str)) - .filter(|key| !key.is_empty()) - .map(ToOwned::to_owned) - .collect::>(); - let primary = keys.first()?.clone(); - self.stream_calls.insert( - primary.clone(), - StreamCallState { - identity, - schema: self - .function_schemas - .get(string_field(item, "name")) - .cloned(), - arguments: String::new(), - passthrough: false, - last_delta: None, - added_payload: None, - }, - ); - for key in keys { - self.stream_keys.insert(key, primary.clone()); - } - Some(primary) - } - - fn stream_call_key(&mut self, payload: &Value) -> Option { - for field in ["item_id", "call_id"] { - if let Some(key) = payload.get(field).and_then(Value::as_str) - && let Some(primary) = self.stream_keys.get(key) - { - return Some(primary.clone()); - } - } - let alias = payload.get("name").and_then(Value::as_str)?; - let identity = self.aliases.get(alias)?.clone(); - let primary = ["item_id", "call_id"] - .into_iter() - .find_map(|field| payload.get(field).and_then(Value::as_str))? - .to_owned(); - self.stream_calls.insert( - primary.clone(), - StreamCallState { - identity, - schema: self.function_schemas.get(alias).cloned(), - arguments: String::new(), - passthrough: false, - last_delta: None, - added_payload: None, - }, - ); - self.stream_keys.insert(primary.clone(), primary.clone()); - Some(primary) - } - - fn primary_stream_call(&self, item: &Map) -> Option { - ["id", "call_id"] - .into_iter() - .filter_map(|field| item.get(field).and_then(Value::as_str)) - .find_map(|key| self.stream_keys.get(key).cloned()) - } - - fn take_stream_call(&mut self, primary: &str) -> Option { - self.stream_keys.retain(|_, owner| owner != primary); - let state = self.stream_calls.remove(primary)?; - self.stream_argument_bytes = self - .stream_argument_bytes - .saturating_sub(state.arguments.len()); - Some(state) - } - - fn rewrite_response_value(&self, value: &mut Value) -> Result<(), GrokRequestEncodeError> { - match value { - Value::Array(values) => { - for value in values { - self.rewrite_response_value(value)?; - } - } - Value::Object(object) => { - for value in object.values_mut() { - self.rewrite_response_value(value)?; - } - match string_field(object, "type") { - "function_call" => self.rewrite_function_call(object)?, - "shell_call" if self.legacy_local_shell => { - rewrite_legacy_local_shell_call(object) - } - _ => {} - } - } - _ => {} - } - Ok(()) - } - - fn rewrite_function_call( - &self, - call: &mut Map, - ) -> Result<(), GrokRequestEncodeError> { - let Some(identity) = self.aliases.get(string_field(call, "name")) else { - return Ok(()); - }; - match identity.kind { - ToolKind::Function => { - let alias = string_field(call, "name").to_owned(); - if let Some(schema) = self.function_schemas.get(&alias) - && let Some(arguments) = call.get("arguments").and_then(Value::as_str) - && let Some(normalized) = normalize_function_arguments(arguments, schema) - { - call.insert("arguments".to_owned(), Value::String(normalized)); - } - call.insert("name".to_owned(), Value::String(identity.name.clone())); - if identity.namespace.is_empty() { - call.remove("namespace"); - } else { - call.insert( - "namespace".to_owned(), - Value::String(identity.namespace.clone()), - ); - } - } - ToolKind::Custom => { - call.insert( - "type".to_owned(), - Value::String("custom_tool_call".to_owned()), - ); - call.insert("name".to_owned(), Value::String(identity.name.clone())); - if identity.namespace.is_empty() { - call.remove("namespace"); - } else { - call.insert( - "namespace".to_owned(), - Value::String(identity.namespace.clone()), - ); - } - let input = decode_custom_tool_input(identity, string_field(call, "arguments")); - call.insert("input".to_owned(), Value::String(input)); - call.remove("arguments"); - } - ToolKind::ToolSearch => { - call.insert( - "type".to_owned(), - Value::String("tool_search_call".to_owned()), - ); - call.insert("execution".to_owned(), Value::String("client".to_owned())); - let arguments = decode_tool_search_arguments(call.get("arguments")); - call.insert("arguments".to_owned(), arguments); - call.remove("name"); - call.remove("namespace"); - } - ToolKind::ApplyPatch => { - let operation = decode_apply_patch_arguments(call.get("arguments"))?; - call.insert( - "type".to_owned(), - Value::String("apply_patch_call".to_owned()), - ); - call.insert("operation".to_owned(), Value::Object(operation)); - call.remove("name"); - call.remove("namespace"); - call.remove("arguments"); - } - } - Ok(()) - } - - fn rewrite_apply_patch_done_event( - &mut self, - mut value: Value, - ) -> Result, GrokRequestEncodeError> { - let original_item = value - .get("item") - .and_then(Value::as_object) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - self.rewrite_response_value(&mut value)?; - let done = value - .as_object() - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let done_item = done - .get("item") - .and_then(Value::as_object) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let state = self - .primary_stream_call(&original_item) - .and_then(|primary| self.take_stream_call(&primary)); - let mut added = state - .and_then(|state| state.added_payload) - .unwrap_or_else(|| { - Map::from_iter([( - "type".to_owned(), - Value::String("response.output_item.added".to_owned()), - )]) - }); - added.insert( - "type".to_owned(), - Value::String("response.output_item.added".to_owned()), - ); - for key in ["output_index", "sequence_number"] { - if !added.contains_key(key) - && let Some(value) = done.get(key) - { - added.insert(key.to_owned(), value.clone()); - } - } - let mut added_item = done_item; - added_item.insert("status".to_owned(), Value::String("in_progress".to_owned())); - added.insert("item".to_owned(), Value::Object(added_item)); - Ok(vec![ - GrokTransformedWireEvent { - event_type: "response.output_item.added".to_owned(), - value: Value::Object(added), - }, - GrokTransformedWireEvent { - event_type: "response.output_item.done".to_owned(), - value: Value::Object(done), - }, - ]) - } - - fn cache_filter_enabled(&self) -> bool { - self.filter_x_search || !self.injected_tool_types.is_empty() - } - - fn is_internal_cache_call(&self, item: &Value) -> bool { - let Some(item) = item.as_object() else { - return false; - }; - let kind = string_field(item, "type").trim(); - if kind == "web_search_call" { - return self.injected_tool_types.contains("web_search"); - } - if !matches!(kind, "custom_tool_call" | "function_call") { - return false; - } - if string_field(item, "call_id").trim().starts_with("xs_call") { - return true; - } - let name = string_field(item, "name").trim(); - matches!( - name, - "x_user_search" | "x_semantic_search" | "x_keyword_search" | "x_thread_fetch" - ) && string_field(item, "namespace").trim().is_empty() - && !self.client_declared_tools.contains(name) - } - - fn record_dropped_cache_item(&mut self, payload: &Value) { - if let Some(index) = payload.get("output_index").and_then(Value::as_u64) { - self.dropped_output_indexes.insert(index); - } - let Some(item) = payload.get("item").and_then(Value::as_object) else { - return; - }; - for field in ["id", "call_id"] { - if let Some(value) = item - .get(field) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - { - self.dropped_item_ids.insert(value.to_owned()); - } - } - } - - fn references_dropped_cache_item(&self, payload: &Value) -> bool { - if payload - .get("output_index") - .and_then(Value::as_u64) - .is_some_and(|index| self.dropped_output_indexes.contains(&index)) - { - return true; - } - ["item_id", "call_id"].into_iter().any(|field| { - payload - .get(field) - .and_then(Value::as_str) - .map(str::trim) - .is_some_and(|value| self.dropped_item_ids.contains(value)) - }) - } - - fn compact_cache_output_index(&self, payload: &mut Value) { - let Some(index) = payload.get("output_index").and_then(Value::as_u64) else { - return; - }; - let removed_before = self - .dropped_output_indexes - .iter() - .filter(|dropped| **dropped < index) - .count() as u64; - if removed_before > 0 - && let Some(payload) = payload.as_object_mut() - { - payload.insert( - "output_index".to_owned(), - Value::from(index.saturating_sub(removed_before)), - ); - } - } - - fn filter_cache_envelope(&self, payload: &mut Value) -> Result<(), GrokRequestEncodeError> { - let Some(payload) = payload.as_object_mut() else { - return Ok(()); - }; - self.filter_cache_output(payload)?; - self.filter_injected_cache_tools(payload)?; - if let Some(response) = payload.get_mut("response").and_then(Value::as_object_mut) { - self.filter_cache_output(response)?; - self.filter_injected_cache_tools(response)?; - } - Ok(()) - } - - fn filter_cache_output( - &self, - envelope: &mut Map, - ) -> Result<(), GrokRequestEncodeError> { - let Some(output) = envelope.get_mut("output") else { - return Ok(()); - }; - let output = output - .as_array_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - output.retain(|item| !self.is_internal_cache_call(item)); - Ok(()) - } - - fn filter_injected_cache_tools( - &self, - envelope: &mut Map, - ) -> Result<(), GrokRequestEncodeError> { - if self.injected_tool_types.is_empty() { - return Ok(()); - } - let Some(tools) = envelope.get_mut("tools") else { - return Ok(()); - }; - let tools = tools - .as_array_mut() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - tools.retain(|tool| { - tool.get("type") - .and_then(Value::as_str) - .is_none_or(|kind| !self.injected_tool_types.contains(kind.trim())) - }); - if tools.is_empty() { - envelope.remove("tools"); - } - Ok(()) - } - - fn restore_visible_tools(&self, response: &mut Map) { - if response.contains_key("tools") { - if self.visible_tools.is_empty() { - response.remove("tools"); - } else { - response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone())); - } - } - } -} - -fn custom_tool_stream_payload( - source: &Map, - kind: &str, - value_key: &str, - value: &str, -) -> Value { - let mut result = Map::from_iter([ - ("type".to_owned(), Value::String(kind.to_owned())), - (value_key.to_owned(), Value::String(value.to_owned())), - ]); - for key in ["item_id", "output_index", "sequence_number"] { - if let Some(value) = source.get(key) { - result.insert(key.to_owned(), value.clone()); - } - } - Value::Object(result) -} - -fn decode_custom_tool_input(identity: &ToolIdentity, arguments: &str) -> String { - serde_json::from_str::(arguments) - .ok() - .and_then(|value| { - value - .get(identity.custom_argument_field()) - .and_then(Value::as_str) - .map(ToOwned::to_owned) - }) - .unwrap_or_else(|| arguments.to_owned()) -} - -fn decode_tool_search_arguments(value: Option<&Value>) -> Value { - let Some(text) = value.and_then(Value::as_str) else { - return value.cloned().unwrap_or(Value::Object(Map::new())); - }; - if text.trim().is_empty() { - return Value::Object(Map::new()); - } - serde_json::from_str(text) - .unwrap_or_else(|_| json_object([("input", Value::String(text.to_owned()))])) -} - -fn decode_apply_patch_arguments( - value: Option<&Value>, -) -> Result, GrokRequestEncodeError> { - let text = value - .and_then(Value::as_str) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let wrapper = serde_json::from_str::(text) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; - validate_apply_patch_operation(wrapper.get("operation")) -} - -fn rewrite_legacy_local_shell_call(call: &mut Map) { - let commands = call - .get("action") - .and_then(Value::as_object) - .and_then(|action| action.get("commands")) - .and_then(Value::as_array) - .map(|commands| { - commands - .iter() - .filter_map(Value::as_str) - .collect::>() - .join("\n") - }) - .unwrap_or_default(); - call.insert( - "type".to_owned(), - Value::String("local_shell_call".to_owned()), - ); - call.insert( - "action".to_owned(), - json_object([ - ("type", Value::String("exec".to_owned())), - ("command", Value::String(commands)), - ]), - ); - call.remove("max_output_length"); -} - -impl fmt::Debug for GrokResponseTransform { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("GrokResponseTransform") - .field("alias_count", &self.aliases.len()) - .field("function_schema_count", &self.function_schemas.len()) - .field("visible_tool_count", &self.visible_tools.len()) - .field("legacy_local_shell", &self.legacy_local_shell) - .field("filter_x_search", &self.filter_x_search) - .field("injected_tool_types", &self.injected_tool_types) - .finish() - } -} - -struct ToolNormalizer { - response: GrokResponseTransform, - identity_aliases: BTreeMap, - deferred_surfaces: Vec, - client_search_tool: Option>, - server_search_eager: bool, - native_shell: bool, -} - -struct NormalizedInputItems { - items: Vec, - loaded_tools: Vec, - visible_tools: Vec, -} - -struct NormalizedToolSearchOutput { - history: Map, - loaded_tools: Vec, - visible_tools: Vec, -} - -impl ToolNormalizer { - fn new() -> Self { - Self { - response: GrokResponseTransform::default(), - identity_aliases: BTreeMap::new(), - deferred_surfaces: Vec::new(), - client_search_tool: None, - server_search_eager: false, - native_shell: false, - } - } - - fn for_replay(response: GrokResponseTransform) -> Self { - let identity_aliases = response - .aliases - .iter() - .map(|(alias, identity)| (identity.clone(), alias.clone())) - .collect(); - Self { - response, - identity_aliases, - deferred_surfaces: Vec::new(), - client_search_tool: None, - server_search_eager: false, - native_shell: false, - } - } - - fn normalize( - mut self, - payload: &mut Map, - ) -> Result { - let (tools, had_tools) = - optional_array(payload.get("tools")).map_err(|error| error.at_field("tools"))?; - if had_tools { - self.response.visible_tools.clone_from(&tools); - } - let client_search = inspect_tool_search(&tools).map_err(|error| error.at_field("tools"))?; - self.normalize_client_search_parallel(payload, client_search) - .map_err(|error| error.at_field("parallel_tool_calls"))?; - - let mut normalized_tools = Vec::with_capacity(tools.len()); - for raw_tool in &tools { - normalized_tools.extend( - self.normalize_tool(raw_tool, "", client_search, false) - .map_err(|error| error.at_field("tools"))?, - ); - } - - if let Some(Value::Array(items)) = payload.get("input") { - let normalized = self - .normalize_input_items(items) - .map_err(|error| error.at_field("input"))?; - normalized_tools.extend(normalized.loaded_tools); - self.response.visible_tools.extend(normalized.visible_tools); - payload.insert("input".to_owned(), Value::Array(normalized.items)); - } else if payload - .get("input") - .is_some_and(|input| !input.is_null() && !input.is_string()) - { - return Err(GrokRequestEncodeError::InvalidRequestField { field: "input" }); - } - - if self.client_search_tool.is_some() { - normalized_tools.push(Value::Object( - self.build_client_search_function() - .map_err(|error| error.at_field("tools"))?, - )); - } - normalized_tools = dedupe_normalized_tools(normalized_tools); - if normalized_tools.is_empty() { - if had_tools { - payload.remove("tools"); - payload.remove("parallel_tool_calls"); - } - } else { - payload.insert("tools".to_owned(), Value::Array(normalized_tools.clone())); - } - self.normalize_tool_choice(payload, &normalized_tools) - .map_err(|error| error.at_field("tool_choice"))?; - Ok(self.response) - } - - fn alias(&mut self, identity: ToolIdentity) -> String { - if let Some(alias) = self.identity_aliases.get(&identity) { - return alias.clone(); - } - let base = match identity.kind { - ToolKind::ToolSearch => "xai_proxy_tool_search".to_owned(), - ToolKind::ApplyPatch => "xai_proxy_apply_patch".to_owned(), - ToolKind::Function | ToolKind::Custom if !identity.namespace.is_empty() => { - format!("{}__{}", identity.namespace, identity.name) - } - ToolKind::Function | ToolKind::Custom => identity.name.clone(), - }; - let key = format!( - "{}\0{}\0{}", - identity.kind as u8, identity.namespace, identity.name - ); - let mut alias = truncate_tool_alias(&base, &key); - if self - .response - .aliases - .get(&alias) - .is_some_and(|existing| existing != &identity) - { - alias = hashed_tool_alias(&base, &key); - } - self.response - .aliases - .insert(alias.clone(), identity.clone()); - self.identity_aliases.insert(identity, alias.clone()); - alias - } - - fn normalize_client_search_parallel( - &self, - payload: &mut Map, - client_search: bool, - ) -> Result<(), GrokRequestEncodeError> { - if !client_search { - return Ok(()); - } - match payload.get("parallel_tool_calls") { - None | Some(Value::Null) => { - payload.insert("parallel_tool_calls".to_owned(), Value::Bool(false)); - } - Some(Value::Bool(true)) => { - payload.insert("parallel_tool_calls".to_owned(), Value::Bool(false)); - } - Some(Value::Bool(false)) => {} - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - Ok(()) - } - - fn normalize_tool( - &mut self, - raw: &Value, - namespace: &str, - client_search: bool, - force: bool, - ) -> Result, GrokRequestEncodeError> { - let tool = raw - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let kind = string_field(tool, "type").trim(); - match kind { - "function" => self.normalize_function_tool(tool, namespace, client_search, force), - "namespace" => self.normalize_namespace_tool(tool, client_search, force), - "tool_search" => self.normalize_tool_search(tool, force), - "custom" => self.normalize_custom_tool(tool, namespace), - "web_search" - | "web_search_preview" - | "web_search_preview_2025_03_11" - | "web_search_2025_08_26" => self.normalize_web_search_tool(tool, kind), - "mcp" => self.normalize_mcp_tool(tool, client_search, force), - "shell" => self.normalize_shell_tool(tool), - "local_shell" => self.normalize_legacy_local_shell_tool(tool), - "apply_patch" => self.normalize_apply_patch_tool(tool), - "x_search" | "collections_search" | "file_search" | "code_execution" - | "code_interpreter" => Ok(vec![Value::Object(without_defer_loading(tool))]), - // 对未获 xAI 接受的工具做静默过滤，并在过滤后同步收敛 tool_choice； - // 不要把客户端可选扩展升级成整单 400。 - "" | "computer_use_preview" | "image_generation" => Ok(Vec::new()), - _ => Ok(Vec::new()), - } - } - - fn normalize_function_tool( - &mut self, - tool: &Map, - namespace: &str, - client_search: bool, - force: bool, - ) -> Result, GrokRequestEncodeError> { - let name = string_field(tool, "name").trim(); - if name.is_empty() { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let deferred = tool - .get("defer_loading") - .and_then(Value::as_bool) - .unwrap_or(false); - if deferred && client_search && !force { - if namespace.is_empty() { - self.deferred_surfaces.push(describe_deferred_tool( - name, - string_field(tool, "description"), - )); - } - return Ok(Vec::new()); - } - let mut converted = without_defer_loading(tool); - let function_schema = tool - .get("parameters") - .filter(|schema| schema_contains_integer(schema, 0)) - .cloned(); - match converted.get("parameters").cloned() { - None | Some(Value::Null) => { - converted.insert( - "parameters".to_owned(), - json_object([ - ("type", Value::String("object".to_owned())), - ("properties", Value::Object(Map::new())), - ]), - ); - } - Some(parameters) => { - if let Some(normalized) = normalize_function_parameters_root(¶meters)? { - converted.insert("parameters".to_owned(), normalized); - } - } - } - let alias = self.alias(ToolIdentity::new(ToolKind::Function, namespace, name)); - if let Some(schema) = function_schema { - self.response.function_schemas.insert(alias.clone(), schema); - } - converted.insert("name".to_owned(), Value::String(alias)); - Ok(vec![Value::Object(converted)]) - } - - fn normalize_namespace_tool( - &mut self, - tool: &Map, - client_search: bool, - force: bool, - ) -> Result, GrokRequestEncodeError> { - let name = string_field(tool, "name").trim(); - let children = tool - .get("tools") - .and_then(Value::as_array) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - if name.is_empty() { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - if client_search && !force && namespace_has_deferred_functions(children) { - self.deferred_surfaces.push(describe_deferred_tool( - name, - string_field(tool, "description"), - )); - } - let mut converted = Vec::new(); - for child in children { - if child.pointer("/type").and_then(Value::as_str) != Some("function") { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - converted.extend(self.normalize_tool(child, name, client_search, force)?); - } - Ok(converted) - } - - fn normalize_tool_search( - &mut self, - tool: &Map, - force: bool, - ) -> Result, GrokRequestEncodeError> { - if force { - return Ok(Vec::new()); - } - match string_field(tool, "execution") - .trim() - .to_ascii_lowercase() - .as_str() - { - "" | "server" => { - self.server_search_eager = true; - } - "client" => { - self.client_search_tool = Some(tool.clone()); - } - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - Ok(Vec::new()) - } - - fn normalize_custom_tool( - &mut self, - tool: &Map, - namespace: &str, - ) -> Result, GrokRequestEncodeError> { - let name = string_field(tool, "name").trim(); - if name.is_empty() || tool.get("format").is_some_and(|value| !value.is_object()) { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let identity = ToolIdentity::new(ToolKind::Custom, namespace, name); - let argument_field = identity.custom_argument_field(); - let description = if identity.is_root_custom_apply_patch() { - "The apply_patch tool edits files using Codex patch format. Provide the complete raw patch text in the patch string field.".to_owned() - } else { - let mut description = string_field(tool, "description").trim().to_owned(); - if !description.is_empty() { - description.push('\n'); - } - description.push_str("Provide the custom tool input in the input string field."); - description - }; - let alias = self.alias(identity); - Ok(vec![json_object([ - ("type", Value::String("function".to_owned())), - ("name", Value::String(alias)), - ("description", Value::String(description)), - ( - "parameters", - json_object([ - ("type", Value::String("object".to_owned())), - ( - "properties", - json_object([( - argument_field, - json_object([("type", Value::String("string".to_owned()))]), - )]), - ), - ( - "required", - Value::Array(vec![Value::String(argument_field.to_owned())]), - ), - ("additionalProperties", Value::Bool(false)), - ]), - ), - ])]) - } - - fn build_client_search_function( - &mut self, - ) -> Result, GrokRequestEncodeError> { - let tool = self - .client_search_tool - .as_ref() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let mut description = string_field(tool, "description").trim().to_owned(); - if description.is_empty() { - description.push_str("Search for tools needed to continue the task."); - } - if !self.deferred_surfaces.is_empty() { - description.push_str("\nDeferred tool surfaces available to search:\n- "); - description.push_str(&self.deferred_surfaces.join("\n- ")); - } - description.truncate(description.len().min(MAX_TOOL_SEARCH_DESCRIPTION_BYTES)); - let parameters = match tool.get("parameters") { - None => json_object([ - ("type", Value::String("object".to_owned())), - ("properties", Value::Object(Map::new())), - ("additionalProperties", Value::Bool(true)), - ]), - Some(Value::Object(parameters)) => Value::Object(parameters.clone()), - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("function".to_owned())), - ("name".to_owned(), Value::String(alias)), - ("description".to_owned(), Value::String(description)), - ("parameters".to_owned(), parameters), - ])) - } -} - -fn optional_array(value: Option<&Value>) -> Result<(Vec, bool), GrokRequestEncodeError> { - match value { - None | Some(Value::Null) => Ok((Vec::new(), false)), - Some(Value::Array(values)) => Ok((values.clone(), true)), - Some(_) => Err(GrokRequestEncodeError::InvalidRequestNormalization), - } -} - -/// Grok Build 要求 function parameters 的根节点必为 object；Codex 会为可选参数 -/// 生成 `object | null`，这里只移除根节点 nullability，嵌套字段保持原样。 -fn normalize_function_parameters_root( - value: &Value, -) -> Result, GrokRequestEncodeError> { - let Some(schema) = value.as_object() else { - return Ok(None); - }; - let mut normalized = schema.clone(); - let mut changed = false; - - if let Some(types) = normalized.get("type").and_then(Value::as_array).cloned() { - let removed_null = types.iter().any(|value| value.as_str() == Some("null")); - if removed_null { - let remaining = types - .into_iter() - .filter(|value| value.as_str() != Some("null")) - .collect::>(); - if remaining.len() != 1 || remaining[0].as_str() != Some("object") { - return Err(invalid_function_parameters_root()); - } - normalized.insert("type".to_owned(), Value::String("object".to_owned())); - changed = true; - } - } - - for keyword in ["anyOf", "oneOf"] { - let Some(branches) = normalized.get(keyword).and_then(Value::as_array).cloned() else { - continue; - }; - let removed_null = branches.iter().any(is_null_only_schema); - if !removed_null { - continue; - } - let remaining = branches - .into_iter() - .filter(|branch| !is_null_only_schema(branch)) - .collect::>(); - if remaining.is_empty() - || remaining.iter().any(|branch| { - branch.as_object().is_none_or(|branch| { - !is_object_root_schema(branch, &normalized, &mut BTreeSet::new()) - }) - }) - { - return Err(invalid_function_parameters_root()); - } - if remaining.len() == 1 && normalized.len() == 1 { - normalized = remaining[0] - .as_object() - .cloned() - .ok_or_else(invalid_function_parameters_root)?; - normalized.insert("type".to_owned(), Value::String("object".to_owned())); - } else { - normalized.insert(keyword.to_owned(), Value::Array(remaining)); - normalized.insert("type".to_owned(), Value::String("object".to_owned())); - } - changed = true; - } - - Ok(changed.then_some(Value::Object(normalized))) -} - -fn is_null_only_schema(value: &Value) -> bool { - let Some(schema) = value.as_object() else { - return false; - }; - match schema.get("type") { - Some(Value::String(kind)) => kind == "null", - Some(Value::Array(types)) if !types.is_empty() => { - types.iter().all(|kind| kind.as_str() == Some("null")) - } - _ => false, - } -} - -fn is_object_root_schema( - schema: &Map, - root: &Map, - visited: &mut BTreeSet, -) -> bool { - match schema.get("type") { - Some(Value::String(kind)) => return kind == "object", - Some(Value::Array(types)) if !types.is_empty() => { - return types.iter().all(|kind| kind.as_str() == Some("object")); - } - Some(_) => return false, - None => {} - } - if schema.contains_key("properties") { - return true; - } - let Some(reference) = schema.get("$ref").and_then(Value::as_str) else { - return false; - }; - if !visited.insert(reference.to_owned()) { - return false; - } - resolve_local_schema_ref(root, reference) - .is_some_and(|resolved| is_object_root_schema(resolved, root, visited)) -} - -fn resolve_local_schema_ref<'a>( - root: &'a Map, - reference: &str, -) -> Option<&'a Map> { - if reference == "#" { - return Some(root); - } - let path = reference.strip_prefix("#/")?; - let mut segments = path.split('/'); - let first = decode_json_pointer_segment(segments.next()?); - let mut current = root.get(&first)?; - for segment in segments { - let segment = decode_json_pointer_segment(segment); - current = current.as_object()?.get(&segment)?; - } - current.as_object() -} - -fn decode_json_pointer_segment(segment: &str) -> String { - segment.replace("~1", "/").replace("~0", "~") -} - -const fn invalid_function_parameters_root() -> GrokRequestEncodeError { - GrokRequestEncodeError::InvalidRequestField { - field: "tools[].parameters", - } -} - -fn inspect_tool_search(tools: &[Value]) -> Result { - let mut client_search = false; - let mut server_search = false; - for raw_tool in tools { - let tool = raw_tool - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - if string_field(tool, "type") != "tool_search" { - continue; - } - match string_field(tool, "execution") - .trim() - .to_ascii_lowercase() - .as_str() - { - "" | "server" if !client_search => server_search = true, - "client" if !client_search && !server_search => client_search = true, - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - } - Ok(client_search) -} - -fn string_field<'a>(value: &'a Map, key: &str) -> &'a str { - value.get(key).and_then(Value::as_str).unwrap_or_default() -} - -fn json_object(entries: [(&str, Value); N]) -> Value { - Value::Object(Map::from_iter( - entries - .into_iter() - .map(|(key, value)| (key.to_owned(), value)), - )) -} - -fn without_defer_loading(tool: &Map) -> Map { - let mut converted = tool.clone(); - converted.remove("defer_loading"); - converted -} - -fn namespace_has_deferred_functions(children: &[Value]) -> bool { - children.iter().any(|child| { - child.pointer("/type").and_then(Value::as_str) == Some("function") - && child - .pointer("/defer_loading") - .and_then(Value::as_bool) - .unwrap_or(false) - }) -} - -fn describe_deferred_tool(name: &str, description: &str) -> String { - let description = description.trim(); - if description.is_empty() { - return name.to_owned(); - } - let description = description.chars().take(240).collect::(); - format!("{name}: {description}") -} - -fn truncate_tool_alias(base: &str, key: &str) -> String { - if base.len() <= MAX_BUILD_TOOL_ALIAS_LENGTH { - base.to_owned() - } else { - hashed_tool_alias(base, key) - } -} - -fn hashed_tool_alias(base: &str, key: &str) -> String { - let suffix = format!("__{}", short_tool_hash(key)); - let limit = MAX_BUILD_TOOL_ALIAS_LENGTH.saturating_sub(suffix.len()); - let mut end = limit.min(base.len()); - while !base.is_char_boundary(end) { - end = end.saturating_sub(1); - } - format!("{}{suffix}", &base[..end]) -} - -fn short_tool_hash(value: &str) -> String { - Sha256::digest(value) - .iter() - .take(5) - .map(|byte| format!("{byte:02x}")) - .collect::()[..9] - .to_owned() -} - -fn dedupe_normalized_tools(tools: Vec) -> Vec { - let mut result = Vec::with_capacity(tools.len()); - let mut seen = BTreeSet::new(); - for tool in tools { - if normalized_tool_dedupe_key(&tool).is_none_or(|key| seen.insert(key)) { - result.push(tool); - } - } - result -} - -fn normalized_tool_dedupe_key(tool: &Value) -> Option { - let object = tool.as_object()?; - let kind = string_field(object, "type").trim(); - if !kind.is_empty() { - if let Some(name) = object - .get("name") - .and_then(Value::as_str) - .map(str::trim) - .filter(|name| !name.is_empty()) - { - return Some(format!("type:{kind}\0name:{name}")); - } - if kind == "mcp" - && let Some(label) = object - .get("server_label") - .and_then(Value::as_str) - .map(str::trim) - .filter(|label| !label.is_empty()) - { - return Some(format!("type:mcp\0server_label:{label}")); - } - } - serde_json::to_string(tool) - .ok() - .map(|encoded| format!("json:{encoded}")) -} - -impl ToolNormalizer { - fn normalize_web_search_tool( - &mut self, - tool: &Map, - kind: &str, - ) -> Result, GrokRequestEncodeError> { - let filters = normalize_web_search_filters(tool)?; - if let Some(content_types) = tool.get("search_content_types") { - content_types - .as_array() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - } - if kind == "web_search" && tool.len() == 1 { - return Ok(vec![Value::Object(tool.clone())]); - } - let mut converted = - Map::from_iter([("type".to_owned(), Value::String("web_search".to_owned()))]); - if let Some(filters) = filters { - converted.insert("filters".to_owned(), Value::Object(filters)); - } - Ok(vec![Value::Object(converted)]) - } - - fn normalize_mcp_tool( - &mut self, - tool: &Map, - client_search: bool, - force: bool, - ) -> Result, GrokRequestEncodeError> { - let deferred = tool - .get("defer_loading") - .and_then(Value::as_bool) - .unwrap_or(false); - if deferred && client_search && !force { - let label = ["server_label", "name"] - .into_iter() - .find_map(|field| tool.get(field).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - self.deferred_surfaces.push(describe_deferred_tool( - label, - string_field(tool, "description"), - )); - return Ok(Vec::new()); - } - Ok(vec![Value::Object(without_defer_loading(tool))]) - } - - fn normalize_shell_tool( - &mut self, - tool: &Map, - ) -> Result, GrokRequestEncodeError> { - if self.response.legacy_local_shell { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - self.native_shell = true; - Ok(vec![Value::Object(without_defer_loading(tool))]) - } - - fn normalize_legacy_local_shell_tool( - &mut self, - _tool: &Map, - ) -> Result, GrokRequestEncodeError> { - if self.native_shell || self.response.legacy_local_shell { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - self.response.legacy_local_shell = true; - Ok(vec![json_object([ - ("type", Value::String("shell".to_owned())), - ( - "environment", - json_object([("type", Value::String("local".to_owned()))]), - ), - ])]) - } - - fn normalize_apply_patch_tool( - &mut self, - _tool: &Map, - ) -> Result, GrokRequestEncodeError> { - let alias = self.alias(ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")); - let operation = json_object([ - ("type", Value::String("object".to_owned())), - ( - "properties", - json_object([ - ( - "type", - json_object([ - ("type", Value::String("string".to_owned())), - ( - "enum", - Value::Array( - ["create_file", "update_file", "delete_file"] - .into_iter() - .map(|value| Value::String(value.to_owned())) - .collect(), - ), - ), - ]), - ), - ( - "path", - json_object([ - ("type", Value::String("string".to_owned())), - ("minLength", Value::from(1)), - ]), - ), - ( - "diff", - json_object([("type", Value::String("string".to_owned()))]), - ), - ]), - ), - ( - "required", - Value::Array( - ["type", "path"] - .into_iter() - .map(|value| Value::String(value.to_owned())) - .collect(), - ), - ), - ("additionalProperties", Value::Bool(false)), - ]); - Ok(vec![json_object([ - ("type", Value::String("function".to_owned())), - ("name", Value::String(alias)), - ( - "description", - Value::String( - "Create, update, or delete one file using a structured V4A patch operation. create_file and update_file require path and diff; delete_file requires path." - .to_owned(), - ), - ), - ( - "parameters", - json_object([ - ("type", Value::String("object".to_owned())), - ("properties", json_object([("operation", operation)])), - ( - "required", - Value::Array(vec![Value::String("operation".to_owned())]), - ), - ("additionalProperties", Value::Bool(false)), - ]), - ), - ("strict", Value::Bool(true)), - ])]) - } - - fn normalize_tool_choice( - &mut self, - payload: &mut Map, - normalized_tools: &[Value], - ) -> Result<(), GrokRequestEncodeError> { - let Some(choice) = payload.get("tool_choice").cloned() else { - return Ok(()); - }; - if choice.is_null() { - return Ok(()); - } - if normalized_tools.is_empty() { - payload.remove("tool_choice"); - return Ok(()); - } - let Some(mut object) = choice.as_object().cloned() else { - return choice - .as_str() - .filter(|value| matches!(*value, "none" | "auto" | "required")) - .map(|_| ()) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization); - }; - let kind = string_field(&object, "type").to_owned(); - match kind.as_str() { - "tool_search" => { - if self.client_search_tool.is_none() { - if self.server_search_eager { - payload.insert("tool_choice".to_owned(), Value::String("auto".to_owned())); - return Ok(()); - } - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); - payload.insert( - "tool_choice".to_owned(), - json_object([ - ("type", Value::String("function".to_owned())), - ("name", Value::String(alias)), - ]), - ); - } - "custom" => { - let identity = ToolIdentity::new( - ToolKind::Custom, - string_field(&object, "namespace").trim(), - string_field(&object, "name").trim(), - ); - let alias = self - .identity_aliases - .get(&identity) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - object.insert("type".to_owned(), Value::String("function".to_owned())); - object.insert("name".to_owned(), Value::String(alias)); - object.remove("namespace"); - payload.insert("tool_choice".to_owned(), Value::Object(object)); - } - "apply_patch" => { - let identity = ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch"); - let alias = self - .identity_aliases - .get(&identity) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - payload.insert( - "tool_choice".to_owned(), - json_object([ - ("type", Value::String("function".to_owned())), - ("name", Value::String(alias)), - ]), - ); - } - "allowed_tools" => { - payload.remove("tool_choice"); - } - "function" => { - self.normalize_function_tool_choice(payload, object, normalized_tools)?; - } - _ => { - if let Some(hosted_kind) = normalize_hosted_tool_choice_kind(&kind) { - let matching = tools_of_type(normalized_tools, hosted_kind); - if matching.is_empty() { - payload.remove("tool_choice"); - return Ok(()); - } - object.insert("type".to_owned(), Value::String(hosted_kind.to_owned())); - payload.insert("tool_choice".to_owned(), Value::Object(object)); - } else { - payload.remove("tool_choice"); - } - } - } - Ok(()) - } - - fn normalize_function_tool_choice( - &self, - payload: &mut Map, - mut object: Map, - normalized_tools: &[Value], - ) -> Result<(), GrokRequestEncodeError> { - if let Some(Value::Object(function)) = object.get_mut("function") { - rewrite_namespace_choice(function, &self.identity_aliases)?; - let name = string_field(function, "name").trim(); - if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { - payload.remove("tool_choice"); - return Ok(()); - } - payload.insert("tool_choice".to_owned(), Value::Object(object)); - return Ok(()); - } - rewrite_namespace_choice(&mut object, &self.identity_aliases)?; - let name = string_field(&object, "name").trim(); - if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { - payload.remove("tool_choice"); - return Ok(()); - } - payload.insert("tool_choice".to_owned(), Value::Object(object)); - Ok(()) - } -} - -fn normalize_web_search_filters( - tool: &Map, -) -> Result>, GrokRequestEncodeError> { - let nested = match tool.get("filters") { - None | Some(Value::Null) => None, - Some(Value::Object(filters)) => filters - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?, - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let top_level = tool - .get("allowed_domains") - .map(normalize_allowed_domains) - .transpose()?; - if nested.is_some() && top_level.is_some() && nested != top_level { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - Ok(nested - .or(top_level) - .map(|domains| Map::from_iter([("allowed_domains".to_owned(), Value::Array(domains))]))) -} - -fn normalize_allowed_domains(value: &Value) -> Result, GrokRequestEncodeError> { - let domains = value - .as_array() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - if domains - .iter() - .any(|domain| domain.as_str().map(str::trim).is_none_or(str::is_empty)) - { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - Ok(domains.clone()) -} - -fn normalize_hosted_tool_choice_kind(kind: &str) -> Option<&'static str> { - match kind { - "web_search" - | "web_search_preview" - | "web_search_preview_2025_03_11" - | "web_search_2025_08_26" => Some("web_search"), - "x_search" => Some("x_search"), - "image_generation" => Some("image_generation"), - "collections_search" => Some("collections_search"), - "file_search" => Some("file_search"), - "code_execution" => Some("code_execution"), - "code_interpreter" => Some("code_interpreter"), - "mcp" => Some("mcp"), - "shell" | "local_shell" => Some("shell"), - _ => None, - } -} - -fn has_tool_type(tools: &[Value], kind: &str) -> bool { - tools - .iter() - .any(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind)) -} - -fn has_named_tool(tools: &[Value], kind: &str, name: &str) -> bool { - tools.iter().any(|tool| { - tool.pointer("/type").and_then(Value::as_str) == Some(kind) - && tool.pointer("/name").and_then(Value::as_str) == Some(name) - }) -} - -fn tools_of_type(tools: &[Value], kind: &str) -> Vec { - tools - .iter() - .filter(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind)) - .cloned() - .collect() -} - -fn rewrite_namespace_choice( - object: &mut Map, - aliases: &BTreeMap, -) -> Result<(), GrokRequestEncodeError> { - let name = string_field(object, "name").trim(); - let namespace = string_field(object, "namespace").trim(); - if name.is_empty() || namespace.is_empty() { - return Ok(()); - } - let identity = ToolIdentity::new(ToolKind::Function, namespace, name); - let alias = aliases - .get(&identity) - .cloned() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - object.insert("name".to_owned(), Value::String(alias)); - object.remove("namespace"); - Ok(()) -} - -impl ToolNormalizer { - fn normalize_input_items( - &mut self, - items: &[Value], - ) -> Result { - let mut rewritten = Vec::with_capacity(items.len()); - let mut loaded_tools = Vec::new(); - let mut visible_tools = Vec::new(); - for raw_item in items { - let Some(item) = raw_item.as_object() else { - rewritten.push(raw_item.clone()); - continue; - }; - let mut item_type = string_field(item, "type").trim(); - if item_type.is_empty() && !string_field(item, "role").trim().is_empty() { - item_type = "message"; - } - match item_type { - "message" => rewritten.push(Value::Object(self.normalize_message_input(item)?)), - "function_call" => { - rewritten.push(Value::Object(self.normalize_function_call_input(item)?)); - } - "function_call_output" => rewritten.push(Value::Object( - self.normalize_function_call_output_input(item, true)?, - )), - "reasoning" => rewritten.push(Value::Object(sanitize_reasoning_input(item))), - "file_search_call" - | "web_search_call" - | "image_generation_call" - | "code_interpreter_call" - | "shell_call" - | "mcp_list_tools" - | "mcp_approval_request" - | "mcp_approval_response" - | "mcp_call" => { - rewritten.push(Value::Object(sanitize_native_history_input( - item, item_type, - ))); - } - "compaction" | "compaction_summary" => { - rewritten.extend(normalize_compaction_input(item)?); - } - "tool_search_call" => { - rewritten.push(Value::Object(self.normalize_tool_search_call(item)?)); - } - "tool_search_output" => { - let normalized = self.normalize_tool_search_output(item)?; - rewritten.push(Value::Object(normalized.history)); - loaded_tools.extend(normalized.loaded_tools); - visible_tools.extend(normalized.visible_tools); - } - "custom_tool_call" => { - rewritten.push(Value::Object(self.normalize_custom_tool_call_input(item)?)); - } - "custom_tool_call_output" => rewritten.push(Value::Object( - self.normalize_function_call_output_input(item, false)?, - )), - "apply_patch_call" => { - rewritten.push(Value::Object(self.normalize_apply_patch_call_input(item)?)); - } - "apply_patch_call_output" => { - rewritten.push(Value::Object(normalize_apply_patch_output_input(item)?)); - } - "agent_message" => rewritten.push(normalize_agent_message_input(item)), - "local_shell_call" => { - rewritten.push(Value::Object(normalize_legacy_local_shell_call_input( - item, - )?)); - } - "local_shell_call_output" => rewritten.push(Value::Object( - normalize_legacy_local_shell_output_input(item)?, - )), - "shell_call_output" => { - rewritten.push(Value::Object(normalize_shell_call_output_input(item)?)); - } - "mcp_tool_call_output" => rewritten.push(normalize_mcp_output_input(item)?), - "compaction_trigger" => { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - "additional_tools" => { - let (tools, visible) = self.normalize_additional_tools_input(item)?; - loaded_tools.extend(tools); - visible_tools.extend(visible); - } - "" => rewritten.push(raw_item.clone()), - unsupported => { - rewritten.push(unsupported_input_history_boundary(item, unsupported)) - } - } - } - Ok(NormalizedInputItems { - items: rewritten, - loaded_tools, - visible_tools, - }) - } - - fn normalize_message_input( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let role = match string_field(item, "role").trim() { - "" => "assistant", - "model" => "assistant", - role => role, - }; - let content = self.normalize_message_content(item.get("content"), role)?; - // 透明代理合法形态：只覆盖归一化字段并剥离 Grok 注入键，未知官方 - // 字段原样保留。 - let mut converted = item.clone(); - strip_grok_internal_keys(&mut converted); - converted.insert("type".to_owned(), Value::String("message".to_owned())); - converted.insert("role".to_owned(), Value::String(role.to_owned())); - converted.insert("content".to_owned(), content); - Ok(converted) - } - - fn normalize_message_content( - &mut self, - value: Option<&Value>, - role: &str, - ) -> Result { - if let Some(Value::String(text)) = value { - return Ok(Value::String(text.clone())); - } - let items = value - .and_then(Value::as_array) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - if role == "assistant" { - let texts = items - .iter() - .map(|raw| { - let item = raw.as_object()?; - match string_field(item, "type") { - "text" | "input_text" | "output_text" => { - item.get("text").and_then(Value::as_str) - } - "refusal" => item.get("refusal").and_then(Value::as_str), - _ => None, - } - }) - .collect::>>(); - if let Some(texts) = texts { - return Ok(Value::String(texts.join("\n"))); - } - } - let text_part_type = if role == "assistant" { - "output_text" - } else { - "input_text" - }; - let mut normalized = Vec::with_capacity(items.len()); - for raw in items { - let item = raw - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let converted = match string_field(item, "type") { - "text" | "input_text" | "output_text" => json_object([ - ("type", Value::String(text_part_type.to_owned())), - ("text", Value::String(string_field(item, "text").to_owned())), - ]), - "refusal" => json_object([ - ("type", Value::String(text_part_type.to_owned())), - ( - "text", - Value::String(string_field(item, "refusal").to_owned()), - ), - ]), - "input_image" => Value::Object(self.normalize_input_image_part(item)?), - "input_file" => Value::Object(normalize_input_file_part(item)), - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - normalized.push(converted); - } - Ok(Value::Array(normalized)) - } - - fn normalize_input_image_part( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let detail = match item.get("detail") { - None | Some(Value::Null) => "auto", - Some(Value::String(detail)) if detail.trim().is_empty() => "auto", - Some(Value::String(detail)) => detail.trim(), - Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let detail = match detail { - "auto" | "low" | "high" => detail, - "original" => "high", - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let mut converted = Map::from_iter([ - ("type".to_owned(), Value::String("input_image".to_owned())), - ("detail".to_owned(), Value::String(detail.to_owned())), - ]); - if let Some(value) = item.get("image_url").or_else(|| item.get("url")) - && !value.is_null() - { - converted.insert("image_url".to_owned(), value.clone()); - } - if let Some(value) = item.get("file_id") - && !value.is_null() - { - converted.insert("file_id".to_owned(), value.clone()); - } - Ok(converted) - } - - fn normalize_function_call_input( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let mut name = required_trimmed_string(item, "name")?.to_owned(); - let call_id = required_trimmed_string(item, "call_id")?; - let arguments = encode_function_arguments(item.get("arguments"))?; - let namespace = string_field(item, "namespace").trim(); - if !namespace.is_empty() { - name = self.alias(ToolIdentity::new(ToolKind::Function, namespace, &name)); - } - let mut converted = item.clone(); - strip_grok_internal_keys(&mut converted); - converted.remove("namespace"); - converted.insert("type".to_owned(), Value::String("function_call".to_owned())); - converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); - converted.insert("name".to_owned(), Value::String(name)); - converted.insert("arguments".to_owned(), Value::String(arguments)); - Ok(converted) - } - - fn normalize_custom_tool_call_input( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let name = required_trimmed_string(item, "name")?; - let input = item - .get("input") - .and_then(Value::as_str) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let call_id = required_trimmed_string(item, "call_id")?; - let namespace = string_field(item, "namespace").trim(); - let identity = ToolIdentity::new(ToolKind::Custom, namespace, name); - let arguments = serde_json::to_string(&json_object([( - identity.custom_argument_field(), - Value::String(input.to_owned()), - )])) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; - let alias = self.alias(identity); - let mut converted = item.clone(); - strip_grok_internal_keys(&mut converted); - converted.remove("namespace"); - converted.remove("input"); - converted.insert("type".to_owned(), Value::String("function_call".to_owned())); - converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); - converted.insert("name".to_owned(), Value::String(alias)); - converted.insert("arguments".to_owned(), Value::String(arguments)); - Ok(converted) - } - - fn normalize_function_call_output_input( - &mut self, - item: &Map, - allow_content_blocks: bool, - ) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let output = match item.get("output") { - Some(Value::Array(blocks)) - if allow_content_blocks && is_function_output_content_array(blocks) => - { - Value::Array(self.normalize_function_output_blocks(blocks)?) - } - output => Value::String(encode_tool_output(output)?), - }; - let mut converted = item.clone(); - strip_grok_internal_keys(&mut converted); - converted.insert( - "type".to_owned(), - Value::String("function_call_output".to_owned()), - ); - converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); - converted.insert("output".to_owned(), output); - Ok(converted) - } - - fn normalize_function_output_blocks( - &mut self, - blocks: &[Value], - ) -> Result, GrokRequestEncodeError> { - blocks - .iter() - .map(|raw| { - let block = raw - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - match string_field(block, "type") { - "input_text" => Ok(json_object([ - ("type", Value::String("input_text".to_owned())), - ( - "text", - Value::String( - block - .get("text") - .and_then(Value::as_str) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? - .to_owned(), - ), - ), - ])), - "input_image" => { - require_content_source(block, &["image_url", "file_id"])?; - self.normalize_input_image_part(block).map(Value::Object) - } - "input_file" => { - require_content_source(block, &["file_data", "file_id", "file_url"])?; - Ok(Value::Object(normalize_input_file_part(block))) - } - _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - }) - .collect() - } - - fn normalize_tool_search_call( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let execution = string_field(item, "execution").trim().to_ascii_lowercase(); - if execution.is_empty() || execution == "server" { - self.server_search_eager = true; - return Ok(boundary_message( - "A server-side tool search occurred here; selected tools are made available directly.", - )); - } - if execution != "client" { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let arguments = encode_function_arguments(item.get("arguments"))?; - let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("function_call".to_owned())), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("name".to_owned(), Value::String(alias)), - ("arguments".to_owned(), Value::String(arguments)), - ])) - } - - fn normalize_tool_search_output( - &mut self, - item: &Map, - ) -> Result { - let execution = string_field(item, "execution").trim().to_ascii_lowercase(); - if !matches!(execution.as_str(), "" | "client" | "server") { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let call_id = required_trimmed_string(item, "call_id")?; - let tools = item - .get("tools") - .and_then(Value::as_array) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let mut normalized = Vec::new(); - for tool in tools { - normalized.extend(self.normalize_tool(tool, "", false, true)?); - } - let message = format!( - "Tool search completed; {} selected tool definitions are now available.", - tools.len() - ); - let history = if execution == "client" { - Map::from_iter([ - ( - "type".to_owned(), - Value::String("function_call_output".to_owned()), - ), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("output".to_owned(), Value::String(message)), - ]) - } else { - self.server_search_eager = true; - boundary_message(&message) - }; - Ok(NormalizedToolSearchOutput { - history, - loaded_tools: normalized, - visible_tools: tools.clone(), - }) - } - - fn normalize_apply_patch_call_input( - &mut self, - item: &Map, - ) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let operation = validate_apply_patch_operation(item.get("operation"))?; - let arguments = - serde_json::to_string(&json_object([("operation", Value::Object(operation))])) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; - let alias = self.alias(ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")); - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("function_call".to_owned())), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("name".to_owned(), Value::String(alias)), - ("arguments".to_owned(), Value::String(arguments)), - ])) - } - - fn normalize_additional_tools_input( - &mut self, - item: &Map, - ) -> Result<(Vec, Vec), GrokRequestEncodeError> { - let tools = item - .get("tools") - .and_then(Value::as_array) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let mut normalized = Vec::new(); - for raw_tool in tools { - normalized.extend(self.normalize_tool(raw_tool, "", false, true)?); - } - Ok((normalized, tools.clone())) - } -} - -fn normalize_input_file_part(item: &Map) -> Map { - let mut converted = - Map::from_iter([("type".to_owned(), Value::String("input_file".to_owned()))]); - for key in ["file_data", "file_id", "filename", "file_url"] { - if let Some(value) = item.get(key) - && !value.is_null() - { - converted.insert(key.to_owned(), value.clone()); - } - } - converted -} - -fn required_trimmed_string<'a>( - item: &'a Map, - field: &str, -) -> Result<&'a str, GrokRequestEncodeError> { - item.get(field) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) -} - -fn encode_function_arguments(value: Option<&Value>) -> Result { - match value { - Some(Value::String(value)) => Ok(value.clone()), - None | Some(Value::Null) => Ok("{}".to_owned()), - Some(value) => serde_json::to_string(value) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization), - } -} - -fn normalize_function_arguments(arguments: &str, schema: &Value) -> Option { - if arguments.trim().is_empty() || !schema.is_object() { - return None; - } - let mut value = serde_json::from_str::(arguments).ok()?; - if !normalize_argument_value(&mut value, schema, schema, 0) { - return None; - } - serde_json::to_string(&value).ok() -} - -fn normalize_argument_value(value: &mut Value, schema: &Value, root: &Value, depth: usize) -> bool { - if depth > 64 { - return false; - } - let Some(schema) = schema.as_object() else { - return false; - }; - let mut changed = false; - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) - && let Some(resolved) = resolve_local_schema_value_ref(root, reference) - { - changed |= normalize_argument_value(value, resolved, root, depth + 1); - } - for keyword in ["allOf", "anyOf", "oneOf"] { - if let Some(branches) = schema.get(keyword).and_then(Value::as_array) { - for branch in branches { - changed |= normalize_argument_value(value, branch, root, depth + 1); - } - } - } - if schema_requires_integer(schema) - && let Value::Number(number) = value - && let Some(normalized) = normalize_integral_number(number) - { - *number = normalized; - return true; - } - match value { - Value::Object(object) => { - let properties = schema.get("properties").and_then(Value::as_object); - let additional = schema - .get("additionalProperties") - .filter(|value| value.is_object()); - for (key, item) in object { - let property = properties - .and_then(|properties| properties.get(key)) - .or(additional); - if let Some(property) = property { - changed |= normalize_argument_value(item, property, root, depth + 1); - } - } - } - Value::Array(items) => { - let prefix = schema.get("prefixItems").and_then(Value::as_array); - let item_schema = schema.get("items").filter(|value| value.is_object()); - for (index, item) in items.iter_mut().enumerate() { - let schema = prefix - .and_then(|prefix| prefix.get(index)) - .filter(|value| value.is_object()) - .or(item_schema); - if let Some(schema) = schema { - changed |= normalize_argument_value(item, schema, root, depth + 1); - } - } - } - _ => {} - } - changed -} - -fn schema_requires_integer(schema: &Map) -> bool { - match schema.get("type") { - Some(Value::String(kind)) => kind == "integer", - Some(Value::Array(kinds)) => { - let mut integer = false; - for kind in kinds.iter().filter_map(Value::as_str) { - if kind == "number" { - return false; - } - integer |= kind == "integer"; - } - integer - } - _ => false, - } -} - -fn normalize_integral_number(number: &serde_json::Number) -> Option { - let raw = number.to_string(); - if raw.len() > MAX_JSON_NUMBER_TEXT_BYTES || !raw.contains(['.', 'e', 'E']) { - return None; - } - let (mantissa, exponent_text) = raw.find(['e', 'E']).map_or((raw.as_str(), ""), |index| { - (&raw[..index], &raw[index + 1..]) - }); - let (negative, mantissa) = mantissa - .strip_prefix('-') - .map_or((false, mantissa), |value| (true, value)); - let (whole, fraction) = mantissa.split_once('.').unwrap_or((mantissa, "")); - let mut digits = format!("{whole}{fraction}") - .trim_start_matches('0') - .to_owned(); - if digits.is_empty() { - return "0".parse().ok().filter(|_| raw != "0"); - } - let exponent = parse_bounded_decimal_exponent(exponent_text)?; - let decimal_shift = exponent - i64::try_from(fraction.len()).ok()?; - if decimal_shift < 0 { - let fractional_digits = usize::try_from(-decimal_shift).ok()?; - if fractional_digits > digits.len() - || !digits[digits.len() - fractional_digits..] - .bytes() - .all(|byte| byte == b'0') - { - return None; - } - digits.truncate(digits.len() - fractional_digits); - let trimmed = digits.trim_start_matches('0'); - if trimmed.is_empty() { - return "0".parse().ok(); - } - digits = trimmed.to_owned(); - } else if decimal_shift > 0 { - let shift = usize::try_from(decimal_shift).ok()?; - if shift - > MAX_EXACT_JSON_INTEGER_TEXT - .len() - .saturating_sub(digits.len()) - { - return None; - } - digits.extend(std::iter::repeat_n('0', shift)); - } - if digits.len() > MAX_EXACT_JSON_INTEGER_TEXT.len() - || (digits.len() == MAX_EXACT_JSON_INTEGER_TEXT.len() - && digits.as_str() > MAX_EXACT_JSON_INTEGER_TEXT) - { - return None; - } - let normalized = if negative { - format!("-{digits}") - } else { - digits - }; - if normalized == raw { - None - } else { - normalized.parse().ok() - } -} - -fn exact_nonnegative_sequence(value: &Value) -> Option { - let number = value.as_number()?; - let sequence = number - .as_u64() - .or_else(|| normalize_integral_number(number)?.as_u64())?; - (sequence <= MAX_EXACT_JSON_INTEGER).then_some(sequence) -} - -fn parse_bounded_decimal_exponent(raw: &str) -> Option { - if raw.is_empty() { - return Some(0); - } - let (sign, digits) = if let Some(value) = raw.strip_prefix('+') { - (1_i64, value) - } else if let Some(value) = raw.strip_prefix('-') { - (-1_i64, value) - } else { - (1_i64, raw) - }; - if digits.is_empty() { - return None; - } - let digits = digits.trim_start_matches('0'); - if digits.is_empty() { - return Some(0); - } - if digits.len() > 9 { - return None; - } - digits.parse::().ok().map(|value| sign * value) -} - -fn resolve_local_schema_value_ref<'a>(root: &'a Value, reference: &str) -> Option<&'a Value> { - if reference == "#" { - return Some(root); - } - let pointer = reference.strip_prefix('#')?; - if !pointer.starts_with('/') { - return None; - } - root.pointer(pointer) -} - -fn schema_contains_integer(schema: &Value, depth: usize) -> bool { - let mut visited = BTreeSet::new(); - schema_contains_reachable_integer(schema, schema, &mut visited, depth) -} - -fn schema_contains_reachable_integer( - schema: &Value, - root: &Value, - visited_refs: &mut BTreeSet, - depth: usize, -) -> bool { - if depth > 64 { - return false; - } - let Some(schema) = schema.as_object() else { - return false; - }; - if schema_requires_integer(schema) { - return true; - } - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) - && visited_refs.insert(reference.to_owned()) - && let Some(resolved) = resolve_local_schema_value_ref(root, reference) - && schema_contains_reachable_integer(resolved, root, visited_refs, depth + 1) - { - return true; - } - for keyword in ["allOf", "anyOf", "oneOf", "prefixItems"] { - if schema - .get(keyword) - .and_then(Value::as_array) - .is_some_and(|branches| { - branches.iter().any(|branch| { - schema_contains_reachable_integer(branch, root, visited_refs, depth + 1) - }) - }) - { - return true; - } - } - for keyword in ["items", "additionalProperties"] { - if schema.get(keyword).is_some_and(|child| { - schema_contains_reachable_integer(child, root, visited_refs, depth + 1) - }) { - return true; - } - } - schema - .get("properties") - .and_then(Value::as_object) - .is_some_and(|properties| { - properties.values().any(|property| { - schema_contains_reachable_integer(property, root, visited_refs, depth + 1) - }) - }) -} - -fn encode_tool_output(value: Option<&Value>) -> Result { - match value { - None | Some(Value::Null) => Ok(String::new()), - Some(Value::String(value)) => Ok(value.clone()), - Some(value) => serde_json::to_string(value) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization), - } -} - -fn is_function_output_content_array(blocks: &[Value]) -> bool { - blocks.iter().any(|raw| { - raw.pointer("/type") - .and_then(Value::as_str) - .is_some_and(|kind| kind.starts_with("input_")) - }) -} - -fn require_content_source( - block: &Map, - fields: &[&str], -) -> Result<(), GrokRequestEncodeError> { - let has_source = fields.iter().any(|field| { - block - .get(*field) - .and_then(Value::as_str) - .is_some_and(|value| !value.trim().is_empty()) - }); - has_source - .then_some(()) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) -} - -fn sanitize_reasoning_input(item: &Map) -> Map { - let mut converted = - copy_non_null_history_fields(item, &["id", "summary", "content", "encrypted_content"]); - // Grok CLI 会在 summary/content 条目上注入 phase 等内部键。 - strip_grok_internal_entry_keys(&mut converted, &["summary", "content"]); - converted.insert("type".to_owned(), Value::String("reasoning".to_owned())); - if has_portable_reasoning_content(&converted) { - converted - } else { - boundary_message( - "A prior model reasoning item was omitted because it has no portable content for Grok Build.", - ) - } -} - -fn normalize_compaction_input( - item: &Map, -) -> Result, GrokRequestEncodeError> { - let encrypted = item - .get("encrypted_content") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()); - let has_structured_compaction_shape = item.contains_key("summary") - || item.contains_key("id") - || item.contains_key("status") - || string_field(item, "type").trim() == "compaction_summary"; - - // 旧版 adapter 曾将明文摘要直接写入 encrypted_content，且不带 - // id/status/summary。仅对这一可识别的历史形态保留明文恢复；新形态严格回放为 - // xAI reasoning 密文 + 可见 conversation summary。 - if !has_structured_compaction_shape { - let Some(summary) = encrypted else { - return Ok(Vec::new()); - }; - let continuation = format!( - "This session is being continued from a previous conversation that ran out of context. \ - The summary below covers the earlier portion of the conversation.\n\n{summary}" - ); - return Ok(vec![Value::Object(compaction_summary_message( - &continuation, - ))]); - } - - let mut converted = Vec::with_capacity(2); - if let Some(encrypted) = encrypted { - converted.push(json!({ - "type": "reasoning", - "summary": [], - "encrypted_content": encrypted, - })); - } - if let Some(summary) = compact_summary_text(item.get("summary")) { - converted.push(Value::Object(compaction_summary_message(&format!( - "\n{summary}\n" - )))); - } - Ok(converted) -} - -fn compact_summary_text(value: Option<&Value>) -> Option { - let text = value? - .as_array()? - .iter() - .filter_map(Value::as_object) - .filter_map(|part| part.get("text").and_then(Value::as_str)) - .map(str::trim) - .filter(|text| !text.is_empty()) - .collect::>() - .join("\n"); - (!text.is_empty()).then_some(text) -} - -fn compaction_summary_message(text: &str) -> Map { - Map::from_iter([ - ("type".to_owned(), Value::String("message".to_owned())), - ("role".to_owned(), Value::String("user".to_owned())), - ( - "content".to_owned(), - Value::Array(vec![json_object([ - ("type", Value::String("input_text".to_owned())), - ("text", Value::String(text.to_owned())), - ])]), - ), - ]) -} - -pub(super) fn strip_invalid_encrypted_reasoning_from_body(body: &mut Map) -> bool { - let Some(input) = body.get_mut("input").and_then(Value::as_array_mut) else { - return false; - }; - let has_encrypted_reasoning = input.iter().any(|item| { - item.as_object().is_some_and(|item| { - string_field(item, "type").trim() == "reasoning" - && item - .get("encrypted_content") - .and_then(Value::as_str) - .is_some_and(|value| !value.trim().is_empty()) - }) - }); - if !has_encrypted_reasoning { - return false; - } - - input.retain_mut(|item| { - let Some(item) = item.as_object_mut() else { - return true; - }; - if string_field(item, "type").trim() != "reasoning" { - return true; - } - item.remove("encrypted_content"); - if item.get("content").is_some_and(Value::is_null) { - item.remove("content"); - } - item.len() > 1 - }); - true -} - -fn has_portable_reasoning_content(item: &Map) -> bool { - item.get("encrypted_content") - .and_then(Value::as_str) - .is_some_and(|value| !value.trim().is_empty()) - || ["summary", "content"].into_iter().any(|field| { - item.get(field) - .and_then(Value::as_array) - .is_some_and(|values| !values.is_empty()) - }) -} - -fn sanitize_native_history_input(item: &Map, item_type: &str) -> Map { - let fields: &[&str] = match item_type { - "file_search_call" => &["id", "queries", "status", "results"], - "web_search_call" => &["action", "id", "status"], - "image_generation_call" => &["id", "result", "status"], - "code_interpreter_call" => &["code", "container_id", "id", "outputs", "status"], - "shell_call" => &["id", "call_id", "action", "status", "environment"], - "mcp_list_tools" => &["id", "server_label", "tools", "error"], - "mcp_approval_request" => &["arguments", "id", "name", "server_label"], - "mcp_approval_response" => &["approval_request_id", "approve", "id", "reason"], - "mcp_call" => &[ - "arguments", - "id", - "name", - "server_label", - "approval_request_id", - "error", - "output", - "status", - ], - _ => &[], - }; - let mut converted = copy_non_null_history_fields(item, fields); - // Grok CLI 会在 shell_call 的 action object 里注入内部键与 null 占位 - // 字段（如 `timeout_ms: null`）；只在这一层剥离，深层内容原样保留。 - if item_type == "shell_call" - && let Some(Value::Object(action)) = converted.get_mut("action") - { - strip_grok_internal_keys(action); - action.retain(|_, value| !value.is_null()); - } - converted.insert("type".to_owned(), Value::String(item_type.to_owned())); - converted -} - -fn copy_non_null_history_fields(item: &Map, fields: &[&str]) -> Map { - fields - .iter() - .filter_map(|field| { - item.get(*field) - .filter(|value| !value.is_null()) - .map(|value| ((*field).to_owned(), value.clone())) - }) - .collect() -} - -fn strip_grok_internal_keys(object: &mut Map) { - for key in GROK_INTERNAL_HISTORY_KEYS { - object.remove(*key); - } -} - -fn strip_grok_internal_entry_keys(converted: &mut Map, fields: &[&str]) { - for field in fields { - let Some(Value::Array(entries)) = converted.get_mut(*field) else { - continue; - }; - for entry in entries { - if let Some(object) = entry.as_object_mut() { - strip_grok_internal_keys(object); - } - } - } -} - -fn unsupported_input_history_boundary(item: &Map, kind: &str) -> Value { - let mut lines = vec![ - "A prior Responses history item was omitted because Grok Build cannot deserialize this Codex item type." - .to_owned(), - format!("Type: {kind}"), - ]; - for key in ["id", "call_id", "name", "status"] { - let value = string_field(item, key).trim(); - if !value.is_empty() { - lines.push(format!("{}: {value}", key.replace('_', " "))); - } - } - Value::Object(boundary_message(&lines.join("\n"))) -} - -fn boundary_message(text: &str) -> Map { - Map::from_iter([ - ("type".to_owned(), Value::String("message".to_owned())), - ("role".to_owned(), Value::String("developer".to_owned())), - ( - "content".to_owned(), - Value::Array(vec![json_object([ - ("type", Value::String("input_text".to_owned())), - ("text", Value::String(text.to_owned())), - ])]), - ), - ]) -} - -fn normalize_agent_message_input(item: &Map) -> Value { - let Some(content) = item.get("content").and_then(text_input_content) else { - return Value::Object(boundary_message( - "An encrypted inter-agent message occurred here but is not portable to the Grok Build account.", - )); - }; - let author = non_empty_or(string_field(item, "author"), "agent"); - let recipient = non_empty_or(string_field(item, "recipient"), "recipient"); - Value::Object(boundary_message(&format!( - "Agent message ({author} -> {recipient}):\n{content}" - ))) -} - -fn text_input_content(value: &Value) -> Option { - match value { - Value::String(text) => Some(text.clone()), - Value::Array(items) => items - .iter() - .map(|item| { - let item = item.as_object()?; - matches!( - string_field(item, "type"), - "input_text" | "output_text" | "text" - ) - .then(|| string_field(item, "text").to_owned()) - }) - .collect::>>() - .map(|parts| parts.join("\n")), - _ => None, - } -} - -fn non_empty_or<'a>(value: &'a str, fallback: &'a str) -> &'a str { - let value = value.trim(); - if value.is_empty() { fallback } else { value } -} - -fn normalize_mcp_output_input(item: &Map) -> Result { - let output = serde_json::to_string(item.get("output").unwrap_or(&Value::Null)) - .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; - let call_id = non_empty_or(string_field(item, "call_id"), "unknown"); - Ok(Value::Object(boundary_message(&format!( - "MCP tool output for call {call_id}: {output}" - )))) -} - -fn validate_apply_patch_operation( - value: Option<&Value>, -) -> Result, GrokRequestEncodeError> { - let operation = value - .and_then(Value::as_object) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let kind = required_trimmed_string(operation, "type")?; - required_trimmed_string(operation, "path")?; - match kind { - "create_file" | "update_file" if operation.get("diff").is_some_and(Value::is_string) => {} - "delete_file" => {} - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - Ok(operation.clone()) -} - -fn normalize_apply_patch_output_input( - item: &Map, -) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let status = match string_field(item, "status").trim() { - "" | "completed" => "completed", - "failed" => "failed", - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let output = encode_tool_output(item.get("output"))?; - let mut message = format!("Apply patch status: {status}"); - if !output.is_empty() { - message.push('\n'); - message.push_str(&output); - } - Ok(Map::from_iter([ - ( - "type".to_owned(), - Value::String("function_call_output".to_owned()), - ), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("output".to_owned(), Value::String(message)), - ])) -} - -fn normalize_legacy_local_shell_call_input( - item: &Map, -) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let action = legacy_shell_action(item.get("action"))?; - let mut converted = Map::from_iter([ - ("type".to_owned(), Value::String("shell_call".to_owned())), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("action".to_owned(), Value::Object(action)), - ]); - for key in ["id", "status", "timeout_ms", "max_output_length"] { - if let Some(value) = item.get(key) { - converted.insert(key.to_owned(), value.clone()); - } - } - Ok(converted) -} - -fn legacy_shell_action( - value: Option<&Value>, -) -> Result, GrokRequestEncodeError> { - let action = value - .and_then(Value::as_object) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - if !matches!(string_field(action, "type").trim(), "" | "exec") { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let command = legacy_shell_command(action)?; - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("exec".to_owned())), - ( - "commands".to_owned(), - Value::Array(vec![Value::String(command)]), - ), - ])) -} - -fn legacy_shell_command(action: &Map) -> Result { - let mut command = match action.get("command") { - Some(Value::String(value)) => value.trim().to_owned(), - Some(Value::Array(parts)) => parts - .iter() - .map(|part| { - part.as_str() - .map(quote_shell_argument) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) - }) - .collect::, _>>()? - .join(" "), - _ => action - .get("commands") - .and_then(Value::as_array) - .map(|commands| { - commands - .iter() - .map(|command| { - command - .as_str() - .map(ToOwned::to_owned) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) - }) - .collect::, _>>() - .map(|commands| commands.join("\n")) - }) - .transpose()? - .unwrap_or_default(), - }; - if command.is_empty() { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - if let Some(environment) = action.get("env").and_then(Value::as_object) - && !environment.is_empty() - { - let assignments = environment - .iter() - .map(|(name, value)| { - if !valid_environment_name(name) { - return Err(GrokRequestEncodeError::InvalidRequestNormalization); - } - let value = value - .as_str() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - Ok(format!("{name}={}", quote_shell_argument(value))) - }) - .collect::, _>>()?; - command = format!("env {} {command}", assignments.join(" ")); - } - if let Some(directory) = action - .get("working_directory") - .and_then(Value::as_str) - .map(str::trim) - .filter(|directory| !directory.is_empty()) - { - command = format!("cd {} && {command}", quote_shell_argument(directory)); - } - Ok(command) -} - -fn normalize_legacy_local_shell_output_input( - item: &Map, -) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let output = match item.get("output") { - Some(Value::Array(output)) => output.clone(), - Some(Value::String(output)) => { - let exit_code = item - .get("exit_code") - .and_then(Value::as_i64) - .unwrap_or_else(|| { - i64::from(string_field(item, "status").eq_ignore_ascii_case("failed")) - }); - vec![shell_output_block(output, "", "exit", Some(exit_code))] - } - _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), - }; - let mut converted = Map::from_iter([ - ( - "type".to_owned(), - Value::String("shell_call_output".to_owned()), - ), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("output".to_owned(), Value::Array(output)), - ]); - if let Some(value) = item.get("max_output_length") - && !value.is_null() - { - converted.insert("max_output_length".to_owned(), value.clone()); - } - Ok(converted) -} - -fn normalize_shell_call_output_input( - item: &Map, -) -> Result, GrokRequestEncodeError> { - let call_id = required_trimmed_string(item, "call_id")?; - let output = normalize_shell_output_blocks(item.get("output"), item.get("status"))?; - let mut converted = Map::from_iter([ - ( - "type".to_owned(), - Value::String("shell_call_output".to_owned()), - ), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("output".to_owned(), Value::Array(output)), - ]); - if let Some(value) = item.get("max_output_length") - && !value.is_null() - { - converted.insert("max_output_length".to_owned(), value.clone()); - } - Ok(converted) -} - -fn normalize_shell_output_blocks( - value: Option<&Value>, - status: Option<&Value>, -) -> Result, GrokRequestEncodeError> { - match value { - Some(Value::Array(blocks)) => blocks - .iter() - .map(|raw| { - let block = raw - .as_object() - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - let stdout = string_field(block, "stdout"); - let stderr = string_field(block, "stderr"); - let outcome = block - .get("outcome") - .and_then(Value::as_object) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - match string_field(outcome, "type").trim() { - "exit" => { - let exit_code = outcome - .get("exit_code") - .or_else(|| outcome.get("exitCode")) - .and_then(Value::as_i64) - .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; - Ok(shell_output_block(stdout, stderr, "exit", Some(exit_code))) - } - "timeout" => Ok(shell_output_block(stdout, stderr, "timeout", None)), - _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), - } - }) - .collect(), - Some(Value::String(output)) => { - let failed = status - .and_then(Value::as_str) - .is_some_and(|status| status.eq_ignore_ascii_case("failed")); - Ok(vec![shell_output_block( - output, - "", - "exit", - Some(i64::from(failed)), - )]) - } - _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), - } -} - -fn shell_output_block( - stdout: &str, - stderr: &str, - outcome_type: &str, - exit_code: Option, -) -> Value { - let mut outcome = Map::from_iter([("type".to_owned(), Value::String(outcome_type.to_owned()))]); - if let Some(exit_code) = exit_code { - outcome.insert("exit_code".to_owned(), Value::from(exit_code)); - } - json_object([ - ("stdout", Value::String(stdout.to_owned())), - ("stderr", Value::String(stderr.to_owned())), - ("outcome", Value::Object(outcome)), - ]) -} - -fn quote_shell_argument(value: &str) -> String { - if value.is_empty() { - return "''".to_owned(); - } - if value - .chars() - .all(|character| character.is_ascii_alphanumeric() || "_@%+=:,./-".contains(character)) - { - return value.to_owned(); - } - format!("'{}'", value.replace('\'', "'\"'\"'")) -} - -fn valid_environment_name(value: &str) -> bool { - value.chars().enumerate().all(|(index, character)| { - character.is_ascii_alphabetic() - || character == '_' - || (index > 0 && character.is_ascii_digit()) - }) && !value.is_empty() -} diff --git a/backend/crates/providers/xai/src/transport/request/history.rs b/backend/crates/providers/xai/src/transport/request/history.rs new file mode 100644 index 000000000..d7cb2fb42 --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/history.rs @@ -0,0 +1,578 @@ +//! OpenAI history 到 Grok Build history 的重建与清理。 + +use super::*; + +pub(super) fn require_content_source( + block: &Map, + fields: &[&str], +) -> Result<(), GrokRequestEncodeError> { + let has_source = fields.iter().any(|field| { + block + .get(*field) + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()) + }); + has_source + .then_some(()) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) +} + +pub(super) fn sanitize_reasoning_input(item: &Map) -> Map { + let mut converted = + copy_non_null_history_fields(item, &["id", "summary", "content", "encrypted_content"]); + // Grok CLI 会在 summary/content 条目上注入 phase 等内部键。 + strip_grok_internal_entry_keys(&mut converted, &["summary", "content"]); + converted.insert("type".to_owned(), Value::String("reasoning".to_owned())); + if has_portable_reasoning_content(&converted) { + converted + } else { + boundary_message( + "A prior model reasoning item was omitted because it has no portable content for Grok Build.", + ) + } +} + +pub(super) fn normalize_compaction_input( + item: &Map, +) -> Result, GrokRequestEncodeError> { + let encrypted = item + .get("encrypted_content") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + let has_structured_compaction_shape = item.contains_key("summary") + || item.contains_key("id") + || item.contains_key("status") + || string_field(item, "type").trim() == "compaction_summary"; + + // 旧版 adapter 曾将明文摘要直接写入 encrypted_content，且不带 + // id/status/summary。仅对这一可识别的历史形态保留明文恢复；新形态严格回放为 + // xAI reasoning 密文 + 可见 conversation summary。 + if !has_structured_compaction_shape { + let Some(summary) = encrypted else { + return Ok(Vec::new()); + }; + let continuation = format!( + "This session is being continued from a previous conversation that ran out of context. \ + The summary below covers the earlier portion of the conversation.\n\n{summary}" + ); + return Ok(vec![Value::Object(compaction_summary_message( + &continuation, + ))]); + } + + let mut converted = Vec::with_capacity(2); + if let Some(encrypted) = encrypted { + converted.push(json!({ + "type": "reasoning", + "summary": [], + "encrypted_content": encrypted, + })); + } + if let Some(summary) = compact_summary_text(item.get("summary")) { + converted.push(Value::Object(compaction_summary_message(&format!( + "\n{summary}\n" + )))); + } + Ok(converted) +} + +pub(super) fn compact_summary_text(value: Option<&Value>) -> Option { + let text = value? + .as_array()? + .iter() + .filter_map(Value::as_object) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .map(str::trim) + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n"); + (!text.is_empty()).then_some(text) +} + +pub(super) fn compaction_summary_message(text: &str) -> Map { + Map::from_iter([ + ("type".to_owned(), Value::String("message".to_owned())), + ("role".to_owned(), Value::String("user".to_owned())), + ( + "content".to_owned(), + Value::Array(vec![json_object([ + ("type", Value::String("input_text".to_owned())), + ("text", Value::String(text.to_owned())), + ])]), + ), + ]) +} + +pub(in crate::transport) fn strip_invalid_encrypted_reasoning_from_body( + body: &mut Map, +) -> bool { + let Some(input) = body.get_mut("input").and_then(Value::as_array_mut) else { + return false; + }; + let has_encrypted_reasoning = input.iter().any(|item| { + item.as_object().is_some_and(|item| { + string_field(item, "type").trim() == "reasoning" + && item + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()) + }) + }); + if !has_encrypted_reasoning { + return false; + } + + input.retain_mut(|item| { + let Some(item) = item.as_object_mut() else { + return true; + }; + if string_field(item, "type").trim() != "reasoning" { + return true; + } + item.remove("encrypted_content"); + if item.get("content").is_some_and(Value::is_null) { + item.remove("content"); + } + item.len() > 1 + }); + true +} + +pub(super) fn has_portable_reasoning_content(item: &Map) -> bool { + item.get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()) + || ["summary", "content"].into_iter().any(|field| { + item.get(field) + .and_then(Value::as_array) + .is_some_and(|values| !values.is_empty()) + }) +} + +pub(super) fn sanitize_native_history_input( + item: &Map, + item_type: &str, +) -> Map { + let fields: &[&str] = match item_type { + "file_search_call" => &["id", "queries", "status", "results"], + "web_search_call" => &["action", "id", "status"], + "image_generation_call" => &["id", "result", "status"], + "code_interpreter_call" => &["code", "container_id", "id", "outputs", "status"], + "shell_call" => &["id", "call_id", "action", "status", "environment"], + "mcp_list_tools" => &["id", "server_label", "tools", "error"], + "mcp_approval_request" => &["arguments", "id", "name", "server_label"], + "mcp_approval_response" => &["approval_request_id", "approve", "id", "reason"], + "mcp_call" => &[ + "arguments", + "id", + "name", + "server_label", + "approval_request_id", + "error", + "output", + "status", + ], + _ => &[], + }; + let mut converted = copy_non_null_history_fields(item, fields); + // Grok CLI 会在 shell_call 的 action object 里注入内部键与 null 占位 + // 字段（如 `timeout_ms: null`）；只在这一层剥离，深层内容原样保留。 + if item_type == "shell_call" + && let Some(Value::Object(action)) = converted.get_mut("action") + { + strip_grok_internal_keys(action); + action.retain(|_, value| !value.is_null()); + } + converted.insert("type".to_owned(), Value::String(item_type.to_owned())); + converted +} + +pub(super) fn copy_non_null_history_fields( + item: &Map, + fields: &[&str], +) -> Map { + fields + .iter() + .filter_map(|field| { + item.get(*field) + .filter(|value| !value.is_null()) + .map(|value| ((*field).to_owned(), value.clone())) + }) + .collect() +} + +pub(super) fn strip_grok_internal_keys(object: &mut Map) { + for key in GROK_INTERNAL_HISTORY_KEYS { + object.remove(*key); + } +} + +pub(super) fn strip_grok_internal_entry_keys(converted: &mut Map, fields: &[&str]) { + for field in fields { + let Some(Value::Array(entries)) = converted.get_mut(*field) else { + continue; + }; + for entry in entries { + if let Some(object) = entry.as_object_mut() { + strip_grok_internal_keys(object); + } + } + } +} + +pub(super) fn unsupported_input_history_boundary(item: &Map, kind: &str) -> Value { + let mut lines = vec![ + "A prior Responses history item was omitted because Grok Build cannot deserialize this Codex item type." + .to_owned(), + format!("Type: {kind}"), + ]; + for key in ["id", "call_id", "name", "status"] { + let value = string_field(item, key).trim(); + if !value.is_empty() { + lines.push(format!("{}: {value}", key.replace('_', " "))); + } + } + Value::Object(boundary_message(&lines.join("\n"))) +} + +pub(super) fn boundary_message(text: &str) -> Map { + Map::from_iter([ + ("type".to_owned(), Value::String("message".to_owned())), + ("role".to_owned(), Value::String("developer".to_owned())), + ( + "content".to_owned(), + Value::Array(vec![json_object([ + ("type", Value::String("input_text".to_owned())), + ("text", Value::String(text.to_owned())), + ])]), + ), + ]) +} + +pub(super) fn normalize_agent_message_input(item: &Map) -> Value { + let Some(content) = item.get("content").and_then(text_input_content) else { + return Value::Object(boundary_message( + "An encrypted inter-agent message occurred here but is not portable to the Grok Build account.", + )); + }; + let author = non_empty_or(string_field(item, "author"), "agent"); + let recipient = non_empty_or(string_field(item, "recipient"), "recipient"); + Value::Object(boundary_message(&format!( + "Agent message ({author} -> {recipient}):\n{content}" + ))) +} + +pub(super) fn text_input_content(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Array(items) => items + .iter() + .map(|item| { + let item = item.as_object()?; + matches!( + string_field(item, "type"), + "input_text" | "output_text" | "text" + ) + .then(|| string_field(item, "text").to_owned()) + }) + .collect::>>() + .map(|parts| parts.join("\n")), + _ => None, + } +} + +pub(super) fn non_empty_or<'a>(value: &'a str, fallback: &'a str) -> &'a str { + let value = value.trim(); + if value.is_empty() { fallback } else { value } +} + +pub(super) fn normalize_mcp_output_input( + item: &Map, +) -> Result { + let output = serde_json::to_string(item.get("output").unwrap_or(&Value::Null)) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; + let call_id = non_empty_or(string_field(item, "call_id"), "unknown"); + Ok(Value::Object(boundary_message(&format!( + "MCP tool output for call {call_id}: {output}" + )))) +} + +pub(super) fn validate_apply_patch_operation( + value: Option<&Value>, +) -> Result, GrokRequestEncodeError> { + let operation = value + .and_then(Value::as_object) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let kind = required_trimmed_string(operation, "type")?; + required_trimmed_string(operation, "path")?; + match kind { + "create_file" | "update_file" if operation.get("diff").is_some_and(Value::is_string) => {} + "delete_file" => {} + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + Ok(operation.clone()) +} + +pub(super) fn normalize_apply_patch_output_input( + item: &Map, +) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let status = match string_field(item, "status").trim() { + "" | "completed" => "completed", + "failed" => "failed", + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let output = encode_tool_output(item.get("output"))?; + let mut message = format!("Apply patch status: {status}"); + if !output.is_empty() { + message.push('\n'); + message.push_str(&output); + } + Ok(Map::from_iter([ + ( + "type".to_owned(), + Value::String("function_call_output".to_owned()), + ), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("output".to_owned(), Value::String(message)), + ])) +} + +pub(super) fn normalize_legacy_local_shell_call_input( + item: &Map, +) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let action = legacy_shell_action(item.get("action"))?; + let mut converted = Map::from_iter([ + ("type".to_owned(), Value::String("shell_call".to_owned())), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("action".to_owned(), Value::Object(action)), + ]); + for key in ["id", "status", "timeout_ms", "max_output_length"] { + if let Some(value) = item.get(key) { + converted.insert(key.to_owned(), value.clone()); + } + } + Ok(converted) +} + +pub(super) fn legacy_shell_action( + value: Option<&Value>, +) -> Result, GrokRequestEncodeError> { + let action = value + .and_then(Value::as_object) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if !matches!(string_field(action, "type").trim(), "" | "exec") { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let command = legacy_shell_command(action)?; + Ok(Map::from_iter([ + ("type".to_owned(), Value::String("exec".to_owned())), + ( + "commands".to_owned(), + Value::Array(vec![Value::String(command)]), + ), + ])) +} + +pub(super) fn legacy_shell_command( + action: &Map, +) -> Result { + let mut command = match action.get("command") { + Some(Value::String(value)) => value.trim().to_owned(), + Some(Value::Array(parts)) => parts + .iter() + .map(|part| { + part.as_str() + .map(quote_shell_argument) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) + }) + .collect::, _>>()? + .join(" "), + _ => action + .get("commands") + .and_then(Value::as_array) + .map(|commands| { + commands + .iter() + .map(|command| { + command + .as_str() + .map(ToOwned::to_owned) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) + }) + .collect::, _>>() + .map(|commands| commands.join("\n")) + }) + .transpose()? + .unwrap_or_default(), + }; + if command.is_empty() { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + if let Some(environment) = action.get("env").and_then(Value::as_object) + && !environment.is_empty() + { + let assignments = environment + .iter() + .map(|(name, value)| { + if !valid_environment_name(name) { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let value = value + .as_str() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + Ok(format!("{name}={}", quote_shell_argument(value))) + }) + .collect::, _>>()?; + command = format!("env {} {command}", assignments.join(" ")); + } + if let Some(directory) = action + .get("working_directory") + .and_then(Value::as_str) + .map(str::trim) + .filter(|directory| !directory.is_empty()) + { + command = format!("cd {} && {command}", quote_shell_argument(directory)); + } + Ok(command) +} + +pub(super) fn normalize_legacy_local_shell_output_input( + item: &Map, +) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let output = match item.get("output") { + Some(Value::Array(output)) => output.clone(), + Some(Value::String(output)) => { + let exit_code = item + .get("exit_code") + .and_then(Value::as_i64) + .unwrap_or_else(|| { + i64::from(string_field(item, "status").eq_ignore_ascii_case("failed")) + }); + vec![shell_output_block(output, "", "exit", Some(exit_code))] + } + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let mut converted = Map::from_iter([ + ( + "type".to_owned(), + Value::String("shell_call_output".to_owned()), + ), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("output".to_owned(), Value::Array(output)), + ]); + if let Some(value) = item.get("max_output_length") + && !value.is_null() + { + converted.insert("max_output_length".to_owned(), value.clone()); + } + Ok(converted) +} + +pub(super) fn normalize_shell_call_output_input( + item: &Map, +) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let output = normalize_shell_output_blocks(item.get("output"), item.get("status"))?; + let mut converted = Map::from_iter([ + ( + "type".to_owned(), + Value::String("shell_call_output".to_owned()), + ), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("output".to_owned(), Value::Array(output)), + ]); + if let Some(value) = item.get("max_output_length") + && !value.is_null() + { + converted.insert("max_output_length".to_owned(), value.clone()); + } + Ok(converted) +} + +pub(super) fn normalize_shell_output_blocks( + value: Option<&Value>, + status: Option<&Value>, +) -> Result, GrokRequestEncodeError> { + match value { + Some(Value::Array(blocks)) => blocks + .iter() + .map(|raw| { + let block = raw + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let stdout = string_field(block, "stdout"); + let stderr = string_field(block, "stderr"); + let outcome = block + .get("outcome") + .and_then(Value::as_object) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + match string_field(outcome, "type").trim() { + "exit" => { + let exit_code = outcome + .get("exit_code") + .or_else(|| outcome.get("exitCode")) + .and_then(Value::as_i64) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + Ok(shell_output_block(stdout, stderr, "exit", Some(exit_code))) + } + "timeout" => Ok(shell_output_block(stdout, stderr, "timeout", None)), + _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + }) + .collect(), + Some(Value::String(output)) => { + let failed = status + .and_then(Value::as_str) + .is_some_and(|status| status.eq_ignore_ascii_case("failed")); + Ok(vec![shell_output_block( + output, + "", + "exit", + Some(i64::from(failed)), + )]) + } + _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), + } +} + +pub(super) fn shell_output_block( + stdout: &str, + stderr: &str, + outcome_type: &str, + exit_code: Option, +) -> Value { + let mut outcome = Map::from_iter([("type".to_owned(), Value::String(outcome_type.to_owned()))]); + if let Some(exit_code) = exit_code { + outcome.insert("exit_code".to_owned(), Value::from(exit_code)); + } + json_object([ + ("stdout", Value::String(stdout.to_owned())), + ("stderr", Value::String(stderr.to_owned())), + ("outcome", Value::Object(outcome)), + ]) +} + +pub(super) fn quote_shell_argument(value: &str) -> String { + if value.is_empty() { + return "''".to_owned(); + } + if value + .chars() + .all(|character| character.is_ascii_alphanumeric() || "_@%+=:,./-".contains(character)) + { + return value.to_owned(); + } + format!("'{}'", value.replace('\'', "'\"'\"'")) +} + +pub(super) fn valid_environment_name(value: &str) -> bool { + value.chars().enumerate().all(|(index, character)| { + character.is_ascii_alphabetic() + || character == '_' + || (index > 0 && character.is_ascii_digit()) + }) && !value.is_empty() +} diff --git a/backend/crates/providers/xai/src/transport/request/mod.rs b/backend/crates/providers/xai/src/transport/request/mod.rs new file mode 100644 index 000000000..edbf3be62 --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/mod.rs @@ -0,0 +1,923 @@ +//! OpenAI Responses 请求到官方 Grok Build wire 的转换边界。 + +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt; + +use gateway_core::operation::GenerateRequest; +use gateway_core::policy::ClientApiKeyId; +use serde_json::{Map, Value, json}; +use sha2::{Digest as _, Sha256}; + +use super::{GrokSessionAffinityKey, XAI_PROVIDER_NAME}; + +mod history; +mod response; +mod tools; + +pub(super) use history::strip_invalid_encrypted_reasoning_from_body; +pub(crate) use response::GrokResponseTransform; + +use history::*; +use tools::*; + +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", + "previousResponseId", + "previous_response_id", + "responseId", + "response_id", + "sessionId", + "session_id", + "x-grok-agent-id", + "x-grok-conv-id", + "x-grok-session-id", +]; +const SESSION_FIELDS: &[&str] = &[ + "prompt_cache_key", + "session_id", + "sessionId", + "conversation_id", + "conversationId", +]; +const PROTOCOL_CONTEXT_SESSION_FIELDS: &[&str] = &["conversation_id", "session_id", "thread_id"]; +const MAX_SESSION_SEED_BYTES: usize = 1_024; +const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"]; +/// Grok CLI 在历史条目上注入、Grok Build 无法反序列化的内部键；只在已知 +/// 注入位置剥离，避免误伤工具 schema 或输出中恰好同名的键。 +const GROK_INTERNAL_HISTORY_KEYS: &[&str] = + &["phase", "internal_chat_message_metadata_passthrough"]; + +/// 保留客户端 OpenAI Responses object 的 xAI 上游请求。 +pub struct GrokResponsesRequest { + body: Map, + session_id: Option, + reasoning_replay_session_id: Option, + affinity: Option, + response_transform: GrokResponseTransform, +} + +impl GrokResponsesRequest { + /// 返回发送到 `/v1/responses` 的 JSON object。 + #[must_use] + pub const fn body(&self) -> &Map { + &self.body + } + + /// 返回按下游租户隔离后的稳定 Grok 会话 UUID。 + #[must_use] + pub fn session_id(&self) -> Option<&str> { + self.session_id.as_deref() + } + + /// 返回只有显式会话身份才启用的 reasoning replay scope。 + pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> { + self.reasoning_replay_session_id.as_deref() + } + + /// 返回归一化后的 xAI wire 模型。 + pub(crate) fn upstream_model(&self) -> Option<&str> { + self.body.get("model").and_then(Value::as_str) + } + + pub(crate) fn has_previous_response_id(&self) -> bool { + self.body + .get("previous_response_id") + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| !value.is_empty()) + } + + pub(crate) fn replay_input_items(&self) -> Option> { + self.body.get("input").and_then(Value::as_array).cloned() + } + + /// 返回与会话一致、额外绑定模型的账号亲和键。 + #[must_use] + pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> { + self.affinity.as_ref() + } + + pub(crate) fn response_transform(&self) -> GrokResponseTransform { + self.response_transform.clone() + } + + pub(crate) fn input_items(&self) -> Vec { + match self.body.get("input") { + Some(Value::Array(items)) => items.clone(), + Some(Value::String(input)) => vec![json_object([ + ("type", Value::String("message".to_owned())), + ("role", Value::String("user".to_owned())), + ("content", Value::String(input.clone())), + ])], + _ => Vec::new(), + } + } + + pub(crate) fn set_replay_input( + &mut self, + input: Vec, + ) -> Result<(), GrokRequestEncodeError> { + let mut normalizer = ToolNormalizer::for_replay(self.response_transform.clone()); + let mut normalized = Vec::with_capacity(input.len()); + for item in input { + match item { + Value::Object(item) if string_field(&item, "type") == "custom_tool_call" => { + normalized.push(Value::Object( + normalizer.normalize_custom_tool_call_input(&item)?, + )); + } + item => normalized.push(item), + } + } + self.response_transform = normalizer.response; + self.body + .insert("input".to_owned(), Value::Array(normalized)); + Ok(()) + } + + /// 为同账号的一次 xAI `invalid_encrypted_content` 恢复请求移除被拒绝的密文。 + /// + /// 只在至少一个 reasoning item 含非空密文时改写；可读 summary/content、ID 与 + /// 其他历史保持不变，剥离后只剩 `type` 的空壳一并删除。 + pub(crate) fn strip_invalid_encrypted_reasoning(&mut self) -> bool { + strip_invalid_encrypted_reasoning_from_body(&mut self.body) + } + + pub(crate) fn set_previous_response_id(&mut self, response_id: Option) { + match response_id { + Some(response_id) => { + self.body.insert( + "previous_response_id".to_owned(), + Value::String(response_id), + ); + } + None => { + self.body.remove("previous_response_id"); + } + } + } + + pub(crate) fn inherit_session(&mut self, session_id: Option<&str>) { + let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) else { + self.clear_session(); + return; + }; + self.session_id = Some(session_id.to_owned()); + if self.reasoning_replay_session_id.is_some() { + self.reasoning_replay_session_id = Some(session_id.to_owned()); + } + self.affinity = None; + self.body.insert( + "prompt_cache_key".to_owned(), + Value::String(session_id.to_owned()), + ); + } + + pub(crate) fn clear_session(&mut self) { + self.body.remove("prompt_cache_key"); + self.session_id = None; + self.reasoning_replay_session_id = None; + self.affinity = None; + } + + pub fn encode( + request: &GenerateRequest, + upstream_model: &str, + client_api_key_ref: &ClientApiKeyId, + ) -> Result { + Self::encode_inner(request, upstream_model, client_api_key_ref, true, false) + } + + pub(super) fn encode_compaction_source( + request: &GenerateRequest, + upstream_model: &str, + client_api_key_ref: &ClientApiKeyId, + ) -> Result { + Self::encode_inner(request, upstream_model, client_api_key_ref, false, true) + } + + fn encode_inner( + request: &GenerateRequest, + upstream_model: &str, + client_api_key_ref: &ClientApiKeyId, + allow_cache_route: bool, + should_consume_terminal_compaction_trigger: bool, + ) -> Result { + let payload = request.protocol_payload(); + if payload.protocol() != "openai" { + return Err(GrokRequestEncodeError::InvalidProtocolPayload); + } + let mut body = payload.body().clone(); + if should_consume_terminal_compaction_trigger { + consume_terminal_compaction_trigger(&mut body)?; + } + let upstream_model = resolve_grok_text_responses_model_id(upstream_model); + // 这些字段属于 Codex/OpenAI 侧请求控制，不是 xAI 上游协议字段。 + // OpenAI 透明路径会保留未知字段；这里只在 xAI adapter 内做最小剥离。 + body.remove("provider_options"); + body.remove("service_tier"); + let session_seed = explicit_session_seed(request, &body); + let enable_cache_route = allow_cache_route && session_seed.is_some(); + let identity = resolve_session_identity( + client_api_key_ref.as_str(), + &upstream_model, + session_seed.as_deref(), + &body, + ); + sanitize_account_identity(&mut body); + sanitize_client_metadata(&mut body); + normalize_build_request(&mut body, &upstream_model)?; + let mut response_transform = normalize_responses_request(&mut body)?; + response_transform.observe_client_cache_tools(); + if enable_cache_route { + enable_grok_prompt_cache_route(&mut body, &upstream_model, &mut response_transform); + } + let (session_id, affinity) = identity.map_or((None, None), |(session_id, affinity)| { + (Some(session_id), Some(affinity)) + }); + let reasoning_replay_session_id = + session_seed.is_some().then(|| session_id.clone()).flatten(); + match session_id.as_ref() { + Some(session_id) => { + body.insert( + "prompt_cache_key".to_owned(), + Value::String(session_id.clone()), + ); + } + None => { + body.remove("prompt_cache_key"); + } + } + body.insert("model".to_owned(), Value::String(upstream_model)); + body.insert("stream".to_owned(), Value::Bool(true)); + Ok(Self { + body, + session_id, + reasoning_replay_session_id, + affinity, + response_transform, + }) + } + + pub(crate) fn to_json_bytes(&self) -> Result, GrokRequestEncodeError> { + serde_json::to_vec(&self.body).map_err(|_| GrokRequestEncodeError::Serialization) + } +} + +fn resolve_grok_text_responses_model_id(model: &str) -> String { + let stripped = strip_grok_provider_prefix(model); + let normalized = stripped.to_ascii_lowercase(); + match normalized.as_str() { + "grok" | "grok-latest" | "grok-4.5" | "grok-4.5-latest" | "grok-build-latest" => { + "grok-4.5".to_owned() + } + "grok-4.6" | "grok-4.6-latest" => "grok-4.6".to_owned(), + "grok-4.3" | "grok-4.3-latest" => "grok-4.3".to_owned(), + "grok-3-mini" + | "grok-3-mini-fast" + | "grok-build-0.1" + | "grok-composer-2.5-fast" + | "grok-4.20-0309-reasoning" + | "grok-4.20-0309-non-reasoning" + | "grok-4.20-multi-agent-0309" => normalized, + "grok-build" => "grok-build-0.1".to_owned(), + "grok-composer" | "composer-2.5" => "grok-composer-2.5-fast".to_owned(), + "grok-4.20-reasoning" => "grok-4.20-0309-reasoning".to_owned(), + "grok-4.20-non-reasoning" => "grok-4.20-0309-non-reasoning".to_owned(), + "grok-4.20-multi-agent" | "grok-4.20-multi-agent-latest" => { + "grok-4.20-multi-agent-0309".to_owned() + } + _ => stripped.to_owned(), + } +} + +fn consume_terminal_compaction_trigger( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + let Some(Value::Array(input)) = body.get_mut("input") else { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + }; + let Some(Value::Object(last)) = input.last() else { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + }; + if last.get("type").and_then(Value::as_str) != Some("compaction_trigger") { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + input.pop(); + Ok(()) +} + +impl fmt::Debug for GrokResponsesRequest { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("GrokResponsesRequest") + .field("body_keys", &self.body.keys().collect::>()) + .field("has_session", &self.session_id.is_some()) + .field( + "has_response_transform", + &!self.response_transform.is_empty(), + ) + .field("body", &"") + .finish() + } +} + +/// Generate 到 Responses 的编码错误，不保留 option 与 prompt 值。 +#[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, + /// JSON 序列化意外失败。 + #[error("Grok Build request serialization failed")] + Serialization, + /// Responses 兼容字段无法安全归一化。 + #[error("Grok Build request normalization failed")] + InvalidRequestNormalization, + /// 请求中的具体字段无法安全转换为 Grok Build 接受的形态。 + #[error("Grok Build request field `{field}` could not be normalized safely")] + InvalidRequestField { field: &'static str }, +} + +impl GrokRequestEncodeError { + fn at_field(self, field: &'static str) -> Self { + match self { + Self::InvalidRequestNormalization => Self::InvalidRequestField { field }, + error => error, + } + } +} + +fn sanitize_account_identity(body: &mut Map) { + for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { + body.remove(*field); + } + let Some(Value::Object(metadata)) = body.get_mut("metadata") else { + return; + }; + for field in IDENTITY_FIELDS.iter().chain(ACCOUNT_BOUND_FIELDS) { + metadata.remove(*field); + } +} + +fn sanitize_client_metadata(body: &mut Map) { + // Codex 的 client_metadata 是本地 transport envelope，可能包含工作目录、仓库地址、 + // installation/session 标识。会话亲和信息已在调用本函数前提取，整个 envelope 都不能 + // 越过 Grok Build 边界。 + body.remove("client_metadata"); +} + +fn enable_grok_prompt_cache_route( + body: &mut Map, + upstream_model: &str, + response: &mut GrokResponseTransform, +) { + let tools = body + .entry("tools".to_owned()) + .or_insert_with(|| Value::Array(Vec::new())); + let Value::Array(tools) = tools else { + return; + }; + response.observe_upstream_cache_tools(tools); + if is_cache_media_model(upstream_model) || has_tool_type(tools, "image_generation") { + return; + } + if tools.is_empty() { + for tool in GROK_CACHE_ROUTE_TOOLS { + tools.push(json_object([("type", Value::String((*tool).to_owned()))])); + response.mark_injected_cache_tool(tool); + } + // 无客户端工具时以 none 选中缓存路由，同时不授予搜索能力。 + body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); + } else if !has_tool_type(tools, "x_search") { + tools.push(json_object([( + "type", + Value::String("x_search".to_owned()), + )])); + response.mark_injected_cache_tool("x_search"); + } +} + +fn is_cache_media_model(model: &str) -> bool { + let model = model.trim().to_ascii_lowercase(); + ["image", "imagine", "video"] + .into_iter() + .any(|marker| model.contains(marker)) +} + +fn explicit_session_seed(request: &GenerateRequest, body: &Map) -> Option { + request + .prompt_cache_key() + .and_then(valid_session_seed) + .map(ToOwned::to_owned) + .or_else(|| first_session_value(body)) + .or_else(|| { + body.get("metadata") + .and_then(Value::as_object) + .and_then(first_session_value) + }) + .or_else(|| { + PROTOCOL_CONTEXT_SESSION_FIELDS.iter().find_map(|field| { + request + .protocol_payload() + .context() + .get(*field) + .and_then(Value::as_str) + .and_then(valid_session_seed) + .map(ToOwned::to_owned) + }) + }) +} + +fn first_session_value(body: &Map) -> Option { + let prompt_cache_key = body + .get("prompt_cache_key") + .and_then(Value::as_str) + .and_then(valid_session_seed) + .map(ToOwned::to_owned); + prompt_cache_key + .or_else(|| { + let metadata = body.get("metadata")?.as_object()?; + ["session_id", "sessionId"] + .into_iter() + .find_map(|field| metadata.get(field).and_then(Value::as_str)) + .and_then(valid_session_seed) + .map(ToOwned::to_owned) + .or_else(|| { + metadata + .get("user_id") + .and_then(Value::as_str) + .and_then(session_seed_from_user_id) + }) + }) + .or_else(|| { + SESSION_FIELDS[1..] + .iter() + .find_map(|field| body.get(*field).and_then(Value::as_str)) + .and_then(valid_session_seed) + .map(ToOwned::to_owned) + }) +} + +fn valid_session_seed(value: &str) -> Option<&str> { + let value = value.trim(); + (!value.is_empty() + && value.len() <= MAX_SESSION_SEED_BYTES + && !value.chars().any(char::is_control)) + .then_some(value) +} + +fn session_seed_from_user_id(value: &str) -> Option { + let value = value.trim(); + if let Ok(Value::Object(embedded)) = serde_json::from_str::(value) + && let Some(seed) = ["session_id", "sessionId"] + .into_iter() + .find_map(|field| embedded.get(field).and_then(Value::as_str)) + .and_then(valid_session_seed) + { + return Some(seed.to_owned()); + } + value + .rfind("_session_") + .and_then(|index| valid_session_seed(&value[index + "_session_".len()..])) + .map(ToOwned::to_owned) +} + +fn resolve_session_identity( + client_api_key_ref: &str, + upstream_model: &str, + explicit_seed: Option<&str>, + body: &Map, +) -> Option<(String, GrokSessionAffinityKey)> { + let model = upstream_model.trim().to_ascii_lowercase(); + if client_api_key_ref.is_empty() || model.is_empty() { + return None; + } + if let Some(seed) = explicit_seed { + let upstream_source = + format!("xai:build-session:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{seed}"); + let affinity_source = format!( + "xai:build-affinity:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{model}:{seed}" + ); + return Some(( + digest_uuid(&upstream_source), + GrokSessionAffinityKey::from_digest(Sha256::digest(affinity_source).into()), + )); + } + let (system, first_user) = message_anchors(body); + let first_user = truncate_anchor(&first_user, 200); + if first_user.is_empty() { + return None; + } + let system = truncate_anchor(&system, 100); + let upstream_source = format!( + "xai:build-soft-session:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{system}:{first_user}" + ); + let affinity_source = format!( + "xai:build-soft-affinity:v2:{client_api_key_ref}:{XAI_PROVIDER_NAME}:{model}:{system}:{first_user}" + ); + Some(( + digest_uuid(&upstream_source), + GrokSessionAffinityKey::from_digest(Sha256::digest(affinity_source).into()), + )) +} + +fn digest_uuid(source: &str) -> String { + let digest = Sha256::digest(source); + let value = digest[..16] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + format!( + "{}-{}-{}-{}-{}", + &value[0..8], + &value[8..12], + &value[12..16], + &value[16..20], + &value[20..32] + ) +} + +fn truncate_anchor(value: &str, max_chars: usize) -> String { + value.trim().chars().take(max_chars).collect() +} + +fn message_anchors(body: &Map) -> (String, String) { + let mut system = body + .get("instructions") + .map(flatten_message_content) + .filter(|value| !value.is_empty()) + .or_else(|| body.get("system").map(flatten_message_content)) + .unwrap_or_default(); + let mut first_user = String::new(); + if let Some(messages) = body.get("messages").and_then(Value::as_array) { + collect_role_anchors(messages, &mut system, &mut first_user); + } + if first_user.is_empty() { + match body.get("input") { + Some(Value::String(value)) => first_user = value.trim().to_owned(), + Some(Value::Array(items)) => { + collect_role_anchors(items, &mut system, &mut first_user); + } + _ => {} + } + } + (system, first_user) +} + +fn collect_role_anchors(items: &[Value], system: &mut String, first_user: &mut String) { + for item in items { + let Some(item) = item.as_object() else { + continue; + }; + let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default(); + if !item_type.is_empty() && item_type != "message" { + continue; + } + let role = item + .get("role") + .and_then(Value::as_str) + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + let content = item + .get("content") + .map(flatten_message_content) + .filter(|value| !value.is_empty()) + .or_else(|| { + item.get("text") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + }) + .unwrap_or_default(); + if content.is_empty() { + continue; + } + match role.as_str() { + "system" | "developer" if system.is_empty() => *system = content, + "user" if first_user.is_empty() => *first_user = content, + "" if first_user.is_empty() => *first_user = content, + _ => {} + } + if !first_user.is_empty() && !system.is_empty() { + break; + } + } +} + +fn flatten_message_content(value: &Value) -> String { + match value { + Value::String(value) => value.trim().to_owned(), + Value::Array(parts) => parts + .iter() + .filter_map(Value::as_object) + .filter(|part| { + matches!( + part.get("type").and_then(Value::as_str).unwrap_or_default(), + "" | "text" | "input_text" | "output_text" + ) + }) + .filter_map(|part| part.get("text").and_then(Value::as_str).map(str::trim)) + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n"), + _ => String::new(), + } +} + +fn normalize_responses_request( + body: &mut Map, +) -> Result { + if let Some(response_format) = body.remove("response_format") { + let text = body + .entry("text".to_owned()) + .or_insert_with(|| Value::Object(Map::new())); + if text.is_null() { + *text = Value::Object(Map::new()); + } + let text = text + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "text" })?; + if text.get("format").is_none_or(Value::is_null) { + text.insert( + "format".to_owned(), + normalize_response_format(response_format) + .map_err(|error| error.at_field("response_format"))?, + ); + } + } + patch_reasoning_text_types(body); + ToolNormalizer::new().normalize(body) +} + +fn normalize_build_request( + body: &mut Map, + upstream_model: &str, +) -> Result<(), GrokRequestEncodeError> { + apply_build_response_defaults(body)?; + normalize_build_reasoning_effort(body, upstream_model); + sanitize_build_model_fields(body, upstream_model); + Ok(()) +} + +fn apply_build_response_defaults( + body: &mut Map, +) -> Result<(), GrokRequestEncodeError> { + if body.get("store").is_none_or(Value::is_null) { + body.insert("store".to_owned(), Value::Bool(false)); + } + + let include = body + .entry("include".to_owned()) + .or_insert_with(|| Value::Array(Vec::new())); + if include.is_null() { + *include = Value::Array(Vec::new()); + } + let include = include + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestField { field: "include" })?; + if include.iter().any(|value| !value.is_string()) { + return Err(GrokRequestEncodeError::InvalidRequestField { field: "include" }); + } + if !include + .iter() + .any(|value| value.as_str() == Some("reasoning.encrypted_content")) + { + include.push(Value::String("reasoning.encrypted_content".to_owned())); + } + Ok(()) +} + +fn normalize_build_reasoning_effort(body: &mut Map, upstream_model: &str) { + let supports_effort = grok_model_supports_reasoning_effort(upstream_model); + if let Some(reasoning) = body.get_mut("reasoning").and_then(Value::as_object_mut) { + if let Some(effort) = reasoning.remove("effort") + && supports_effort + && let Some(effort) = normalize_grok_reasoning_effort_value(&effort) + { + reasoning.insert("effort".to_owned(), Value::String(effort.to_owned())); + } + if reasoning.is_empty() { + body.remove("reasoning"); + } + } + + if let Some(effort) = body.remove("reasoning_effort") + && supports_effort + && let Some(effort) = normalize_grok_reasoning_effort_value(&effort) + { + body.insert( + "reasoning_effort".to_owned(), + Value::String(effort.to_owned()), + ); + } + + let camel_effort = body.remove("reasoningEffort"); + if !body.contains_key("reasoning_effort") + && supports_effort + && let Some(effort) = camel_effort + .as_ref() + .and_then(normalize_grok_reasoning_effort_value) + { + body.insert( + "reasoning_effort".to_owned(), + Value::String(effort.to_owned()), + ); + } + + if is_grok_composer_model(upstream_model) { + body.remove("reasoning"); + body.remove("reasoning_effort"); + body.remove("reasoningEffort"); + } +} + +fn normalize_grok_reasoning_effort_value(value: &Value) -> Option<&'static str> { + let normalized = value + .as_str()? + .trim() + .to_ascii_lowercase() + .chars() + .filter(|character| !matches!(character, '-' | '_' | ' ')) + .collect::(); + match normalized.as_str() { + "none" => Some("none"), + "minimal" | "low" => Some("low"), + "medium" => Some("medium"), + "high" | "xhigh" | "extrahigh" | "max" | "ultra" => Some("high"), + _ => None, + } +} + +fn strip_grok_provider_prefix(model: &str) -> &str { + let model = model.trim(); + let lower = model.to_ascii_lowercase(); + ["xai/", "x-ai/", "grok/"] + .into_iter() + .find_map(|prefix| { + lower + .starts_with(prefix) + .then(|| model[prefix.len()..].trim()) + }) + .unwrap_or(model) +} + +fn grok_model_last_slug(model: &str) -> String { + model + .trim() + .rsplit_once('/') + .map_or_else(|| model.trim(), |(_, slug)| slug.trim()) + .to_ascii_lowercase() +} + +fn is_grok_composer_model(model: &str) -> bool { + matches!( + grok_model_last_slug(model).as_str(), + "grok-composer" | "grok-composer-2.5-fast" | "composer-2.5" + ) +} + +fn grok_model_supports_reasoning_effort(model: &str) -> bool { + matches!( + strip_grok_provider_prefix(model) + .to_ascii_lowercase() + .as_str(), + "grok-4.5" + | "grok-4.5-latest" + | "grok-4.6" + | "grok-4.6-latest" + | "grok-4.3" + | "grok-4.3-latest" + | "grok-3-mini" + | "grok-3-mini-fast" + | "grok-4.20-0309-reasoning" + | "grok-4.20-reasoning" + | "grok-4.20-multi-agent-0309" + ) +} + +fn sanitize_build_model_fields(body: &mut Map, upstream_model: &str) { + for field in ["prompt_cache_retention", "safety_identifier"] { + body.remove(field); + } + if upstream_model.trim().eq_ignore_ascii_case("grok-4.5") { + for field in [ + "presence_penalty", + "presencePenalty", + "frequency_penalty", + "frequencyPenalty", + "stop", + ] { + body.remove(field); + } + } + if grok_model_last_slug(upstream_model).starts_with("grok-4.20") { + body.remove("logprobs"); + body.remove("top_logprobs"); + } + delete_json_field_recursive(body, "external_web_access"); +} + +fn delete_json_field_recursive(object: &mut Map, field: &str) { + object.remove(field); + for value in object.values_mut() { + delete_json_field_from_value(value, field); + } +} + +fn delete_json_field_from_value(value: &mut Value, field: &str) { + match value { + Value::Object(child) => delete_json_field_recursive(child, field), + Value::Array(children) => { + for child in children { + delete_json_field_from_value(child, field); + } + } + _ => {} + } +} + +fn normalize_response_format(value: Value) -> Result { + let Some(format) = value.as_object() else { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + }; + if format.get("type").and_then(Value::as_str) != Some("json_schema") { + return Ok(value); + } + let Some(schema) = format.get("json_schema").and_then(Value::as_object) else { + return Ok(value); + }; + let mut normalized = Map::new(); + normalized.insert("type".to_owned(), Value::String("json_schema".to_owned())); + normalized.extend( + schema + .iter() + .filter(|(key, _)| key.as_str() != "type") + .map(|(key, value)| (key.clone(), value.clone())), + ); + Ok(Value::Object(normalized)) +} + +fn patch_reasoning_text_types(body: &mut Map) { + let Some(items) = body.get_mut("input").and_then(Value::as_array_mut) else { + return; + }; + for item in items { + let Some(item) = item.as_object_mut() else { + continue; + }; + if item.get("type").and_then(Value::as_str) != Some("reasoning") { + continue; + } + let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else { + continue; + }; + for part in content { + let Some(part) = part.as_object_mut() else { + continue; + }; + part.entry("type".to_owned()) + .or_insert_with(|| Value::String("reasoning_text".to_owned())); + } + } +} + +const MAX_BUILD_TOOL_ALIAS_LENGTH: usize = 128; +const MAX_TOOL_SEARCH_DESCRIPTION_BYTES: usize = 16 << 10; +const MAX_BUFFERED_FUNCTION_ARGUMENTS_BYTES: usize = 1 << 20; +const MAX_TOTAL_BUFFERED_FUNCTION_ARGUMENTS_BYTES: usize = 4 << 20; +const MAX_JSON_NUMBER_TEXT_BYTES: usize = 256; +const MAX_EXACT_JSON_INTEGER_TEXT: &str = "9007199254740991"; +const MAX_EXACT_JSON_INTEGER: u64 = 9_007_199_254_740_991; diff --git a/backend/crates/providers/xai/src/transport/request/response.rs b/backend/crates/providers/xai/src/transport/request/response.rs new file mode 100644 index 000000000..cc65adcae --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/response.rs @@ -0,0 +1,839 @@ +//! Grok Build 输出恢复为 OpenAI Responses wire 所需的状态。 + +use super::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(super) enum ToolKind { + Function, + Custom, + ToolSearch, + ApplyPatch, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(super) struct ToolIdentity { + pub(super) kind: ToolKind, + pub(super) namespace: String, + pub(super) name: String, +} + +impl ToolIdentity { + pub(super) fn new(kind: ToolKind, namespace: &str, name: &str) -> Self { + Self { + kind, + namespace: namespace.to_owned(), + name: name.to_owned(), + } + } + + pub(super) fn is_root_custom_apply_patch(&self) -> bool { + self.kind == ToolKind::Custom && self.namespace.is_empty() && self.name == "apply_patch" + } + + pub(super) fn custom_argument_field(&self) -> &'static str { + if self.is_root_custom_apply_patch() { + "patch" + } else { + "input" + } + } +} + +#[derive(Debug, Clone)] +pub(super) struct StreamCallState { + identity: ToolIdentity, + schema: Option, + arguments: String, + passthrough: bool, + last_delta: Option>, + added_payload: Option>, +} + +#[derive(Clone, Default)] +pub(crate) struct GrokResponseTransform { + pub(super) aliases: BTreeMap, + pub(super) function_schemas: BTreeMap, + pub(super) visible_tools: Vec, + pub(super) legacy_local_shell: bool, + pub(super) filter_x_search: bool, + pub(super) injected_tool_types: BTreeSet, + pub(super) client_declared_tools: BTreeSet, + pub(super) dropped_output_indexes: BTreeSet, + pub(super) dropped_item_ids: BTreeSet, + pub(super) stream_calls: BTreeMap, + pub(super) stream_keys: BTreeMap, + pub(super) stream_argument_bytes: usize, + pub(super) stream_sequence_next: Option, +} + +pub(crate) struct GrokTransformedWireEvent { + event_type: String, + value: Value, +} + +impl GrokTransformedWireEvent { + #[must_use] + pub(crate) fn event_type(&self) -> &str { + &self.event_type + } + + #[must_use] + pub(crate) fn into_value(self) -> Value { + self.value + } +} + +impl GrokResponseTransform { + #[must_use] + pub(crate) fn is_empty(&self) -> bool { + self.aliases.is_empty() + && self.function_schemas.is_empty() + && self.visible_tools.is_empty() + && !self.legacy_local_shell + && !self.filter_x_search + && self.injected_tool_types.is_empty() + } + + pub(super) fn observe_client_cache_tools(&mut self) { + for tool in &self.visible_tools { + let kind = tool.get("type").and_then(Value::as_str).unwrap_or_default(); + if kind == "x_search" { + self.filter_x_search = true; + } + if matches!(kind, "function" | "custom") + && let Some(name) = tool + .get("name") + .and_then(Value::as_str) + .map(str::trim) + .filter(|name| !name.is_empty()) + { + self.client_declared_tools.insert(name.to_owned()); + } + } + } + + pub(super) fn observe_upstream_cache_tools(&mut self, tools: &[Value]) { + if has_tool_type(tools, "x_search") { + self.filter_x_search = true; + } + } + + pub(super) fn mark_injected_cache_tool(&mut self, kind: &str) { + self.injected_tool_types.insert(kind.to_owned()); + if kind == "x_search" { + self.filter_x_search = true; + } + } + + pub(crate) fn resequence_stream_value(&mut self, value: &mut Value) { + let Some(payload) = value.as_object_mut() else { + return; + }; + let Some(raw_sequence) = payload.get("sequence_number") else { + return; + }; + if self.stream_sequence_next.is_none() { + let Some(sequence) = exact_nonnegative_sequence(raw_sequence) else { + return; + }; + self.stream_sequence_next = Some(sequence); + } + let Some(next) = self.stream_sequence_next.as_mut() else { + return; + }; + payload.insert( + "sequence_number".to_owned(), + Value::Number(serde_json::Number::from(*next)), + ); + *next = next.saturating_add(1); + } + + pub(crate) fn rewrite_stream_event( + &mut self, + event_type: &str, + mut value: Value, + ) -> Result, GrokRequestEncodeError> { + if self.cache_filter_enabled() { + if value + .get("item") + .is_some_and(|item| self.is_internal_cache_call(item)) + { + self.record_dropped_cache_item(&value); + return Ok(Vec::new()); + } + self.filter_cache_envelope(&mut value)?; + if self.references_dropped_cache_item(&value) { + return Ok(Vec::new()); + } + self.compact_cache_output_index(&mut value); + } + if event_type == "response.output_item.added" + && let Some(item) = value.get("item").and_then(Value::as_object) + && let Some(primary) = self.remember_stream_call(item) + && self + .stream_calls + .get(&primary) + .is_some_and(|state| state.identity.kind == ToolKind::ApplyPatch) + { + if let Some(state) = self.stream_calls.get_mut(&primary) { + state.added_payload = value.as_object().cloned(); + } + return Ok(Vec::new()); + } + + if event_type == "response.function_call_arguments.delta" + && let Some(primary) = self.stream_call_key(&value) + { + let delta = value + .get("delta") + .and_then(Value::as_str) + .unwrap_or_default(); + let state = self + .stream_calls + .get_mut(&primary) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + match state.identity.kind { + ToolKind::Custom => { + state.arguments.push_str(delta); + state.last_delta = value.as_object().cloned().map(|mut payload| { + payload.remove("delta"); + payload + }); + return Ok(Vec::new()); + } + ToolKind::ToolSearch | ToolKind::ApplyPatch => return Ok(Vec::new()), + ToolKind::Function if state.schema.is_some() => { + let call_limit_exceeded = state.arguments.len().saturating_add(delta.len()) + > MAX_BUFFERED_FUNCTION_ARGUMENTS_BYTES; + let total_limit_exceeded = + self.stream_argument_bytes.saturating_add(delta.len()) + > MAX_TOTAL_BUFFERED_FUNCTION_ARGUMENTS_BYTES; + if !state.passthrough && (call_limit_exceeded || total_limit_exceeded) { + let buffered = std::mem::take(&mut state.arguments); + self.stream_argument_bytes = + self.stream_argument_bytes.saturating_sub(buffered.len()); + state.passthrough = true; + state.last_delta = None; + let mut output = Vec::with_capacity(2); + if !buffered.is_empty() { + let mut flushed = value.clone(); + flushed + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? + .insert("delta".to_owned(), Value::String(buffered)); + output.push(GrokTransformedWireEvent { + event_type: event_type.to_owned(), + value: flushed, + }); + } + output.push(GrokTransformedWireEvent { + event_type: event_type.to_owned(), + value, + }); + return Ok(output); + } + if !state.passthrough { + state.arguments.push_str(delta); + self.stream_argument_bytes = + self.stream_argument_bytes.saturating_add(delta.len()); + state.last_delta = value.as_object().cloned().map(|mut payload| { + payload.remove("delta"); + payload + }); + return Ok(Vec::new()); + } + } + ToolKind::Function => {} + } + } + + if event_type == "response.function_call_arguments.done" + && let Some(primary) = self.stream_call_key(&value) + { + let state = self + .stream_calls + .get(&primary) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + match state.identity.kind { + ToolKind::ToolSearch | ToolKind::ApplyPatch => return Ok(Vec::new()), + ToolKind::Custom => { + let arguments = value + .get("arguments") + .and_then(Value::as_str) + .filter(|arguments| !arguments.is_empty()) + .unwrap_or(&state.arguments); + let input = decode_custom_tool_input(&state.identity, arguments); + let mut output = Vec::with_capacity(2); + if let Some(delta) = state.last_delta.as_ref() { + output.push(GrokTransformedWireEvent { + event_type: "response.custom_tool_call_input.delta".to_owned(), + value: custom_tool_stream_payload( + delta, + "response.custom_tool_call_input.delta", + "delta", + &input, + ), + }); + } + output.push(GrokTransformedWireEvent { + event_type: "response.custom_tool_call_input.done".to_owned(), + value: custom_tool_stream_payload( + value + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?, + "response.custom_tool_call_input.done", + "input", + &input, + ), + }); + if let Some(state) = self.stream_calls.get_mut(&primary) { + state.arguments.clear(); + state.last_delta = None; + } + return Ok(output); + } + ToolKind::Function if state.schema.is_some() => { + let schema = state + .schema + .clone() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let passthrough = state.passthrough; + let buffered = state.arguments.clone(); + let last_delta = state.last_delta.clone(); + let arguments = value + .get("arguments") + .and_then(Value::as_str) + .filter(|arguments| !arguments.is_empty()) + .unwrap_or(&buffered); + let normalized = normalize_function_arguments(arguments, &schema) + .unwrap_or_else(|| arguments.to_owned()); + if passthrough { + value + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? + .insert("arguments".to_owned(), Value::String(normalized)); + return Ok(vec![GrokTransformedWireEvent { + event_type: event_type.to_owned(), + value, + }]); + } + + let mut output = Vec::with_capacity(2); + if let Some(mut delta) = last_delta { + delta.insert("delta".to_owned(), Value::String(normalized.clone())); + output.push(GrokTransformedWireEvent { + event_type: "response.function_call_arguments.delta".to_owned(), + value: Value::Object(delta), + }); + } + value + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? + .insert("arguments".to_owned(), Value::String(normalized)); + output.push(GrokTransformedWireEvent { + event_type: event_type.to_owned(), + value, + }); + if let Some(state) = self.stream_calls.get_mut(&primary) { + self.stream_argument_bytes = self + .stream_argument_bytes + .saturating_sub(state.arguments.len()); + state.arguments.clear(); + state.last_delta = None; + } + return Ok(output); + } + ToolKind::Function => {} + } + } + + if event_type == "response.output_item.done" + && let Some(item) = value.get("item").and_then(Value::as_object) + && self + .aliases + .get(string_field(item, "name")) + .is_some_and(|identity| identity.kind == ToolKind::ApplyPatch) + { + return self.rewrite_apply_patch_done_event(value); + } + + let completed_call = (event_type == "response.output_item.done") + .then(|| value.get("item").and_then(Value::as_object)) + .flatten() + .and_then(|item| self.primary_stream_call(item)); + self.rewrite_response_value(&mut value)?; + if let Some(response) = value.get_mut("response").and_then(Value::as_object_mut) { + self.restore_visible_tools(response); + } + if let Some(primary) = completed_call { + self.take_stream_call(&primary); + } + Ok(vec![GrokTransformedWireEvent { + event_type: event_type.to_owned(), + value, + }]) + } + + fn remember_stream_call(&mut self, item: &Map) -> Option { + if string_field(item, "type") != "function_call" { + return None; + } + let identity = self.aliases.get(string_field(item, "name"))?.clone(); + let keys = ["id", "call_id"] + .into_iter() + .filter_map(|field| item.get(field).and_then(Value::as_str)) + .filter(|key| !key.is_empty()) + .map(ToOwned::to_owned) + .collect::>(); + let primary = keys.first()?.clone(); + self.stream_calls.insert( + primary.clone(), + StreamCallState { + identity, + schema: self + .function_schemas + .get(string_field(item, "name")) + .cloned(), + arguments: String::new(), + passthrough: false, + last_delta: None, + added_payload: None, + }, + ); + for key in keys { + self.stream_keys.insert(key, primary.clone()); + } + Some(primary) + } + + fn stream_call_key(&mut self, payload: &Value) -> Option { + for field in ["item_id", "call_id"] { + if let Some(key) = payload.get(field).and_then(Value::as_str) + && let Some(primary) = self.stream_keys.get(key) + { + return Some(primary.clone()); + } + } + let alias = payload.get("name").and_then(Value::as_str)?; + let identity = self.aliases.get(alias)?.clone(); + let primary = ["item_id", "call_id"] + .into_iter() + .find_map(|field| payload.get(field).and_then(Value::as_str))? + .to_owned(); + self.stream_calls.insert( + primary.clone(), + StreamCallState { + identity, + schema: self.function_schemas.get(alias).cloned(), + arguments: String::new(), + passthrough: false, + last_delta: None, + added_payload: None, + }, + ); + self.stream_keys.insert(primary.clone(), primary.clone()); + Some(primary) + } + + fn primary_stream_call(&self, item: &Map) -> Option { + ["id", "call_id"] + .into_iter() + .filter_map(|field| item.get(field).and_then(Value::as_str)) + .find_map(|key| self.stream_keys.get(key).cloned()) + } + + fn take_stream_call(&mut self, primary: &str) -> Option { + self.stream_keys.retain(|_, owner| owner != primary); + let state = self.stream_calls.remove(primary)?; + self.stream_argument_bytes = self + .stream_argument_bytes + .saturating_sub(state.arguments.len()); + Some(state) + } + + fn rewrite_response_value(&self, value: &mut Value) -> Result<(), GrokRequestEncodeError> { + match value { + Value::Array(values) => { + for value in values { + self.rewrite_response_value(value)?; + } + } + Value::Object(object) => { + for value in object.values_mut() { + self.rewrite_response_value(value)?; + } + match string_field(object, "type") { + "function_call" => self.rewrite_function_call(object)?, + "shell_call" if self.legacy_local_shell => { + rewrite_legacy_local_shell_call(object) + } + _ => {} + } + } + _ => {} + } + Ok(()) + } + + fn rewrite_function_call( + &self, + call: &mut Map, + ) -> Result<(), GrokRequestEncodeError> { + let Some(identity) = self.aliases.get(string_field(call, "name")) else { + return Ok(()); + }; + match identity.kind { + ToolKind::Function => { + let alias = string_field(call, "name").to_owned(); + if let Some(schema) = self.function_schemas.get(&alias) + && let Some(arguments) = call.get("arguments").and_then(Value::as_str) + && let Some(normalized) = normalize_function_arguments(arguments, schema) + { + call.insert("arguments".to_owned(), Value::String(normalized)); + } + call.insert("name".to_owned(), Value::String(identity.name.clone())); + if identity.namespace.is_empty() { + call.remove("namespace"); + } else { + call.insert( + "namespace".to_owned(), + Value::String(identity.namespace.clone()), + ); + } + } + ToolKind::Custom => { + call.insert( + "type".to_owned(), + Value::String("custom_tool_call".to_owned()), + ); + call.insert("name".to_owned(), Value::String(identity.name.clone())); + if identity.namespace.is_empty() { + call.remove("namespace"); + } else { + call.insert( + "namespace".to_owned(), + Value::String(identity.namespace.clone()), + ); + } + let input = decode_custom_tool_input(identity, string_field(call, "arguments")); + call.insert("input".to_owned(), Value::String(input)); + call.remove("arguments"); + } + ToolKind::ToolSearch => { + call.insert( + "type".to_owned(), + Value::String("tool_search_call".to_owned()), + ); + call.insert("execution".to_owned(), Value::String("client".to_owned())); + let arguments = decode_tool_search_arguments(call.get("arguments")); + call.insert("arguments".to_owned(), arguments); + call.remove("name"); + call.remove("namespace"); + } + ToolKind::ApplyPatch => { + let operation = decode_apply_patch_arguments(call.get("arguments"))?; + call.insert( + "type".to_owned(), + Value::String("apply_patch_call".to_owned()), + ); + call.insert("operation".to_owned(), Value::Object(operation)); + call.remove("name"); + call.remove("namespace"); + call.remove("arguments"); + } + } + Ok(()) + } + + fn rewrite_apply_patch_done_event( + &mut self, + mut value: Value, + ) -> Result, GrokRequestEncodeError> { + let original_item = value + .get("item") + .and_then(Value::as_object) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + self.rewrite_response_value(&mut value)?; + let done = value + .as_object() + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let done_item = done + .get("item") + .and_then(Value::as_object) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let state = self + .primary_stream_call(&original_item) + .and_then(|primary| self.take_stream_call(&primary)); + let mut added = state + .and_then(|state| state.added_payload) + .unwrap_or_else(|| { + Map::from_iter([( + "type".to_owned(), + Value::String("response.output_item.added".to_owned()), + )]) + }); + added.insert( + "type".to_owned(), + Value::String("response.output_item.added".to_owned()), + ); + for key in ["output_index", "sequence_number"] { + if !added.contains_key(key) + && let Some(value) = done.get(key) + { + added.insert(key.to_owned(), value.clone()); + } + } + let mut added_item = done_item; + added_item.insert("status".to_owned(), Value::String("in_progress".to_owned())); + added.insert("item".to_owned(), Value::Object(added_item)); + Ok(vec![ + GrokTransformedWireEvent { + event_type: "response.output_item.added".to_owned(), + value: Value::Object(added), + }, + GrokTransformedWireEvent { + event_type: "response.output_item.done".to_owned(), + value: Value::Object(done), + }, + ]) + } + + fn cache_filter_enabled(&self) -> bool { + self.filter_x_search || !self.injected_tool_types.is_empty() + } + + fn is_internal_cache_call(&self, item: &Value) -> bool { + let Some(item) = item.as_object() else { + return false; + }; + let kind = string_field(item, "type").trim(); + if kind == "web_search_call" { + return self.injected_tool_types.contains("web_search"); + } + if !matches!(kind, "custom_tool_call" | "function_call") { + return false; + } + if string_field(item, "call_id").trim().starts_with("xs_call") { + return true; + } + let name = string_field(item, "name").trim(); + matches!( + name, + "x_user_search" | "x_semantic_search" | "x_keyword_search" | "x_thread_fetch" + ) && string_field(item, "namespace").trim().is_empty() + && !self.client_declared_tools.contains(name) + } + + fn record_dropped_cache_item(&mut self, payload: &Value) { + if let Some(index) = payload.get("output_index").and_then(Value::as_u64) { + self.dropped_output_indexes.insert(index); + } + let Some(item) = payload.get("item").and_then(Value::as_object) else { + return; + }; + for field in ["id", "call_id"] { + if let Some(value) = item + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + self.dropped_item_ids.insert(value.to_owned()); + } + } + } + + fn references_dropped_cache_item(&self, payload: &Value) -> bool { + if payload + .get("output_index") + .and_then(Value::as_u64) + .is_some_and(|index| self.dropped_output_indexes.contains(&index)) + { + return true; + } + ["item_id", "call_id"].into_iter().any(|field| { + payload + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| self.dropped_item_ids.contains(value)) + }) + } + + fn compact_cache_output_index(&self, payload: &mut Value) { + let Some(index) = payload.get("output_index").and_then(Value::as_u64) else { + return; + }; + let removed_before = self + .dropped_output_indexes + .iter() + .filter(|dropped| **dropped < index) + .count() as u64; + if removed_before > 0 + && let Some(payload) = payload.as_object_mut() + { + payload.insert( + "output_index".to_owned(), + Value::from(index.saturating_sub(removed_before)), + ); + } + } + + fn filter_cache_envelope(&self, payload: &mut Value) -> Result<(), GrokRequestEncodeError> { + let Some(payload) = payload.as_object_mut() else { + return Ok(()); + }; + self.filter_cache_output(payload)?; + self.filter_injected_cache_tools(payload)?; + if let Some(response) = payload.get_mut("response").and_then(Value::as_object_mut) { + self.filter_cache_output(response)?; + self.filter_injected_cache_tools(response)?; + } + Ok(()) + } + + fn filter_cache_output( + &self, + envelope: &mut Map, + ) -> Result<(), GrokRequestEncodeError> { + let Some(output) = envelope.get_mut("output") else { + return Ok(()); + }; + let output = output + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + output.retain(|item| !self.is_internal_cache_call(item)); + Ok(()) + } + + fn filter_injected_cache_tools( + &self, + envelope: &mut Map, + ) -> Result<(), GrokRequestEncodeError> { + if self.injected_tool_types.is_empty() { + return Ok(()); + } + let Some(tools) = envelope.get_mut("tools") else { + return Ok(()); + }; + let tools = tools + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + tools.retain(|tool| { + tool.get("type") + .and_then(Value::as_str) + .is_none_or(|kind| !self.injected_tool_types.contains(kind.trim())) + }); + if tools.is_empty() { + envelope.remove("tools"); + } + Ok(()) + } + + fn restore_visible_tools(&self, response: &mut Map) { + if response.contains_key("tools") { + if self.visible_tools.is_empty() { + response.remove("tools"); + } else { + response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone())); + } + } + } +} + +pub(super) fn custom_tool_stream_payload( + source: &Map, + kind: &str, + value_key: &str, + value: &str, +) -> Value { + let mut result = Map::from_iter([ + ("type".to_owned(), Value::String(kind.to_owned())), + (value_key.to_owned(), Value::String(value.to_owned())), + ]); + for key in ["item_id", "output_index", "sequence_number"] { + if let Some(value) = source.get(key) { + result.insert(key.to_owned(), value.clone()); + } + } + Value::Object(result) +} + +pub(super) fn decode_custom_tool_input(identity: &ToolIdentity, arguments: &str) -> String { + serde_json::from_str::(arguments) + .ok() + .and_then(|value| { + value + .get(identity.custom_argument_field()) + .and_then(Value::as_str) + .map(ToOwned::to_owned) + }) + .unwrap_or_else(|| arguments.to_owned()) +} + +pub(super) fn decode_tool_search_arguments(value: Option<&Value>) -> Value { + let Some(text) = value.and_then(Value::as_str) else { + return value.cloned().unwrap_or(Value::Object(Map::new())); + }; + if text.trim().is_empty() { + return Value::Object(Map::new()); + } + serde_json::from_str(text) + .unwrap_or_else(|_| json_object([("input", Value::String(text.to_owned()))])) +} + +pub(super) fn decode_apply_patch_arguments( + value: Option<&Value>, +) -> Result, GrokRequestEncodeError> { + let text = value + .and_then(Value::as_str) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let wrapper = serde_json::from_str::(text) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; + validate_apply_patch_operation(wrapper.get("operation")) +} + +pub(super) fn rewrite_legacy_local_shell_call(call: &mut Map) { + let commands = call + .get("action") + .and_then(Value::as_object) + .and_then(|action| action.get("commands")) + .and_then(Value::as_array) + .map(|commands| { + commands + .iter() + .filter_map(Value::as_str) + .collect::>() + .join("\n") + }) + .unwrap_or_default(); + call.insert( + "type".to_owned(), + Value::String("local_shell_call".to_owned()), + ); + call.insert( + "action".to_owned(), + json_object([ + ("type", Value::String("exec".to_owned())), + ("command", Value::String(commands)), + ]), + ); + call.remove("max_output_length"); +} + +impl fmt::Debug for GrokResponseTransform { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("GrokResponseTransform") + .field("alias_count", &self.aliases.len()) + .field("function_schema_count", &self.function_schemas.len()) + .field("visible_tool_count", &self.visible_tools.len()) + .field("legacy_local_shell", &self.legacy_local_shell) + .field("filter_x_search", &self.filter_x_search) + .field("injected_tool_types", &self.injected_tool_types) + .finish() + } +} diff --git a/backend/crates/providers/xai/src/transport/request/tools.rs b/backend/crates/providers/xai/src/transport/request/tools.rs new file mode 100644 index 000000000..9c9485436 --- /dev/null +++ b/backend/crates/providers/xai/src/transport/request/tools.rs @@ -0,0 +1,1752 @@ +//! Codex/OpenAI 工具声明与历史调用到 Grok 工具合同的归一化。 + +use super::response::{ToolIdentity, ToolKind}; +use super::*; + +pub(super) struct ToolNormalizer { + pub(super) response: GrokResponseTransform, + identity_aliases: BTreeMap, + deferred_surfaces: Vec, + client_search_tool: Option>, + server_search_eager: bool, + native_shell: bool, +} + +pub(super) struct NormalizedInputItems { + items: Vec, + loaded_tools: Vec, + visible_tools: Vec, +} + +pub(super) struct NormalizedToolSearchOutput { + history: Map, + loaded_tools: Vec, + visible_tools: Vec, +} + +impl ToolNormalizer { + pub(super) fn new() -> Self { + Self { + response: GrokResponseTransform::default(), + identity_aliases: BTreeMap::new(), + deferred_surfaces: Vec::new(), + client_search_tool: None, + server_search_eager: false, + native_shell: false, + } + } + + pub(super) fn for_replay(response: GrokResponseTransform) -> Self { + let identity_aliases = response + .aliases + .iter() + .map(|(alias, identity)| (identity.clone(), alias.clone())) + .collect(); + Self { + response, + identity_aliases, + deferred_surfaces: Vec::new(), + client_search_tool: None, + server_search_eager: false, + native_shell: false, + } + } + + pub(super) fn normalize( + mut self, + payload: &mut Map, + ) -> Result { + let (tools, had_tools) = + optional_array(payload.get("tools")).map_err(|error| error.at_field("tools"))?; + if had_tools { + self.response.visible_tools.clone_from(&tools); + } + let client_search = inspect_tool_search(&tools).map_err(|error| error.at_field("tools"))?; + self.normalize_client_search_parallel(payload, client_search) + .map_err(|error| error.at_field("parallel_tool_calls"))?; + + let mut normalized_tools = Vec::with_capacity(tools.len()); + for raw_tool in &tools { + normalized_tools.extend( + self.normalize_tool(raw_tool, "", client_search, false) + .map_err(|error| error.at_field("tools"))?, + ); + } + + if let Some(Value::Array(items)) = payload.get("input") { + let normalized = self + .normalize_input_items(items) + .map_err(|error| error.at_field("input"))?; + normalized_tools.extend(normalized.loaded_tools); + self.response.visible_tools.extend(normalized.visible_tools); + payload.insert("input".to_owned(), Value::Array(normalized.items)); + } else if payload + .get("input") + .is_some_and(|input| !input.is_null() && !input.is_string()) + { + return Err(GrokRequestEncodeError::InvalidRequestField { field: "input" }); + } + + if self.client_search_tool.is_some() { + normalized_tools.push(Value::Object( + self.build_client_search_function() + .map_err(|error| error.at_field("tools"))?, + )); + } + normalized_tools = dedupe_normalized_tools(normalized_tools); + if normalized_tools.is_empty() { + if had_tools { + payload.remove("tools"); + payload.remove("parallel_tool_calls"); + } + } else { + payload.insert("tools".to_owned(), Value::Array(normalized_tools.clone())); + } + self.normalize_tool_choice(payload, &normalized_tools) + .map_err(|error| error.at_field("tool_choice"))?; + Ok(self.response) + } + + fn alias(&mut self, identity: ToolIdentity) -> String { + if let Some(alias) = self.identity_aliases.get(&identity) { + return alias.clone(); + } + let base = match identity.kind { + ToolKind::ToolSearch => "xai_proxy_tool_search".to_owned(), + ToolKind::ApplyPatch => "xai_proxy_apply_patch".to_owned(), + ToolKind::Function | ToolKind::Custom if !identity.namespace.is_empty() => { + format!("{}__{}", identity.namespace, identity.name) + } + ToolKind::Function | ToolKind::Custom => identity.name.clone(), + }; + let key = format!( + "{}\0{}\0{}", + identity.kind as u8, identity.namespace, identity.name + ); + let mut alias = truncate_tool_alias(&base, &key); + if self + .response + .aliases + .get(&alias) + .is_some_and(|existing| existing != &identity) + { + alias = hashed_tool_alias(&base, &key); + } + self.response + .aliases + .insert(alias.clone(), identity.clone()); + self.identity_aliases.insert(identity, alias.clone()); + alias + } + + fn normalize_client_search_parallel( + &self, + payload: &mut Map, + client_search: bool, + ) -> Result<(), GrokRequestEncodeError> { + if !client_search { + return Ok(()); + } + match payload.get("parallel_tool_calls") { + None | Some(Value::Null) => { + payload.insert("parallel_tool_calls".to_owned(), Value::Bool(false)); + } + Some(Value::Bool(true)) => { + payload.insert("parallel_tool_calls".to_owned(), Value::Bool(false)); + } + Some(Value::Bool(false)) => {} + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + Ok(()) + } + + fn normalize_tool( + &mut self, + raw: &Value, + namespace: &str, + client_search: bool, + force: bool, + ) -> Result, GrokRequestEncodeError> { + let tool = raw + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let kind = string_field(tool, "type").trim(); + match kind { + "function" => self.normalize_function_tool(tool, namespace, client_search, force), + "namespace" => self.normalize_namespace_tool(tool, client_search, force), + "tool_search" => self.normalize_tool_search(tool, force), + "custom" => self.normalize_custom_tool(tool, namespace), + "web_search" + | "web_search_preview" + | "web_search_preview_2025_03_11" + | "web_search_2025_08_26" => self.normalize_web_search_tool(tool, kind), + "mcp" => self.normalize_mcp_tool(tool, client_search, force), + "shell" => self.normalize_shell_tool(tool), + "local_shell" => self.normalize_legacy_local_shell_tool(tool), + "apply_patch" => self.normalize_apply_patch_tool(tool), + "x_search" | "collections_search" | "file_search" | "code_execution" + | "code_interpreter" => Ok(vec![Value::Object(without_defer_loading(tool))]), + // 对未获 xAI 接受的工具做静默过滤，并在过滤后同步收敛 tool_choice； + // 不要把客户端可选扩展升级成整单 400。 + "" | "computer_use_preview" | "image_generation" => Ok(Vec::new()), + _ => Ok(Vec::new()), + } + } + + fn normalize_function_tool( + &mut self, + tool: &Map, + namespace: &str, + client_search: bool, + force: bool, + ) -> Result, GrokRequestEncodeError> { + let name = string_field(tool, "name").trim(); + if name.is_empty() { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let deferred = tool + .get("defer_loading") + .and_then(Value::as_bool) + .unwrap_or(false); + if deferred && client_search && !force { + if namespace.is_empty() { + self.deferred_surfaces.push(describe_deferred_tool( + name, + string_field(tool, "description"), + )); + } + return Ok(Vec::new()); + } + let mut converted = without_defer_loading(tool); + let function_schema = tool + .get("parameters") + .filter(|schema| schema_contains_integer(schema, 0)) + .cloned(); + match converted.get("parameters").cloned() { + None | Some(Value::Null) => { + converted.insert( + "parameters".to_owned(), + json_object([ + ("type", Value::String("object".to_owned())), + ("properties", Value::Object(Map::new())), + ]), + ); + } + Some(parameters) => { + if let Some(normalized) = normalize_function_parameters_root(¶meters)? { + converted.insert("parameters".to_owned(), normalized); + } + } + } + let alias = self.alias(ToolIdentity::new(ToolKind::Function, namespace, name)); + if let Some(schema) = function_schema { + self.response.function_schemas.insert(alias.clone(), schema); + } + converted.insert("name".to_owned(), Value::String(alias)); + Ok(vec![Value::Object(converted)]) + } + + fn normalize_namespace_tool( + &mut self, + tool: &Map, + client_search: bool, + force: bool, + ) -> Result, GrokRequestEncodeError> { + let name = string_field(tool, "name").trim(); + let children = tool + .get("tools") + .and_then(Value::as_array) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if name.is_empty() { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + if client_search && !force && namespace_has_deferred_functions(children) { + self.deferred_surfaces.push(describe_deferred_tool( + name, + string_field(tool, "description"), + )); + } + let mut converted = Vec::new(); + for child in children { + if child.pointer("/type").and_then(Value::as_str) != Some("function") { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + converted.extend(self.normalize_tool(child, name, client_search, force)?); + } + Ok(converted) + } + + fn normalize_tool_search( + &mut self, + tool: &Map, + force: bool, + ) -> Result, GrokRequestEncodeError> { + if force { + return Ok(Vec::new()); + } + match string_field(tool, "execution") + .trim() + .to_ascii_lowercase() + .as_str() + { + "" | "server" => { + self.server_search_eager = true; + } + "client" => { + self.client_search_tool = Some(tool.clone()); + } + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + Ok(Vec::new()) + } + + fn normalize_custom_tool( + &mut self, + tool: &Map, + namespace: &str, + ) -> Result, GrokRequestEncodeError> { + let name = string_field(tool, "name").trim(); + if name.is_empty() || tool.get("format").is_some_and(|value| !value.is_object()) { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let identity = ToolIdentity::new(ToolKind::Custom, namespace, name); + let argument_field = identity.custom_argument_field(); + let description = if identity.is_root_custom_apply_patch() { + "The apply_patch tool edits files using Codex patch format. Provide the complete raw patch text in the patch string field.".to_owned() + } else { + let mut description = string_field(tool, "description").trim().to_owned(); + if !description.is_empty() { + description.push('\n'); + } + description.push_str("Provide the custom tool input in the input string field."); + description + }; + let alias = self.alias(identity); + Ok(vec![json_object([ + ("type", Value::String("function".to_owned())), + ("name", Value::String(alias)), + ("description", Value::String(description)), + ( + "parameters", + json_object([ + ("type", Value::String("object".to_owned())), + ( + "properties", + json_object([( + argument_field, + json_object([("type", Value::String("string".to_owned()))]), + )]), + ), + ( + "required", + Value::Array(vec![Value::String(argument_field.to_owned())]), + ), + ("additionalProperties", Value::Bool(false)), + ]), + ), + ])]) + } + + fn build_client_search_function( + &mut self, + ) -> Result, GrokRequestEncodeError> { + let tool = self + .client_search_tool + .as_ref() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let mut description = string_field(tool, "description").trim().to_owned(); + if description.is_empty() { + description.push_str("Search for tools needed to continue the task."); + } + if !self.deferred_surfaces.is_empty() { + description.push_str("\nDeferred tool surfaces available to search:\n- "); + description.push_str(&self.deferred_surfaces.join("\n- ")); + } + description.truncate(description.len().min(MAX_TOOL_SEARCH_DESCRIPTION_BYTES)); + let parameters = match tool.get("parameters") { + None => json_object([ + ("type", Value::String("object".to_owned())), + ("properties", Value::Object(Map::new())), + ("additionalProperties", Value::Bool(true)), + ]), + Some(Value::Object(parameters)) => Value::Object(parameters.clone()), + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); + Ok(Map::from_iter([ + ("type".to_owned(), Value::String("function".to_owned())), + ("name".to_owned(), Value::String(alias)), + ("description".to_owned(), Value::String(description)), + ("parameters".to_owned(), parameters), + ])) + } +} + +pub(super) fn optional_array( + value: Option<&Value>, +) -> Result<(Vec, bool), GrokRequestEncodeError> { + match value { + None | Some(Value::Null) => Ok((Vec::new(), false)), + Some(Value::Array(values)) => Ok((values.clone(), true)), + Some(_) => Err(GrokRequestEncodeError::InvalidRequestNormalization), + } +} + +/// Grok Build 要求 function parameters 的根节点必为 object；Codex 会为可选参数 +/// 生成 `object | null`，这里只移除根节点 nullability，嵌套字段保持原样。 +pub(super) fn normalize_function_parameters_root( + value: &Value, +) -> Result, GrokRequestEncodeError> { + let Some(schema) = value.as_object() else { + return Ok(None); + }; + let mut normalized = schema.clone(); + let mut changed = false; + + if let Some(types) = normalized.get("type").and_then(Value::as_array).cloned() { + let removed_null = types.iter().any(|value| value.as_str() == Some("null")); + if removed_null { + let remaining = types + .into_iter() + .filter(|value| value.as_str() != Some("null")) + .collect::>(); + if remaining.len() != 1 || remaining[0].as_str() != Some("object") { + return Err(invalid_function_parameters_root()); + } + normalized.insert("type".to_owned(), Value::String("object".to_owned())); + changed = true; + } + } + + for keyword in ["anyOf", "oneOf"] { + let Some(branches) = normalized.get(keyword).and_then(Value::as_array).cloned() else { + continue; + }; + let removed_null = branches.iter().any(is_null_only_schema); + if !removed_null { + continue; + } + let remaining = branches + .into_iter() + .filter(|branch| !is_null_only_schema(branch)) + .collect::>(); + if remaining.is_empty() + || remaining.iter().any(|branch| { + branch.as_object().is_none_or(|branch| { + !is_object_root_schema(branch, &normalized, &mut BTreeSet::new()) + }) + }) + { + return Err(invalid_function_parameters_root()); + } + if remaining.len() == 1 && normalized.len() == 1 { + normalized = remaining[0] + .as_object() + .cloned() + .ok_or_else(invalid_function_parameters_root)?; + normalized.insert("type".to_owned(), Value::String("object".to_owned())); + } else { + normalized.insert(keyword.to_owned(), Value::Array(remaining)); + normalized.insert("type".to_owned(), Value::String("object".to_owned())); + } + changed = true; + } + + Ok(changed.then_some(Value::Object(normalized))) +} + +pub(super) fn is_null_only_schema(value: &Value) -> bool { + let Some(schema) = value.as_object() else { + return false; + }; + match schema.get("type") { + Some(Value::String(kind)) => kind == "null", + Some(Value::Array(types)) if !types.is_empty() => { + types.iter().all(|kind| kind.as_str() == Some("null")) + } + _ => false, + } +} + +pub(super) fn is_object_root_schema( + schema: &Map, + root: &Map, + visited: &mut BTreeSet, +) -> bool { + match schema.get("type") { + Some(Value::String(kind)) => return kind == "object", + Some(Value::Array(types)) if !types.is_empty() => { + return types.iter().all(|kind| kind.as_str() == Some("object")); + } + Some(_) => return false, + None => {} + } + if schema.contains_key("properties") { + return true; + } + let Some(reference) = schema.get("$ref").and_then(Value::as_str) else { + return false; + }; + if !visited.insert(reference.to_owned()) { + return false; + } + resolve_local_schema_ref(root, reference) + .is_some_and(|resolved| is_object_root_schema(resolved, root, visited)) +} + +pub(super) fn resolve_local_schema_ref<'a>( + root: &'a Map, + reference: &str, +) -> Option<&'a Map> { + if reference == "#" { + return Some(root); + } + let path = reference.strip_prefix("#/")?; + let mut segments = path.split('/'); + let first = decode_json_pointer_segment(segments.next()?); + let mut current = root.get(&first)?; + for segment in segments { + let segment = decode_json_pointer_segment(segment); + current = current.as_object()?.get(&segment)?; + } + current.as_object() +} + +pub(super) fn decode_json_pointer_segment(segment: &str) -> String { + segment.replace("~1", "/").replace("~0", "~") +} + +pub(super) const fn invalid_function_parameters_root() -> GrokRequestEncodeError { + GrokRequestEncodeError::InvalidRequestField { + field: "tools[].parameters", + } +} + +pub(super) fn inspect_tool_search(tools: &[Value]) -> Result { + let mut client_search = false; + let mut server_search = false; + for raw_tool in tools { + let tool = raw_tool + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if string_field(tool, "type") != "tool_search" { + continue; + } + match string_field(tool, "execution") + .trim() + .to_ascii_lowercase() + .as_str() + { + "" | "server" if !client_search => server_search = true, + "client" if !client_search && !server_search => client_search = true, + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + } + Ok(client_search) +} + +pub(super) fn string_field<'a>(value: &'a Map, key: &str) -> &'a str { + value.get(key).and_then(Value::as_str).unwrap_or_default() +} + +pub(super) fn json_object(entries: [(&str, Value); N]) -> Value { + Value::Object(Map::from_iter( + entries + .into_iter() + .map(|(key, value)| (key.to_owned(), value)), + )) +} + +pub(super) fn without_defer_loading(tool: &Map) -> Map { + let mut converted = tool.clone(); + converted.remove("defer_loading"); + converted +} + +pub(super) fn namespace_has_deferred_functions(children: &[Value]) -> bool { + children.iter().any(|child| { + child.pointer("/type").and_then(Value::as_str) == Some("function") + && child + .pointer("/defer_loading") + .and_then(Value::as_bool) + .unwrap_or(false) + }) +} + +pub(super) fn describe_deferred_tool(name: &str, description: &str) -> String { + let description = description.trim(); + if description.is_empty() { + return name.to_owned(); + } + let description = description.chars().take(240).collect::(); + format!("{name}: {description}") +} + +pub(super) fn truncate_tool_alias(base: &str, key: &str) -> String { + if base.len() <= MAX_BUILD_TOOL_ALIAS_LENGTH { + base.to_owned() + } else { + hashed_tool_alias(base, key) + } +} + +pub(super) fn hashed_tool_alias(base: &str, key: &str) -> String { + let suffix = format!("__{}", short_tool_hash(key)); + let limit = MAX_BUILD_TOOL_ALIAS_LENGTH.saturating_sub(suffix.len()); + let mut end = limit.min(base.len()); + while !base.is_char_boundary(end) { + end = end.saturating_sub(1); + } + format!("{}{suffix}", &base[..end]) +} + +pub(super) fn short_tool_hash(value: &str) -> String { + Sha256::digest(value) + .iter() + .take(5) + .map(|byte| format!("{byte:02x}")) + .collect::()[..9] + .to_owned() +} + +pub(super) fn dedupe_normalized_tools(tools: Vec) -> Vec { + let mut result = Vec::with_capacity(tools.len()); + let mut seen = BTreeSet::new(); + for tool in tools { + if normalized_tool_dedupe_key(&tool).is_none_or(|key| seen.insert(key)) { + result.push(tool); + } + } + result +} + +pub(super) fn normalized_tool_dedupe_key(tool: &Value) -> Option { + let object = tool.as_object()?; + let kind = string_field(object, "type").trim(); + if !kind.is_empty() { + if let Some(name) = object + .get("name") + .and_then(Value::as_str) + .map(str::trim) + .filter(|name| !name.is_empty()) + { + return Some(format!("type:{kind}\0name:{name}")); + } + if kind == "mcp" + && let Some(label) = object + .get("server_label") + .and_then(Value::as_str) + .map(str::trim) + .filter(|label| !label.is_empty()) + { + return Some(format!("type:mcp\0server_label:{label}")); + } + } + serde_json::to_string(tool) + .ok() + .map(|encoded| format!("json:{encoded}")) +} + +impl ToolNormalizer { + fn normalize_web_search_tool( + &mut self, + tool: &Map, + kind: &str, + ) -> Result, GrokRequestEncodeError> { + let filters = normalize_web_search_filters(tool)?; + if let Some(content_types) = tool.get("search_content_types") { + content_types + .as_array() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + } + if kind == "web_search" && tool.len() == 1 { + return Ok(vec![Value::Object(tool.clone())]); + } + let mut converted = + Map::from_iter([("type".to_owned(), Value::String("web_search".to_owned()))]); + if let Some(filters) = filters { + converted.insert("filters".to_owned(), Value::Object(filters)); + } + Ok(vec![Value::Object(converted)]) + } + + fn normalize_mcp_tool( + &mut self, + tool: &Map, + client_search: bool, + force: bool, + ) -> Result, GrokRequestEncodeError> { + let deferred = tool + .get("defer_loading") + .and_then(Value::as_bool) + .unwrap_or(false); + if deferred && client_search && !force { + let label = ["server_label", "name"] + .into_iter() + .find_map(|field| tool.get(field).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + self.deferred_surfaces.push(describe_deferred_tool( + label, + string_field(tool, "description"), + )); + return Ok(Vec::new()); + } + Ok(vec![Value::Object(without_defer_loading(tool))]) + } + + fn normalize_shell_tool( + &mut self, + tool: &Map, + ) -> Result, GrokRequestEncodeError> { + if self.response.legacy_local_shell { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + self.native_shell = true; + Ok(vec![Value::Object(without_defer_loading(tool))]) + } + + fn normalize_legacy_local_shell_tool( + &mut self, + _tool: &Map, + ) -> Result, GrokRequestEncodeError> { + if self.native_shell || self.response.legacy_local_shell { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + self.response.legacy_local_shell = true; + Ok(vec![json_object([ + ("type", Value::String("shell".to_owned())), + ( + "environment", + json_object([("type", Value::String("local".to_owned()))]), + ), + ])]) + } + + fn normalize_apply_patch_tool( + &mut self, + _tool: &Map, + ) -> Result, GrokRequestEncodeError> { + let alias = self.alias(ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")); + let operation = json_object([ + ("type", Value::String("object".to_owned())), + ( + "properties", + json_object([ + ( + "type", + json_object([ + ("type", Value::String("string".to_owned())), + ( + "enum", + Value::Array( + ["create_file", "update_file", "delete_file"] + .into_iter() + .map(|value| Value::String(value.to_owned())) + .collect(), + ), + ), + ]), + ), + ( + "path", + json_object([ + ("type", Value::String("string".to_owned())), + ("minLength", Value::from(1)), + ]), + ), + ( + "diff", + json_object([("type", Value::String("string".to_owned()))]), + ), + ]), + ), + ( + "required", + Value::Array( + ["type", "path"] + .into_iter() + .map(|value| Value::String(value.to_owned())) + .collect(), + ), + ), + ("additionalProperties", Value::Bool(false)), + ]); + Ok(vec![json_object([ + ("type", Value::String("function".to_owned())), + ("name", Value::String(alias)), + ( + "description", + Value::String( + "Create, update, or delete one file using a structured V4A patch operation. create_file and update_file require path and diff; delete_file requires path." + .to_owned(), + ), + ), + ( + "parameters", + json_object([ + ("type", Value::String("object".to_owned())), + ("properties", json_object([("operation", operation)])), + ( + "required", + Value::Array(vec![Value::String("operation".to_owned())]), + ), + ("additionalProperties", Value::Bool(false)), + ]), + ), + ("strict", Value::Bool(true)), + ])]) + } + + fn normalize_tool_choice( + &mut self, + payload: &mut Map, + normalized_tools: &[Value], + ) -> Result<(), GrokRequestEncodeError> { + let Some(choice) = payload.get("tool_choice").cloned() else { + return Ok(()); + }; + if choice.is_null() { + return Ok(()); + } + if normalized_tools.is_empty() { + payload.remove("tool_choice"); + return Ok(()); + } + let Some(mut object) = choice.as_object().cloned() else { + return choice + .as_str() + .filter(|value| matches!(*value, "none" | "auto" | "required")) + .map(|_| ()) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization); + }; + let kind = string_field(&object, "type").to_owned(); + match kind.as_str() { + "tool_search" => { + if self.client_search_tool.is_none() { + if self.server_search_eager { + payload.insert("tool_choice".to_owned(), Value::String("auto".to_owned())); + return Ok(()); + } + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); + payload.insert( + "tool_choice".to_owned(), + json_object([ + ("type", Value::String("function".to_owned())), + ("name", Value::String(alias)), + ]), + ); + } + "custom" => { + let identity = ToolIdentity::new( + ToolKind::Custom, + string_field(&object, "namespace").trim(), + string_field(&object, "name").trim(), + ); + let alias = self + .identity_aliases + .get(&identity) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + object.insert("type".to_owned(), Value::String("function".to_owned())); + object.insert("name".to_owned(), Value::String(alias)); + object.remove("namespace"); + payload.insert("tool_choice".to_owned(), Value::Object(object)); + } + "apply_patch" => { + let identity = ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch"); + let alias = self + .identity_aliases + .get(&identity) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + payload.insert( + "tool_choice".to_owned(), + json_object([ + ("type", Value::String("function".to_owned())), + ("name", Value::String(alias)), + ]), + ); + } + "allowed_tools" => { + payload.remove("tool_choice"); + } + "function" => { + self.normalize_function_tool_choice(payload, object, normalized_tools)?; + } + _ => { + if let Some(hosted_kind) = normalize_hosted_tool_choice_kind(&kind) { + let matching = tools_of_type(normalized_tools, hosted_kind); + if matching.is_empty() { + payload.remove("tool_choice"); + return Ok(()); + } + object.insert("type".to_owned(), Value::String(hosted_kind.to_owned())); + payload.insert("tool_choice".to_owned(), Value::Object(object)); + } else { + payload.remove("tool_choice"); + } + } + } + Ok(()) + } + + fn normalize_function_tool_choice( + &self, + payload: &mut Map, + mut object: Map, + normalized_tools: &[Value], + ) -> Result<(), GrokRequestEncodeError> { + if let Some(Value::Object(function)) = object.get_mut("function") { + rewrite_namespace_choice(function, &self.identity_aliases)?; + let name = string_field(function, "name").trim(); + if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { + payload.remove("tool_choice"); + return Ok(()); + } + payload.insert("tool_choice".to_owned(), Value::Object(object)); + return Ok(()); + } + rewrite_namespace_choice(&mut object, &self.identity_aliases)?; + let name = string_field(&object, "name").trim(); + if !name.is_empty() && !has_named_tool(normalized_tools, "function", name) { + payload.remove("tool_choice"); + return Ok(()); + } + payload.insert("tool_choice".to_owned(), Value::Object(object)); + Ok(()) + } +} + +pub(super) fn normalize_web_search_filters( + tool: &Map, +) -> Result>, GrokRequestEncodeError> { + let nested = match tool.get("filters") { + None | Some(Value::Null) => None, + Some(Value::Object(filters)) => filters + .get("allowed_domains") + .map(normalize_allowed_domains) + .transpose()?, + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let top_level = tool + .get("allowed_domains") + .map(normalize_allowed_domains) + .transpose()?; + if nested.is_some() && top_level.is_some() && nested != top_level { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + Ok(nested + .or(top_level) + .map(|domains| Map::from_iter([("allowed_domains".to_owned(), Value::Array(domains))]))) +} + +pub(super) fn normalize_allowed_domains( + value: &Value, +) -> Result, GrokRequestEncodeError> { + let domains = value + .as_array() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if domains + .iter() + .any(|domain| domain.as_str().map(str::trim).is_none_or(str::is_empty)) + { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + Ok(domains.clone()) +} + +pub(super) fn normalize_hosted_tool_choice_kind(kind: &str) -> Option<&'static str> { + match kind { + "web_search" + | "web_search_preview" + | "web_search_preview_2025_03_11" + | "web_search_2025_08_26" => Some("web_search"), + "x_search" => Some("x_search"), + "image_generation" => Some("image_generation"), + "collections_search" => Some("collections_search"), + "file_search" => Some("file_search"), + "code_execution" => Some("code_execution"), + "code_interpreter" => Some("code_interpreter"), + "mcp" => Some("mcp"), + "shell" | "local_shell" => Some("shell"), + _ => None, + } +} + +pub(super) fn has_tool_type(tools: &[Value], kind: &str) -> bool { + tools + .iter() + .any(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind)) +} + +pub(super) fn has_named_tool(tools: &[Value], kind: &str, name: &str) -> bool { + tools.iter().any(|tool| { + tool.pointer("/type").and_then(Value::as_str) == Some(kind) + && tool.pointer("/name").and_then(Value::as_str) == Some(name) + }) +} + +pub(super) fn tools_of_type(tools: &[Value], kind: &str) -> Vec { + tools + .iter() + .filter(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind)) + .cloned() + .collect() +} + +pub(super) fn rewrite_namespace_choice( + object: &mut Map, + aliases: &BTreeMap, +) -> Result<(), GrokRequestEncodeError> { + let name = string_field(object, "name").trim(); + let namespace = string_field(object, "namespace").trim(); + if name.is_empty() || namespace.is_empty() { + return Ok(()); + } + let identity = ToolIdentity::new(ToolKind::Function, namespace, name); + let alias = aliases + .get(&identity) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + object.insert("name".to_owned(), Value::String(alias)); + object.remove("namespace"); + Ok(()) +} + +impl ToolNormalizer { + fn normalize_input_items( + &mut self, + items: &[Value], + ) -> Result { + let mut rewritten = Vec::with_capacity(items.len()); + let mut loaded_tools = Vec::new(); + let mut visible_tools = Vec::new(); + for raw_item in items { + let Some(item) = raw_item.as_object() else { + rewritten.push(raw_item.clone()); + continue; + }; + let mut item_type = string_field(item, "type").trim(); + if item_type.is_empty() && !string_field(item, "role").trim().is_empty() { + item_type = "message"; + } + match item_type { + "message" => rewritten.push(Value::Object(self.normalize_message_input(item)?)), + "function_call" => { + rewritten.push(Value::Object(self.normalize_function_call_input(item)?)); + } + "function_call_output" => rewritten.push(Value::Object( + self.normalize_function_call_output_input(item, true)?, + )), + "reasoning" => rewritten.push(Value::Object(sanitize_reasoning_input(item))), + "file_search_call" + | "web_search_call" + | "image_generation_call" + | "code_interpreter_call" + | "shell_call" + | "mcp_list_tools" + | "mcp_approval_request" + | "mcp_approval_response" + | "mcp_call" => { + rewritten.push(Value::Object(sanitize_native_history_input( + item, item_type, + ))); + } + "compaction" | "compaction_summary" => { + rewritten.extend(normalize_compaction_input(item)?); + } + "tool_search_call" => { + rewritten.push(Value::Object(self.normalize_tool_search_call(item)?)); + } + "tool_search_output" => { + let normalized = self.normalize_tool_search_output(item)?; + rewritten.push(Value::Object(normalized.history)); + loaded_tools.extend(normalized.loaded_tools); + visible_tools.extend(normalized.visible_tools); + } + "custom_tool_call" => { + rewritten.push(Value::Object(self.normalize_custom_tool_call_input(item)?)); + } + "custom_tool_call_output" => rewritten.push(Value::Object( + self.normalize_function_call_output_input(item, false)?, + )), + "apply_patch_call" => { + rewritten.push(Value::Object(self.normalize_apply_patch_call_input(item)?)); + } + "apply_patch_call_output" => { + rewritten.push(Value::Object(normalize_apply_patch_output_input(item)?)); + } + "agent_message" => rewritten.push(normalize_agent_message_input(item)), + "local_shell_call" => { + rewritten.push(Value::Object(normalize_legacy_local_shell_call_input( + item, + )?)); + } + "local_shell_call_output" => rewritten.push(Value::Object( + normalize_legacy_local_shell_output_input(item)?, + )), + "shell_call_output" => { + rewritten.push(Value::Object(normalize_shell_call_output_input(item)?)); + } + "mcp_tool_call_output" => rewritten.push(normalize_mcp_output_input(item)?), + "compaction_trigger" => { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + "additional_tools" => { + let (tools, visible) = self.normalize_additional_tools_input(item)?; + loaded_tools.extend(tools); + visible_tools.extend(visible); + } + "" => rewritten.push(raw_item.clone()), + unsupported => { + rewritten.push(unsupported_input_history_boundary(item, unsupported)) + } + } + } + Ok(NormalizedInputItems { + items: rewritten, + loaded_tools, + visible_tools, + }) + } + + fn normalize_message_input( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let role = match string_field(item, "role").trim() { + "" => "assistant", + "model" => "assistant", + role => role, + }; + let content = self.normalize_message_content(item.get("content"), role)?; + // 透明代理合法形态：只覆盖归一化字段并剥离 Grok 注入键，未知官方 + // 字段原样保留。 + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.insert("type".to_owned(), Value::String("message".to_owned())); + converted.insert("role".to_owned(), Value::String(role.to_owned())); + converted.insert("content".to_owned(), content); + Ok(converted) + } + + fn normalize_message_content( + &mut self, + value: Option<&Value>, + role: &str, + ) -> Result { + if let Some(Value::String(text)) = value { + return Ok(Value::String(text.clone())); + } + let items = value + .and_then(Value::as_array) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if role == "assistant" { + let texts = items + .iter() + .map(|raw| { + let item = raw.as_object()?; + match string_field(item, "type") { + "text" | "input_text" | "output_text" => { + item.get("text").and_then(Value::as_str) + } + "refusal" => item.get("refusal").and_then(Value::as_str), + _ => None, + } + }) + .collect::>>(); + if let Some(texts) = texts { + return Ok(Value::String(texts.join("\n"))); + } + } + let text_part_type = if role == "assistant" { + "output_text" + } else { + "input_text" + }; + let mut normalized = Vec::with_capacity(items.len()); + for raw in items { + let item = raw + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let converted = match string_field(item, "type") { + "text" | "input_text" | "output_text" => json_object([ + ("type", Value::String(text_part_type.to_owned())), + ("text", Value::String(string_field(item, "text").to_owned())), + ]), + "refusal" => json_object([ + ("type", Value::String(text_part_type.to_owned())), + ( + "text", + Value::String(string_field(item, "refusal").to_owned()), + ), + ]), + "input_image" => Value::Object(self.normalize_input_image_part(item)?), + "input_file" => Value::Object(normalize_input_file_part(item)), + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + normalized.push(converted); + } + Ok(Value::Array(normalized)) + } + + fn normalize_input_image_part( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let detail = match item.get("detail") { + None | Some(Value::Null) => "auto", + Some(Value::String(detail)) if detail.trim().is_empty() => "auto", + Some(Value::String(detail)) => detail.trim(), + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let detail = match detail { + "auto" | "low" | "high" => detail, + "original" => "high", + _ => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + }; + let mut converted = Map::from_iter([ + ("type".to_owned(), Value::String("input_image".to_owned())), + ("detail".to_owned(), Value::String(detail.to_owned())), + ]); + if let Some(value) = item.get("image_url").or_else(|| item.get("url")) + && !value.is_null() + { + converted.insert("image_url".to_owned(), value.clone()); + } + if let Some(value) = item.get("file_id") + && !value.is_null() + { + converted.insert("file_id".to_owned(), value.clone()); + } + Ok(converted) + } + + fn normalize_function_call_input( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let mut name = required_trimmed_string(item, "name")?.to_owned(); + let call_id = required_trimmed_string(item, "call_id")?; + let arguments = encode_function_arguments(item.get("arguments"))?; + let namespace = string_field(item, "namespace").trim(); + if !namespace.is_empty() { + name = self.alias(ToolIdentity::new(ToolKind::Function, namespace, &name)); + } + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.remove("namespace"); + converted.insert("type".to_owned(), Value::String("function_call".to_owned())); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("name".to_owned(), Value::String(name)); + converted.insert("arguments".to_owned(), Value::String(arguments)); + Ok(converted) + } + + pub(super) fn normalize_custom_tool_call_input( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let name = required_trimmed_string(item, "name")?; + let input = item + .get("input") + .and_then(Value::as_str) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let call_id = required_trimmed_string(item, "call_id")?; + let namespace = string_field(item, "namespace").trim(); + let identity = ToolIdentity::new(ToolKind::Custom, namespace, name); + let arguments = serde_json::to_string(&json_object([( + identity.custom_argument_field(), + Value::String(input.to_owned()), + )])) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; + let alias = self.alias(identity); + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.remove("namespace"); + converted.remove("input"); + converted.insert("type".to_owned(), Value::String("function_call".to_owned())); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("name".to_owned(), Value::String(alias)); + converted.insert("arguments".to_owned(), Value::String(arguments)); + Ok(converted) + } + + fn normalize_function_call_output_input( + &mut self, + item: &Map, + allow_content_blocks: bool, + ) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let output = match item.get("output") { + Some(Value::Array(blocks)) + if allow_content_blocks && is_function_output_content_array(blocks) => + { + Value::Array(self.normalize_function_output_blocks(blocks)?) + } + output => Value::String(encode_tool_output(output)?), + }; + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.insert( + "type".to_owned(), + Value::String("function_call_output".to_owned()), + ); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("output".to_owned(), output); + Ok(converted) + } + + fn normalize_function_output_blocks( + &mut self, + blocks: &[Value], + ) -> Result, GrokRequestEncodeError> { + blocks + .iter() + .map(|raw| { + let block = raw + .as_object() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + match string_field(block, "type") { + "input_text" => Ok(json_object([ + ("type", Value::String("input_text".to_owned())), + ( + "text", + Value::String( + block + .get("text") + .and_then(Value::as_str) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)? + .to_owned(), + ), + ), + ])), + "input_image" => { + require_content_source(block, &["image_url", "file_id"])?; + self.normalize_input_image_part(block).map(Value::Object) + } + "input_file" => { + require_content_source(block, &["file_data", "file_id", "file_url"])?; + Ok(Value::Object(normalize_input_file_part(block))) + } + _ => Err(GrokRequestEncodeError::InvalidRequestNormalization), + } + }) + .collect() + } + + fn normalize_tool_search_call( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let execution = string_field(item, "execution").trim().to_ascii_lowercase(); + if execution.is_empty() || execution == "server" { + self.server_search_eager = true; + return Ok(boundary_message( + "A server-side tool search occurred here; selected tools are made available directly.", + )); + } + if execution != "client" { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let arguments = encode_function_arguments(item.get("arguments"))?; + let alias = self.alias(ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")); + Ok(Map::from_iter([ + ("type".to_owned(), Value::String("function_call".to_owned())), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("name".to_owned(), Value::String(alias)), + ("arguments".to_owned(), Value::String(arguments)), + ])) + } + + fn normalize_tool_search_output( + &mut self, + item: &Map, + ) -> Result { + let execution = string_field(item, "execution").trim().to_ascii_lowercase(); + if !matches!(execution.as_str(), "" | "client" | "server") { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let call_id = required_trimmed_string(item, "call_id")?; + let tools = item + .get("tools") + .and_then(Value::as_array) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let mut normalized = Vec::new(); + for tool in tools { + normalized.extend(self.normalize_tool(tool, "", false, true)?); + } + let message = format!( + "Tool search completed; {} selected tool definitions are now available.", + tools.len() + ); + let history = if execution == "client" { + Map::from_iter([ + ( + "type".to_owned(), + Value::String("function_call_output".to_owned()), + ), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("output".to_owned(), Value::String(message)), + ]) + } else { + self.server_search_eager = true; + boundary_message(&message) + }; + Ok(NormalizedToolSearchOutput { + history, + loaded_tools: normalized, + visible_tools: tools.clone(), + }) + } + + fn normalize_apply_patch_call_input( + &mut self, + item: &Map, + ) -> Result, GrokRequestEncodeError> { + let call_id = required_trimmed_string(item, "call_id")?; + let operation = validate_apply_patch_operation(item.get("operation"))?; + let arguments = + serde_json::to_string(&json_object([("operation", Value::Object(operation))])) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; + let alias = self.alias(ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")); + Ok(Map::from_iter([ + ("type".to_owned(), Value::String("function_call".to_owned())), + ("call_id".to_owned(), Value::String(call_id.to_owned())), + ("name".to_owned(), Value::String(alias)), + ("arguments".to_owned(), Value::String(arguments)), + ])) + } + + fn normalize_additional_tools_input( + &mut self, + item: &Map, + ) -> Result<(Vec, Vec), GrokRequestEncodeError> { + let tools = item + .get("tools") + .and_then(Value::as_array) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let mut normalized = Vec::new(); + for raw_tool in tools { + normalized.extend(self.normalize_tool(raw_tool, "", false, true)?); + } + Ok((normalized, tools.clone())) + } +} + +pub(super) fn normalize_input_file_part(item: &Map) -> Map { + let mut converted = + Map::from_iter([("type".to_owned(), Value::String("input_file".to_owned()))]); + for key in ["file_data", "file_id", "filename", "file_url"] { + if let Some(value) = item.get(key) + && !value.is_null() + { + converted.insert(key.to_owned(), value.clone()); + } + } + converted +} + +pub(super) fn required_trimmed_string<'a>( + item: &'a Map, + field: &str, +) -> Result<&'a str, GrokRequestEncodeError> { + item.get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization) +} + +pub(super) fn encode_function_arguments( + value: Option<&Value>, +) -> Result { + match value { + Some(Value::String(value)) => Ok(value.clone()), + None | Some(Value::Null) => Ok("{}".to_owned()), + Some(value) => serde_json::to_string(value) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization), + } +} + +pub(super) fn normalize_function_arguments(arguments: &str, schema: &Value) -> Option { + if arguments.trim().is_empty() || !schema.is_object() { + return None; + } + let mut value = serde_json::from_str::(arguments).ok()?; + if !normalize_argument_value(&mut value, schema, schema, 0) { + return None; + } + serde_json::to_string(&value).ok() +} + +pub(super) fn normalize_argument_value( + value: &mut Value, + schema: &Value, + root: &Value, + depth: usize, +) -> bool { + if depth > 64 { + return false; + } + let Some(schema) = schema.as_object() else { + return false; + }; + let mut changed = false; + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) + && let Some(resolved) = resolve_local_schema_value_ref(root, reference) + { + changed |= normalize_argument_value(value, resolved, root, depth + 1); + } + for keyword in ["allOf", "anyOf", "oneOf"] { + if let Some(branches) = schema.get(keyword).and_then(Value::as_array) { + for branch in branches { + changed |= normalize_argument_value(value, branch, root, depth + 1); + } + } + } + if schema_requires_integer(schema) + && let Value::Number(number) = value + && let Some(normalized) = normalize_integral_number(number) + { + *number = normalized; + return true; + } + match value { + Value::Object(object) => { + let properties = schema.get("properties").and_then(Value::as_object); + let additional = schema + .get("additionalProperties") + .filter(|value| value.is_object()); + for (key, item) in object { + let property = properties + .and_then(|properties| properties.get(key)) + .or(additional); + if let Some(property) = property { + changed |= normalize_argument_value(item, property, root, depth + 1); + } + } + } + Value::Array(items) => { + let prefix = schema.get("prefixItems").and_then(Value::as_array); + let item_schema = schema.get("items").filter(|value| value.is_object()); + for (index, item) in items.iter_mut().enumerate() { + let schema = prefix + .and_then(|prefix| prefix.get(index)) + .filter(|value| value.is_object()) + .or(item_schema); + if let Some(schema) = schema { + changed |= normalize_argument_value(item, schema, root, depth + 1); + } + } + } + _ => {} + } + changed +} + +pub(super) fn schema_requires_integer(schema: &Map) -> bool { + match schema.get("type") { + Some(Value::String(kind)) => kind == "integer", + Some(Value::Array(kinds)) => { + let mut integer = false; + for kind in kinds.iter().filter_map(Value::as_str) { + if kind == "number" { + return false; + } + integer |= kind == "integer"; + } + integer + } + _ => false, + } +} + +pub(super) fn normalize_integral_number(number: &serde_json::Number) -> Option { + let raw = number.to_string(); + if raw.len() > MAX_JSON_NUMBER_TEXT_BYTES || !raw.contains(['.', 'e', 'E']) { + return None; + } + let (mantissa, exponent_text) = raw.find(['e', 'E']).map_or((raw.as_str(), ""), |index| { + (&raw[..index], &raw[index + 1..]) + }); + let (negative, mantissa) = mantissa + .strip_prefix('-') + .map_or((false, mantissa), |value| (true, value)); + let (whole, fraction) = mantissa.split_once('.').unwrap_or((mantissa, "")); + let mut digits = format!("{whole}{fraction}") + .trim_start_matches('0') + .to_owned(); + if digits.is_empty() { + return "0".parse().ok().filter(|_| raw != "0"); + } + let exponent = parse_bounded_decimal_exponent(exponent_text)?; + let decimal_shift = exponent - i64::try_from(fraction.len()).ok()?; + if decimal_shift < 0 { + let fractional_digits = usize::try_from(-decimal_shift).ok()?; + if fractional_digits > digits.len() + || !digits[digits.len() - fractional_digits..] + .bytes() + .all(|byte| byte == b'0') + { + return None; + } + digits.truncate(digits.len() - fractional_digits); + let trimmed = digits.trim_start_matches('0'); + if trimmed.is_empty() { + return "0".parse().ok(); + } + digits = trimmed.to_owned(); + } else if decimal_shift > 0 { + let shift = usize::try_from(decimal_shift).ok()?; + if shift + > MAX_EXACT_JSON_INTEGER_TEXT + .len() + .saturating_sub(digits.len()) + { + return None; + } + digits.extend(std::iter::repeat_n('0', shift)); + } + if digits.len() > MAX_EXACT_JSON_INTEGER_TEXT.len() + || (digits.len() == MAX_EXACT_JSON_INTEGER_TEXT.len() + && digits.as_str() > MAX_EXACT_JSON_INTEGER_TEXT) + { + return None; + } + let normalized = if negative { + format!("-{digits}") + } else { + digits + }; + if normalized == raw { + None + } else { + normalized.parse().ok() + } +} + +pub(super) fn exact_nonnegative_sequence(value: &Value) -> Option { + let number = value.as_number()?; + let sequence = number + .as_u64() + .or_else(|| normalize_integral_number(number)?.as_u64())?; + (sequence <= MAX_EXACT_JSON_INTEGER).then_some(sequence) +} + +pub(super) fn parse_bounded_decimal_exponent(raw: &str) -> Option { + if raw.is_empty() { + return Some(0); + } + let (sign, digits) = if let Some(value) = raw.strip_prefix('+') { + (1_i64, value) + } else if let Some(value) = raw.strip_prefix('-') { + (-1_i64, value) + } else { + (1_i64, raw) + }; + if digits.is_empty() { + return None; + } + let digits = digits.trim_start_matches('0'); + if digits.is_empty() { + return Some(0); + } + if digits.len() > 9 { + return None; + } + digits.parse::().ok().map(|value| sign * value) +} + +pub(super) fn resolve_local_schema_value_ref<'a>( + root: &'a Value, + reference: &str, +) -> Option<&'a Value> { + if reference == "#" { + return Some(root); + } + let pointer = reference.strip_prefix('#')?; + if !pointer.starts_with('/') { + return None; + } + root.pointer(pointer) +} + +pub(super) fn schema_contains_integer(schema: &Value, depth: usize) -> bool { + let mut visited = BTreeSet::new(); + schema_contains_reachable_integer(schema, schema, &mut visited, depth) +} + +pub(super) fn schema_contains_reachable_integer( + schema: &Value, + root: &Value, + visited_refs: &mut BTreeSet, + depth: usize, +) -> bool { + if depth > 64 { + return false; + } + let Some(schema) = schema.as_object() else { + return false; + }; + if schema_requires_integer(schema) { + return true; + } + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) + && visited_refs.insert(reference.to_owned()) + && let Some(resolved) = resolve_local_schema_value_ref(root, reference) + && schema_contains_reachable_integer(resolved, root, visited_refs, depth + 1) + { + return true; + } + for keyword in ["allOf", "anyOf", "oneOf", "prefixItems"] { + if schema + .get(keyword) + .and_then(Value::as_array) + .is_some_and(|branches| { + branches.iter().any(|branch| { + schema_contains_reachable_integer(branch, root, visited_refs, depth + 1) + }) + }) + { + return true; + } + } + for keyword in ["items", "additionalProperties"] { + if schema.get(keyword).is_some_and(|child| { + schema_contains_reachable_integer(child, root, visited_refs, depth + 1) + }) { + return true; + } + } + schema + .get("properties") + .and_then(Value::as_object) + .is_some_and(|properties| { + properties.values().any(|property| { + schema_contains_reachable_integer(property, root, visited_refs, depth + 1) + }) + }) +} + +pub(super) fn encode_tool_output(value: Option<&Value>) -> Result { + match value { + None | Some(Value::Null) => Ok(String::new()), + Some(Value::String(value)) => Ok(value.clone()), + Some(value) => serde_json::to_string(value) + .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization), + } +} + +pub(super) fn is_function_output_content_array(blocks: &[Value]) -> bool { + blocks.iter().any(|raw| { + raw.pointer("/type") + .and_then(Value::as_str) + .is_some_and(|kind| kind.starts_with("input_")) + }) +} diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index d0455a6a8..8a8e67968 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -1235,7 +1235,7 @@ async fn transport_error_observation_should_retain_available_metrics() { } #[tokio::test] -async fn compaction_stream_should_emit_openai_wire_and_keep_only_accounting_canonical() { +async fn compaction_stream_should_emit_openai_wire_and_keep_only_metering_canonical() { let summary = valid_compaction_summary("validated summary"); let transport = StubInferenceTransport::sequence([InferenceMode::SuccessBody(compaction_sse( &summary, diff --git a/backend/crates/providers/xai/tests/transport/request.rs b/backend/crates/providers/xai/tests/transport/request/mod.rs similarity index 100% rename from backend/crates/providers/xai/tests/transport/request.rs rename to backend/crates/providers/xai/tests/transport/request/mod.rs diff --git a/docs/architecture.md b/docs/architecture.md index 91c8e9198..e448e91dc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -72,8 +72,8 @@ flowchart LR 3. API 只调用 Core/Admin 抽象；Store 只实现端口。 4. 具体实现只在组合根相遇。 -`gateway-api` 与 `apps/gateway` 的 architecture tests 冻结关键源码树和依赖禁区；其余边界由 manifest、 -集成测试和评审共同维护。 +组合根的 workspace architecture tests 冻结成员清单、依赖 DAG、公开模块面、源码纪律以及生产/测试模块 +镜像关系；具体行为边界由各 crate 的集成测试维护。 ### 3.1 Rust 模块组织约定 @@ -84,12 +84,18 @@ flowchart LR `目录/mod.rs`，叶子模块使用 `名称.rs`，不混用 `名称.rs` 与 `名称/` 两套入口。 2. 生产源码不使用内联 `mod name { ... }`、`#[path]` 或 `include!` 隐藏文件归属。`lib.rs` 和各级 `mod.rs` 是模块门面，负责声明子模块、选择性 re-export 与少量同层协调。 -3. 子模块默认私有；只有真实的跨 crate 生产合同才能使用 `pub`。当前 adapter/provider 的存量公开模块由 - allowlist 冻结，集成测试组织完成调整后只允许缩小，不能无审计扩张。 +3. 子模块默认私有；只有真实的跨 crate 生产合同才能使用 `pub`。adapter/provider 的公开根模块由 + allowlist 冻结，变更必须同步完成边界审计，不能为测试或兼容路径创建第二 owner。 4. owner 级依赖保持单向。两个模块互相引用时，应把共享值对象、端口或生命周期能力提升到共同 owner， 不能依靠同一 crate 内可见性掩盖循环边界。 +5. 每个 crate 的 `src/` 与 `tests/` 平级，生产源码不承载测试。`tests/` 镜像 `src/` 的模块目录形态：例如 + `src/foo/mod.rs` 对应 `tests/foo/mod.rs`，`src/foo/bar.rs` 对应 `tests/foo/bar.rs`。一个生产模块可以没有 + 测试；额外场景测试必须放在最近的生产 owner 目录下。根级 `support` 和冻结的 crate/workspace 架构场景 + 是明确例外。 -当前遵守情况、越界点与收敛顺序见 [后端 Rust 模块架构审计](backend-module-audit.md)。 +较大的模块门面只负责组合和 re-export：账号 Admin HTTP 边界按 `wire`、`credentials`、`handlers`、 +`presenter` 划分；Provider 执行按 continuation、stream、failure、observation 和 worker 等职责拆分；xAI +请求转换按 response、tools 与 history 拆分。各子模块之间只使用 owner 内最小可见性。 ### 3.2 `gateway-core` 内部 owner @@ -102,7 +108,8 @@ flowchart LR | `upstream` | 跨 Engine、Event、Error 与 Provider 共用的 transport 名称和发送状态 | | `lifecycle` | 取消信号、连接注册与 drain 合同 | | `engine` | attempt、发送/提交屏障、执行编排和持久化调用时序 | -| `routing` | 冻结路由事实、请求计划和运行时快照 | +| `routing` | 冻结路由事实、请求计划以及运行时快照的表示与编译 | +| `runtime` | 当前快照的发布、读取、revision 订阅与周期对账任务 | ## 4. 数据面请求生命周期 @@ -263,7 +270,7 @@ credential 与 quota 是两组独立事实：credential refresh 不等于 quota Worker 由各 Bundle 贡献、由 Host 统一监督： - Store：过期请求恢复、历史保留和 PostgreSQL/Redis 观测队列； -- Core：RuntimeSnapshot 周期对账和 Redis change 订阅； +- Core：`runtime` owner 的 RuntimeSnapshot 周期对账和 Redis change 订阅； - Admin：S3/R2 备份 daemon，负责调度、执行、删除收敛与保留清理； - Provider：credential refresh、quota/catalog 健康和官方版本/etag 检查。