From 7cb898b40a5fb10544a09857313c7122f02b69cf Mon Sep 17 00:00:00 2001 From: zyycn Date: Wed, 22 Jul 2026 19:51:12 +0800 Subject: [PATCH] refactor(gateway): finalize production API audit Remove unused full-response transports and unconnected domain reservations while keeping test-only setup inside integration tests. Restore image-generation accounting and OpenAI WebSocket pool observations across the engine, store, and admin boundaries without changing the transparent OpenAI WS/SSE path. Co-authored-by: Codex --- .../gateway-admin/src/model/accounts.rs | 8 + .../gateway-admin/src/model/observability.rs | 13 + .../crates/gateway-api/src/admin/accounts.rs | 46 +++ .../gateway-api/src/admin/observability.rs | 22 ++ .../src/openai/responses/request.rs | 10 + .../gateway-api/tests/admin/accounts.rs | 8 + .../gateway-api/tests/admin/observability.rs | 30 +- .../gateway-api/tests/openai/responses/mod.rs | 11 + .../crates/gateway-core/src/accounting/mod.rs | 36 +-- .../gateway-core/src/engine/continuation.rs | 93 ------ .../gateway-core/src/engine/coordinator.rs | 27 +- .../gateway-core/src/engine/execution.rs | 2 + backend/crates/gateway-core/src/engine/mod.rs | 3 + backend/crates/gateway-core/src/error.rs | 6 - backend/crates/gateway-core/src/event.rs | 56 ++-- .../crates/gateway-core/src/operation/mod.rs | 58 ++-- .../gateway-core/src/routing/snapshot.rs | 5 - backend/crates/gateway-core/src/task/mod.rs | 12 +- .../gateway-core/tests/accounting/mod.rs | 27 -- .../gateway-core/tests/engine/continuation.rs | 30 +- .../gateway-core/tests/engine/coordinator.rs | 203 +++++++++--- .../gateway-core/tests/operation/mod.rs | 7 +- backend/crates/gateway-core/tests/task/mod.rs | 19 +- .../crates/gateway-host/tests/workers/mod.rs | 30 +- .../crates/gateway-protocol/src/openai/mod.rs | 2 - .../gateway-protocol/src/openai/schema.rs | 129 -------- .../gateway-protocol/tests/openai/mod.rs | 1 - .../gateway-protocol/tests/openai/schema.rs | 89 ------ backend/crates/gateway-store/src/lib.rs | 4 - .../gateway-store/src/postgres/execution.rs | 56 +++- .../gateway-store/src/postgres/history.rs | 12 +- .../src/postgres/observability.rs | 59 +++- .../src/postgres/provider_accounts.rs | 8 + .../gateway-store/src/redis/runtime_change.rs | 50 +-- .../gateway-store/tests/postgres/execution.rs | 153 ++++++++-- .../gateway-store/tests/postgres/history.rs | 43 +-- .../tests/postgres/observability.rs | 42 ++- .../crates/gateway-store/tests/redis/mod.rs | 1 - .../tests/redis/runtime_change.rs | 12 - .../gateway-store/tests/redis/worker_lease.rs | 51 +++- .../providers/openai/src/credential/admin.rs | 23 -- .../providers/openai/src/credential/mod.rs | 2 +- .../openai/src/credential/refresh.rs | 8 - .../crates/providers/openai/src/provider.rs | 12 +- .../openai/src/transport/canonical.rs | 2 + .../providers/openai/src/transport/client.rs | 38 +-- .../openai/src/transport/client_sse.rs | 227 +------------- .../openai/src/transport/diagnostics.rs | 29 -- .../providers/openai/src/transport/mod.rs | 8 +- .../src/transport/protocol/responses.rs | 190 +----------- .../openai/src/transport/response_meta.rs | 44 +-- .../providers/openai/src/transport/time.rs | 6 - .../openai/src/transport/websocket/breaker.rs | 3 +- .../src/transport/websocket/coordinator.rs | 74 +---- .../transport/websocket/exchange/collect.rs | 133 -------- .../src/transport/websocket/exchange/mod.rs | 68 +---- .../openai/src/transport/websocket/mod.rs | 5 +- .../src/transport/websocket/pool/mod.rs | 5 +- .../openai/tests/credential/admin.rs | 28 +- .../openai/tests/provider/contract.rs | 4 +- .../openai/tests/transport/canonical.rs | 23 ++ .../openai/tests/transport/latency.rs | 288 ++++++++++-------- .../providers/openai/tests/transport/mod.rs | 159 +++++++++- .../openai/tests/transport/websocket.rs | 53 ++-- .../openai/tests/transport/websocket_pool.rs | 3 +- .../providers/xai/src/credential/client.rs | 12 +- .../providers/xai/src/credential/http.rs | 6 - .../providers/xai/src/credential/refresh.rs | 8 - .../xai/src/credential/repository.rs | 20 -- .../providers/xai/src/credential/token.rs | 18 +- .../providers/xai/src/transport/canonical.rs | 2 + .../providers/xai/src/transport/headers.rs | 6 - .../providers/xai/src/transport/network.rs | 12 +- .../providers/xai/src/transport/request.rs | 27 -- .../xai/tests/credential/contract.rs | 36 --- .../providers/xai/tests/credential/refresh.rs | 55 +--- .../providers/xai/tests/provider/contract.rs | 8 +- .../xai/tests/transport/canonical.rs | 17 ++ backend/migrations/0001_initial.sql | 23 ++ 79 files changed, 1220 insertions(+), 1939 deletions(-) delete mode 100644 backend/crates/gateway-protocol/src/openai/schema.rs delete mode 100644 backend/crates/gateway-protocol/tests/openai/schema.rs delete mode 100644 backend/crates/gateway-store/tests/redis/runtime_change.rs delete mode 100644 backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs diff --git a/backend/crates/gateway-admin/src/model/accounts.rs b/backend/crates/gateway-admin/src/model/accounts.rs index 96fd5fd3e..b279dd975 100644 --- a/backend/crates/gateway-admin/src/model/accounts.rs +++ b/backend/crates/gateway-admin/src/model/accounts.rs @@ -111,6 +111,10 @@ pub struct AccountModelUsage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: super::observability::CostCoverage, pub costs: Vec, @@ -128,6 +132,10 @@ pub struct AccountUsage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: super::observability::CostCoverage, pub costs: Vec, diff --git a/backend/crates/gateway-admin/src/model/observability.rs b/backend/crates/gateway-admin/src/model/observability.rs index 91cfa8ff8..a467e982a 100644 --- a/backend/crates/gateway-admin/src/model/observability.rs +++ b/backend/crates/gateway-admin/src/model/observability.rs @@ -443,6 +443,10 @@ pub struct DashboardAccountModelUsage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: CostCoverage, pub costs: Vec, @@ -467,6 +471,10 @@ pub struct DashboardAccountUsage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: CostCoverage, pub costs: Vec, @@ -506,6 +514,7 @@ pub struct UsageRecord { pub upstream_model_id: Option, pub upstream_transport: Option, pub http_version: Option, + pub websocket_pool: Option, pub attempt_count: u32, pub upstream_send_state: String, pub downstream_committed_at: Option>, @@ -524,6 +533,8 @@ pub struct UsageRecord { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, pub total_tokens: Option, pub cost_source: String, pub cost_amount: Option, @@ -545,6 +556,8 @@ pub struct UsageRecord { pub request_kind: Option, pub subagent_kind: Option, pub compact: bool, + pub image_generation_requested: bool, + pub image_generation_succeeded: Option, pub started_at: DateTime, pub deadline_at: DateTime, pub completed_at: Option>, diff --git a/backend/crates/gateway-api/src/admin/accounts.rs b/backend/crates/gateway-api/src/admin/accounts.rs index 94ae96453..1be7d6e7b 100644 --- a/backend/crates/gateway-api/src/admin/accounts.rs +++ b/backend/crates/gateway-api/src/admin/accounts.rs @@ -236,6 +236,14 @@ pub struct AccountUsageView { 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, @@ -267,6 +275,14 @@ pub struct ModelUsageView { 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, @@ -879,6 +895,12 @@ fn quota_local_usage(usage: &AccountUsage) -> Value { "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_tokens(total_tokens), }) @@ -908,6 +930,14 @@ fn account_usage_view(usage: Option, now: DateTime) -> Accoun 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, @@ -967,6 +997,14 @@ fn account_model_usage_view(usage: AccountModelUsage, now: DateTime) -> Mod 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()), @@ -1006,6 +1044,14 @@ fn empty_account_usage() -> AccountUsageView { 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, diff --git a/backend/crates/gateway-api/src/admin/observability.rs b/backend/crates/gateway-api/src/admin/observability.rs index 5ab7834fd..22fdd5c88 100644 --- a/backend/crates/gateway-api/src/admin/observability.rs +++ b/backend/crates/gateway-api/src/admin/observability.rs @@ -521,12 +521,16 @@ pub struct TokenDetailsView { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, pub total_tokens: Option, pub input_tokens_display: String, pub output_tokens_display: String, pub cached_tokens_display: String, pub cache_write_tokens_display: String, pub reasoning_tokens_display: String, + pub image_input_tokens_display: String, + pub image_output_tokens_display: String, pub total_tokens_display: String, } @@ -595,6 +599,8 @@ pub struct UsageRecordView { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, pub message: String, pub metadata: UsageRecordMetadataView, pub created_at: DateTime, @@ -623,6 +629,9 @@ pub struct UsageRecordMetadataView { pub protocol: String, pub logical_outcome: String, pub attempt_count: u64, + pub websocket_pool: Option, + pub image_generation_requested: bool, + pub image_generation_succeeded: Option, } /// 单次上游尝试展示。 @@ -1520,6 +1529,8 @@ fn token_details(record: &domain::UsageRecord) -> TokenDetailsView { cached_tokens: record.cached_tokens, cache_write_tokens: record.cache_write_tokens, reasoning_tokens: record.reasoning_tokens, + image_input_tokens: record.image_input_tokens, + image_output_tokens: record.image_output_tokens, total_tokens: record.total_tokens, input_tokens_display: record .input_tokens @@ -1536,6 +1547,12 @@ fn token_details(record: &domain::UsageRecord) -> TokenDetailsView { reasoning_tokens_display: record .reasoning_tokens .map_or_else(|| "-".to_owned(), format_number), + image_input_tokens_display: record + .image_input_tokens + .map_or_else(|| "-".to_owned(), format_number), + image_output_tokens_display: record + .image_output_tokens + .map_or_else(|| "-".to_owned(), format_number), total_tokens_display: record .total_tokens .map_or_else(|| "-".to_owned(), format_number), @@ -1687,11 +1704,16 @@ fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView { cached_tokens: record.cached_tokens, cache_write_tokens: record.cache_write_tokens, reasoning_tokens: record.reasoning_tokens, + image_input_tokens: record.image_input_tokens, + image_output_tokens: record.image_output_tokens, message, metadata: UsageRecordMetadataView { protocol: record.protocol, logical_outcome: outcome.clone(), attempt_count: u64::from(record.attempt_count), + websocket_pool: record.websocket_pool, + image_generation_requested: record.image_generation_requested, + image_generation_succeeded: record.image_generation_succeeded, }, created_at: record.started_at, created_at_display: china_datetime(&record.started_at), diff --git a/backend/crates/gateway-api/src/openai/responses/request.rs b/backend/crates/gateway-api/src/openai/responses/request.rs index e1af6fa91..63b51c01e 100644 --- a/backend/crates/gateway-api/src/openai/responses/request.rs +++ b/backend/crates/gateway-api/src/openai/responses/request.rs @@ -411,6 +411,15 @@ pub(super) fn decode_request_inner( .get("tools") .and_then(Value::as_array) .is_some_and(|tools| !tools.is_empty()); + let image_generation_requested = + object + .get("tools") + .and_then(Value::as_array) + .is_some_and(|tools| { + tools.iter().any(|tool| { + tool.get("type").and_then(Value::as_str) == Some("image_generation") + }) + }); let vision_requested = contains_type(object.get("input"), "input_image"); let reasoning = canonical_reasoning(&object); let reasoning_requested = object.get("reasoning").is_some_and(Value::is_object); @@ -438,6 +447,7 @@ pub(super) fn decode_request_inner( })?; let mut request = GenerateRequest::from_protocol_payload(messages, payload) .with_estimated_context_tokens(input_token_estimate) + .with_image_generation_requested(image_generation_requested) .with_provider_options(provider_options) .with_response_persistence(if store { ResponsePersistence::Store diff --git a/backend/crates/gateway-api/tests/admin/accounts.rs b/backend/crates/gateway-api/tests/admin/accounts.rs index b3b722f53..ae20b29af 100644 --- a/backend/crates/gateway-api/tests/admin/accounts.rs +++ b/backend/crates/gateway-api/tests/admin/accounts.rs @@ -71,6 +71,14 @@ mod response { 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, diff --git a/backend/crates/gateway-api/tests/admin/observability.rs b/backend/crates/gateway-api/tests/admin/observability.rs index 7f0998e99..1814c8dee 100644 --- a/backend/crates/gateway-api/tests/admin/observability.rs +++ b/backend/crates/gateway-api/tests/admin/observability.rs @@ -172,12 +172,16 @@ mod response { cached_tokens: None, cache_write_tokens: None, reasoning_tokens: None, + image_input_tokens: None, + image_output_tokens: None, total_tokens: Some(3), input_tokens_display: "1".to_owned(), output_tokens_display: "2".to_owned(), cached_tokens_display: "-".to_owned(), cache_write_tokens_display: "-".to_owned(), reasoning_tokens_display: "-".to_owned(), + image_input_tokens_display: "-".to_owned(), + image_output_tokens_display: "-".to_owned(), total_tokens_display: "3".to_owned(), }; let cursor = CursorWire { @@ -257,7 +261,7 @@ async fn usage_route_should_forward_a_bounded_unknown_outcome_filter() { } #[tokio::test] -async fn usage_route_should_return_the_recorded_endpoint_as_route() { +async fn usage_route_should_expose_image_and_websocket_facts() { use axum::{ body::{Body, to_bytes}, http::{Request, StatusCode, header}, @@ -293,6 +297,7 @@ async fn usage_route_should_return_the_recorded_endpoint_as_route() { upstream_model_id: Some("grok-4.5".to_owned()), upstream_transport: Some("http_sse".to_owned()), http_version: Some("h2".to_owned()), + websocket_pool: Some("reuse".to_owned()), attempt_count: 1, upstream_send_state: "sent".to_owned(), downstream_committed_at: None, @@ -311,6 +316,8 @@ async fn usage_route_should_return_the_recorded_endpoint_as_route() { cached_tokens: Some(0), cache_write_tokens: Some(0), reasoning_tokens: Some(0), + image_input_tokens: Some(31), + image_output_tokens: Some(9), total_tokens: Some(2), cost_source: "unavailable".to_owned(), cost_amount: None, @@ -332,6 +339,8 @@ async fn usage_route_should_return_the_recorded_endpoint_as_route() { request_kind: None, subagent_kind: None, compact: false, + image_generation_requested: true, + image_generation_succeeded: Some(true), started_at: Utc::now(), deadline_at: Utc::now(), completed_at: Some(Utc::now()), @@ -354,5 +363,22 @@ async fn usage_route_should_return_the_recorded_endpoint_as_route() { .expect("usage response body"); let value: serde_json::Value = serde_json::from_slice(&body).expect("usage response JSON"); - assert_eq!(value["data"]["items"][0]["route"], "/v1/responses"); + assert_eq!( + serde_json::json!({ + "route": value["data"]["items"][0]["route"], + "imageInputTokens": value["data"]["items"][0]["tokenDetails"]["imageInputTokens"], + "imageOutputTokens": value["data"]["items"][0]["tokenDetails"]["imageOutputTokens"], + "websocketPool": value["data"]["items"][0]["metadata"]["websocketPool"], + "imageGenerationRequested": value["data"]["items"][0]["metadata"]["imageGenerationRequested"], + "imageGenerationSucceeded": value["data"]["items"][0]["metadata"]["imageGenerationSucceeded"], + }), + serde_json::json!({ + "route": "/v1/responses", + "imageInputTokens": 31, + "imageOutputTokens": 9, + "websocketPool": "reuse", + "imageGenerationRequested": true, + "imageGenerationSucceeded": true, + }) + ); } diff --git a/backend/crates/gateway-api/tests/openai/responses/mod.rs b/backend/crates/gateway-api/tests/openai/responses/mod.rs index 983c46fc7..e74274952 100644 --- a/backend/crates/gateway-api/tests/openai/responses/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/mod.rs @@ -212,6 +212,17 @@ fn decoder_should_preserve_roles_tools_reasoning_schema_and_provider_options() { ); } +#[test] +fn decoder_should_mark_image_generation_tool_intent() { + let decoded = generate_request(json!({ + "model": "smart-code", + "input": "draw a circuit diagram", + "tools": [{"type": "image_generation"}] + })); + + assert!(generate_operation(&decoded).image_generation_requested()); +} + #[test] fn decoder_should_preserve_only_explicit_xai_provider_options() { let decoded = generate_request(json!({ diff --git a/backend/crates/gateway-core/src/accounting/mod.rs b/backend/crates/gateway-core/src/accounting/mod.rs index 41c2ea72c..caaae04ae 100644 --- a/backend/crates/gateway-core/src/accounting/mod.rs +++ b/backend/crates/gateway-core/src/accounting/mod.rs @@ -152,7 +152,7 @@ impl Money { } } -/// `model_requests` 的六个公共 Token 事实。 +/// `model_requests` 的公共 Token 事实。 #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct Usage { pub input_tokens: Option, @@ -160,6 +160,8 @@ pub struct Usage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, /// Provider/协议报告的独立事实，不从其他列相加推导。 pub total_tokens: Option, } @@ -173,6 +175,8 @@ impl Usage { cached_tokens: None, cache_write_tokens: None, reasoning_tokens: None, + image_input_tokens: None, + image_output_tokens: None, total_tokens: None, } } @@ -194,6 +198,12 @@ impl Usage { if newer.reasoning_tokens.is_some() { self.reasoning_tokens = newer.reasoning_tokens; } + if newer.image_input_tokens.is_some() { + self.image_input_tokens = newer.image_input_tokens; + } + if newer.image_output_tokens.is_some() { + self.image_output_tokens = newer.image_output_tokens; + } if newer.total_tokens.is_some() { self.total_tokens = newer.total_tokens; } @@ -223,7 +233,6 @@ impl CostSource { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum CostEstimateStatus { Known, - Partial, Unknown, } @@ -232,7 +241,6 @@ impl CostEstimateStatus { pub const fn as_str(self) -> &'static str { match self { Self::Known => "known", - Self::Partial => "partial", Self::Unknown => "unknown", } } @@ -486,28 +494,6 @@ impl CostEstimate { } } - /// 创建有金额的 known/partial 结果。 - /// - /// # Errors - /// - /// `Unavailable` 不能携带金额，`Unknown` 不能伪装成可信金额。 - pub fn priced( - status: CostEstimateStatus, - source: CostSource, - total: Money, - ) -> Result { - if source == CostSource::Unavailable || status == CostEstimateStatus::Unknown { - return Err(AccountingError::InvalidCostEstimate { - status: status.as_str(), - }); - } - Ok(Self { - status, - source, - total: Some(total), - }) - } - #[must_use] pub const fn status(&self) -> CostEstimateStatus { self.status diff --git a/backend/crates/gateway-core/src/engine/continuation.rs b/backend/crates/gateway-core/src/engine/continuation.rs index cd103a4f8..329c26bb2 100644 --- a/backend/crates/gateway-core/src/engine/continuation.rs +++ b/backend/crates/gateway-core/src/engine/continuation.rs @@ -6,7 +6,6 @@ use std::fmt; use futures::future::BoxFuture; -use thiserror::Error; use crate::engine::credential::ProviderAccountId; use crate::error::{IdentifierError, SafeUpstreamValue, validate_text}; @@ -35,13 +34,6 @@ impl PreviousResponseId { } } -/// Provider 声明的 native handle 复用方式。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum NativeContinuationReuse { - Reusable, - SingleUse, -} - /// Provider 原生 response handle 的续接范围。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum NativeContinuationScope { @@ -64,7 +56,6 @@ pub struct NativeContinuationPin { provider: ProviderKind, instance: ProviderInstanceId, account: ProviderAccountId, - reuse: NativeContinuationReuse, scope: NativeContinuationScope, } @@ -76,7 +67,6 @@ impl NativeContinuationPin { provider: ProviderKind, instance: ProviderInstanceId, account: ProviderAccountId, - reuse: NativeContinuationReuse, ) -> Self { Self { previous_response_id, @@ -84,7 +74,6 @@ impl NativeContinuationPin { provider, instance, account, - reuse, scope: NativeContinuationScope::ConnectionLocal, } } @@ -122,11 +111,6 @@ impl NativeContinuationPin { &self.account } - #[must_use] - pub const fn reuse(&self) -> NativeContinuationReuse { - self.reuse - } - #[must_use] pub const fn scope(&self) -> NativeContinuationScope { self.scope @@ -153,7 +137,6 @@ impl fmt::Debug for NativeContinuationPin { .field("provider", &self.provider) .field("instance", &self.instance) .field("account", &self.account) - .field("reuse", &self.reuse) .field("scope", &self.scope) .finish() } @@ -199,82 +182,6 @@ impl fmt::Debug for ContinuationBinding { } } -/// Request-local 的 single-use claim 状态。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum NativeClaimState { - Available, - Claimed, - Consumed, - Ambiguous, -} - -/// 不落库的 native claim 状态机。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct NativeClaim { - reuse: NativeContinuationReuse, - state: NativeClaimState, -} - -impl NativeClaim { - #[must_use] - pub const fn new(reuse: NativeContinuationReuse) -> Self { - Self { - reuse, - state: NativeClaimState::Available, - } - } - - #[must_use] - pub const fn state(self) -> NativeClaimState { - self.state - } - - /// 在发送前领取 handle。 - /// - /// # Errors - /// - /// Single-use handle 已经被领取或终结时返回错误。 - pub fn claim(&mut self) -> Result<(), ContinuationError> { - if self.reuse == NativeContinuationReuse::Reusable { - return Ok(()); - } - if self.state != NativeClaimState::Available { - return Err(ContinuationError::AlreadyClaimed); - } - self.state = NativeClaimState::Claimed; - Ok(()) - } - - /// 明确未发送时释放 single-use claim。 - pub fn release_not_sent(&mut self) { - if self.reuse == NativeContinuationReuse::SingleUse - && self.state == NativeClaimState::Claimed - { - self.state = NativeClaimState::Available; - } - } - - /// 上游明确接收后消费 handle。 - pub fn consume(&mut self) { - if self.reuse == NativeContinuationReuse::SingleUse { - self.state = NativeClaimState::Consumed; - } - } - - /// 发送结果不确定时禁止再次使用 handle。 - pub fn mark_ambiguous(&mut self) { - if self.reuse == NativeContinuationReuse::SingleUse { - self.state = NativeClaimState::Ambiguous; - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] -pub enum ContinuationError { - #[error("native continuation handle is already claimed or terminal")] - AlreadyClaimed, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] #[error("native continuation store is unavailable")] pub struct NativeContinuationStoreError; diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 093fa6a41..c9a927eed 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -68,6 +68,7 @@ where .map(ProviderAccountStateOwner::from_continuation); let continuation_attempt = initial_continuation_attempt(&operation, &plan, continuation.as_ref()); + let image_generation_requested = operation.image_generation_requested(); let commit_policy = if continuation_attempt == ContinuationAttempt::None { StreamCommitPolicy::FirstCanonicalEvent } else { @@ -100,6 +101,7 @@ where upstream_complete: false, finalized: false, usage: Usage::new(), + image_generation_requested, cost: CostEstimate::unavailable(), timings: ModelRequestTimings::default(), gateway_response_id, @@ -177,6 +179,7 @@ pub struct ResponseExecutionSession { upstream_complete: bool, finalized: bool, usage: Usage, + image_generation_requested: bool, cost: CostEstimate, timings: ModelRequestTimings, gateway_response_id: String, @@ -955,7 +958,8 @@ where .and_then(|current| current.response_observation.as_ref()) .and_then(ProviderResponseObservation::status_code) .or(Some(200)); - let (upstream_transport, http_version) = self.current_transport_observation(); + let (upstream_transport, http_version, websocket_pool) = + self.current_transport_observation(); self.engine .store() .finalize_model_request(ModelRequestFinalization { @@ -971,10 +975,12 @@ where upstream_response_id: self.upstream_response_id.clone(), upstream_transport, http_version, + websocket_pool, error: None, provider_error_code: None, retry_after_ms: None, usage: self.usage.clone(), + image_generation_succeeded: self.image_generation_succeeded(), cost: self.cost.clone(), timings: self.timings.clone(), completed_at, @@ -1094,7 +1100,8 @@ where .as_ref() .and_then(|current| current.response_observation.as_ref()) .and_then(ProviderResponseObservation::status_code); - let (upstream_transport, http_version) = self.current_transport_observation(); + let (upstream_transport, http_version, websocket_pool) = + self.current_transport_observation(); self.engine .store() .finalize_model_request(ModelRequestFinalization { @@ -1112,10 +1119,12 @@ where .or_else(|| self.upstream_response_id.clone()), upstream_transport, http_version, + websocket_pool, error: Some(finalization.error), provider_error_code: finalization.provider_error_code, retry_after_ms: finalization.retry_after_ms, usage: self.usage.clone(), + image_generation_succeeded: self.image_generation_succeeded(), cost: self.cost.clone(), timings: self.timings.clone(), completed_at, @@ -1125,21 +1134,29 @@ where Ok(()) } - fn current_transport_observation(&self) -> (Option, Option) { + fn current_transport_observation(&self) -> (Option, Option, Option) { let Some(current) = self.current.as_ref() else { - return (None, None); + return (None, None, None); }; let Some(observation) = current.response_observation.as_ref() else { - return (None, None); + return (None, None, None); }; ( Some(observation.transport().as_str().to_owned()), observation .http_version() .map(|version| version.as_str().to_owned()), + observation + .websocket_pool() + .map(|kind| kind.as_str().to_owned()), ) } + fn image_generation_succeeded(&self) -> Option { + self.image_generation_requested + .then(|| self.usage.image_output_tokens.unwrap_or_default() > 0) + } + fn current_send_state(&self) -> UpstreamSendState { if self.attempts == 0 { UpstreamSendState::NotSent diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index fc972bfcc..139e1f417 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -341,6 +341,7 @@ impl DefaultExecutionService { request_kind: observation.request_kind, subagent_kind: observation.subagent_kind, compact: observation.compact, + image_generation_requested: request.operation.image_generation_requested(), started_at, deadline_at, }; @@ -488,6 +489,7 @@ impl DefaultExecutionService { request_kind: Some("account_connection_test".to_owned()), subagent_kind: None, compact: false, + image_generation_requested: false, started_at, deadline_at, }; diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 8239bb6b7..46847f8f5 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -523,6 +523,7 @@ pub struct NewModelRequest { pub request_kind: Option, pub subagent_kind: Option, pub compact: bool, + pub image_generation_requested: bool, pub started_at: SystemTime, pub deadline_at: SystemTime, } @@ -587,10 +588,12 @@ pub struct ModelRequestFinalization { pub upstream_response_id: Option, pub upstream_transport: Option, pub http_version: Option, + pub websocket_pool: Option, pub error: Option, pub provider_error_code: Option, pub retry_after_ms: Option, pub usage: Usage, + pub image_generation_succeeded: Option, pub cost: CostEstimate, pub timings: ModelRequestTimings, pub completed_at: SystemTime, diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index ea2e1708b..29305c086 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -115,12 +115,6 @@ pub enum AccountingError { /// 货币代码无效。 #[error("currency must be a three-letter uppercase ASCII code")] InvalidCurrency, - /// 价格估算状态与金额字段不一致。 - #[error("cost estimate fields do not match status `{status}`")] - InvalidCostEstimate { - /// 估算状态。 - status: &'static str, - }, } /// 跨 Provider 稳定的上游失败分类。 diff --git a/backend/crates/gateway-core/src/event.rs b/backend/crates/gateway-core/src/event.rs index 38afbd8d8..a287385c5 100644 --- a/backend/crates/gateway-core/src/event.rs +++ b/backend/crates/gateway-core/src/event.rs @@ -45,16 +45,6 @@ impl ResponseMeta { &self.response_id } - /// 以客户端协议层冻结的网关 response ID 替换 Provider 临时 ID。 - /// - /// Provider ID 只用于当前 attempt 内校验事件关联；跨请求 continuation 必须 - /// 使用网关拥有且可按调用方隔离解析的 ID。 - #[must_use] - pub fn with_gateway_response_id(mut self, response_id: impl Into) -> Self { - self.response_id = response_id.into(); - self - } - /// 返回对客户端公开的模型名。 #[must_use] pub fn model(&self) -> &str { @@ -477,6 +467,7 @@ impl fmt::Debug for ProviderResponseHeader { pub struct ProviderResponseObservation { transport: UpstreamTransport, http_version: Option, + websocket_pool: Option, status_code: Option, request_id: Option, timings: ProviderResponseTimings, @@ -489,6 +480,7 @@ impl ProviderResponseObservation { Self { transport, http_version: None, + websocket_pool: None, status_code: None, request_id: None, timings: ProviderResponseTimings::default(), @@ -502,6 +494,12 @@ impl ProviderResponseObservation { self } + #[must_use] + pub const fn with_websocket_pool(mut self, kind: WebSocketPoolKind) -> Self { + self.websocket_pool = Some(kind); + self + } + #[must_use] pub fn with_status_code(mut self, status_code: u16) -> Self { self.status_code = (100..=599).contains(&status_code).then_some(status_code); @@ -536,6 +534,11 @@ impl ProviderResponseObservation { self.http_version } + #[must_use] + pub const fn websocket_pool(&self) -> Option { + self.websocket_pool + } + #[must_use] pub const fn status_code(&self) -> Option { self.status_code @@ -557,6 +560,23 @@ impl ProviderResponseObservation { } } +/// WebSocket 请求使用新连接或复用池中连接的事实。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum WebSocketPoolKind { + New, + Reuse, +} + +impl WebSocketPoolKind { + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::New => "new", + Self::Reuse => "reuse", + } + } +} + /// 所有 Provider 都必须输出的稳定事件集合。 #[derive(Debug, Clone, PartialEq)] #[non_exhaustive] @@ -640,13 +660,6 @@ impl ProviderEvent { } } - /// 附着只由同协议客户端连接保存、并在下一轮原样交还 Provider 的状态。 - #[must_use] - pub fn with_session_update(mut self, state: ProviderSessionState) -> Self { - self.session_update = Some(Box::new(state)); - self - } - /// 把本事件标记为 Provider 连接内状态的提交边界。 pub fn attach_session_update(&mut self, state: ProviderSessionState) { self.session_update = Some(Box::new(state)); @@ -732,15 +745,6 @@ impl From for ProviderEvent { } impl GatewayEvent { - /// Canonical event 均为客户端 encoder 可交付事件。 - /// - /// Attempt Coordinator 会把首个事件标记为 downstream commit 候选， - /// adapter 可以先编码，但在 commit 持久化成功前不得把结果交给客户端。 - #[must_use] - pub const fn is_downstream_deliverable(&self) -> bool { - true - } - /// 返回该 fact 是否足以冻结下游 commit barrier。 /// /// 生命周期、结构与结算 facts 可以在换号前安全丢弃；首个可见增量或 diff --git a/backend/crates/gateway-core/src/operation/mod.rs b/backend/crates/gateway-core/src/operation/mod.rs index 12b50d76f..6e578b8a7 100644 --- a/backend/crates/gateway-core/src/operation/mod.rs +++ b/backend/crates/gateway-core/src/operation/mod.rs @@ -656,7 +656,7 @@ pub struct GenerateRequest { continuation: Option, max_output_tokens: Option, estimated_context_tokens: u64, - retry_safety: RetrySafety, + image_generation_requested: bool, response_persistence: ResponsePersistence, } @@ -713,7 +713,7 @@ impl GenerateRequest { continuation: None, max_output_tokens: None, estimated_context_tokens: 0, - retry_safety: RetrySafety::NonIdempotent, + image_generation_requested: false, response_persistence: ResponsePersistence::Store, } } @@ -767,10 +767,10 @@ impl GenerateRequest { self } - /// 显式声明重放安全性。 + /// 记录协议 adapter 已识别的图片生成工具意图。 #[must_use] - pub const fn with_retry_safety(mut self, retry_safety: RetrySafety) -> Self { - self.retry_safety = retry_safety; + pub const fn with_image_generation_requested(mut self, requested: bool) -> Self { + self.image_generation_requested = requested; self } @@ -855,10 +855,10 @@ impl GenerateRequest { self.estimated_context_tokens } - /// 返回显式重放安全性。 + /// 返回客户端是否请求了图片生成工具。 #[must_use] - pub const fn retry_safety(&self) -> RetrySafety { - self.retry_safety + pub const fn image_generation_requested(&self) -> bool { + self.image_generation_requested } /// 返回客户端的响应持久化意图。 @@ -939,7 +939,10 @@ impl fmt::Debug for GenerateRequest { &self.payload.prompt_cache_key.as_ref().map(|_| ""), ) .field("estimated_context_tokens", &self.estimated_context_tokens) - .field("retry_safety", &self.retry_safety) + .field( + "image_generation_requested", + &self.image_generation_requested, + ) .field("response_persistence", &self.response_persistence) .field("provider_options", &self.payload.provider_options) .field( @@ -1104,7 +1107,6 @@ impl fmt::Debug for RerankRequest { pub struct ImageRequest { prompt: String, count: u32, - retry_safety: RetrySafety, } impl ImageRequest { @@ -1118,11 +1120,7 @@ impl ImageRequest { if prompt.is_empty() { return Err(OperationError::EmptyField { field: "prompt" }); } - Ok(Self { - prompt, - count: 1, - retry_safety: RetrySafety::NonIdempotent, - }) + Ok(Self { prompt, count: 1 }) } /// 设置图像数量。 @@ -1139,7 +1137,6 @@ impl fmt::Debug for ImageRequest { .debug_struct("ImageRequest") .field("prompt", &"") .field("count", &self.count) - .field("retry_safety", &self.retry_safety) .finish() } } @@ -1149,7 +1146,6 @@ impl fmt::Debug for ImageRequest { pub struct SpeechRequest { input: String, voice: String, - retry_safety: RetrySafety, } impl SpeechRequest { @@ -1167,11 +1163,7 @@ impl SpeechRequest { if voice.is_empty() { return Err(OperationError::EmptyField { field: "voice" }); } - Ok(Self { - input, - voice, - retry_safety: RetrySafety::NonIdempotent, - }) + Ok(Self { input, voice }) } } @@ -1181,7 +1173,6 @@ impl fmt::Debug for SpeechRequest { .debug_struct("SpeechRequest") .field("input", &"") .field("voice", &self.voice) - .field("retry_safety", &self.retry_safety) .finish() } } @@ -1247,11 +1238,24 @@ impl Operation { #[must_use] pub const fn retry_safety(&self) -> RetrySafety { match self { - Self::Generate(request) => request.retry_safety, - Self::CompactConversation(request) => request.generation.retry_safety, Self::Embed(_) | Self::Rerank(_) => RetrySafety::Idempotent, - Self::GenerateImage(request) => request.retry_safety, - Self::Speech(request) => request.retry_safety, + Self::Generate(_) + | Self::CompactConversation(_) + | Self::GenerateImage(_) + | Self::Speech(_) => RetrySafety::NonIdempotent, + } + } + + /// 返回该 operation 是否代表一次图片生成请求。 + #[must_use] + pub const fn image_generation_requested(&self) -> bool { + match self { + Self::Generate(request) => request.image_generation_requested(), + Self::CompactConversation(_) + | Self::Embed(_) + | Self::Rerank(_) + | Self::GenerateImage(_) + | Self::Speech(_) => false, } } diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index 5e81561c1..16ab2116c 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -529,11 +529,6 @@ impl RuntimeSnapshot { self.client_policies.values() } - #[must_use] - pub fn client_policy(&self, key_id: &ClientApiKeyId) -> Option<&ClientPolicy> { - self.client_policies.get(key_id) - } - pub fn plan( &self, public_model: &PublicModelId, diff --git a/backend/crates/gateway-core/src/task/mod.rs b/backend/crates/gateway-core/src/task/mod.rs index 4c82e34f9..1e2bfa27d 100644 --- a/backend/crates/gateway-core/src/task/mod.rs +++ b/backend/crates/gateway-core/src/task/mod.rs @@ -18,19 +18,17 @@ pub enum WorkerKind { QuotaCatalogHealth, RuntimeSnapshotReconciliation, RuntimeChangeSubscription, - NativeClaimRecovery, StaleModelRequestRecovery, Retention, OpsFlush, } impl WorkerKind { - pub const ALL: [Self; 8] = [ + pub const ALL: [Self; 7] = [ Self::OAuthRefresh, Self::QuotaCatalogHealth, Self::RuntimeSnapshotReconciliation, Self::RuntimeChangeSubscription, - Self::NativeClaimRecovery, Self::StaleModelRequestRecovery, Self::Retention, Self::OpsFlush, @@ -43,7 +41,6 @@ impl WorkerKind { Self::QuotaCatalogHealth => "quota_catalog_health", Self::RuntimeSnapshotReconciliation => "runtime_snapshot_reconciliation", Self::RuntimeChangeSubscription => "runtime_change_subscription", - Self::NativeClaimRecovery => "native_claim_recovery", Self::StaleModelRequestRecovery => "stale_model_request_recovery", Self::Retention => "retention", Self::OpsFlush => "ops_flush", @@ -402,10 +399,9 @@ impl WorkerRegistration { } } -/// Final DB 明确没有相应持久化状态时允许的两种禁用原因。 +/// Final DB 明确没有相应持久化状态时允许的禁用原因。 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum WorkerDisabledReason { - NoPersistentNativeClaimState, NoBufferedOpsEvents, } @@ -413,7 +409,6 @@ impl WorkerDisabledReason { #[must_use] pub const fn kind(self) -> WorkerKind { match self { - Self::NoPersistentNativeClaimState => WorkerKind::NativeClaimRecovery, Self::NoBufferedOpsEvents => WorkerKind::OpsFlush, } } @@ -421,9 +416,6 @@ impl WorkerDisabledReason { #[must_use] pub const fn as_str(self) -> &'static str { match self { - Self::NoPersistentNativeClaimState => { - "native continuation has no persistent claim state" - } Self::NoBufferedOpsEvents => "ops events have no flush buffer", } } diff --git a/backend/crates/gateway-core/tests/accounting/mod.rs b/backend/crates/gateway-core/tests/accounting/mod.rs index 6c6d5df1a..e0865fb7a 100644 --- a/backend/crates/gateway-core/tests/accounting/mod.rs +++ b/backend/crates/gateway-core/tests/accounting/mod.rs @@ -77,33 +77,6 @@ fn unknown_cost_should_not_fabricate_zero() { assert_eq!(CostEstimate::unavailable().total(), None); } -#[test] -fn provider_reported_cost_should_preserve_exact_total() { - let total = Money::new( - Decimal::from_str("0.00125").expect("valid amount"), - CurrencyCode::new("USD").expect("valid currency"), - ); - let estimate = CostEstimate::priced( - CostEstimateStatus::Known, - CostSource::ProviderReported, - total, - ) - .expect("reported total is valid"); - - assert_eq!(estimate.total(), Some(total)); -} - -#[test] -fn unavailable_source_should_reject_amount() { - let total = Money::new( - Decimal::ZERO, - CurrencyCode::new("USD").expect("valid currency"), - ); - assert!( - CostEstimate::priced(CostEstimateStatus::Known, CostSource::Unavailable, total).is_err() - ); -} - #[test] fn usage_total_should_remain_independent() { let usage = Usage { diff --git a/backend/crates/gateway-core/tests/engine/continuation.rs b/backend/crates/gateway-core/tests/engine/continuation.rs index 7eabaab9c..5f7da7b2a 100644 --- a/backend/crates/gateway-core/tests/engine/continuation.rs +++ b/backend/crates/gateway-core/tests/engine/continuation.rs @@ -1,6 +1,5 @@ use gateway_core::engine::continuation::{ - ContinuationBinding, NativeClaim, NativeClaimState, NativeContinuationPin, - NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, PreviousResponseId, }; use gateway_core::engine::credential::ProviderAccountId; use gateway_core::error::SafeUpstreamValue; @@ -13,7 +12,6 @@ fn pin() -> NativeContinuationPin { ProviderKind::new("openai").expect("valid provider"), ProviderInstanceId::new("inst_openai").expect("valid instance"), ProviderAccountId::new("acct_codex").expect("valid account"), - NativeContinuationReuse::SingleUse, ) } @@ -41,29 +39,3 @@ fn native_pin_should_reject_different_account() { &ProviderAccountId::new("acct_other").expect("valid account"), )); } - -#[test] -fn single_use_claim_should_not_be_claimed_twice() { - let mut claim = NativeClaim::new(NativeContinuationReuse::SingleUse); - claim.claim().expect("first claim succeeds"); - - assert!(claim.claim().is_err()); -} - -#[test] -fn proven_not_sent_should_release_single_use_claim() { - let mut claim = NativeClaim::new(NativeContinuationReuse::SingleUse); - claim.claim().expect("claim succeeds"); - claim.release_not_sent(); - - assert_eq!(claim.state(), NativeClaimState::Available); -} - -#[test] -fn ambiguous_send_should_make_single_use_claim_terminal() { - let mut claim = NativeClaim::new(NativeContinuationReuse::SingleUse); - claim.claim().expect("claim succeeds"); - claim.mark_ambiguous(); - - assert_eq!(claim.state(), NativeClaimState::Ambiguous); -} diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 3da902b01..3b79888ee 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -10,7 +10,7 @@ use futures::executor::block_on; use gateway_core::accounting::{CalculatedCost, CostSource, ProviderReportedCost, Usage}; use gateway_core::engine::continuation::{ - ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, PreviousResponseId, }; use gateway_core::engine::credential::{ AccountSelectionPolicy, CredentialRevision, ProviderAccountId, RotationStrategy, @@ -30,10 +30,10 @@ use gateway_core::error::{ }; use gateway_core::event::{ GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseObservation, - ProviderResponseTimings, ResponseMeta, UpstreamHttpVersion, + ProviderResponseTimings, ResponseMeta, UpstreamHttpVersion, WebSocketPoolKind, }; use gateway_core::operation::{ - ContentPart, GenerateRequest, Message, MessageRole, Operation, OperationKind, + ContentPart, EmbedRequest, GenerateRequest, Message, MessageRole, Operation, ProviderSessionState, RetrySafety, }; use gateway_core::policy::ClientApiKeyId; @@ -52,6 +52,9 @@ struct FinalState { committed: bool, client_status_code: Option, total_tokens: Option, + image_input_tokens: Option, + image_output_tokens: Option, + image_generation_succeeded: Option, provider_error_code: Option, retry_after_ms: Option, latency_ms: Option, @@ -59,6 +62,7 @@ struct FinalState { upstream_response_id: Option, upstream_transport: Option, http_version: Option, + websocket_pool: Option, upstream_request_id: Option, upstream_status_code: Option, transport_decision_wait_ms: Option, @@ -172,6 +176,9 @@ impl ExecutionStore for FakeStore { committed: finalization.downstream_committed_at.is_some(), client_status_code: finalization.client_status_code, total_tokens: finalization.usage.total_tokens, + image_input_tokens: finalization.usage.image_input_tokens, + image_output_tokens: finalization.usage.image_output_tokens, + image_generation_succeeded: finalization.image_generation_succeeded, provider_error_code: finalization.provider_error_code, retry_after_ms: finalization.retry_after_ms, latency_ms: finalization.timings.latency_ms, @@ -179,6 +186,7 @@ impl ExecutionStore for FakeStore { upstream_response_id: finalization.upstream_response_id, upstream_transport: finalization.upstream_transport, http_version: finalization.http_version, + websocket_pool: finalization.websocket_pool, upstream_request_id: finalization.upstream_request_id, upstream_status_code: finalization.upstream_status_code, transport_decision_wait_ms: finalization.timings.transport_decision_wait_ms, @@ -304,16 +312,28 @@ impl Provider for ScriptedProvider { } fn operation(retry_safety: RetrySafety) -> Operation { + if retry_safety == RetrySafety::Idempotent { + return Operation::Embed( + EmbedRequest::new(vec!["hello".to_owned()]).expect("embedding request"), + ); + } + generate_operation() +} + +fn generate_operation() -> Operation { let message = Message::new( MessageRole::User, vec![ContentPart::Text("hello".to_owned())], ) .expect("message"); - Operation::Generate( - GenerateRequest::new(vec![message]) - .expect("generate request") - .with_retry_safety(retry_safety), - ) + Operation::Generate(GenerateRequest::new(vec![message]).expect("generate request")) +} + +fn image_generate_operation() -> Operation { + let Operation::Generate(request) = generate_operation() else { + unreachable!("generate_operation must return Generate") + }; + Operation::Generate(request.with_image_generation_requested(true)) } fn complete_stream(total_tokens: Option) -> Vec> { @@ -335,14 +355,31 @@ fn complete_stream(total_tokens: Option) -> Vec) -> Vec> { + let mut usage = Usage::new(); + usage.input_tokens = Some(12); + usage.output_tokens = Some(5); + usage.image_input_tokens = Some(31); + usage.image_output_tokens = image_output_tokens; + usage.total_tokens = Some(17); + vec![ + Ok(GatewayEvent::Started(ResponseMeta::new( + "response-image", + "gpt-5", + ))), + Ok(GatewayEvent::Usage(usage)), + Ok(GatewayEvent::Completed(ResponseMeta::new( + "response-image", + "gpt-5", + ))), + ] +} + fn plan(operation: &Operation, instance_count: u32) -> RoutingPlan { let provider = ProviderKind::new("openai").expect("provider"); let public_model = PublicModelId::new("gpt-5").expect("public model"); - let capabilities = ModelCapabilities::new( - BTreeSet::from([OperationKind::Generate]), - 128_000, - Some(32_000), - ); + let capabilities = + ModelCapabilities::new(BTreeSet::from([operation.kind()]), 128_000, Some(32_000)); let mut instances = Vec::new(); let mut provider_models = Vec::new(); for index in 1..=instance_count { @@ -384,7 +421,7 @@ fn plan(operation: &Operation, instance_count: u32) -> RoutingPlan { .expect("routing plan") } -fn model_request(deadline: SystemTime) -> NewModelRequest { +fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest { let client_key = ClientApiKeyId::new("key_client_1").expect("client key id"); NewModelRequest { id: ModelRequestId::new("req_core_1").expect("request id"), @@ -392,7 +429,7 @@ fn model_request(deadline: SystemTime) -> NewModelRequest { client_api_key_ref: client_key, config_revision: ConfigRevision::new(1).expect("config revision"), protocol: "openai".to_owned(), - operation: OperationKind::Generate, + operation: operation.kind(), endpoint: "responses".to_owned(), client_transport: "http_sse".to_owned(), requested_model: PublicModelId::new("gpt-5").expect("model"), @@ -404,6 +441,7 @@ fn model_request(deadline: SystemTime) -> NewModelRequest { request_kind: Some("responses".to_owned()), subagent_kind: None, compact: false, + image_generation_requested: operation.image_generation_requested(), started_at: SystemTime::now(), deadline_at: deadline, } @@ -443,7 +481,7 @@ fn terminal_non_idempotent_failure( }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -469,7 +507,7 @@ fn success_updates_one_model_request_and_persists_usage() { }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -521,6 +559,78 @@ fn success_updates_one_model_request_and_persists_usage() { ); } +fn finalized_image_request(image_output_tokens: Option) -> FinalState { + let operation = image_generate_operation(); + let route_plan = plan(&operation, 1); + let (coordinator, store, _) = coordinator(vec![Script::Stream { + account_id: "acct_image", + items: image_stream(image_output_tokens), + }]); + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start image execution"); + block_on(session.collect_uncommitted()).expect("collect image response"); + block_on(session.commit_downstream(Some(200))).expect("commit image response"); + + store.state.lock().expect("store lock").finalizations[0].clone() +} + +#[test] +fn image_request_with_output_should_persist_success_and_image_tokens() { + let finalization = finalized_image_request(Some(9)); + + assert_eq!( + ( + finalization.image_generation_succeeded, + finalization.image_input_tokens, + finalization.image_output_tokens, + ), + (Some(true), Some(31), Some(9)) + ); +} + +#[test] +fn image_request_without_output_should_persist_failure() { + let finalization = finalized_image_request(None); + + assert_eq!(finalization.image_generation_succeeded, Some(false)); +} + +#[test] +fn failed_image_request_should_persist_failure() { + let operation = image_generate_operation(); + let route_plan = plan(&operation, 1); + let (coordinator, store, _) = coordinator(vec![Script::Stream { + account_id: "acct_image", + items: vec![Err(ProviderError::new( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + ))], + }]); + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start image execution"); + block_on(session.collect_uncommitted()).expect_err("image request should fail"); + let state = store.state.lock().expect("store lock"); + + assert_eq!( + state.finalizations[0].image_generation_succeeded, + Some(false) + ); +} + #[test] fn response_observation_is_persisted_but_never_delivered() { let operation = operation(RetrySafety::Idempotent); @@ -529,6 +639,7 @@ fn response_observation_is_persisted_but_never_delivered() { UpstreamTransport::new("http_sse").expect("actual transport"), ) .with_http_version(UpstreamHttpVersion::Http2) + .with_websocket_pool(WebSocketPoolKind::New) .with_status_code(200) .with_request_id(SafeUpstreamValue::new("upstream-observed").expect("request id")) .with_timings(ProviderResponseTimings { @@ -548,7 +659,7 @@ fn response_observation_is_persisted_but_never_delivered() { items, }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -565,6 +676,7 @@ fn response_observation_is_persisted_but_never_delivered() { let finalization = &state.finalizations[0]; assert_eq!(finalization.upstream_transport.as_deref(), Some("http_sse")); assert_eq!(finalization.http_version.as_deref(), Some("HTTP/2")); + assert_eq!(finalization.websocket_pool.as_deref(), Some("new")); assert_eq!( finalization.upstream_request_id.as_deref(), Some("upstream-observed") @@ -603,7 +715,7 @@ fn unknown_wire_event_before_response_identity_is_discarded_with_retried_attempt }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -667,7 +779,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -712,7 +824,7 @@ fn websocket_success_keeps_client_http_status_absent() { items: complete_stream(None), }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -731,7 +843,7 @@ fn websocket_success_keeps_client_http_status_absent() { #[test] fn authenticated_native_continuation_reaches_every_attempt_context() { - let operation = operation(RetrySafety::Idempotent); + let operation = generate_operation(); let route_plan = plan(&operation, 1); let (coordinator, _, provider) = coordinator(vec![Script::Stream { account_id: "acct_one", @@ -743,11 +855,10 @@ fn authenticated_native_continuation_reaches_every_attempt_context() { ProviderKind::new("openai").expect("provider"), ProviderInstanceId::new("inst_openai_1").expect("instance"), ProviderAccountId::new("acct_one").expect("account"), - NativeContinuationReuse::Reusable, ); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -791,7 +902,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() { #[test] fn native_continuation_replays_owner_before_safely_switching_account() { - let Operation::Generate(generate) = operation(RetrySafety::Idempotent) else { + let Operation::Generate(generate) = generate_operation() else { panic!("generate operation"); }; let mut payload = Map::new(); @@ -834,11 +945,10 @@ fn native_continuation_replays_owner_before_safely_switching_account() { ProviderKind::new("openai").expect("provider"), ProviderInstanceId::new("inst_openai_1").expect("instance"), ProviderAccountId::new("acct_one").expect("account"), - NativeContinuationReuse::Reusable, ); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -880,7 +990,7 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() { }]); let required = ProviderAccountId::new("acct_required").expect("account id"); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, Some(required.clone()), @@ -911,7 +1021,7 @@ fn provider_metadata_for_another_account_fails_closed() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, Some(ProviderAccountId::new("acct_required").expect("account id")), @@ -949,7 +1059,7 @@ fn required_account_disables_instance_fallback_before_stream() { ]); let required = ProviderAccountId::new("acct_required").expect("account id"); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, Some(required.clone()), @@ -988,7 +1098,7 @@ fn local_selection_failure_falls_back_without_instance_failure_observation() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1034,7 +1144,7 @@ fn required_account_disables_account_retry_after_stream_creation() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, Some(ProviderAccountId::new("acct_required").expect("account id")), @@ -1089,7 +1199,7 @@ fn latest_provider_reported_cost_is_persisted_as_known_usd_total() { items: events, }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1129,7 +1239,7 @@ fn calculated_cost_is_persisted_when_provider_does_not_report_cost() { items: events, }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1173,7 +1283,7 @@ fn provider_reported_cost_should_not_be_replaced_by_calculated_cost() { items: events, }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1224,7 +1334,7 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1264,7 +1374,7 @@ fn pre_commit_failure_excludes_account_and_retries_same_target() { ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1326,7 +1436,7 @@ fn recovered_credential_retries_the_same_account_exactly_once() { ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1370,7 +1480,7 @@ fn non_idempotent_explicit_429_rejection_rotates_account_before_output() { ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1454,7 +1564,7 @@ fn explicit_429_after_structural_event_should_retry_before_commit() { }, ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1481,7 +1591,6 @@ fn native_continuation_explicit_429_is_not_retried() { ProviderKind::new("openai").expect("provider"), ProviderInstanceId::new("inst_openai_1").expect("instance"), ProviderAccountId::new("acct_first").expect("account"), - NativeContinuationReuse::Reusable, ); let (store, provider) = terminal_non_idempotent_failure( vec![Err(ProviderError::new( @@ -1539,7 +1648,7 @@ fn non_idempotent_not_sent_failure_can_fallback_target() { ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1581,7 +1690,7 @@ fn ambiguous_send_state_stops_retry() { ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1633,7 +1742,7 @@ fn structural_event_before_replay_safe_failure_should_switch_account_before_comm ]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1676,7 +1785,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() { items: complete_stream(None), }]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1712,7 +1821,7 @@ fn local_unavailability_before_stream_does_not_create_attempt_or_instance_failur )]); let mut session = block_on(coordinator.start( - model_request(SystemTime::now() + Duration::from_secs(30)), + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, None, @@ -1745,7 +1854,7 @@ fn expired_deadline_finalizes_without_calling_provider() { let (coordinator, store, provider) = coordinator(vec![]); let error = match block_on(coordinator.start( - model_request(SystemTime::UNIX_EPOCH), + model_request(&operation, SystemTime::UNIX_EPOCH), operation, route_plan, None, diff --git a/backend/crates/gateway-core/tests/operation/mod.rs b/backend/crates/gateway-core/tests/operation/mod.rs index ed347f71f..4551daed8 100644 --- a/backend/crates/gateway-core/tests/operation/mod.rs +++ b/backend/crates/gateway-core/tests/operation/mod.rs @@ -322,10 +322,13 @@ fn embed_and_rerank_operations_should_be_idempotent() { } #[test] -fn compact_conversation_should_inherit_generation_retry_safety() { +fn generate_and_compact_conversation_should_be_non_idempotent() { let generate = GenerateRequest::new(vec![text_message(MessageRole::User, "history")]) .expect("compaction history is valid"); - assert_eq!(generate.retry_safety(), RetrySafety::NonIdempotent); + assert_eq!( + Operation::Generate(generate.clone()).retry_safety(), + RetrySafety::NonIdempotent + ); let compact = Operation::CompactConversation(CompactConversationRequest::new(generate)); assert_eq!(compact.retry_safety(), RetrySafety::NonIdempotent); diff --git a/backend/crates/gateway-core/tests/task/mod.rs b/backend/crates/gateway-core/tests/task/mod.rs index ec85a5e3c..a4a21ab8a 100644 --- a/backend/crates/gateway-core/tests/task/mod.rs +++ b/backend/crates/gateway-core/tests/task/mod.rs @@ -4,9 +4,9 @@ use std::time::Duration; use futures::future::BoxFuture; use gateway_core::engine::CancellationToken; use gateway_core::task::{ - DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext, - WorkerDefinitionError, WorkerDisabledReason, WorkerFencingToken, WorkerId, WorkerKind, - WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule, WorkerTaskError, + DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerCycleContext, WorkerDefinitionError, + WorkerFencingToken, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration, + WorkerRunnable, WorkerSchedule, WorkerTaskError, }; struct NoopScheduledTask; @@ -116,19 +116,6 @@ fn registration_rejects_lease_ttl_that_differs_from_schedule() { )); } -#[test] -fn disabled_contribution_rejects_reason_for_another_kind() { - let contribution = WorkerContribution::Disabled { - kind: WorkerKind::OpsFlush, - reason: WorkerDisabledReason::NoPersistentNativeClaimState, - }; - - assert_eq!( - contribution.validate(), - Err(WorkerDefinitionError::DisabledReasonMismatch) - ); -} - #[test] fn daemon_restart_policy_rejects_reverse_backoff_range() { assert_eq!( diff --git a/backend/crates/gateway-host/tests/workers/mod.rs b/backend/crates/gateway-host/tests/workers/mod.rs index 83ee519f0..81a33f6ac 100644 --- a/backend/crates/gateway-host/tests/workers/mod.rs +++ b/backend/crates/gateway-host/tests/workers/mod.rs @@ -344,10 +344,10 @@ async fn registry_composes_multiple_owners_but_rejects_duplicate_identity() { } #[tokio::test] -async fn only_database_fact_reasons_can_disable_the_two_nonexistent_workers() { +async fn only_database_fact_reasons_can_disable_nonexistent_workers() { let mut plan = complete_plan(Vec::new()); plan.push(WorkerContribution::Registration(registration_value( - WorkerKind::NativeClaimRecovery, + WorkerKind::OpsFlush, "store", CountingTask::new(0), test_schedule(), @@ -356,9 +356,7 @@ async fn only_database_fact_reasons_can_disable_the_two_nonexistent_workers() { assert!(matches!( supervisor().start(plan, Arc::new(FakeLeasePort::default())), - Err(WorkerStartError::KindDisabled( - WorkerKind::NativeClaimRecovery - )) + Err(WorkerStartError::KindDisabled(WorkerKind::OpsFlush)) )); } @@ -386,15 +384,15 @@ async fn disabled_database_facts_never_execute_or_acquire_a_lease() { let snapshot = supervisor.health_source().snapshot(); assert!(snapshot.iter().any(|state| { - state.key == WorkerHealthKey::Disabled(WorkerKind::NativeClaimRecovery) + state.key == WorkerHealthKey::Disabled(WorkerKind::OpsFlush) && state.state == WorkerRuntimeState::Disabled })); - assert!(leases.requests().iter().all(|request| { - !matches!( - request.worker().kind(), - WorkerKind::NativeClaimRecovery | WorkerKind::OpsFlush - ) - })); + assert!( + leases + .requests() + .iter() + .all(|request| { request.worker().kind() != WorkerKind::OpsFlush }) + ); supervisor.shutdown(Duration::from_secs(1)).await; } @@ -704,10 +702,6 @@ fn complete_plan(mut custom: Vec) -> Vec )); } } - custom.push(WorkerContribution::Disabled { - kind: WorkerKind::NativeClaimRecovery, - reason: WorkerDisabledReason::NoPersistentNativeClaimState, - }); custom.push(WorkerContribution::Disabled { kind: WorkerKind::OpsFlush, reason: WorkerDisabledReason::NoBufferedOpsEvents, @@ -733,10 +727,6 @@ fn all_active_plan(task: impl ScheduledTask + Clone + 'static) -> Vec>(); - plan.push(WorkerContribution::Disabled { - kind: WorkerKind::NativeClaimRecovery, - reason: WorkerDisabledReason::NoPersistentNativeClaimState, - }); plan.push(WorkerContribution::Disabled { kind: WorkerKind::OpsFlush, reason: WorkerDisabledReason::NoBufferedOpsEvents, diff --git a/backend/crates/gateway-protocol/src/openai/mod.rs b/backend/crates/gateway-protocol/src/openai/mod.rs index 676b5ab82..92ee03347 100644 --- a/backend/crates/gateway-protocol/src/openai/mod.rs +++ b/backend/crates/gateway-protocol/src/openai/mod.rs @@ -11,7 +11,5 @@ pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request"; /// OpenAI/Codex 事件语义、用量与限流字段编解码。 pub mod events; -/// OpenAI/Codex JSON schema 纯转换。 -pub mod schema; /// Server-Sent Events 帧的解析与编码。 pub mod sse; diff --git a/backend/crates/gateway-protocol/src/openai/schema.rs b/backend/crates/gateway-protocol/src/openai/schema.rs deleted file mode 100644 index 9c4f7ace6..000000000 --- a/backend/crates/gateway-protocol/src/openai/schema.rs +++ /dev/null @@ -1,129 +0,0 @@ -//! Codex 协议所需的 JSON schema 纯处理逻辑。 - -use serde_json::{Map, Value}; - -pub fn reconvert_tuple_values(data: Value, schema: &Value) -> Value { - reconvert_tuple_values_with_root(data, schema, schema) -} - -fn reconvert_tuple_values_with_root(data: Value, schema: &Value, root_schema: &Value) -> Value { - let Value::Object(schema_object) = schema else { - return data; - }; - - if let Some(reference) = schema_object.get("$ref").and_then(Value::as_str) { - let Some(resolved) = resolve_ref(reference, root_schema) else { - return data; - }; - return reconvert_tuple_values_with_root(data, resolved, root_schema); - } - - if let Some(Value::Array(prefix_items)) = schema_object.get("prefixItems") { - let mut data_object = match data { - Value::Object(data_object) => data_object, - other => return other, - }; - let values = prefix_items - .iter() - .enumerate() - .map(|(index, item_schema)| { - let value = data_object - .remove(&index.to_string()) - .unwrap_or(Value::Null); - reconvert_tuple_values_with_root(value, item_schema, root_schema) - }) - .collect(); - return Value::Array(values); - } - - if let Some(Value::Object(properties)) = schema_object.get("properties") { - let mut data_object = match data { - Value::Object(data_object) => data_object, - other => return other, - }; - for (key, property_schema) in properties { - if let Some(value) = data_object.remove(key) { - data_object.insert( - key.clone(), - reconvert_tuple_values_with_root(value, property_schema, root_schema), - ); - } - } - return Value::Object(data_object); - } - - if let Some(items_schema) = schema_object.get("items") { - let values = match data { - Value::Array(values) => values, - other => return other, - }; - return Value::Array( - values - .into_iter() - .map(|value| reconvert_tuple_values_with_root(value, items_schema, root_schema)) - .collect(), - ); - } - - for key in ["oneOf", "anyOf", "allOf"] { - let Some(Value::Array(branches)) = schema_object.get(key) else { - continue; - }; - let Some(branch) = branches.iter().find(|branch| has_tuple_schemas(branch)) else { - continue; - }; - return reconvert_tuple_values_with_root(data, branch, root_schema); - } - - data -} - -fn resolve_ref<'a>(reference: &str, root_schema: &'a Value) -> Option<&'a Value> { - let path = reference - .strip_prefix("#/$defs/") - .map(|path| ("$defs", path)) - .or_else(|| { - reference - .strip_prefix("#/definitions/") - .map(|path| ("definitions", path)) - })?; - root_schema.get(path.0)?.get(path.1) -} - -fn has_tuple_schemas(value: &Value) -> bool { - let Value::Object(object) = value else { - return false; - }; - if object - .get("prefixItems") - .is_some_and(serde_json::Value::is_array) - { - return true; - } - - object_values_have_tuple_schema(object, "properties") - || object.get("items").is_some_and(has_tuple_schemas) - || schema_array_has_tuple_schema(object, "oneOf") - || schema_array_has_tuple_schema(object, "anyOf") - || schema_array_has_tuple_schema(object, "allOf") - || object_values_have_tuple_schema(object, "$defs") - || object_values_have_tuple_schema(object, "definitions") - || object.get("if").is_some_and(has_tuple_schemas) - || object.get("then").is_some_and(has_tuple_schemas) - || object.get("else").is_some_and(has_tuple_schemas) - || object.get("not").is_some_and(has_tuple_schemas) -} - -fn object_values_have_tuple_schema(object: &Map, key: &str) -> bool { - let Some(Value::Object(values)) = object.get(key) else { - return false; - }; - values.values().any(has_tuple_schemas) -} - -fn schema_array_has_tuple_schema(object: &Map, key: &str) -> bool { - let Some(Value::Array(values)) = object.get(key) else { - return false; - }; - values.iter().any(has_tuple_schemas) -} diff --git a/backend/crates/gateway-protocol/tests/openai/mod.rs b/backend/crates/gateway-protocol/tests/openai/mod.rs index b8ba697cb..395b02310 100644 --- a/backend/crates/gateway-protocol/tests/openai/mod.rs +++ b/backend/crates/gateway-protocol/tests/openai/mod.rs @@ -1,3 +1,2 @@ mod events; -mod schema; mod sse; diff --git a/backend/crates/gateway-protocol/tests/openai/schema.rs b/backend/crates/gateway-protocol/tests/openai/schema.rs deleted file mode 100644 index a689c44be..000000000 --- a/backend/crates/gateway-protocol/tests/openai/schema.rs +++ /dev/null @@ -1,89 +0,0 @@ -use gateway_protocol::openai::schema::reconvert_tuple_values; -use serde_json::json; - -#[test] -fn tuple_reconversion_should_restore_nested_arrays_without_touching_other_properties() { - let schema = json!({ - "type": "object", - "properties": { - "point": { - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ] - }, - "label": {"type": "string"} - } - }); - let encoded = json!({"point": {"0": 12.5, "1": 9.0}, "label": "A"}); - - assert_eq!( - reconvert_tuple_values(encoded, &schema), - json!({"point": [12.5, 9.0], "label": "A"}) - ); -} - -#[test] -fn tuple_reconversion_should_resolve_local_defs_reference() { - let schema = json!({ - "$defs": { - "pair": { - "prefixItems": [ - {"type": "string"}, - {"type": "integer"} - ] - } - }, - "properties": { - "entry": {"$ref": "#/$defs/pair"} - } - }); - - assert_eq!( - reconvert_tuple_values(json!({"entry": {"0": "age", "1": 42}}), &schema), - json!({"entry": ["age", 42]}) - ); -} - -#[test] -fn tuple_reconversion_should_recurse_through_homogeneous_arrays() { - let schema = json!({ - "type": "array", - "items": { - "prefixItems": [ - {"type": "integer"}, - {"type": "integer"} - ] - } - }); - - assert_eq!( - reconvert_tuple_values(json!([{"0": 1, "1": 2}, {"0": 3, "1": 4}]), &schema,), - json!([[1, 2], [3, 4]]) - ); -} - -#[test] -fn tuple_reconversion_should_select_union_branch_containing_tuple_schema() { - let schema = json!({ - "anyOf": [ - {"type": "string"}, - {"prefixItems": [{"type": "boolean"}, {"type": "string"}]} - ] - }); - - assert_eq!( - reconvert_tuple_values(json!({"0": true, "1": "ready"}), &schema), - json!([true, "ready"]) - ); -} - -#[test] -fn tuple_reconversion_should_preserve_values_when_schema_does_not_describe_a_tuple() { - let value = json!({"0": "not", "1": "a tuple", "extra": true}); - - assert_eq!( - reconvert_tuple_values(value.clone(), &json!({"type": "object"})), - value - ); -} diff --git a/backend/crates/gateway-store/src/lib.rs b/backend/crates/gateway-store/src/lib.rs index 4cab7bf0e..6ff30ae0c 100644 --- a/backend/crates/gateway-store/src/lib.rs +++ b/backend/crates/gateway-store/src/lib.rs @@ -346,10 +346,6 @@ fn store_worker_contributions( let retention_id = WorkerId::try_new(WorkerKind::Retention, "postgres").map_err(worker_definition_error)?; Ok(vec![ - WorkerContribution::Disabled { - kind: WorkerKind::NativeClaimRecovery, - reason: WorkerDisabledReason::NoPersistentNativeClaimState, - }, WorkerContribution::Registration(scheduled_worker( stale_id, Duration::from_secs(30), diff --git a/backend/crates/gateway-store/src/postgres/execution.rs b/backend/crates/gateway-store/src/postgres/execution.rs index d4c36e20e..2080ce8f8 100644 --- a/backend/crates/gateway-store/src/postgres/execution.rs +++ b/backend/crates/gateway-store/src/postgres/execution.rs @@ -95,6 +95,7 @@ pub struct NewModelRequest { pub request_kind: Option, pub subagent_kind: Option, pub compact: bool, + pub image_generation_requested: bool, pub started_at: DateTime, pub deadline_at: DateTime, } @@ -171,6 +172,8 @@ pub struct ModelRequestUsage { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, pub total_tokens: Option, } @@ -222,11 +225,13 @@ pub struct ModelRequestFinalization { pub upstream_response_id: Option, pub upstream_transport: Option, pub http_version: Option, + pub websocket_pool: Option, pub error_kind: Option, pub provider_error_code: Option, pub error_message: Option, pub retry_after_ms: Option, pub usage: ModelRequestUsage, + pub image_generation_succeeded: Option, pub cost_source: CostSource, pub cost_amount: Option, pub cost_currency: Option, @@ -265,6 +270,13 @@ impl ModelRequestFinalization { require_nonempty(ENTITY, field, value)?; } } + if self + .websocket_pool + .as_deref() + .is_some_and(|kind| !matches!(kind, "new" | "reuse")) + { + return Err(invalid("websocket pool must be new or reuse")); + } self.timings.validate() } } @@ -333,11 +345,11 @@ impl ModelRequestRepository for PgExecutionStore { id, client_api_key_id, client_api_key_ref, config_revision, protocol, operation, endpoint, client_transport, requested_model_id, input_token_estimate, client_ip, user_agent, reasoning_effort, - reasoning_preset, request_kind, subagent_kind, compact, started_at, - deadline_at + reasoning_preset, request_kind, subagent_kind, compact, + image_generation_requested, started_at, deadline_at ) values ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::inet, $12, $13, - $14, $15, $16, $17, $18, $19 + $14, $15, $16, $17, $18, $19, $20 )", ) .bind(request.id) @@ -360,6 +372,7 @@ impl ModelRequestRepository for PgExecutionStore { .bind(request.request_kind) .bind(request.subagent_kind) .bind(request.compact) + .bind(request.image_generation_requested) .bind(request.started_at) .bind(request.deadline_at) .execute(&self.pool) @@ -486,13 +499,15 @@ impl ModelRequestRepository for PgExecutionStore { error_kind = $11, provider_error_code = $12, error_message = $13, retry_after_ms = $14, input_tokens = $15, output_tokens = $16, cached_tokens = $17, cache_write_tokens = $18, reasoning_tokens = $19, - total_tokens = $20, cost_source = $21, cost_amount = $22::numeric, - cost_currency = $23, transport_decision_wait_ms = $24, connect_ms = $25, - headers_ms = $26, first_event_ms = $27, first_reasoning_ms = $28, - first_text_ms = $29, first_token_ms = $30, provider_processing_ms = $31, - latency_ms = $32, completed_at = $33, - upstream_transport = coalesce($34, upstream_transport), - http_version = coalesce($35, http_version) + image_input_tokens = $20, image_output_tokens = $21, total_tokens = $22, + image_generation_succeeded = $23, cost_source = $24, + cost_amount = $25::numeric, cost_currency = $26, + transport_decision_wait_ms = $27, connect_ms = $28, + headers_ms = $29, first_event_ms = $30, first_reasoning_ms = $31, + first_text_ms = $32, first_token_ms = $33, provider_processing_ms = $34, + latency_ms = $35, completed_at = $36, + upstream_transport = coalesce($37, upstream_transport), + http_version = coalesce($38, http_version), websocket_pool = $39 where id = $1 and outcome = 'running'", ) .bind(&finalization.model_request_id) @@ -532,10 +547,19 @@ impl ModelRequestRepository for PgExecutionStore { finalization.usage.reasoning_tokens, "reasoning_tokens", )?) + .bind(optional_i64( + finalization.usage.image_input_tokens, + "image_input_tokens", + )?) + .bind(optional_i64( + finalization.usage.image_output_tokens, + "image_output_tokens", + )?) .bind(optional_i64( finalization.usage.total_tokens, "total_tokens", )?) + .bind(finalization.image_generation_succeeded) .bind(finalization.cost_source.as_str()) .bind(finalization.cost_amount.map(|amount| amount.to_string())) .bind(finalization.cost_currency) @@ -569,6 +593,7 @@ impl ModelRequestRepository for PgExecutionStore { .bind(finalization.completed_at) .bind(finalization.upstream_transport) .bind(finalization.http_version) + .bind(finalization.websocket_pool) .execute(&self.pool) .await .map_err(|_| postgres_unavailable("finalize model request"))?; @@ -582,7 +607,11 @@ impl ModelRequestRepository for PgExecutionStore { let result = sqlx::query( "update model_requests set outcome = 'incomplete', error_kind = 'process_interrupted', - error_message = 'request did not reach a terminal state', completed_at = $1 + error_message = 'request did not reach a terminal state', + image_generation_succeeded = case + when image_generation_requested then false else null + end, + completed_at = $1 where outcome = 'running' and deadline_at <= $1", ) .bind(now) @@ -622,6 +651,7 @@ impl ExecutionStore for PgExecutionStore { request_kind: request.request_kind, subagent_kind: request.subagent_kind, compact: request.compact, + image_generation_requested: request.image_generation_requested, started_at: DateTime::::from(request.started_at), deadline_at: DateTime::::from(request.deadline_at), }) @@ -790,11 +820,13 @@ impl ExecutionStore for PgExecutionStore { upstream_response_id: finalization.upstream_response_id, upstream_transport: finalization.upstream_transport, http_version: finalization.http_version, + websocket_pool: finalization.websocket_pool, error_kind, provider_error_code: finalization.provider_error_code, error_message, retry_after_ms: finalization.retry_after_ms, usage: usage_from_core(finalization.usage), + image_generation_succeeded: finalization.image_generation_succeeded, cost_source, cost_amount, cost_currency, @@ -838,6 +870,8 @@ fn usage_from_core(usage: CoreUsage) -> ModelRequestUsage { cached_tokens: usage.cached_tokens, cache_write_tokens: usage.cache_write_tokens, reasoning_tokens: usage.reasoning_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, total_tokens: usage.total_tokens, } } diff --git a/backend/crates/gateway-store/src/postgres/history.rs b/backend/crates/gateway-store/src/postgres/history.rs index 7955c4dfb..dc2576d87 100644 --- a/backend/crates/gateway-store/src/postgres/history.rs +++ b/backend/crates/gateway-store/src/postgres/history.rs @@ -6,10 +6,7 @@ use gateway_core::engine::continuation::{NativeContinuationPort, NativeContinuat use gateway_core::policy::ClientApiKeyId; use gateway_core::{ engine::{ - continuation::{ - NativeContinuationPin, NativeContinuationReuse, NativeContinuationScope, - PreviousResponseId, - }, + continuation::{NativeContinuationPin, NativeContinuationScope, PreviousResponseId}, credential::ProviderAccountId, }, error::SafeUpstreamValue, @@ -52,7 +49,6 @@ pub trait ModelRequestHistoryRepository: Send + Sync { &self, client_response_id: &str, caller_client_api_key_ref: &str, - reuse: NativeContinuationReuse, ) -> StoreResult>; } @@ -105,7 +101,6 @@ impl ModelRequestHistoryRepository for PgHistoryRepository { &self, client_response_id: &str, caller_client_api_key_ref: &str, - reuse: NativeContinuationReuse, ) -> StoreResult> { require_nonempty( "native continuation", @@ -139,7 +134,7 @@ impl ModelRequestHistoryRepository for PgHistoryRepository { .fetch_optional(&self.pool) .await .map_err(|_| postgres_unavailable("resolve native continuation pin"))?; - row.map(|row| native_pin_from_row(client_response_id, row, reuse)) + row.map(|row| native_pin_from_row(client_response_id, row)) .transpose() } } @@ -157,7 +152,6 @@ impl NativeContinuationPort for PgHistoryRepository { self.resolve_native_continuation_pin( previous_response_id.as_str(), client_api_key_id.as_str(), - NativeContinuationReuse::Reusable, ) .await .map_err(|_| NativeContinuationStoreError) @@ -245,7 +239,6 @@ fn history_from_row(row: HistoryRow) -> ModelRequestHistoryRecord { fn native_pin_from_row( client_response_id: &str, row: (String, String, String, String), - reuse: NativeContinuationReuse, ) -> StoreResult { let previous_response_id = PreviousResponseId::new(client_response_id.to_owned()) .map_err(|_| invalid_native_pin("invalid client response ID"))?; @@ -263,7 +256,6 @@ fn native_pin_from_row( provider, instance, account, - reuse, ) .with_scope(NativeContinuationScope::Persisted)) } diff --git a/backend/crates/gateway-store/src/postgres/observability.rs b/backend/crates/gateway-store/src/postgres/observability.rs index 08109c60c..8a5518fd1 100644 --- a/backend/crates/gateway-store/src/postgres/observability.rs +++ b/backend/crates/gateway-store/src/postgres/observability.rs @@ -391,6 +391,10 @@ pub struct ProviderAccountUsageObservation { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: CostCoverage, pub costs: Vec, @@ -408,6 +412,10 @@ pub struct ProviderAccountModelUsageObservation { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, + pub image_request_count: u64, + pub image_request_failed_count: u64, pub total_tokens: Option, pub cost_coverage: CostCoverage, pub costs: Vec, @@ -479,6 +487,7 @@ pub struct UsageRecord { pub upstream_model_id: Option, pub upstream_transport: Option, pub http_version: Option, + pub websocket_pool: Option, pub attempt_count: u32, pub upstream_send_state: String, pub downstream_committed_at: Option>, @@ -497,6 +506,8 @@ pub struct UsageRecord { pub cached_tokens: Option, pub cache_write_tokens: Option, pub reasoning_tokens: Option, + pub image_input_tokens: Option, + pub image_output_tokens: Option, pub total_tokens: Option, pub cost_source: String, pub cost_amount: Option, @@ -517,6 +528,8 @@ pub struct UsageRecord { pub request_kind: Option, pub subagent_kind: Option, pub compact: bool, + pub image_generation_requested: bool, + pub image_generation_succeeded: Option, pub started_at: DateTime, pub deadline_at: DateTime, pub completed_at: Option>, @@ -1190,6 +1203,10 @@ fn admin_dashboard_account_usage( cached_tokens: usage.cached_tokens, cache_write_tokens: usage.cache_write_tokens, reasoning_tokens: usage.reasoning_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, total_tokens: usage.total_tokens, cost_coverage: admin_cost_coverage(usage.cost_coverage), costs: admin_currency_costs(usage.costs)?, @@ -1214,6 +1231,10 @@ fn admin_dashboard_account_model_usage( cached_tokens: usage.cached_tokens, cache_write_tokens: usage.cache_write_tokens, reasoning_tokens: usage.reasoning_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, total_tokens: usage.total_tokens, cost_coverage: admin_cost_coverage(usage.cost_coverage), costs: admin_currency_costs(usage.costs)?, @@ -1297,6 +1318,7 @@ fn admin_usage_record(record: UsageRecord) -> AdminStoreResult AdminStoreResult AdminStoreResult StoreResult StoreResult StoreResult AdminStoreResu cached_tokens: usage.cached_tokens, cache_write_tokens: usage.cache_write_tokens, reasoning_tokens: usage.reasoning_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, total_tokens: usage.total_tokens, cost_coverage: admin_account_cost_coverage(usage.cost_coverage), costs: admin_account_costs(usage.costs)?, @@ -1788,6 +1792,10 @@ fn admin_account_model_usage( cached_tokens: usage.cached_tokens, cache_write_tokens: usage.cache_write_tokens, reasoning_tokens: usage.reasoning_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, total_tokens: usage.total_tokens, cost_coverage: admin_account_cost_coverage(usage.cost_coverage), costs: admin_account_costs(usage.costs)?, diff --git a/backend/crates/gateway-store/src/redis/runtime_change.rs b/backend/crates/gateway-store/src/redis/runtime_change.rs index ff212d74b..99633b979 100644 --- a/backend/crates/gateway-store/src/redis/runtime_change.rs +++ b/backend/crates/gateway-store/src/redis/runtime_change.rs @@ -11,34 +11,13 @@ use gateway_core::routing::{ use redis::AsyncCommands; use serde::{Deserialize, Serialize}; -use crate::{Revision, StoreError, StoreResult, redis_unavailable, require_nonempty}; +use crate::{Revision, StoreError, StoreResult, redis_unavailable}; -use super::{namespace, resource_fingerprint}; +use super::namespace; #[derive(Debug, Clone, PartialEq, Eq)] pub enum RuntimeChange { - SnapshotPublished { - config_revision: Revision, - }, - ProviderAccountChanged { - provider_account_fingerprint: String, - credential_revision: Revision, - }, -} - -impl RuntimeChange { - pub fn provider_account_changed( - provider_account_id: &str, - credential_revision: Revision, - ) -> StoreResult { - Ok(Self::ProviderAccountChanged { - provider_account_fingerprint: resource_fingerprint( - "runtime change", - provider_account_id, - )?, - credential_revision, - }) - } + SnapshotPublished { config_revision: Revision }, } pub type RuntimeChangeSubscription = @@ -134,7 +113,6 @@ impl SnapshotSubscriptionPort for RedisRuntimeChangeRepository { ConfigRevision::new(config_revision.get()) .map_err(|_| SnapshotSubscriptionError::unavailable()), ), - Ok(RuntimeChange::ProviderAccountChanged { .. }) => None, Err(_) => Some(Err(SnapshotSubscriptionError::unavailable())), } }); @@ -147,7 +125,6 @@ impl SnapshotSubscriptionPort for RedisRuntimeChangeRepository { struct RuntimeChangeWire { kind: String, revision: u64, - account: Option, } impl From<&RuntimeChange> for RuntimeChangeWire { @@ -156,15 +133,6 @@ impl From<&RuntimeChange> for RuntimeChangeWire { RuntimeChange::SnapshotPublished { config_revision } => Self { kind: "snapshot".to_owned(), revision: config_revision.get(), - account: None, - }, - RuntimeChange::ProviderAccountChanged { - provider_account_fingerprint, - credential_revision, - } => Self { - kind: "provider_account".to_owned(), - revision: credential_revision.get(), - account: Some(provider_account_fingerprint.clone()), }, } } @@ -176,19 +144,9 @@ impl TryFrom for RuntimeChange { fn try_from(value: RuntimeChangeWire) -> Result { let revision = Revision::new(value.revision)?; match value.kind.as_str() { - "snapshot" if value.account.is_none() => Ok(Self::SnapshotPublished { + "snapshot" => Ok(Self::SnapshotPublished { config_revision: revision, }), - "provider_account" => { - let account = value - .account - .ok_or_else(|| invalid("provider account fingerprint is missing"))?; - require_nonempty("runtime change", "account", &account)?; - Ok(Self::ProviderAccountChanged { - provider_account_fingerprint: account, - credential_revision: revision, - }) - } _ => Err(invalid("runtime change kind is invalid")), } } diff --git a/backend/crates/gateway-store/tests/postgres/execution.rs b/backend/crates/gateway-store/tests/postgres/execution.rs index 068a44001..5a02a0631 100644 --- a/backend/crates/gateway-store/tests/postgres/execution.rs +++ b/backend/crates/gateway-store/tests/postgres/execution.rs @@ -1,5 +1,5 @@ use chrono::{Duration, Utc}; -use gateway_core::accounting::{CalculatedCost, Usage}; +use gateway_core::accounting::{CalculatedCost, CostEstimate, Usage}; use gateway_core::engine::{ ExecutionOutcome, ExecutionStore, ModelRequestFinalization as CoreModelRequestFinalization, ModelRequestId, ModelRequestTimings as CoreModelRequestTimings, UpstreamSendState, @@ -35,6 +35,7 @@ fn model_request_rejects_mismatched_client_key_live_id() { request_kind: None, subagent_kind: None, compact: false, + image_generation_requested: false, started_at, deadline_at: started_at + Duration::seconds(30), }; @@ -136,36 +137,17 @@ async fn core_adapter_should_persist_calculated_cost_exactly() { .expect("seed model request attempt"); let repository = PgExecutionStore::new(database.pool.clone()); - ExecutionStore::finalize_model_request( - &repository, - CoreModelRequestFinalization { - request_id: ModelRequestId::new("req_calculated_cost").expect("request id"), - outcome: ExecutionOutcome::Succeeded, - send_state: UpstreamSendState::Sent, - attempt_count: 1, - downstream_committed_at: None, - client_status_code: Some(200), - client_response_id: None, - upstream_status_code: Some(200), - upstream_request_id: None, - upstream_response_id: None, - upstream_transport: Some("websocket".to_owned()), - http_version: Some("HTTP/2".to_owned()), - error: None, - provider_error_code: None, - retry_after_ms: None, - usage: Usage::new(), - cost: CalculatedCost::from_usd_ticks(12_345) - .expect("calculated cost") - .into_estimate(), - timings: CoreModelRequestTimings::default(), - completed_at: std::time::SystemTime::now(), - }, - ) - .await - .expect("persist calculated cost"); - let persisted: (String, String, String, String, String) = sqlx::query_as( - "select cost_source, cost_amount::text, cost_currency, upstream_transport, http_version + let mut finalization = successful_core_finalization("req_calculated_cost"); + finalization.websocket_pool = Some("reuse".to_owned()); + finalization.cost = CalculatedCost::from_usd_ticks(12_345) + .expect("calculated cost") + .into_estimate(); + ExecutionStore::finalize_model_request(&repository, finalization) + .await + .expect("persist calculated cost"); + let persisted: (String, String, String, String, String, String) = sqlx::query_as( + "select cost_source, cost_amount::text, cost_currency, upstream_transport, http_version, + websocket_pool from model_requests where id = 'req_calculated_cost'", ) .fetch_one(&database.pool) @@ -180,11 +162,120 @@ async fn core_adapter_should_persist_calculated_cost_exactly() { "USD".to_owned(), "websocket".to_owned(), "HTTP/2".to_owned(), + "reuse".to_owned(), ) ); database.close().await; } +#[tokio::test] +async fn core_adapter_should_persist_image_result_and_new_websocket_pool() { + let Some(database) = TestDatabase::create("execution_image_usage").await else { + return; + }; + seed_running_request(&database.pool, "req_image_usage") + .await + .expect("seed image request"); + sqlx::query( + "update model_requests + set image_generation_requested = true, provider_kind = 'openai', + provider_account_ref = 'acct_openai', upstream_model_id = 'gpt-image', + upstream_transport = 'websocket', attempt_count = 1 + where id = 'req_image_usage'", + ) + .execute(&database.pool) + .await + .expect("seed image attempt"); + let repository = PgExecutionStore::new(database.pool.clone()); + let mut usage = Usage::new(); + usage.image_input_tokens = Some(31); + usage.image_output_tokens = Some(9); + let mut finalization = successful_core_finalization("req_image_usage"); + finalization.usage = usage; + finalization.image_generation_succeeded = Some(true); + finalization.websocket_pool = Some("new".to_owned()); + + ExecutionStore::finalize_model_request(&repository, finalization) + .await + .expect("persist image usage"); + let persisted: (bool, Option, Option, Option, Option) = sqlx::query_as( + "select image_generation_requested, image_input_tokens, image_output_tokens, + image_generation_succeeded, websocket_pool + from model_requests where id = 'req_image_usage'", + ) + .fetch_one(&database.pool) + .await + .expect("load image usage"); + + assert_eq!( + persisted, + (true, Some(31), Some(9), Some(true), Some("new".to_owned())) + ); + database.close().await; +} + +#[tokio::test] +async fn expired_image_request_should_be_recovered_as_failed() { + let Some(database) = TestDatabase::create("execution_expired_image").await else { + return; + }; + seed_running_request(&database.pool, "req_expired_image") + .await + .expect("seed image request"); + sqlx::query( + "update model_requests + set image_generation_requested = true, + started_at = now() - interval '2 seconds', + deadline_at = now() - interval '1 second' + where id = 'req_expired_image'", + ) + .execute(&database.pool) + .await + .expect("expire image request"); + let repository = PgExecutionStore::new(database.pool.clone()); + + repository + .recover_expired_model_requests(Utc::now()) + .await + .expect("recover expired image request"); + let persisted: (String, Option) = sqlx::query_as( + "select outcome, image_generation_succeeded + from model_requests where id = 'req_expired_image'", + ) + .fetch_one(&database.pool) + .await + .expect("load recovered image request"); + + assert_eq!(persisted, ("incomplete".to_owned(), Some(false))); + database.close().await; +} + +fn successful_core_finalization(id: &str) -> CoreModelRequestFinalization { + CoreModelRequestFinalization { + request_id: ModelRequestId::new(id).expect("request id"), + outcome: ExecutionOutcome::Succeeded, + send_state: UpstreamSendState::Sent, + attempt_count: 1, + downstream_committed_at: None, + client_status_code: Some(200), + client_response_id: None, + upstream_status_code: Some(200), + upstream_request_id: None, + upstream_response_id: None, + upstream_transport: Some("websocket".to_owned()), + http_version: Some("HTTP/2".to_owned()), + websocket_pool: None, + error: None, + provider_error_code: None, + retry_after_ms: None, + usage: Usage::new(), + image_generation_succeeded: None, + cost: CostEstimate::unavailable(), + timings: CoreModelRequestTimings::default(), + completed_at: std::time::SystemTime::now(), + } +} + async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx::Error> { sqlx::query( "insert into model_requests ( diff --git a/backend/crates/gateway-store/tests/postgres/history.rs b/backend/crates/gateway-store/tests/postgres/history.rs index 7fcf895ab..97ac27b2b 100644 --- a/backend/crates/gateway-store/tests/postgres/history.rs +++ b/backend/crates/gateway-store/tests/postgres/history.rs @@ -1,5 +1,4 @@ use chrono::Utc; -use gateway_core::engine::continuation::NativeContinuationReuse; use gateway_store::postgres::{ ModelRequestHistoryRecord, ModelRequestHistoryRepository, PgHistoryRepository, }; @@ -36,11 +35,7 @@ async fn native_pin_is_caller_scoped_and_uses_live_account_and_upstream_handle() let repository = PgHistoryRepository::new(database.pool.clone()); let pin = repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner") .await .expect("resolve native pin") .expect("eligible pin"); @@ -52,11 +47,7 @@ async fn native_pin_is_caller_scoped_and_uses_live_account_and_upstream_handle() assert!(!debug.contains("resp_upstream_native")); let other_client = repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_other", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_other") .await .expect("resolve other caller"); assert!(other_client.is_none()); @@ -66,11 +57,7 @@ async fn native_pin_is_caller_scoped_and_uses_live_account_and_upstream_handle() .await .expect("delete live account while preserving historical ref"); let deleted_account = repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner") .await .expect("resolve deleted account"); assert!(deleted_account.is_none()); @@ -103,11 +90,7 @@ async fn native_pin_requires_committed_success_but_not_a_specific_transport() { .expect("clear downstream commit"); assert!( repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) .await .expect("resolve uncommitted request") .is_none() @@ -123,11 +106,7 @@ async fn native_pin_requires_committed_success_but_not_a_specific_transport() { .expect("mark request failed"); assert!( repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) .await .expect("resolve failed request") .is_none() @@ -139,11 +118,7 @@ async fn native_pin_requires_committed_success_but_not_a_specific_transport() { .expect("restore successful request"); assert!( repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) .await .expect("resolve restored request") .is_some() @@ -157,11 +132,7 @@ async fn native_pin_requires_committed_success_but_not_a_specific_transport() { .expect("mark response connection-independent"); assert!( repository - .resolve_native_continuation_pin( - "resp_gateway_visible", - "key_owner", - NativeContinuationReuse::Reusable, - ) + .resolve_native_continuation_pin("resp_gateway_visible", "key_owner",) .await .expect("resolve HTTP response") .is_some() diff --git a/backend/crates/gateway-store/tests/postgres/observability.rs b/backend/crates/gateway-store/tests/postgres/observability.rs index 1a47cb157..7322d652b 100644 --- a/backend/crates/gateway-store/tests/postgres/observability.rs +++ b/backend/crates/gateway-store/tests/postgres/observability.rs @@ -406,6 +406,17 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact assert_eq!(account_usage[0].cost_coverage.provider_reported_count, 1); assert_eq!(account_usage[0].cost_coverage.unavailable_count, 1); assert_eq!(account_usage[0].costs[0].amount.as_str(), "1.25"); + assert_eq!( + ( + account_usage[0].image_input_tokens, + account_usage[0].image_output_tokens, + account_usage[0].image_request_count, + account_usage[0].image_request_failed_count, + account_usage[0].models[0].image_request_count, + account_usage[0].models[0].image_request_failed_count, + ), + (Some(31), Some(9), 1, 1, 1, 1) + ); let usage_page = repository .list_usage_records(UsageRecordQuery { @@ -421,6 +432,21 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact .await .expect("usage records"); assert_eq!(usage_page.total, 2); + let successful_image = usage_page + .items + .iter() + .find(|record| record.id == "req_observe_success") + .expect("successful image usage record"); + assert_eq!( + ( + successful_image.websocket_pool.as_deref(), + successful_image.image_input_tokens, + successful_image.image_output_tokens, + successful_image.image_generation_requested, + successful_image.image_generation_succeeded, + ), + (Some("reuse"), Some(31), Some(9), true, Some(true)) + ); let detail = repository .usage_record_detail("req_observe_failed") @@ -522,19 +548,22 @@ async fn seed_observability_facts( id, client_api_key_ref, config_revision, protocol, operation, endpoint, client_transport, requested_model_id, input_token_estimate, provider_instance_id, provider_kind, provider_account_id, - provider_account_ref, upstream_model_id, upstream_transport, attempt_count, + provider_account_ref, upstream_model_id, upstream_transport, websocket_pool, + attempt_count, upstream_send_state, downstream_committed_at, outcome, client_status_code, upstream_status_code, client_response_id, upstream_request_id, upstream_response_id, input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens, - total_tokens, cost_source, cost_amount, cost_currency, first_token_ms, latency_ms, + image_input_tokens, image_output_tokens, total_tokens, + image_generation_requested, image_generation_succeeded, + cost_source, cost_amount, cost_currency, first_token_ms, latency_ms, started_at, deadline_at, completed_at ) values ( 'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses', 'http_sse', 'public-model', 100, 'inst_observe', 'openai', 'acct_observe', 'acct_observe', 'upstream-model', - 'http_sse', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200, + 'http_sse', 'reuse', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200, 'resp_observe_success', 'upstream_req_success', 'upstream_resp_success', - 100, 20, 40, 3, 5, 120, + 100, 20, 40, 3, 5, 31, 9, 120, true, true, 'provider_reported', 1.25, 'USD', 120, 900, $1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes' )", @@ -550,14 +579,15 @@ async fn seed_observability_facts( provider_account_ref, upstream_model_id, upstream_transport, attempt_count, upstream_send_state, outcome, client_status_code, upstream_status_code, error_kind, provider_error_code, error_message, retry_after_ms, - input_tokens, cached_tokens, cost_source, latency_ms, + input_tokens, cached_tokens, image_generation_requested, + image_generation_succeeded, cost_source, latency_ms, started_at, deadline_at, completed_at ) values ( 'req_observe_failed', 'key_observe', 1, 'openai', 'responses', '/v1/responses', 'http_sse', 'public-model', 80, 'inst_observe', 'openai', 'acct_observe', 'acct_observe', 'upstream-model', 'http_sse', 2, 'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit', - 'upstream limited', 1000, 0, 0, 'unavailable', 700, + 'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700, $1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes' )", ) diff --git a/backend/crates/gateway-store/tests/redis/mod.rs b/backend/crates/gateway-store/tests/redis/mod.rs index bf9c76cba..a39cc628b 100644 --- a/backend/crates/gateway-store/tests/redis/mod.rs +++ b/backend/crates/gateway-store/tests/redis/mod.rs @@ -4,7 +4,6 @@ mod credential_leases; mod credential_state; mod oauth_pending; mod provider_circuit; -mod runtime_change; mod worker_lease; use chrono::Utc; diff --git a/backend/crates/gateway-store/tests/redis/runtime_change.rs b/backend/crates/gateway-store/tests/redis/runtime_change.rs deleted file mode 100644 index c9af1daf8..000000000 --- a/backend/crates/gateway-store/tests/redis/runtime_change.rs +++ /dev/null @@ -1,12 +0,0 @@ -use gateway_store::{Revision, redis::RuntimeChange}; - -#[test] -fn runtime_change_does_not_expose_raw_account_id() { - let account_id = "real-account-id"; - let change = RuntimeChange::provider_account_changed( - account_id, - Revision::new(1).expect("positive revision"), - ) - .expect("valid change"); - assert!(!format!("{change:?}").contains(account_id)); -} diff --git a/backend/crates/gateway-store/tests/redis/worker_lease.rs b/backend/crates/gateway-store/tests/redis/worker_lease.rs index c8b877f60..bc55bbdb9 100644 --- a/backend/crates/gateway-store/tests/redis/worker_lease.rs +++ b/backend/crates/gateway-store/tests/redis/worker_lease.rs @@ -1,10 +1,15 @@ //! Worker leader lease 通过 StoreBundle 的中立能力进行集成验证。 use std::collections::BTreeSet; +use std::str::FromStr; use std::time::Duration; use gateway_core::task::{WorkerId, WorkerKind, WorkerLeaseAcquisition, WorkerLeaseRequest}; use gateway_store::{StoreConfig, initialize}; +use sqlx::{ + ConnectOptions as _, PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; use uuid::Uuid; #[tokio::test] @@ -15,7 +20,8 @@ async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() { ) else { return; }; - let config = store_config(&database_url, &redis_url); + let database = TestDatabase::create(&database_url).await; + let config = store_config(&database.url, &redis_url); let mut first = initialize(config.clone()) .await .expect("first Store bundle"); @@ -29,7 +35,6 @@ async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() { assert_eq!( kinds, BTreeSet::from([ - WorkerKind::NativeClaimRecovery, WorkerKind::StaleModelRequestRecovery, WorkerKind::Retention, WorkerKind::OpsFlush, @@ -75,6 +80,48 @@ async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() { }; assert!(second_guard.fencing_token() > first_token); second_guard.release().await.expect("second lease release"); + + drop(first); + drop(second); + database.close().await; +} + +struct TestDatabase { + admin: PgPool, + name: String, + url: String, +} + +impl TestDatabase { + async fn create(database_url: &str) -> Self { + let name = format!("cpr_store_worker_{}", Uuid::new_v4().simple()); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(database_url) + .await + .expect("connect test PostgreSQL"); + sqlx::raw_sql(sqlx::AssertSqlSafe(format!("create database \"{name}\""))) + .execute(&admin) + .await + .expect("create worker test database"); + let url = PgConnectOptions::from_str(database_url) + .expect("parse test PostgreSQL URL") + .database(&name) + .to_url_lossy() + .to_string(); + Self { admin, name, url } + } + + async fn close(self) { + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "drop database \"{}\" with (force)", + self.name + ))) + .execute(&self.admin) + .await + .expect("drop worker test database"); + self.admin.close().await; + } } fn store_config(database_url: &str, redis_url: &str) -> StoreConfig { diff --git a/backend/crates/providers/openai/src/credential/admin.rs b/backend/crates/providers/openai/src/credential/admin.rs index 4bcc3c45d..00db6804c 100644 --- a/backend/crates/providers/openai/src/credential/admin.rs +++ b/backend/crates/providers/openai/src/credential/admin.rs @@ -88,10 +88,6 @@ impl std::fmt::Debug for ImportCodexOAuthCredential { } } -pub struct ImportCodexOAuthCredentialBatch { - pub items: Vec, -} - /// Provider-owned 文档归一后的唯一 Core 写入批次。 pub struct PreparedCodexAccountImport { accounts: Vec, @@ -390,25 +386,6 @@ impl CodexCredentialAdmin { }) } - pub fn prepare_import_batch( - &self, - batch: ImportCodexOAuthCredentialBatch, - ) -> Result, CodexCredentialAdminError> { - if batch.items.is_empty() || batch.items.len() > MAX_BATCH { - return Err(CodexCredentialAdminError::InvalidInput); - } - let mut ids = BTreeSet::new(); - let mut prepared = Vec::with_capacity(batch.items.len()); - for item in batch.items { - let account = self.prepare_import(item)?; - if !ids.insert(account.account.id().clone()) { - return Err(CodexCredentialAdminError::InvalidInput); - } - prepared.push(account); - } - Ok(prepared) - } - /// 严格输出可被旧 CPR import 直接读取的 canonical 文档。 pub fn format_cpr_export( &self, diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 57003f7a7..a9a829f54 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -19,7 +19,7 @@ pub(crate) use oauth::oauth_owner_ref; pub use admin::{ CodexCprExportDocument, CodexCredentialAdmin, CodexCredentialAdminError, CodexCredentialAdminService, ExportManagedCodexCredential, ImportCodexOAuthCredential, - ImportCodexOAuthCredentialBatch, PreparedCodexAccountImport, PreparedCodexCredentialRotation, + PreparedCodexAccountImport, PreparedCodexCredentialRotation, PreparedCodexCredentialRotationGuard, RotateManagedCodexCredential, }; pub use catalog::{ diff --git a/backend/crates/providers/openai/src/credential/refresh.rs b/backend/crates/providers/openai/src/credential/refresh.rs index 3c4cadf2d..ee0b9d252 100644 --- a/backend/crates/providers/openai/src/credential/refresh.rs +++ b/backend/crates/providers/openai/src/credential/refresh.rs @@ -140,14 +140,6 @@ impl CodexCredentialRefreshService { Ok(outcomes) } - pub async fn refresh_one( - &self, - due: DueCodexCredential, - ) -> Result { - let policy = self.runtime_policy.load_refresh_policy().await?; - self.refresh_one_with_policy(due, policy).await - } - async fn refresh_one_with_policy( &self, due: DueCodexCredential, diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 9125aaffe..2b5353064 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -19,7 +19,7 @@ use gateway_core::error::{ }; use gateway_core::event::{ GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, - ProviderResponseTimings, UpstreamHttpVersion, + ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind, }; use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProviderSessionState}; use gateway_core::provider_ports::ProviderInstanceCatalogPort; @@ -760,6 +760,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { &response.diagnostics, &response.response_metadata, &response.transport_metrics, + response.websocket_pool_decision, )?); synchronize_passive_quota("a, &lease, &response.rate_limit_headers).await; if let Some(etag) = response.response_metadata.models_etag.as_deref() @@ -910,6 +911,7 @@ fn codex_response_observation( diagnostics: &CodexUpstreamDiagnostics, response_metadata: &CodexResponseMetadata, metrics: &CodexTransportMetrics, + websocket_pool_decision: Option, ) -> Result { let mut observation = ProviderResponseObservation::new( UpstreamTransport::new(actual_transport_name(transport)) @@ -928,6 +930,13 @@ fn codex_response_observation( { 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 + }); + } if let Some(status_code) = diagnostics.status_code { observation = observation.with_status_code(status_code); } @@ -968,6 +977,7 @@ fn codex_error_observation(error: &CodexClientError) -> Option Usage { normalized.cached_tokens = Some(usage.cached_tokens); normalized.cache_write_tokens = Some(usage.cache_write_tokens); normalized.reasoning_tokens = Some(usage.reasoning_tokens); + normalized.image_input_tokens = Some(usage.image_input_tokens); + normalized.image_output_tokens = Some(usage.image_output_tokens); normalized.total_tokens = Some(usage.total_tokens); normalized } diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index c2625c10a..04e9afbab 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -12,8 +12,7 @@ use crate::transport::profile::CodexWireProfileState; use bytes::Bytes; use futures::{Stream, StreamExt}; use gateway_protocol::openai::{ - WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY, - events::{TokenUsage, retry_after_seconds_from_body}, + WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY, events::retry_after_seconds_from_body, sse::SseError, }; use reqwest::{ @@ -303,39 +302,6 @@ impl fmt::Debug for CodexRequestContext<'_> { } } -/// Codex Responses 上游响应。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexBackendResponse { - /// 完整 SSE 文本。 - pub body: String, - /// 实际使用的上游传输。 - pub transport: CodexBackendTransport, - /// 从 SSE 中提取出的最终 usage。 - pub usage: Option, - /// 响应头里的最新 turn state。 - pub turn_state: Option, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游透传的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 首个有效上游 SSE/WebSocket 事件到达代理的耗时。 - pub first_token_ms: Option, - /// 首个 reasoning 输出事件到达代理的耗时。 - pub first_reasoning_ms: Option, - /// 首个正文输出事件到达代理的耗时。 - pub first_text_ms: Option, - /// WebSocket 连接池决策。 - pub websocket_pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, - /// 安全响应元数据。 - pub response_metadata: CodexResponseMetadata, - /// 传输选择与低延迟阶段耗时。 - pub transport_metrics: CodexTransportMetrics, - /// 当前响应是否由池中 WebSocket 保留 connection-local continuation。 - pub connection_local_continuation: bool, -} - /// Codex Responses 实际使用的上游传输。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexBackendTransport { @@ -431,8 +397,6 @@ pub struct CodexBackendClient { pub(super) base_url: String, pub(super) profile: CodexWireProfileState, pub(super) websocket_pool: Option>, - pub(super) websocket_initial_event_timeout: Option, - pub(super) websocket_fast_path_budget: Duration, pub(super) websocket_origin_breaker: WebSocketOriginBreaker, pub(super) websocket_origin_key: String, } diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 9749533e2..f026654f1 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -1,13 +1,9 @@ -use std::{ - sync::Arc, - time::{Duration, Instant}, -}; +use std::{sync::Arc, time::Instant}; use futures::{StreamExt, TryStreamExt}; use gateway_protocol::openai::{ X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER, - events::{extract_sse_usage, retry_after_seconds_from_body}, - sse::SseEventDecoder, + events::retry_after_seconds_from_body, }; use reqwest::{ Client, Response as ReqwestResponse, @@ -34,9 +30,9 @@ use crate::transport::{ websocket::{ CodexWebSocketConnection, CodexWebSocketPool, CodexWebSocketPoolKey, DEFAULT_INITIAL_EVENT_TIMEOUT, WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker, - WebSocketPoolDecision, execute_prepared_response_create_request, - execute_prepared_response_create_request_stream, post_send_ambiguous, - prepare_response_create_request_with_pool, write_websocket_audit_artifact_from_env, + WebSocketPoolDecision, execute_prepared_response_create_request_stream, + post_send_ambiguous, prepare_response_create_request_with_pool, + write_websocket_audit_artifact_from_env, }, }; @@ -56,27 +52,10 @@ impl CodexBackendClient { base_url, profile, websocket_pool: None, - websocket_initial_event_timeout: Some(DEFAULT_INITIAL_EVENT_TIMEOUT), - websocket_fast_path_budget: WEBSOCKET_FAST_PATH_BUDGET, websocket_origin_breaker: WebSocketOriginBreaker::default(), } } - pub fn with_websocket_initial_event_timeout(mut self, timeout: Option) -> Self { - self.websocket_initial_event_timeout = timeout.filter(|timeout| !timeout.is_zero()); - self - } - - pub fn with_websocket_fast_path_budget(mut self, budget: Duration) -> Self { - self.websocket_fast_path_budget = budget.max(Duration::from_millis(1)); - self - } - - pub fn with_websocket_origin_breaker(mut self, breaker: WebSocketOriginBreaker) -> Self { - self.websocket_origin_breaker = breaker; - self - } - /// 为 Responses WebSocket 请求启用连接池。 pub fn with_websocket_pool(mut self, pool: Arc) -> Self { self.websocket_pool = Some(pool); @@ -90,97 +69,6 @@ impl CodexBackendClient { } } - /// 发送 Responses SSE 请求并读取完整响应。 - /// HTTP POST + SSE fallback (when WebSocket pool is disabled). - async fn create_response_http_sse( - &self, - upstream_request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - started_at: Instant, - ) -> CodexClientResult { - let headers = self.request_headers_for_http_response(upstream_request, context)?; - let headers_started_at = Instant::now(); - let response = self - .client - .post(endpoint_url(&self.base_url, CODEX_RESPONSES_PATH)) - .headers(headers) - .json(&upstream_request) - .send() - .await?; - let upstream_headers_ms = elapsed_duration_millis(headers_started_at.elapsed()); - let http_version = http_version_name(response.version()).to_string(); - let status = response.status(); - let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers()); - let turn_state = response_meta::turn_state(response.headers()); - let set_cookie_headers = response_meta::set_cookie_headers(response.headers()); - let rate_limit_headers = response_meta::rate_limit_headers(response.headers()); - let response_metadata = response_meta::response_metadata(response.headers()); - let retry_after_seconds = retry_after_seconds(response.headers(), None); - - if !status.is_success() { - let body = read_capped_error_body(response).await?; - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - diagnostics: Box::new(diagnostics), - set_cookie_headers, - rate_limit_headers, - transport: CodexBackendTransport::HttpSse, - transport_metrics: Box::new(CodexTransportMetrics { - upstream_headers_ms: Some(upstream_headers_ms), - http_version: Some(http_version), - ..CodexTransportMetrics::default() - }), - send_phase: CodexUpstreamSendPhase::AfterPayload, - }); - } - - let mut body_bytes = Vec::new(); - let mut first_token_ms = None; - let mut first_reasoning_ms = None; - let mut first_text_ms = None; - let mut first_event_ms = None; - let mut output_decoder = SseEventDecoder::default(); - let mut stream = http_sse_stream(response); - while let Some(chunk) = stream.try_next().await? { - first_event_ms.get_or_insert_with(|| elapsed_duration_millis(started_at.elapsed())); - body_bytes.extend_from_slice(&chunk); - response_meta::update_response_timing_ms( - started_at, - &mut output_decoder, - &chunk, - &mut first_token_ms, - &mut first_reasoning_ms, - &mut first_text_ms, - ); - } - let body = String::from_utf8_lossy(&body_bytes).into_owned(); - let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?; - Ok(CodexBackendResponse { - body, - transport: CodexBackendTransport::HttpSse, - usage, - turn_state, - set_cookie_headers, - rate_limit_headers, - first_token_ms, - first_reasoning_ms, - first_text_ms, - websocket_pool_decision: None, - diagnostics, - response_metadata, - transport_metrics: CodexTransportMetrics { - upstream_headers_ms: Some(upstream_headers_ms), - first_event_ms, - http_version: Some(http_version), - ..CodexTransportMetrics::default() - }, - connection_local_continuation: false, - }) - } - /// 发送 Responses SSE 请求并返回 live SSE 流（HTTP SSE fallback）。 pub(crate) async fn create_response_stream_http_sse( &self, @@ -246,49 +134,6 @@ impl CodexBackendClient { }) } - pub async fn create_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - self.create_response_started_at(request, context, Instant::now()) - .await - } - - pub async fn create_response_started_at( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - started_at: Instant, - ) -> CodexClientResult { - self.create_response_with_pool_account_started_at(request, context, None, started_at) - .await - } - - pub async fn create_response_with_pool_account_started_at( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - pool_account_id: Option<&str>, - started_at: Instant, - ) -> CodexClientResult { - let prepared = self - .prepare_response_transport_with_pool_account(request, context, pool_account_id) - .await?; - self.create_response_with_prepared(request, context, prepared, started_at) - .await - } - - /// 发送 Responses SSE 请求并返回 live SSE 流。 - pub async fn create_response_stream( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - self.create_response_stream_with_pool_account(request, context, None) - .await - } - pub async fn create_response_stream_with_pool_account( &self, request: &CodexResponsesRequest, @@ -347,7 +192,7 @@ impl CodexBackendClient { let fast_path_budget = match requirement { TransportRequirement::PersistedContinuation | TransportRequirement::ExternalUnknown - | TransportRequirement::NewChain => Some(self.websocket_fast_path_budget), + | TransportRequirement::NewChain => Some(WEBSOCKET_FAST_PATH_BUDGET), TransportRequirement::ExplicitWebSocketWarmup | TransportRequirement::ExactWebSocketContinuation => None, TransportRequirement::HttpRequired => None, @@ -360,7 +205,7 @@ impl CodexBackendClient { &self.websocket_origin_key, fast_path_budget, requirement.requires_websocket(), - self.websocket_initial_event_timeout, + Some(DEFAULT_INITIAL_EVENT_TIMEOUT), ) .await; let prepared = match prepared { @@ -429,64 +274,6 @@ impl CodexBackendClient { }) } - #[doc(hidden)] - pub(crate) async fn create_response_with_prepared( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - prepared: PreparedResponseTransport, - started_at: Instant, - ) -> CodexClientResult { - let PreparedResponseTransport { - requirement, - route, - metrics, - } = prepared; - let result = match route { - PreparedResponseRoute::Http => self - .create_response_http_sse(request, context, started_at) - .await - .map(|mut response| { - merge_preparation_metrics(&mut response.transport_metrics, metrics); - response - }), - PreparedResponseRoute::WebSocket(route) => { - let PreparedWebSocketRoute { request, prepared } = *route; - execute_prepared_response_create_request(&request, prepared, started_at) - .await - .map_err(websocket_exchange_error_to_client_error) - .map(|exchange| CodexBackendResponse { - body: exchange.body, - transport: CodexBackendTransport::WebSocket, - usage: exchange.usage, - turn_state: exchange.turn_state, - set_cookie_headers: exchange.set_cookie_headers, - rate_limit_headers: exchange.rate_limit_headers, - first_token_ms: exchange.first_token_ms, - first_reasoning_ms: exchange.first_reasoning_ms, - first_text_ms: exchange.first_text_ms, - websocket_pool_decision: exchange.pool_decision, - diagnostics: exchange.diagnostics, - response_metadata: exchange.response_metadata, - transport_metrics: CodexTransportMetrics { - first_event_ms: exchange.first_event_ms, - ..metrics - }, - connection_local_continuation: exchange.connection_local_continuation, - }) - } - }; - result.inspect_err(|error| { - tracing::warn!( - request_id = %context.request_id, - transport_requirement = requirement.as_str(), - failure_phase = "post_send_or_explicit_response", - error = %error, - "Responses transport failed after preparation; automatic fallback is disabled" - ); - }) - } - #[doc(hidden)] pub(crate) async fn create_response_stream_with_prepared( &self, diff --git a/backend/crates/providers/openai/src/transport/diagnostics.rs b/backend/crates/providers/openai/src/transport/diagnostics.rs index c0bc4c78e..65d753186 100644 --- a/backend/crates/providers/openai/src/transport/diagnostics.rs +++ b/backend/crates/providers/openai/src/transport/diagnostics.rs @@ -287,26 +287,6 @@ impl CodexUpstreamDiagnostics { } } - pub fn from_pairs(status_code: Option, headers: &[(String, String)]) -> Self { - let request_id = UPSTREAM_REQUEST_ID_HEADERS - .iter() - .find_map(|name| pair_value(headers, name)); - let trace_headers = UPSTREAM_TRACE_HEADERS - .iter() - .filter_map(|name| pair_value(headers, name).map(|value| ((*name).to_string(), value))) - .collect(); - let identity_authorization_error = pair_value(headers, IDENTITY_AUTHORIZATION_ERROR_HEADER); - let identity_error_code = pair_value(headers, IDENTITY_ERROR_JSON_HEADER) - .and_then(|encoded| decode_identity_error_code(&encoded)); - Self { - status_code, - request_id, - identity_authorization_error, - identity_error_code, - trace_headers, - } - } - pub fn is_empty(&self) -> bool { self.status_code.is_none() && self.request_id.is_none() @@ -316,15 +296,6 @@ impl CodexUpstreamDiagnostics { } } -fn pair_value(headers: &[(String, String)], name: &str) -> Option { - headers - .iter() - .find(|(candidate, _)| candidate.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.trim()) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) -} - fn first_header(headers: &HeaderMap, names: &[&str]) -> Option { names.iter().find_map(|name| header_value(headers, name)) } diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index 95dd299a8..a3b54b624 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -24,10 +24,10 @@ pub use self::{ CodexModelCatalogSnapshot, MAX_CODEX_MODEL_CATALOG_BYTES, parse_codex_model_catalog, }, client::{ - CodexBackendClient, CodexBackendResponse, CodexBackendSseStream, - CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult, - CodexRateLimitHeaderUpdates, CodexRequestContext, CodexTransportDecision, - CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client, + CodexBackendClient, CodexBackendSseStream, CodexBackendStreamingResponse, + CodexBackendTransport, CodexClientError, CodexClientResult, CodexRateLimitHeaderUpdates, + CodexRequestContext, CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate, + build_reqwest_client, }, diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase}, endpoints::{ diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index 39e5b87fd..a20b70f20 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -1,10 +1,6 @@ -use std::{collections::BTreeSet, fmt}; +use std::fmt; -use gateway_protocol::openai::{ - events, - schema::reconvert_tuple_values, - sse::{SseError, SseEvent, parse_sse_events}, -}; +use gateway_protocol::openai::events; use serde::Serialize; use serde_json::{Map, Value}; @@ -24,8 +20,6 @@ const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegatio pub struct CodexResponsesRequest { /// 上游请求体（唯一真相源）。 body: Map, - /// tuple schema 原始定义，仅供响应重构时使用。 - pub tuple_schema: Option, /// 是否由客户端显式提供了 prompt cache key。 pub explicit_prompt_cache_key: bool, /// 客户端会话 ID。 @@ -79,7 +73,6 @@ impl fmt::Debug for CodexResponsesRequest { formatter .debug_struct("CodexResponsesRequest") .field("body", &"") - .field("has_tuple_schema", &self.tuple_schema.is_some()) .field("explicit_prompt_cache_key", &self.explicit_prompt_cache_key) .field( "has_local_conversation_id", @@ -424,164 +417,6 @@ impl ResponsesSseFailure { } } -/// 从完成 SSE 中提取会话亲和性和 replay 元数据。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CompletedResponseMetadata { - /// 上游 response id。 - pub response_id: String, - /// 完成响应中的 function call ids。 - pub function_call_ids: Vec, - /// 完成响应的完整 output items。 - pub output: Vec, -} - -/// 从 Codex Responses SSE 中提取完成响应元数据。 -pub fn completed_response_metadata( - body: &str, -) -> Result, SseError> { - let events = parse_sse_events(body)?; - let mut response_id = None; - let mut function_call_ids = BTreeSet::new(); - let mut output = Vec::new(); - let mut completed_items = Vec::new(); - - for event in events { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - match event.event.as_deref() { - Some("response.output_item.done") => { - if let Some(item) = value.get("item") { - completed_items.push(item.clone()); - } - if let Some(call_id) = value.pointer("/item/call_id").and_then(Value::as_str) { - function_call_ids.insert(call_id.to_string()); - } - } - Some("response.completed") => { - response_id = value - .pointer("/response/id") - .and_then(Value::as_str) - .map(ToString::to_string); - collect_response_function_call_ids(&value, &mut function_call_ids); - if let Some(completed_output) = - value.pointer("/response/output").and_then(Value::as_array) - { - completed_output.clone_into(&mut output); - } - if output.is_empty() { - output.clone_from(&completed_items); - } - } - _ => {} - } - } - - Ok(response_id.map(|response_id| CompletedResponseMetadata { - response_id, - function_call_ids: function_call_ids.into_iter().collect(), - output, - })) -} - -/// 对单个 Responses SSE 事件的数据执行 tuple schema 回转换。 -pub fn reconvert_responses_sse_event_tuple_values( - event_name: Option<&str>, - mut data: Value, - tuple_schema: &Value, -) -> Value { - match responses_event_type(event_name, Some(&data)) { - Some("response.output_text.delta") => { - reconvert_output_text_delta_tuple_values(&mut data, tuple_schema); - } - Some("response.output_item.done") => { - if let Some(item) = data.get_mut("item") { - reconvert_output_item_tuple_values(item, tuple_schema); - } - } - Some("response.completed" | "response.incomplete") => { - if let Some(response) = data.get_mut("response") { - reconvert_completed_response_tuple_values(response, Some(tuple_schema)); - } - } - _ => {} - } - data -} - -/// 判断 Responses SSE 事件是否为终止事件。 -pub fn response_sse_event_is_terminal(event: &SseEvent) -> bool { - let value = serde_json::from_str::(&event.data).ok(); - matches!( - responses_event_type(event.event.as_deref(), value.as_ref()), - Some("response.completed" | "response.incomplete" | "response.failed" | "error") - ) -} - -fn responses_event_type<'a>( - event_name: Option<&'a str>, - data: Option<&'a Value>, -) -> Option<&'a str> { - event_name.or_else(|| { - data.and_then(|data| data.get("type")) - .and_then(Value::as_str) - }) -} - -fn reconvert_output_text_delta_tuple_values(data: &mut Value, tuple_schema: &Value) { - let Some(delta) = data.get("delta").and_then(Value::as_str) else { - return; - }; - let Ok(parsed) = serde_json::from_str::(delta) else { - return; - }; - let reconverted = reconvert_tuple_values(parsed, tuple_schema); - data["delta"] = Value::String(reconverted.to_string()); -} - -fn reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) { - let Some(tuple_schema) = tuple_schema else { - return; - }; - let Some(items) = response.get_mut("output").and_then(Value::as_array_mut) else { - return; - }; - for item in items { - reconvert_output_item_tuple_values(item, tuple_schema); - } -} - -fn reconvert_output_item_tuple_values(item: &mut Value, tuple_schema: &Value) { - let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else { - return; - }; - for part in content { - let part_type = part.get("type").and_then(Value::as_str); - if !matches!(part_type, Some("output_text" | "text")) { - continue; - } - let Some(text) = part.get("text").and_then(Value::as_str) else { - continue; - }; - let Ok(parsed) = serde_json::from_str::(text) else { - continue; - }; - let reconverted = reconvert_tuple_values(parsed, tuple_schema); - part["text"] = Value::String(reconverted.to_string()); - } -} - -fn collect_response_function_call_ids(value: &Value, function_call_ids: &mut BTreeSet) { - let Some(output) = value.pointer("/response/output").and_then(Value::as_array) else { - return; - }; - for item in output { - if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { - function_call_ids.insert(call_id.to_string()); - } - } -} - fn failure_message(value: &Value) -> Option { value .pointer("/response/error/message") @@ -632,7 +467,6 @@ impl CodexResponsesRequest { pub fn from_body(body: Map) -> Self { Self { body, - tuple_schema: None, explicit_prompt_cache_key: false, client_conversation_id: None, client_session_id: None, @@ -679,12 +513,6 @@ impl CodexResponsesRequest { .unwrap_or_default() } - /// 设置模型名（模型后缀路由归一）。 - pub fn set_model(&mut self, model: impl Into) { - self.body - .insert("model".to_string(), Value::String(model.into())); - } - /// 指令文本（缺省空串）。 pub fn instructions(&self) -> &str { self.body @@ -714,11 +542,6 @@ impl CodexResponsesRequest { .unwrap_or(true) } - /// 设置流式标志。 - pub fn set_stream(&mut self, stream: bool) { - self.body.insert("stream".to_string(), Value::Bool(stream)); - } - /// 是否要求上游存储响应。 pub fn store(&self) -> bool { self.body @@ -847,15 +670,6 @@ impl CodexResponsesRequest { } } } - - /// 判断请求是否声明了图片生成工具。 - pub fn expects_image_generation(&self) -> bool { - self.tools().is_some_and(|tools| { - tools - .iter() - .any(|tool| tool.get("type").and_then(Value::as_str) == Some("image_generation")) - }) - } } fn non_empty_owned_string(value: &str) -> Option { diff --git a/backend/crates/providers/openai/src/transport/response_meta.rs b/backend/crates/providers/openai/src/transport/response_meta.rs index 93b229d2c..9bacce0a7 100644 --- a/backend/crates/providers/openai/src/transport/response_meta.rs +++ b/backend/crates/providers/openai/src/transport/response_meta.rs @@ -1,12 +1,8 @@ //! 上游响应元数据提取辅助。 -use std::time::Instant; - -use gateway_protocol::openai::{events::is_rate_limit_header_name, sse::SseEventDecoder}; +use gateway_protocol::openai::events::is_rate_limit_header_name; use reqwest::header::{HeaderMap, SET_COOKIE}; -use crate::transport::protocol::responses::response_event_signals; - use super::diagnostics::CodexUpstreamDiagnostics; /// Codex Responses 上游响应元数据。 @@ -117,41 +113,3 @@ fn upsert_client_header(headers: &mut Vec<(String, String)>, header: (String, St headers.push(header); } } - -pub(super) fn update_response_timing_ms( - started_at: Instant, - decoder: &mut SseEventDecoder, - chunk: &[u8], - first_token_ms: &mut Option, - first_reasoning_ms: &mut Option, - first_text_ms: &mut Option, -) { - let Ok(events) = decoder.push(chunk) else { - return; - }; - for event in events { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - let event_type = event.event.as_deref().or_else(|| { - value - .get("type") - .and_then(serde_json::Value::as_str) - .filter(|value| !value.trim().is_empty()) - }); - let signals = response_event_signals(event_type, &value); - if !signals.semantic_output { - continue; - } - let elapsed_ms = crate::transport::time::elapsed_millis_i64(started_at).max(1); - if first_token_ms.is_none() { - *first_token_ms = Some(elapsed_ms); - } - if signals.reasoning_output && first_reasoning_ms.is_none() { - *first_reasoning_ms = Some(elapsed_ms); - } - if signals.text_output && first_text_ms.is_none() { - *first_text_ms = Some(elapsed_ms); - } - } -} diff --git a/backend/crates/providers/openai/src/transport/time.rs b/backend/crates/providers/openai/src/transport/time.rs index bfe1387f1..03214875c 100644 --- a/backend/crates/providers/openai/src/transport/time.rs +++ b/backend/crates/providers/openai/src/transport/time.rs @@ -1,15 +1,9 @@ //! Codex transport 私有时间辅助。 -use std::time::Instant; - use chrono::{DateTime, FixedOffset, Utc}; const CHINA_OFFSET_SECONDS: i32 = 8 * 60 * 60; -pub(crate) fn elapsed_millis_i64(started_at: Instant) -> i64 { - started_at.elapsed().as_millis().min(i64::MAX as u128) as i64 -} - pub(crate) fn china_filename_timestamp_millis(value: &DateTime) -> String { value .with_timezone(&china_offset()) diff --git a/backend/crates/providers/openai/src/transport/websocket/breaker.rs b/backend/crates/providers/openai/src/transport/websocket/breaker.rs index aec6bfb0e..5ca180c6d 100644 --- a/backend/crates/providers/openai/src/transport/websocket/breaker.rs +++ b/backend/crates/providers/openai/src/transport/websocket/breaker.rs @@ -6,9 +6,10 @@ use std::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, }, - time::{Duration, Instant}, + time::Duration, }; +use tokio::time::Instant; use uuid::Uuid; const DEFAULT_FAILURE_THRESHOLD: usize = 3; diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index ef50fd9ab..435d8b0b5 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -11,9 +11,7 @@ use super::{ }, error::CodexWebSocketExchangeError, exchange::{ - CodexWebSocketExchange, CodexWebSocketStreamingExchange, CollectedWebSocket, - WebSocketStreamPoolReturn, WebSocketTerminalKind, collect_websocket_response, - reusable_websocket_metadata, stream_websocket_response, + CodexWebSocketStreamingExchange, WebSocketStreamPoolReturn, stream_websocket_response, }, handshake::{connect_pumped_websocket, send_websocket_request, websocket_connection_metadata}, model::{ @@ -491,76 +489,6 @@ fn continuation_unavailable( CodexWebSocketExchangeError::ContinuationUnavailable { reason } } -/// 发送 payload 后只等待该次 exchange；此边界之后不允许 transport fallback。 -pub(crate) async fn execute_prepared_response_create_request( - request: &CodexWebSocketRequest, - prepared: PreparedWebSocket, - started_at: Instant, -) -> Result { - let PreparedWebSocket { - connection, - binding, - initial_event_timeout, - .. - } = prepared; - let (lease, reused, pool_decision) = binding.into_parts(); - let PooledWebSocketConnection { - websocket, - metadata, - continuation, - created_at, - } = connection; - if let Err(error) = send_websocket_request(&websocket, request.payload_text()).await { - discard_after_send(websocket, lease).await; - return Err(post_send_ambiguous(error)); - } - let connection_local_available = lease.is_some(); - let collected = collect_websocket_response( - websocket, - metadata, - continuation, - reused, - started_at, - initial_event_timeout, - ) - .await; - let CollectedWebSocket { - mut exchange, - websocket, - metadata, - continuation, - terminal, - } = match collected { - Ok(collected) => collected, - Err(error) => { - if let Some(lease) = lease { - lease.discard().await; - } - return Err(post_send_ambiguous(error)); - } - }; - exchange.pool_decision = pool_decision; - exchange.connection_local_continuation = connection_local_available; - match (terminal, lease) { - (WebSocketTerminalKind::Completed, Some(lease)) => { - lease - .put(PooledWebSocketConnection { - websocket, - metadata: reusable_websocket_metadata(metadata), - continuation, - created_at, - }) - .await; - } - (_, Some(lease)) => { - lease.discard().await; - websocket.close().await; - } - (_, None) => websocket.close().await, - } - Ok(exchange) -} - pub(crate) async fn execute_prepared_response_create_request_stream( request: &CodexWebSocketRequest, prepared: PreparedWebSocket, diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs deleted file mode 100644 index eb85031da..000000000 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/collect.rs +++ /dev/null @@ -1,133 +0,0 @@ -//! 聚合式 WebSocket exchange。 - -use std::time::{Duration, Instant}; - -use gateway_protocol::openai::{events, sse::SseEventDecoder}; -use tungstenite::Message; - -use crate::transport::response_meta; - -use super::super::{ - pool::{CodexWebSocketConnectionMetadata, WebSocketContinuationState}, - pump::PumpedWebSocket, -}; -use super::io::{ - next_websocket_message, receive_idle_timeout, reused_connection_died_before_first_event, -}; -use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event}; -use super::{CodexWebSocketExchange, CodexWebSocketExchangeError}; - -pub(in crate::transport::websocket) struct CollectedWebSocket { - pub(in crate::transport::websocket) exchange: CodexWebSocketExchange, - pub(in crate::transport::websocket) websocket: PumpedWebSocket, - pub(in crate::transport::websocket) metadata: CodexWebSocketConnectionMetadata, - pub(in crate::transport::websocket) continuation: WebSocketContinuationState, - pub(in crate::transport::websocket) terminal: WebSocketTerminalKind, -} - -pub(in crate::transport::websocket) async fn collect_websocket_response( - mut websocket: PumpedWebSocket, - mut metadata: CodexWebSocketConnectionMetadata, - mut continuation: WebSocketContinuationState, - reused_connection: bool, - started_at: Instant, - initial_event_timeout: Option, -) -> Result { - let mut body = String::new(); - let mut saw_upstream_activity = false; - let mut first_token_ms = None; - let mut first_reasoning_ms = None; - let mut first_text_ms = None; - let mut first_event_ms = None; - let mut output_decoder = SseEventDecoder::default(); - - loop { - let receive_timeout = receive_idle_timeout(saw_upstream_activity, initial_event_timeout); - let message = match next_websocket_message(&mut websocket, receive_timeout).await { - Ok(message) => message, - Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { timeout }) - if !saw_upstream_activity => - { - if reused_connection { - return Err(reused_connection_died_before_first_event( - &CodexWebSocketExchangeError::InitialEventTimeout { timeout }, - )); - } - return Err(CodexWebSocketExchangeError::InitialEventTimeout { timeout }); - } - Err(error) => return Err(error), - }; - let Some(message) = message else { - break; - }; - let text = match message { - Message::Text(text) => { - saw_upstream_activity = true; - first_event_ms.get_or_insert_with(|| { - crate::transport::time::elapsed_millis_i64(started_at).max(1) - }); - text - } - Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), - Message::Close(_) if reused_connection && !saw_upstream_activity => { - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { - message: "websocket closed".to_string(), - }, - ); - } - Message::Close(_) => break, - _ => continue, - }; - let raw = text.to_string(); - let ExchangeAction::Forward { frame, terminal } = - reduce_websocket_event(&raw, &mut metadata, &mut continuation)? - else { - continue; - }; - body.push_str(&frame); - response_meta::update_response_timing_ms( - started_at, - &mut output_decoder, - frame.as_bytes(), - &mut first_token_ms, - &mut first_reasoning_ms, - &mut first_text_ms, - ); - if let Some(terminal) = terminal { - let usage = events::extract_sse_usage(&body)?; - let exchange = CodexWebSocketExchange { - body, - usage, - turn_state: metadata.turn_state.clone(), - set_cookie_headers: metadata.set_cookie_headers.clone(), - rate_limit_headers: metadata.rate_limit_headers.clone(), - first_token_ms, - first_reasoning_ms, - first_text_ms, - first_event_ms, - pool_decision: None, - connection_local_continuation: false, - diagnostics: metadata.diagnostics.clone(), - response_metadata: metadata.response_metadata.clone(), - }; - return Ok(CollectedWebSocket { - exchange, - websocket, - metadata, - continuation, - terminal, - }); - } - } - - if reused_connection && !saw_upstream_activity { - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { - message: "websocket closed before terminal event".to_string(), - }, - ); - } - - Err(CodexWebSocketExchangeError::ClosedBeforeTerminal) -} diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs index d938d3499..9e7440349 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs @@ -1,19 +1,13 @@ //! WebSocket 帧交换、流转发与结果类型。 -mod collect; mod io; mod reducer; mod stream; -use std::{ - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; +use std::{pin::Pin, sync::Arc, time::Duration}; use bytes::Bytes; use futures::Stream; -use gateway_protocol::openai::events::TokenUsage; use tokio::sync::Mutex; use uuid::Uuid; @@ -21,54 +15,14 @@ use crate::transport::{ diagnostics::CodexUpstreamDiagnostics, response_meta::CodexResponseMetadata, }; -use super::breaker::WebSocketOriginBreaker; use super::error::CodexWebSocketExchangeError; use super::pool::{CodexWebSocketConnectionMetadata, WebSocketPoolDecision}; -use super::{ - CodexWebSocketRequest, execute_prepared_response_create_request, - prepare_response_create_request_with_pool, -}; const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); const WEBSOCKET_STREAM_BUFFER: usize = 16; -pub(super) use self::{ - collect::{CollectedWebSocket, collect_websocket_response}, - reducer::WebSocketTerminalKind, - stream::{WebSocketStreamPoolReturn, stream_websocket_response}, -}; - -/// Responses WebSocket exchange result. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketExchange { - /// 完整 SSE 文本。 - pub body: String, - /// 从 SSE 中提取的 usage。 - pub usage: Option, - /// 上游 metadata 帧中的最新 turn state。 - pub turn_state: Option, - /// 上游握手响应里的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游握手响应里的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 首个有效上游 WebSocket 事件到达代理的耗时。 - pub first_token_ms: Option, - /// 首个 reasoning 输出事件到达代理的耗时。 - pub first_reasoning_ms: Option, - /// 首个正文输出事件到达代理的耗时。 - pub first_text_ms: Option, - /// 首个上游协议事件到达代理的耗时。 - pub first_event_ms: Option, - /// WebSocket 连接池决策。 - pub pool_decision: Option, - /// terminal completed 后该 socket 是否会保留 connection-local continuation。 - pub connection_local_continuation: bool, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, - /// 安全响应元数据。 - pub response_metadata: CodexResponseMetadata, -} +pub(super) use self::stream::{WebSocketStreamPoolReturn, stream_websocket_response}; /// Responses WebSocket live SSE exchange result. pub struct CodexWebSocketStreamingExchange { @@ -104,24 +58,6 @@ pub type CodexWebSocketRateLimitHeaderUpdates = Arc> /// WebSocket live stream turn-state update. pub type CodexWebSocketTurnStateUpdate = Arc>>; -/// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 -pub async fn execute_response_create_request( - request: &CodexWebSocketRequest, -) -> Result { - let breaker = WebSocketOriginBreaker::default(); - let prepared = prepare_response_create_request_with_pool( - request, - None, - &breaker, - request.connection().endpoint(), - None, - false, - None, - ) - .await?; - execute_prepared_response_create_request(request, prepared, Instant::now()).await -} - pub(super) fn reusable_websocket_metadata( mut metadata: CodexWebSocketConnectionMetadata, ) -> CodexWebSocketConnectionMetadata { diff --git a/backend/crates/providers/openai/src/transport/websocket/mod.rs b/backend/crates/providers/openai/src/transport/websocket/mod.rs index ad31ea0d2..6d1d892c3 100644 --- a/backend/crates/providers/openai/src/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/mod.rs @@ -21,9 +21,8 @@ pub use self::{ }, error::{CodexWebSocketExchangeError, CodexWebSocketUpstreamError}, exchange::{ - CodexWebSocketExchange, CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream, + CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate, - execute_response_create_request, }, handshake::responses_websocket_endpoint, model::{ @@ -37,7 +36,7 @@ pub use self::{ }; pub(crate) use self::{ coordinator::{ - PreparedWebSocket, WEBSOCKET_FAST_PATH_BUDGET, execute_prepared_response_create_request, + PreparedWebSocket, WEBSOCKET_FAST_PATH_BUDGET, execute_prepared_response_create_request_stream, post_send_ambiguous, prepare_response_create_request_with_pool, }, diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs index 1d0a3f155..88217440e 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs @@ -10,7 +10,6 @@ use std::{ time::Duration, }; -use serde_json::{Value, json}; use tokio::{sync::Semaphore, time::Instant}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; use uuid::Uuid; @@ -393,8 +392,8 @@ impl WebSocketPoolDecision { self.kind.as_str() } - pub fn metadata_value(self) -> Value { - json!({ "kind": self.kind() }) + pub const fn is_reuse(self) -> bool { + matches!(self.kind, WebSocketPoolDecisionKind::Reuse) } } diff --git a/backend/crates/providers/openai/tests/credential/admin.rs b/backend/crates/providers/openai/tests/credential/admin.rs index a2528db95..c5b81f299 100644 --- a/backend/crates/providers/openai/tests/credential/admin.rs +++ b/backend/crates/providers/openai/tests/credential/admin.rs @@ -18,8 +18,7 @@ use provider_openai::credential::{ CodexCredentialAdminError, CodexCredentialAdminService, CodexCredentialCodec, CodexIdentityExpectation, CodexIdentityVerification, CodexIdentityVerificationError, CodexJwtIdentityVerifier, ExportManagedCodexCredential, ImportCodexOAuthCredential, - ImportCodexOAuthCredentialBatch, ReqwestCodexAuthenticatedAccountSource, - ReqwestOpenAiJwksSource, RotateManagedCodexCredential, + ReqwestCodexAuthenticatedAccountSource, ReqwestOpenAiJwksSource, RotateManagedCodexCredential, }; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; use secrecy::{ExposeSecret, SecretString}; @@ -62,31 +61,6 @@ fn prepare_import_returns_core_account_with_plaintext_json() { ); } -#[test] -fn prepare_batch_rejects_duplicate_ids_before_store_mutation() { - let error = CodexCredentialAdmin - .prepare_import_batch(ImportCodexOAuthCredentialBatch { - items: vec![ - import("acct_duplicate", "first"), - import("acct_duplicate", "second"), - ], - }) - .expect_err("duplicate IDs must fail"); - assert_eq!(error, CodexCredentialAdminError::InvalidInput); -} - -#[test] -fn prepare_batch_is_all_or_nothing_validation() { - let mut invalid = import("acct_invalid", "second"); - invalid.provider_instance_id.clear(); - let error = CodexCredentialAdmin - .prepare_import_batch(ImportCodexOAuthCredentialBatch { - items: vec![import("acct_valid", "first"), invalid], - }) - .expect_err("invalid item must reject the whole prepared batch"); - assert_eq!(error, CodexCredentialAdminError::InvalidInput); -} - #[test] fn prepare_rotation_preserves_provider_owned_cookie_data() { let account = codex_account("acct_rotate"); diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index a154fb02d..7862d56e9 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -6,7 +6,7 @@ use std::time::{Duration, SystemTime}; use chrono::Utc; use futures::{SinkExt, StreamExt}; use gateway_core::engine::continuation::{ - ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, PreviousResponseId, }; use gateway_core::engine::credential::{ AccountAvailability, ProviderAccountId, ProviderAccountStore, @@ -1714,7 +1714,6 @@ async fn native_continuation_sends_upstream_handle_never_gateway_client_id() { ProviderKind::new("openai").expect("provider"), instance_id(), ProviderAccountId::new("acct_provider").expect("account id"), - NativeContinuationReuse::Reusable, ); let second_operation = identity_tainted_operation(); let mut second = provider @@ -1810,7 +1809,6 @@ async fn connection_replay_restores_full_input_without_reusing_native_handle() { ProviderKind::new("openai").expect("provider"), instance_id(), ProviderAccountId::new("acct_provider").expect("account id"), - NativeContinuationReuse::Reusable, ); let second_operation = identity_tainted_operation() .with_provider_session_state(provider_state.expect("connection replay state")); diff --git a/backend/crates/providers/openai/tests/transport/canonical.rs b/backend/crates/providers/openai/tests/transport/canonical.rs index c7a9a4bc2..edfbdc1a4 100644 --- a/backend/crates/providers/openai/tests/transport/canonical.rs +++ b/backend/crates/providers/openai/tests/transport/canonical.rs @@ -81,6 +81,29 @@ fn decoder_should_normalize_text_usage_and_completion() { )); } +#[test] +fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_image_usage\",\"model\":\"gpt-test\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_usage\",\"model\":\"gpt-test\",\"status\":\"completed\",\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"total_tokens\":17},\"tool_usage\":{\"image_gen\":{\"input_tokens\":31,\"output_tokens\":9}}}}\n\n", + ); + let events = CodexCanonicalDecoder::new("fallback") + .push(body.as_bytes()) + .expect("canonical image usage response"); + let image_usage = canonical_facts(&events) + .into_iter() + .find_map(|event| match event { + GatewayEvent::Usage(usage) => { + Some((usage.image_input_tokens, usage.image_output_tokens)) + } + _ => None, + }); + + assert_eq!(image_usage, Some((Some(31), Some(9)))); +} + #[test] fn decoder_should_accept_official_codex_metadata_lifecycle_events() { let body = concat!( diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index 36f78d076..6e4432408 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -27,7 +27,45 @@ fn explicit_websocket_warmup_request(conversation_id: &str) -> CodexResponsesReq request } -#[tokio::test] +async fn reject_websocket_openings(listener: &TcpListener, attempts: usize) { + for _ in 0..attempts { + let (mut opening, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut opening).await; + assert!(request.starts_with("GET /codex/responses HTTP/1.1")); + opening + .write_all( + b"HTTP/1.1 503 Service Unavailable\r\ncontent-length: 0\r\nconnection: close\r\n\r\n", + ) + .await + .unwrap(); + } +} + +async fn record_websocket_origin_failures( + backend: &CodexBackendClient, + attempts: usize, + scenario: &str, +) { + for attempt in 0..attempts { + let conversation_id = format!("conversation-{scenario}-{attempt}"); + let request_id = format!("req_{scenario}_{attempt}"); + let error = backend + .create_response( + &explicit_websocket_warmup_request(&conversation_id), + request_context(&request_id, Some("chatgpt-account")), + ) + .await + .expect_err("seed WebSocket opening should fail"); + match error { + CodexClientError::Upstream { status, .. } + if status == reqwest::StatusCode::SERVICE_UNAVAILABLE => {} + error => panic!("{scenario} opening {attempt} failed unexpectedly: {error}"), + } + tokio::task::yield_now().await; + } +} + +#[tokio::test(start_paused = true)] async fn cold_websocket_should_fall_back_without_recording_a_successful_connect() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -47,8 +85,7 @@ async fn cold_websocket_should_fall_back_without_recording_a_successful_connect( format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(30)); + .with_websocket_pool(Arc::clone(&pool)); let response = backend .create_response( @@ -70,7 +107,7 @@ async fn cold_websocket_should_fall_back_without_recording_a_successful_connect( pool.shutdown().await; } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_request() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -78,7 +115,7 @@ async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_requ let server = tokio::spawn(async move { let (websocket_stream, _) = listener.accept().await.unwrap(); let websocket_server = tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(60)).await; + tokio::time::sleep(Duration::from_secs(1)).await; let mut websocket = accept_codex_test_websocket(websocket_stream).await; websocket_ready_tx.send(()).unwrap(); let _payload = websocket.next().await.unwrap().unwrap(); @@ -102,8 +139,7 @@ async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_requ format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(30)); + .with_websocket_pool(Arc::clone(&pool)); let request = new_chain_request("conversation-background-ready"); let first = backend @@ -157,8 +193,7 @@ async fn shared_websocket_opening_should_keep_the_original_fast_path_deadline() format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(200)); + .with_websocket_pool(Arc::clone(&pool)); let request = new_chain_request("conversation-original-deadline"); let first = backend @@ -168,20 +203,23 @@ async fn shared_websocket_opening_should_keep_the_original_fast_path_deadline() ) .await .expect("first request should use HTTP after the opening budget"); - let second = timeout( - Duration::from_millis(100), - backend.create_response( + let second_started_at = tokio::time::Instant::now(); + let second = backend + .create_response( &request, request_context("req_original_deadline_second", Some("chatgpt-account")), - ), - ) - .await - .expect("shared opening must not grant the second request another full budget") - .expect("second request should use HTTP"); + ) + .await + .expect("second request should use HTTP"); + let second_elapsed = second_started_at.elapsed(); server.await.unwrap(); assert_eq!(first.transport, CodexBackendTransport::HttpSse); assert_eq!(second.transport, CodexBackendTransport::HttpSse); + assert!( + second_elapsed < Duration::from_millis(400), + "shared opening waited for a new budget: {second_elapsed:?}" + ); assert_eq!( second.transport_metrics.decision, Some(CodexTransportDecision::Http2WebSocketBudgetExhausted) @@ -204,10 +242,7 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() { write_completed_sse_response(&mut http).await; let mut byte = [0_u8; 1]; - let read = timeout(Duration::from_secs(1), stale_opening.read(&mut byte)) - .await - .expect("account eviction should close the opening socket") - .unwrap(); + let read = stale_opening.read(&mut byte).await.unwrap(); assert_eq!(read, 0); let (fresh_stream, _) = listener.accept().await.unwrap(); @@ -226,8 +261,7 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() { format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(30)); + .with_websocket_pool(Arc::clone(&pool)); let request = new_chain_request("conversation-evict-opening"); let first = backend @@ -238,33 +272,23 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() { .await .expect("first request should use HTTP after the opening budget"); pool.evict_account("chatgpt-account").await; - let second = timeout( - Duration::from_secs(2), - backend - .clone() - .with_websocket_fast_path_budget(Duration::from_millis(500)) - .create_response( - &request, - request_context("req_evict_opening_second", Some("chatgpt-account")), - ), - ) - .await - .expect("fresh websocket request should finish") - .expect("the next request should build a fresh websocket"); - timeout(Duration::from_secs(2), server) + let second = backend + .clone() + .create_response( + &request, + request_context("req_evict_opening_second", Some("chatgpt-account")), + ) .await - .expect("test server should finish") - .unwrap(); + .expect("the next request should build a fresh websocket"); + server.await.unwrap(); assert_eq!(first.transport, CodexBackendTransport::HttpSse); assert_eq!(second.transport, CodexBackendTransport::WebSocket); assert!(second.body.contains("resp_after_eviction")); - timeout(Duration::from_secs(2), pool.shutdown()) - .await - .expect("pool shutdown should finish"); + pool.shutdown().await; } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -291,8 +315,7 @@ async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() { format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(30)); + .with_websocket_pool(Arc::clone(&pool)); backend .create_response( @@ -851,10 +874,7 @@ async fn concurrent_same_key_should_singleflight_websocket_opening() { format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::new(CodexWebSocketPool::default())) - // 该用例验证 same-key singleflight，不验证 200ms deadline。并行全套测试下 - // 调度停顿可能让首请求也越过 200ms，制造两个 HTTP fallback 的假失败。 - .with_websocket_fast_path_budget(Duration::from_secs(2)); + .with_websocket_pool(Arc::new(CodexWebSocketPool::default())); let request = new_chain_request("conversation-singleflight"); let first_backend = backend.clone(); let first_request = request.clone(); @@ -970,19 +990,12 @@ async fn origin_breaker_should_open_then_allow_only_one_half_open_probe() { async fn cancelled_half_open_opening_should_allow_another_probe() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let origin_key = format!("http://{addr}"); - let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { - failure_threshold: 1, - failure_window: Duration::from_secs(1), - open_duration: Duration::ZERO, - }); - let WebSocketOriginBreakerDecision::Allowed(permit) = breaker.try_acquire(&origin_key) else { - panic!("closed breaker should grant the initial permit"); - }; - permit.fast_timeout(); + let breaker_config = WebSocketOriginBreakerConfig::default(); let (opening_started_tx, opening_started_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { + reject_websocket_openings(&listener, breaker_config.failure_threshold).await; + let (mut opening, _) = listener.accept().await.unwrap(); let request = read_http_request(&mut opening).await; assert!(request.starts_with("GET /codex/responses HTTP/1.1")); @@ -994,17 +1007,37 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { .expect("account eviction should close the half-open probe") .unwrap(); assert_eq!(read, 0); + + let (fresh_stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(fresh_stream).await; + let _payload = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response("resp_after_cancelled_probe", 2, 1).into(), + )) + .await + .unwrap(); }); let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), - origin_key.clone(), + format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_origin_breaker(breaker.clone()); + .with_websocket_pool(Arc::clone(&pool)); + record_websocket_origin_failures( + &backend, + breaker_config.failure_threshold, + "cancelled-half-open-seed", + ) + .await; + tokio::time::pause(); + tokio::time::advance(breaker_config.open_duration).await; + tokio::time::resume(); + + let probe_backend = backend.clone(); let attempt = tokio::spawn(async move { - backend + probe_backend .create_response( &explicit_websocket_warmup_request("conversation-cancelled-half-open"), request_context("req_cancelled_half_open", Some("chatgpt-account")), @@ -1019,17 +1052,18 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { .expect("cancelled request should finish") .unwrap() .expect_err("cancelled half-open opening should fail the request"); + + let response = backend + .create_response( + &explicit_websocket_warmup_request("conversation-next-half-open-probe"), + request_context("req_next_half_open_probe", Some("chatgpt-account")), + ) + .await + .expect("cancelled half-open opening should allow another probe"); server.await.unwrap(); - let WebSocketOriginBreakerDecision::Allowed(next_probe) = breaker.try_acquire(&origin_key) - else { - panic!("cancelled half-open probe should release its ownership"); - }; - assert!(matches!( - breaker.try_acquire(&origin_key), - WebSocketOriginBreakerDecision::HalfOpenBusy - )); - next_probe.succeed(); + assert_eq!(response.transport, CodexBackendTransport::WebSocket); + assert!(response.body.contains("resp_after_cancelled_probe")); pool.shutdown().await; } @@ -1057,14 +1091,19 @@ fn origin_breaker_should_count_hard_opening_failures() { async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); + let breaker_config = WebSocketOriginBreakerConfig::default(); + assert!(breaker_config.failure_threshold >= 3); + let seed_failures = breaker_config.failure_threshold - 2; let (late_failure_tx, late_failure_rx) = tokio::sync::oneshot::channel(); let (second_used_websocket_tx, second_used_websocket_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { + reject_websocket_openings(&listener, seed_failures).await; + let (mut delayed_opening, _) = listener.accept().await.unwrap(); let opening = read_http_request(&mut delayed_opening).await; assert!(opening.starts_with("GET /codex/responses HTTP/1.1")); let delayed_failure = tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(80)).await; + tokio::time::sleep(Duration::from_secs(1)).await; delayed_opening .write_all( b"HTTP/1.1 503 Service Unavailable\r\ncontent-length: 0\r\nconnection: close\r\n\r\n", @@ -1114,20 +1153,14 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() { write_completed_sse_response(&mut second_http).await; } }); - let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { - failure_threshold: 2, - failure_window: Duration::from_secs(1), - open_duration: Duration::from_secs(1), - }); let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_millis(30)) - .with_websocket_origin_breaker(breaker); + .with_websocket_pool(Arc::clone(&pool)); + record_websocket_origin_failures(&backend, seed_failures, "late-failure-seed").await; let first = backend .create_response( @@ -1136,58 +1169,34 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() { ) .await .expect("first request should use HTTP after the opening budget"); - timeout(Duration::from_secs(2), late_failure_rx) - .await - .expect("delayed opening should receive its 503 response") - .unwrap(); + late_failure_rx.await.unwrap(); tokio::time::sleep(Duration::from_millis(20)).await; - let second = timeout( - Duration::from_secs(2), - backend - .clone() - .with_websocket_fast_path_budget(Duration::from_millis(500)) - .create_response( - &new_chain_request("conversation-late-failure-second"), - request_context("req_late_failure_second", Some("chatgpt-account")), - ), - ) - .await - .expect("second request should finish") - .expect("one degraded opening must not open a threshold-two breaker"); - let second_used_websocket = timeout(Duration::from_secs(2), second_used_websocket_rx) + let second = backend + .create_response( + &new_chain_request("conversation-late-failure-second"), + request_context("req_late_failure_second", Some("chatgpt-account")), + ) .await - .expect("test server should observe the second transport") - .unwrap(); - timeout(Duration::from_secs(2), server) - .await - .expect("test server should finish") - .unwrap(); + .expect("one degraded opening must not open the default breaker"); + let second_used_websocket = second_used_websocket_rx.await.unwrap(); + server.await.unwrap(); assert_eq!(first.transport, CodexBackendTransport::HttpSse); assert!(second_used_websocket); assert_eq!(second.transport, CodexBackendTransport::WebSocket); assert!(second.body.contains("resp_after_late_failure")); - timeout(Duration::from_secs(2), pool.shutdown()) - .await - .expect("pool shutdown should finish"); + pool.shutdown().await; } #[tokio::test] async fn half_open_upstream_response_should_close_origin_breaker() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let origin_key = format!("http://{addr}"); - let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { - failure_threshold: 1, - failure_window: Duration::from_secs(1), - open_duration: Duration::ZERO, - }); - let WebSocketOriginBreakerDecision::Allowed(permit) = breaker.try_acquire(&origin_key) else { - panic!("closed breaker should grant the initial permit"); - }; - permit.fast_timeout(); + let breaker_config = WebSocketOriginBreakerConfig::default(); let server = tokio::spawn(async move { + reject_websocket_openings(&listener, breaker_config.failure_threshold).await; + let (mut probe, _) = listener.accept().await.unwrap(); let opening = read_http_request(&mut probe).await; assert!(opening.starts_with("GET /codex/responses HTTP/1.1")); @@ -1197,15 +1206,33 @@ async fn half_open_upstream_response_should_close_origin_breaker() { ) .await .unwrap(); + + let (next_stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(next_stream).await; + let _payload = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response("resp_after_half_open_success", 2, 1).into(), + )) + .await + .unwrap(); }); + let pool = Arc::new(CodexWebSocketPool::default()); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), - origin_key.clone(), + format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::new(CodexWebSocketPool::default())) - .with_websocket_fast_path_budget(Duration::from_secs(1)) - .with_websocket_origin_breaker(breaker.clone()); + .with_websocket_pool(Arc::clone(&pool)); + record_websocket_origin_failures( + &backend, + breaker_config.failure_threshold, + "half-open-response-seed", + ) + .await; + tokio::time::pause(); + tokio::time::advance(breaker_config.open_duration).await; + tokio::time::resume(); let error = backend .create_response( @@ -1214,7 +1241,6 @@ async fn half_open_upstream_response_should_close_origin_breaker() { ) .await .expect_err("half-open account response should remain explicit"); - server.await.unwrap(); std::assert_matches!( error, @@ -1224,8 +1250,16 @@ async fn half_open_upstream_response_should_close_origin_breaker() { .. } if status == reqwest::StatusCode::UNAUTHORIZED ); - assert!(matches!( - breaker.try_acquire(&origin_key), - WebSocketOriginBreakerDecision::Allowed(_) - )); + let response = backend + .create_response( + &explicit_websocket_warmup_request("conversation-after-half-open-response"), + request_context("req_after_half_open_response", Some("chatgpt-account")), + ) + .await + .expect("half-open upstream response should close the origin breaker"); + server.await.unwrap(); + + assert_eq!(response.transport, CodexBackendTransport::WebSocket); + assert!(response.body.contains("resp_after_half_open_success")); + pool.shutdown().await; } diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index da9d74552..5f5e16bf7 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -4,22 +4,24 @@ use std::{ Arc, Mutex, atomic::{AtomicUsize, Ordering}, }, - time::Duration, + time::{Duration, Instant}, }; use chrono::Utc; use futures::{SinkExt, StreamExt}; +use gateway_protocol::openai::events::{TokenUsage, extract_sse_usage}; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; use provider_openai::transport::protocol::responses::{ CodexResponsesRequest, PreviousResponseScope, }; use provider_openai::transport::websocket::{ CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketRequest, - execute_response_create_request, responses_websocket_endpoint, + responses_websocket_endpoint, }; use provider_openai::transport::{ - CodexBackendClient, CodexBackendTransport, CodexClientError, CodexRequestContext, - CodexTransportDecision, CodexWebSocketPool, CodexWebSocketPoolConfig, + CodexBackendClient, CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, + CodexClientResult, CodexRequestContext, CodexResponseMetadata, CodexTransportDecision, + CodexTransportMetrics, CodexWebSocketPool, CodexWebSocketPoolConfig, WebSocketPoolDecision, }; use serde_json::{Map, Value, json}; use tokio::{ @@ -146,6 +148,155 @@ fn request_context<'a>( } } +#[derive(Debug)] +struct CollectedBackendResponse { + body: String, + transport: CodexBackendTransport, + usage: Option, + turn_state: Option, + set_cookie_headers: Vec, + rate_limit_headers: Vec<(String, String)>, + websocket_pool_decision: Option, + response_metadata: CodexResponseMetadata, + transport_metrics: CodexTransportMetrics, + connection_local_continuation: bool, +} + +trait CodexBackendClientTestExt { + async fn create_response( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult; + + async fn create_response_stream( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult; + + async fn create_response_with_pool_account_started_at( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + pool_account_id: Option<&str>, + started_at: Instant, + ) -> CodexClientResult; +} + +impl CodexBackendClientTestExt for CodexBackendClient { + async fn create_response( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let started_at = Instant::now(); + let response = self + .create_response_stream_with_pool_account(request, context, None) + .await?; + collect_backend_response(response, started_at).await + } + + async fn create_response_stream( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + self.create_response_stream_with_pool_account(request, context, None) + .await + } + + async fn create_response_with_pool_account_started_at( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + pool_account_id: Option<&str>, + started_at: Instant, + ) -> CodexClientResult { + let response = self + .create_response_stream_with_pool_account(request, context, pool_account_id) + .await?; + collect_backend_response(response, started_at).await + } +} + +async fn collect_backend_response( + response: CodexBackendStreamingResponse, + started_at: Instant, +) -> CodexClientResult { + let CodexBackendStreamingResponse { + mut body, + transport, + mut turn_state, + set_cookie_headers, + mut rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + websocket_pool_decision, + diagnostics: _, + response_metadata, + mut transport_metrics, + connection_local_continuation, + } = response; + let mut body_bytes = Vec::new(); + while let Some(chunk) = body.next().await { + let chunk = chunk?; + transport_metrics.first_event_ms.get_or_insert_with(|| { + i64::try_from(started_at.elapsed().as_millis()) + .unwrap_or(i64::MAX) + .max(1) + }); + body_bytes.extend_from_slice(&chunk); + } + if let Some(updates) = rate_limit_header_updates { + rate_limit_headers.extend(updates.lock().await.iter().cloned()); + } + if let Some(update) = turn_state_update { + turn_state = update.lock().await.clone().or(turn_state); + } + let body = String::from_utf8_lossy(&body_bytes).into_owned(); + let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?; + Ok(CollectedBackendResponse { + body, + transport, + usage, + turn_state, + set_cookie_headers, + rate_limit_headers, + websocket_pool_decision, + response_metadata, + transport_metrics, + connection_local_continuation, + }) +} + +async fn execute_response_create_request( + request: &CodexWebSocketRequest, +) -> CodexClientResult { + let endpoint = request.connection().endpoint(); + let base_url = endpoint + .strip_suffix("/codex/responses") + .unwrap_or(endpoint) + .replacen("ws://", "http://", 1) + .replacen("wss://", "https://", 1); + let mut body = serde_json::from_str::>(request.payload_text()) + .expect("test WebSocket payload should be valid JSON"); + body.remove("type"); + let mut upstream_request = CodexResponsesRequest::from_body(body); + upstream_request.use_websocket = true; + let client = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + base_url, + test_wire_profile(), + ); + client + .create_response( + &upstream_request, + request_context("req_test_websocket", Some("chatgpt-account")), + ) + .await +} + fn codex_request( model: impl Into, instructions: impl Into, diff --git a/backend/crates/providers/openai/tests/transport/websocket.rs b/backend/crates/providers/openai/tests/transport/websocket.rs index d9922104e..5b531c90c 100644 --- a/backend/crates/providers/openai/tests/transport/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/websocket.rs @@ -409,12 +409,18 @@ async fn websocket_execute_response_create_request_should_preserve_opening_error .expect_err("failed opening should surface upstream status"); server.await.unwrap(); - let CodexWebSocketExchangeError::Upstream(error) = error else { + let CodexClientError::Upstream { + status, + retry_after_seconds, + body, + .. + } = error + else { panic!("expected upstream opening error"); }; - assert_eq!(error.status_code, 429); - assert_eq!(error.retry_after_seconds, Some(33)); - assert!(error.body.contains("rate limited")); + assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); + assert_eq!(retry_after_seconds, Some(33)); + assert!(body.contains("rate limited")); } #[tokio::test] @@ -438,14 +444,14 @@ async fn websocket_execute_response_create_request_should_reject_binary_event() .expect_err("binary websocket events should be rejected"); server.await.unwrap(); - assert_eq!( - std::error::Error::source(&error).map(ToString::to_string), - Some("unexpected binary websocket event".to_string()) - ); std::assert_matches!( error, - CodexWebSocketExchangeError::PostSendAmbiguous { message, .. } + CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { + message, + source: Some(source), + }) if message.contains("unexpected binary websocket event") + && matches!(*source, CodexWebSocketExchangeError::UnexpectedBinaryEvent) ); } @@ -834,7 +840,10 @@ async fn websocket_execute_response_create_request_should_reject_invalid_complet .expect_err("invalid response.completed should be rejected"); server.await.unwrap(); - let CodexWebSocketExchangeError::PostSendAmbiguous { message, .. } = error else { + let CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { + message, .. + }) = error + else { panic!("expected post-send ambiguous websocket response"); }; assert!(message.contains("failed to parse ResponseCompleted")); @@ -879,7 +888,10 @@ async fn websocket_execute_response_create_request_should_reject_completed_witho std::assert_matches!( error, - CodexWebSocketExchangeError::PostSendAmbiguous { message, .. } + CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { + message, + .. + }) if message.contains("response.completed is missing response") ); } @@ -931,7 +943,7 @@ async fn websocket_execute_response_create_request_should_return_error_terminal_ assert!(response.body.contains("No tool output found")); } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn codex_backend_client_should_timeout_when_upstream_is_silent() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -949,8 +961,7 @@ async fn codex_backend_client_should_timeout_when_upstream_is_silent() { reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), test_wire_profile(), - ) - .with_websocket_initial_event_timeout(Some(Duration::from_millis(30))); + ); let mut request = codex_request("gpt-5.5", "be brief", Vec::new()); request.set_previous_response_id(Some("resp_previous".to_string())); request.previous_response_scope = Some(PreviousResponseScope::Persisted); @@ -970,11 +981,11 @@ async fn codex_backend_client_should_timeout_when_upstream_is_silent() { message, .. }) - if message.contains("30ms") + if message.contains("20s") ); } -#[tokio::test(start_paused = true)] +#[tokio::test] async fn websocket_response_created_should_switch_to_active_stream_idle_timeout() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -1009,8 +1020,9 @@ async fn websocket_response_created_should_switch_to_active_stream_idle_timeout( tokio::spawn(async move { execute_response_create_request(&prepared).await }); created_rx.await.unwrap(); - tokio::time::advance(Duration::from_secs(30)).await; tokio::task::yield_now().await; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(30)).await; let response = response_task .await .expect("websocket task should finish") @@ -1132,7 +1144,7 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() { ); } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn codex_backend_client_stream_should_keep_socket_after_structural_activity() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -1158,7 +1170,7 @@ async fn codex_backend_client_stream_should_keep_socket_after_structural_activit .await .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; + tokio::time::sleep(Duration::from_secs(30)).await; first_websocket .send(Message::Text( json!({ @@ -1182,8 +1194,7 @@ async fn codex_backend_client_stream_should_keep_socket_after_structural_activit reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), test_wire_profile(), - ) - .with_websocket_initial_event_timeout(Some(Duration::from_millis(30))); + ); let request = pooled_websocket_request("conversation-structural-no-pool"); let response = backend diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index de6ede43f..dbdb66718 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -1068,8 +1068,7 @@ async fn concurrent_pool_transport_counts( format!("http://{addr}"), test_wire_profile(), ) - .with_websocket_pool(Arc::clone(&pool)) - .with_websocket_fast_path_budget(Duration::from_secs(5)); + .with_websocket_pool(Arc::clone(&pool)); let start = Arc::new(tokio::sync::Barrier::new(REQUEST_COUNT + 1)); let mut requests = Vec::new(); for index in 0..REQUEST_COUNT { diff --git a/backend/crates/providers/xai/src/credential/client.rs b/backend/crates/providers/xai/src/credential/client.rs index ec774365e..675e81dc0 100644 --- a/backend/crates/providers/xai/src/credential/client.rs +++ b/backend/crates/providers/xai/src/credential/client.rs @@ -228,21 +228,11 @@ impl GrokOAuthClient { discovery: &DiscoveryDocument, grant: &RefreshTokenGrant, ) -> Result { - let mut form = vec![ + let form = vec![ FormField::public("grant_type", "refresh_token"), FormField::secret("refresh_token", grant.refresh_token().clone()), FormField::public("client_id", self.config.client_id()), ]; - if let Some(principal) = grant.principal() { - form.push(FormField::public( - "principal_type", - principal.principal_type(), - )); - form.push(FormField::secret( - "principal_id", - crate::SecretValue::new(principal.principal_id().to_owned()), - )); - } let response = self .execute( OAuthOperation::RefreshToken, diff --git a/backend/crates/providers/xai/src/credential/http.rs b/backend/crates/providers/xai/src/credential/http.rs index ddd3cb295..62a38a969 100644 --- a/backend/crates/providers/xai/src/credential/http.rs +++ b/backend/crates/providers/xai/src/credential/http.rs @@ -64,12 +64,6 @@ impl FormValue { Self::Secret(value) => value.expose(), } } - - /// Reports whether debug and telemetry layers must redact this value. - #[must_use] - pub const fn is_secret(&self) -> bool { - matches!(self, Self::Secret(_)) - } } impl fmt::Debug for FormValue { diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index 5547355c3..786a48cd7 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -479,14 +479,6 @@ impl GrokCredentialRefreshService { Ok(outcomes) } - pub async fn refresh_one( - &self, - credential: DueGrokCredential, - ) -> Result { - let policy = self.runtime_policy.load_refresh_policy().await?; - self.refresh_one_with_policy(credential, policy).await - } - async fn refresh_one_with_policy( &self, credential: DueGrokCredential, diff --git a/backend/crates/providers/xai/src/credential/repository.rs b/backend/crates/providers/xai/src/credential/repository.rs index bac64d364..62b7928a8 100644 --- a/backend/crates/providers/xai/src/credential/repository.rs +++ b/backend/crates/providers/xai/src/credential/repository.rs @@ -181,26 +181,6 @@ impl GrokCredentialAdmin { }) } - /// 批量验证 OAuth account；任一输入失败或 ID 重复时整批拒绝。 - pub fn prepare_import_batch( - &self, - inputs: &[CreateGrokCredential], - ) -> Result, GrokCredentialRepositoryError> { - if inputs.is_empty() || inputs.len() > MAX_IMPORT_BATCH { - return Err(GrokCredentialRepositoryError::InvalidInput("batch")); - } - let mut ids = BTreeSet::new(); - let mut prepared = Vec::with_capacity(inputs.len()); - for input in inputs { - let account = self.prepare_import(input)?; - if !ids.insert(account.account.id().clone()) { - return Err(GrokCredentialRepositoryError::InvalidInput("account_id")); - } - prepared.push(account); - } - Ok(prepared) - } - /// 把已验证 token 的 Provider 生命周期事实一次性投影为 Core 创建命令。 pub fn prepare_verified_account( &self, diff --git a/backend/crates/providers/xai/src/credential/token.rs b/backend/crates/providers/xai/src/credential/token.rs index 14c0b8f93..f95cf8e57 100644 --- a/backend/crates/providers/xai/src/credential/token.rs +++ b/backend/crates/providers/xai/src/credential/token.rs @@ -69,33 +69,18 @@ impl OAuthPrincipal { #[derive(Clone)] pub struct RefreshTokenGrant { refresh_token: SecretValue, - principal: Option, } impl RefreshTokenGrant { /// Creates a refresh grant for an existing verified credential. #[must_use] pub fn new(refresh_token: SecretValue) -> Self { - Self { - refresh_token, - principal: None, - } - } - - /// Attaches optional team principal metadata. - #[must_use] - pub fn with_principal(mut self, principal: OAuthPrincipal) -> Self { - self.principal = Some(principal); - self + Self { refresh_token } } pub(crate) fn refresh_token(&self) -> &SecretValue { &self.refresh_token } - - pub(crate) fn principal(&self) -> Option<&OAuthPrincipal> { - self.principal.as_ref() - } } impl fmt::Debug for RefreshTokenGrant { @@ -103,7 +88,6 @@ impl fmt::Debug for RefreshTokenGrant { formatter .debug_struct("RefreshTokenGrant") .field("refresh_token", &"[REDACTED]") - .field("principal", &self.principal) .finish() } } diff --git a/backend/crates/providers/xai/src/transport/canonical.rs b/backend/crates/providers/xai/src/transport/canonical.rs index b4fef4810..80332f266 100644 --- a/backend/crates/providers/xai/src/transport/canonical.rs +++ b/backend/crates/providers/xai/src/transport/canonical.rs @@ -638,6 +638,8 @@ fn core_usage(usage: TokenUsage) -> Usage { normalized.cached_tokens = Some(usage.cached_tokens); normalized.cache_write_tokens = Some(usage.cache_write_tokens); normalized.reasoning_tokens = Some(usage.reasoning_tokens); + normalized.image_input_tokens = Some(usage.image_input_tokens); + normalized.image_output_tokens = Some(usage.image_output_tokens); normalized.total_tokens = Some(usage.total_tokens); normalized } diff --git a/backend/crates/providers/xai/src/transport/headers.rs b/backend/crates/providers/xai/src/transport/headers.rs index de15e77bb..2b70362b3 100644 --- a/backend/crates/providers/xai/src/transport/headers.rs +++ b/backend/crates/providers/xai/src/transport/headers.rs @@ -53,12 +53,6 @@ impl GrokHeaderValue { Self::Sensitive(value) => value.expose(), } } - - /// Reports whether the value must be excluded from logs and diagnostics. - #[must_use] - pub const fn is_sensitive(&self) -> bool { - matches!(self, Self::Sensitive(_)) - } } impl fmt::Debug for GrokHeaderValue { diff --git a/backend/crates/providers/xai/src/transport/network.rs b/backend/crates/providers/xai/src/transport/network.rs index c1a199875..1b5da13b7 100644 --- a/backend/crates/providers/xai/src/transport/network.rs +++ b/backend/crates/providers/xai/src/transport/network.rs @@ -177,14 +177,14 @@ impl GrokEndpointPolicy for OfficialGrokEndpointPolicy { &self, timeout: Option, ) -> Result { - build_official_client(OFFICIAL_OAUTH_HOST, timeout) + build_official_client(GrokDnsResolutionPolicy::official_oauth(), timeout) } fn build_inference_client( &self, timeout: Option, ) -> Result { - build_official_client(OFFICIAL_INFERENCE_HOST, timeout) + build_official_client(GrokDnsResolutionPolicy::official_inference(), timeout) } fn validate_oauth(&self, url: &Url) -> bool { @@ -579,7 +579,7 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport { } fn build_official_client( - host: &'static str, + dns_policy: GrokDnsResolutionPolicy, timeout: Option, ) -> Result { let mut builder = Client::builder() @@ -591,7 +591,7 @@ fn build_official_client( .http2_adaptive_window(true) .tcp_nodelay(true) .https_only(true) - .dns_resolver(Arc::new(StrictDnsResolver::new(host)?)); + .dns_resolver(Arc::new(StrictDnsResolver::new(dns_policy)?)); if let Some(timeout) = timeout { builder = builder.timeout(timeout); } @@ -633,9 +633,9 @@ struct StrictDnsResolver { } impl StrictDnsResolver { - fn new(allowed_host: &'static str) -> Result { + fn new(policy: GrokDnsResolutionPolicy) -> Result { Ok(Self { - policy: GrokDnsResolutionPolicy { allowed_host }, + policy, trusted_doh: TrustedDohResolver::new()?, }) } diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs index 9630039fb..2416d5fda 100644 --- a/backend/crates/providers/xai/src/transport/request.rs +++ b/backend/crates/providers/xai/src/transport/request.rs @@ -168,33 +168,6 @@ impl GrokResponsesRequest { self.turn_index = None; } - pub fn strip_reasoning_encrypted_content(&mut self) -> bool { - let Some(items) = self.body.get_mut("input").and_then(Value::as_array_mut) else { - return false; - }; - let original_len = items.len(); - let mut changed = false; - items.retain_mut(|item| { - let Some(item) = item.as_object_mut() else { - return true; - }; - if string_field(item, "type") != "reasoning" { - return true; - } - for field in ["encrypted_content", "id", "status"] { - changed |= item.remove(field).is_some(); - } - let portable = ["summary", "content"].into_iter().any(|field| { - item.get(field) - .and_then(Value::as_array) - .is_some_and(|values| !values.is_empty()) - }); - changed |= !portable; - portable - }); - changed || items.len() != original_len - } - pub fn encode( request: &GenerateRequest, upstream_model: &str, diff --git a/backend/crates/providers/xai/tests/credential/contract.rs b/backend/crates/providers/xai/tests/credential/contract.rs index aba87fa89..71c8a1bd6 100644 --- a/backend/crates/providers/xai/tests/credential/contract.rs +++ b/backend/crates/providers/xai/tests/credential/contract.rs @@ -396,39 +396,3 @@ fn oauth_secret_debug_never_exposes_plaintext() { assert!(debug.contains("REDACTED")); assert!(!debug.contains("visible-only-to-provider")); } - -#[test] -fn admin_batch_preparation_is_all_or_nothing_and_rejects_duplicate_ids() { - let valid = [ - create_input("batch-a", "subject-batch-a"), - create_input("batch-b", "subject-batch-b"), - ]; - let prepared = GrokCredentialAdmin - .prepare_import_batch(&valid) - .expect("valid batch"); - assert_eq!(prepared.len(), 2); - - let duplicate = [ - create_input("batch-duplicate", "subject-one"), - create_input("batch-duplicate", "subject-two"), - ]; - assert!(matches!( - GrokCredentialAdmin.prepare_import_batch(&duplicate), - Err(GrokCredentialRepositoryError::InvalidInput("account_id")) - )); -} - -#[test] -fn admin_batch_preparation_enforces_empty_and_upper_boundaries() { - assert!(matches!( - GrokCredentialAdmin.prepare_import_batch(&[]), - Err(GrokCredentialRepositoryError::InvalidInput("batch")) - )); - let oversized = (0..201) - .map(|index| create_input(&format!("batch-{index}"), &format!("subject-{index}"))) - .collect::>(); - assert!(matches!( - GrokCredentialAdmin.prepare_import_batch(&oversized), - Err(GrokCredentialRepositoryError::InvalidInput("batch")) - )); -} diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index 48a5d90a2..0d1c7aa65 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -12,16 +12,15 @@ use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError, }; use provider_xai::{ - GrokCredentialAdmin, GrokCredentialCatalogCache, GrokCredentialRecovery, - GrokCredentialRecoveryOutcome, GrokCredentialRefreshError, GrokCredentialRefreshOutcome, - GrokCredentialRefreshService, GrokCredentialRefresher, GrokCredentialRepository, - GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportFuture, - GrokModelCatalogTransportResponse, GrokRefreshFailure, GrokRefreshTokens, - RotateManagedGrokCredential, SecretValue, + GrokCredentialCatalogCache, GrokCredentialRecovery, GrokCredentialRecoveryOutcome, + GrokCredentialRefreshError, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, + GrokCredentialRefresher, GrokCredentialRepository, GrokModelCatalogRequest, + GrokModelCatalogTransport, GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, + GrokRefreshFailure, GrokRefreshTokens, SecretValue, }; use crate::support::{ - MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input, credential_object, profile, + MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input, credential_object, runtime_policy, seed_input, }; @@ -420,48 +419,6 @@ async fn pre_send_transient_failure_applies_bounded_cooldown() { ); } -#[tokio::test] -async fn stale_due_snapshot_cannot_overwrite_newer_revision() { - let input = due_input("stale"); - let id = input.account_id.clone(); - let (store, repository, _, service) = - fixture(input, [Err(GrokRefreshFailure::InvalidGrant)], true).await; - let due = repository - .list_due_refresh() - .await - .expect("due list") - .pop() - .expect("due account"); - let current = store - .load_credential(&id, CredentialRevision::new(1).expect("revision")) - .await - .expect("current credential"); - let prepared = GrokCredentialAdmin - .prepare_rotation(&RotateManagedGrokCredential { - current, - secret: provider_xai::GrokOAuthSecret { - access_token: SecretValue::new("admin-access"), - refresh_token: SecretValue::new("admin-refresh"), - id_token: None, - scope: provider_xai::OFFICIAL_SCOPES.join(" "), - }, - verified_account: profile("subject-stale"), - next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), - }) - .expect("admin rotation"); - assert!(matches!( - store - .compare_and_swap_credential(prepared.credential) - .await - .expect("persist admin rotation"), - CredentialCasOutcome::Updated(revision) if revision.get() == 2 - )); - assert!(matches!( - service.refresh_one(due).await.expect("stale outcome"), - GrokCredentialRefreshOutcome::Stale { account_id } if account_id == id - )); -} - #[tokio::test] async fn invalid_refresh_lifetime_is_rejected_without_cas_write() { let input = due_input("bad-lifetime"); diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index 03e80ca05..6f332e1a3 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -6,7 +6,7 @@ use std::time::{Duration, SystemTime}; use futures::{StreamExt, stream}; use gateway_core::engine::continuation::{ - ContinuationBinding, NativeContinuationPin, NativeContinuationReuse, PreviousResponseId, + ContinuationBinding, NativeContinuationPin, PreviousResponseId, }; use gateway_core::engine::credential::{ AccountSelectionPolicy, CredentialRevision, ProviderAccountStore, RotationStrategy, @@ -1154,7 +1154,6 @@ async fn native_previous_response_is_rejected_before_selection() { ProviderKind::new("openai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let error = match provider .execute( @@ -1189,7 +1188,6 @@ async fn native_previous_response_pins_account_and_sends_upstream_handle() { ProviderKind::new("xai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let mut stream = provider .execute( @@ -1236,7 +1234,6 @@ async fn native_previous_response_does_not_allow_quota_or_rate_limit_account_rot ProviderKind::new("xai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let mut stream = provider .execute( @@ -1330,7 +1327,6 @@ async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_acco ProviderKind::new("xai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let continued_operation = operation_with_state( json!({ @@ -1457,7 +1453,6 @@ async fn replay_owner_should_reencode_custom_apply_patch_call_for_grok() { ProviderKind::new("xai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let replay_operation = operation_with_state( json!({ @@ -1534,7 +1529,6 @@ async fn missing_native_response_should_be_replay_safe_for_the_same_account() { ProviderKind::new("xai").expect("provider"), instance_id(), account_id("provider"), - NativeContinuationReuse::Reusable, ); let operation = operation_with_state( json!({ diff --git a/backend/crates/providers/xai/tests/transport/canonical.rs b/backend/crates/providers/xai/tests/transport/canonical.rs index 7f037ea91..2179c6e9d 100644 --- a/backend/crates/providers/xai/tests/transport/canonical.rs +++ b/backend/crates/providers/xai/tests/transport/canonical.rs @@ -115,6 +115,23 @@ fn decoder_should_normalize_text_usage_and_completion() { )); } +#[test] +fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() { + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_image_usage\",\"model\":\"grok-4.5\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_usage\",\"model\":\"grok-4.5\",\"status\":\"completed\",\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"total_tokens\":17},\"tool_usage\":{\"image_gen\":{\"input_tokens\":31,\"output_tokens\":9}}}}\n\n", + ); + let events = decode_canonical(body.as_bytes()).expect("canonical image usage response"); + let image_usage = events.iter().find_map(|event| match event { + GatewayEvent::Usage(usage) => Some((usage.image_input_tokens, usage.image_output_tokens)), + _ => None, + }); + + assert_eq!(image_usage, Some((Some(31), Some(9)))); +} + #[test] fn decoder_should_leave_cost_unavailable_when_upstream_omits_it() { let body = concat!( diff --git a/backend/migrations/0001_initial.sql b/backend/migrations/0001_initial.sql index 11b35e3de..3b310d44c 100644 --- a/backend/migrations/0001_initial.sql +++ b/backend/migrations/0001_initial.sql @@ -244,6 +244,7 @@ create table model_requests ( provider_account_ref text, upstream_transport text, http_version text, + websocket_pool text, attempt_count integer not null default 0, upstream_send_state text not null default 'not_sent', downstream_committed_at timestamptz, @@ -262,7 +263,10 @@ create table model_requests ( cached_tokens bigint, cache_write_tokens bigint, reasoning_tokens bigint, + image_input_tokens bigint, + image_output_tokens bigint, total_tokens bigint, + image_generation_succeeded boolean, cost_source text not null default 'unavailable', cost_amount numeric(20, 10), cost_currency text, @@ -282,6 +286,7 @@ create table model_requests ( request_kind text, subagent_kind text, compact boolean not null default false, + image_generation_requested boolean not null default false, started_at timestamptz not null, deadline_at timestamptz not null, completed_at timestamptz, @@ -326,8 +331,26 @@ create table model_requests ( and (cached_tokens is null or cached_tokens >= 0) and (cache_write_tokens is null or cache_write_tokens >= 0) and (reasoning_tokens is null or reasoning_tokens >= 0) + and (image_input_tokens is null or image_input_tokens >= 0) + and (image_output_tokens is null or image_output_tokens >= 0) and (total_tokens is null or total_tokens >= 0) ), + constraint model_requests_image_generation_ck check ( + ( + not image_generation_requested + and image_generation_succeeded is null + ) + or ( + image_generation_requested + and ( + (outcome = 'running' and image_generation_succeeded is null) + or (outcome <> 'running' and image_generation_succeeded is not null) + ) + ) + ), + constraint model_requests_websocket_pool_ck check ( + websocket_pool is null or websocket_pool in ('new', 'reuse') + ), constraint model_requests_cost_ck check ( cost_source in ('provider_reported', 'calculated', 'unavailable') and (