starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): preserve root session account affinity

Browse Source 
Share the preferred account across derived child executions while keeping their continuation and WebSocket transport identities isolated.

Claim and migrate affinity atomically, clear account-bound state on fallback, reject cross-account previous-response continuations, and preserve complete replay transcripts.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-25 14:22:36 +08:00
co-authored by Codex
parent 75f9d2a290
commit 4f50c67332
13 changed files with 749 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/provider_ports/mod.rs
+26
View File
@@ -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<ProviderAccountId, ProviderStoreError>>;
/// 仅当当前绑定等于 `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<ProviderAccountId, ProviderStoreError>>;
fn clear<'a>(
&'a self,
provider_kind: &'a ProviderKind,
backend/crates/gateway-store/src/redis/provider_session_affinity.rs
+76 -5
View File
@@ -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<ProviderAccountId, ProviderStoreError>> {
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::<String>(&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<ProviderAccountId, ProviderStoreError>> {
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::<String>(&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<u64, ProviderStoreError> {
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, ProviderStoreError> {
ProviderAccountId::new(account_id)
.map_err(|_| provider_invalid("decode provider session affinity"))
}
fn provider_unavailable(operation: &'static str) -> ProviderStoreError {
ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation)
}
backend/crates/gateway-store/tests/redis/provider_session_affinity.rs
+49
View File
@@ -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 {
backend/crates/providers/openai/src/credential/affinity.rs
+30 -29
View File
@@ -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<ProviderSessionAffinityKey> {
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<ProviderSessionAffinityKey> {
let mut hasher = Sha256::new();
hasher.update(b"codex-session-affinity-v1\0");
backend/crates/providers/openai/src/credential/selector.rs
+107 -36
View File
@@ -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<ProviderAccountId> {
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<CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
affinity_telemetry: AffinityTelemetry,
affinity_expected_account_id: ProviderAccountId,
_guard: Box<dyn ProviderLeaseGuard>,
}
@@ -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 {
backend/crates/providers/openai/src/provider.rs
+39 -11
View File
@@ -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);
backend/crates/providers/openai/src/transport/protocol/responses.rs
+3 -7
View File
@@ -525,11 +525,6 @@ impl CodexResponsesRequest {
.map_or(&[], Vec::as_slice)
}
/// 替换输入条目（仅用于代理持有完整历史时的换号重放）。
pub fn set_input(&mut self, input: Vec<Value>) {
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<String> {
self.semantics()
.subagent_kind
backend/crates/providers/openai/src/transport/protocol/websocket.rs
+3 -1
View File
@@ -5,7 +5,9 @@ use serde_json::{Map, Value};
use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement};
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
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 握手审计快照。
backend/crates/providers/openai/src/transport/request.rs
+49 -81
View File
@@ -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<Value> {
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<String> {
@@ -672,3 +614,29 @@ fn context_string(context: &Map<String, Value>, field: &str) -> Option<String> {
.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"))
);
}
}
backend/crates/providers/openai/tests/credential/contract.rs
+74 -4
View File
@@ -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(&current, None));
block_on(selector.record_success(&current, 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]
backend/crates/providers/openai/tests/provider/contract.rs
+223 -64
View File
@@ -344,6 +344,8 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"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<FrozenAccountScope> {
"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::<Vec<_>>();
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");
backend/crates/providers/openai/tests/support.rs
+49
View File
@@ -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<ProviderAccountId, ProviderStoreError>> {
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<ProviderAccountId, ProviderStoreError>> {
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,
backend/crates/providers/xai/tests/support.rs
+21
View File
@@ -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<ProviderAccountId, ProviderStoreError>> {
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<ProviderAccountId, ProviderStoreError>> {
Box::pin(async move { Ok(replacement_account_id.clone()) })
}
fn clear<'a>(
&'a self,
_provider_kind: &'a ProviderKind,