diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 2e67cfaa3..7e0d561b6 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -267,6 +267,38 @@ impl CodexProvider { } } + async fn remember_generate_turn( + &self, + generate: &GenerateRequest, + headers: &[MiddlewareHeader], + context: &AttemptContext, + affinity: Option<&CodexSessionAffinity>, + ) -> Result<(), ProviderError> { + if context.is_diagnostic_required_account() { + return Ok(()); + } + if let Some(affinity) = affinity + && let Some(turn) = turn_id_with_headers( + generate.protocol_payload().body(), + generate.protocol_payload().context(), + headers, + ) + .and_then(|turn| derive_turn_alias(&turn, context.client_api_key_ref())) + { + self.selector + .remember_turn( + &turn, + affinity, + context + .account_selection_policy() + .openai_session_affinity_ttl(), + ) + .await + .map_err(map_selection_error)?; + } + Ok(()) + } + fn client_for_request( &self, context: &AttemptContext, @@ -622,10 +654,20 @@ impl Provider for CodexProvider { .transpose() .map_err(map_request_error)?; // 固定调度约束:Guardian 分类先于选号,不能随原生/适配器发送路径改变 - // 复用编码器的权威 metadata 解析,但只为原生路径准备会话与传输状态 + // 复用编码器的权威 metadata 解析,但只为原生路径准备续接与传输状态 let guardian = upstream .as_ref() .is_some_and(CodexResponsesRequest::is_guardian); + let adapter_affinity = upstream + .as_ref() + .filter(|_| adapter.is_some()) + .map(|request| { + ( + derive_codex_session_affinity(request, context.client_api_key_ref()), + derive_codex_cyber_policy_session_key(request, context.client_api_key_ref()), + false, + ) + }); let preselection = upstream .filter(|_| adapter.is_none()) .map(|upstream| self.prepare_generate_request(generate, upstream, &context)); @@ -633,7 +675,7 @@ impl Provider for CodexProvider { prepared.checked_native_continuation_scope(&context)?; } let (selection_session_affinity, selection_cyber_policy_key, requires_websocket) = - preselection.map_or((None, None, false), |prepared| { + preselection.map_or(adapter_affinity.unwrap_or((None, None, false)), |prepared| { let requires_websocket = transport_requirement(&prepared.upstream).requires_websocket() || (context.continuation_attempt() == ContinuationAttempt::Native @@ -680,6 +722,16 @@ impl Provider for CodexProvider { }; let account_selection_wait_ms = u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + if adapter.is_some() { + // 适配器自行解释业务改写,宿主轮次关联只使用本次选号已确认的会话 + self.remember_generate_turn( + generate, + &[], + &context, + selection_session_affinity.as_ref(), + ) + .await?; + } let provider_kind = ProviderKind::new(PROVIDER_NAME) .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?; let account_id = lease.account_id().clone(); @@ -781,39 +833,24 @@ impl CodexProvider { derive_codex_transport_key(&upstream_request, context.client_api_key_ref()); let previous_session = processed.previous_session; let continuation_requested = processed.continuation_requested; - let session_affinity = processed.session_affinity; - let cyber_policy_session_key = processed.cyber_policy_session_key; if !context.is_diagnostic_required_account() { self.selector .validate_translated_selection( &mut lease, - session_affinity.as_ref(), - cyber_policy_session_key.as_ref(), + processed.session_affinity.as_ref(), + processed.cyber_policy_session_key.as_ref(), context.account_selection_policy(), ) .await .map_err(map_selection_error)?; } - if !context.is_diagnostic_required_account() - && let Some(affinity) = session_affinity.as_ref() - && let Some(turn) = turn_id_with_headers( - generate.protocol_payload().body(), - generate.protocol_payload().context(), - &middleware_headers, - ) - .and_then(|turn| derive_turn_alias(&turn, context.client_api_key_ref())) - { - self.selector - .remember_turn( - &turn, - affinity, - context - .account_selection_policy() - .openai_session_affinity_ttl(), - ) - .await - .map_err(map_selection_error)?; - } + self.remember_generate_turn( + &generate, + &middleware_headers, + &context, + processed.session_affinity.as_ref(), + ) + .await?; let lease = Arc::new(lease); if previous_session.as_ref().is_some_and(|state| { state diff --git a/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs b/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs index 06ab16d2f..e95efe7f7 100644 --- a/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs +++ b/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs @@ -7,6 +7,174 @@ use gateway_core::engine::upstream_adapter::{ UpstreamAccountConnection, UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan, }; +#[derive(Debug, Default)] +struct UnmatchedAdapterPlan { + selections: std::sync::atomic::AtomicUsize, +} + +impl UpstreamAdapterPlan for UnmatchedAdapterPlan { + fn select( + &self, + _: &AttemptContext, + provider: &ProviderKind, + model: &UpstreamModelId, + ) -> Result>, gateway_core::error::ProviderError> { + assert_eq!(provider.as_str(), "openai"); + assert_eq!(model.as_str(), "gpt-5.4"); + self.selections + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(None) + } +} + +#[tokio::test] +async fn unmatched_adapter_preserves_native_oauth_and_api_key_http_wire() { + for api_key in [false, true] { + let store = Arc::new(MemoryAccountStore::default()); + let server = MockServer::start().await; + if api_key { + store + .seed_api_key( + "acct_provider_contract", + server.uri(), + provider_openai::credential::ResponsesTransport::Http, + ) + .await; + } else { + create_account(&store, "acct_provider_contract").await; + } + let endpoint = if api_key { + "/responses" + } else { + "/codex/responses" + }; + let response = concat!( + "event: response.output_text.delta\n", + "data: { \"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"native\",\"future\":{\"keep\":true} }\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_wire\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2},\"future\":true}}\n\n" + ); + Mock::given(method("POST")) + .and(path(endpoint)) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(response), + ) + .expect(2) + .mount(&server) + .await; + let provider = provider_with_base_url(&store, server.uri()); + let plan = Arc::new(UnmatchedAdapterPlan::default()); + // 不透明内容仅为离线夹具,不向真实上游发送或尝试解码 + let body = json!({ + "model": "gpt-5.4", + "stream": true, + "store": false, + "instructions": "Keep the original history.", + "input": [ + {"type":"reasoning","id":"rs_fixture","summary":[],"encrypted_content":"opaque-offline-fixture"}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"Third-party visible reply"}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"Continue"}]} + ], + "tools": [{"type":"function","name":"lookup_fixture","parameters":{"type":"object","properties":{}}}], + "future_extension": {"nested":[1,true,"preserve"]} + }); + for with_plan in [false, true] { + let frozen_plan = with_plan.then(|| { + gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( + plan.clone(), + ExtensionSetReference::new( + ExtensionSetId::new("unmatched-wire-test".to_owned()).unwrap(), + Arc::new(TestExtensionLease), + ), + ) + }); + let context = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_unmatched_wire").unwrap(), + ClientApiKeyId::new("key_openai_contract").unwrap(), + ) + .with_upstream_adapters(frozen_plan) + .with_middleware( + None, + Arc::from([]), + "/v1/responses".into(), + ClientTransport::HttpSse, + ), + NonZeroU32::MIN, + SystemTime::now() + Duration::from_secs(30), + account_policy(), + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(contract_account_scope()), + None, + CancellationToken::new(), + ); + let payload = ProtocolPayload::json_object("openai", body.as_object().unwrap().clone()) + .unwrap() + .with_context(Map::from_iter([ + ("use_websocket".into(), json!(false)), + ( + "opaque_request_headers".into(), + json!([ + ["x-future-business", STANDARD.encode(b"first")], + ["x-future-business", STANDARD.encode(b"second")] + ]), + ), + ])); + let mut stream = Arc::clone(&provider) + .execute( + planned_request( + "openai", + Operation::Generate(GenerateRequest::from_protocol_payload(payload)), + ), + context, + ) + .await + .unwrap(); + let mut wire = Vec::new(); + while let Some(event) = stream.next().await { + if let Some(frame) = event + .unwrap() + .wire_event() + .and_then(|wire| wire.raw_sse_frame()) + { + wire.extend_from_slice(frame); + } + } + assert_eq!( + wire, + response.as_bytes(), + "response changed: api_key={api_key}, with_plan={with_plan}" + ); + } + assert!(plan.selections.load(std::sync::atomic::Ordering::SeqCst) > 0); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].method, requests[1].method); + assert_eq!(requests[0].url, requests[1].url); + assert_eq!( + requests[0].headers, requests[1].headers, + "headers changed: api_key={api_key}" + ); + assert_eq!( + requests[0].body, requests[1].body, + "body changed: api_key={api_key}" + ); + let sent = captured_request_body(&requests[1]); + for field in ["input", "tools", "future_extension", "instructions"] { + assert_eq!( + sent[field], body[field], + "native field changed: {field}, api_key={api_key}" + ); + } + assert_eq!( + captured_header_values(&requests[1], "x-future-business"), + vec![b"first".to_vec(), b"second".to_vec()] + ); + } +} + #[derive(Debug)] struct AdapterProbe { polls: Arc, @@ -170,6 +338,388 @@ async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_an type SelectedConnection = Arc; +type AdapterInvocations = Arc>>; + +struct AffinityAdapterProbe { + invocations: Option, +} + +impl std::fmt::Debug for AffinityAdapterProbe { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AffinityAdapterProbe") + .finish_non_exhaustive() + } +} + +impl UpstreamAdapterPlan for AffinityAdapterProbe { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + ) -> Result>, gateway_core::error::ProviderError> { + Ok(Some(Arc::new(Self { + invocations: self.invocations.clone(), + }))) + } +} + +impl UpstreamAdapter for AffinityAdapterProbe { + fn transport(&self) -> &str { + "http_sse" + } + + fn execute( + self: Arc, + invocation: UpstreamAdapterInvocation, + ) -> gateway_core::engine::provider::EventStream { + Box::pin(futures::stream::once(async move { + if let Some(invocations) = &self.invocations { + invocations.lock().unwrap().push(invocation); + } + Err(gateway_core::error::ProviderError::new( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )) + })) + } +} + +fn affinity_adapter_context( + request_id: &str, + client_key: &str, + middleware: Option>, + invocations: Option, +) -> AttemptContext { + let extension = ExtensionSetReference::new( + ExtensionSetId::new("adapter-affinity-test".to_owned()).unwrap(), + Arc::new(TestExtensionLease), + ); + let plan = gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( + Arc::new(AffinityAdapterProbe { invocations }), + extension.clone(), + ); + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new(request_id).unwrap(), + ClientApiKeyId::new(client_key).unwrap(), + ) + .with_upstream_adapters(Some(plan)) + .with_middleware( + middleware.map(|plan| FrozenMiddlewarePlan::new(plan, extension)), + Arc::from([]), + "/v1/responses".to_owned(), + ClientTransport::HttpSse, + ), + NonZeroU32::MIN, + SystemTime::now() + Duration::from_secs(5), + account_policy(), + AccountAttemptContext::new(BTreeSet::new(), None, None) + .with_account_scope(contract_account_scope()), + None, + CancellationToken::new(), + ) +} + +#[tokio::test] +async fn upstream_adapter_reuses_session_account_and_isolates_client_keys() { + 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 request = || planned_request("openai", session_binding::turn_request("session", "turn")); + drop( + provider + .clone() + .execute( + request(), + affinity_adapter_context("req_seed", "key_a", None, None), + ) + .await + .unwrap(), + ); + assert_eq!(affinity.binding_count(), 1); + assert!(!affinity.alias_ttls().is_empty()); + create_account(&store, "acct_subagent_b").await; + store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap()); + for (key, expected) in [("key_a", "acct_subagent_a"), ("key_b", "acct_subagent_b")] { + let stream = provider + .clone() + .execute( + request(), + affinity_adapter_context("req_follow", key, None, None), + ) + .await + .unwrap(); + assert_eq!(stream.metadata().provider_account_id().as_str(), expected); + } + assert_eq!(affinity.binding_count(), 2); + store + .set_enabled(&ProviderAccountId::new("acct_subagent_a").unwrap(), false) + .await + .unwrap(); + let stream = provider + .execute( + request(), + affinity_adapter_context("req_migrate", "key_a", None, None), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn upstream_adapter_preserves_plugin_rewrites_without_changing_admitted_affinity() { + use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAffinityPort}; + use sha2::Digest; + + 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 client = ClientApiKeyId::new("key_a").unwrap(); + let provider_kind = ProviderKind::new("openai").unwrap(); + let operation = session_binding::turn_request("session-a", "client-turn"); + let session_key = ProviderSessionAffinityKey::try_new( + provider + .request_observation(&operation, &client) + .continuation + .affinity_hash + .unwrap(), + ) + .unwrap(); + let turn_key = |turn: &str| { + ProviderSessionAffinityKey::try_new(hex::encode(sha2::Sha256::digest(format!( + "codex-session-affinity-v1\0client-turn\0key_a\0{turn}" + )))) + .unwrap() + }; + drop( + provider + .clone() + .execute( + planned_request( + "openai", + session_binding::turn_request("session-a", "seed-a"), + ), + affinity_adapter_context("req_seed_a", "key_a", None, None), + ) + .await + .unwrap(), + ); + create_account(&store, "acct_subagent_b").await; + store.set_scheduling("acct_subagent_b", None, AccountWeight::new(100).unwrap()); + drop( + provider + .clone() + .execute( + planned_request( + "openai", + session_binding::turn_request("session-b", "seed-b"), + ), + affinity_adapter_context("req_seed_b", "key_a", None, None), + ) + .await + .unwrap(), + ); + let other_turn = affinity + .load_alias(&provider_kind, &turn_key("seed-b")) + .await + .unwrap() + .unwrap(); + // 插件可使用已有会话或自己的会话标识,宿主不据此重绑或新增亲和状态 + for plugin_session in ["session-b", "plugin-session"] { + let metadata = json!({ + "session_id": plugin_session, "thread_id": plugin_session, "turn_id": "plugin-turn" + }) + .to_string(); + let replacement = json!({"x-codex-turn-metadata": metadata}); + let headers = vec![MiddlewareHeader::new( + "x-codex-turn-metadata", + Bytes::from(metadata), + )]; + let middleware = RecordingMiddleware { + observed: Default::default(), + replacement: ("client_metadata".to_owned(), replacement.clone()), + request_headers: headers.clone(), + }; + let invocations = Arc::new(Mutex::new(Vec::new())); + let mut stream = provider + .clone() + .execute( + planned_request("openai", operation.clone()), + affinity_adapter_context( + "req_changed", + "key_a", + Some(Arc::new(middleware)), + Some(invocations.clone()), + ), + ) + .await + .expect("plugin rewrites must reach the adapter with the selected account"); + assert!( + invocations.lock().unwrap().is_empty(), + "adapter must stay cold" + ); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!( + stream.next().await.unwrap().unwrap_err().kind(), + ProviderErrorKind::Cancelled + ); + let invocation = invocations + .lock() + .unwrap() + .pop() + .expect("adapter was polled"); + assert_eq!(invocation.account.account_id().as_str(), "acct_subagent_a"); + assert_eq!( + invocation.metadata.provider_account_id(), + invocation.account.account_id() + ); + assert_eq!(invocation.headers, headers); + let Operation::Generate(original) = &operation else { + unreachable!() + }; + let mut expected = original.protocol_payload().body().clone(); + expected.insert("client_metadata".to_owned(), replacement); + let Operation::Generate(delivered) = &invocation.operation else { + unreachable!() + }; + assert_eq!(delivered.protocol_payload().protocol(), "openai"); + assert_eq!(delivered.protocol_payload().body(), &expected); + assert_eq!( + delivered.protocol_payload().context(), + original.protocol_payload().context() + ); + assert_eq!(affinity.binding_count(), 2); + assert_eq!( + affinity + .load(&provider_kind, &session_key) + .await + .unwrap() + .unwrap() + .as_str(), + "acct_subagent_a" + ); + assert_eq!( + affinity + .load(&provider_kind, &other_turn.session_key) + .await + .unwrap() + .unwrap() + .as_str(), + "acct_subagent_b" + ); + assert_eq!( + affinity + .load_alias(&provider_kind, &turn_key("seed-b")) + .await + .unwrap(), + Some(other_turn.clone()) + ); + let alias = affinity + .load_alias(&provider_kind, &turn_key("client-turn")) + .await + .unwrap() + .expect("client turn remains associated"); + assert_eq!(alias.session_key, session_key); + assert!(!alias.follow_only); + assert!( + affinity + .load_alias(&provider_kind, &turn_key("plugin-turn")) + .await + .unwrap() + .is_none() + ); + } + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn upstream_adapter_child_cannot_migrate_a_strict_session() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let server = MockServer::start().await; + let provider = provider_with_affinity_and_base_url( + &store, + Arc::new(MemorySessionAffinity::default()), + server.uri(), + ); + let request = |thread| { + planned_request( + "openai", + Operation::Generate(generate_with_session_context("root", Some(thread), None)), + ) + }; + drop( + provider + .clone() + .execute( + request("root"), + affinity_adapter_context("req_root", "key_a", None, None), + ) + .await + .unwrap(), + ); + create_account(&store, "acct_subagent_b").await; + store + .set_enabled(&ProviderAccountId::new("acct_subagent_a").unwrap(), false) + .await + .unwrap(); + let error = provider + .execute( + request("child"), + affinity_adapter_context("req_child", "key_a", None, None), + ) + .await + .err() + .expect("strict child cannot replace the session owner"); + assert!(error.retry_is_prohibited()); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn upstream_adapter_does_not_infer_sessions_from_unknown_protocols_or_missing_identity() { + for (protocol, body) in [ + ( + "example-source", + json!({"model":"gpt-5.4", "input":"hello", "session_id":"opaque-session"}), + ), + ("openai", json!({"model":"gpt-5.4", "input":"hello"})), + ] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_provider_contract").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 operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object(protocol, body.as_object().unwrap().clone()).unwrap(), + )); + drop( + provider + .execute( + planned_request("openai", operation), + affinity_adapter_context("req_unassociated", "key_a", None, None), + ) + .await + .unwrap(), + ); + assert_eq!(affinity.binding_count(), 0); + assert!(affinity.alias_ttls().is_empty()); + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + struct ConnectionProbe { selected: Arc>>>, } diff --git a/docs/architecture.md b/docs/architecture.md index 603d5537e..a23f37c2b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -727,6 +727,8 @@ Store 只保存带版本的绑定事实,Provider 在取得租约后、发送 官方独立 Images 仅携带 `turn_id`,通过已观测 Responses 轮次关联到会话,关联同时保留根线程或后代线程的调度权限。 轮次关联保存不透明会话键,不复制账号;显式会话身份只用于校验轮次归属。 绑定和轮次关联按同一请求策略中的可配置时长保留,均按 Client Key 隔离。 +受管上游适配器收到 OpenAI 协议请求时,按进入 attempt 中间件前的会话事实选号,并在成功准入后记录轮次关联。 +插件改写的正文与头部原样交给适配器,不重新解释为宿主会话绑定;其他输入协议不推测 OpenAI 会话身份。 成功准入按冻结时长续期,已关联的 Images 同时续期轮次关联;更新配置不扫描旧键,响应完成不续期。 未知轮次不推测会话归属,已有轮次不能改指其他会话或更改调度权限 @@ -754,7 +756,7 @@ Core 拥有容量池、上限与等待合同,OpenAI Provider 负责识别请 - OpenAI Provider 在选号前通过统一协议解析识别 `guardian`;原生发送与插件上游适配器使用相同分类。 请求头形成的协议上下文与正文 metadata 沿用编码器的既有优先顺序;其他协议不得仅凭同名字段使用审批池。 - 适配器不因此继承原生会话亲和、续接或 WebSocket 约束,选号后的 attempt 中间件也不追溯改变已取得的租约。 + 会话亲和沿用上述选号合同,适配器不继承原生续接或 WebSocket 约束;选号后的 attempt 中间件不追溯改变已取得的租约。 - `openaiGuardianReservedConcurrency = R > 0` 时,普通池使用账号原上限,审批池使用独立上限 R。 选号信号、负载评分、策略投影、请求容量观测和原子租约使用同一池的计数与上限,数值边界见 [运行设置 API](api.md#8-运行设置) - Store 按账号和容量池隔离租约,复用续期、取消与释放生命周期,不复制总计数。