From 16aa08091cf1d5c2a3d4f12e416cc3424434e309 Mon Sep 17 00:00:00 2001 From: Ke Wang Date: Sat, 24 Oct 2026 20:20:47 +0800 Subject: [PATCH] feat(plugins): select upstream adapters by account (#509) * feat(plugins): select upstream adapters within account scopes * fix(plugins): defer adapter decisions until account selection Keep scoped adapters pending until an account is known so unavailable fallbacks cannot mask a healthy earlier match. Filter native xAI accounts when only an adapter can accept the request, preserving normal selection and cold-stream behavior. Cover fallback precedence and mixed-pool prewarm eligibility regressions. Co-authored-by: Codex --------- Co-authored-by: Zyy Co-authored-by: Codex --- backend/crates/gateway-core/src/engine/mod.rs | 39 +- .../src/engine/upstream_adapter.rs | 19 +- .../gateway-core/tests/engine/extensions.rs | 93 ++- .../runtime/plugin-host-compatibility.json | 2 +- .../src/adapter/upstream_adapter/execution.rs | 5 + .../src/adapter/upstream_adapter/mod.rs | 38 +- .../tests/generation/upstream_adapter.rs | 617 ++++++++++++------ .../runtime/tests/support/native.rs | 3 +- .../gateway-plugin/sdk/docs/manifest.md | 2 +- .../sdk/docs/upstream-adapters.md | 14 +- .../sdk/src/call/upstream_adapter/mod.rs | 2 + .../gateway-plugin/sdk/src/capability.rs | 2 +- .../sdk/tests/call/upstream_adapter/mod.rs | 24 +- .../gateway-plugin/sdk/tests/manifest.rs | 2 +- .../openai/src/credential/selector.rs | 30 + .../providers/openai/src/provider/mod.rs | 28 +- .../provider/contract/upstream_adapter.rs | 378 +++++++++-- .../providers/xai/src/credential/selector.rs | 18 + .../crates/providers/xai/src/provider/mod.rs | 174 +++-- .../providers/xai/src/transport/session.rs | 4 + .../xai/tests/provider/contract/mod.rs | 48 +- .../xai/tests/provider/contract/scheduling.rs | 284 ++++++++ 22 files changed, 1466 insertions(+), 360 deletions(-) diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 0709840e5..d338986bf 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -575,11 +575,48 @@ impl AttemptContext { &self, provider: &ProviderKind, model: &UpstreamModelId, + ) -> Result { + let account = self + .required_account() + .or_else(|| match self.continuation_attempt() { + ContinuationAttempt::Native => self + .continuation() + .and_then(continuation::ContinuationBinding::pinned) + .map(|pin| pin.account()), + ContinuationAttempt::ReplayOwner => self + .account_state_owner() + .filter(|owner| owner.provider() == provider) + .map(|owner| owner.account()), + ContinuationAttempt::None | ContinuationAttempt::ReplayAny => None, + }); + self.request.upstream_adapters.as_ref().map_or( + Ok(upstream_adapter::UpstreamAdapterSelection::Native), + |plan| plan.select(self, provider, model, account), + ) + } + + /// 账号候选资格检查和最终执行解析使用同一冻结计划，不调用插件进程 + pub fn upstream_adapter_for_account( + &self, + provider: &ProviderKind, + model: &UpstreamModelId, + account: &ProviderAccountId, ) -> Result>, ProviderError> { self.request .upstream_adapters .as_ref() - .map_or(Ok(None), |plan| plan.select(self, provider, model)) + .map_or(Ok(None), |plan| { + match plan.select(self, provider, model, Some(account))? { + upstream_adapter::UpstreamAdapterSelection::Native => Ok(None), + upstream_adapter::UpstreamAdapterSelection::Adapter(adapter) => { + Ok(Some(adapter)) + } + // 已知账号的解析必须给出最终执行端，不能绕过资格检查进入原生路径 + upstream_adapter::UpstreamAdapterSelection::PendingAccount => Err( + ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::NotSent), + ), + } + }) } #[must_use] diff --git a/backend/crates/gateway-core/src/engine/upstream_adapter.rs b/backend/crates/gateway-core/src/engine/upstream_adapter.rs index 836dc1f02..43518f5bf 100644 --- a/backend/crates/gateway-core/src/engine/upstream_adapter.rs +++ b/backend/crates/gateway-core/src/engine/upstream_adapter.rs @@ -31,7 +31,16 @@ pub trait UpstreamAdapterPlan: fmt::Debug + Send + Sync { context: &AttemptContext, provider: &ProviderKind, model: &UpstreamModelId, - ) -> Result>, ProviderError>; + account: Option<&ProviderAccountId>, + ) -> Result; +} + +/// 未确定账号时不能把账号范围声明视为原生执行或继续命中后置兜底 +#[derive(Debug)] +pub enum UpstreamAdapterSelection { + Native, + Adapter(Arc), + PendingAccount, } /// 请求同时持有计划及发布集合，防止旧集合先排空进程而中断在途调用 @@ -58,8 +67,9 @@ impl FrozenUpstreamAdapterPlan { context: &AttemptContext, provider: &ProviderKind, model: &UpstreamModelId, - ) -> Result>, ProviderError> { - self.plan.select(context, provider, model) + account: Option<&ProviderAccountId>, + ) -> Result { + self.plan.select(context, provider, model, account) } } @@ -71,6 +81,9 @@ pub trait UpstreamAdapter: fmt::Debug + Send + Sync { /// 注册时声明的上游传输，用于宿主登记 attempt fn transport(&self) -> &str; + + /// 选号前复核注册的认证资格，发送前仍由执行端校验当前凭据 + fn supports_authentication(&self, kind: &str) -> bool; } /// 一次执行的宿主事实；账号及 Provider 归属不可由插件返回值覆盖 diff --git a/backend/crates/gateway-core/tests/engine/extensions.rs b/backend/crates/gateway-core/tests/engine/extensions.rs index 4fead6150..2afa34a03 100644 --- a/backend/crates/gateway-core/tests/engine/extensions.rs +++ b/backend/crates/gateway-core/tests/engine/extensions.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use gateway_core::engine::extensions::{ExecutionExtensionIndex, ExecutionExtensionPlans}; -use gateway_core::engine::upstream_adapter::{UpstreamAdapter, UpstreamAdapterPlan}; +use gateway_core::engine::upstream_adapter::{UpstreamAdapterPlan, UpstreamAdapterSelection}; use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; use gateway_core::{ engine::AttemptContext, error::ProviderError, identity::ProviderKind, routing::UpstreamModelId, @@ -31,8 +31,9 @@ impl UpstreamAdapterPlan for EmptyAdapterPlan { _: &AttemptContext, _: &ProviderKind, _: &UpstreamModelId, - ) -> Result>, ProviderError> { - Ok(None) + _: Option<&gateway_core::account::ProviderAccountId>, + ) -> Result { + Ok(UpstreamAdapterSelection::Native) } } @@ -89,3 +90,89 @@ fn execution_index_does_not_own_registered_generations() { assert!(index.middleware(&reference).is_none()); assert!(index.upstream_adapters(&reference).is_none()); } + +#[derive(Debug, Default)] +struct AccountSelectionProbe(std::sync::Mutex>>); + +impl UpstreamAdapterPlan for AccountSelectionProbe { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + account: Option<&gateway_core::account::ProviderAccountId>, + ) -> Result { + self.0 + .lock() + .unwrap() + .push(account.map(|account| account.as_str().into())); + Ok(UpstreamAdapterSelection::Native) + } +} + +#[test] +fn adapter_preselection_follows_only_the_current_attempt_account_owner() { + use gateway_core::{ + account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy}, + engine::continuation::{ContinuationBinding, NativeContinuationPin, PreviousResponseId}, + engine::upstream_adapter::FrozenUpstreamAdapterPlan, + engine::{ + AccountAttemptContext, ContinuationAttempt, ModelRequestId, RequestAttemptContext, + }, + lifecycle::CancellationToken, + policy::ClientApiKeyId, + }; + use std::{ + num::NonZeroU32, + time::{Duration, SystemTime}, + }; + + let provider = ProviderKind::new("openai").unwrap(); + let model = UpstreamModelId::new("model").unwrap(); + let key = ClientApiKeyId::new("key").unwrap(); + let pin = NativeContinuationPin::new( + PreviousResponseId::new("public"), + PreviousResponseId::new("upstream"), + key.clone(), + provider.clone(), + ProviderAccountId::new("acct_pinned").unwrap(), + ); + for (attempt, required, expected) in [ + (ContinuationAttempt::Native, None, Some("acct_pinned")), + (ContinuationAttempt::ReplayAny, None, None), + (ContinuationAttempt::None, None, None), + ( + ContinuationAttempt::Native, + Some("acct_required"), + Some("acct_required"), + ), + ] { + let probe = Arc::new(AccountSelectionProbe::default()); + let context = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_account_plan").unwrap(), + key.clone(), + ) + .with_upstream_adapters(Some(FrozenUpstreamAdapterPlan::new( + probe.clone(), + reference("account-plan"), + ))), + NonZeroU32::MIN, + SystemTime::now() + Duration::from_secs(5), + AccountSelectionPolicy::new(RotationStrategy::RoundRobin, NonZeroU32::MIN), + AccountAttemptContext::new( + Default::default(), + required.map(|id| ProviderAccountId::new(id).unwrap()), + None, + ), + Some(ContinuationBinding::Pinned(pin.clone())), + CancellationToken::new(), + ) + .with_continuation_attempt(attempt); + assert!(matches!( + context.upstream_adapter(&provider, &model).unwrap(), + UpstreamAdapterSelection::Native + )); + assert_eq!(*probe.0.lock().unwrap(), vec![expected.map(str::to_owned)]); + } +} diff --git a/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json b/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json index 4c26ac476..140e18855 100644 --- a/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json +++ b/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json @@ -9,7 +9,7 @@ { "capability": "model_catalog", "versions": [1] }, { "capability": "retry_policy", "versions": [2] }, { "capability": "middleware", "versions": [5] }, - { "capability": "upstream_adapter", "versions": [4] }, + { "capability": "upstream_adapter", "versions": [5] }, { "capability": "observer", "versions": [2] }, { "capability": "command_line", "versions": [1] }, { "capability": "management", "versions": [2] }, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs index d5eb382b4..9fc2a2b30 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs @@ -152,6 +152,11 @@ impl Execution { if invocation.operation.protocol() != declaration.protocol || invocation.metadata.provider().as_str() != declaration.provider.as_str() || invocation.metadata.provider_account_id() != invocation.account.account_id() + || (!declaration.account_ids.is_empty() + && !declaration + .account_ids + .iter() + .any(|id| id == invocation.account.account_id().as_str())) || !declaration .authentication_kinds .iter() diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/mod.rs index b95566b68..078a66647 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/mod.rs @@ -1,4 +1,4 @@ -//! 将插件上游适配贡献编译为可按请求选择的执行计划 +//! 将插件上游适配贡献编译为按请求与账号范围解析的执行计划 mod event; mod execution; @@ -10,9 +10,13 @@ pub(crate) mod target; use std::{fmt, sync::Arc}; use gateway_core::{ + account::ProviderAccountId, engine::{ AttemptContext, - upstream_adapter::{UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan}, + upstream_adapter::{ + UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan, + UpstreamAdapterSelection, + }, }, error::{ProviderError, ProviderErrorKind}, identity::ProviderKind, @@ -58,6 +62,13 @@ impl UpstreamAdapter for PluginUpstreamAdapter { self.declaration.transport.as_str() } + fn supports_authentication(&self, kind: &str) -> bool { + self.declaration + .authentication_kinds + .iter() + .any(|candidate| candidate == kind) + } + fn execute( self: Arc, invocation: UpstreamAdapterInvocation, @@ -91,7 +102,8 @@ impl UpstreamAdapterPlan for PluginUpstreamAdapterPlan { context: &AttemptContext, provider: &ProviderKind, model: &UpstreamModelId, - ) -> Result>, ProviderError> { + account: Option<&ProviderAccountId>, + ) -> Result { let model = context .requested_model() .map_or(model.as_str(), |model| model.as_str()); @@ -115,13 +127,29 @@ impl UpstreamAdapterPlan for PluginUpstreamAdapterPlan { .iter() .any(|candidate| candidate == model)) { + if !adapter.declaration.account_ids.is_empty() { + let Some(account) = account else { + // 后置声明只覆盖未命中此账号范围的候选，不能提前以其故障拒绝整个请求 + return Ok(UpstreamAdapterSelection::PendingAccount); + }; + if !adapter + .declaration + .account_ids + .iter() + .any(|id| id == account.as_str()) + { + continue; + } + } if !adapter.session.is_ready() { return Err(unavailable()); } - return Ok(Some(Arc::clone(adapter) as Arc)); + return Ok(UpstreamAdapterSelection::Adapter( + Arc::clone(adapter) as Arc + )); } } - Ok(None) + Ok(UpstreamAdapterSelection::Native) } } diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs index c546995e4..e4ac2e4b5 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs @@ -18,7 +18,9 @@ use gateway_core::{ AccountAttemptContext, AttemptContext, ModelRequestId, RequestAttemptContext, middleware::MiddlewareHeader, provider::{EventStream, ProviderCallMetadata}, - upstream_adapter::{UpstreamAccountConnection, UpstreamAdapterInvocation}, + upstream_adapter::{ + UpstreamAccountConnection, UpstreamAdapterInvocation, UpstreamAdapterSelection, + }, }, error::{ProviderError, ProviderErrorKind}, event::GatewayEvent, @@ -86,7 +88,7 @@ impl UpstreamAccountConnection for Account { fn configuration(base: &str, streaming: bool) -> Value { json!({ - "upstream_registration":{"adapters":[{"id":"example","provider":"openai","base_url":format!("{base}/adapter/"),"paths":[{"path":"responses","purpose":"inference"}],"authentication_kinds":["oauth"],"transport":if streaming {"http_sse"} else {"http_json"},"protocol":"openai"}]}, + "upstream_registration":{"adapters":[{"id":"example","provider":"openai","base_url":format!("{base}/adapter/"),"paths":[{"path":"responses","purpose":"inference"}],"authentication_kinds":["oauth"],"account_ids":[],"transport":if streaming {"http_sse"} else {"http_json"},"protocol":"openai"}]}, "upstream_callbacks":[{"method":if streaming {"host.upstream.http.do_stream"} else {"host.upstream.http.do"},"params":{"method":"POST","path":"responses","headers":[["content-type",b"application/json".to_vec()]]},"body":"request"}], "upstream_events":[ {"event":{"facts":[{"type":"started","id":"response-one","model":"native-model"}]}}, @@ -113,12 +115,13 @@ async fn setup_with_limits( Arc, gateway_plugin_runtime::PluginRuntime, ) { - let contributes = Contributions::from([crate::support::contribution( + let (capability, contribution) = crate::support::contribution( Capability::UpstreamAdapter, vec![Stage::Upstream], vec!["openai".into()], vec!["openai".into()], - )]); + ); + let contributes = Contributions::from([(capability, contribution)]); let (cache, store, runtime) = super::setup_with_limits(contributes, Default::default(), limits).await; { @@ -209,7 +212,7 @@ fn execute_with_context( let provider = ProviderKind::new("openai").unwrap(); let model = UpstreamModelId::new("native-model").unwrap(); let adapter = context - .upstream_adapter(&provider, &model) + .upstream_adapter_for_account(&provider, &model, &account.id) .unwrap() .unwrap(); let metadata = ProviderCallMetadata::new( @@ -240,6 +243,144 @@ fn execute_with_context( }) } +#[tokio::test] +async fn account_scoped_preselection_keeps_fallback_failures_within_matching_accounts() { + for fallback_first in [false, true] { + let mut config = configuration("https://example.invalid", false); + config["upstream_registration"]["adapters"][0]["account_ids"] = json!(["acct_one"]); + let (cache, store, runtime) = setup(config).await; + { + let mut snapshot = store.snapshot.lock().unwrap(); + let mut fallback = snapshot.instances[0].clone(); + fallback.id = "instance-two".into(); + fallback.name = "Fallback".into(); + fallback.configuration = configuration("https://fallback.invalid", false); + fallback.configuration["upstream_registration"] = json!(false); + if fallback_first { + snapshot.instances.insert(0, fallback); + } else { + snapshot.instances.push(fallback); + } + } + let generation = + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + { + let context = context(&runtime, &generation, "key-one"); + let provider = ProviderKind::new("openai").unwrap(); + let model = UpstreamModelId::new("native-model").unwrap(); + let preselection = context.upstream_adapter(&provider, &model); + let selected = context.upstream_adapter_for_account( + &provider, + &model, + &ProviderAccountId::new("acct_one").unwrap(), + ); + if !fallback_first { + assert!(matches!( + preselection.unwrap(), + UpstreamAdapterSelection::PendingAccount + )); + assert!(selected.unwrap().is_some()); + } else { + assert_eq!( + preselection.unwrap_err().kind(), + ProviderErrorKind::Unavailable + ); + assert_eq!(selected.unwrap_err().kind(), ProviderErrorKind::Unavailable); + } + assert_eq!( + context + .upstream_adapter_for_account( + &provider, + &model, + &ProviderAccountId::new("acct_native").unwrap() + ) + .unwrap_err() + .kind(), + ProviderErrorKind::Unavailable + ); + } + drop(generation); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; + } +} + +#[tokio::test] +async fn account_scoped_adapters_share_a_model_without_sharing_their_targets() { + let first = MockServer::start().await; + let second = MockServer::start().await; + for server in [&first, &second] { + Mock::given(method("POST")) + .and(path("/adapter/responses")) + .respond_with(ResponseTemplate::new(200).set_body_string("{}")) + .expect(1) + .mount(server) + .await; + } + let mut config = configuration(&first.uri(), false); + let mut second_declaration = config["upstream_registration"]["adapters"][0].clone(); + config["upstream_registration"]["adapters"][0]["account_ids"] = json!(["acct_one"]); + second_declaration["id"] = json!("second"); + second_declaration["base_url"] = json!(format!("{}/adapter/", second.uri())); + second_declaration["account_ids"] = json!(["acct_two"]); + config["upstream_registration"]["adapters"] + .as_array_mut() + .unwrap() + .push(second_declaration); + let (cache, _, runtime) = setup(config).await; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + let provider = ProviderKind::new("openai").unwrap(); + let model = UpstreamModelId::new("native-model").unwrap(); + { + let context = context(&runtime, &generation, "key-one"); + assert!(matches!( + context.upstream_adapter(&provider, &model).unwrap(), + UpstreamAdapterSelection::PendingAccount + )); + assert!( + context + .upstream_adapter_for_account( + &provider, + &model, + &ProviderAccountId::new("acct_native").unwrap() + ) + .unwrap() + .is_none() + ); + } + for (account, expected) in [("acct_one", "example"), ("acct_two", "second")] { + let mut stream = execute( + &runtime, + &generation, + Account::with_identity(account, 1), + "key-one", + None, + ); + let mut completed = 0; + let mut adapter = None; + while let Some(event) = stream.next().await { + let event = event.unwrap(); + completed += event + .canonical_facts() + .iter() + .filter(|fact| matches!(fact, GatewayEvent::Completed(_))) + .count(); + if let Some(state) = event.session_update() { + adapter = Some(state.extension_owner().unwrap().adapter_id.clone()); + } + } + assert_eq!(completed, 1); + assert_eq!(adapter.as_deref(), Some(expected)); + } + drop(generation); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; +} + #[tokio::test] async fn adapter_preserves_three_state_fast_in_request_metadata() { use gateway_core::account::FastMode; @@ -852,32 +993,38 @@ async fn adapter_does_not_publish_completed_when_rpc_fails_after_terminal_frame( #[tokio::test] async fn overlapping_adapter_declarations_keep_first_match_without_blocking_preparation() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/adapter/responses")) - .respond_with(ResponseTemplate::new(200).set_body_string("{}")) - .expect(1) - .mount(&server) - .await; - let mut config = configuration(&server.uri(), false); - let mut duplicate = config["upstream_registration"]["adapters"][0].clone(); - duplicate["base_url"] = json!("https://unused.invalid/adapter/"); - config["upstream_registration"]["adapters"] - .as_array_mut() - .unwrap() - .push(duplicate); - let (cache, _, runtime) = setup(config).await; - let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); - while let Some(event) = stream.next().await { - event.unwrap(); + for scoped in [false, true] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/adapter/responses")) + .respond_with(ResponseTemplate::new(200).set_body_string("{}")) + .expect(1) + .mount(&server) + .await; + let mut config = configuration(&server.uri(), false); + if scoped { + config["upstream_registration"]["adapters"][0]["account_ids"] = json!(["acct_one"]); + } + let mut duplicate = config["upstream_registration"]["adapters"][0].clone(); + duplicate["base_url"] = json!("https://unused.invalid/adapter/"); + config["upstream_registration"]["adapters"] + .as_array_mut() + .unwrap() + .push(duplicate); + let (cache, _, runtime) = setup(config).await; + let generation = + ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + let mut stream = execute(&runtime, &generation, Account::new(), "key-one", None); + while let Some(event) = stream.next().await { + event.unwrap(); + } + drop(stream); + drop(generation); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; } - drop(stream); - drop(generation); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; } #[tokio::test] @@ -912,6 +1059,78 @@ async fn adapter_registration_preserves_plugin_targets_with_typed_provider_decod runtime.shutdown().await; } +#[tokio::test] +async fn adapter_registration_requires_explicit_account_scope() { + let mut config = configuration("https://example.invalid", false); + config["upstream_registration"]["adapters"][0] + .as_object_mut() + .unwrap() + .remove("account_ids"); + let (cache, store, runtime) = setup(config).await; + let snapshot = store.snapshot.lock().unwrap().clone(); + assert!( + gateway_admin::ports::plugins::PluginPreparation::prepare(&runtime, snapshot) + .await + .is_err() + ); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn account_scoped_adapter_rechecks_the_invocation_account_before_sending() { + let server = MockServer::start().await; + let mut config = configuration(&server.uri(), false); + config["upstream_registration"]["adapters"][0]["account_ids"] = json!(["acct_one"]); + let (cache, _, runtime) = setup(config).await; + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + let context = context(&runtime, &generation, "key-one"); + let provider = ProviderKind::new("openai").unwrap(); + let model = UpstreamModelId::new("native-model").unwrap(); + let adapter = context + .upstream_adapter_for_account( + &provider, + &model, + &ProviderAccountId::new("acct_one").unwrap(), + ) + .unwrap() + .unwrap(); + let account = Account::with_identity("acct_other", 1); + let metadata = ProviderCallMetadata::new( + provider, + model, + account.id.clone(), + UpstreamTransport::new(adapter.transport()).unwrap(), + ); + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + json!({"model":"native-model","input":"hello"}) + .as_object() + .unwrap() + .clone(), + ) + .unwrap(), + )); + let mut stream = adapter.execute(UpstreamAdapterInvocation { + operation, + headers: vec![], + context, + metadata, + account: account.clone(), + }); + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + assert_eq!(account.failures.load(Ordering::SeqCst), 1); + assert!(server.received_requests().await.unwrap().is_empty()); + drop(stream); + drop(generation); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; +} + #[tokio::test] #[expect( clippy::result_large_err, @@ -1172,177 +1391,191 @@ async fn usage_without_response_model_still_uses_native_account_pricing() { #[tokio::test] async fn core_published_adapter_uses_normal_attempt_usage_and_cost_ledger() { use crate::support::environment::Environment; - let Some(mut environment) = Environment::create_command().await else { - eprintln!("SKIP: plugin integration environment absent"); - return; - }; - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(header("authorization", "Bearer fixture-native-token")) - .respond_with(ResponseTemplate::new(200).set_body_string("{}")) - .expect(1) - .mount(&server) - .await; - let account = environment.account(None).await; - let group = environment.account_group_with_account(&account).await; - let key_id = "key_adapter_ledger"; - environment - .client_key_with_limits_and_groups( - key_id, - "sk-adapter-ledger-fixture", - gateway_core::policy::RateLimits::unlimited(), - vec![group], - ) - .await; - environment - .install_plugin(configuration(&server.uri(), false)) - .await; - environment.install_plugin(json!({ - "plugin_id":"test.adapter-ledger-caller", - "command_registration":{"commands":[{"name":"run","description":"执行模型"}]}, - "command_nested_model_fixture":{ - "request":{"client_key_id":key_id,"model":crate::support::native::MODEL,"protocol":"openai","operation":"generate","provider":"openai","account_id":account.as_str()}, - "body":{"model":crate::support::native::MODEL,"input":"adapter ledger"} - }, - "command_result":{"stdout":"ok","stderr":"","exit_code":0} - })).await; - let (runtime, core) = environment.command_plane().await; - let snapshot = environment - .store - .admin_ports() - .plugins() - .load_instances() - .await - .unwrap(); - let caller = snapshot - .instances - .iter() - .find(|instance| instance.configuration.get("command_registration").is_some()) - .unwrap(); - let commands = runtime.prepare_command_line().await.unwrap(); - environment.store.start_command_line_writes().unwrap(); - let result = commands.execute(&caller.id, "run", &[]).await.unwrap(); - assert_eq!(result.stdout, "ok"); - commands.shutdown().await; - environment - .store - .shutdown_command_line_writes() - .await - .unwrap(); - environment - .assert_upstream_accounting(key_id, account.as_str()) - .await; - drop(core); - drop(runtime); - environment.close().await; + for scoped in [false, true] { + let Some(mut environment) = Environment::create_command().await else { + eprintln!("SKIP: plugin integration environment absent"); + return; + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(header("authorization", "Bearer fixture-native-token")) + .respond_with(ResponseTemplate::new(200).set_body_string("{}")) + .expect(1) + .mount(&server) + .await; + let account = environment.account(None).await; + let group = environment.account_group_with_account(&account).await; + let key_id = "key_adapter_ledger"; + environment + .client_key_with_limits_and_groups( + key_id, + "sk-adapter-ledger-fixture", + gateway_core::policy::RateLimits::unlimited(), + vec![group], + ) + .await; + let mut adapter_config = configuration(&server.uri(), false); + if scoped { + adapter_config["upstream_registration"]["adapters"][0]["account_ids"] = + json!([account.as_str()]); + } + environment.install_plugin(adapter_config).await; + environment.install_plugin(json!({ + "plugin_id":"test.adapter-ledger-caller", + "command_registration":{"commands":[{"name":"run","description":"执行模型"}]}, + "command_nested_model_fixture":{ + "request":{"client_key_id":key_id,"model":crate::support::native::MODEL,"protocol":"openai","operation":"generate","provider":"openai","account_id":account.as_str()}, + "body":{"model":crate::support::native::MODEL,"input":"adapter ledger"} + }, + "command_result":{"stdout":"ok","stderr":"","exit_code":0} + })).await; + let (runtime, core) = environment.command_plane().await; + let snapshot = environment + .store + .admin_ports() + .plugins() + .load_instances() + .await + .unwrap(); + let caller = snapshot + .instances + .iter() + .find(|instance| instance.configuration.get("command_registration").is_some()) + .unwrap(); + let commands = runtime.prepare_command_line().await.unwrap(); + environment.store.start_command_line_writes().unwrap(); + let result = commands.execute(&caller.id, "run", &[]).await.unwrap(); + assert_eq!(result.stdout, "ok"); + commands.shutdown().await; + environment + .store + .shutdown_command_line_writes() + .await + .unwrap(); + environment + .assert_upstream_accounting(key_id, account.as_str()) + .await; + drop(core); + drop(runtime); + environment.close().await; + } } #[tokio::test] async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_and_reap_workers() { - let original = MockServer::start().await; - let replacement = MockServer::start().await; - for (server, expected) in [(&original, 2), (&replacement, 1)] { - Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(200).set_body_string("{}")) - .expect(expected) - .mount(server) - .await; - } - let (cache, store, runtime) = setup(configuration(&original.uri(), false)).await; - let old = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) - .await - .unwrap(); - let account = Account::new(); - let mut old_stream = execute(&runtime, &old, account.clone(), "key-one", None); - old_stream.next().await.unwrap().unwrap(); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = gateway_admin::model::Revision::new(2).unwrap(); - snapshot.instances[0].revision = snapshot.config_revision; - snapshot.instances[0].configuration = configuration(&replacement.uri(), false); - } - let new = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) - .await - .unwrap(); - assert_ne!(old.id(), new.id()); - let mut new_stream = execute(&runtime, &new, account.clone(), "key-one", None); - new_stream.next().await.unwrap().unwrap(); - drop(old); - let old_state = old_stream - .next() - .await - .unwrap() - .unwrap() - .take_session_update() - .unwrap(); - assert!(old_stream.next().await.is_none()); - drop(old_stream); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = gateway_admin::model::Revision::new(3).unwrap(); - snapshot.instances[0].revision = snapshot.config_revision; - snapshot.instances[0].enabled = false; - } - let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(3).unwrap()) - .await - .unwrap(); - assert!( - context(&runtime, &disabled, "key-one") - .upstream_adapter( - &ProviderKind::new("openai").unwrap(), - &UpstreamModelId::new("native-model").unwrap() - ) - .unwrap() - .is_none() - ); - drop(new); - assert!( - new_stream + for scoped in [false, true] { + let scoped_configuration = |base: &str| { + let mut config = configuration(base, false); + if scoped { + config["upstream_registration"]["adapters"][0]["account_ids"] = json!(["acct_one"]); + } + config + }; + let original = MockServer::start().await; + let replacement = MockServer::start().await; + for (server, expected) in [(&original, 2), (&replacement, 1)] { + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_string("{}")) + .expect(expected) + .mount(server) + .await; + } + let (cache, store, runtime) = setup(scoped_configuration(&original.uri())).await; + let old = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + let account = Account::new(); + let mut old_stream = execute(&runtime, &old, account.clone(), "key-one", None); + old_stream.next().await.unwrap().unwrap(); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = gateway_admin::model::Revision::new(2).unwrap(); + snapshot.instances[0].revision = snapshot.config_revision; + snapshot.instances[0].configuration = scoped_configuration(&replacement.uri()); + } + let new = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap()) + .await + .unwrap(); + assert_ne!(old.id(), new.id()); + let mut new_stream = execute(&runtime, &new, account.clone(), "key-one", None); + new_stream.next().await.unwrap().unwrap(); + drop(old); + let old_state = old_stream .next() .await .unwrap() .unwrap() - .canonical_facts() - .iter() - .any(|fact| matches!(fact, GatewayEvent::Completed(_))) - ); - assert!(new_stream.next().await.is_none()); - drop(new_stream); - drop(disabled); - super::wait_until_empty(cache.path()).await; - assert_eq!(Arc::strong_count(&account), 1); - { - let mut snapshot = store.snapshot.lock().unwrap(); - snapshot.config_revision = gateway_admin::model::Revision::new(4).unwrap(); - snapshot.instances[0].revision = snapshot.config_revision; - snapshot.instances[0].enabled = true; - snapshot.instances[0].configuration = configuration(&original.uri(), false); - } - let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(4).unwrap()) + .take_session_update() + .unwrap(); + assert!(old_stream.next().await.is_none()); + drop(old_stream); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = gateway_admin::model::Revision::new(3).unwrap(); + snapshot.instances[0].revision = snapshot.config_revision; + snapshot.instances[0].enabled = false; + } + let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(3).unwrap()) + .await + .unwrap(); + assert!(matches!( + context(&runtime, &disabled, "key-one") + .upstream_adapter( + &ProviderKind::new("openai").unwrap(), + &UpstreamModelId::new("native-model").unwrap() + ) + .unwrap(), + UpstreamAdapterSelection::Native + )); + drop(new); + assert!( + new_stream + .next() + .await + .unwrap() + .unwrap() + .canonical_facts() + .iter() + .any(|fact| matches!(fact, GatewayEvent::Completed(_))) + ); + assert!(new_stream.next().await.is_none()); + drop(new_stream); + drop(disabled); + super::wait_until_empty(cache.path()).await; + assert_eq!(Arc::strong_count(&account), 1); + { + let mut snapshot = store.snapshot.lock().unwrap(); + snapshot.config_revision = gateway_admin::model::Revision::new(4).unwrap(); + snapshot.instances[0].revision = snapshot.config_revision; + snapshot.instances[0].enabled = true; + snapshot.instances[0].configuration = scoped_configuration(&original.uri()); + } + let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(4).unwrap()) + .await + .unwrap(); + // 恢复同一配置仍是新代次，不能复活旧进程中的续接身份 + let error = execute( + &runtime, + &restored, + account.clone(), + "key-one", + Some(old_state), + ) + .next() .await - .unwrap(); - // 恢复同一配置仍是新代次，不能复活旧进程中的续接身份 - let error = execute( - &runtime, - &restored, - account.clone(), - "key-one", - Some(old_state), - ) - .next() - .await - .unwrap() - .unwrap_err(); - assert_eq!(error.send_state(), UpstreamSendState::NotSent); - let mut fresh = execute(&runtime, &restored, account.clone(), "key-one", None); - while let Some(event) = fresh.next().await { - event.unwrap(); + .unwrap() + .unwrap_err(); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + let mut fresh = execute(&runtime, &restored, account.clone(), "key-one", None); + while let Some(event) = fresh.next().await { + event.unwrap(); + } + drop(fresh); + drop(restored); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; + assert_eq!(Arc::strong_count(&account), 1); } - drop(fresh); - drop(restored); - runtime.shutdown().await; - super::wait_until_empty(cache.path()).await; - assert_eq!(Arc::strong_count(&account), 1); } #[tokio::test] @@ -1514,7 +1747,7 @@ async fn large_adapter_requests_use_the_frozen_budget_for_both_http_capabilities let provider = ProviderKind::new("openai").unwrap(); let model = UpstreamModelId::new("native-model").unwrap(); let adapter = context - .upstream_adapter(&provider, &model) + .upstream_adapter_for_account(&provider, &model, &account.id) .unwrap() .unwrap(); let metadata = ProviderCallMetadata::new( diff --git a/backend/crates/gateway-plugin/runtime/tests/support/native.rs b/backend/crates/gateway-plugin/runtime/tests/support/native.rs index b4b7deebb..2e1eb750f 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/native.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/native.rs @@ -75,7 +75,8 @@ impl Provider for NativeProvider { let model = candidate.upstream_model().cloned().ok_or_else(|| { ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) })?; - let adapter = context.upstream_adapter(candidate.provider(), &model)?; + let adapter = + context.upstream_adapter_for_account(candidate.provider(), &model, &account)?; let metadata = ProviderCallMetadata::new( candidate.provider().clone(), model, diff --git a/backend/crates/gateway-plugin/sdk/docs/manifest.md b/backend/crates/gateway-plugin/sdk/docs/manifest.md index e57eebb9d..242f794a0 100644 --- a/backend/crates/gateway-plugin/sdk/docs/manifest.md +++ b/backend/crates/gateway-plugin/sdk/docs/manifest.md @@ -55,7 +55,7 @@ RPC v5 的帧布局、编码与会话生命周期见[通信边界](../README.md# | --- | --- | | `middleware` | `5` | | `scheduler` | `3` | -| `upstream_adapter` | `4` | +| `upstream_adapter` | `5` | | `frontend_authentication`、`model_router`、`retry_policy`、`observer`、`management` | `2` | | `model_catalog`、`command_line`、`maintenance` | `1` | diff --git a/backend/crates/gateway-plugin/sdk/docs/upstream-adapters.md b/backend/crates/gateway-plugin/sdk/docs/upstream-adapters.md index a603fb9f5..a040d58b4 100644 --- a/backend/crates/gateway-plugin/sdk/docs/upstream-adapters.md +++ b/backend/crates/gateway-plugin/sdk/docs/upstream-adapters.md @@ -2,13 +2,13 @@ [返回 SDK](../README.md) · [清单声明](manifest.md#扩展项简写) · [能力与回调](capabilities.md) -`upstream_adapter` v4 是现有生成请求洋葱链的终端，只挂到内置 `openai` / `xai`，不替换独立图像或压缩端点。 +`upstream_adapter` v5 是现有生成请求洋葱链的终端，只挂到内置 `openai` / `xai`，不替换独立图像或压缩端点。 插件定义 Base URL、业务路径、请求编码和响应解析；宿主继续持有身份、ACL、选号、租约、凭据、代理、重试和账本。 普通请求或响应加工使用 [middleware](capabilities.md#洋葱中间件) ## 声明与绑定 -清单声明 `upstream_adapter`、版本 `4`，以及匹配的 `inputFormats` / `outputFormats`；固定阶段 `upstream` 由作者清单规范化入口生成。 +清单声明 `upstream_adapter`、版本 `5`，以及匹配的 `inputFormats` / `outputFormats`；固定阶段 `upstream` 由作者清单规范化入口生成。 使用 `PluginBuilder::on` 注册 `methods::UPSTREAM_ADAPTER_REGISTER` 与 `methods::UPSTREAM_ADAPTER_EXECUTE` @@ -24,11 +24,21 @@ | `transport` | `http_json`、`http_sse` 或 `websocket` | | `protocol` | 请求和客户端 wire 使用的协议 | | `models` | 公开请求模型的精确匹配，空集合沿用绑定范围 | +| `account_ids` | 上游账号 ID 范围，必填；空数组匹配所有账号 | 实例绑定沿用 Key、账号组、Provider、模型范围，只允许一次 `upstream` 绑定和 `reject` 故障策略。 多个声明命中时按发布顺序选择首个适配器，插件负责声明的业务含义。适配器未命中时仍使用原生上游， 已命中但不可用时拒绝请求。安装、配置、启停和升级使用已有插件管理 +账号范围只收窄绑定，不能扩大 Key、分组或模型权限。宿主先排除适配器不可用或认证类型不兼容的账号， +再沿用原调度、亲和和租约；取得账号后从同一发布代次解析适配器。不同账号可以为同一公开模型选择不同适配器， +未匹配账号保持原生执行，重叠声明仍按发布顺序取首个匹配。 +前置账号范围已匹配的可用适配器不受后置兜底声明故障影响。 +xAI 原生编码不支持的请求仍可选择匹配的适配器账号；原生账号会在取得租约前被排除 + +账号范围由插件注册结果提供，不改变管理端 binding 中 `accountGroupIds` 的请求作用域语义。 +停用或卸载实例会移除其绑定；维护时须先停用依赖该适配器的账号并排空请求，账号范围不提供永久依赖约束。 + 能力、清单、进程 RPC 和宿主兼容声明分别校验各自版本，见[清单](manifest.md)。 Provider、业务协议和传输名称不能写入宿主兼容清单的 capability / RPC 版本位置 diff --git a/backend/crates/gateway-plugin/sdk/src/call/upstream_adapter/mod.rs b/backend/crates/gateway-plugin/sdk/src/call/upstream_adapter/mod.rs index 1d13612a3..a39e073fc 100644 --- a/backend/crates/gateway-plugin/sdk/src/call/upstream_adapter/mod.rs +++ b/backend/crates/gateway-plugin/sdk/src/call/upstream_adapter/mod.rs @@ -66,6 +66,8 @@ pub struct UpstreamAdapterDeclaration { /// 公开请求模型的匹配范围；空集合沿用功能绑定范围 #[serde(default)] pub models: Vec, + /// 显式上游账号范围，空集合匹配所有账号，不扩大请求的账号权限 + pub account_ids: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] diff --git a/backend/crates/gateway-plugin/sdk/src/capability.rs b/backend/crates/gateway-plugin/sdk/src/capability.rs index 601779c6f..c736e9a09 100644 --- a/backend/crates/gateway-plugin/sdk/src/capability.rs +++ b/backend/crates/gateway-plugin/sdk/src/capability.rs @@ -33,7 +33,7 @@ impl Capability { | Self::Observer | Self::Management => &[2], Self::Scheduler => &[3], - Self::UpstreamAdapter => &[4], + Self::UpstreamAdapter => &[5], Self::ModelCatalog | Self::CommandLine | Self::Maintenance => &[1], } } diff --git a/backend/crates/gateway-plugin/sdk/tests/call/upstream_adapter/mod.rs b/backend/crates/gateway-plugin/sdk/tests/call/upstream_adapter/mod.rs index d22e18864..383719e5f 100644 --- a/backend/crates/gateway-plugin/sdk/tests/call/upstream_adapter/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/call/upstream_adapter/mod.rs @@ -1,3 +1,25 @@ -//! 上游适配器调用合同的测试入口 +//! 验证上游适配器注册与调用合同 mod codec; + +#[test] +fn registration_round_trip_preserves_explicit_account_scopes() { + use gateway_plugin_sdk::call::upstream_adapter::UpstreamAdapterRegistration; + use serde_json::json; + + for accounts in [json!([]), json!(["acct_one", "acct_two"])] { + let wire = json!({"adapters":[{ + "id":"example", "provider":"openai", "base_url":"https://example.invalid/", + "paths":[{"path":"responses","purpose":"inference"}], + "authentication_kinds":["oauth"], "transport":"http_sse", "protocol":"openai", + "account_ids":accounts + }]}); + let registration: UpstreamAdapterRegistration = serde_json::from_value(wire).unwrap(); + let serialized = serde_json::to_value(registration).unwrap(); + let round_trip: UpstreamAdapterRegistration = serde_json::from_value(serialized).unwrap(); + assert_eq!( + serde_json::to_value(&round_trip.adapters[0].account_ids).unwrap(), + accounts + ); + } +} diff --git a/backend/crates/gateway-plugin/sdk/tests/manifest.rs b/backend/crates/gateway-plugin/sdk/tests/manifest.rs index 42d460311..0e6500d1c 100644 --- a/backend/crates/gateway-plugin/sdk/tests/manifest.rs +++ b/backend/crates/gateway-plugin/sdk/tests/manifest.rs @@ -135,7 +135,7 @@ fn capability_contract_versions_are_checked_independently_of_rpc_version() { for (capability, stage, supported, needs_formats) in [ (Capability::Middleware, Stage::Request, 5, true), (Capability::Management, Stage::Management, 2, false), - (Capability::UpstreamAdapter, Stage::Upstream, 4, true), + (Capability::UpstreamAdapter, Stage::Upstream, 5, true), (Capability::Scheduler, Stage::Scheduling, 3, false), ] { for mut manifest in [source_manifest(), packaged_manifest()] { diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 586f43537..91deb6984 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -165,6 +165,7 @@ impl<'a> CodexSelectionModel<'a> { } struct CredentialSelectionInput<'a> { + managed_upstream: bool, requires_websocket: bool, requires_oauth: bool, request_url: &'a Url, @@ -350,6 +351,7 @@ impl CodexCredentialSelector { request: &SelectCodexCredential<'_>, ) -> Result { let input = CredentialSelectionInput { + managed_upstream: false, requires_websocket: false, requires_oauth: false, request_url: request.request_url, @@ -375,6 +377,7 @@ impl CodexCredentialSelector { guardian: bool, ) -> Result { let input = CredentialSelectionInput { + managed_upstream: true, requires_websocket, requires_oauth: false, request_url: request.request_url, @@ -450,6 +453,7 @@ impl CodexCredentialSelector { request: &SelectCodexProviderEndpointCredential<'_>, ) -> Result { let input = CredentialSelectionInput { + managed_upstream: false, requires_websocket: false, requires_oauth: request.requires_oauth, request_url: request.request_url, @@ -506,6 +510,13 @@ impl CodexCredentialSelector { model: CodexSelectionModel<'_>, follow_only: bool, ) -> Result { + let adapter_model = match model { + CodexSelectionModel::Requested(model) if request.managed_upstream => Some( + gateway_core::routing::UpstreamModelId::new(model) + .map_err(|_| CredentialSelectionError::NoEligibleCredential)?, + ), + _ => None, + }; let diagnostic = request.attempt.is_diagnostic_required_account(); let binding_key = request.session_affinity_key.filter(|_| !diagnostic); let affinity_mode = request @@ -633,7 +644,26 @@ impl CodexCredentialSelector { .collect::>(); let mut eligible = Vec::with_capacity(accounts.len()); for account in accounts { + let adapter = if let Some(model) = adapter_model.as_ref() { + match request.attempt.upstream_adapter_for_account( + &self.provider_kind, + model, + account.id(), + ) { + Ok(adapter) => adapter, + // 已匹配但不可用的适配器不能回退到该账号的原生目标 + Err(_) => continue, + } + } else { + None + }; + if adapter.as_ref().is_some_and(|adapter| { + !adapter.supports_authentication(account.authentication_kind()) + }) { + continue; + } if request.requires_websocket + && adapter.is_none() && pinned_account.as_ref().is_none_or(|id| id == account.id()) { let runtime = match self.repository.reload_selected_credential(&account).await { diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 7e0d561b6..6603b9322 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -16,6 +16,7 @@ use gateway_core::engine::provider::{ ContinuationRequestObservation, EventStream, Provider, ProviderCallMetadata, ProviderRequest, ProviderRequestObservation, ProviderSelectionObservation, ProviderStream, }; +use gateway_core::engine::upstream_adapter::UpstreamAdapterSelection; use gateway_core::engine::{AttemptContext, AttemptTransport, ContinuationAttempt}; use gateway_core::error::{ ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure, @@ -640,7 +641,8 @@ impl Provider for CodexProvider { unreachable!("generate settings keep the operation kind") }; let adapter = context.upstream_adapter(candidate.provider(), upstream_model)?; - if adapter.is_none() + let adapted = matches!(adapter, UpstreamAdapterSelection::Adapter(_)); + if !adapted && generate .provider_session_state(PROVIDER_NAME) .is_some_and(|state| state.extension_owner().is_some()) @@ -658,18 +660,15 @@ impl Provider for CodexProvider { 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 adapter_affinity = upstream.as_ref().filter(|_| adapted).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()) + .filter(|_| !adapted) .map(|upstream| self.prepare_generate_request(generate, upstream, &context)); if let Some(prepared) = preselection.as_ref() { prepared.checked_native_continuation_scope(&context)?; @@ -722,6 +721,11 @@ impl Provider for CodexProvider { }; let account_selection_wait_ms = u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); + let adapter = context.upstream_adapter_for_account( + candidate.provider(), + upstream_model, + lease.account_id(), + )?; if adapter.is_some() { // 适配器自行解释业务改写，宿主轮次关联只使用本次选号已确认的会话 self.remember_generate_turn( 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 e95efe7f7..2ca3eddbb 100644 --- a/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs +++ b/backend/crates/providers/openai/tests/provider/contract/upstream_adapter.rs @@ -5,6 +5,7 @@ use gateway_core::account::AccountRuntimeSignals; use gateway_core::account::FastMode; use gateway_core::engine::upstream_adapter::{ UpstreamAccountConnection, UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan, + UpstreamAdapterSelection, }; #[derive(Debug, Default)] @@ -18,12 +19,13 @@ impl UpstreamAdapterPlan for UnmatchedAdapterPlan { _: &AttemptContext, provider: &ProviderKind, model: &UpstreamModelId, - ) -> Result>, gateway_core::error::ProviderError> { + _: Option<&ProviderAccountId>, + ) -> Result { 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) + Ok(UpstreamAdapterSelection::Native) } } @@ -186,12 +188,10 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AdapterProb _: &AttemptContext, provider: &ProviderKind, _: &UpstreamModelId, - ) -> Result< - Option>, - gateway_core::error::ProviderError, - > { + _: Option<&ProviderAccountId>, + ) -> Result { assert_eq!(provider.as_str(), "openai"); - Ok(Some(Arc::new(Self { + Ok(UpstreamAdapterSelection::Adapter(Arc::new(Self { polls: self.polls.clone(), }))) } @@ -202,6 +202,10 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe { "http_sse" } + fn supports_authentication(&self, kind: &str) -> bool { + kind == "oauth" + } + fn execute( self: Arc, invocation: gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation, @@ -341,6 +345,7 @@ type SelectedConnection = Arc; type AdapterInvocations = Arc>>; struct AffinityAdapterProbe { + account_scoped: bool, invocations: Option, } @@ -357,8 +362,18 @@ impl UpstreamAdapterPlan for AffinityAdapterProbe { _: &AttemptContext, _: &ProviderKind, _: &UpstreamModelId, - ) -> Result>, gateway_core::error::ProviderError> { - Ok(Some(Arc::new(Self { + account: Option<&ProviderAccountId>, + ) -> Result { + if self.account_scoped { + let Some(account) = account else { + return Ok(UpstreamAdapterSelection::PendingAccount); + }; + if !["acct_subagent_a", "acct_subagent_b"].contains(&account.as_str()) { + return Ok(UpstreamAdapterSelection::Native); + } + } + Ok(UpstreamAdapterSelection::Adapter(Arc::new(Self { + account_scoped: self.account_scoped, invocations: self.invocations.clone(), }))) } @@ -369,6 +384,10 @@ impl UpstreamAdapter for AffinityAdapterProbe { "http_sse" } + fn supports_authentication(&self, kind: &str) -> bool { + kind == "oauth" + } + fn execute( self: Arc, invocation: UpstreamAdapterInvocation, @@ -387,6 +406,7 @@ impl UpstreamAdapter for AffinityAdapterProbe { fn affinity_adapter_context( request_id: &str, + account_scoped: bool, client_key: &str, middleware: Option>, invocations: Option, @@ -396,7 +416,10 @@ fn affinity_adapter_context( Arc::new(TestExtensionLease), ); let plan = gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( - Arc::new(AffinityAdapterProbe { invocations }), + Arc::new(AffinityAdapterProbe { + account_scoped, + invocations, + }), extension.clone(), ); AttemptContext::new( @@ -423,54 +446,57 @@ fn affinity_adapter_context( #[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), - ) + for account_scoped in [false, true] { + 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", account_scoped, "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", account_scoped, 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(), - ); - 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")] { + .unwrap(); let stream = provider - .clone() .execute( request(), - affinity_adapter_context("req_follow", key, None, None), + affinity_adapter_context("req_migrate", account_scoped, "key_a", None, None), ) .await .unwrap(); - assert_eq!(stream.metadata().provider_account_id().as_str(), expected); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); + assert!(server.received_requests().await.unwrap().is_empty()); } - 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] @@ -508,7 +534,7 @@ async fn upstream_adapter_preserves_plugin_rewrites_without_changing_admitted_af "openai", session_binding::turn_request("session-a", "seed-a"), ), - affinity_adapter_context("req_seed_a", "key_a", None, None), + affinity_adapter_context("req_seed_a", false, "key_a", None, None), ) .await .unwrap(), @@ -523,7 +549,7 @@ async fn upstream_adapter_preserves_plugin_rewrites_without_changing_admitted_af "openai", session_binding::turn_request("session-b", "seed-b"), ), - affinity_adapter_context("req_seed_b", "key_a", None, None), + affinity_adapter_context("req_seed_b", false, "key_a", None, None), ) .await .unwrap(), @@ -556,6 +582,7 @@ async fn upstream_adapter_preserves_plugin_rewrites_without_changing_admitted_af planned_request("openai", operation.clone()), affinity_adapter_context( "req_changed", + false, "key_a", Some(Arc::new(middleware)), Some(invocations.clone()), @@ -665,7 +692,7 @@ async fn upstream_adapter_child_cannot_migrate_a_strict_session() { .clone() .execute( request("root"), - affinity_adapter_context("req_root", "key_a", None, None), + affinity_adapter_context("req_root", false, "key_a", None, None), ) .await .unwrap(), @@ -678,7 +705,7 @@ async fn upstream_adapter_child_cannot_migrate_a_strict_session() { let error = provider .execute( request("child"), - affinity_adapter_context("req_child", "key_a", None, None), + affinity_adapter_context("req_child", false, "key_a", None, None), ) .await .err() @@ -709,7 +736,7 @@ async fn upstream_adapter_does_not_infer_sessions_from_unknown_protocols_or_miss provider .execute( planned_request("openai", operation), - affinity_adapter_context("req_unassociated", "key_a", None, None), + affinity_adapter_context("req_unassociated", false, "key_a", None, None), ) .await .unwrap(), @@ -720,6 +747,238 @@ async fn upstream_adapter_does_not_infer_sessions_from_unknown_protocols_or_miss } } +#[derive(Debug, Clone)] +struct AccountScopedProbe { + available: bool, + oauth: bool, +} + +impl UpstreamAdapterPlan for AccountScopedProbe { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + account: Option<&ProviderAccountId>, + ) -> Result { + let Some(account) = account else { + return Ok(UpstreamAdapterSelection::PendingAccount); + }; + if account.as_str() != "acct_subagent_a" { + return Ok(UpstreamAdapterSelection::Native); + } + if !self.available { + return Err(gateway_core::error::ProviderError::new( + ProviderErrorKind::Unavailable, + UpstreamSendState::NotSent, + )); + } + Ok(UpstreamAdapterSelection::Adapter(Arc::new(self.clone()))) + } +} + +impl UpstreamAdapter for AccountScopedProbe { + fn transport(&self) -> &str { + "http_json" + } + fn supports_authentication(&self, kind: &str) -> bool { + self.oauth && kind == "oauth" + } + fn execute( + self: Arc, + invocation: UpstreamAdapterInvocation, + ) -> gateway_core::engine::provider::EventStream { + assert_eq!(invocation.account.account_id().as_str(), "acct_subagent_a"); + Box::pin(futures::stream::once(async { + Err(gateway_core::error::ProviderError::new( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )) + })) + } +} + +#[tokio::test] +async fn account_scoped_adapter_selection_reuses_native_pool_and_filters_ineligible_accounts() { + for (available, oauth) in [(true, true), (true, false), (false, true)] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + store.set_scheduling("acct_subagent_a", None, AccountWeight::new(100).unwrap()); + let server = MockServer::start().await; + store + .seed_api_key( + "acct_subagent_b", + server.uri(), + provider_openai::credential::ResponsesTransport::Http, + ) + .await; + Mock::given(method("POST")) + .and(path("/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CAPTURE_COMPLETED_SSE), + ) + .expect(1) + .mount(&server) + .await; + let provider = provider_with_base_url(&store, server.uri()); + let request_context = || { + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_account_adapter").unwrap(), + ClientApiKeyId::new("key_openai_contract").unwrap(), + ) + .with_upstream_adapters(Some( + gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( + Arc::new(AccountScopedProbe { available, oauth }), + ExtensionSetReference::new( + ExtensionSetId::new("account-adapters".to_owned()).unwrap(), + Arc::new(TestExtensionLease), + ), + ), + )), + 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(), + ) + }; + if available && oauth { + let mut stream = provider + .clone() + .execute( + planned_request("openai", generate_operation()), + request_context(), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!( + stream.next().await.unwrap().unwrap_err().kind(), + ProviderErrorKind::Cancelled + ); + drop(stream); + assert!(server.received_requests().await.unwrap().is_empty()); + store + .set_enabled(&ProviderAccountId::new("acct_subagent_a").unwrap(), false) + .await + .unwrap(); + } + let mut stream = provider + .execute( + planned_request("openai", generate_operation()), + request_context(), + ) + .await + .unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_b" + ); + while let Some(event) = stream.next().await { + event.unwrap(); + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } +} + +#[tokio::test] +async fn account_scoped_adapter_preselection_keeps_extension_continuations_out_of_native_checks() { + use gateway_core::operation::ExtensionSessionOwner; + + for attempt in [ + ContinuationAttempt::Native, + ContinuationAttempt::ReplayOwner, + ContinuationAttempt::ReplayAny, + ] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_subagent_a").await; + let server = MockServer::start().await; + let provider = provider_with_base_url(&store, server.uri()); + let kind = ProviderKind::new("openai").unwrap(); + let account = ProviderAccountId::new("acct_subagent_a").unwrap(); + let key = ClientApiKeyId::new("key_openai_contract").unwrap(); + let pin = NativeContinuationPin::new( + PreviousResponseId::new("public-response"), + PreviousResponseId::new("adapter-response"), + key.clone(), + kind.clone(), + account.clone(), + ); + let context = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_scoped_continuation").unwrap(), + key, + ) + .with_upstream_adapters(Some( + gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( + Arc::new(AccountScopedProbe { + available: true, + oauth: true, + }), + ExtensionSetReference::new( + ExtensionSetId::new("account-adapters".to_owned()).unwrap(), + Arc::new(TestExtensionLease), + ), + ), + )), + NonZeroU32::MIN, + SystemTime::now() + Duration::from_secs(5), + account_policy(), + AccountAttemptContext::new( + Default::default(), + None, + Some(ProviderAccountStateOwner::new(kind, account)), + ) + .with_account_scope(contract_account_scope()), + Some(ContinuationBinding::Pinned(pin)), + CancellationToken::new(), + ) + .with_continuation_attempt(attempt); + let state = ProviderSessionState::new( + "openai", + Map::from_iter([("opaque".into(), json!("adapter-only"))]), + ) + .unwrap() + .with_extension_owner(ExtensionSessionOwner { + instance_id: "instance".into(), + contribution_id: "adapter".into(), + adapter_id: "example".into(), + generation: 1, + incarnation: "process".into(), + connection_local: true, + }); + let operation = generate_operation().with_provider_session_state(state); + let result = provider + .execute(planned_request("openai", operation), context) + .await; + if attempt == ContinuationAttempt::ReplayAny { + let error = result + .err() + .expect("unowned extension state cannot migrate"); + assert_eq!(error.send_state(), UpstreamSendState::NotSent); + } else { + let mut stream = result.unwrap(); + assert_eq!( + stream.metadata().provider_account_id().as_str(), + "acct_subagent_a" + ); + assert_eq!( + stream.next().await.unwrap().unwrap_err().kind(), + ProviderErrorKind::Cancelled + ); + } + assert!(server.received_requests().await.unwrap().is_empty()); + } +} + struct ConnectionProbe { selected: Arc>>>, } @@ -738,8 +997,9 @@ impl UpstreamAdapterPlan for ConnectionProbe { _: &AttemptContext, _: &ProviderKind, _: &UpstreamModelId, - ) -> Result>, gateway_core::error::ProviderError> { - Ok(Some(Arc::new(Self { + _: Option<&ProviderAccountId>, + ) -> Result { + Ok(UpstreamAdapterSelection::Adapter(Arc::new(Self { selected: Arc::clone(&self.selected), }))) } @@ -750,6 +1010,10 @@ impl UpstreamAdapter for ConnectionProbe { "http_sse" } + fn supports_authentication(&self, kind: &str) -> bool { + kind == "oauth" + } + fn execute( self: Arc, invocation: UpstreamAdapterInvocation, diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index 9a442ed16..dd9a79db0 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -141,6 +141,24 @@ impl GrokAccountSessionSelector { }) .collect() }; + let accounts = accounts + .into_iter() + .filter(|account| { + request.upstream_context.as_ref().is_none_or(|context| { + context + .upstream_adapter_for_account( + &self.provider_kind, + request.upstream_model(), + account.id(), + ) + .is_ok_and(|adapter| { + adapter.map_or(!request.adapter_required, |adapter| { + adapter.supports_authentication(account.authentication_kind()) + }) + }) + }) + }) + .collect::>(); let catalog_eligible = if diagnostic { accounts } else { diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs index 94825b187..21394882a 100644 --- a/backend/crates/providers/xai/src/provider/mod.rs +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -16,6 +16,7 @@ use gateway_core::engine::provider::{ EventStream, Provider, ProviderCallMetadata, ProviderRequest, ProviderRequestObservation, ProviderSelectionObservation, ProviderStream, }; +use gateway_core::engine::upstream_adapter::UpstreamAdapterSelection; use gateway_core::engine::{AttemptContext, ContinuationAttempt}; use gateway_core::error::{ ClientVisibleUpstreamError, ContinuationFailure, ContinuationRecoveryDisposition, @@ -267,7 +268,7 @@ impl GrokBuildProvider { ) -> Result { let adapter = context.upstream_adapter(candidate.provider(), candidate_upstream_model(candidate)?)?; - if adapter.is_none() + if !matches!(adapter, UpstreamAdapterSelection::Adapter(_)) && generate .provider_session_state(XAI_PROVIDER_NAME) .is_some_and(|state| state.extension_owner().is_some()) @@ -275,11 +276,16 @@ impl GrokBuildProvider { return Err(invalid_continuation()); } let selected = self - .select_grok_attempt(generate, candidate, context, adapter.is_some()) + .select_grok_attempt(generate, candidate, context, adapter) .await?; let provider_kind = ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?; let account_id = selected.selected.account_id().clone(); + let adapter = selected.context.upstream_adapter_for_account( + candidate.provider(), + &selected.upstream_model, + &account_id, + )?; let model = selected.upstream_model.as_str().to_owned(); let middleware_context = selected.context.clone(); let provider = Arc::clone(&self); @@ -313,74 +319,43 @@ impl GrokBuildProvider { generate: &GenerateRequest, candidate: &ProviderCandidate, context: AttemptContext, - adapted: bool, + adapter: UpstreamAdapterSelection, ) -> Result { let upstream_model = candidate_upstream_model(candidate)?; - let previous_session = if adapted { - None - } else { - decode_xai_session_state(generate)? + let (facts, native_error) = match grok_selection_facts( + generate, + upstream_model, + &context, + matches!(adapter, UpstreamAdapterSelection::Adapter(_)), + ) { + Ok(facts) => (facts, None), + Err(error) if matches!(adapter, UpstreamAdapterSelection::PendingAccount) => { + // 原生编码失败只排除原生候选，插件仍拥有自身协议；Core 的账号 owner 继续生效 + ( + grok_selection_facts(generate, upstream_model, &context, true)?, + Some(error), + ) + } + Err(error) => return Err(error), }; - let native_source = !adapted && generate.protocol_payload().protocol() == "openai"; - let native_compaction = native_source - && crate::transport::compaction::has_terminal_compaction_trigger(generate); - let (selection_model, operation_account, affinity) = if native_compaction { - validate_compaction_context(&context)?; - let operation_account = previous_session_account(previous_session.as_ref())?; - let request = GrokCompactionRequest::encode( - generate, - upstream_model.as_str(), - context.client_api_key_ref(), - ) - .map_err(map_request_error)?; - let selection_model = request - .upstream_model() - .ok_or_else(protocol_not_sent)? - .to_owned(); - ( - selection_model, - merge_frozen_account_owner(&context, operation_account)?, - request.affinity().cloned(), - ) - } else if native_source { - let operation_account = continuation_account(&context, previous_session.as_ref())?; - let request = GrokResponsesRequest::encode( - generate, - upstream_model.as_str(), - context.client_api_key_ref(), - ) - .map_err(map_request_error)?; - let selection_model = request - .upstream_model() - .ok_or_else(protocol_not_sent)? - .to_owned(); - ( - selection_model, - merge_frozen_account_owner(&context, operation_account)?, - request.affinity().cloned(), - ) - } else { - // 源协议由转换插件拥有；这里只使用冻结的候选模型与 Core 账号 owner， - // 不把未知 JSON 中的同名字段解释为 xAI 会话或亲和事实 - let operation_account = continuation_account(&context, previous_session.as_ref())?; - ( - upstream_model.as_str().to_owned(), - merge_frozen_account_owner(&context, operation_account)?, - None, - ) - }; - let selection_model = - UpstreamModelId::new(selection_model).map_err(|_| protocol_not_sent())?; let selection_started_at = Instant::now(); let selected = select_grok_session( self.selector.as_ref(), candidate, - &selection_model, + &facts.model, &context, - operation_account, - affinity, + facts.account, + facts.affinity, + native_error.is_some(), ) - .await?; + .await + .map_err(|error| { + if error.kind() == ProviderErrorKind::NoEligibleAccount { + native_error.unwrap_or(error) + } else { + error + } + })?; let account_selection_wait_ms = u64::try_from(selection_started_at.elapsed().as_millis()).unwrap_or(u64::MAX); Ok(SelectedGrokAttempt { @@ -647,6 +622,78 @@ fn native_request_requirements(request: &GenerateRequest) -> CapabilityRequireme .capability_requirements() } +struct GrokSelectionFacts { + model: UpstreamModelId, + account: Option, + affinity: Option, +} + +fn grok_selection_facts( + generate: &GenerateRequest, + upstream_model: &UpstreamModelId, + context: &AttemptContext, + adapted: bool, +) -> Result { + let previous_session = if adapted { + None + } else { + decode_xai_session_state(generate)? + }; + let native_source = !adapted && generate.protocol_payload().protocol() == "openai"; + let native_compaction = + native_source && crate::transport::compaction::has_terminal_compaction_trigger(generate); + let (selection_model, operation_account, affinity) = if native_compaction { + validate_compaction_context(context)?; + let operation_account = previous_session_account(previous_session.as_ref())?; + let request = GrokCompactionRequest::encode( + generate, + upstream_model.as_str(), + context.client_api_key_ref(), + ) + .map_err(map_request_error)?; + let selection_model = request + .upstream_model() + .ok_or_else(protocol_not_sent)? + .to_owned(); + ( + selection_model, + merge_frozen_account_owner(context, operation_account)?, + request.affinity().cloned(), + ) + } else if native_source { + let operation_account = continuation_account(context, previous_session.as_ref())?; + let request = GrokResponsesRequest::encode( + generate, + upstream_model.as_str(), + context.client_api_key_ref(), + ) + .map_err(map_request_error)?; + let selection_model = request + .upstream_model() + .ok_or_else(protocol_not_sent)? + .to_owned(); + ( + selection_model, + merge_frozen_account_owner(context, operation_account)?, + request.affinity().cloned(), + ) + } else { + // 源协议由转换插件拥有；这里只使用冻结的候选模型与 Core 账号 owner， + // 不把未知 JSON 中的同名字段解释为 xAI 会话或亲和事实 + let operation_account = continuation_account(context, previous_session.as_ref())?; + ( + upstream_model.as_str().to_owned(), + merge_frozen_account_owner(context, operation_account)?, + None, + ) + }; + Ok(GrokSelectionFacts { + model: UpstreamModelId::new(selection_model).map_err(|_| protocol_not_sent())?, + account: operation_account, + affinity, + }) +} + async fn select_grok_session( selector: &dyn GrokSessionSelector, candidate: &ProviderCandidate, @@ -654,9 +701,10 @@ async fn select_grok_session( context: &AttemptContext, operation_account: Option, affinity: Option, + adapter_required: bool, ) -> Result { let required_account = context.required_account().cloned().or(operation_account); - let selection = GrokSessionSelection::new( + let mut selection = GrokSessionSelection::new( wire_upstream_model.clone(), context.excluded_accounts().clone(), required_account.clone(), @@ -678,6 +726,8 @@ async fn select_grok_session( AccountEligibilityPolicy::Enforce }) .with_affinity(affinity); + selection.upstream_context = Some(context.clone()); + selection.adapter_required = adapter_required; if context.deadline().is_elapsed() { return Err(provider_error( ProviderErrorKind::Timeout, diff --git a/backend/crates/providers/xai/src/transport/session.rs b/backend/crates/providers/xai/src/transport/session.rs index 3ff9fa4bd..086da01e8 100644 --- a/backend/crates/providers/xai/src/transport/session.rs +++ b/backend/crates/providers/xai/src/transport/session.rs @@ -247,6 +247,8 @@ impl fmt::Debug for SelectedGrokSession { /// 单次 selector 调用的入参，持有冻结且不含密钥的尝试视图 #[derive(Debug, Clone)] pub struct GrokSessionSelection { + pub(crate) upstream_context: Option, + pub(crate) adapter_required: bool, pub(crate) lease_metadata: Option>, upstream_model: UpstreamModelId, excluded_accounts: BTreeSet, @@ -297,6 +299,8 @@ impl GrokSessionSelection { ) -> Self { Self { lease_metadata: None, + upstream_context: None, + adapter_required: false, upstream_model, excluded_accounts, required_account, diff --git a/backend/crates/providers/xai/tests/provider/contract/mod.rs b/backend/crates/providers/xai/tests/provider/contract/mod.rs index 226c27d63..b5c9ecf42 100644 --- a/backend/crates/providers/xai/tests/provider/contract/mod.rs +++ b/backend/crates/providers/xai/tests/provider/contract/mod.rs @@ -15,7 +15,7 @@ use bytes::Bytes; use futures::{StreamExt, future::BoxFuture, stream}; use gateway_core::account::{ AccountFeedbackStats, AccountRuntimeSignals, AccountSelectionPolicy, CredentialRevision, - CredentialState, ProviderAccountStore, RotationStrategy, + CredentialState, ProviderAccountId, ProviderAccountStore, RotationStrategy, }; use gateway_core::engine::continuation::{ ContinuationBinding, NativeContinuationPin, PreviousResponseId, @@ -26,6 +26,7 @@ use gateway_core::engine::middleware::{ MiddlewareNext, MiddlewarePlan, MiddlewareRequest, MiddlewareResponse, }; use gateway_core::engine::provider::{Provider, ProviderRequest}; +use gateway_core::engine::upstream_adapter::UpstreamAdapterSelection; use gateway_core::engine::{ AccountAttemptContext, AttemptContext, ContinuationAttempt, ModelRequestId, }; @@ -4395,6 +4396,7 @@ async fn configured_request_profile_reaches_inference_headers() { #[derive(Debug)] struct AdapterProbe { polls: Arc, + account_scoped: bool, } impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AdapterProbe { @@ -4403,13 +4405,18 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AdapterProb _: &AttemptContext, provider: &ProviderKind, _: &UpstreamModelId, - ) -> Result< - Option>, - gateway_core::error::ProviderError, - > { + account: Option<&ProviderAccountId>, + ) -> Result { assert_eq!(provider.as_str(), "xai"); - Ok(Some(Arc::new(Self { + if self.account_scoped && account.is_none() { + return Ok(UpstreamAdapterSelection::PendingAccount); + } + if let Some(account) = account { + assert_eq!(account.as_str(), "acct_provider"); + } + Ok(UpstreamAdapterSelection::Adapter(Arc::new(Self { polls: self.polls.clone(), + account_scoped: self.account_scoped, }))) } } @@ -4419,6 +4426,10 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe { "http_sse" } + fn supports_authentication(&self, kind: &str) -> bool { + kind == "oauth" + } + fn execute( self: Arc, invocation: gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation, @@ -4476,20 +4487,22 @@ impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe { #[tokio::test] async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_and_stays_cold() { - assert_upstream_adapter_stays_cold(operation()).await; - for protocol in ["openai", "example-source"] { - let operation = Operation::Generate(GenerateRequest::from_protocol_payload( - ProtocolPayload::json_object( - protocol, - object(json!({"model": MODEL, "input": "hello", "generate": false})), - ) - .unwrap(), - )); - assert_upstream_adapter_stays_cold(operation).await; + for account_scoped in [false, true] { + assert_upstream_adapter_stays_cold(operation(), account_scoped).await; + for protocol in ["openai", "example-source"] { + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + protocol, + object(json!({"model": MODEL, "input": "hello", "generate": false})), + ) + .unwrap(), + )); + assert_upstream_adapter_stays_cold(operation, account_scoped).await; + } } } -async fn assert_upstream_adapter_stays_cold(operation: Operation) { +async fn assert_upstream_adapter_stays_cold(operation: Operation, account_scoped: bool) { let selector = StubSelector::success(); let transport = StubInferenceTransport::success(); let provider = provider(selector.clone(), transport.clone()).await; @@ -4515,6 +4528,7 @@ async fn assert_upstream_adapter_stays_cold(operation: Operation) { gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( Arc::new(AdapterProbe { polls: polls.clone(), + account_scoped, }), lease.clone(), ), diff --git a/backend/crates/providers/xai/tests/provider/contract/scheduling.rs b/backend/crates/providers/xai/tests/provider/contract/scheduling.rs index 8700a7fd9..7da370b34 100644 --- a/backend/crates/providers/xai/tests/provider/contract/scheduling.rs +++ b/backend/crates/providers/xai/tests/provider/contract/scheduling.rs @@ -38,6 +38,106 @@ impl RequestPolicyPlan for Filter { } } +#[derive(Debug)] +struct PreferAdapterAccount(ProviderAccountId); + +impl RequestPolicyPlan for PreferAdapterAccount { + fn route_model( + &self, + _: ModelRouteInput, + ) -> BoxFuture<'static, Result> { + Box::pin(async { Ok(ModelRouteDecision::Unhandled) }) + } + + fn schedule_account( + &self, + input: AccountScheduleInput, + ) -> BoxFuture<'static, Result> { + let preferred = self.0.clone(); + let decision = if input + .candidates() + .iter() + .any(|candidate| candidate.account_id() == &preferred) + { + AccountScheduleDecision::Pick(preferred) + } else { + AccountScheduleDecision::Delegate + }; + Box::pin(async move { Ok(decision) }) + } +} + +#[derive(Debug, Clone, Copy)] +enum AdapterState { + Available, + Unavailable, + AuthenticationMismatch, +} + +#[derive(Debug)] +struct AccountEligibilityPlan { + state: AdapterState, + polls: Arc, +} + +impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AccountEligibilityPlan { + fn select( + &self, + _: &AttemptContext, + _: &ProviderKind, + _: &UpstreamModelId, + account: Option<&ProviderAccountId>, + ) -> Result { + let Some(account) = account else { + return Ok(UpstreamAdapterSelection::PendingAccount); + }; + if account != &account_id("adapter-filter-a") { + return Ok(UpstreamAdapterSelection::Native); + } + match self.state { + AdapterState::Unavailable => Err(ProviderError::new( + ProviderErrorKind::Unavailable, + UpstreamSendState::NotSent, + )), + AdapterState::Available | AdapterState::AuthenticationMismatch => Ok( + UpstreamAdapterSelection::Adapter(Arc::new(AccountAdapter { + supports_authentication: matches!(self.state, AdapterState::Available), + polls: Arc::clone(&self.polls), + })), + ), + } + } +} + +#[derive(Debug)] +struct AccountAdapter { + supports_authentication: bool, + polls: Arc, +} + +impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AccountAdapter { + fn transport(&self) -> &str { + "http_sse" + } + + fn supports_authentication(&self, _: &str) -> bool { + self.supports_authentication + } + + fn execute( + self: Arc, + _: gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation, + ) -> gateway_core::engine::provider::EventStream { + Box::pin(futures::stream::once(async move { + self.polls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Err(ProviderError::new( + ProviderErrorKind::Cancelled, + UpstreamSendState::NotSent, + )) + })) + } +} + #[derive(Default)] struct Leases { busy: Mutex, @@ -270,3 +370,187 @@ async fn filtered_scheduling_covers_lease_retry_waiting_and_pinned_terminal_reje assert!(store.account(&account_id("filter-a")).unwrap().enabled()); } } + +#[tokio::test] +async fn account_scoped_adapter_filters_candidates_by_adapter_and_native_request_eligibility() { + for (state, prewarm) in [ + (AdapterState::Available, false), + (AdapterState::Unavailable, false), + (AdapterState::AuthenticationMismatch, false), + (AdapterState::Available, true), + (AdapterState::Unavailable, true), + (AdapterState::AuthenticationMismatch, true), + ] { + let store = MemoryProviderAccountStore::shared(); + let repository = GrokCredentialRepository::new(store.clone()); + let mut accounts = BTreeMap::new(); + let kind = ProviderKind::new("xai").unwrap(); + for suffix in ["adapter-filter-a", "adapter-filter-b"] { + let input = create_input(suffix, suffix); + seed_input(&store, &input).await.unwrap(); + repository + .update_state(&UpdateGrokCredentialState { + account_id: input.account_id.clone(), + expected_revision: CredentialRevision::new(1).unwrap(), + credential_state: CredentialState::Ready, + error_reason: None, + error_message: None, + observed_at: chrono::Utc::now(), + }) + .await + .unwrap(); + accounts.insert( + input.account_id, + RuntimeAccount::new(kind.clone(), BTreeSet::new()), + ); + } + let leases = Arc::new(Leases::default()); + let cache = MemoryGrokCatalogCache::shared(); + let selector = GrokAccountSessionSelector::new( + kind.clone(), + repository.clone(), + cache.clone(), + Arc::new(crate::support::grok_quota_service( + repository.clone(), + Arc::new(UnavailableBillingTransport), + )), + leases.clone(), + Arc::new(MemoryCooldownPort::default()), + Arc::new(AccountFeedbackStats::default()), + ); + let provider = Arc::new( + GrokBuildProvider::new( + Arc::new(selector), + StubInferenceTransport::success(), + Arc::new(crate::support::grok_catalog_service( + repository, + Arc::new(StaticCatalogTransport), + cache, + )), + StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), + Arc::new(AccountFeedbackStats::default()), + crate::support::xai_wire_profile(), + ) + .unwrap(), + ); + let id = ModelRequestId::new("req_xai_adapter_filter").unwrap(); + let key = ClientApiKeyId::new("key_xai_contract").unwrap(); + let policy = RequestPolicyContext::new( + Arc::new(PreferAdapterAccount(account_id( + if matches!(state, AdapterState::Available) { + "adapter-filter-b" + } else { + "adapter-filter-a" + }, + ))), + ExtensionSetReference::new( + ExtensionSetId::new("adapter-filter-policy".into()).unwrap(), + Arc::new(TestExtensionLease), + ), + id.clone(), + key.clone(), + vec![], + ); + let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let adapter_plan = gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new( + Arc::new(AccountEligibilityPlan { + state, + polls: Arc::clone(&polls), + }), + ExtensionSetReference::new( + ExtensionSetId::new("adapter-filter-plan".into()).unwrap(), + Arc::new(TestExtensionLease), + ), + ); + let scope = Arc::new(FrozenAccountScope::new( + Arc::new(RuntimeAccountDirectory::new(accounts)), + ClientRoutingScope::all_accounts(), + )); + let operation = if prewarm { + Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + object(json!({"model": MODEL, "input": "hello", "generate": false})), + ) + .unwrap(), + )) + } else { + operation() + }; + let snapshot = RuntimeSnapshot::new( + ConfigRevision::new(1).unwrap(), + gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None), + vec![kind.clone()], + vec![ProviderModel::new( + kind, + UpstreamModelId::new(MODEL).unwrap(), + ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(131_072)) + .with_feature(Feature::NativeContinuation, SupportLevel::Native), + )], + vec![], + ) + .unwrap(); + let plan = snapshot + .plan( + &PublicModelId::new(MODEL).unwrap(), + &operation, + scope.clone(), + &RoutingContext::default(), + ) + .unwrap(); + let request = ProviderRequest::new(operation, plan.candidates()[0].clone()); + let attempt = AttemptContext::new( + gateway_core::engine::RequestAttemptContext::new(id, key) + .with_request_policy(Some(policy)) + .with_upstream_adapters(Some(adapter_plan)), + NonZeroU32::MIN, + SystemTime::now() + Duration::from_secs(5), + selection_policy(), + AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope(scope), + None, + CancellationToken::new(), + ); + + let result = provider.execute(request, attempt).await; + if prewarm && !matches!(state, AdapterState::Available) { + assert_eq!( + result + .err() + .expect("no adapter can accept the prewarm request") + .kind(), + ProviderErrorKind::Unsupported + ); + assert!(leases.acquired.lock().unwrap().is_empty()); + assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0); + continue; + } + let mut stream = result.expect("eligible account remains available"); + let selected_account = account_id(if prewarm { + "adapter-filter-a" + } else { + "adapter-filter-b" + }); + assert_eq!( + stream.metadata().provider_account_id(), + &selected_account, + "{state:?}" + ); + assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!( + leases + .acquired + .lock() + .unwrap() + .iter() + .all(|account| account == &selected_account) + ); + if prewarm { + assert_eq!( + stream.next().await.unwrap().unwrap_err().kind(), + ProviderErrorKind::Cancelled + ); + assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + drop(stream); + } +}