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

fix(openai): prioritize session affinity and isolate child failover

Browse Source 
Prefer client-scoped session bindings over ordinary account weights while
keeping existing availability checks and failover behavior. Let child
threads inherit the root account on first selection and update only their
own binding when they migrate.

Share session and thread identity extraction across Responses, Search and
image endpoints. Renew existing Redis bindings for 24 hours and preserve
account-bound state isolation without adding a table or scheduler.

Validation: workspace tests (1838 passed, none ignored), rustfmt, and strict
workspace Clippy with all targets and features.

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