diff --git a/backend/crates/gateway-api/tests/admin/settings.rs b/backend/crates/gateway-api/tests/admin/settings.rs index e4abfe63d..b14d8d93a 100644 --- a/backend/crates/gateway-api/tests/admin/settings.rs +++ b/backend/crates/gateway-api/tests/admin/settings.rs @@ -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(); diff --git a/backend/crates/gateway-core/src/account/selection.rs b/backend/crates/gateway-core/src/account/selection.rs index 00b83cc93..79c03c3c6 100644 --- a/backend/crates/gateway-core/src/account/selection.rs +++ b/backend/crates/gateway-core/src/account/selection.rs @@ -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 { 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 { diff --git a/backend/crates/gateway-core/src/settings/values.rs b/backend/crates/gateway-core/src/settings/values.rs index fe5753b3f..a75c075c9 100644 --- a/backend/crates/gateway-core/src/settings/values.rs +++ b/backend/crates/gateway-core/src/settings/values.rs @@ -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, diff --git a/backend/crates/gateway-core/tests/settings/mod.rs b/backend/crates/gateway-core/tests/settings/mod.rs index 3dc945919..c3b6f0291 100644 --- a/backend/crates/gateway-core/tests/settings/mod.rs +++ b/backend/crates/gateway-core/tests/settings/mod.rs @@ -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 diff --git a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs index a400c4922..2c331d64f 100644 --- a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs @@ -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; + } +} diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 664c3b405..739ec6077 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -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, 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 { 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 { - 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 { let mut body = serde_json::from_slice::>(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, context: &Map, 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 { 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, diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index b8df5b4dc..435044416 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -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; diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index ac792f094..591f6888c 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -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 { 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 { 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::>(); - 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, 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( diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 8bc35628d..145ec52eb 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -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)?; diff --git a/backend/crates/providers/openai/src/provider/live.rs b/backend/crates/providers/openai/src/provider/live.rs index 02acfd786..e49c37e38 100644 --- a/backend/crates/providers/openai/src/provider/live.rs +++ b/backend/crates/providers/openai/src/provider/live.rs @@ -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)?; diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index f857155c8..4066a523f 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -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)?; diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index 3811e652d..79051508f 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -44,7 +44,6 @@ pub struct CodexResponsesRequest { pub client_session_id: Option, /// 官方元数据中的逻辑会话身份，独立于缓存路由键 pub client_account_session_id: Option, - pub client_account_thread_id: Option, /// 后代线程身份，用于严格模式下只跟随会话账号 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, diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index d62b833ec..6d0cad009 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -733,8 +733,6 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map) -> AttemptContext { - relaxed_context_with_ttl(request_id, scope, Duration::from_secs(24 * 3600)) +fn preferred_context(request_id: &str, scope: Arc) -> 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, 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, + 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::>(); + 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) + ); +} diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index 9fbdf1363..639071c0a 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -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 diff --git a/backend/migrations/0024_preferred_account_affinity.sql b/backend/migrations/0024_preferred_account_affinity.sql new file mode 100644 index 000000000..1c116d8a7 --- /dev/null +++ b/backend/migrations/0024_preferred_account_affinity.sql @@ -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; diff --git a/docs/api.md b/docs/api.md index da306f791..0276b96d5 100644 --- a/docs/api.md +++ b/docs/api.md @@ -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 优先模式沿用可用的会话主账号，不因权重回切而分流。 +没有可用亲和时，在最高可用权重层内按配置评分。 原生续写账号绑定仍是硬约束，不因回切而主动换号或触发历史重放。配置随运行设置原子保存和发布， 新请求使用新值，已开始请求及其重试沿用原快照，无需重启 diff --git a/docs/architecture.md b/docs/architecture.md index 204d59737..b88bce5df 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -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 同时续期轮次关联；更新配置不扫描旧键，响应完成不续期。 未知轮次不推测会话归属，已有轮次不能改指其他会话或更改调度权限 diff --git a/frontend/src/api/modules/settings/index.ts b/frontend/src/api/modules/settings/index.ts index fea626987..cb7a09070 100644 --- a/frontend/src/api/modules/settings/index.ts +++ b/frontend/src/api/modules/settings/index.ts @@ -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' diff --git a/frontend/src/views/settings/components/AccountAffinityCard.vue b/frontend/src/views/settings/components/AccountAffinityCard.vue index faa9819a8..19dd2b13d 100644 --- a/frontend/src/views/settings/components/AccountAffinityCard.vue +++ b/frontend/src/views/settings/components/AccountAffinityCard.vue @@ -7,7 +7,7 @@ const props = defineProps<{ disabled: boolean }>() const accountAffinity = defineModel({ required: true }) const affinityTtlHours = defineModel('ttlHours', { required: true }) const maxAccountRotations = defineModel('maxAccountRotations', { required: true }) -const affinityOptions = [{ label: '宽松', value: 'relaxed' }, { label: '严格', value: 'strict' }] +const affinityOptions = [{ label: '宽松', value: 'relaxed' }, { label: '优先', value: 'preferred' }, { label: '严格', value: 'strict' }]
-

宽松：各线程独立绑定账号，子线程优先沿用主账号，可按模型、可用性和并发容量分流

+

宽松：会话内请求直接按调度策略选号，不优先主账号

+

优先：会话内请求优先沿用主账号，繁忙或不可用时临时分流，不改变会话绑定

严格：同一会话共用账号，子线程等待当前账号，并跟随主线程换号

重新选号仍遵循当前调度策略