mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(openai): preserve session affinity for managed upstream adapters (#508)
* fix(openai): preserve session affinity for managed upstream adapters * test(openai): avoid coupling wire invariance to adapter lookup count * fix(openai): preserve plugin control over adapter request identity Keep affinity and turn aliases tied to the admitted request before attempt middleware. Pass plugin rewrites to the adapter without native identity revalidation. Co-authored-by: Codex <noreply@openai.com> --------- Co-authored-by: Zyy <zyycn@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Ke Wang
Codex
Zyy
@@ -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
|
||||
|
||||
@@ -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<Option<Arc<dyn UpstreamAdapter>>, 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<std::sync::atomic::AtomicUsize>,
|
||||
@@ -170,6 +338,388 @@ async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_an
|
||||
|
||||
type SelectedConnection = Arc<dyn UpstreamAccountConnection>;
|
||||
|
||||
type AdapterInvocations = Arc<Mutex<Vec<UpstreamAdapterInvocation>>>;
|
||||
|
||||
struct AffinityAdapterProbe {
|
||||
invocations: Option<AdapterInvocations>,
|
||||
}
|
||||
|
||||
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<Option<Arc<dyn UpstreamAdapter>>, 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<Self>,
|
||||
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<Arc<dyn MiddlewarePlan>>,
|
||||
invocations: Option<AdapterInvocations>,
|
||||
) -> 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<Mutex<Option<oneshot::Sender<SelectedConnection>>>>,
|
||||
}
|
||||
|
||||
@@ -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 按账号和容量池隔离租约，复用续期、取消与释放生命周期，不复制总计数。
|
||||
|
||||
Reference in New Issue
Block a user