diff --git a/backend/crates/gateway-api/src/openai/images/http.rs b/backend/crates/gateway-api/src/openai/images/http.rs index 8d9b5d3a3..45eabbe13 100644 --- a/backend/crates/gateway-api/src/openai/images/http.rs +++ b/backend/crates/gateway-api/src/openai/images/http.rs @@ -10,14 +10,14 @@ use axum::{ }; use gateway_core::error::{GatewayError, GatewayErrorKind}; use gateway_core::operation::{ImageRequest, ImageRequestKind, Operation, RawJsonPayload}; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::ApiState; use crate::openai::{ auth::{authenticate_client, client_access_error_response}, error::gateway_error_response, provider_endpoint::collect_raw_json_response, - responses::request_client_context, + responses::{OpenAiRequestHeaders, request_client_context}, }; const OPENAI_PROTOCOL: &str = "openai"; @@ -95,7 +95,7 @@ fn image_operation( headers: &HeaderMap, kind: ImageRequestKind, ) -> Result { - let mut context = Map::new(); + let mut context = OpenAiRequestHeaders::from_headers(headers).session_context(); if let Some(turn_id) = headers .get("x-codex-image-turn-id") .and_then(|value| value.to_str().ok()) diff --git a/backend/crates/gateway-api/src/openai/responses/request.rs b/backend/crates/gateway-api/src/openai/responses/request.rs index f2aadfd54..43574a2d0 100644 --- a/backend/crates/gateway-api/src/openai/responses/request.rs +++ b/backend/crates/gateway-api/src/openai/responses/request.rs @@ -93,6 +93,15 @@ impl OpenAiRequestHeaders { } } + /// 独立端点只消费会话身份与 turn metadata，不携带 Responses 的连接状态。 + pub(crate) fn session_context(&self) -> Map { + let mut context = Map::new(); + insert_protocol_context(&mut context, "session_id", self.session_id.as_ref()); + insert_protocol_context(&mut context, "thread_id", self.thread_id.as_ref()); + insert_protocol_context(&mut context, "turn_metadata", self.turn_metadata.as_ref()); + context + } + fn protocol_context( &self, use_websocket: Option, diff --git a/backend/crates/gateway-api/src/openai/search/http.rs b/backend/crates/gateway-api/src/openai/search/http.rs index c5ba5db02..ad7f21bbf 100644 --- a/backend/crates/gateway-api/src/openai/search/http.rs +++ b/backend/crates/gateway-api/src/openai/search/http.rs @@ -10,18 +10,16 @@ use axum::{ }; use gateway_core::error::{GatewayError, GatewayErrorKind}; use gateway_core::operation::{Operation, RawJsonPayload, StandaloneSearchRequest}; -use serde_json::{Map, Value}; use crate::ApiState; use crate::openai::{ auth::{authenticate_client, client_access_error_response}, error::gateway_error_response, provider_endpoint::collect_raw_json_response, - responses::request_client_context, + responses::{OpenAiRequestHeaders, request_client_context}, }; const OPENAI_PROTOCOL: &str = "openai"; -const TURN_METADATA_CONTEXT_KEY: &str = "turn_metadata"; /// `POST /v1/alpha/search`。 pub(crate) async fn standalone_search( @@ -54,16 +52,7 @@ pub(crate) async fn standalone_search( } fn search_operation(body: Bytes, headers: &HeaderMap) -> Result { - let mut context = Map::new(); - if let Some(turn_metadata) = headers - .get("x-codex-turn-metadata") - .and_then(|value| value.to_str().ok()) - { - context.insert( - TURN_METADATA_CONTEXT_KEY.to_owned(), - Value::String(turn_metadata.to_owned()), - ); - } + let context = OpenAiRequestHeaders::from_headers(headers).session_context(); let payload = RawJsonPayload::new(OPENAI_PROTOCOL, body) .map_err(|_| { GatewayError::new( diff --git a/backend/crates/gateway-api/tests/openai/images/mod.rs b/backend/crates/gateway-api/tests/openai/images/mod.rs index 093748556..71797d09f 100644 --- a/backend/crates/gateway-api/tests/openai/images/mod.rs +++ b/backend/crates/gateway-api/tests/openai/images/mod.rs @@ -260,6 +260,8 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio .header(AUTHORIZATION, "Bearer sk_images_test") .header("content-type", "application/json") .header("x-codex-image-turn-id", "turn_image_route") + .header("session-id", "root-image-session") + .header("thread-id", "child-image-thread") .body(Body::from(body.to_vec())) .expect("image request"), ) @@ -286,7 +288,7 @@ async fn image_routes_should_not_decode_bodies_and_should_preserve_both_directio assert_eq!(captured.body.as_ref(), *body); assert_eq!( captured.context, - json!({"image_turn_id": "turn_image_route"}) + json!({"image_turn_id": "turn_image_route", "session_id": "root-image-session", "thread_id": "child-image-thread"}) ); } assert_eq!( diff --git a/backend/crates/gateway-api/tests/openai/search/mod.rs b/backend/crates/gateway-api/tests/openai/search/mod.rs index afc684068..f63fc9f10 100644 --- a/backend/crates/gateway-api/tests/openai/search/mod.rs +++ b/backend/crates/gateway-api/tests/openai/search/mod.rs @@ -162,6 +162,8 @@ async fn search_route_should_preserve_request_and_success_response_bytes() { .header(AUTHORIZATION, "Bearer sk_search_test") .header("content-type", "application/json") .header("x-codex-turn-metadata", turn_metadata) + .header("session-id", "root-session") + .header("thread-id", "child-thread") .body(Body::from(request_body.to_vec())) .expect("search request"), ) @@ -184,7 +186,10 @@ async fn search_route_should_preserve_request_and_success_response_bytes() { assert_eq!(captured[0].endpoint, "/v1/alpha/search"); assert_eq!(captured[0].transport, ClientTransport::HttpJson); assert_eq!(captured[0].body.as_ref(), request_body); - assert_eq!(captured[0].context, json!({"turn_metadata": turn_metadata})); + assert_eq!( + captured[0].context, + json!({"turn_metadata": turn_metadata, "session_id": "root-session", "thread_id": "child-thread"}) + ); assert_eq!( *execution .committed_statuses diff --git a/backend/crates/gateway-core/src/account/selection.rs b/backend/crates/gateway-core/src/account/selection.rs index eae42abe5..2ef92ce41 100644 --- a/backend/crates/gateway-core/src/account/selection.rs +++ b/backend/crates/gateway-core/src/account/selection.rs @@ -386,6 +386,8 @@ pub struct AccountSelectionContext { pub now: SystemTime, pub excluded_accounts: BTreeSet, pub preferred_account: Option, + /// 优先账号可调度时，是否优先于其它账号的持久权重。 + pub preferred_account_overrides_weight: bool, pub round_robin_cursor: u64, pub eligibility: AccountEligibilityPolicy, pub account_scope: Option>, @@ -516,7 +518,9 @@ impl AccountSelector { { Some(candidate) => match self.scheduling_blocker(candidate, context) { Some(blocker) => PreferredAccountSelection::Blocked(blocker), - None if candidate.account.weight() < highest_weight => { + None if !context.preferred_account_overrides_weight + && candidate.account.weight() < highest_weight => + { PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight) } None => { diff --git a/backend/crates/gateway-core/tests/account/mod.rs b/backend/crates/gateway-core/tests/account/mod.rs index 73ba9bfaf..e4e44dde2 100644 --- a/backend/crates/gateway-core/tests/account/mod.rs +++ b/backend/crates/gateway-core/tests/account/mod.rs @@ -100,6 +100,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext { now: SystemTime::now(), excluded_accounts: BTreeSet::new(), preferred_account: None, + preferred_account_overrides_weight: false, round_robin_cursor: 0, eligibility: AccountEligibilityPolicy::Enforce, account_scope: None, @@ -417,6 +418,84 @@ fn selector_should_not_let_soft_affinity_bypass_a_higher_weight_account() { ); } +#[test] +fn selector_should_prioritize_affinity_over_weight_when_requested() { + let candidates = [ + weighted_candidate("acct_preferred", 1, 0), + weighted_candidate("acct_higher", 100, 0), + ]; + for strategy in [ + RotationStrategy::Smart, + RotationStrategy::QuotaResetPriority, + RotationStrategy::RoundRobin, + RotationStrategy::Sticky, + ] { + let mut selection = context(strategy); + selection.preferred_account = Some(candidates[0].account.id().clone()); + selection.preferred_account_overrides_weight = true; + let selected = AccountSelector + .select(&candidates, &selection) + .expect("preferred account"); + assert_eq!( + ( + selected.candidate().account.id().as_str(), + selected.preferred() + ), + ("acct_preferred", PreferredAccountSelection::Hit), + "{strategy:?}" + ); + } +} + +#[test] +fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints() { + for blocker in [ + AccountSchedulingBlocker::Excluded, + AccountSchedulingBlocker::ConcurrencyLimit, + AccountSchedulingBlocker::RequestInterval, + AccountSchedulingBlocker::LocalAvailability, + ] { + let mut candidates = [ + weighted_candidate("acct_preferred", 1, 0), + weighted_candidate("acct_higher", 100, 0), + ]; + let mut selection = context(RotationStrategy::Smart); + selection.preferred_account = Some(candidates[0].account.id().clone()); + selection.preferred_account_overrides_weight = true; + match blocker { + AccountSchedulingBlocker::Excluded => { + selection + .excluded_accounts + .insert(candidates[0].account.id().clone()); + } + AccountSchedulingBlocker::ConcurrencyLimit => candidates[0].signals.in_flight = 3, + AccountSchedulingBlocker::RequestInterval => { + candidates[0].signals.last_started_at = Some(selection.now); + selection.policy = AccountSelectionPolicy::new( + RotationStrategy::Smart, + NonZeroU32::new(3).expect("limit"), + Duration::from_secs(1), + ); + } + AccountSchedulingBlocker::LocalAvailability => { + candidates[0].signals.rate_limited_until = + Some(selection.now + Duration::from_secs(60)) + } + _ => unreachable!("test cases"), + } + let selected = AccountSelector + .select(&candidates, &selection) + .expect("fallback account"); + assert_eq!( + ( + selected.candidate().account.id().as_str(), + selected.preferred() + ), + ("acct_higher", PreferredAccountSelection::Blocked(blocker)) + ); + } +} + #[test] fn selector_should_fall_back_when_a_higher_weight_account_reaches_its_override() { let mut saturated = weighted_candidate("acct_high_weight", 100, 1); diff --git a/backend/crates/gateway-protocol/src/openai/codex.rs b/backend/crates/gateway-protocol/src/openai/codex.rs index c314eb0e0..c191be21f 100644 --- a/backend/crates/gateway-protocol/src/openai/codex.rs +++ b/backend/crates/gateway-protocol/src/openai/codex.rs @@ -4,6 +4,41 @@ const MULTI_AGENT_MODE_OPEN_TAG: &str = ""; const MULTI_AGENT_MODE_CLOSE_TAG: &str = ""; const PROACTIVE_MULTI_AGENT_MODE_PREFIX: &str = "Proactive multi-agent delegation is active."; +/// 各 Codex 端点共用的显式根会话身份。连接边界的 session_id 优先， +/// 正文及官方 metadata 只做回退；不从 turn、window 或请求内容猜测归属。 +#[must_use] +pub fn codex_session_id(body: &Map, context: &Map) -> Option { + codex_identity_field(body, context, "session_id") +} + +/// 各 Codex 端点共用的显式线程身份，与根会话采用相同的来源优先级。 +#[must_use] +pub fn codex_thread_id(body: &Map, context: &Map) -> Option { + codex_identity_field(body, context, "thread_id") +} + +fn codex_identity_field( + body: &Map, + context: &Map, + field: &str, +) -> Option { + non_empty_string(context.get(field)) + .or_else(|| non_empty_string(body.get(field))) + .or_else(|| client_metadata_string(body, field)) + .map(str::to_owned) + .or_else(|| { + let metadata = non_empty_string(context.get("turn_metadata")) + .or_else(|| request_turn_metadata(body))?; + let metadata: Value = serde_json::from_str(metadata).ok()?; + non_empty_string(metadata.get(field)).map(str::to_owned) + }) + .or_else(|| { + body.get("context") + .and_then(|context| non_empty_string(context.get(field))) + .map(str::to_owned) + }) +} + /// 从 OpenAI Responses 请求中提取的稳定 Codex 请求语义。 #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct CodexResponsesRequestSemantics { diff --git a/backend/crates/gateway-protocol/src/openai/mod.rs b/backend/crates/gateway-protocol/src/openai/mod.rs index 3583e73dc..e16623724 100644 --- a/backend/crates/gateway-protocol/src/openai/mod.rs +++ b/backend/crates/gateway-protocol/src/openai/mod.rs @@ -18,5 +18,5 @@ pub mod sse; pub use codex::{ CodexResponsesRequestSemantics, codex_responses_request_semantics, - codex_responses_request_semantics_with_turn_metadata, + codex_responses_request_semantics_with_turn_metadata, codex_session_id, codex_thread_id, }; diff --git a/backend/crates/gateway-protocol/tests/openai/codex.rs b/backend/crates/gateway-protocol/tests/openai/codex.rs index 052ebf9ab..f20b939f3 100644 --- a/backend/crates/gateway-protocol/tests/openai/codex.rs +++ b/backend/crates/gateway-protocol/tests/openai/codex.rs @@ -1,6 +1,101 @@ -use gateway_protocol::openai::codex_responses_request_semantics; +use gateway_protocol::openai::{ + codex_responses_request_semantics, codex_session_id, codex_thread_id, +}; use serde_json::{Map, Value, json}; +#[test] +fn codex_session_should_prefer_the_connection_identity_over_body_metadata() { + let body = object( + json!({"session_id":"body-session", "client_metadata":{"session_id":"metadata-session"}}), + ); + let context = object( + json!({"session_id":" root-session ", "turn_metadata":"{\"session_id\":\"turn-session\"}"}), + ); + assert_eq!( + codex_session_id(&body, &context).as_deref(), + Some("root-session") + ); +} + +#[test] +fn codex_session_should_accept_explicit_metadata_and_ignore_unrelated_ids() { + for (body, context) in [ + (json!({"session_id":"root"}), json!({})), + (json!({"client_metadata":{"session_id":"root"}}), json!({})), + ( + json!({"client_metadata":{"x-codex-turn-metadata":"{\"session_id\":\"root\"}"}}), + json!({}), + ), + ( + json!({}), + json!({"turn_metadata":"{\"session_id\":\"root\"}"}), + ), + (json!({"context":{"session_id":"root"}}), json!({})), + ] { + assert_eq!( + codex_session_id(&object(body), &object(context)).as_deref(), + Some("root") + ); + } + for body in [ + json!({"session_id":" ", "thread_id":"child", "turn_id":"turn", "id":"request"}), + json!({"client_metadata":{"x-codex-turn-metadata":"invalid JSON"}}), + ] { + assert!(codex_session_id(&object(body), &Map::new()).is_none()); + } +} + +#[test] +fn codex_thread_should_use_explicit_identity_and_never_guess_from_other_ids() { + for (body, context, expected) in [ + ( + json!({"thread_id":"body", "client_metadata":{"thread_id":"metadata"}}), + json!({"thread_id":" child "}), + Some("child"), + ), + ( + json!({"thread_id":"body", "client_metadata":{"thread_id":"metadata"}}), + json!({}), + Some("body"), + ), + ( + json!({"client_metadata":{"thread_id":"metadata"}}), + json!({}), + Some("metadata"), + ), + ( + json!({"client_metadata":{"x-codex-turn-metadata":"{\"thread_id\":\"metadata-thread\"}"}}), + json!({}), + Some("metadata-thread"), + ), + ( + json!({"turnMetadata":"{\"thread_id\":\"body-turn\"}"}), + json!({"turn_metadata":"{\"thread_id\":\"header-turn\"}"}), + Some("header-turn"), + ), + ( + json!({"context":{"thread_id":"context-thread"}}), + json!({}), + Some("context-thread"), + ), + ( + json!({"thread_id":" ", "session_id":"root", "turn_id":"turn", "agent_name":"agent"}), + json!({}), + None, + ), + ( + json!({"client_metadata":{"x-codex-turn-metadata":"invalid JSON"}}), + json!({}), + None, + ), + ] { + assert_eq!( + codex_thread_id(&object(body), &object(context)).as_deref(), + expected + ); + } +} + #[test] fn codex_semantics_should_use_protocol_context_before_body_metadata() { let body = json!({ diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index a5f6b8149..662f75644 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -1,20 +1,23 @@ -//! OpenAI 根会话锚点到 Store 不透明账号亲和键及诊断上下文的单向派生。 +//! OpenAI 会话及子线程到 Store 不透明账号亲和键及诊断上下文的单向派生。 use std::time::Duration; +use gateway_core::operation::RawJsonPayload; use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::ProviderSessionAffinityKey; +use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; use crate::transport::protocol::responses::CodexResponsesRequest; use crate::transport::request::derive_conversation_anchor; const AFFINITY_KEY_HASH_LENGTH: usize = 12; -pub(crate) const CODEX_ROOT_SESSION_TTL: Duration = Duration::from_secs(4 * 60 * 60); +pub(crate) const CODEX_ROOT_SESSION_TTL: Duration = Duration::from_secs(24 * 60 * 60); /// 一次请求派生出的账号亲和键及其结构化日志上下文。 pub(crate) struct CodexSessionAffinity { key: ProviderSessionAffinityKey, + root_key: Option, key_hash: String, anchor_source: &'static str, anchor: String, @@ -27,6 +30,30 @@ impl CodexSessionAffinity { &self.key } + pub(crate) fn root_key(&self) -> Option<&ProviderSessionAffinityKey> { + self.root_key.as_ref() + } + + /// 子线程首次选号继承根会话偏好，之后仅更新自己的绑定。 + fn with_thread(mut self, thread_id: Option<&str>) -> Option { + if let Some(thread_id) = non_empty(thread_id) + && self + .session_id + .as_deref() + .is_some_and(|root| root != thread_id) + { + let child_key = opaque_affinity_key( + "child-thread", + &format!("{}\0{thread_id}", self.key.expose_to_store()), + )?; + self.root_key = Some(std::mem::replace(&mut self.key, child_key)); + self.key_hash = short_key_hash(&self.key); + self.anchor_source = "child-thread"; + self.anchor = thread_id.to_owned(); + } + Some(self) + } + #[must_use] pub(crate) fn key_hash(&self) -> &str { &self.key_hash @@ -84,6 +111,37 @@ pub(crate) fn derive_codex_session_affinity( ) -> Option { let session_id = non_empty(request.client_session_id.as_deref()).map(str::to_owned); let (anchor_source, anchor) = derive_account_affinity_anchor(request)?; + session_affinity(anchor_source, anchor, session_id, client_api_key_id)? + .with_thread(request.client_thread_id.as_deref()) +} + +/// 原始 JSON 端点只读取会话身份，发送时仍保留原始字节。Search 的 `id` 是官方 +/// 根 session_id，必须与 Responses 共用命名空间，不能另建一份账号亲和。 +pub(crate) fn derive_codex_endpoint_session_affinity( + payload: &RawJsonPayload, + client_api_key_id: &ClientApiKeyId, + body_session_field: &str, +) -> Option { + let body = serde_json::from_slice::>(payload.body()).unwrap_or_default(); + let session_id = + gateway_protocol::openai::codex_session_id(&body, payload.context()).or_else(|| { + non_empty(body.get(body_session_field).and_then(Value::as_str)).map(str::to_owned) + })?; + session_affinity( + "root-session", + session_id.clone(), + Some(session_id), + client_api_key_id, + )? + .with_thread(gateway_protocol::openai::codex_thread_id(&body, payload.context()).as_deref()) +} + +fn session_affinity( + anchor_source: &'static str, + anchor: String, + session_id: Option, + client_api_key_id: &ClientApiKeyId, +) -> Option { let session_key = opaque_affinity_key(anchor_source, &anchor)?; let key = opaque_affinity_key( "client-session", @@ -93,14 +151,10 @@ pub(crate) fn derive_codex_session_affinity( session_key.expose_to_store() ), )?; - // 亲和键本身已经是 SHA-256；日志沿用 WebSocket 诊断的 12 位短哈希长度。 - let key_hash = key - .expose_to_store() - .chars() - .take(AFFINITY_KEY_HASH_LENGTH) - .collect(); + let key_hash = short_key_hash(&key); Some(CodexSessionAffinity { key, + root_key: None, key_hash, anchor_source, anchor, @@ -108,9 +162,15 @@ pub(crate) fn derive_codex_session_affinity( }) } -/// 账号选择优先使用根会话事实；`local_conversation_id` 仍可按 child thread 隔离 -/// WebSocket 与 continuation 状态，但 child 身份不再拆分账号首选项。账号不可调度时 -/// selector 仍可按既有策略切换，并由 request scope 清理账号绑定状态。 +fn short_key_hash(key: &ProviderSessionAffinityKey) -> String { + // 亲和键本身已经是 SHA-256；日志沿用 WebSocket 诊断的 12 位短哈希长度。 + key.expose_to_store() + .chars() + .take(AFFINITY_KEY_HASH_LENGTH) + .collect() +} + +/// 先确定根会话锚点，显式子线程在此基础上派生自己的绑定。 fn derive_account_affinity_anchor( request: &CodexResponsesRequest, ) -> Option<(&'static str, String)> { diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 3daefc374..26c8371a7 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -16,7 +16,8 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_session_affinity, + CodexSessionAffinity, derive_codex_cyber_policy_session_key, + derive_codex_endpoint_session_affinity, derive_codex_session_affinity, derive_previous_response_id_hash, }; pub(crate) use oauth::oauth_owner_ref; diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 6ab07c734..c6040607c 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -96,6 +96,7 @@ pub struct SelectCodexCredential<'a> { pub(crate) struct SelectCodexProviderEndpointCredential<'a> { pub request_url: &'a Url, pub attempt: &'a AttemptContext, + pub session_affinity: Option<&'a CodexSessionAffinity>, } struct CredentialSelectionInput<'a> { @@ -160,6 +161,7 @@ struct AffinitySelection { bound_account: Option, preferred_account: Option, escape_reason: Option, + inherited: bool, } impl AffinitySelection { @@ -168,6 +170,7 @@ impl AffinitySelection { bound_account: Some(account_id.clone()), preferred_account: Some(account_id), escape_reason: None, + inherited: false, } } @@ -176,6 +179,7 @@ impl AffinitySelection { bound_account: Some(account_id), preferred_account: None, escape_reason: Some(reason), + inherited: false, } } @@ -231,7 +235,8 @@ impl AffinitySelection { fn telemetry(&self, selected_account: &ProviderAccountId) -> AffinityTelemetry { AffinityTelemetry { - affinity_hit: self.bound_account.as_ref() == Some(selected_account) + affinity_hit: !self.inherited + && self.bound_account.as_ref() == Some(selected_account) && self.escape_reason.is_none(), escape_reason: self.escape_reason, account_switch: self @@ -333,8 +338,8 @@ impl CodexCredentialSelector { let input = CredentialSelectionInput { request_url: request.request_url, attempt: request.attempt, - session_affinity_key: None, - session_affinity_observation: None, + session_affinity_key: request.session_affinity.map(CodexSessionAffinity::key), + session_affinity_observation: request.session_affinity, }; self.select_inner(&input, None, ModelCatalogEligibility::NotApplicable) .await @@ -447,6 +452,9 @@ impl CodexCredentialSelector { } else { self.resolve_session_affinity( request.session_affinity_key, + request + .session_affinity_observation + .and_then(CodexSessionAffinity::root_key), &candidates, SystemTime::now(), ) @@ -476,7 +484,12 @@ impl CodexCredentialSelector { { affinity.escape(AffinityEscapeReason::PinnedAccount); } - let mut observed_affinity_account = affinity.bound_account().cloned(); + // 根绑定只提供默认偏好，不能作为子线程 CAS 的旧值，也不能跳过子线程首绑。 + let mut observed_affinity_account = if affinity.inherited { + None + } else { + affinity.bound_account().cloned() + }; let mut shortest_retry = None; let policy = request.attempt.account_selection_policy(); @@ -489,6 +502,7 @@ impl CodexCredentialSelector { now: SystemTime::now(), excluded_accounts: excluded.clone(), preferred_account: preferred.clone(), + preferred_account_overrides_weight: true, round_robin_cursor, eligibility: if diagnostic { AccountEligibilityPolicy::BypassForDiagnostic @@ -542,8 +556,8 @@ impl CodexCredentialSelector { } else { None }; - // 原生 continuation/required account 比根亲和绑定更严格；它可以 - // 使用 owner 账号，但不能把已经迁移的根会话拉回旧号。 + // 原生 continuation/required account 比亲和绑定更严格；它可以 + // 使用 owner 账号，但不能把已经迁移的会话拉回旧号。 if pinned_account.is_none() && let Some(effective_account) = initial_affinity_claim && &effective_account != account.id() @@ -576,6 +590,7 @@ impl CodexCredentialSelector { rotation_strategy = policy.strategy().as_str(), account_id = %account.id(), affinity_hit = affinity_telemetry.affinity_hit, + affinity_inherited = affinity.inherited, escape_reason = affinity_telemetry .escape_reason .map_or("", AffinityEscapeReason::as_str), @@ -610,6 +625,15 @@ impl CodexCredentialSelector { ) }) .collect(); + if !diagnostic + && observed_affinity_account.as_ref() == Some(account.id()) + && let Some(key) = request.session_affinity_key + { + // 命中即续期，避免长请求或客户端取消导致活跃会话提前过期。 + // CAS 防止并行请求把已经迁移的绑定改回旧账号。 + self.update_session_affinity(key, account.id(), account.id()) + .await; + } return Ok(CodexCredentialLease { installation_id: runtime.installation_id, account, @@ -630,17 +654,32 @@ impl CodexCredentialSelector { async fn resolve_session_affinity( &self, key: Option<&ProviderSessionAffinityKey>, + root_key: Option<&ProviderSessionAffinityKey>, candidates: &[AccountCandidate], now: SystemTime, ) -> AffinitySelection { let Some(key) = key else { return AffinitySelection::default(); }; - let SessionAffinityLookup::Bound(account_id) = self.lookup_session_affinity(key).await - else { - return AffinitySelection::default(); - }; - affinity_selection_for_bound_account(account_id, candidates, now) + match self.lookup_session_affinity(key).await { + SessionAffinityLookup::Bound(account_id) => { + affinity_selection_for_bound_account(account_id, candidates, now) + } + SessionAffinityLookup::Missing => { + if let Some(root_key) = root_key + && let SessionAffinityLookup::Bound(account_id) = + self.lookup_session_affinity(root_key).await + { + let mut inherited = + affinity_selection_for_bound_account(account_id, candidates, now); + inherited.inherited = true; + inherited + } else { + AffinitySelection::default() + } + } + SessionAffinityLookup::Unavailable => AffinitySelection::default(), + } } async fn lookup_session_affinity( @@ -947,36 +986,46 @@ impl CodexCredentialSelector { let Some(key) = session_affinity_key else { return; }; + self.update_session_affinity(key, expected_affinity_account_id, account.id()) + .await; + } + + pub(crate) async fn update_session_affinity( + &self, + key: &ProviderSessionAffinityKey, + expected_account_id: &ProviderAccountId, + selected_account_id: &ProviderAccountId, + ) { match tokio::time::timeout( SESSION_AFFINITY_TIMEOUT, self.session_affinity.compare_and_bind( &self.provider_kind, key, - expected_affinity_account_id, - account.id(), + expected_account_id, + selected_account_id, CODEX_ROOT_SESSION_TTL, ), ) .await { - Ok(Ok(effective_account)) if &effective_account == account.id() => {} + Ok(Ok(effective_account)) if &effective_account == selected_account_id => {} Ok(Ok(effective_account)) => { tracing::debug!( - account_id = %account.id(), + account_id = %selected_account_id, effective_account_id = %effective_account, - "OpenAI late success left the newer session affinity binding unchanged" + "OpenAI session affinity update left the newer binding unchanged" ); } Ok(Err(error)) => { tracing::warn!( - account_id = %account.id(), + account_id = %selected_account_id, error = %error, "OpenAI session affinity write failed open" ); } Err(_) => { tracing::warn!( - account_id = %account.id(), + account_id = %selected_account_id, timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(), "OpenAI session affinity write timed out" ); diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index ad44affe2..da1afa09c 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -28,6 +28,11 @@ impl CodexProvider { .get("image_turn_id") .and_then(Value::as_str) .map(str::to_owned); + let session_affinity = derive_codex_endpoint_session_affinity( + image.payload(), + context.client_api_key_ref(), + "session_id", + ); self.execute_raw_json_endpoint( context, RawJsonEndpointRequest { @@ -36,6 +41,7 @@ impl CodexProvider { body: image.payload().body().clone(), image_turn_id, turn_metadata: None, + session_affinity, }, ) .await @@ -59,6 +65,11 @@ impl CodexProvider { .get("turn_metadata") .and_then(Value::as_str) .map(str::to_owned); + let session_affinity = derive_codex_endpoint_session_affinity( + search.payload(), + context.client_api_key_ref(), + "id", + ); self.execute_raw_json_endpoint( context, RawJsonEndpointRequest { @@ -67,6 +78,7 @@ impl CodexProvider { body: search.payload().body().clone(), image_turn_id: None, turn_metadata, + session_affinity, }, ) .await @@ -83,6 +95,7 @@ impl CodexProvider { .select_for_provider_endpoint(&SelectCodexProviderEndpointCredential { request_url: &request.response_origin, attempt: &context, + session_affinity: request.session_affinity.as_ref(), }) .await .map_err(map_selection_error)?; @@ -120,6 +133,7 @@ impl CodexProvider { quota: Arc::clone(&self.quota), lease: Arc::clone(&lease), output_started_at: Instant::now(), + session_affinity_key: request.session_affinity.map(CodexSessionAffinity::into_key), }); let stream = ProviderStream::new(metadata, events, lease); Ok(if allows_account_state_mutation { @@ -139,6 +153,7 @@ struct RawJsonEndpointRequest { body: Bytes, image_turn_id: Option, turn_metadata: Option, + session_affinity: Option, } pub(super) struct ColdResponse { @@ -173,6 +188,7 @@ pub(super) struct ColdJsonResponse { pub(super) quota: Arc, pub(super) lease: Arc, pub(super) output_started_at: Instant, + pub(super) session_affinity_key: Option, } #[derive(Clone, Serialize, Deserialize)] @@ -424,6 +440,14 @@ pub(super) fn cold_json_response_stream(request: ColdJsonResponse) -> EventStrea } }; + if allows_account_state_mutation && let Some(key) = request.session_affinity_key.as_ref() { + // JSON 已完整接收；在首个 yield 前提交亲和迁移，避免下游取消漏掉更新。 + request.selector.update_session_affinity( + key, + request.lease.affinity_expected_account_id(), + active_account.id(), + ).await; + } let mut metrics = response.transport_metrics.clone(); metrics.first_event_ms = Some( i64::try_from(request.output_started_at.elapsed().as_millis()).unwrap_or(i64::MAX), diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 797fb4d60..c4a45de89 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -58,8 +58,8 @@ use crate::credential::{ CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, SelectCodexProviderEndpointCredential, - derive_codex_cyber_policy_session_key, derive_codex_session_affinity, - derive_previous_response_id_hash, + derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity, + derive_codex_session_affinity, derive_previous_response_id_hash, }; use crate::session_transport::CodexSessionTransportRecovery; use crate::transport::canonical::{ @@ -223,7 +223,26 @@ impl Provider for CodexProvider { client_api_key_id: &gateway_core::policy::ClientApiKeyId, ) -> ProviderRequestObservation { let Operation::Generate(request) = operation else { - return ProviderRequestObservation::default(); + let affinity = match operation { + Operation::Search(request) => derive_codex_endpoint_session_affinity( + request.payload(), + client_api_key_id, + "id", + ), + Operation::GenerateImage(request) => derive_codex_endpoint_session_affinity( + request.payload(), + client_api_key_id, + "session_id", + ), + _ => None, + }; + return ProviderRequestObservation { + continuation: ContinuationRequestObservation { + affinity_hash: affinity.map(|affinity| affinity.persistence_hash().to_owned()), + ..Default::default() + }, + ..Default::default() + }; }; let Ok(encoded) = encode_generate_request(request, "observability") else { return ProviderRequestObservation::default(); diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 55c57b66f..5af84ca78 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -120,8 +120,6 @@ fn extract_request_context(request: &mut CodexResponsesRequest) { request.codex_window_id = context.codex_window_id; request.parent_thread_id = context.parent_thread_id; request.client_conversation_id = context.conversation_id; - request.client_session_id = context.session_id; - request.client_thread_id = context.thread_id; request.client_request_id = context.client_request_id; request.client_turn_id = context.turn_id; request.responses_lite = context.responses_lite; @@ -137,8 +135,6 @@ struct ExtractedRequestContext { codex_window_id: Option, parent_thread_id: Option, conversation_id: Option, - session_id: Option, - thread_id: Option, client_request_id: Option, turn_id: Option, responses_lite: Option, @@ -164,8 +160,6 @@ impl ExtractedRequestContext { codex_window_id: body_string(body, "codexWindowId"), parent_thread_id: body_string(body, "parentThreadId"), conversation_id: body_string(body, "conversation_id"), - session_id: body_string(body, "session_id"), - thread_id: body_string(body, "thread_id"), client_request_id: body_string(body, "x-client-request-id"), turn_id: body_string(body, "turn_id").or_else(|| { client_metadata @@ -221,8 +215,7 @@ pub(crate) fn derive_conversation_anchor( } /// `thread_spawn` 与父任务共享根会话，但它本身是独立的子任务执行。 -/// 因此子线程必须拥有独立的 continuation/WebSocket 传输身份；账号亲和则从共享的 -/// 根会话单独派生。 +/// 因此子线程必须拥有独立的 continuation/WebSocket 传输身份；账号亲和另行派生。 fn thread_spawn_conversation_anchor( request: &CodexResponsesRequest, ) -> Option<(&'static str, String)> { @@ -529,10 +522,8 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map) { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_affinity_switch_a").await; + create_account(&store, "acct_affinity_switch_b").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let leases = Arc::new(TestLeaseCoordinator::default()); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/alpha/search")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"output":"result"}))) + .expect(2) + .mount(&server) + .await; + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::clone(&affinity), + server.uri(), + Arc::clone(&leases), + ); + let root = Operation::Generate(generate_with_session_context("shared-root", None, None)); + let root_stream = provider + .execute( + planned_request("openai", root.clone()), + context("req_cross_endpoint_root", CancellationToken::new()), + ) + .await + .expect("root selection"); + let root_account = root_stream.metadata().provider_account_id().clone(); + drop(root_stream); + let generation = Operation::Generate(generate_with_session_context( + "shared-root", + thread_id, + None, + )); + let first = provider + .execute( + planned_request("openai", generation.clone()), + context("req_cross_endpoint_first", CancellationToken::new()), + ) + .await + .expect("first selection"); + let first_account = first.metadata().provider_account_id().clone(); + assert_eq!( + first_account, root_account, + "new child inherits the root account" + ); + drop(first); + let other_account = if first_account.as_str() == "acct_affinity_switch_a" { + "acct_affinity_switch_b" + } else { + "acct_affinity_switch_a" + }; + store.set_scheduling( + other_account, + None, + AccountWeight::new(100).expect("weight"), + ); + + // 官方 Search 正文的 id 与 Responses 的 session-id 是同一个根身份。 + let search = Operation::Search(StandaloneSearchRequest::from_raw_json( + RawJsonPayload::new( + "openai", + Bytes::from_static(br#"{ "id":"shared-root", "commands":{}, "future":1, "future":2 }"#), + ) + .expect("search payload") + .with_context(thread_id.map_or_else(Map::new, |thread_id| { + Map::from_iter([( + "turn_metadata".to_owned(), + json!(json!({"thread_id":thread_id}).to_string()), + )]) + })), + )); + let client_key = ClientApiKeyId::new("key_openai_contract").expect("client key"); + assert_eq!( + provider + .request_observation(&search, &client_key) + .continuation + .affinity_hash, + provider + .request_observation(&generation, &client_key) + .continuation + .affinity_hash + ); + let mut same = provider + .execute( + planned_provider_endpoint_request("openai", search.clone()), + context("req_cross_endpoint_search", CancellationToken::new()), + ) + .await + .expect("search selection"); + assert_eq!( + same.metadata().provider_account_id(), + &first_account, + "session affinity outranks the other account's weight" + ); + while let Some(event) = same.next().await { + event.expect("search response"); + } + drop(same); + + // 仍由旧租约流程报告繁忙并换号，亲和不绕过并发限制。 + leases + .busy_accounts + .lock() + .expect("busy accounts") + .insert(first_account.clone()); + let mut fallback = provider + .execute( + planned_provider_endpoint_request("openai", search), + context("req_cross_endpoint_busy", CancellationToken::new()), + ) + .await + .expect("busy fallback"); + assert_eq!( + fallback.metadata().provider_account_id().as_str(), + other_account + ); + fallback + .next() + .await + .expect("first event") + .expect("successful JSON response"); + drop(fallback); // 只消费一个事件也必须完成绑定迁移。 + leases.busy_accounts.lock().expect("busy accounts").clear(); + store.set_scheduling( + first_account.as_str(), + None, + AccountWeight::new(100).expect("weight"), + ); + store.set_scheduling(other_account, None, AccountWeight::new(1).expect("weight")); + + let resumed = provider + .execute( + planned_request("openai", generation.clone()), + context("req_cross_endpoint_resumed", CancellationToken::new()), + ) + .await + .expect("resumed responses"); + assert_eq!( + resumed.metadata().provider_account_id().as_str(), + other_account, + "successful failover is shared back to Responses" + ); + drop(resumed); + for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] { + let image = Operation::GenerateImage(ImageRequest::from_raw_json( + kind, + RawJsonPayload::new("openai", Bytes::from_static(br#"{"prompt":"image"}"#)) + .expect("image payload") + .with_context({ + let mut context = + Map::from_iter([("session_id".to_owned(), json!("shared-root"))]); + if let Some(thread_id) = thread_id { + context.insert("thread_id".to_owned(), json!(thread_id)); + } + context + }), + )); + assert_eq!( + provider + .request_observation(&image, &client_key) + .continuation + .affinity_hash, + provider + .request_observation(&generation, &client_key) + .continuation + .affinity_hash + ); + let selected_image = provider + .execute( + planned_provider_endpoint_request("openai", image), + context("req_cross_endpoint_image", CancellationToken::new()), + ) + .await + .expect("image selection"); + assert_eq!( + selected_image.metadata().provider_account_id().as_str(), + other_account + ); + drop(selected_image); + } + if thread_id.is_some() { + let root_stream = provider + .execute( + planned_request("openai", root), + context( + "req_cross_endpoint_root_after_child", + CancellationToken::new(), + ), + ) + .await + .expect("root after child migration"); + assert_eq!(root_stream.metadata().provider_account_id(), &root_account); + drop(root_stream); + assert_eq!(affinity.binding_count(), 2, "only root and child bindings"); + } else { + let keys = affinity.lookup_keys(); + assert!( + keys.iter().all(|key| key == &keys[0]), + "all endpoints must use the same affinity key" + ); + } + assert_ne!( + provider + .request_observation( + &generation, + &ClientApiKeyId::new("different-client").expect("client key") + ) + .continuation + .affinity_hash, + provider + .request_observation(&generation, &client_key) + .continuation + .affinity_hash + ); + let unrelated = Operation::Generate(generate_with_session_context("another-root", None, None)); + assert_ne!( + provider + .request_observation(&unrelated, &client_key) + .continuation + .affinity_hash, + provider + .request_observation(&generation, &client_key) + .continuation + .affinity_hash + ); + server.verify().await; +} + #[tokio::test] async fn standalone_search_preserves_wire_and_scopes_turn_metadata_to_the_selected_account() { let store = Arc::new(MemoryAccountStore::default()); @@ -3423,14 +3662,18 @@ async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_st } #[tokio::test] -async fn thread_spawn_children_should_share_the_parent_account_affinity_key() { +async fn thread_spawn_children_should_inherit_the_root_with_independent_scoped_bindings() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_thread_spawn_affinity").await; let affinity = Arc::new(MemorySessionAffinity::default()); let provider = provider_with_affinity(&store, Arc::clone(&affinity)); let thread_spawn = r#"{"subagent_kind":"thread_spawn"}"#; + let client_key = ClientApiKeyId::new("key_openai_contract").expect("client key"); + let mut keys = Vec::new(); for (request_id, thread_id, turn_metadata) in [ + ("req_thread_spawn_parent", None, None), + ("req_thread_spawn_root_thread", Some("parent-session"), None), ( "req_thread_spawn_first", Some("child-one"), @@ -3447,31 +3690,334 @@ async fn thread_spawn_children_should_share_the_parent_account_affinity_key() { Some(thread_spawn), ), ("req_thread_spawn_fallback", None, Some(thread_spawn)), - ("req_thread_spawn_parent", None, None), ] { + let operation = Operation::Generate(generate_with_session_context( + "parent-session", + thread_id, + turn_metadata, + )); + keys.push( + provider + .request_observation(&operation, &client_key) + .continuation + .affinity_hash + .expect("affinity hash"), + ); let stream = provider .execute( - planned_request( - "openai", - Operation::Generate(generate_with_session_context( - "parent-session", - thread_id, - turn_metadata, - )), - ), + planned_request("openai", operation), context(request_id, CancellationToken::new()), ) .await .expect("prepare provider stream"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_thread_spawn_affinity" + ); drop(stream); } - let keys = affinity.lookup_keys(); - assert_eq!(keys.len(), 5); - assert!( - keys.iter().all(|key| key == &keys[0]), - "child thread transport identities must not split parent account affinity" + assert_eq!(keys[0], keys[1], "root thread uses the existing root key"); + assert_eq!( + keys[0], keys[5], + "unidentified thread retains root fallback" ); + assert_ne!(keys[0], keys[2]); + assert_ne!(keys[2], keys[3]); + assert_eq!(keys[2], keys[4], "same child reuses its own key"); + assert_eq!(affinity.binding_count(), 3); + for (root, client_key) in [ + ("other-root", client_key), + ( + "parent-session", + ClientApiKeyId::new("other-client").expect("client key"), + ), + ] { + let operation = Operation::Generate(generate_with_session_context( + root, + Some("child-one"), + Some(thread_spawn), + )); + assert_ne!( + provider + .request_observation(&operation, &client_key) + .continuation + .affinity_hash + .as_ref(), + Some(&keys[2]), + "child identity stays scoped to both root and client" + ); + } +} + +#[tokio::test] +async fn child_busy_failover_should_leave_the_running_parent_and_siblings_on_the_root_account() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let leases = Arc::new(TestLeaseCoordinator::default()); + let operation = |thread_id: &str| { + Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ( + "input".to_owned(), + json!([{"role":"user", "content":"keep the complete transcript"}]), + ), + ("session_id".to_owned(), json!("parent-session")), + ("thread_id".to_owned(), json!(thread_id)), + ("turnState".to_owned(), json!("old-account-turn-state")), + ]), + ) + .expect("payload") + .with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])), + )) + }; + let created = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_scope_capture\",\"status\":\"in_progress\"}}\n\n"; + let (parent_url, release_parent, parent_started, parent_server) = + paused_chunked_sse_server(created.to_owned(), CAPTURE_COMPLETED_SSE.to_owned()).await; + let parent_provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::clone(&affinity), + parent_url, + Arc::clone(&leases), + ); + let mut parent = parent_provider + .execute( + planned_request("openai", operation("parent-session")), + context("req_parent_running", CancellationToken::new()), + ) + .await + .expect("parent selection"); + let root_account = parent.metadata().provider_account_id().clone(); + let parent_task = tokio::spawn(async move { + let mut completed = false; + while let Some(event) = parent.next().await { + let event = event.expect("parent remains successful on its original account"); + completed |= event + .canonical_facts() + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + } + assert!(completed, "parent completion must be observed"); + }); + timeout(Duration::from_secs(5), parent_started) + .await + .expect("parent start timeout") + .expect("parent started"); + + create_account(&store, "acct_subagent_b").await; + store.set_scheduling( + "acct_subagent_b", + None, + AccountWeight::new(100).expect("weight"), + ); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + format!("{created}{CAPTURE_COMPLETED_SSE}"), + "text/event-stream", + )) + .expect(2) + .mount(&server) + .await; + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::clone(&affinity), + server.uri(), + Arc::clone(&leases), + ); + let child = provider + .execute( + planned_request("openai", operation("child-one")), + context("req_child_inherit", CancellationToken::new()), + ) + .await + .expect("initial child selection"); + assert_eq!( + child.metadata().provider_account_id(), + &root_account, + "root preference outranks B's weight" + ); + drop(child); + + // 租约层报告 A 已满；既有子线程和首次出现的子线程都能独立选择 B。 + leases + .busy_accounts + .lock() + .expect("busy accounts") + .insert(root_account.clone()); + for thread in ["child-one", "cold-child"] { + let mut child = provider + .execute( + planned_request("openai", operation(thread)), + context("req_child_busy", CancellationToken::new()), + ) + .await + .expect("child busy fallback"); + assert_eq!( + child.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); + let mut completed = false; + while let Some(event) = child.next().await { + let event = event.expect("child completes on B"); + completed |= event + .canonical_facts() + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + } + assert!(completed, "child completion must be observed"); + } + assert!( + !parent_task.is_finished(), + "child failover must not interrupt the active parent" + ); + release_parent + .send(()) + .expect("release parent after child success"); + timeout(Duration::from_secs(5), parent_task) + .await + .expect("parent completion timeout") + .expect("parent task"); + parent_server.await.expect("parent server"); + leases.busy_accounts.lock().expect("busy accounts").clear(); + + for (thread, expected) in [ + ("parent-session", "acct_subagent_a"), + ("child-one", "acct_subagent_b"), + ("cold-child", "acct_subagent_b"), + ("new-sibling", "acct_subagent_a"), + ] { + let stream = provider + .execute( + planned_request("openai", operation(thread)), + context("req_after_parent_completion", CancellationToken::new()), + ) + .await + .expect("selection after parent completion"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + expected, + "thread {thread}" + ); + drop(stream); + } + assert_eq!(affinity.binding_count(), 4); + let requests = server.received_requests().await.expect("child requests"); + assert_eq!(requests.len(), 2); + for request in requests { + assert!(captured_header_values(&request, "x-codex-turn-state").is_empty()); + assert_eq!( + captured_request_body(&request).get("input"), + Some(&json!([{"role":"user", "content":"keep the complete transcript"}])) + ); + assert_eq!( + request + .headers + .get("chatgpt-account-id") + .and_then(|header| header.to_str().ok()), + Some("chatgpt-acct_subagent_b") + ); + } + server.verify().await; +} + +#[tokio::test] +async fn failed_child_failover_should_preserve_both_existing_bindings() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let leases = Arc::new(TestLeaseCoordinator::default()); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/alpha/search")) + .respond_with( + ResponseTemplate::new(500).set_body_json(json!({"error":{"message":"upstream busy"}})), + ) + .expect(1) + .mount(&server) + .await; + let provider = provider_with_affinity_and_base_url_and_leases( + &store, + Arc::clone(&affinity), + server.uri(), + Arc::clone(&leases), + ); + for thread in [None, Some("child")] { + let stream = provider + .execute( + planned_request( + "openai", + Operation::Generate(generate_with_session_context("root", thread, None)), + ), + context("req_seed_child_failure", CancellationToken::new()), + ) + .await + .expect("initial binding"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + drop(stream); + } + create_account(&store, "acct_subagent_b").await; + leases + .busy_accounts + .lock() + .expect("busy accounts") + .insert(ProviderAccountId::new("acct_subagent_a").expect("account ID")); + let search = Operation::Search(StandaloneSearchRequest::from_raw_json( + RawJsonPayload::new( + "openai", + Bytes::from_static(br#"{"id":"root","thread_id":"child","commands":{}}"#), + ) + .expect("search payload"), + )); + let mut child = provider + .execute( + planned_provider_endpoint_request("openai", search), + context("req_failed_child", CancellationToken::new()), + ) + .await + .expect("child fallback selection"); + assert_eq!( + child.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); + let error = loop { + match child.next().await { + Some(Err(error)) => break error, + Some(Ok(_)) => {} + None => panic!("expected upstream failure"), + } + }; + assert_eq!(error.upstream_status(), Some(500)); + drop(child); + leases.busy_accounts.lock().expect("busy accounts").clear(); + for thread in [None, Some("child")] { + let stream = provider + .execute( + planned_request( + "openai", + Operation::Generate(generate_with_session_context("root", thread, None)), + ), + context("req_after_child_failure", CancellationToken::new()), + ) + .await + .expect("binding after failure"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a", + "failure must not migrate either binding" + ); + drop(stream); + } + assert_eq!(affinity.binding_count(), 2); + server.verify().await; } #[tokio::test] diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 8cdf676a8..9066d92fa 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -610,9 +610,14 @@ impl ProviderLeasePort for TestLeaseCoordinator { pub(crate) struct MemorySessionAffinity { bindings: Mutex>, lookups: Mutex>, + renewal_ttls: Mutex>, } impl MemorySessionAffinity { + pub(crate) fn renewal_ttls(&self) -> Vec { + self.renewal_ttls.lock().expect("affinity TTL lock").clone() + } + pub(crate) fn lookup_keys(&self) -> Vec { self.lookups.lock().expect("session affinity lock").clone() } @@ -701,7 +706,7 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity { key: &'a ProviderSessionAffinityKey, expected_account_id: &'a ProviderAccountId, replacement_account_id: &'a ProviderAccountId, - _ttl: Duration, + ttl: Duration, ) -> BoxFuture<'a, Result> { Box::pin(async move { let binding_key = ( @@ -714,6 +719,10 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity { .is_none_or(|current| current == expected_account_id) { bindings.insert(binding_key, replacement_account_id.clone()); + self.renewal_ttls + .lock() + .expect("affinity TTL lock") + .push(ttl); return Ok(replacement_account_id.clone()); } bindings.get(&binding_key).cloned().ok_or_else(|| { diff --git a/backend/crates/providers/openai/tests/transport/request.rs b/backend/crates/providers/openai/tests/transport/request.rs index 943d29b3d..df4880793 100644 --- a/backend/crates/providers/openai/tests/transport/request.rs +++ b/backend/crates/providers/openai/tests/transport/request.rs @@ -163,7 +163,7 @@ fn encoder_should_restore_conversation_fallback_after_body_and_context_values() } #[test] -fn encoder_should_ignore_legacy_context_aliases_and_metadata_fallbacks() { +fn encoder_should_accept_metadata_session_and_thread_ids_but_ignore_legacy_context_aliases() { let request = request(Map::from_iter([ ("model".to_owned(), json!("client-model")), ("input".to_owned(), json!("prompt")), @@ -222,7 +222,17 @@ fn encoder_should_ignore_legacy_context_aliases_and_metadata_fallbacks() { encoded.client_turn_id, ), ( - None, None, None, None, None, None, None, None, None, None, None, + None, + None, + None, + None, + None, + None, + None, + Some("metadata-session".to_owned()), + Some("metadata-thread".to_owned()), + None, + None, ) ); } diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index 0913faf63..bcf2d2806 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -194,6 +194,7 @@ impl GrokAccountSessionSelector { now: SystemTime::now(), excluded_accounts: request.excluded_accounts().clone(), preferred_account: request.required_account().cloned().or(affinity_account), + preferred_account_overrides_weight: false, round_robin_cursor: scheduling.round_robin_cursor(), eligibility: request.eligibility(), account_scope: (!diagnostic).then(|| Arc::clone(request.account_scope())),