From a940eb599fddab92fc2140f7511972e067069c96 Mon Sep 17 00:00:00 2001 From: Zyy Date: Wed, 7 Oct 2026 20:18:03 +0800 Subject: [PATCH] fix(openai): allow atomic first binding from any session request Allow requests with a reliable session identity to claim a missing account binding, including child threads and independent background tasks. Remove the missing-binding wait queue while preserving existing migration rights and bound-account waiting behavior. Cover concurrent root/child claims, memory prewarm and turn reuse, middleware identity changes, cancellation, and bound-account timeouts. Update the documented strict-affinity contract. Co-authored-by: Codex --- .../openai/src/credential/selector.rs | 33 +- .../provider/contract/session_binding.rs | 313 ++++++++++++++++-- .../crates/providers/openai/tests/support.rs | 14 + docs/api.md | 6 +- docs/architecture.md | 4 +- 5 files changed, 318 insertions(+), 52 deletions(-) diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 591f6888c..7b2c5e0ef 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -126,7 +126,6 @@ pub(crate) struct CodexCyberPolicyScope { pub struct CodexCredentialSelector { waiting: ConcurrencyWaitQueue, - binding_waiting: ConcurrencyWaitQueue, provider_kind: ProviderKind, repository: CodexCredentialRepository, leases: Arc, @@ -286,7 +285,6 @@ impl CodexCredentialSelector { cookie_policy, risk_recovery: Mutex::new(HashMap::new()), waiting: ConcurrencyWaitQueue::default(), - binding_waiting: ConcurrencyWaitQueue::default(), account_feedback, } } @@ -349,12 +347,9 @@ impl CodexCredentialSelector { && lease.admitted_session.as_ref() != Some(affinity.key()) { let current = self.lookup_session_affinity(affinity.key()).await?; - 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) + if current + .as_ref() + .is_some_and(|binding| binding.account_id() != &selected_account) || !self .admit_session( affinity.key(), @@ -470,7 +465,6 @@ impl CodexCredentialSelector { queue_policy.timeout = Duration::from_secs(30); } } - let mut binding_wait = None; // 预留名额只对普通请求生效；Guardian 可用满全部名额，并在账号队列中排在普通请求之前 let reserve = request .attempt @@ -749,7 +743,8 @@ impl CodexCredentialSelector { } }; // 只约束内置调度；插件的显式选号保留原裁决和完整候选 - let selection = if (follow_only + // 首次绑定由任意会话请求原子认领；后代只在已有绑定时限制为跟随当前账号 + let selection = if ((follow_only && binding.is_some()) || (affinity_mode == AccountAffinity::Preferred && affinity.preferred_account().is_some())) && selection.is_none_or(|selection| !selection.is_policy_choice()) @@ -796,25 +791,7 @@ impl CodexCredentialSelector { selection.as_ref(), ); let Some(selection) = selection else { - if follow_only - && binding.is_none() - && let Some(key) = binding_key - { - let waiting = binding_wait.get_or_insert_with(|| { - CapacityWait::new( - &self.binding_waiting, - queue_policy, - request.attempt.deadline().at(), - request.attempt.concurrency_wait_budget(), - ) - .with_periodic_recheck(true) - }); - waiting.wait(std::slice::from_ref(key)).await?; - continue 'capacity; - } - if follow_only && let Some(binding) = binding.as_ref() { - binding_wait = None; // 冷却、排除、停用或当前账号不在本次范围时也只等待根线程恢复或迁移 waiting .wait(std::slice::from_ref(binding.account_id())) diff --git a/backend/crates/providers/openai/tests/provider/contract/session_binding.rs b/backend/crates/providers/openai/tests/provider/contract/session_binding.rs index a50b1f066..b3bc1a642 100644 --- a/backend/crates/providers/openai/tests/provider/contract/session_binding.rs +++ b/backend/crates/providers/openai/tests/provider/contract/session_binding.rs @@ -98,6 +98,205 @@ fn image_for_turn(kind: ImageRequestKind, turn: &str) -> Operation { )) } +#[tokio::test] +async fn independent_requests_claim_and_reuse_their_session_account() { + for source in [ + None, + Some("memory_consolidation"), + Some("future_background_task"), + ] { + assert_independent_session_can_claim(source).await; + } +} + +async fn assert_independent_session_can_claim(source: Option<&str>) { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200).set_body_raw( + format!("event: response.created\ndata: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_scope_capture\",\"status\":\"in_progress\"}}}}\n\n{CAPTURE_COMPLETED_SSE}"), + "text/event-stream", + ), + ) + .expect(1) + .mount(&server) + .await; + let provider = provider_with_affinity_and_base_url(&store, affinity.clone(), server.uri()); + + // 独立后台任务无需另一根线程先发请求，预热与正式请求共用自身会话绑定 + for prewarm in [true, false] { + let metadata = json!({ + "session_id": "independent-session", + "thread_id": "independent-worker", + "thread_source": source, + "request_kind": if prewarm { "prewarm" } else { "turn" }, + "turn_id": "independent-turn" + }); + let mut body = json!({ + "model": "gpt-5.4", + "input": "run background task", + "client_metadata": { + "session_id": "independent-session", + "thread_id": "independent-worker", + "x-openai-subagent": source, + "x-codex-turn-metadata": metadata.to_string() + } + }); + if prewarm { + body["generate"] = json!(false); + } + let payload = ProtocolPayload::json_object("openai", body.as_object().unwrap().clone()) + .unwrap() + .with_context(Map::from_iter([("use_websocket".into(), json!(prewarm))])); + let mut stream = timeout( + Duration::from_secs(1), + provider.clone().execute( + planned_request( + "openai", + Operation::Generate(GenerateRequest::from_protocol_payload(payload)), + ), + context( + if prewarm { + "req_independent_prewarm" + } else { + "req_independent_turn" + }, + CancellationToken::new(), + ), + ), + ) + .await + .expect("independent request must not wait for a nonexistent root") + .expect("independent request obtains an account"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!(affinity.binding_count(), 1); + if prewarm { + drop(stream); + create_account(&store, "acct_subagent_b").await; + store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap()); + } else { + let mut completed = false; + while let Some(event) = stream.next().await { + completed |= event + .unwrap() + .canonical_facts() + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + } + assert!( + completed, + "independent turn must reach the upstream and complete" + ); + } + } +} + +#[tokio::test] +async fn concurrent_root_and_child_requests_claim_one_session_account() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + create_account(&store, "acct_subagent_b").await; + let affinity = Arc::new(MemorySessionAffinity::with_initial_claim_barrier(2)); + let provider = provider_with_affinity(&store, affinity.clone()); + let select = |thread| { + provider.clone().execute( + planned_request( + "openai", + Operation::Generate(generate_with_session_context( + "req_root", + Some(thread), + Some(if thread == "req_root" { + "{}" + } else { + r#"{"subagent_kind":"thread_spawn"}"# + }), + )), + ), + context(thread, CancellationToken::new()), + ) + }; + let (first, second) = timeout(Duration::from_secs(2), async { + tokio::join!(select("req_root"), select("req_child")) + }) + .await + .expect("concurrent claims settle without waiting for a parent"); + let first = first.unwrap(); + let second = second.unwrap(); + assert_eq!( + first.metadata().provider_account_id(), + second.metadata().provider_account_id() + ); + assert_eq!(affinity.binding_count(), 1); +} + +#[tokio::test] +async fn parent_metadata_without_a_session_does_not_invent_a_binding() { + 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 generate = GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + json!({ + "model":"gpt-5.4", "input":"hello", + "client_metadata":{"x-codex-turn-metadata":json!({ + "thread_id":"child", "parent_thread_id":"parent", "subagent_kind":"thread_spawn" + }).to_string()} + }) + .as_object() + .unwrap() + .clone(), + ) + .unwrap(), + ); + let stream = timeout( + Duration::from_secs(1), + provider.execute( + planned_request("openai", Operation::Generate(generate)), + context("req_no_session", CancellationToken::new()), + ), + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!(affinity.binding_count(), 0); +} + +#[tokio::test] +async fn final_middleware_child_identity_can_claim_a_missing_binding() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let server = MockServer::start().await; + let provider = provider_with_affinity_and_base_url(&store, affinity.clone(), server.uri()); + let stream = provider.execute( + planned_request("openai", generate_operation()), + context_with_middleware( + "req_final_child_claim", + Arc::new(ChangeTurnMetadata(br#"{"session_id":"new-session","thread_id":"child","subagent_kind":"thread_spawn"}"#)), + FastMode::Default, + ), + ).await.expect("final child identity can atomically claim its session"); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!(affinity.binding_count(), 1); + assert!(server.received_requests().await.unwrap().is_empty()); +} + #[tokio::test] async fn official_image_turn_uses_current_session_owner_even_after_migration() { let store = Arc::new(MemoryAccountStore::default()); @@ -170,42 +369,85 @@ async fn unknown_image_turn_does_not_infer_a_conversation() { } #[tokio::test] -async fn child_waits_for_first_root_claim_and_cancellation_releases_waiter() { +async fn child_claims_before_root_and_then_follows_the_bound_account() { 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 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 child = || { planned_request( "openai", - Operation::Generate(generate_with_session_context("root", Some("child"), None)), + Operation::Generate(generate_with_session_context( + "root", + Some("child"), + Some(r#"{"subagent_kind":"thread_spawn"}"#), + )), ) }; + let stream = timeout( + Duration::from_secs(1), + provider.clone().execute( + child(), + context("req_child_first", CancellationToken::new()), + ), + ) + .await + .expect("child can claim before the root arrives") + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + drop(stream); + assert_eq!(affinity.binding_count(), 1); + create_account(&store, "acct_subagent_b").await; + store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap()); + let root = provider + .clone() + .execute( + planned_request("openai", turn_request("root", "root-turn")), + context("req_root_after_child", CancellationToken::new()), + ) + .await + .unwrap(); + assert_eq!( + root.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + drop(root); + + // 子线程可首次认领，但已有绑定后不能因为主账号繁忙而自行迁移 + leases + .busy_accounts + .lock() + .unwrap() + .insert(ProviderAccountId::new("acct_subagent_a").unwrap()); let cancel = CancellationToken::new(); let mut pending = Box::pin( provider .clone() - .execute(child(), context("req_wait_before_root", cancel.clone())), + .execute(child(), context("req_child_wait", cancel.clone())), ); assert!( timeout(Duration::from_millis(150), pending.as_mut()) .await .is_err() ); - assert_eq!( - affinity.binding_count(), - 0, - "a child must not claim the first account" - ); cancel.cancel(); let error = pending.await.err().expect("cancelled child"); assert_eq!(error.kind(), ProviderErrorKind::Cancelled); assert!(error.retry_is_prohibited()); - let mut pending = Box::pin( - provider - .clone() - .execute(child(), context("req_wait_again", CancellationToken::new())), - ); + let mut pending = Box::pin(provider.clone().execute( + child(), + context("req_child_wait_again", CancellationToken::new()), + )); assert!( timeout(Duration::from_millis(150), pending.as_mut()) .await @@ -214,19 +456,26 @@ async fn child_waits_for_first_root_claim_and_cancellation_releases_waiter() { let root = provider .clone() .execute( - planned_request("openai", turn_request("root", "root-turn")), - context("req_first_root", CancellationToken::new()), + planned_request("openai", turn_request("root", "root-migrate")), + context("req_root_migrate", CancellationToken::new()), ) .await .unwrap(); - let selected = root.metadata().provider_account_id().clone(); + assert_eq!( + root.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); drop(root); let resumed = timeout(Duration::from_secs(2), pending) .await .unwrap() .unwrap(); - assert_eq!(resumed.metadata().provider_account_id(), &selected); + assert_eq!( + resumed.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); assert_eq!(affinity.binding_count(), 1); + assert!(server.received_requests().await.unwrap().is_empty()); } #[tokio::test] @@ -339,7 +588,30 @@ async fn child_queue_timeout_does_not_rebind_or_allow_provider_fallback() { 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 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(), + ); + drop( + provider + .clone() + .execute( + planned_request("openai", turn_request("root", "root-turn")), + context("req_timeout_seed_root", CancellationToken::new()), + ) + .await + .unwrap(), + ); + create_account(&store, "acct_subagent_b").await; + leases + .busy_accounts + .lock() + .unwrap() + .insert(ProviderAccountId::new("acct_subagent_a").unwrap()); let attempt = AttemptContext::new( RequestAttemptContext::new( ModelRequestId::new("req_child_timeout").unwrap(), @@ -369,7 +641,8 @@ async fn child_queue_timeout_does_not_rebind_or_allow_provider_fallback() { .expect("queue timeout"); assert_eq!(error.kind(), ProviderErrorKind::ConcurrencyQueueTimeout); assert!(error.retry_is_prohibited()); - assert_eq!(affinity.binding_count(), 0); + assert_eq!(affinity.binding_count(), 1); + assert!(server.received_requests().await.unwrap().is_empty()); } #[derive(Debug)] diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 8380a77c7..22b885be1 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -745,6 +745,7 @@ impl ProviderLeasePort for TestLeaseCoordinator { #[derive(Default)] pub(crate) struct MemorySessionAffinity { + initial_claim_barrier: Option>, aliases: Mutex>, bindings: Mutex>, @@ -754,6 +755,13 @@ pub(crate) struct MemorySessionAffinity { } impl MemorySessionAffinity { + pub(crate) fn with_initial_claim_barrier(participants: usize) -> Self { + Self { + initial_claim_barrier: Some(Arc::new(tokio::sync::Barrier::new(participants))), + ..Self::default() + } + } + pub(crate) fn alias_ttls(&self) -> Vec { self.alias_ttls.lock().unwrap().clone() } @@ -845,6 +853,12 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity { Result, ProviderStoreError>, > { Box::pin(async move { + // 让并发首请求都读到空绑定，再验证发送前的原子认领冲突 + if expected.is_none() + && let Some(barrier) = &self.initial_claim_barrier + { + barrier.wait().await; + } let key = ( provider.as_str().to_owned(), key.expose_to_store().to_owned(), diff --git a/docs/api.md b/docs/api.md index 2b243ec22..6de8899eb 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1254,10 +1254,10 @@ HTTP 请求头及新建 WS 的握手提示按当时的最终出站档位构造 修改 Key 策略对既有 WebSocket 连接的下一次请求同样生效，已开始的请求保持原有快照 运行设置可以分别启用 Key 与账号的有界排队。Key 并发满时按 Key 等待；OpenAI 优先亲和下会话请求优先主账号， -繁忙时先按调度策略选择可用账号，临时分流不改写会话绑定；没有可立即使用的账号时才按配置排队。严格亲和下后代线程只等待当前会话账号， -根线程迁移后跟随新账号，没有绑定时等待根线程首次认领 +繁忙时先按调度策略选择可用账号，临时分流不改写会话绑定；没有可立即使用的账号时才按配置排队。严格亲和下已有绑定的后代线程只等待当前会话账号， +根线程迁移后跟随新账号。绑定缺失时，任何有可靠会话身份的请求都可按调度策略原子首绑，后续请求沿用该绑定 -严格亲和下的后代线程始终排队，沿用 `concurrencyWaitTimeoutSeconds` 和请求截止时间； +严格亲和下的后代线程始终启用账号排队，沿用 `concurrencyWaitTimeoutSeconds` 和请求截止时间； `maxWaitingPerAccount` 为 0 时使用每队列 1,000 人上限，否则沿用配置的上限。 可识别来源的 Search、Images、Live 创建请求采用同样规则，身份与轮次关联见[会话绑定](architecture.md#6-路由账号范围与-continuation)。 插件显式选号保留原有行为。普通请求关闭账号排队且没有可用容量时返回 `503` / `account_capacity_unavailable`， diff --git a/docs/architecture.md b/docs/architecture.md index f0f8d3fd1..caf23f38f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -624,12 +624,14 @@ OpenAI Provider 解释官方逻辑 `session_id` 与 `thread_id`，按 Client Key - 宽松：会话内请求直接按调度策略选号，不使用主账号偏好，成功准入后更新当前账号 - 优先：会话内所有可关联请求优先使用主账号，并发已满、请求间隔未到、模型权限或可用性不允许时按调度策略临时分流； 首次成功准入认领会话主账号，分流只改变本次租约，不改写绑定，后续请求继续优先主账号 -- 严格：根线程及后代共享当前账号，后代没有绑定时等待根线程认领，账号不可用时只排队，根迁移后重读共享绑定 +- 严格：根线程及后代共享当前账号；绑定缺失时任意会话请求均可按调度策略首次认领， + 已有绑定的后代只跟随当前账号，账号不可用时排队，根迁移后重读共享绑定 可关联的 Search、Images、Live 创建请求共用会话绑定，宽松和优先模式不依赖线程身份分流；严格模式下后代请求保留跟随约束。 需要选号时统一使用配置的调度策略；插件显式选号沿用原有裁决，native continuation 仍受独立的状态 owner 与安全重放边界约束 Store 只保存带版本的绑定事实，Provider 在取得租约后、发送前原子认领、续期或按模式迁移；快照冲突释放租约并重新选择。 +首次认领不依赖根线程先发请求，独立后台任务也使用自身会话绑定；没有可靠会话身份时正常选号，不推测会话关系。 响应完成不写回账号，迁移前已准入的请求可以完成，不能把绑定改回旧账号。绑定存储不可用时不绕过绑定发送 官方独立 Images 仅携带 `turn_id`，通过已观测 Responses 轮次关联到会话，关联同时保留根线程或后代线程的调度权限。