diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs index c20c44802..6154e4edb 100644 --- a/backend/crates/gateway-core/src/provider_ports/mod.rs +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -231,6 +231,32 @@ pub trait ProviderSessionAffinityPort: Send + Sync { ttl: Duration, ) -> BoxFuture<'a, Result<(), ProviderStoreError>>; + /// 仅在亲和键尚未绑定时写入候选账号，并返回原子操作后的实际绑定。 + /// + /// 已存在的绑定绝不会被候选账号覆盖；同一根会话的并发首次请求据此收敛到 + /// 单一账号。TTL 只在首次写入时设置，已有绑定由成功反馈负责刷新。 + fn claim_or_load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + candidate_account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result>; + + /// 仅当当前绑定等于 `expected_account_id`（或键已过期）时写入新账号， + /// 并返回原子操作后的实际绑定。 + /// + /// Provider 用它迁移不可用账号，以及在成功后以 `expected == replacement` + /// 刷新 TTL；迟到的旧账号成功不能覆盖较新的会话 winner。 + fn compare_and_bind<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_account_id: &'a ProviderAccountId, + replacement_account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> BoxFuture<'a, Result>; + fn clear<'a>( &'a self, provider_kind: &'a ProviderKind, diff --git a/backend/crates/gateway-store/src/redis/provider_session_affinity.rs b/backend/crates/gateway-store/src/redis/provider_session_affinity.rs index 1df736d5c..58139db82 100644 --- a/backend/crates/gateway-store/src/redis/provider_session_affinity.rs +++ b/backend/crates/gateway-store/src/redis/provider_session_affinity.rs @@ -16,6 +16,24 @@ use super::{namespace, resource_fingerprint}; const MAX_SESSION_AFFINITY_TTL: Duration = Duration::from_secs(24 * 60 * 60); +const CLAIM_OR_LOAD_SCRIPT: &str = r#" +local current = redis.call('GET', KEYS[1]) +if current then + return current +end +redis.call('PSETEX', KEYS[1], tonumber(ARGV[2]), ARGV[1]) +return ARGV[1] +"#; + +const COMPARE_AND_BIND_SCRIPT: &str = r#" +local current = redis.call('GET', KEYS[1]) +if not current or current == ARGV[1] then + redis.call('PSETEX', KEYS[1], tonumber(ARGV[3]), ARGV[2]) + return ARGV[2] +end +return current +"#; + #[derive(Clone)] pub struct RedisProviderSessionAffinityRepository { connection: ConnectionManager, @@ -79,11 +97,7 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository { ttl: Duration, ) -> futures::future::BoxFuture<'a, Result<(), ProviderStoreError>> { Box::pin(async move { - if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL { - return Err(provider_invalid("validate provider session affinity TTL")); - } - let ttl_millis = u64::try_from(ttl.as_millis()) - .map_err(|_| provider_invalid("validate provider session affinity TTL"))?; + let ttl_millis = session_affinity_ttl_millis(ttl)?; let mut connection = self.connection.clone(); redis::cmd("PSETEX") .arg(self.key(provider_kind, key)?) @@ -95,6 +109,50 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository { }) } + fn claim_or_load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + candidate_account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + let ttl_millis = session_affinity_ttl_millis(ttl)?; + let mut connection = self.connection.clone(); + let effective_account_id = redis::Script::new(CLAIM_OR_LOAD_SCRIPT) + .key(self.key(provider_kind, key)?) + .arg(candidate_account_id.as_str()) + .arg(ttl_millis) + .invoke_async::(&mut connection) + .await + .map_err(|_| provider_unavailable("claim provider session affinity"))?; + decode_account_id(effective_account_id) + }) + } + + fn compare_and_bind<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_account_id: &'a ProviderAccountId, + replacement_account_id: &'a ProviderAccountId, + ttl: Duration, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { + let ttl_millis = session_affinity_ttl_millis(ttl)?; + let mut connection = self.connection.clone(); + let effective_account_id = redis::Script::new(COMPARE_AND_BIND_SCRIPT) + .key(self.key(provider_kind, key)?) + .arg(expected_account_id.as_str()) + .arg(replacement_account_id.as_str()) + .arg(ttl_millis) + .invoke_async::(&mut connection) + .await + .map_err(|_| provider_unavailable("compare provider session affinity"))?; + decode_account_id(effective_account_id) + }) + } + fn clear<'a>( &'a self, provider_kind: &'a ProviderKind, @@ -112,6 +170,19 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository { } } +fn session_affinity_ttl_millis(ttl: Duration) -> Result { + if ttl.is_zero() || ttl > MAX_SESSION_AFFINITY_TTL { + return Err(provider_invalid("validate provider session affinity TTL")); + } + u64::try_from(ttl.as_millis()) + .map_err(|_| provider_invalid("validate provider session affinity TTL")) +} + +fn decode_account_id(account_id: String) -> Result { + ProviderAccountId::new(account_id) + .map_err(|_| provider_invalid("decode provider session affinity")) +} + fn provider_unavailable(operation: &'static str) -> ProviderStoreError { ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation) } diff --git a/backend/crates/gateway-store/tests/redis/provider_session_affinity.rs b/backend/crates/gateway-store/tests/redis/provider_session_affinity.rs index d5a239994..3f5d5d8bf 100644 --- a/backend/crates/gateway-store/tests/redis/provider_session_affinity.rs +++ b/backend/crates/gateway-store/tests/redis/provider_session_affinity.rs @@ -65,6 +65,55 @@ async fn session_affinity_should_overwrite_the_previous_account() { ); } +#[tokio::test] +async fn session_affinity_claim_should_keep_the_first_account() { + let Some((repository, _connection, _namespace)) = affinity_repository().await else { + return; + }; + let provider = ProviderKind::new("openai").expect("provider"); + let key = ProviderSessionAffinityKey::try_new("claim-session").expect("affinity key"); + let first = ProviderAccountId::new("acct_first").expect("first account"); + let second = ProviderAccountId::new("acct_second").expect("second account"); + + let first_winner = repository + .claim_or_load(&provider, &key, &first, Duration::from_secs(60)) + .await + .expect("claim first affinity"); + let second_winner = repository + .claim_or_load(&provider, &key, &second, Duration::from_secs(60)) + .await + .expect("load existing affinity"); + + assert_eq!((first_winner, second_winner), (first.clone(), first)); +} + +#[tokio::test] +async fn session_affinity_compare_should_not_replace_a_newer_winner() { + let Some((repository, _connection, _namespace)) = affinity_repository().await else { + return; + }; + let provider = ProviderKind::new("openai").expect("provider"); + let key = ProviderSessionAffinityKey::try_new("compare-session").expect("affinity key"); + let first = ProviderAccountId::new("acct_first").expect("first account"); + let second = ProviderAccountId::new("acct_second").expect("second account"); + let stale = ProviderAccountId::new("acct_stale").expect("stale account"); + repository + .bind(&provider, &key, &first, Duration::from_secs(60)) + .await + .expect("seed affinity"); + + let migrated = repository + .compare_and_bind(&provider, &key, &first, &second, Duration::from_secs(60)) + .await + .expect("migrate expected affinity"); + let stale_result = repository + .compare_and_bind(&provider, &key, &first, &stale, Duration::from_secs(60)) + .await + .expect("reject stale affinity migration"); + + assert_eq!((migrated, stale_result), (second.clone(), second)); +} + #[tokio::test] async fn session_affinity_should_apply_ttl_and_support_explicit_clear() { let Some((repository, mut connection, namespace)) = affinity_repository().await else { diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 5ec43dbb3..3c03727c2 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -1,4 +1,4 @@ -//! OpenAI 会话锚点到 Store 不透明亲和键的单向派生。 +//! OpenAI 根会话锚点到 Store 不透明账号亲和键的单向派生。 use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::ProviderSessionAffinityKey; @@ -11,34 +11,8 @@ pub(crate) fn derive_codex_session_affinity_key( request: &CodexResponsesRequest, client_api_key_id: &ClientApiKeyId, ) -> Option { - let session_key = if let Some(conversation_id) = request - .local_conversation_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - if let Ok(key) = ProviderSessionAffinityKey::try_new(conversation_id.to_owned()) { - Some(key) - } else { - opaque_affinity_key("local-conversation", conversation_id) - } - } else { - let (domain, value) = derive_conversation_anchor(request)?; - opaque_affinity_key(domain, &value) - }; - - let session_discriminator = request.subagent_kind(); - let session_key = - session_key.and_then(|session_key| match session_discriminator.as_deref() { - None => Some(session_key), - Some(discriminator) => { - // 仅扩展原会话锚点的键空间；后续查找、绑定和调度仍使用同一套亲和逻辑。 - opaque_affinity_key( - "subagent-session", - &format!("{discriminator}\0{}", session_key.expose_to_store()), - ) - } - })?; + let (domain, value) = derive_account_affinity_anchor(request)?; + let session_key = opaque_affinity_key(domain, &value)?; opaque_affinity_key( "client-session", &format!( @@ -49,6 +23,33 @@ pub(crate) fn derive_codex_session_affinity_key( ) } +/// 账号选择优先使用根会话事实；`local_conversation_id` 仍可按 child thread 隔离 +/// WebSocket 与 continuation 状态，但 child 身份不再拆分账号首选项。账号不可调度时 +/// selector 仍可按既有策略切换，并由 request scope 清理账号绑定状态。 +fn derive_account_affinity_anchor( + request: &CodexResponsesRequest, +) -> Option<(&'static str, String)> { + non_empty(request.client_session_id.as_deref()) + .map(|value| ("root-session", value.to_owned())) + .or_else(|| { + non_empty(request.client_conversation_id.as_deref()) + .map(|value| ("root-conversation", value.to_owned())) + }) + .or_else(|| { + request + .explicit_prompt_cache_key + .then(|| request.prompt_cache_key()) + .flatten() + .and_then(|value| non_empty(Some(value))) + .map(|value| ("root-prompt-cache", value.to_owned())) + }) + .or_else(|| { + non_empty(request.local_conversation_id.as_deref()) + .map(|value| ("local-conversation", value.to_owned())) + }) + .or_else(|| derive_conversation_anchor(request)) +} + fn opaque_affinity_key(domain: &str, value: &str) -> Option { let mut hasher = Sha256::new(); hasher.update(b"codex-session-affinity-v1\0"); diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index b7a8067c9..a381abcca 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -476,14 +476,16 @@ impl CodexCredentialSelector { { affinity.escape(AffinityEscapeReason::PinnedAccount); } - let preferred = pinned_account - .clone() - .or_else(|| affinity.preferred_account().cloned()); + let mut observed_affinity_account = affinity.bound_account().cloned(); let mut shortest_retry = None; + let policy = request.attempt.account_selection_policy(); loop { + let preferred = pinned_account + .clone() + .or_else(|| affinity.preferred_account().cloned()); let context = AccountSelectionContext { - policy: request.attempt.account_selection_policy(), + policy, now: SystemTime::now(), excluded_accounts: excluded.clone(), preferred_account: preferred.clone(), @@ -511,7 +513,6 @@ impl CodexCredentialSelector { .map(|candidate| candidate.account.clone()) .ok_or(CredentialSelectionError::InvalidCredential)?; let allows_account_state_mutation = !diagnostic || account.enabled(); - let policy = request.attempt.account_selection_policy(); match self .leases .try_acquire(ProviderLeaseRequest::Scheduling( @@ -532,6 +533,40 @@ impl CodexCredentialSelector { excluded.insert(account.id().clone()); } ProviderLeaseAcquisition::Acquired(guard) => { + let initial_affinity_claim = if !diagnostic + && observed_affinity_account.is_none() + && let Some(key) = request.session_affinity_key + { + self.claim_initial_session_affinity(key, account.id()).await + } else { + None + }; + // 原生 continuation/required account 比根亲和绑定更严格；它可以 + // 使用 owner 账号，但不能把已经迁移的根会话拉回旧号。 + if pinned_account.is_none() + && let Some(effective_account) = initial_affinity_claim + && &effective_account != account.id() + { + drop(guard); + observed_affinity_account = Some(effective_account.clone()); + affinity = affinity_selection_for_bound_account( + effective_account, + &candidates, + SystemTime::now(), + ); + continue; + } + let affinity_expected_account_id = if pinned_account + .as_ref() + .zip(observed_affinity_account.as_ref()) + .is_some_and(|(pinned, bound)| pinned != bound) + { + account.id().clone() + } else { + observed_affinity_account + .clone() + .unwrap_or_else(|| account.id().clone()) + }; let affinity_telemetry = affinity.telemetry(account.id()); tracing::info!( request_id = %request.attempt.request_id(), @@ -583,6 +618,7 @@ impl CodexCredentialSelector { cyber_policy_scope, allows_account_state_mutation, affinity_telemetry, + affinity_expected_account_id, _guard: guard, }); } @@ -603,24 +639,7 @@ impl CodexCredentialSelector { else { return AffinitySelection::default(); }; - let Some(candidate) = candidates - .iter() - .find(|candidate| candidate.account.id() == &account_id) - else { - self.clear_session_affinity(key).await; - return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable); - }; - if candidate - .account - .status_projection(now, candidate.signals.rate_limited_until) - .status - == AccountStatus::Normal - { - return AffinitySelection::preferred(account_id); - } - let reason = affinity_unavailable_reason(candidate, now); - self.clear_session_affinity(key).await; - AffinitySelection::escaped(account_id, reason) + affinity_selection_for_bound_account(account_id, candidates, now) } async fn lookup_session_affinity( @@ -649,22 +668,38 @@ impl CodexCredentialSelector { } } - async fn clear_session_affinity(&self, key: &ProviderSessionAffinityKey) { + async fn claim_initial_session_affinity( + &self, + key: &ProviderSessionAffinityKey, + selected_account_id: &ProviderAccountId, + ) -> Option { match tokio::time::timeout( SESSION_AFFINITY_TIMEOUT, - self.session_affinity.clear(&self.provider_kind, key), + self.session_affinity.claim_or_load( + &self.provider_kind, + key, + selected_account_id, + SESSION_AFFINITY_TTL, + ), ) .await { - Ok(Ok(_)) => {} + Ok(Ok(account_id)) => Some(account_id), Ok(Err(error)) => { - tracing::warn!(error = %error, "OpenAI stale session affinity clear failed open"); + tracing::warn!( + account_id = %selected_account_id, + error = %error, + "OpenAI initial session affinity claim failed open" + ); + None } Err(_) => { tracing::warn!( + account_id = %selected_account_id, timeout_ms = SESSION_AFFINITY_TIMEOUT.as_millis(), - "OpenAI stale session affinity clear timed out" + "OpenAI initial session affinity claim timed out" ); + None } } } @@ -901,6 +936,7 @@ impl CodexCredentialSelector { &self, account: &ProviderAccount, session_affinity_key: Option<&ProviderSessionAffinityKey>, + expected_affinity_account_id: &ProviderAccountId, ) { self.restore_recoverable_account_state(account).await; self.risk_recovery @@ -910,24 +946,26 @@ impl CodexCredentialSelector { let Some(key) = session_affinity_key else { return; }; - if matches!( - self.lookup_session_affinity(key).await, - SessionAffinityLookup::Unavailable - ) { - return; - } match tokio::time::timeout( SESSION_AFFINITY_TIMEOUT, - self.session_affinity.bind( + self.session_affinity.compare_and_bind( &self.provider_kind, key, + expected_affinity_account_id, account.id(), SESSION_AFFINITY_TTL, ), ) .await { - Ok(Ok(())) => {} + Ok(Ok(effective_account)) if &effective_account == account.id() => {} + Ok(Ok(effective_account)) => { + tracing::debug!( + account_id = %account.id(), + effective_account_id = %effective_account, + "OpenAI late success left the newer session affinity binding unchanged" + ); + } Ok(Err(error)) => { tracing::warn!( account_id = %account.id(), @@ -1104,6 +1142,32 @@ impl CodexCredentialSelector { } } +fn affinity_selection_for_bound_account( + account_id: ProviderAccountId, + candidates: &[AccountCandidate], + now: SystemTime, +) -> AffinitySelection { + let Some(candidate) = candidates + .iter() + .find(|candidate| candidate.account.id() == &account_id) + else { + return AffinitySelection::escaped(account_id, AffinityEscapeReason::HardUnavailable); + }; + match candidate + .account + .status_projection(now, candidate.signals.rate_limited_until) + .status + { + AccountStatus::Normal => AffinitySelection::preferred(account_id), + AccountStatus::QuotaExhausted + | AccountStatus::RateLimited + | AccountStatus::Disabled + | AccountStatus::Error => { + AffinitySelection::escaped(account_id, affinity_unavailable_reason(candidate, now)) + } + } +} + fn affinity_unavailable_reason( candidate: &AccountCandidate, now: SystemTime, @@ -1159,6 +1223,7 @@ pub struct CodexCredentialLease { cyber_policy_scope: Option, allows_account_state_mutation: bool, affinity_telemetry: AffinityTelemetry, + affinity_expected_account_id: ProviderAccountId, _guard: Box, } @@ -1216,6 +1281,12 @@ impl CodexCredentialLease { pub const fn account_switch(&self) -> bool { self.affinity_telemetry.account_switch } + + /// 成功反馈只能从选择时观察到的绑定迁移，避免迟到请求覆盖较新 winner。 + #[must_use] + pub(crate) const fn affinity_expected_account_id(&self) -> &ProviderAccountId { + &self.affinity_expected_account_id + } } impl fmt::Debug for CodexCredentialLease { diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index ba10e5cd2..aad0a62f0 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -71,6 +71,9 @@ use crate::transport::profile::{ use crate::transport::protocol::responses::{ CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals, }; +use crate::transport::protocol::websocket::{ + PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, +}; use crate::transport::request::{ CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account, }; @@ -318,8 +321,14 @@ impl Provider for CodexProvider { .is_some_and(|state| state.account_id != lease.account_id().as_str()); let has_explicit_state_owner = context.account_state_owner().is_some() || previous_session.is_some(); - let cross_account = - state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()); + let account_scope = + if state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch()) { + RequestAccountScope::Different + } else if has_explicit_state_owner || lease.affinity_hit() { + RequestAccountScope::Same + } else { + RequestAccountScope::Unknown + }; if context.continuation_attempt() != ContinuationAttempt::None && let Some(continuation) = context.continuation() { @@ -366,13 +375,11 @@ impl Provider for CodexProvider { } } } - let account_scope = if !cross_account { - RequestAccountScope::SameAccount - } else if upstream_request.previous_response_id().is_some() { - RequestAccountScope::CrossAccountContinuationProbe - } else { - RequestAccountScope::CrossAccountFullRequest - }; + if upstream_request.previous_response_id().is_some() + && !account_scope.can_reuse_account_state() + { + return Err(continuation_replay_required_error()); + } scope_request_to_account( &mut upstream_request, lease.installation_id(), @@ -1178,7 +1185,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { // 完成事件一旦交给下游，Core 可以立刻停止轮询 Provider stream； // 在此之前持久化亲和关系，保证成功请求不会因流被提前 drop 而丢失绑定。 selector - .record_success(&active_account, session_affinity_key.as_ref()) + .record_success( + &active_account, + session_affinity_key.as_ref(), + lease.affinity_expected_account_id(), + ) .await; selector .observe_cyber_policy_success(cyber_policy_scope.as_ref()) @@ -1282,7 +1293,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { if allows_account_state_mutation && completed && terminal_failure.is_none() { // 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。 selector - .record_success(&active_account, session_affinity_key.as_ref()) + .record_success( + &active_account, + session_affinity_key.as_ref(), + lease.affinity_expected_account_id(), + ) .await; selector .observe_cyber_policy_success(cyber_policy_scope.as_ref()) @@ -2364,6 +2379,19 @@ fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { map_client_error(error, UpstreamSendState::Ambiguous, true) } +fn continuation_replay_required_error() -> ProviderError { + ProviderError::new( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + ) + .with_continuation_failure(ContinuationFailure::HistoryUnavailable) + .with_client_visible_upstream_error(ClientVisibleUpstreamError::new( + PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, + Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()), + Some("invalid_request_error".to_owned()), + )) +} + fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error); let mut failure = map_client_error(error, UpstreamSendState::Sent, false); diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index e991fd430..bb8149c85 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -525,11 +525,6 @@ impl CodexResponsesRequest { .map_or(&[], Vec::as_slice) } - /// 替换输入条目（仅用于代理持有完整历史时的换号重放）。 - pub fn set_input(&mut self, input: Vec) { - self.body.insert("input".to_string(), Value::Array(input)); - } - /// 是否流式返回。 pub fn stream(&self) -> bool { self.body @@ -617,11 +612,12 @@ impl CodexResponsesRequest { ) } - /// 返回会话亲和所用的子代理区分值。 + /// 返回请求语义与 child transport 隔离所用的子代理区分值。 /// /// Codex 原生请求在 turn metadata 中声明 `subagent_kind`；网关的 /// `/responses/review` 等入口则通过 `client_metadata.x-openai-subagent` - /// 传递给上游。该可选值仅参与会话亲和键派生。 + /// 传递给上游。它不拆分根会话的账号首选项；`thread_spawn` 仍用它派生独立 + /// WebSocket/continuation transport identity。 pub fn subagent_kind(&self) -> Option { self.semantics() .subagent_kind diff --git a/backend/crates/providers/openai/src/transport/protocol/websocket.rs b/backend/crates/providers/openai/src/transport/protocol/websocket.rs index 9682ee3f8..fe7b4b121 100644 --- a/backend/crates/providers/openai/src/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/src/transport/protocol/websocket.rs @@ -5,7 +5,9 @@ use serde_json::{Map, Value}; use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement}; const REDACTED_PAYLOAD_VALUE: &str = ""; -const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found"; +pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found"; +pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str = + "Previous response was not found. Retrying the full request."; const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`."; /// WebSocket 握手审计快照。 diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index b69b61aae..ccd884269 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -63,18 +63,14 @@ const INSTALLATION_ID_KEYS: &[&str] = &[ #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum RequestAccountScope { - SameAccount, - CrossAccountFullRequest, - CrossAccountContinuationProbe, + Same, + Different, + Unknown, } impl RequestAccountScope { - const fn is_cross_account(self) -> bool { - !matches!(self, Self::SameAccount) - } - - const fn preserves_continuation(self) -> bool { - matches!(self, Self::CrossAccountContinuationProbe) + pub(crate) const fn can_reuse_account_state(self) -> bool { + matches!(self, Self::Same) } } @@ -205,8 +201,9 @@ pub(crate) fn derive_conversation_anchor( }) } -/// A `thread_spawn` shares its parent session but is a separate child execution. -/// Its child thread must therefore own the affinity and WebSocket pool identity. +/// `thread_spawn` 与父任务共享根会话，但它本身是独立的子任务执行。 +/// 因此子线程必须拥有独立的 continuation/WebSocket 传输身份；账号亲和则从共享的 +/// 根会话单独派生。 fn thread_spawn_conversation_anchor( request: &CodexResponsesRequest, ) -> Option<(&'static str, String)> { @@ -306,21 +303,17 @@ fn normalize_conversation_anchor_text(text: &str) -> String { /// /// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入 /// 请求头；正文只替换客户端原本声明过的 installation 字段，绝不接受客户端 -/// 提供的 token、cookie 或账号身份。Continuation probe 保留 canonical -/// previous-response 和原始增量 input，不能把它们误当成完整请求清理。 +/// 提供的 token、cookie 或账号身份。`input` 是 Responses 的可回放会话正文， +/// item ID、encrypted content 和 compaction 都必须原样保留。 pub(crate) fn scope_request_to_account( request: &mut CodexResponsesRequest, installation_id: &str, account_scope: RequestAccountScope, ) { - let cross_account = account_scope.is_cross_account(); - let continuation = account_scope - .preserves_continuation() - .then(|| request.previous_response_id().map(ToOwned::to_owned)) - .flatten(); + let reset_account_state = !account_scope.can_reuse_account_state(); let client_metadata_turn_state = metadata_string(request, "x-codex-turn-state"); - let preserve_turn_state = - !cross_account && (request.turn_state.is_some() || client_metadata_turn_state.is_some()); + let preserve_turn_state = !reset_account_state + && (request.turn_state.is_some() || client_metadata_turn_state.is_some()); let turn_state = preserve_turn_state .then(|| request.turn_state.clone()) .flatten(); @@ -329,38 +322,22 @@ pub(crate) fn scope_request_to_account( } else { None }; - let turn_metadata = (!account_scope.preserves_continuation()) - .then(|| { - request - .turn_metadata - .as_deref() - .and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account)) - }) - .flatten(); - let client_metadata_turn_metadata = (!account_scope.preserves_continuation()) - .then(|| { - metadata_string(request, "x-codex-turn-metadata") - .and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account)) - }) - .flatten(); + let turn_metadata = request + .turn_metadata + .as_deref() + .and_then(|metadata| scope_turn_metadata(metadata, installation_id, reset_account_state)); + let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata") + .and_then(|metadata| scope_turn_metadata(&metadata, installation_id, reset_account_state)); - if cross_account { + if reset_account_state { request.passthrough_headers.remove("x-codex-turn-state"); request.passthrough_headers.remove("x-codex-turn-metadata"); - if !account_scope.preserves_continuation() { - sanitize_cross_account_input(request); - } for key in CROSS_ACCOUNT_IDENTITY_KEYS .iter() .chain(ACCOUNT_BOUND_STATE_KEYS) { request.body_mut().remove(*key); } - if account_scope.preserves_continuation() { - for key in TURN_METADATA_KEYS { - request.body_mut().remove(*key); - } - } } for key in INSTALLATION_ID_KEYS { @@ -371,7 +348,7 @@ pub(crate) fn scope_request_to_account( .body() .get(*key) .and_then(Value::as_str) - .and_then(|value| scope_turn_metadata(value, installation_id, cross_account)); + .and_then(|value| scope_turn_metadata(value, installation_id, reset_account_state)); replace_existing_body_string(request, key, scoped.as_deref()); } @@ -383,22 +360,17 @@ pub(crate) fn scope_request_to_account( ( key, metadata.get(key).and_then(Value::as_str).and_then(|value| { - scope_turn_metadata(value, installation_id, cross_account) + scope_turn_metadata(value, installation_id, reset_account_state) }), ) }); - if cross_account { + if reset_account_state { for key in CROSS_ACCOUNT_IDENTITY_KEYS .iter() .chain(ACCOUNT_BOUND_STATE_KEYS) { metadata.remove(*key); } - if account_scope.preserves_continuation() { - for key in TURN_METADATA_KEYS { - metadata.remove(*key); - } - } } metadata.insert( "x-codex-installation-id".to_owned(), @@ -436,36 +408,6 @@ pub(crate) fn scope_request_to_account( request.turn_state = turn_state; request.turn_metadata = turn_metadata; - if let Some(continuation) = continuation { - request.set_previous_response_id(Some(continuation)); - } -} - -fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) { - if request.input().is_empty() { - return; - } - let input = request - .input() - .iter() - .cloned() - .filter_map(sanitize_cross_account_item) - .collect(); - request.set_input(input); -} - -pub(crate) fn sanitize_cross_account_item(mut item: Value) -> Option { - if let Value::Object(object) = &mut item { - if matches!( - object.get("type").and_then(Value::as_str), - Some("compaction" | "compaction_summary" | "context_compaction") - ) { - return None; - } - object.remove("id"); - object.remove("encrypted_content"); - } - Some(item) } fn metadata_string(request: &CodexResponsesRequest, key: &str) -> Option { @@ -672,3 +614,29 @@ fn context_string(context: &Map, field: &str) -> Option { .and_then(Value::as_str) .map(ToOwned::to_owned) } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn thread_spawn_children_should_keep_distinct_transport_anchors() { + let request = |thread_id: &str| { + let mut request = CodexResponsesRequest::from_body(Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), json!("child task")), + ])); + request.client_session_id = Some("root-session".to_owned()); + request.client_thread_id = Some(thread_id.to_owned()); + request.turn_metadata = Some(r#"{"subagent_kind":"thread_spawn"}"#.to_owned()); + request + }; + + assert_ne!( + derive_conversation_anchor(&request("child-one")), + derive_conversation_anchor(&request("child-two")) + ); + } +} diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index 1a4ce0cf4..461a19dc2 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -445,6 +445,74 @@ fn selector_round_robin_cursor_advances_across_requests() { ); } +#[tokio::test] +async fn selector_should_claim_the_initial_session_account_before_upstream_send() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_first", "at-first"); + create_account(&store, "acct_second", "at-second"); + let affinity = Arc::new(MemorySessionAffinity::default()); + let selector = selector_with_affinity( + &store, + Arc::new(TestLeaseCoordinator::default()), + Arc::clone(&affinity), + ); + let provider = ProviderKind::new("openai").expect("provider"); + let key = ProviderSessionAffinityKey::try_new("initial-claim").expect("affinity key"); + let request_attempt = attempt(BTreeSet::new()); + + let selected = selector + .select(&SelectCodexCredential { + upstream_model: "gpt-5.4", + request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses") + .expect("request URL"), + attempt: &request_attempt, + session_affinity_key: Some(&key), + }) + .await + .expect("select initial account"); + + assert_eq!( + affinity + .load(&provider, &key) + .await + .expect("load claimed affinity"), + Some(selected.account_id().clone()) + ); +} + +#[tokio::test] +async fn record_success_should_not_overwrite_a_newer_session_winner() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_first", "at-first"); + create_account(&store, "acct_second", "at-second"); + let affinity = Arc::new(MemorySessionAffinity::default()); + let selector = selector_with_affinity( + &store, + Arc::new(TestLeaseCoordinator::default()), + Arc::clone(&affinity), + ); + let provider = ProviderKind::new("openai").expect("provider"); + let key = ProviderSessionAffinityKey::try_new("late-success").expect("affinity key"); + let first = store.account("acct_first").expect("first account"); + let second = ProviderAccountId::new("acct_second").expect("second account"); + affinity + .bind(&provider, &key, &second, Duration::from_secs(60)) + .await + .expect("seed newer affinity"); + + selector + .record_success(&first, Some(&key), first.id()) + .await; + + assert_eq!( + affinity + .load(&provider, &key) + .await + .expect("load preserved affinity"), + Some(second) + ); +} + #[tokio::test] async fn selector_should_reuse_the_account_bound_to_the_same_session() { let store = Arc::new(MemoryAccountStore::default()); @@ -469,7 +537,9 @@ async fn selector_should_reuse_the_account_bound_to_the_same_session() { }) .await .expect("select first account"); - selector.record_success(first.account(), Some(&key)).await; + selector + .record_success(first.account(), Some(&key), first.account_id()) + .await; let first_account = first.account_id().clone(); let second_attempt = attempt(BTreeSet::new()); @@ -535,7 +605,7 @@ async fn selector_should_replace_a_busy_affinity_binding_after_the_fallback_succ .await .expect("select fallback account"); selector - .record_success(selected.account(), Some(&key)) + .record_success(selected.account(), Some(&key), &bound) .await; assert_eq!( @@ -1005,7 +1075,7 @@ fn successful_upstream_response_recovers_non_quota_terminal_states() { let current = store.account("acct_primary").expect("stale account"); let selector = selector(&store, Arc::new(TestLeaseCoordinator::default())); - block_on(selector.record_success(¤t, None)); + block_on(selector.record_success(¤t, None, current.id())); assert_eq!( store @@ -1289,7 +1359,7 @@ fn cloudflare_challenge_does_not_change_persisted_account_facts() { .credential_state(), CredentialState::Ready ); - block_on(selector.record_success(lease.account(), None)); + block_on(selector.record_success(lease.account(), None, lease.account_id())); } #[test] diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 86386d59b..78e5d6aff 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -344,6 +344,8 @@ fn contract_account_scope() -> Arc { "acct_local_affinity", "acct_metadata_new", "acct_metadata_old", + "acct_unknown_continuation", + "acct_unknown_turn_state", "acct_prefetch_limit", "acct_presentation", "acct_provider_contract", @@ -352,7 +354,8 @@ fn contract_account_scope() -> Arc { "acct_scope_same", "acct_semantic_failure", "acct_session_affinity", - "acct_subagent", + "acct_subagent_a", + "acct_subagent_b", "acct_success_exhausted", "acct_thread_spawn_affinity", "acct_truncated_stream", @@ -456,6 +459,24 @@ fn replay_any_context(request_id: &str, owner_account_id: &str) -> AttemptContex .with_continuation_attempt(ContinuationAttempt::ReplayAny) } +fn external_continuation_context(request_id: &str) -> AttemptContext { + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new(request_id).expect("request id"), + ClientApiKeyId::new("key_openai_contract").expect("client key id"), + ), + NonZeroU32::new(1).expect("attempt"), + SystemTime::now() + Duration::from_secs(30), + account_policy(), + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(contract_account_scope()), + Some(ContinuationBinding::External(PreviousResponseId::new( + "external-previous-response", + ))), + CancellationToken::new(), + ) +} + async fn capture_scoped_http_request( request_id: &str, selected_account_id: &str, @@ -1319,6 +1340,53 @@ async fn cross_account_scope_removes_only_account_bound_body_fields() { ); } +#[tokio::test] +async fn cross_account_full_replay_should_preserve_the_complete_transcript() { + let transcript = json!([ + { + "type": "reasoning", + "id": "reasoning-item-id", + "encrypted_content": "reasoning-ciphertext", + "summary": [] + }, + { + "type": "compaction", + "id": "compaction-item-id", + "encrypted_content": "compaction-ciphertext" + }, + { + "type": "agent_message", + "id": "agent-message-id", + "content": [ + {"type": "output_text", "text": "visible"}, + {"type": "encrypted_content", "encrypted_content": "nested-ciphertext"} + ] + }, + { + "type": "function_call", + "id": "function-call-id", + "call_id": "call-id", + "name": "tool", + "arguments": "{}", + "encrypted_function_args": "function-args-ciphertext" + } + ]); + let request = capture_scoped_http_request( + "req_scope_transcript", + "acct_scope_new", + "acct_scope_old", + Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), transcript.clone()), + ]), + Map::new(), + ) + .await; + let body: Value = captured_request_body(&request); + + assert_eq!(body.get("input"), Some(&transcript)); +} + #[tokio::test] async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() { let request = capture_scoped_http_request( @@ -1359,40 +1427,19 @@ async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() { } #[tokio::test] -async fn cross_account_continuation_probe_should_preserve_delta_input_and_previous_response() { +async fn cross_account_continuation_should_require_client_replay_without_an_upstream_probe() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_scope_new").await; let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CAPTURE_COMPLETED_SSE), - ) - .expect(1) - .mount(&server) - .await; - let delta_input = json!([ - { - "type": "message", - "id": "delta-message-id", - "encrypted_content": "opaque-delta-ciphertext", - "role": "user", - "content": "current turn only" - }, - { - "type": "compaction", - "id": "delta-compaction-id", - "encrypted_content": "opaque-compaction-ciphertext" - } - ]); let generation = GenerateRequest::from_protocol_payload( ProtocolPayload::json_object( "openai", Map::from_iter([ ("model".to_owned(), json!("gpt-5.4")), - ("input".to_owned(), delta_input.clone()), + ( + "input".to_owned(), + json!([{"role": "user", "content": "delta"}]), + ), ( "previous_response_id".to_owned(), json!("client-previous-response"), @@ -1432,41 +1479,137 @@ async fn cross_account_continuation_probe_should_preserve_delta_input_and_previo ) .expect("provider session state"), ); - let mut stream = provider_with_base_url(&store, server.uri()) + let result = provider_with_base_url(&store, server.uri()) .execute( planned_request("openai", Operation::Generate(generation)), replay_any_context("req_cross_account_probe", "acct_scope_old"), ) + .await; + let Err(error) = result else { + panic!("cross-account continuation must require a full client replay"); + }; + let detail = error + .client_visible_upstream_error() + .expect("client replay error detail"); + let requests = server + .received_requests() .await - .expect("prepare continuation probe stream"); + .expect("captured upstream requests"); + + assert_eq!( + ( + error.kind(), + error.send_state(), + error.continuation_failure(), + detail.code(), + detail.error_type(), + ), + ( + ProviderErrorKind::InvalidRequest, + UpstreamSendState::NotSent, + Some(ContinuationFailure::HistoryUnavailable), + Some("previous_response_not_found"), + Some("invalid_request_error"), + ) + ); + assert!(requests.is_empty()); +} + +#[tokio::test] +async fn missing_affinity_continuation_should_require_client_replay_without_an_upstream_probe() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_unknown_continuation").await; + let server = MockServer::start().await; + let generation = 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": "delta"}]), + ), + ( + "previous_response_id".to_owned(), + json!("external-previous-response"), + ), + ]), + ) + .expect("OpenAI payload") + .with_context(Map::from_iter([ + ("use_websocket".to_owned(), json!(false)), + ("session_id".to_owned(), json!("expired-affinity-session")), + ])), + ); + let result = provider_with_base_url(&store, server.uri()) + .execute( + planned_request("openai", Operation::Generate(generation)), + external_continuation_context("req_unknown_continuation"), + ) + .await; + let Err(error) = result else { + panic!("unowned continuation must require a full client replay"); + }; + let requests = server + .received_requests() + .await + .expect("captured upstream requests"); + + assert_eq!( + error + .client_visible_upstream_error() + .and_then(|detail| detail.code()), + Some("previous_response_not_found") + ); + assert!(requests.is_empty()); +} + +#[tokio::test] +async fn missing_affinity_full_request_should_clear_unowned_turn_state() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_unknown_turn_state").await; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CAPTURE_COMPLETED_SSE), + ) + .expect(1) + .mount(&server) + .await; + let generation = GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), json!("full transcript")), + ]), + ) + .expect("OpenAI payload") + .with_context(Map::from_iter([ + ("use_websocket".to_owned(), json!(false)), + ("session_id".to_owned(), json!("expired-turn-state-session")), + ("turn_id".to_owned(), json!("turn-after-affinity-expiry")), + ("turn_state".to_owned(), json!("unowned-turn-state")), + ])), + ); + let mut stream = provider_with_base_url(&store, server.uri()) + .execute( + planned_request("openai", Operation::Generate(generation)), + context("req_unknown_turn_state", CancellationToken::new()), + ) + .await + .expect("prepare full replay stream"); while let Some(event) = stream.next().await { - event.expect("continuation probe response"); + event.expect("full replay response"); } let requests = server .received_requests() .await - .expect("captured continuation probe request"); - let body: Value = captured_request_body(&requests[0]); + .expect("captured full replay request"); - assert_eq!(body.get("input"), Some(&delta_input)); - assert_eq!( - body.get("previous_response_id"), - Some(&json!("upstream-previous-response")) - ); - assert!(body.get("authorization").is_none()); - assert!(body.get("turnMetadata").is_none()); - assert!( - body.pointer("/client_metadata/x-codex-turn-state") - .is_none() - ); - assert!( - body.pointer("/client_metadata/x-codex-turn-metadata") - .is_none() - ); - assert_eq!( - body.pointer("/client_metadata/future"), - Some(&json!("keep")) - ); assert!(captured_header_values(&requests[0], "x-codex-turn-state").is_empty()); } @@ -1637,9 +1780,10 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key() } #[tokio::test] -async fn subagent_requests_should_use_an_affinity_key_isolated_from_the_root_session() { +async fn subagent_requests_should_share_the_root_session_account_affinity_key() { let store = Arc::new(MemoryAccountStore::default()); - create_account(&store, "acct_subagent").await; + create_account(&store, "acct_subagent_a").await; + create_account(&store, "acct_subagent_b").await; let affinity = Arc::new(MemorySessionAffinity::default()); let server = MockServer::start().await; Mock::given(method("POST")) @@ -1689,10 +1833,25 @@ async fn subagent_requests_should_use_an_affinity_key_isolated_from_the_root_ses let keys = affinity.lookup_keys(); assert_eq!(keys.len(), 3); - assert_ne!(keys[0], keys[1], "subagent must not inherit the root key"); - assert_eq!( - keys[1], keys[2], - "the same subagent should retain its own key" + assert!( + keys.iter().all(|key| key == &keys[0]), + "root and derived subagent requests must prefer the same account" + ); + let requests = server + .received_requests() + .await + .expect("captured root and subagent requests"); + let selected_accounts = requests + .iter() + .filter_map(|request| request.headers.get("chatgpt-account-id")) + .filter_map(|value| value.to_str().ok()) + .collect::>(); + assert_eq!(selected_accounts.len(), 3); + assert!( + selected_accounts + .iter() + .all(|account| account == &selected_accounts[0]), + "root and subagents should route to the same preferred account" ); } @@ -1867,7 +2026,7 @@ async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_st } #[tokio::test] -async fn thread_spawn_children_should_have_independent_session_affinity_keys() { +async fn thread_spawn_children_should_share_the_parent_account_affinity_key() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_thread_spawn_affinity").await; let affinity = Arc::new(MemorySessionAffinity::default()); @@ -1912,10 +2071,10 @@ async fn thread_spawn_children_should_have_independent_session_affinity_keys() { let keys = affinity.lookup_keys(); assert_eq!(keys.len(), 5); - assert_ne!(keys[0], keys[1]); - assert_eq!(keys[0], keys[2]); - assert_ne!(keys[3], keys[4]); - assert_ne!(keys[0], keys[3]); + assert!( + keys.iter().all(|key| key == &keys[0]), + "child thread transport identities must not split parent account affinity" + ); } #[tokio::test] @@ -2539,7 +2698,7 @@ async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() { } #[tokio::test] -async fn previous_response_not_found_should_remain_client_visible_without_internal_replay_proof() { +async fn same_account_previous_response_not_found_should_remain_client_visible() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_client_history").await; let server = MockServer::start().await; @@ -2578,7 +2737,7 @@ async fn previous_response_not_found_should_remain_client_visible_without_intern "openai", Operation::Generate(GenerateRequest::from_protocol_payload(payload)), ), - context("req_history_missing", CancellationToken::new()), + context_with_state_owner("req_history_missing", "acct_client_history"), ) .await .expect("prepare missing-history stream"); diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index e8ed6abe9..38522def5 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -678,6 +678,55 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity { }) } + fn claim_or_load<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + candidate_account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut bindings = self.bindings.lock().expect("session affinity lock"); + Ok(bindings + .entry(( + provider_kind.as_str().to_owned(), + key.expose_to_store().to_owned(), + )) + .or_insert_with(|| candidate_account_id.clone()) + .clone()) + }) + } + + fn compare_and_bind<'a>( + &'a self, + provider_kind: &'a ProviderKind, + key: &'a ProviderSessionAffinityKey, + expected_account_id: &'a ProviderAccountId, + replacement_account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let binding_key = ( + provider_kind.as_str().to_owned(), + key.expose_to_store().to_owned(), + ); + let mut bindings = self.bindings.lock().expect("session affinity lock"); + if bindings + .get(&binding_key) + .is_none_or(|current| current == expected_account_id) + { + bindings.insert(binding_key, replacement_account_id.clone()); + return Ok(replacement_account_id.clone()); + } + bindings.get(&binding_key).cloned().ok_or_else(|| { + ProviderStoreError::new( + gateway_core::provider_ports::ProviderStoreErrorKind::Unavailable, + "resolve in-memory provider session affinity", + ) + }) + }) + } + fn clear<'a>( &'a self, provider_kind: &'a ProviderKind, diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index 2b8ef8665..d1e297298 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -713,6 +713,27 @@ impl ProviderSessionAffinityPort for TestSessionAffinity { Box::pin(async { Ok(()) }) } + fn claim_or_load<'a>( + &'a self, + _provider_kind: &'a ProviderKind, + _key: &'a ProviderSessionAffinityKey, + candidate_account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { Ok(candidate_account_id.clone()) }) + } + + fn compare_and_bind<'a>( + &'a self, + _provider_kind: &'a ProviderKind, + _key: &'a ProviderSessionAffinityKey, + _expected_account_id: &'a ProviderAccountId, + replacement_account_id: &'a ProviderAccountId, + _ttl: Duration, + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { Ok(replacement_account_id.clone()) }) + } + fn clear<'a>( &'a self, _provider_kind: &'a ProviderKind,