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

fix(openai): allow atomic first binding from any session request

Browse Source 
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 <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-07 20:18:03 +08:00
co-authored by Codex
parent 22a848b360
commit a940eb599f
5 changed files with 318 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/credential/selector.rs
+5 -28
View File
@@ -126,7 +126,6 @@ pub(crate) struct CodexCyberPolicyScope {
pub struct CodexCredentialSelector {
waiting: ConcurrencyWaitQueue<ProviderAccountId>,
binding_waiting: ConcurrencyWaitQueue<ProviderSessionAffinityKey>,
provider_kind: ProviderKind,
repository: CodexCredentialRepository,
leases: Arc<dyn ProviderLeasePort>,
@@ -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()))
backend/crates/providers/openai/tests/provider/contract/session_binding.rs
+293 -20
View File
@@ -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)]
backend/crates/providers/openai/tests/support.rs
+14
View File
@@ -745,6 +745,7 @@ impl ProviderLeasePort for TestLeaseCoordinator {
#[derive(Default)]
pub(crate) struct MemorySessionAffinity {
initial_claim_barrier: Option<Arc<tokio::sync::Barrier>>,
aliases: Mutex<BTreeMap<(String, String), gateway_core::provider_ports::ProviderSessionAlias>>,
bindings:
Mutex<BTreeMap<(String, String), gateway_core::provider_ports::ProviderSessionBinding>>,
@@ -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<Duration> {
self.alias_ttls.lock().unwrap().clone()
}
@@ -845,6 +853,12 @@ impl ProviderSessionAffinityPort for MemorySessionAffinity {
Result<Option<gateway_core::provider_ports::ProviderSessionBinding>, 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(),
docs/api.md
+3 -3
View File
@@ -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`，
docs/architecture.md
+3 -1
View File
@@ -624,12 +624,14 @@ OpenAI Provider 解释官方逻辑 `session_id` 与 `thread_id`，按 Client Key
- 宽松：会话内请求直接按调度策略选号，不使用主账号偏好，成功准入后更新当前账号
- 优先：会话内所有可关联请求优先使用主账号，并发已满、请求间隔未到、模型权限或可用性不允许时按调度策略临时分流；
首次成功准入认领会话主账号，分流只改变本次租约，不改写绑定，后续请求继续优先主账号
- 严格：根线程及后代共享当前账号，后代没有绑定时等待根线程认领，账号不可用时只排队，根迁移后重读共享绑定
- 严格：根线程及后代共享当前账号；绑定缺失时任意会话请求均可按调度策略首次认领，
已有绑定的后代只跟随当前账号，账号不可用时排队，根迁移后重读共享绑定
可关联的 Search、Images、Live 创建请求共用会话绑定，宽松和优先模式不依赖线程身份分流；严格模式下后代请求保留跟随约束。
需要选号时统一使用配置的调度策略；插件显式选号沿用原有裁决，native continuation 仍受独立的状态 owner 与安全重放边界约束
Store 只保存带版本的绑定事实，Provider 在取得租约后、发送前原子认领、续期或按模式迁移；快照冲突释放租约并重新选择。
首次认领不依赖根线程先发请求，独立后台任务也使用自身会话绑定；没有可靠会话身份时正常选号，不推测会话关系。
响应完成不写回账号，迁移前已准入的请求可以完成，不能把绑定改回旧账号。绑定存储不可用时不绕过绑定发送
官方独立 Images 仅携带 `turn_id`，通过已观测 Responses 轮次关联到会话，关联同时保留根线程或后代线程的调度权限。