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

feat(settings): add preferred account affinity mode

Browse Source 
Route relaxed requests through the configured scheduler and keep the session primary during preferred-mode spillover. Default new settings to strict while preserving saved modes.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-07 03:50:59 +08:00
co-authored by Codex
parent 3d1dc79f5d
commit 119032c067
21 changed files with 722 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-api/tests/admin/settings.rs
+1 -1
View File
@@ -1392,7 +1392,7 @@ async fn guardian_reservation_round_trips_and_rejects_invalid_values() {
#[tokio::test]
async fn account_affinity_and_rotation_budget_round_trip_and_reject_invalid_values() {
for (mode, rotations, ttl) in [("relaxed", 0, 1), ("strict", 31, 720)] {
for (mode, rotations, ttl) in [("relaxed", 0, 1), ("preferred", 3, 24), ("strict", 31, 720)] {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
backend/crates/gateway-core/src/account/selection.rs
+6 -3
View File
@@ -52,12 +52,13 @@ pub const MAX_ACCOUNT_ROTATIONS: u32 = 31;
/// 会话绑定与轮次关联共用的保留时长上限，单位小时
pub const MAX_SESSION_AFFINITY_TTL_HOURS: u32 = 720;
/// OpenAI 会话账号亲和；具体线程身份由 Provider 解释
/// OpenAI 会话账号亲和；请求身份和模式行为由 Provider 解释
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AccountAffinity {
#[default]
Relaxed,
Preferred,
#[default]
Strict,
}
@@ -66,6 +67,7 @@ impl AccountAffinity {
pub const fn as_str(self) -> &'static str {
match self {
Self::Relaxed => "relaxed",
Self::Preferred => "preferred",
Self::Strict => "strict",
}
}
@@ -74,6 +76,7 @@ impl AccountAffinity {
pub fn parse(value: &str) -> Option<Self> {
match value {
"relaxed" => Some(Self::Relaxed),
"preferred" => Some(Self::Preferred),
"strict" => Some(Self::Strict),
_ => None,
}
@@ -107,7 +110,7 @@ impl AccountSelectionPolicy {
max_concurrent_per_account: max_concurrent_per_account.into(),
request_interval,
openai_guardian_reserved_concurrency: 0,
openai_account_affinity: AccountAffinity::Relaxed,
openai_account_affinity: AccountAffinity::default(),
max_account_rotations: 3,
openai_session_affinity_ttl: Duration::from_secs(24 * 3600),
queue_policy: ConcurrencyQueuePolicy {
backend/crates/gateway-core/src/settings/values.rs
+1 -1
View File
@@ -138,7 +138,7 @@ impl SettingsValues {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
openai_guardian_reserved_concurrency: 0,
openai_account_affinity: crate::account::AccountAffinity::Relaxed,
openai_account_affinity: crate::account::AccountAffinity::default(),
max_account_rotations: 3,
openai_session_affinity_ttl_hours: 24,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
backend/crates/gateway-core/tests/settings/mod.rs
+8 -4
View File
@@ -200,21 +200,25 @@ fn affinity_and_rotation_settings_are_validated_without_mutating_the_snapshot()
use gateway_core::account::AccountAffinity;
let original = snapshot(1, "host");
let facts = serde_json::to_value(original.settings()).unwrap();
assert_eq!(facts["openai_account_affinity"], "relaxed");
assert_eq!(facts["openai_account_affinity"], "strict");
assert_eq!(facts["max_account_rotations"], 3);
for rotations in [0, 31] {
for (mode, rotations) in [
(AccountAffinity::Relaxed, 0),
(AccountAffinity::Preferred, 3),
(AccountAffinity::Strict, 31),
] {
let changed = original
.with_settings(
&original
.settings()
.clone()
.with_openai_account_affinity(AccountAffinity::Strict)
.with_openai_account_affinity(mode)
.with_max_account_rotations(rotations),
)
.unwrap();
let changed = serde_json::to_value(changed.settings()).unwrap();
assert_eq!(changed["max_account_rotations"], rotations);
assert_eq!(changed["openai_account_affinity"], "strict");
assert_eq!(changed["openai_account_affinity"], mode.as_str());
}
assert!(
original
backend/crates/gateway-store/tests/postgres/runtime_settings.rs
+42 -1
View File
@@ -1188,7 +1188,7 @@ async fn account_affinity_upgrade_defaults_and_updates_reach_the_snapshot() {
let Some(database) = TestDatabase::create_through("account_affinity_upgrade", 22).await else {
return;
};
sqlx::query("update runtime_settings set rotation_strategy = 'sticky' where id = 1")
sqlx::query("update runtime_settings set rotation_strategy = 'sticky', config_revision = config_revision + 1 where id = 1")
.execute(&database.pool)
.await
.unwrap();
@@ -1203,6 +1203,7 @@ async fn account_affinity_upgrade_defaults_and_updates_reach_the_snapshot() {
assert_eq!(before.values.max_account_rotations, 3);
assert_eq!(before.values.openai_session_affinity_ttl_hours, 24);
for (mode, budget, ttl) in [
(AccountAffinity::Preferred, 3, 24),
(AccountAffinity::Strict, 31, 168),
(AccountAffinity::Relaxed, 0, 720),
] {
@@ -1261,3 +1262,43 @@ async fn account_affinity_upgrade_defaults_and_updates_reach_the_snapshot() {
);
database.close().await;
}
#[tokio::test]
async fn preferred_affinity_migration_defaults_to_strict_and_preserves_saved_modes() {
use gateway_core::account::AccountAffinity;
let Some(fresh) = TestDatabase::create("affinity_default").await else {
return;
};
let repository = PgRuntimeSettingsRepository::new(fresh.pool.clone());
assert_eq!(
repository
.load_runtime_settings()
.await
.unwrap()
.values
.openai_account_affinity,
AccountAffinity::Strict
);
fresh.close().await;
for mode in [AccountAffinity::Relaxed, AccountAffinity::Strict] {
let database = TestDatabase::create_through("affinity_preserve", 23)
.await
.unwrap();
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
let mut update = settings_with_margin(3600);
update.values.openai_account_affinity = mode;
repository.update_runtime_settings(update).await.unwrap();
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
assert_eq!(
repository
.load_runtime_settings()
.await
.unwrap()
.values
.openai_account_affinity,
mode
);
database.close().await;
}
}
backend/crates/providers/openai/src/credential/affinity.rs
+18 -68
View File
@@ -1,6 +1,5 @@
//! 逻辑会话账号绑定与线程传输隔离键的单向派生
use gateway_core::account::AccountAffinity;
use gateway_core::operation::RawJsonPayload;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAlias};
@@ -15,7 +14,6 @@ const AFFINITY_KEY_HASH_LENGTH: usize = 12;
/// 一次请求派生出的账号亲和键及其结构化日志上下文
pub(crate) struct CodexSessionAffinity {
key: ProviderSessionAffinityKey,
root_key: Option<ProviderSessionAffinityKey>,
turn_alias: Option<(ProviderSessionAffinityKey, ProviderSessionAlias)>,
key_hash: String,
anchor_source: &'static str,
@@ -28,26 +26,17 @@ impl CodexSessionAffinity {
pub(crate) fn from_turn_alias(
turn: ProviderSessionAffinityKey,
alias: ProviderSessionAlias,
mode: AccountAffinity,
) -> Self {
let strict_root = alias
.root_session_key
.as_ref()
.filter(|_| mode == AccountAffinity::Strict);
// 模式切换对新请求生效，旧轮次的子线程身份在严格模式下也归入根绑定
let key = strict_root.unwrap_or(&alias.session_key).clone();
// 旧轮次可能登记在线程绑定下，统一还原为会话身份并保留原记录用于续期
let root = alias.root_session_key.as_ref();
let key = root.unwrap_or(&alias.session_key).clone();
Self {
key_hash: short_key_hash(&key),
key,
root_key: if strict_root.is_some() {
None
} else {
alias.root_session_key.clone()
},
anchor_source: "turn-session",
anchor: String::new(),
session_id: None,
follow_only: strict_root.is_some() || alias.follow_only,
follow_only: root.is_some() || alias.follow_only,
turn_alias: Some((turn, alias)),
}
}
@@ -64,38 +53,10 @@ impl CodexSessionAffinity {
.unwrap_or_else(|| ProviderSessionAlias {
session_key: self.key.clone(),
follow_only: self.follow_only,
root_session_key: self.root_key.clone(),
root_session_key: None,
})
}
pub(crate) const fn root_key(&self) -> Option<&ProviderSessionAffinityKey> {
self.root_key.as_ref()
}
fn with_policy(
mut self,
mode: AccountAffinity,
thread: Option<&str>,
follow_only: bool,
) -> Self {
self.follow_only = follow_only;
if mode == AccountAffinity::Relaxed
&& let Some(thread) = non_empty(thread)
&& self.session_id.as_deref() != Some(thread)
&& let Some(key) = opaque_affinity_key(
"child-thread",
&format!("{}\0{thread}", self.key.expose_to_store()),
)
{
self.root_key = Some(self.key);
self.key = key;
self.key_hash = short_key_hash(&self.key);
self.follow_only = false;
}
// 无法识别子线程时保留跟随约束，避免子请求改写根绑定
self
}
pub(crate) const fn follow_only(&self) -> bool {
self.follow_only
}
@@ -164,7 +125,6 @@ pub(crate) fn derive_previous_response_id_hash(
pub(crate) fn derive_codex_session_affinity(
request: &CodexResponsesRequest,
client_api_key_id: &ClientApiKeyId,
mode: AccountAffinity,
) -> Option<CodexSessionAffinity> {
let session_id = non_empty(request.client_account_session_id.as_deref())?;
session_affinity(
@@ -173,13 +133,7 @@ pub(crate) fn derive_codex_session_affinity(
Some(session_id.to_owned()),
client_api_key_id,
)
.map(|affinity| {
affinity.with_policy(
mode,
request.client_account_thread_id.as_deref(),
request.client_account_follow_only,
)
})
.map(|affinity| affinity.with_follow_only(request.client_account_follow_only))
}
/// 原生 JSON 端点与 Responses 共用绑定命名空间，Search id 是逻辑会话身份
@@ -187,9 +141,8 @@ pub(crate) fn derive_codex_endpoint_session_affinity(
payload: &RawJsonPayload,
client_api_key_id: &ClientApiKeyId,
body_session_field: &str,
mode: AccountAffinity,
) -> Option<CodexSessionAffinity> {
derive_endpoint_affinity_with_headers(payload, client_api_key_id, body_session_field, &[], mode)
derive_endpoint_affinity_with_headers(payload, client_api_key_id, body_session_field, &[])
}
pub(crate) fn derive_endpoint_affinity_with_headers(
@@ -197,7 +150,6 @@ pub(crate) fn derive_endpoint_affinity_with_headers(
client_api_key_id: &ClientApiKeyId,
body_session_field: &str,
headers: &[gateway_core::engine::middleware::MiddlewareHeader],
mode: AccountAffinity,
) -> Option<CodexSessionAffinity> {
let mut body = serde_json::from_slice::<Map<String, Value>>(payload.body()).unwrap_or_default();
if let Some(session) = body.get(body_session_field).cloned() {
@@ -211,11 +163,11 @@ pub(crate) fn derive_endpoint_affinity_with_headers(
client_api_key_id,
)
.map(|affinity| {
affinity.with_policy(
mode,
account_thread_with_headers(&body, payload.context(), headers).as_deref(),
follows_session_with_headers(&body, payload.context(), headers),
)
affinity.with_follow_only(follows_session_with_headers(
&body,
payload.context(),
headers,
))
})
}
@@ -243,7 +195,7 @@ pub(crate) fn follows_session_with_headers(
.is_some()
}
pub(crate) fn account_thread_with_headers(
fn account_thread_with_headers(
body: &Map<String, Value>,
context: &Map<String, Value>,
headers: &[gateway_core::engine::middleware::MiddlewareHeader],
@@ -297,7 +249,6 @@ pub(crate) fn derive_live_session_affinity(
request: &gateway_core::operation::ProviderHttpRequest,
headers: &[gateway_core::engine::middleware::MiddlewareHeader],
client_api_key_id: &ClientApiKeyId,
mode: AccountAffinity,
) -> Option<CodexSessionAffinity> {
let mut effective = headers.to_vec();
effective.extend(
@@ -322,11 +273,11 @@ pub(crate) fn derive_live_session_affinity(
client_api_key_id,
)
.map(|affinity| {
affinity.with_policy(
mode,
account_thread_with_headers(&Map::new(), &Map::new(), &effective).as_deref(),
follows_session_with_headers(&Map::new(), &Map::new(), &effective),
)
affinity.with_follow_only(follows_session_with_headers(
&Map::new(),
&Map::new(),
&effective,
))
})
}
@@ -359,7 +310,6 @@ fn session_affinity(
let key_hash = short_key_hash(&key);
Some(CodexSessionAffinity {
key,
root_key: None,
turn_alias: None,
key_hash,
anchor_source,
backend/crates/providers/openai/src/credential/mod.rs
+5 -6
View File
@@ -17,12 +17,11 @@ pub mod token_client;
mod types;
pub(crate) use affinity::{
CodexSessionAffinity, account_session_with_headers, account_thread_with_headers,
derive_codex_cyber_policy_session_key, derive_codex_endpoint_session_affinity,
derive_codex_session_affinity, derive_codex_transport_key,
derive_endpoint_affinity_with_headers, derive_live_session_affinity,
derive_previous_response_id_hash, derive_turn_alias, follows_session_with_headers,
turn_id_with_headers,
CodexSessionAffinity, account_session_with_headers, derive_codex_cyber_policy_session_key,
derive_codex_endpoint_session_affinity, derive_codex_session_affinity,
derive_codex_transport_key, derive_endpoint_affinity_with_headers,
derive_live_session_affinity, derive_previous_response_id_hash, derive_turn_alias,
follows_session_with_headers, turn_id_with_headers,
};
pub(crate) use oauth::oauth_owner_ref;
pub(crate) use types::parse_access_token_expiration;
backend/crates/providers/openai/src/credential/selector.rs
+70 -48
View File
@@ -8,10 +8,10 @@ use std::time::{Duration, SystemTime};
use gateway_core::error::ErrorSource;
use gateway_core::account::{
AccountCandidate, AccountCapacitySnapshot, AccountEligibilityPolicy, AccountErrorReason,
AccountFeedbackStats, AccountRuntimeSignals, AccountSchedulingBlocker, AccountSelectionContext,
AccountSelector, AccountStatus, CredentialState, PreferredAccountSelection, ProviderAccount,
ProviderAccountId, QuotaEvidence,
AccountAffinity, AccountCandidate, AccountCapacitySnapshot, AccountEligibilityPolicy,
AccountErrorReason, AccountFeedbackStats, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus,
CredentialState, PreferredAccountSelection, ProviderAccount, ProviderAccountId, QuotaEvidence,
};
use gateway_core::concurrency::{CapacityWait, ConcurrencyWaitQueue, QueueRejection, WaitPriority};
use gateway_core::engine::{AttemptContext, ContinuationAttempt, policy::AccountPolicyError};
@@ -342,14 +342,16 @@ impl CodexCredentialSelector {
lease: &mut CodexCredentialLease,
session_affinity: Option<&CodexSessionAffinity>,
cyber_policy_session_key: Option<&ProviderSessionAffinityKey>,
session_ttl: Duration,
policy: AccountSelectionPolicy,
) -> Result<(), CredentialSelectionError> {
let selected_account = lease.account.id().clone();
if let Some(affinity) = session_affinity
&& lease.admitted_session.as_ref() != Some(affinity.key())
{
let current = self.lookup_session_affinity(affinity.key()).await?;
if (affinity.follow_only() && current.is_none())
if (policy.openai_account_affinity() == AccountAffinity::Strict
&& affinity.follow_only()
&& current.is_none())
|| current
.as_ref()
.is_some_and(|binding| binding.account_id() != &selected_account)
@@ -358,7 +360,7 @@ impl CodexCredentialSelector {
affinity.key(),
current.as_ref(),
&selected_account,
session_ttl,
policy.openai_session_affinity_ttl(),
)
.await?
{
@@ -413,6 +415,11 @@ impl CodexCredentialSelector {
upstream_model: Option<&str>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let follow_only = !request.attempt.is_diagnostic_required_account()
&& request
.attempt
.account_selection_policy()
.openai_account_affinity()
== AccountAffinity::Strict
&& request
.session_affinity_observation
.is_some_and(CodexSessionAffinity::follow_only);
@@ -447,6 +454,10 @@ impl CodexCredentialSelector {
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let diagnostic = request.attempt.is_diagnostic_required_account();
let binding_key = request.session_affinity_key.filter(|_| !diagnostic);
let affinity_mode = request
.attempt
.account_selection_policy()
.openai_account_affinity();
let mut queue_policy = request.attempt.account_selection_policy().queue_policy();
if follow_only {
// 关闭普通账号排队不能放开后代线程换号，沿用队列时间设置并保留有界人数
@@ -508,10 +519,12 @@ impl CodexCredentialSelector {
} else {
None
};
if binding
.as_ref()
.zip(continuation_account.as_ref())
.is_some_and(|(binding, owner)| binding.account_id() != owner)
// 非严格模式的续写可以归属分流账号；严格绑定迁移后才拒绝旧 owner
if affinity_mode == AccountAffinity::Strict
&& binding
.as_ref()
.zip(continuation_account.as_ref())
.is_some_and(|(binding, owner)| binding.account_id() != owner)
{
// 旧账号的增量状态不能把已迁移的会话拉回去，交回客户端完整历史重放
return Err(CredentialSelectionError::ContinuationOwnerChanged);
@@ -640,28 +653,16 @@ impl CodexCredentialSelector {
AccountCandidate { account, signals }
})
.collect::<Vec<_>>();
let inherited = if binding.is_none() && !diagnostic {
if let Some(root) = request
.session_affinity_observation
.and_then(CodexSessionAffinity::root_key)
{
self.lookup_session_affinity(root).await?
} else {
None
}
} else {
None
};
let mut affinity = binding.as_ref().or(inherited.as_ref()).map_or_else(
AffinitySelection::default,
|binding| {
let mut affinity = binding
.as_ref()
.filter(|_| affinity_mode != AccountAffinity::Relaxed)
.map_or_else(AffinitySelection::default, |binding| {
affinity_selection_for_bound_account(
binding.account_id().clone(),
&candidates,
SystemTime::now(),
)
},
);
});
let cyber_policy_scope = self
.prepare_cyber_policy_scope(cyber_policy_session_key)
.await;
@@ -672,7 +673,7 @@ impl CodexCredentialSelector {
{
excluded.extend(state.excluded_accounts().iter().cloned());
}
// 继续遵守既有重试排除与固定账号规则；只迁移本请求所属绑定
// 原生续写归属和显式固定账号优先于会话偏好，不能把增量状态发送给其他账号
if let Some(required) = pinned_account.as_ref() {
excluded.extend(
candidates
@@ -717,7 +718,10 @@ impl CodexCredentialSelector {
let mut wait_candidates =
AccountSelector.wait_candidates(&candidates, &wait_context);
let preferred_wait = preferred.as_ref().filter(|account| {
!diagnostic && queue_policy.max_waiting > 0 && wait_candidates.contains(account)
!diagnostic
&& (affinity_mode == AccountAffinity::Strict || pinned_account.is_some())
&& queue_policy.max_waiting > 0
&& wait_candidates.contains(account)
});
if let Some(preferred) = preferred_wait {
wait_candidates.retain(|account| account == preferred);
@@ -745,22 +749,31 @@ impl CodexCredentialSelector {
}
};
// 只约束内置调度；插件的显式选号保留原裁决和完整候选
let selection = if follow_only
let selection = if (follow_only
|| (affinity_mode == AccountAffinity::Preferred
&& affinity.preferred_account().is_some()))
&& selection.is_none_or(|selection| !selection.is_policy_choice())
{
// 优先模式先尝试会话主账号，繁忙或不可用时才采用调度结果；高权重回切不抢占主账号
let preferred = if follow_only {
binding.as_ref().map(ProviderSessionBinding::account_id)
} else {
preferred.as_ref()
};
let mut follow_context = context.clone();
follow_context.now = SystemTime::now();
follow_context.excluded_accounts.extend(
candidates
.iter()
.filter(|candidate| {
binding.as_ref().is_none_or(|binding| {
candidate.account.id() != binding.account_id()
})
})
.filter(|candidate| preferred != Some(candidate.account.id()))
.map(|candidate| candidate.account.id().clone()),
);
AccountSelector.select(&candidates, &follow_context)
let preferred_selection = AccountSelector.select(&candidates, &follow_context);
if follow_only {
preferred_selection
} else {
preferred_selection.or(selection)
}
} else {
selection
};
@@ -904,7 +917,11 @@ impl CodexCredentialSelector {
.admit_session(
key,
binding.as_ref(),
account.id(),
// 分流只改变本次租约，CAS 续期仍保留会话主账号
binding
.as_ref()
.filter(|_| affinity_mode == AccountAffinity::Preferred)
.map_or(account.id(), ProviderSessionBinding::account_id),
request
.attempt
.account_selection_policy()
@@ -987,14 +1004,22 @@ impl CodexCredentialSelector {
session: &CodexSessionAffinity,
session_ttl: Duration,
) -> Result<(), CredentialSelectionError> {
let mut alias = session.alias_record();
if session.turn_alias().is_none()
&& let Some(existing) = self.session_for_turn(turn).await?
{
// 同一轮次可能保存在线程键下；复验会话与权限后续期原记录，不能因格式不同改写归属
if existing.key() != session.key() || existing.follow_only() != session.follow_only() {
return Err(CredentialSelectionError::SessionBound(Box::new(
CredentialSelectionError::NoEligibleCredential,
)));
}
alias = existing.alias_record();
}
let applied = tokio::time::timeout(
SESSION_AFFINITY_TIMEOUT,
self.session_affinity.bind_alias(
&self.provider_kind,
turn,
&session.alias_record(),
session_ttl,
),
self.session_affinity
.bind_alias(&self.provider_kind, turn, &alias, session_ttl),
)
.await
.map_err(binding_store_error)?
@@ -1010,7 +1035,6 @@ impl CodexCredentialSelector {
pub(crate) async fn session_for_turn(
&self,
turn: &ProviderSessionAffinityKey,
mode: gateway_core::account::AccountAffinity,
) -> Result<Option<CodexSessionAffinity>, CredentialSelectionError> {
tokio::time::timeout(
SESSION_AFFINITY_TIMEOUT,
@@ -1019,9 +1043,7 @@ impl CodexCredentialSelector {
.await
.map_err(binding_store_error)?
.map_err(binding_store_error)
.map(|alias| {
alias.map(|alias| CodexSessionAffinity::from_turn_alias(turn.clone(), alias, mode))
})
.map(|alias| alias.map(|alias| CodexSessionAffinity::from_turn_alias(turn.clone(), alias)))
}
async fn lookup_session_affinity(
backend/crates/providers/openai/src/provider/execution.rs
+3 -13
View File
@@ -63,7 +63,6 @@ impl CodexProvider {
context.client_api_key_ref(),
"session_id",
headers,
context.account_selection_policy().openai_account_affinity(),
);
let turn = headers
.iter()
@@ -79,10 +78,7 @@ impl CodexProvider {
turn.and_then(|turn| derive_turn_alias(turn, context.client_api_key_ref()))
{
self.selector
.session_for_turn(
&alias,
context.account_selection_policy().openai_account_affinity(),
)
.session_for_turn(&alias)
.await
.map_err(map_selection_error)?
} else {
@@ -91,9 +87,7 @@ impl CodexProvider {
if explicit
.as_ref()
.zip(inferred.as_ref())
.is_some_and(|(explicit, inferred)| {
explicit.key() != inferred.key() && Some(explicit.key()) != inferred.root_key()
})
.is_some_and(|(explicit, inferred)| explicit.key() != inferred.key())
{
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
@@ -134,7 +128,6 @@ impl CodexProvider {
search.payload(),
context.client_api_key_ref(),
"id",
context.account_selection_policy().openai_account_affinity(),
);
let response_origin = self.search_url.clone();
self.execute_raw_json_endpoint(
@@ -221,7 +214,6 @@ impl CodexProvider {
context.client_api_key_ref(),
"id",
&middleware_headers,
context.account_selection_policy().openai_account_affinity(),
),
_ => None,
};
@@ -231,9 +223,7 @@ impl CodexProvider {
&mut lease,
affinity.as_ref(),
None,
context
.account_selection_policy()
.openai_session_affinity_ttl(),
context.account_selection_policy(),
)
.await
.map_err(map_selection_error)?;
backend/crates/providers/openai/src/provider/live.rs
+3 -10
View File
@@ -243,12 +243,8 @@ impl CodexProvider {
UpstreamSendState::NotSent,
)
})?;
let session_affinity = derive_live_session_affinity(
&request,
&[],
context.client_api_key_ref(),
context.account_selection_policy().openai_account_affinity(),
);
let session_affinity =
derive_live_session_affinity(&request, &[], context.client_api_key_ref());
let selection_started_at = Instant::now();
let lease = self
.selector
@@ -321,7 +317,6 @@ impl CodexProvider {
&request,
&middleware_headers,
context.client_api_key_ref(),
context.account_selection_policy().openai_account_affinity(),
);
if !context.is_diagnostic_required_account() {
self.selector
@@ -329,9 +324,7 @@ impl CodexProvider {
&mut lease,
session_affinity.as_ref(),
None,
context
.account_selection_policy()
.openai_session_affinity_ttl(),
context.account_selection_policy(),
)
.await
.map_err(map_selection_error)?;
backend/crates/providers/openai/src/provider/mod.rs
+5 -21
View File
@@ -213,11 +213,8 @@ impl CodexProvider {
clear_request_turn_state(&mut upstream);
}
}
let session_affinity = derive_codex_session_affinity(
&upstream,
context.client_api_key_ref(),
context.account_selection_policy().openai_account_affinity(),
);
let session_affinity =
derive_codex_session_affinity(&upstream, context.client_api_key_ref());
let cyber_policy_session_key =
derive_codex_cyber_policy_session_key(&upstream, context.client_api_key_ref());
PreparedGenerateRequest {
@@ -385,13 +382,11 @@ impl Provider for CodexProvider {
request.payload(),
client_api_key_id,
"id",
gateway_core::account::AccountAffinity::Strict,
),
Operation::GenerateImage(request) => derive_codex_endpoint_session_affinity(
request.payload(),
client_api_key_id,
"session_id",
gateway_core::account::AccountAffinity::Strict,
),
_ => None,
};
@@ -419,12 +414,8 @@ impl Provider for CodexProvider {
let reasoning_effort = semantics.reasoning_effort.clone();
let previous_response_id = encoded.previous_response_id();
let continuation = ContinuationRequestObservation {
affinity_hash: derive_codex_session_affinity(
&encoded,
client_api_key_id,
gateway_core::account::AccountAffinity::Strict,
)
.map(|affinity| affinity.persistence_hash().to_owned()),
affinity_hash: derive_codex_session_affinity(&encoded, client_api_key_id)
.map(|affinity| affinity.persistence_hash().to_owned()),
previous_response_id_hash: previous_response_id.map(|response_id| {
derive_previous_response_id_hash(response_id, client_api_key_id)
}),
@@ -723,11 +714,6 @@ impl CodexProvider {
generate.protocol_payload().context(),
&middleware_headers,
);
upstream.client_account_thread_id = crate::credential::account_thread_with_headers(
generate.protocol_payload().body(),
generate.protocol_payload().context(),
&middleware_headers,
);
let processed = self.prepare_generate_request(&generate, upstream, &context);
let mut upstream_request = processed.upstream;
let session_transport_key =
@@ -742,9 +728,7 @@ impl CodexProvider {
&mut lease,
session_affinity.as_ref(),
cyber_policy_session_key.as_ref(),
context
.account_selection_policy()
.openai_session_affinity_ttl(),
context.account_selection_policy(),
)
.await
.map_err(map_selection_error)?;
backend/crates/providers/openai/src/transport/protocol/responses.rs
-2
View File
@@ -44,7 +44,6 @@ pub struct CodexResponsesRequest {
pub client_session_id: Option<String>,
/// 官方元数据中的逻辑会话身份，独立于缓存路由键
pub client_account_session_id: Option<String>,
pub client_account_thread_id: Option<String>,
/// 后代线程身份，用于严格模式下只跟随会话账号
pub client_account_follow_only: bool,
/// 客户端 thread ID，仅保留在受控本地上下文
@@ -506,7 +505,6 @@ impl CodexResponsesRequest {
client_conversation_id: None,
client_session_id: None,
client_account_session_id: None,
client_account_thread_id: None,
client_account_follow_only: false,
client_thread_id: None,
client_request_id: None,
backend/crates/providers/openai/src/transport/request.rs
-2
View File
@@ -733,8 +733,6 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map<Str
crate::credential::follows_session_with_headers(request.body(), context, &[]);
request.client_account_session_id =
gateway_protocol::openai::codex_account_session_id(request.body(), context);
request.client_account_thread_id =
gateway_protocol::openai::codex_account_thread_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());
backend/crates/providers/openai/tests/credential/contract/capacity.rs
+1 -1
View File
@@ -49,7 +49,7 @@ fn capacity_attempt(
.with_request_policy(request_policy),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
policy,
policy.with_openai_account_affinity(gateway_core::account::AccountAffinity::Strict),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
backend/crates/providers/openai/tests/provider/contract/session_binding.rs
+527 -139
View File
@@ -1,4 +1,4 @@
//! 验证最终请求身份、官方图片轮次关联与共享账号迁移
//! 验证会话主账号偏好、严格绑定、最终请求身份与轮次关联
use super::*;
@@ -585,28 +585,37 @@ async fn old_native_continuation_cannot_restore_the_pre_migration_account() {
assert!(server.received_requests().await.unwrap().is_empty());
}
fn relaxed_context(request_id: &str, scope: Arc<FrozenAccountScope>) -> AttemptContext {
relaxed_context_with_ttl(request_id, scope, Duration::from_secs(24 * 3600))
fn preferred_context(request_id: &str, scope: Arc<FrozenAccountScope>) -> AttemptContext {
preferred_context_with_ttl(request_id, scope, Duration::from_secs(24 * 3600))
}
fn relaxed_context_with_ttl(
fn preferred_context_with_ttl(
request_id: &str,
scope: Arc<FrozenAccountScope>,
ttl: Duration,
) -> AttemptContext {
preferred_context_with_policy(
request_id,
scope,
account_policy()
.with_openai_account_affinity(gateway_core::account::AccountAffinity::Preferred)
.with_openai_session_affinity_ttl(ttl),
)
}
fn preferred_context_with_policy(
request_id: &str,
scope: Arc<FrozenAccountScope>,
policy: gateway_core::account::AccountSelectionPolicy,
) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).unwrap(),
ClientApiKeyId::new("key_openai_contract").unwrap(),
),
NonZeroU32::new(1).unwrap(),
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
gateway_core::account::AccountSelectionPolicy::new(
gateway_core::account::RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_openai_session_affinity_ttl(ttl),
policy,
AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope(scope),
None,
CancellationToken::new(),
@@ -614,114 +623,141 @@ fn relaxed_context_with_ttl(
}
#[tokio::test]
async fn relaxed_children_keep_independent_bindings_across_busy_accounts_and_root_migration() {
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;
let provider = provider_with_affinity_and_base_url_and_leases(
&store,
affinity.clone(),
server.uri(),
leases.clone(),
);
let select = |thread: &'static str| {
let provider = provider.clone();
async move {
provider
.execute(
planned_request(
"openai",
Operation::Generate(generate_with_session_context(
"root",
Some(thread),
None,
)),
),
relaxed_context("req_relaxed", contract_account_scope()),
)
.await
.unwrap()
}
async fn relaxed_and_preferred_session_requests_apply_distinct_scheduling_policies() {
use gateway_core::account::{AccountAffinity, SmartSchedulingConfig};
use gateway_core::concurrency::ConcurrencyQueuePolicy;
use gateway_core::operation::{
ProviderHttpHeader, ProviderHttpMethod, ProviderHttpRequest, RawHttpPayload,
};
let root = select("root").await;
create_account(&store, "acct_subagent_b").await;
store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap());
let sibling = select("sibling").await;
assert_eq!(
sibling.metadata().provider_account_id().as_str(),
"acct_subagent_a",
"new child inherits root preference over weight"
);
drop(sibling);
leases
.busy_accounts
.lock()
.unwrap()
.insert(ProviderAccountId::new("acct_subagent_a").unwrap());
let child = select("child").await;
assert_eq!(
child.metadata().provider_account_id().as_str(),
"acct_subagent_b"
);
drop(child);
leases.busy_accounts.lock().unwrap().clear();
for (thread, expected) in [
("root", "acct_subagent_a"),
("sibling", "acct_subagent_a"),
("child", "acct_subagent_b"),
for (mode, max_waiting) in [
(AccountAffinity::Relaxed, 0),
(AccountAffinity::Relaxed, 10),
(AccountAffinity::Preferred, 0),
(AccountAffinity::Preferred, 10),
] {
assert_eq!(
select(thread)
.await
.metadata()
.provider_account_id()
.as_str(),
expected
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;
let provider = provider_with_affinity_and_base_url_and_leases(
&store,
affinity.clone(),
server.uri(),
leases.clone(),
);
let policy = account_policy()
.with_openai_account_affinity(mode)
.with_smart_scheduling(SmartSchedulingConfig::new([1.0; 6], true).unwrap())
.with_queue(ConcurrencyQueuePolicy {
max_waiting,
timeout: Duration::from_secs(1),
});
let select = |operation| {
let request = if matches!(operation, Operation::Generate(_)) {
planned_request("openai", operation)
} else if matches!(operation, Operation::ProviderHttp(_)) {
let provider_kind = ProviderKind::new("openai").unwrap();
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
SettingsValues::new(2, 10, "smart", BTreeMap::new(), None, None),
vec![provider_kind.clone()],
Vec::new(),
Vec::new(),
)
.unwrap();
let plan = snapshot
.plan_provider_endpoint(
&provider_kind,
Some(&UpstreamModelId::new("gpt-live-1-codex").unwrap()),
&operation,
contract_account_scope(),
&RoutingContext::default(),
)
.unwrap();
ProviderRequest::new(operation, plan.candidates()[0].clone())
} else {
planned_provider_endpoint_request("openai", operation)
};
provider.clone().execute(
request,
preferred_context_with_policy(
"req_preferred_session",
contract_account_scope(),
policy,
),
)
};
drop(select(turn_request("root", "shared-turn")).await.unwrap());
create_account(&store, "acct_subagent_b").await;
store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap());
// 后备池中另有低权重账号，验证分流继续遵循配置的调度策略
create_account(&store, "acct_provider_contract").await;
for operation in [
turn_request("root", "next-turn"),
Operation::Generate(generate_with_session_context("root", Some("child"), None)),
Operation::Generate(generate_with_session_context("root", Some("sibling"), None)),
image_for_turn(ImageRequestKind::Generation, "shared-turn"),
image_for_turn(ImageRequestKind::Edit, "shared-turn"),
Operation::ProviderHttp(
ProviderHttpRequest::new(
"realtime-calls",
ProviderHttpMethod::Post,
None,
vec![
ProviderHttpHeader::new("session-id", Bytes::from_static(b"root")),
ProviderHttpHeader::new("thread-id", Bytes::from_static(b"live-child")),
],
RawHttpPayload::new(
"openai",
Bytes::from_static(
br#"{"sdp":"v=0","session":{"model":"gpt-live-1-codex"}}"#,
),
)
.unwrap(),
)
.unwrap(),
),
Operation::Search(StandaloneSearchRequest::from_raw_json(
RawJsonPayload::new(
"openai",
Bytes::from_static(br#"{"id":"root","commands":{"search_query":[]}}"#),
)
.unwrap(),
)),
] {
for busy in [false, true, false] {
if busy {
leases
.busy_accounts
.lock()
.unwrap()
.insert(ProviderAccountId::new("acct_subagent_a").unwrap());
} else {
leases.busy_accounts.lock().unwrap().clear();
}
let stream =
tokio::time::timeout(Duration::from_millis(500), select(operation.clone()))
.await
.expect("an idle fallback must be admitted without waiting for the primary")
.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
if busy || mode == AccountAffinity::Relaxed {
"acct_subagent_b"
} else {
"acct_subagent_a"
}
);
}
}
assert_eq!(affinity.binding_count(), 1);
}
leases
.busy_accounts
.lock()
.unwrap()
.insert(ProviderAccountId::new("acct_subagent_a").unwrap());
assert_eq!(
select("root")
.await
.metadata()
.provider_account_id()
.as_str(),
"acct_subagent_b"
);
leases.busy_accounts.lock().unwrap().clear();
assert_eq!(
select("sibling")
.await
.metadata()
.provider_account_id()
.as_str(),
"acct_subagent_a",
"root migration cannot move an existing child"
);
assert_eq!(
select("new-sibling")
.await
.metadata()
.provider_account_id()
.as_str(),
"acct_subagent_b"
);
assert_eq!(
root.metadata().provider_account_id().as_str(),
"acct_subagent_a",
"the admitted root lease stays on its original account"
);
assert_eq!(affinity.binding_count(), 4);
}
#[tokio::test]
async fn relaxed_child_can_claim_without_a_root_and_model_access_can_split_accounts() {
async fn preferred_child_can_claim_without_a_root_and_model_access_can_split_accounts() {
use gateway_core::account::{AccountModelAccess, AccountModelAccessMode};
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
@@ -734,7 +770,7 @@ async fn relaxed_child_can_claim_without_a_root_and_model_access_can_split_accou
Operation::Generate(generate_with_session_context("root", Some(thread), None)),
model,
),
relaxed_context("req_relaxed_model", scope),
preferred_context("req_preferred_model", scope),
)
};
drop(
@@ -745,7 +781,7 @@ async fn relaxed_child_can_claim_without_a_root_and_model_access_can_split_accou
assert_eq!(
affinity.binding_count(),
1,
"early child must not claim the root binding"
"the first session request claims the shared primary"
);
drop(
select("root", "gpt-5.4", contract_account_scope())
@@ -795,11 +831,11 @@ async fn relaxed_child_can_claim_without_a_root_and_model_access_can_split_accou
.as_str(),
"acct_subagent_a"
);
assert_eq!(affinity.binding_count(), 3);
assert_eq!(affinity.binding_count(), 1);
}
#[tokio::test]
async fn relaxed_image_turn_follows_child_binding_and_accepts_matching_root_identity() {
async fn preferred_image_turn_prefers_the_session_primary_and_accepts_matching_root_identity() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
let affinity = Arc::new(MemorySessionAffinity::default());
@@ -816,7 +852,7 @@ async fn relaxed_image_turn_follows_child_binding_and_accepts_matching_root_iden
.clone()
.execute(
planned_request("openai", turn_request("root", "root-turn")),
relaxed_context("req_relaxed_root", contract_account_scope()),
preferred_context("req_preferred_root", contract_account_scope()),
)
.await
.unwrap(),
@@ -835,7 +871,7 @@ async fn relaxed_image_turn_follows_child_binding_and_accepts_matching_root_iden
.clone()
.execute(
planned_request("openai", child),
relaxed_context("req_relaxed_child", contract_account_scope()),
preferred_context("req_preferred_child", contract_account_scope()),
)
.await
.unwrap(),
@@ -859,21 +895,21 @@ async fn relaxed_image_turn_follows_child_binding_and_accepts_matching_root_iden
.clone()
.execute(
planned_provider_endpoint_request("openai", image),
relaxed_context("req_relaxed_image", contract_account_scope()),
preferred_context("req_preferred_image", contract_account_scope()),
)
.await
.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
"acct_subagent_b"
"acct_subagent_a"
);
}
// 新请求切到严格模式后，旧子线程轮次必须跟随根绑定；再切回宽松仍保留子绑定
// 轮次只关联会话身份，模式切换不能恢复一次临时分流的账号
for strict in [true, false] {
let attempt = if strict {
context("req_strict_old_turn", CancellationToken::new())
} else {
relaxed_context("req_relaxed_old_turn", contract_account_scope())
preferred_context("req_preferred_old_turn", contract_account_scope())
};
let stream = provider
.clone()
@@ -888,11 +924,7 @@ async fn relaxed_image_turn_follows_child_binding_and_accepts_matching_root_iden
.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
if strict {
"acct_subagent_a"
} else {
"acct_subagent_b"
}
"acct_subagent_a"
);
}
}
@@ -911,7 +943,7 @@ async fn configured_ttl_renews_binding_and_image_alias_only_at_successful_admiss
.clone()
.execute(
planned_request("openai", turn_request("root", "ttl-turn")),
relaxed_context_with_ttl("req_ttl_root", contract_account_scope(), ttl),
preferred_context_with_ttl("req_ttl_root", contract_account_scope(), ttl),
)
.await
.unwrap();
@@ -932,7 +964,7 @@ async fn configured_ttl_renews_binding_and_image_alias_only_at_successful_admiss
"openai",
image_for_turn(ImageRequestKind::Generation, "ttl-turn"),
),
relaxed_context_with_ttl("req_ttl_image", contract_account_scope(), week),
preferred_context_with_ttl("req_ttl_image", contract_account_scope(), week),
)
.await
.unwrap(),
@@ -942,7 +974,7 @@ async fn configured_ttl_renews_binding_and_image_alias_only_at_successful_admiss
}
#[tokio::test]
async fn relaxed_final_thread_rewrite_cannot_use_another_threads_account() {
async fn preferred_final_thread_rewrite_within_the_session_preserves_the_primary() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
let affinity = Arc::new(MemorySessionAffinity::default());
@@ -959,7 +991,7 @@ async fn relaxed_final_thread_rewrite_cannot_use_another_threads_account() {
.clone()
.execute(
request("root"),
relaxed_context("req_seed_root", contract_account_scope()),
preferred_context("req_seed_root", contract_account_scope()),
)
.await
.unwrap(),
@@ -972,7 +1004,7 @@ async fn relaxed_final_thread_rewrite_cannot_use_another_threads_account() {
.clone()
.execute(
request("child"),
relaxed_context("req_seed_child", contract_account_scope()),
preferred_context("req_seed_child", contract_account_scope()),
)
.await
.unwrap(),
@@ -1001,17 +1033,373 @@ async fn relaxed_final_thread_rewrite_cannot_use_another_threads_account() {
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
account_policy()
.with_openai_account_affinity(gateway_core::account::AccountAffinity::Relaxed),
.with_openai_account_affinity(gateway_core::account::AccountAffinity::Preferred),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
);
let error = provider
.execute(request("root"), attempt)
.await
.err()
.expect("thread owner mismatch");
assert!(error.retry_is_prohibited());
let stream = provider.execute(request("root"), attempt).await.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
"acct_subagent_a"
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn preferred_spillover_uses_selected_credentials_on_the_wire() {
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/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(3)
.mount(&server)
.await;
let provider = provider_with_affinity_and_base_url_and_leases(
&store,
affinity.clone(),
server.uri(),
leases.clone(),
);
for (turn, busy) in [("first", false), ("spillover", true), ("recovered", false)] {
if busy {
create_account(&store, "acct_subagent_b").await;
leases
.busy_accounts
.lock()
.unwrap()
.insert(ProviderAccountId::new("acct_subagent_a").unwrap());
} else {
leases.busy_accounts.lock().unwrap().clear();
}
let mut stream = provider
.clone()
.execute(
planned_request("openai", turn_request("root", turn)),
preferred_context("req_preferred_wire", contract_account_scope()),
)
.await
.unwrap();
while let Some(event) = stream.next().await {
event.unwrap();
}
}
let sent = server.received_requests().await.unwrap();
let accounts = sent
.iter()
.map(|request| request.headers["chatgpt-account-id"].to_str().unwrap())
.collect::<Vec<_>>();
assert_eq!(
accounts,
[
"chatgpt-acct_subagent_a",
"chatgpt-acct_subagent_b",
"chatgpt-acct_subagent_a"
]
);
assert_eq!(affinity.binding_count(), 1);
server.verify().await;
}
#[tokio::test]
async fn preferred_unavailable_primary_can_recover_without_losing_its_binding() {
use gateway_core::account::{
AccountAffinity, AccountRuntimeSignals, AccountSelectionPolicy, RotationStrategy,
};
use gateway_core::concurrency::ConcurrencyQueuePolicy;
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
let leases = Arc::new(TestLeaseCoordinator::default());
let server = MockServer::start().await;
let provider = provider_with_affinity_and_base_url_and_leases(
&store,
Arc::new(MemorySessionAffinity::default()),
server.uri(),
leases.clone(),
);
let policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::from_secs(30),
)
.with_openai_account_affinity(AccountAffinity::Preferred)
.with_queue(ConcurrencyQueuePolicy {
max_waiting: 10,
timeout: Duration::from_secs(1),
});
let select = || {
provider.clone().execute(
planned_request("openai", turn_request("root", "availability-turn")),
preferred_context_with_policy(
"req_primary_availability",
contract_account_scope(),
policy,
),
)
};
drop(select().await.unwrap());
create_account(&store, "acct_subagent_b").await;
let primary = ProviderAccountId::new("acct_subagent_a").unwrap();
for reason in ["concurrency", "interval", "disabled", "quota"] {
let mut signals = AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
};
match reason {
"concurrency" => signals.in_flight = 2,
"interval" => signals.last_started_at = Some(SystemTime::now()),
"disabled" => {
store.set_enabled(&primary, false).await.unwrap();
}
"quota" => {
store
.apply_quota_access(QuotaAccessChange {
account_id: primary.clone(),
expected_revision: store.account(primary.as_str()).unwrap().revision(),
state: QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
SystemTime::now(),
None,
),
})
.await
.unwrap();
}
_ => unreachable!(),
}
leases
.signals
.lock()
.unwrap()
.insert(primary.clone(), signals);
let fallback = tokio::time::timeout(Duration::from_millis(500), select())
.await
.unwrap()
.unwrap();
assert_eq!(
fallback.metadata().provider_account_id().as_str(),
"acct_subagent_b",
"{reason}"
);
drop(fallback);
leases.signals.lock().unwrap().clear();
store.set_enabled(&primary, true).await.unwrap();
store
.apply_quota_access(QuotaAccessChange {
account_id: primary.clone(),
expected_revision: store.account(primary.as_str()).unwrap().revision(),
state: QuotaState::allowed(SystemTime::now()),
})
.await
.unwrap();
assert_eq!(
select()
.await
.unwrap()
.metadata()
.provider_account_id()
.as_str(),
"acct_subagent_a",
"{reason}"
);
}
}
#[tokio::test]
async fn preferred_native_continuation_keeps_its_actual_owner_without_rebinding_the_primary() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
let provider = provider_with_affinity(&store, Arc::new(MemorySessionAffinity::default()));
let request = || planned_request("openai", turn_request("root", "continuation-turn"));
drop(
provider
.clone()
.execute(
request(),
preferred_context("req_primary", contract_account_scope()),
)
.await
.unwrap(),
);
create_account(&store, "acct_subagent_b").await;
for continuation_attempt in [
ContinuationAttempt::Native,
ContinuationAttempt::ReplayOwner,
ContinuationAttempt::ReplayAny,
] {
let account = ProviderAccountId::new("acct_subagent_b").unwrap();
let provider_kind = ProviderKind::new("openai").unwrap();
let client = ClientApiKeyId::new("key_openai_contract").unwrap();
let owner = ProviderAccountStateOwner::new(provider_kind.clone(), account.clone());
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_spillover"),
PreviousResponseId::new("resp_spillover"),
client.clone(),
provider_kind,
account,
)
.with_scope(gateway_core::engine::continuation::NativeContinuationScope::Persisted);
let attempt = AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_spillover_continuation").unwrap(),
client,
)
.with_request_location(Some(global_request_location())),
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
account_policy()
.with_openai_account_affinity(gateway_core::account::AccountAffinity::Preferred),
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner))
.with_account_scope(contract_account_scope()),
Some(ContinuationBinding::Pinned(pin)),
CancellationToken::new(),
)
.with_continuation_attempt(continuation_attempt);
let result = provider.clone().execute(request(), attempt).await;
if continuation_attempt == ContinuationAttempt::ReplayAny {
// 仍携带原生句柄的请求不能越过状态归属约束切回主账号
assert_eq!(
result.err().unwrap().kind(),
ProviderErrorKind::ContinuationRecoveryRequired
);
} else {
let stream = result.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
"acct_subagent_b"
);
}
assert_eq!(
provider
.clone()
.execute(
request(),
preferred_context("req_after_continuation", contract_account_scope())
)
.await
.unwrap()
.metadata()
.provider_account_id()
.as_str(),
"acct_subagent_a"
);
}
}
#[tokio::test]
async fn legacy_thread_turn_alias_resumes_on_the_session_primary_without_rewriting_identity() {
use gateway_core::provider_ports::{
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionAlias,
};
use sha2::Digest;
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_subagent_a").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let provider = provider_with_affinity(&store, affinity.clone());
let client = ClientApiKeyId::new("key_openai_contract").unwrap();
let provider_kind = ProviderKind::new("openai").unwrap();
let root = turn_request("root", "root-turn");
let root_key = ProviderSessionAffinityKey::try_new(
provider
.request_observation(&root, &client)
.continuation
.affinity_hash
.unwrap(),
)
.unwrap();
drop(
provider
.clone()
.execute(
planned_request("openai", root),
preferred_context("req_legacy_root", contract_account_scope()),
)
.await
.unwrap(),
);
create_account(&store, "acct_subagent_b").await;
let child_key = ProviderSessionAffinityKey::try_new("legacy-child-binding").unwrap();
affinity.seed_binding(
&provider_kind,
child_key.expose_to_store(),
ProviderAccountId::new("acct_subagent_b").unwrap(),
);
// 持久化轮次键沿用既有哈希合同，模拟升级前的 child → root 关联
let turn_key = ProviderSessionAffinityKey::try_new(hex::encode(sha2::Sha256::digest(
b"codex-session-affinity-v1\0client-turn\0key_openai_contract\0legacy-turn",
)))
.unwrap();
let legacy = ProviderSessionAlias {
session_key: child_key,
root_session_key: Some(root_key),
follow_only: false,
};
affinity
.bind_alias(
&provider_kind,
&turn_key,
&legacy,
Duration::from_secs(3600),
)
.await
.unwrap();
let resume = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", json!({"model":"gpt-5.4","input":"resume","client_metadata":{"session_id":"root","thread_id":"child","turn_id":"legacy-turn"}}).as_object().unwrap().clone()).unwrap(),
));
let stream = provider
.clone()
.execute(
planned_request("openai", resume),
preferred_context("req_legacy_resume", contract_account_scope()),
)
.await
.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
"acct_subagent_a"
);
drop(stream);
for attempt in [
preferred_context("req_legacy_image", contract_account_scope()),
context("req_legacy_strict_image", CancellationToken::new()),
] {
let stream = provider
.clone()
.execute(
planned_provider_endpoint_request(
"openai",
image_for_turn(ImageRequestKind::Generation, "legacy-turn"),
),
attempt,
)
.await
.unwrap();
assert_eq!(
stream.metadata().provider_account_id().as_str(),
"acct_subagent_a"
);
}
assert_eq!(
affinity
.load_alias(&provider_kind, &turn_key)
.await
.unwrap(),
Some(legacy)
);
}
backend/migrations/.frozen-sha256
+1
View File
@@ -21,3 +21,4 @@ e8c1339166ac02fb2d4a077424ee5f64f6dc20eb4d3f8ec4fcb09c0f2a0e61f6 0020_plugin_tr
22faf2646168275150b1786c0f463f9ff2ca52463a272479e4fcc6f46c8cfa37 0021_runtime_scheduling.sql
77a140ddfb9dfae9b7d7349e5bfc9c8a7e7927bae6e77c1a627b34fcbb714bdb 0022_refresh_margin_and_group_fast_mode.sql
d431ebf24669451da06bb6671136bf58cddcc713b43a1a705eb33f9cb4f9b577 0023_error_details_and_account_affinity.sql
fa99dde0fe0b3d7da6693d76b2735629dde0dc255ea8c53f0354857a4e80669b 0024_preferred_account_affinity.sql
backend/migrations/0024_preferred_account_affinity.sql
+9
View File
@@ -0,0 +1,9 @@
-- 会话主账号优先作为独立模式，保留已有亲和配置
alter table runtime_settings
drop constraint runtime_settings_openai_account_affinity_check,
add constraint runtime_settings_openai_account_affinity_check
check (openai_account_affinity in ('relaxed', 'preferred', 'strict')),
alter column openai_account_affinity set default 'strict';
-- 尚未发布过配置的初始记录采用严格模式，已经保存的设置保留原选项
update runtime_settings set openai_account_affinity = 'strict' where config_revision = 1;
docs/api.md
+11 -7
View File
@@ -1233,8 +1233,8 @@ HTTP 请求头及新建 WS 的握手提示按当时的最终出站档位构造
`response.create`；空闲连接不占名额，内部重试不重复占用。
修改 Key 策略对既有 WebSocket 连接的下一次请求同样生效，已开始的请求保持原有快照
运行设置可以分别启用 Key 与账号的有界排队。Key 并发满时按 Key 等待；OpenAI 宽松亲和下各线程独立调度，
新子线程优先沿用根账号，必要时可分流；已有子绑定不随根账号迁移。严格亲和下后代线程只等待当前会话账号，
运行设置可以分别启用 Key 与账号的有界排队。Key 并发满时按 Key 等待；OpenAI 优先亲和下会话请求优先主账号，
繁忙时先按调度策略选择可用账号，临时分流不改写会话绑定；没有可立即使用的账号时才按配置排队。严格亲和下后代线程只等待当前会话账号，
根线程迁移后跟随新账号，没有绑定时等待根线程首次认领
严格亲和下的后代线程始终排队，沿用 `concurrencyWaitTimeoutSeconds` 和请求截止时间；
@@ -1348,10 +1348,13 @@ Guardian 以 `subagent_kind` 或 `client_metadata.x-openai-subagent` 值 `guardi
仍受总等待容量与等待时限约束。不限并发的账号和关闭排队时的其余行为不变。
设置更新请求须包含该字段
`openaiAccountAffinity` 控制 OpenAI 账号亲和，默认 `relaxed`（宽松），也可设置为 `strict`（严格）。
宽松模式允许子代理按模型权限、可用性与容量独立分流，各线程保留自己的账号绑定；严格模式让同一会话共用账号，
子代理等待当前账号。需要选号时仍使用 `rotationStrategy` 指定的调度策略。
插件显式选号及原生续写的状态归属约束仍然生效，缺少可识别线程身份的子请求仍按跟随方式处理
`openaiAccountAffinity` 控制 OpenAI 账号亲和，默认 `strict`（严格），已有保存的模式保持不变：
- `relaxed`（宽松）：会话内请求直接按 `rotationStrategy` 选号，不优先主账号
- `preferred`（优先）：会话内所有可关联请求优先主账号，并发已满、请求间隔未到、额度耗尽、停用、限流冷却或不支持当前模型时临时分流，保留会话绑定，后续请求仍优先主账号
- `strict`（严格）：同一会话共用账号，后代请求等待当前账号并跟随根请求换号
插件显式选号及原生续写的状态归属约束在三种模式下仍然生效
`openaiSessionAffinityTtlHours` 是账号绑定及其会话关联的滑动保留时长，单位小时，默认 24，取值 1～720。
各类请求统一在发送前成功准入时按请求冻结值续期。
@@ -1396,7 +1399,8 @@ Guardian 以 `subagent_kind` 或 `client_metadata.x-openai-subagent` 值 `guardi
不影响已入队请求的位置、队内 FIFO、容量上限或等待超时。两项默认均为 `0`
`preferHigherWeight` 默认关闭。开启后，更高权重账号恢复可用时，后续允许重新选号的请求优先回切；
同权重且可用的会话亲和继续保留。没有可用亲和时，在最高可用权重层内按配置评分。
同权重且可用的会话亲和继续保留。OpenAI 优先模式沿用可用的会话主账号，不因权重回切而分流。
没有可用亲和时，在最高可用权重层内按配置评分。
原生续写账号绑定仍是硬约束，不因回切而主动换号或触发历史重放。配置随运行设置原子保存和发布，
新请求使用新值，已开始请求及其重试沿用原快照，无需重启
docs/architecture.md
+7 -5
View File
@@ -354,6 +354,7 @@ Client Key 并发占用、账号调度槽位与请求恢复记录使用可续期
`ConcurrencyWaitQueue` 在同一锁内读取队长并入队；Provider 只传递已限定续写范围的候选。
排队系数为零时保留最短队列规则；已入队位置、队内 FIFO、取消回收、容量上限和请求共享等待预算保持原合同。
高权重回切仅影响智能调度的软亲和：先验证候选资格，再决定是否由更高权重层覆盖亲和；同层亲和保留。
OpenAI 优先模式沿用可用的会话主账号，高权重回切仅参与需要重新选号的请求。
Provider 限定的原生续写账号范围仍是硬约束。回切在请求选号时发生，不新增后台任务或主动历史重放。
插件显式选号沿用既有候选校验和实际亲和结果，委托内置选号时才应用智能调度配置
@@ -614,18 +615,19 @@ Continuation 仍受原请求的 Client Key、账号范围、Provider 和发送/
OpenAI Provider 解释官方逻辑 `session_id` 与 `thread_id`，按 Client Key 隔离账号绑定，
不将缓存路由键或请求内容推断为账号身份。账号亲和模式随请求策略冻结：
- 宽松：根线程与各子线程独立绑定，新子线程优先使用根账号，模型权限、可用性或容量不允许时可独立选号；
根绑定只提供初始偏好，子线程认领、迁移和续写校验只操作自己的绑定，根迁移也不改写已有子绑定
- 宽松：会话内请求直接按调度策略选号，不使用主账号偏好，成功准入后更新当前账号
- 优先：会话内所有可关联请求优先使用主账号，并发已满、请求间隔未到、模型权限或可用性不允许时按调度策略临时分流；
首次成功准入认领会话主账号，分流只改变本次租约，不改写绑定，后续请求继续优先主账号
- 严格：根线程及后代共享当前账号，后代没有绑定时等待根线程认领，账号不可用时只排队，根迁移后重读共享绑定
无法识别线程的后代请求保留跟随约束。可关联的 Search、Images、Live 创建请求使用所属线程的绑定。
可关联的 Search、Images、Live 创建请求共用会话绑定，宽松和优先模式不依赖线程身份分流；严格模式下后代请求保留跟随约束。
需要选号时统一使用配置的调度策略；插件显式选号沿用原有裁决，native continuation 仍受独立的状态 owner 与安全重放边界约束
Store 只保存带版本的绑定事实，Provider 在取得租约后、发送前原子认领或迁移；快照冲突释放租约并重新选择。
Store 只保存带版本的绑定事实，Provider 在取得租约后、发送前原子认领、续期或按模式迁移；快照冲突释放租约并重新选择。
响应完成不写回账号，迁移前已准入的请求可以完成，不能把绑定改回旧账号。绑定存储不可用时不绕过绑定发送
官方独立 Images 仅携带 `turn_id`，通过已观测 Responses 轮次关联到会话，关联同时保留根线程或后代线程的调度权限。
轮次关联保存不透明线程键及可选的根会话键，不复制账号；显式根身份只用于校验轮次归属。
轮次关联保存不透明会话键，不复制账号；显式会话身份只用于校验轮次归属。
绑定和轮次关联按同一请求策略中的可配置时长保留，均按 Client Key 隔离。
成功准入按冻结时长续期，已关联的 Images 同时续期轮次关联；更新配置不扫描旧键，响应完成不续期。
未知轮次不推测会话归属，已有轮次不能改指其他会话或更改调度权限
frontend/src/api/modules/settings/index.ts
+1 -1
View File
@@ -13,7 +13,7 @@ export interface RequestLocation {
timezone: string
}
export type AccountAffinity = 'relaxed' | 'strict'
export type AccountAffinity = 'relaxed' | 'preferred' | 'strict'
export type RotationStrategy = 'smart' | 'quota_reset_priority' | 'round_robin' | 'sticky'
frontend/src/views/settings/components/AccountAffinityCard.vue
+3 -2
View File
@@ -7,7 +7,7 @@ const props = defineProps<{ disabled: boolean }>()
const accountAffinity = defineModel<AccountAffinity | ''>({ required: true })
const affinityTtlHours = defineModel<string>('ttlHours', { required: true })
const maxAccountRotations = defineModel<string>('maxAccountRotations', { required: true })
const affinityOptions = [{ label: '宽松', value: 'relaxed' }, { label: '严格', value: 'strict' }]
const affinityOptions = [{ label: '宽松', value: 'relaxed' }, { label: '优先', value: 'preferred' }, { label: '严格', value: 'strict' }]
</script>
<template>
@@ -22,7 +22,8 @@ const affinityOptions = [{ label: '宽松', value: 'relaxed' }, { label: '严格
</button>
</template>
<div class="max-w-72 space-y-2 px-3 py-2 text-cp-sm leading-relaxed text-cp-text-secondary">
<p>宽松：各线程独立绑定账号，子线程优先沿用主账号，可按模型、可用性和并发容量分流</p>
<p>宽松：会话内请求直接按调度策略选号，不优先主账号</p>
<p>优先：会话内请求优先沿用主账号，繁忙或不可用时临时分流，不改变会话绑定</p>
<p>严格：同一会话共用账号，子线程等待当前账号，并跟随主线程换号</p>
<p>重新选号仍遵循当前调度策略</p>
</div>