diff --git a/backend/crates/gateway-core/src/engine/authentication.rs b/backend/crates/gateway-core/src/engine/authentication.rs index 074fcda58..ec2024d33 100644 --- a/backend/crates/gateway-core/src/engine/authentication.rs +++ b/backend/crates/gateway-core/src/engine/authentication.rs @@ -16,8 +16,9 @@ use crate::{ pub const MAXIMUM_AUTHORIZATION_BYTES: usize = 8 * 1024; /// API 从数据面 Authorization 头构造的有界认证信封 -#[derive(Clone, PartialEq, Eq)] +#[derive(Clone)] pub struct ClientAuthenticationRequest { + invocation: Option>, request_context: Option>, settings: Option, authorization: Arc, @@ -25,6 +26,18 @@ pub struct ClientAuthenticationRequest { } impl ClientAuthenticationRequest { + /// 子认证继承受信调用的来源、取消和副作用归属，不据此跳过认证器 + #[must_use] + pub fn with_invocation(mut self, origin: super::extensions::InvocationOrigin) -> Self { + self.invocation = Some(Arc::new(origin)); + self + } + + #[must_use] + pub fn invocation(&self) -> Option<&super::extensions::InvocationOrigin> { + self.invocation.as_deref() + } + #[must_use] pub fn with_request_context( mut self, @@ -35,12 +48,18 @@ impl ClientAuthenticationRequest { } #[must_use] pub fn request_context(&self) -> Option<&super::extensions::RequestContext> { - self.request_context.as_deref() + self.invocation + .as_ref() + .and_then(|origin| origin.extensions.request()) + .or(self.request_context.as_deref()) } #[must_use] pub fn shared_request_context(&self) -> Option> { - self.request_context.clone() + self.invocation + .as_ref() + .and_then(|origin| origin.extensions.shared_request()) + .or_else(|| self.request_context.clone()) } #[must_use] @@ -79,6 +98,7 @@ impl ClientAuthenticationRequest { .filter(|value| !value.is_empty()) .and_then(|value| PlaintextClientApiKey::new(value.to_owned()).ok()); Ok(Self { + invocation: None, settings: None, request_context: None, authorization: authorization.into(), diff --git a/backend/crates/gateway-core/src/engine/execution/mod.rs b/backend/crates/gateway-core/src/engine/execution/mod.rs index 9addf09e0..214eaf312 100644 --- a/backend/crates/gateway-core/src/engine/execution/mod.rs +++ b/backend/crates/gateway-core/src/engine/execution/mod.rs @@ -12,7 +12,6 @@ pub use contract::{ StartedExecution, }; mod scope; -pub use scope::BoundModelExecutionContext; mod service; pub use failure::gateway_error_from_engine; pub use service::DefaultExecutionService; diff --git a/backend/crates/gateway-core/src/engine/execution/scope.rs b/backend/crates/gateway-core/src/engine/execution/scope.rs index f9bcb3c7c..d4cc87cc1 100644 --- a/backend/crates/gateway-core/src/engine/execution/scope.rs +++ b/backend/crates/gateway-core/src/engine/execution/scope.rs @@ -15,13 +15,10 @@ use crate::{ lifecycle::{CancellationToken, Deadline}, operation::Operation, policy::ClientApiKeyId, - routing::{ - FrozenAccountScope, PublicModelId, UpstreamModelId, request_settings::RequestSettings, - }, + routing::{FrozenAccountScope, PublicModelId, UpstreamModelId}, }; use std::{ collections::BTreeMap, - fmt, sync::{ Arc, Mutex, Weak, atomic::{AtomicUsize, Ordering}, @@ -188,57 +185,6 @@ pub(super) struct ActiveRequestLease { pub(super) authority: Arc, } -/// 管理页或 CLI 一次调用持有的显式模型执行身份 -/// -/// 结构不暴露 Key 明文；Runtime 只能把它原样交回 Core 发起模型请求 -#[derive(Clone)] -pub struct BoundModelExecutionContext { - pub(super) authority: Arc, -} - -impl BoundModelExecutionContext { - /// 与实际执行共用绑定时的快照，不能回退到后续发布的配置 - #[must_use] - pub fn request_body_limit(&self) -> usize { - self.authority - .client - .snapshot - .responses_max_decompressed_body_bytes() - } - - #[must_use] - pub fn request_settings(&self) -> Option<&RequestSettings> { - self.authority.client.request_settings() - } -} - -pub(super) struct BoundModelExecutionAuthority { - pub(super) client: AuthenticatedClient, - pub(super) account_scope: Arc, - pub(super) cancellation: CancellationToken, - pub(super) extension_scope: ExtensionCallScope, - pub(super) initiating_plugin_instance_id: String, -} - -impl fmt::Debug for BoundModelExecutionContext { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("BoundModelExecutionContext") - .field("key_id", &self.authority.client.policy.key_id()) - .field( - "initiating_plugin_instance_id", - &self.authority.initiating_plugin_instance_id, - ) - .finish_non_exhaustive() - } -} - -impl Drop for BoundModelExecutionAuthority { - fn drop(&mut self) { - self.cancellation.cancel(); - } -} - impl ActiveRequestLease { pub(super) fn release(self) { drop(self); diff --git a/backend/crates/gateway-core/src/engine/execution/service.rs b/backend/crates/gateway-core/src/engine/execution/service.rs index 464e0687f..8148aeb43 100644 --- a/backend/crates/gateway-core/src/engine/execution/service.rs +++ b/backend/crates/gateway-core/src/engine/execution/service.rs @@ -12,9 +12,8 @@ use super::contract::{ use super::probe_store::{ProbeObservation, TransientExecutionStore}; use super::scope::{ ActiveRequestAuthority, ActiveRequestLease, ActiveRequestRegistry, AuthorizedExecution, - BoundExecutionFacts, BoundModelExecutionAuthority, BoundModelExecutionContext, - ExecutionStartGuard, ExecutionTarget, NestedExecutionFacts, NestedExecutionGraph, - PendingStartExecution, PreparedExecutionStart, + BoundExecutionFacts, ExecutionStartGuard, ExecutionTarget, NestedExecutionFacts, + NestedExecutionGraph, PendingStartExecution, PreparedExecutionStart, }; use super::{COORDINATION_TIMEOUT, DIAGNOSTIC_TIMEOUT}; use super::{scope::ExecutionAdmission, session::DefaultExecutionSession}; @@ -41,8 +40,8 @@ use crate::{ middleware::FrozenMiddlewarePlan, nested::{ AffinityLookupPort, AffinityLookupRequest, AffinityLookupResult, - BoundModelExecutionBinding, BoundModelExecutionRequest, ExecutionEffects, - NestedModelExecutionPort, NestedModelExecutionRequest, + BoundModelExecutionRequest, ExecutionEffects, NestedModelExecutionPort, + NestedModelExecutionRequest, }, observation::{ FrozenRequestObservationContext, RequestObservationDispatch, RequestObservationScope, @@ -1081,72 +1080,54 @@ impl DefaultExecutionService { }) } - async fn bind_bound_model_inner( - &self, - binding: BoundModelExecutionBinding, - ) -> Result { - let client = self - .client_for_key_id(&binding.client_key_id, binding.settings) - .await?; - if binding.extension_scope.len() >= ExtensionCallScope::MAXIMUM_DEPTH { - return Err(GatewayError::new( - GatewayErrorKind::PolicyDenied, - "plugin model execution depth exceeded", - )); - } - let extension_scope = binding - .extension_scope - .extending(binding.initiating_plugin_instance_id.clone()) - .ok_or_else(|| { - GatewayError::new( - GatewayErrorKind::Internal, - "bound plugin execution scope is invalid", - ) - })?; - let account_scope = Arc::clone(client.policy.account_scope()); - Ok(BoundModelExecutionContext { - authority: Arc::new(BoundModelExecutionAuthority { - client, - account_scope, - cancellation: binding.cancellation, - extension_scope, - initiating_plugin_instance_id: binding.initiating_plugin_instance_id, - }), - }) - } - async fn start_bound_model_inner( &self, request: BoundModelExecutionRequest, ) -> Result { let BoundModelExecutionRequest { - context, + client, + initiating_plugin_instance_id, + cancellation, + extension_scope, public_model, operation, mut metadata, provider, account, } = request; - let authority = &context.authority; - if authority.cancellation.is_cancelled() { + if cancellation.is_cancelled() { return Err(GatewayError::new( GatewayErrorKind::Cancelled, "bound plugin invocation was cancelled", )); } - let now = SystemTime::now(); - let deadline_at = Deadline::from_timeout(now, authority.client.execution_timeout()) + if extension_scope.len() >= ExtensionCallScope::MAXIMUM_DEPTH { + return Err(GatewayError::new( + GatewayErrorKind::PolicyDenied, + "plugin model execution depth exceeded", + )); + } + let extension_scope = extension_scope + .extending(initiating_plugin_instance_id.clone()) .ok_or_else(|| { + GatewayError::new( + GatewayErrorKind::PolicyDenied, + "plugin model execution scope is recursive", + ) + })?; + let now = SystemTime::now(); + let deadline_at = + Deadline::from_timeout(now, client.execution_timeout()).ok_or_else(|| { GatewayError::new(GatewayErrorKind::Internal, "system clock is invalid") })?; let (account_scope, required_provider) = restrict_model_execution_scope( - authority.account_scope.as_ref(), + client.policy.account_scope().as_ref(), provider, account.as_ref(), None, )?; // 连接承载多个独立执行；递归预算和副作用只在本次执行的子调用图内共享 - // 父连接的取消仍通过 authority 传播，不能因新建执行丢失扩展调用链 + // 父调用的取消继续向子执行传播，不能因新建执行丢失扩展调用链 let graph = Arc::new(NestedExecutionGraph::new(Arc::new( ExecutionEffects::default(), ))); @@ -1154,12 +1135,12 @@ impl DefaultExecutionService { let execution_effects_baseline = graph.effects.epoch(); metadata.endpoint = "host.model".to_owned(); metadata.transport = ClientTransport::InternalPlugin; - metadata.inherit_request_origin(&authority.extension_scope); + metadata.inherit_request_origin(&extension_scope); self.client_api_key_usage - .record_used(authority.client.policy.key_id()); + .record_used(client.policy.key_id()); self.start_authorized( PendingStartExecution { - client: authority.client.clone(), + client, target: ExecutionTarget::Model(public_model), operation, metadata, @@ -1168,13 +1149,13 @@ impl DefaultExecutionService { account_scope, response_control: None, deadline_at, - cancellation: authority.cancellation.child_token(), - extension_scope: authority.extension_scope.clone(), + cancellation: cancellation.child_token(), + extension_scope, required_provider, required_account: account, nested: None, bound: Some(BoundExecutionFacts { - initiating_plugin_instance_id: authority.initiating_plugin_instance_id.clone(), + initiating_plugin_instance_id, }), graph: Some(graph), execution_effects_baseline, @@ -1725,22 +1706,54 @@ impl ExecutionService for DefaultExecutionService { } impl NestedModelExecutionPort for DefaultExecutionService { + fn resolve_client( + &self, + client_key_id: ClientApiKeyId, + settings: Option, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { self.client_for_key_id(&client_key_id, settings).await }) + } + + fn verify_client( + &self, + request: ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let cancellation = request + .invocation() + .map(|origin| origin.cancellation.clone()); + let verification = self.verify_request(request).fuse(); + pin_mut!(verification); + match cancellation { + Some(cancellation) => { + let cancelled = cancellation.cancelled().fuse(); + pin_mut!(cancelled); + select_biased! { + () = cancelled => Err(GatewayError::new(GatewayErrorKind::Cancelled, "client verification was cancelled")), + result = verification => result.map_err(authentication_gateway_error), + } + } + None => verification.await.map_err(authentication_gateway_error), + } + }) + } + fn models( &self, - context: BoundModelExecutionContext, + client: AuthenticatedClient, + cancellation: CancellationToken, protocol: String, client_version: String, ) -> BoxFuture<'_, Result, GatewayError>> { Box::pin(async move { - let authority = &context.authority; - let cancellation = authority.cancellation.cancelled().fuse(); + let cancelled = cancellation.cancelled().fuse(); let timeout = Delay::new(DIAGNOSTIC_TIMEOUT).fuse(); let catalog = self - .client_model_catalog(&authority.client, &protocol, &client_version) + .client_model_catalog(&client, &protocol, &client_version) .fuse(); - pin_mut!(cancellation, timeout, catalog); + pin_mut!(cancelled, timeout, catalog); let models = select_biased! { - () = cancellation => return Err(GatewayError::new(GatewayErrorKind::Cancelled, "model catalog call was cancelled")), + () = cancelled => return Err(GatewayError::new(GatewayErrorKind::Cancelled, "model catalog call was cancelled")), () = timeout => return Err(GatewayError::new(GatewayErrorKind::Timeout, "model catalog deadline elapsed")), result = catalog => result.map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "model catalog is unavailable"))?, }; @@ -1754,13 +1767,6 @@ impl NestedModelExecutionPort for DefaultExecutionService { }) } - fn bind( - &self, - binding: BoundModelExecutionBinding, - ) -> BoxFuture<'_, Result> { - Box::pin(async move { self.bind_bound_model_inner(binding).await }) - } - fn start( &self, request: NestedModelExecutionRequest, diff --git a/backend/crates/gateway-core/src/engine/extensions.rs b/backend/crates/gateway-core/src/engine/extensions.rs index a840ba3b2..03c9e6cd1 100644 --- a/backend/crates/gateway-core/src/engine/extensions.rs +++ b/backend/crates/gateway-core/src/engine/extensions.rs @@ -62,6 +62,17 @@ impl ExtensionCallScope { } } +/// 宿主签发的父调用来源；身份裁决、执行链和实际资源仍由各自 owner 持有 +#[derive(Clone, Debug)] +pub struct InvocationOrigin { + pub request_id: String, + pub call_id: String, + pub cancellation: crate::lifecycle::CancellationToken, + pub extensions: ExtensionCallScope, + pub plan: Option, + pub effects: Option>, +} + /// 同一代次的具体执行计划；不拥有发布状态，运行资源由对应发布集合保活 #[derive(Debug, Default)] pub struct ExecutionExtensionPlans { diff --git a/backend/crates/gateway-core/src/engine/nested.rs b/backend/crates/gateway-core/src/engine/nested.rs index 82dede6a3..fb94cd4c7 100644 --- a/backend/crates/gateway-core/src/engine/nested.rs +++ b/backend/crates/gateway-core/src/engine/nested.rs @@ -19,7 +19,7 @@ use crate::{ use super::{ ModelRequestId, - execution::{BoundModelExecutionContext, ExecutionRequestMetadata, StartedExecution}, + execution::{AuthenticatedClient, ExecutionRequestMetadata, StartedExecution}, }; /// 同一请求图内不可逆外部副作用的单调水位 @@ -59,37 +59,14 @@ pub struct NestedModelExecutionRequest { pub parent_account: Option, } -/// 插件为当前调用或连接显式绑定的模型执行身份 -/// -/// Key ID 由插件选择，发起实例由宿主确定 -/// Core 在绑定时冻结当前 -/// Key 策略；每个实际模型请求仍独立经过准入、预算、账本与计费 -pub struct BoundModelExecutionBinding { - pub settings: Option, - pub client_key_id: ClientApiKeyId, - pub initiating_plugin_instance_id: String, - /// 身份随父调用或连接回收；每个模型请求使用独立的执行期限 - pub cancellation: CancellationToken, - /// 宿主保留的扩展调用链，异步观察不能通过新建模型请求清空防递归事实 - pub extension_scope: super::extensions::ExtensionCallScope, -} - -impl fmt::Debug for BoundModelExecutionBinding { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("BoundModelExecutionBinding") - .field("client_key_id", &self.client_key_id) - .field( - "initiating_plugin_instance_id", - &self.initiating_plugin_instance_id, - ) - .finish_non_exhaustive() - } -} - /// 显式身份作用域内的一次模型请求 +/// +/// 身份与权限可先用于查询；执行链只在真正启动模型请求时建立 pub struct BoundModelExecutionRequest { - pub context: BoundModelExecutionContext, + pub client: AuthenticatedClient, + pub initiating_plugin_instance_id: String, + pub cancellation: CancellationToken, + pub extension_scope: super::extensions::ExtensionCallScope, pub public_model: PublicModelId, pub operation: Operation, pub metadata: ExecutionRequestMetadata, @@ -101,7 +78,11 @@ impl fmt::Debug for BoundModelExecutionRequest { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("BoundModelExecutionRequest") - .field("context", &self.context) + .field("client", &self.client) + .field( + "initiating_plugin_instance_id", + &self.initiating_plugin_instance_id, + ) .field("public_model", &self.public_model) .field("operation", &self.operation.kind()) .field("provider", &self.provider) @@ -130,19 +111,27 @@ impl fmt::Debug for NestedModelExecutionRequest { /// 嵌套模型执行只进入 Core 的既有执行会话；Runtime 不实现重试、commit 或计费 pub trait NestedModelExecutionPort: Send + Sync { + /// 只解析身份与冻结权限，不建立模型执行图或记录使用事实 + fn resolve_client( + &self, + client_key_id: ClientApiKeyId, + settings: Option, + ) -> BoxFuture<'_, Result>; + + /// 复用入口认证计划及 principal 映射，不进行模型准入 + fn verify_client( + &self, + request: super::authentication::ClientAuthenticationRequest, + ) -> BoxFuture<'_, Result>; + fn models( &self, - context: BoundModelExecutionContext, + client: AuthenticatedClient, + cancellation: CancellationToken, protocol: String, client_version: String, ) -> BoxFuture<'_, Result, GatewayError>>; - /// 为插件调用冻结显式执行身份；取消父调用会取消其所有模型执行 - fn bind( - &self, - binding: BoundModelExecutionBinding, - ) -> BoxFuture<'_, Result>; - fn start( &self, request: NestedModelExecutionRequest, diff --git a/backend/crates/gateway-core/tests/engine/execution/identity.rs b/backend/crates/gateway-core/tests/engine/execution/identity.rs new file mode 100644 index 000000000..557630fe2 --- /dev/null +++ b/backend/crates/gateway-core/tests/engine/execution/identity.rs @@ -0,0 +1,269 @@ +//! 身份、冻结权限与查询不建立模型执行图；实际执行才检查调用链并占用执行资源 + +use super::*; +use gateway_core::engine::{ + extensions::{ExtensionCallScope, InvocationOrigin}, + nested::{BoundModelExecutionRequest, NestedModelExecutionPort}, +}; +use gateway_core::lifecycle::CancellationToken; + +#[test] +fn identity_queries_keep_frozen_permissions_without_using_execution_resources() { + block_on(async { + let snapshots = RuntimeSnapshotHandle::new(start_snapshot()); + let store = Arc::new(TrackingExecutionStore::default()); + let usage = Arc::new(RecordingClientApiKeyUsage::default()); + let service = DefaultExecutionService::new( + snapshots.clone(), + store.clone(), + ProviderRegistry::new([ + Arc::new(NativeCatalogProvider { fail: false }) as Arc + ]) + .unwrap(), + Arc::new(UnusedAdmissions), + Arc::new(UnusedContinuation), + usage.clone(), + Arc::new(crate::RecordingDiagnostics::default()), + ); + let settings = service.request_settings().unwrap(); + let key = ClientApiKeyId::new("key_start_test").unwrap(); + let client = + NestedModelExecutionPort::resolve_client(&service, key.clone(), Some(settings.clone())) + .await + .unwrap(); + snapshots.publish(start_snapshot_with_policy( + 2, + false, + RateLimits::unlimited(), + FastMode::Disabled, + )); + assert!( + client + .policy() + .account_scope() + .allows(&ProviderAccountId::new("acct_start").unwrap(),) + ); + assert!( + !client + .policy() + .account_scope() + .allows(&ProviderAccountId::new("acct_missing").unwrap(),) + ); + let models = NestedModelExecutionPort::models( + &service, + client.clone(), + CancellationToken::default(), + "codex".into(), + "0.154.0".into(), + ) + .await + .unwrap(); + assert_eq!(models, vec![PublicModelId::new("gpt-start").unwrap()]); + let frozen = + NestedModelExecutionPort::resolve_client(&service, key.clone(), Some(settings)) + .await + .unwrap(); + assert_eq!(frozen.snapshot().revision().get(), 1); + let current = NestedModelExecutionPort::resolve_client(&service, key, None).await; + assert!(matches!(current, Err(error) if error.kind() == GatewayErrorKind::Unauthorized)); + + let cancellation = CancellationToken::default(); + cancellation.cancel(); + let cancelled = NestedModelExecutionPort::models( + &service, + client, + cancellation, + "codex".into(), + "0.154.0".into(), + ) + .await; + assert!(matches!(cancelled, Err(error) if error.kind() == GatewayErrorKind::Cancelled)); + assert!(usage.recorded().is_empty()); + assert!(!store.touched.load(Ordering::SeqCst)); + }); +} + +#[test] +fn execution_scope_limits_apply_only_when_starting_a_model() { + block_on(async { + let store = Arc::new(TrackingExecutionStore::default()); + let usage = Arc::new(RecordingClientApiKeyUsage::default()); + let service = DefaultExecutionService::new( + RuntimeSnapshotHandle::new(start_snapshot()), + store.clone(), + ProviderRegistry::new([ + Arc::new(NativeCatalogProvider { fail: false }) as Arc + ]) + .unwrap(), + Arc::new(UnusedAdmissions), + Arc::new(UnusedContinuation), + usage.clone(), + Arc::new(crate::RecordingDiagnostics::default()), + ); + let repeated = ExtensionCallScope::default() + .extending("entry".into()) + .unwrap(); + let exhausted = (0..ExtensionCallScope::MAXIMUM_DEPTH) + .fold(ExtensionCallScope::default(), |scope, index| { + scope.extending(format!("ancestor-{index}")).unwrap() + }); + for extension_scope in [repeated, exhausted] { + let client = NestedModelExecutionPort::resolve_client( + &service, + ClientApiKeyId::new("key_start_test").unwrap(), + None, + ) + .await + .unwrap(); + let result = NestedModelExecutionPort::start_bound( + &service, + BoundModelExecutionRequest { + client: client.clone(), + initiating_plugin_instance_id: "entry".into(), + cancellation: CancellationToken::default(), + extension_scope, + public_model: PublicModelId::new("gpt-start").unwrap(), + operation: start_operation(), + metadata: ExecutionRequestMetadata { + protocol: "openai".into(), + endpoint: "host.model".into(), + transport: ClientTransport::InternalPlugin, + stream: false, + client_ip: None, + user_agent: None, + previous_response_id: None, + }, + provider: None, + account: None, + }, + ) + .await; + assert!(matches!(result, Err(error) if error.kind() == GatewayErrorKind::PolicyDenied)); + let models = NestedModelExecutionPort::models( + &service, + client, + CancellationToken::default(), + "codex".into(), + "0.154.0".into(), + ) + .await + .unwrap(); + assert_eq!(models, vec![PublicModelId::new("gpt-start").unwrap()]); + } + assert!(usage.recorded().is_empty()); + assert!(!store.touched.load(Ordering::SeqCst)); + }); +} + +struct ParentAwareAuthentication { + calls: AtomicUsize, + mapped: bool, +} + +impl FrontendAuthenticationPlan for ParentAwareAuthentication { + fn authenticate<'a>( + &'a self, + request: &'a ClientAuthenticationRequest, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.calls.fetch_add(1, Ordering::SeqCst); + let origin = request.invocation().expect("parent invocation"); + assert_eq!(origin.request_id, "http-request"); + assert_eq!(origin.call_id, "http-call"); + assert!(origin.extensions.contains("authentication-plugin")); + assert!(origin.effects.is_some()); + assert!(!origin.cancellation.is_cancelled()); + assert_eq!(request.request_context().unwrap().uri, "/plugin/proxy"); + assert!(request.settings().is_some()); + Ok(FrontendAuthenticationDecision::Authenticated { + principal: "external-principal".into(), + }) + }) + } + + fn client_key_id(&self, principal: &str) -> Option { + assert_eq!(principal, "external-principal"); + self.mapped + .then(|| ClientApiKeyId::new("key_start_test").unwrap()) + } + + fn exclusive(&self) -> bool { + true + } +} + +#[test] +fn verification_keeps_parent_origin_and_requires_the_frontend_principal_mapping() { + block_on(async { + for mapped in [true, false] { + let authentication = FrontendAuthenticationExtensionIndex::default(); + let plan = Arc::new(ParentAwareAuthentication { + calls: AtomicUsize::new(0), + mapped, + }); + let snapshot = start_snapshot().with_extensions(Some(authentication_generation( + &authentication, + "identity-verification", + plan.clone(), + ))); + let store = Arc::new(TrackingExecutionStore::default()); + let usage = Arc::new(RecordingClientApiKeyUsage::default()); + let service = DefaultExecutionService::new( + RuntimeSnapshotHandle::new(snapshot), + store.clone(), + ProviderRegistry::default(), + Arc::new(UnusedAdmissions), + Arc::new(UnusedContinuation), + usage.clone(), + Arc::new(crate::RecordingDiagnostics::default()), + ) + .with_frontend_authentication(authentication); + let settings = service.request_settings().unwrap(); + let origin = InvocationOrigin { + request_id: "http-request".into(), + call_id: "http-call".into(), + cancellation: CancellationToken::default(), + extensions: plugin_origin_scope() + .extending("authentication-plugin".into()) + .unwrap(), + plan: None, + effects: Some(Arc::new( + gateway_core::engine::nested::ExecutionEffects::default(), + )), + }; + let mut cancelled_origin = origin.clone(); + cancelled_origin.cancellation = CancellationToken::default(); + cancelled_origin.cancellation.cancel(); + let cancelled = NestedModelExecutionPort::verify_client( + &service, + ClientAuthenticationRequest::new("External fixture") + .unwrap() + .with_invocation(cancelled_origin) + .with_settings(settings.clone()), + ) + .await; + assert!(matches!(cancelled, Err(error) if error.kind() == GatewayErrorKind::Cancelled)); + assert_eq!(plan.calls.load(Ordering::SeqCst), 0); + let result = NestedModelExecutionPort::verify_client( + &service, + ClientAuthenticationRequest::new("External fixture") + .unwrap() + .with_invocation(origin) + .with_settings(settings.clone()), + ) + .await; + if mapped { + let client = result.unwrap(); + assert_eq!(client.policy().key_id().as_str(), "key_start_test"); + assert_eq!(client.request_settings(), Some(&settings)); + } else { + assert!( + matches!(result, Err(error) if error.kind() == GatewayErrorKind::Unauthorized) + ); + } + assert_eq!(plan.calls.load(Ordering::SeqCst), 1); + assert!(usage.recorded().is_empty()); + assert!(!store.touched.load(Ordering::SeqCst)); + } + }); +} diff --git a/backend/crates/gateway-core/tests/engine/execution/mod.rs b/backend/crates/gateway-core/tests/engine/execution/mod.rs index 3e57bbe8c..55f6f353c 100644 --- a/backend/crates/gateway-core/tests/engine/execution/mod.rs +++ b/backend/crates/gateway-core/tests/engine/execution/mod.rs @@ -1,5 +1,6 @@ //! 执行服务的准入、预算、请求记录与终结行为测试 +mod identity; mod outbound; use gateway_core::routing::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}; @@ -1215,9 +1216,7 @@ fn plugin_origin_scope() -> gateway_core::engine::extensions::ExtensionCallScope #[test] fn explicit_key_model_execution_preserves_http_origin_and_leaves_background_origin_unknown() { - use gateway_core::engine::nested::{ - BoundModelExecutionBinding, BoundModelExecutionRequest, NestedModelExecutionPort, - }; + use gateway_core::engine::nested::{BoundModelExecutionRequest, NestedModelExecutionPort}; block_on(async { for has_origin in [true, false] { let store = Arc::new(TrackingExecutionStore::default()); @@ -1234,11 +1233,17 @@ fn explicit_key_model_execution_preserves_http_origin_and_leaves_background_orig Arc::new(crate::RecordingDiagnostics::default()), ); let input = request(&service, ClientTransport::InternalPlugin); - let context = NestedModelExecutionPort::bind( + let client = NestedModelExecutionPort::resolve_client( &service, - BoundModelExecutionBinding { - settings: None, - client_key_id: ClientApiKeyId::new("key_start_test").unwrap(), + ClientApiKeyId::new("key_start_test").unwrap(), + None, + ) + .await + .unwrap(); + let mut started = NestedModelExecutionPort::start_bound( + &service, + BoundModelExecutionRequest { + client, initiating_plugin_instance_id: "origin-fixture".into(), cancellation: Default::default(), extension_scope: if has_origin { @@ -1246,14 +1251,6 @@ fn explicit_key_model_execution_preserves_http_origin_and_leaves_background_orig } else { Default::default() }, - }, - ) - .await - .unwrap(); - let mut started = NestedModelExecutionPort::start_bound( - &service, - BoundModelExecutionRequest { - context: context.clone(), public_model: input.public_model, operation: input.operation, metadata: input.metadata, diff --git a/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json b/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json index 6771f7d4a..c090f8781 100644 --- a/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json +++ b/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json @@ -1,7 +1,7 @@ { "schema_version": 2, "manifest_schema_versions": [2], - "protocol_versions": [6], + "protocol_versions": [7], "capabilities": [ { "capability": "frontend_authentication", "versions": [2] }, { "capability": "scheduler", "versions": [3] }, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs b/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs index 4268dd024..fbbba5ac7 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs @@ -195,21 +195,29 @@ impl FrontendAuthenticationPlan for PluginFrontendAuthenticationPlan { let mut context = session.context(Stage::Authentication, session.budget(Stage::Authentication)); context.request = crate::adapter::request_context(request.request_context()); + context.request_id = request.invocation().map(|origin| origin.request_id.clone()); let _scope = callbacks - .prepare_frontend_authentication(&context, request.shared_request_context()) + .prepare_frontend_authentication(&context, request) .map_err(|_| FrontendAuthenticationError)?; let input = FrontendAuthenticationRequest { authorization: request.authorization().to_owned(), }; - let reply = session - .call( - "frontend_auth.authenticate", - context, - serde_json::json!({}), - serde_json::to_vec(&input).map_err(|_| FrontendAuthenticationError)?, - ) - .await - .map_err(|_| FrontendAuthenticationError)?; + let invocation = session.call( + "frontend_auth.authenticate", + context, + serde_json::json!({}), + serde_json::to_vec(&input).map_err(|_| FrontendAuthenticationError)?, + ); + let reply = if let Some(origin) = request.invocation() { + tokio::select! { + biased; + () = origin.cancellation.cancelled() => return Err(FrontendAuthenticationError), + reply = invocation => reply, + } + } else { + invocation.await + } + .map_err(|_| FrontendAuthenticationError)?; let result: FrontendAuthenticationResult = serde_json::from_slice(&reply.payload).map_err(|_| FrontendAuthenticationError)?; match result { diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs index 80802116a..fa5a0aa13 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs @@ -51,7 +51,9 @@ pub(super) async fn invoke( request_id: context.request_id.clone(), call_id: context.call_id.clone(), cancellation: context.cancellation.clone(), - plan, + plan: Some(plan), + extensions: context.extensions.clone(), + effects: None, }), context.extensions.clone(), invocation.clone(), diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs index 5543fb890..b73a04a44 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs @@ -38,7 +38,9 @@ pub(super) async fn invoke( request_id: context.request_id.clone(), call_id: context.call_id.clone(), cancellation: context.cancellation.clone(), - plan: context.plan.clone(), + plan: Some(context.plan.clone()), + extensions: context.extensions.clone(), + effects: None, }), context.extensions.clone(), invocation, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs index 129d85ccf..88c9112ec 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs @@ -56,7 +56,9 @@ pub(super) async fn invoke( request_id: context.connection_id.clone(), call_id: call.resource_scope_id.clone(), cancellation: context.cancellation.clone(), - plan, + plan: Some(plan), + extensions: context.extensions.clone(), + effects: None, }), context.extensions, invocation.clone(), diff --git a/backend/crates/gateway-plugin/runtime/src/callback/accounts.rs b/backend/crates/gateway-plugin/runtime/src/callback/accounts.rs index 6782dcf84..46d8b6bfc 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/accounts.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/accounts.rs @@ -66,6 +66,32 @@ impl PluginAccounts { Self { slot } } + pub(super) async fn http_proxy( + &self, + account_id: String, + credential_revision: u64, + ) -> Result, PluginFault> { + let id = ProviderAccountId::new(account_id).map_err(|_| invalid())?; + let expected = Revision::new(credential_revision).map_err(|_| invalid())?; + let account = self + .slot + .upgrade() + .map_err(map_admin_error)? + .get_runtime(&id) + .await + .map_err(map_admin_error)?; + if !account.enabled { + return Err(denied()); + } + if account.credential_revision != expected { + return Err(PluginFault::new( + ErrorCode::Conflict, + "account credential revision has changed", + )); + } + Ok(account.effective_outbound_proxy().cloned()) + } + pub(super) async fn call( &self, context: &CallContext, diff --git a/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs index de188d0b9..94acff463 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/call_resources.rs @@ -16,9 +16,8 @@ pub(super) struct CallResources { pub(super) cancellation: CancellationToken, pub(super) scope: Arc, pub(super) http_resources: Arc, - pub(super) model_bindings: tokio::sync::Mutex< - BTreeMap, - >, + pub(super) client_identities: + tokio::sync::Mutex>, pub(super) state: Mutex, } @@ -32,13 +31,18 @@ impl CallResources { Self { deadline: deadline.clone(), operation_timeout: Duration::from_millis(context.timeout_ms), - cancellation: CancellationToken::new(), + cancellation: scope + .origin + .as_ref() + .map_or_else(CancellationToken::new, |origin| { + origin.cancellation.child_token() + }), scope, http_resources: middleware .as_ref() .and_then(|middleware| middleware.http_resources()) .unwrap_or_default(), - model_bindings: tokio::sync::Mutex::new(BTreeMap::new()), + client_identities: tokio::sync::Mutex::new(BTreeMap::new()), state: Mutex::new(CallState { middleware, ..CallState::default() @@ -46,6 +50,40 @@ impl CallResources { } } + /// 子调用共享受信来源；身份解析与执行准入仍留在各能力 owner + pub(super) fn child_origin( + &self, + context: &CallContext, + ) -> Result { + let parent = self.scope.origin.as_ref(); + Ok(gateway_core::engine::extensions::InvocationOrigin { + request_id: parent.map_or_else( + || { + context + .request_id + .clone() + .unwrap_or_else(|| context.resource_scope_id.clone()) + }, + |origin| origin.request_id.clone(), + ), + call_id: parent.map_or_else( + || format!("{}:{}", context.resource_scope_id, context.call_id), + |origin| origin.call_id.clone(), + ), + cancellation: self.cancellation.child_token(), + extensions: if self.scope.extension_scope.contains(&context.instance_id) { + self.scope.extension_scope.clone() + } else { + self.scope + .extension_scope + .extending(context.instance_id.clone()) + .ok_or_else(super::invalid)? + }, + plan: parent.and_then(|origin| origin.plan.clone()), + effects: self.scope.execution_effects.clone(), + }) + } + pub(super) fn close(&self) { self.cancellation.cancel(); let mut state = self @@ -114,7 +152,9 @@ impl CallResources { .unwrap_or_else(std::sync::PoisonError::into_inner) .middleware .clone(); - middleware.and_then(|middleware| middleware.request_settings()) + middleware + .and_then(|middleware| middleware.request_settings()) + .or_else(|| self.scope.settings.clone()) } /// 建立受管流后，网络操作各自计时；连接空闲时间不消耗后续操作的预算 diff --git a/backend/crates/gateway-plugin/runtime/src/callback/data.rs b/backend/crates/gateway-plugin/runtime/src/callback/data.rs index 6df0f493a..2bd4ff19a 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/data.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/data.rs @@ -98,7 +98,7 @@ impl PluginData { next_cursor: page.next_cursor.map(|id| id.as_str().to_owned()), }) } - data::QUOTA_GET | data::QUOTA_REFRESH => { + data::QUOTA_GET | data::QUOTA_REFRESH | data::QUOTA_DETAILS => { let accounts = self.accounts.upgrade().map_err(map_admin_error)?; let query: data::QuotaFactsQuery = serde_json::from_slice(payload).map_err(|_| invalid())?; @@ -109,6 +109,35 @@ impl PluginData { accounts.get_quota(&account_id).await } .map_err(map_admin_error)?; + if method == data::QUOTA_DETAILS { + return super::admin::encode(&data::QuotaDetails { + schema_version: 1, + account_id: account_id.as_str().to_owned(), + observed_at_ms: quota.observed_at.map(|time| time.timestamp_millis()), + limit_reached: quota.limit_reached, + windows: quota + .windows + .into_iter() + .map(|window| data::QuotaWindowDetails { + key: window.key, + group: window.group, + label: window.label, + limit_id: window.limit_id, + limit_name: window.limit_name, + role: window.role.map(|role| role.as_str().to_owned()), + window_seconds: window.window_seconds, + used_percent: window.used_percent, + reset_at_ms: window.reset_at.map(|time| time.timestamp_millis()), + limit_reached: window.limit_reached, + }) + .collect(), + credits: quota.credits.map(|credits| data::QuotaCredits { + has_credits: credits.has_credits, + unlimited: credits.unlimited, + balance: credits.balance, + }), + }); + } serde_json::to_vec(&data::QuotaFacts { schema_version: 1, account_id: account_id.as_str().to_owned(), diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http.rs b/backend/crates/gateway-plugin/runtime/src/callback/http.rs index 206435f98..8ea083253 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http.rs @@ -29,16 +29,42 @@ impl HttpCallbacks<'_> { self.request_limits.body_bytes(self.call) } + pub(super) async fn account( + &self, + accounts: &super::accounts::PluginAccounts, + params: serde_json::Value, + payload: Vec, + ) -> Result { + let request: wire::AccountHttpRequest = decode(params)?; + let maximum_body_bytes = self.body_limit(&payload)?; + let proxy = accounts + .http_proxy(request.account_id, request.credential_revision) + .await?; + self.send( + HttpRequest { + method: request.request.method, + url: request.request.url, + headers: request.request.headers, + body: payload, + maximum_body_bytes, + }, + true, + None, + proxy.as_ref(), + ) + .await + } + pub(super) async fn send( &self, request: HttpRequest, streaming: bool, upstream_purpose: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result { let Self { client, network, - scope, call, maximum_payload, .. @@ -48,9 +74,7 @@ impl HttpCallbacks<'_> { } let timeout = call.timeout()?; let attempt = call.start_upstream(upstream_purpose)?; - let response = client - .open(request, scope.proxy.as_ref(), network, timeout) - .await; + let response = client.open(request, proxy, network, timeout).await; attempt.finish( response .as_ref() @@ -127,6 +151,7 @@ impl HttpCallbacks<'_> { }, method == "host.http.do_stream", None, + self.scope.proxy.as_ref(), ) .await } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs index cc045d2f3..bfaa1d9f2 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/http_dispatch.rs @@ -36,25 +36,23 @@ pub(super) async fn call( request.timeout_ms, &context.instance_id, )?; - let parent = call.scope.origin.as_ref(); + call.scope.check_child_depth(&context.instance_id)?; + let origin = call.child_origin(context)?; let child = core::Context { plugin_instance_id: Some(context.instance_id.clone()), - request_id: context - .request_id - .clone() - .unwrap_or_else(|| context.resource_scope_id.clone()), + request_id: origin.request_id, call_id: uuid::Uuid::now_v7().to_string(), - parent_call_id: Some(parent.map_or_else( - || format!("{}:{}", context.resource_scope_id, context.call_id), - |parent| parent.call_id.clone(), - )), - extensions: call.scope.child_extensions(&context.instance_id)?, - // 句柄可以在父 RPC End 后交给 HTTP 发送端；生命周期属于父调用和响应正文，不能在 RPC End 提前取消 - cancellation: parent.map_or_else( - || call.cancellation.child_token(), - |parent| parent.cancellation.child_token(), - ), - plan: parent.map(|parent| parent.plan.clone()), + parent_call_id: Some(origin.call_id), + extensions: origin.extensions, + // 转交正文后由消费端持有原请求寿命，父 RPC End 只撤销新业务分派 + cancellation: call + .scope + .origin + .as_ref() + .map_or(origin.cancellation, |parent| { + parent.cancellation.child_token() + }), + plan: origin.plan, }; let body = call.http_resources.body(request.body, payload).await?; let mut http = core::Request::new(body); diff --git a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs index 6b2f8f6e4..2e69f4305 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/mod.rs @@ -221,7 +221,9 @@ impl PluginCallbacks { .clone() .unwrap_or_else(|| context.resource_scope_id.clone()), cancellation, - plan, + extensions: scope.extension_scope.clone(), + effects: scope.execution_effects.clone(), + plan: Some(plan), }); self.register_scope(context, scope) } @@ -278,7 +280,7 @@ impl PluginCallbacks { pub(crate) fn prepare_frontend_authentication( &self, context: &CallContext, - request: Option>, + request: &gateway_core::engine::authentication::ClientAuthenticationRequest, ) -> Result, PluginFault> { if context.stage != gateway_plugin_sdk::Stage::Authentication || context.account_id.is_some() @@ -287,7 +289,16 @@ impl PluginCallbacks { { return Err(denied()); } - self.prepare_scope(context, None, request) + let mut scope = CallbackScope::for_call(context); + scope.settings = request.settings().cloned(); + if let Some(origin) = request.invocation() { + scope.extension_scope = origin.extensions.clone(); + scope.execution_effects = origin.effects.clone(); + scope.origin = Some(origin.clone()); + } else if let Some(request) = request.shared_request_context() { + scope.extension_scope = scope.extension_scope.with_request(request); + } + self.register_scope(context, scope) } pub(crate) async fn save_command_account( @@ -500,7 +511,12 @@ impl CallbackHandler for PluginCallbacks { ) { return keys.call(&context, &method, params, &payload).await; } - if method == "host.models.list" { + if matches!( + method.as_str(), + "host.models.list" + | gateway_plugin_sdk::call::data::CLIENT_VERIFY + | gateway_plugin_sdk::call::data::ACCESS_CHECK + ) { return models.call(&context, &call, &method, params, payload).await; } if matches!( @@ -524,6 +540,7 @@ impl CallbackHandler for PluginCallbacks { | gateway_plugin_sdk::call::data::KEYS_GET | gateway_plugin_sdk::call::data::QUOTA_GET | gateway_plugin_sdk::call::data::QUOTA_REFRESH + | gateway_plugin_sdk::call::data::QUOTA_DETAILS ) { return data.call(&method, params, &payload).await; } @@ -543,6 +560,9 @@ impl CallbackHandler for PluginCallbacks { maximum_payload, request_limits: &request_limits, }; + if method == gateway_plugin_sdk::call::host::ACCOUNT_HTTP { + return http.account(&accounts, params, payload).await; + } if method.starts_with("host.upstream.") { return upstream::dispatch(&http, &method, params, payload).await; } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/model.rs b/backend/crates/gateway-plugin/runtime/src/callback/model.rs index 168263db4..5f8dae51a 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/model.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/model.rs @@ -1,4 +1,4 @@ -//! 嵌套模型执行；只驱动 Core 会话，不拥有路由、重试或计费规则 +//! 复用 Core 的身份查询与模型执行入口；不拥有权限、路由、重试或计费规则 use std::sync::{Arc, OnceLock, Weak}; @@ -8,12 +8,12 @@ use gateway_core::{ account::ProviderAccountId, engine::{ CommitRequirement, ModelRequestId, + authentication::ClientAuthenticationRequest, execution::{ - BoundModelExecutionContext, ClientTransport, ExecutionRequestMetadata, ExecutionSession, + AuthenticatedClient, ClientTransport, ExecutionRequestMetadata, ExecutionSession, }, nested::{ - BoundModelExecutionBinding, BoundModelExecutionRequest, NestedModelExecutionPort, - NestedModelExecutionRequest, + BoundModelExecutionRequest, NestedModelExecutionPort, NestedModelExecutionRequest, }, }, event::ProviderEvent, @@ -23,10 +23,11 @@ use gateway_core::{ RawJsonPayload, StandaloneSearchRequest, }, policy::ClientApiKeyId, - routing::PublicModelId, + routing::{PublicModelId, UpstreamModelId}, }; use gateway_plugin_sdk::{ CallContext, ErrorCode, PluginFault, + call::data, call::host::{ ModelEventBatch, ModelExecuteRequest, ModelExecuteResult, ModelListRequest, ModelListResult, ModelOperation, ModelStreamCloseRequest, ModelStreamReadRequest, @@ -38,6 +39,7 @@ use super::{CallResources, denied, invalid}; use crate::RpcReply; const MAX_MODEL_STREAMS_PER_CALL: usize = 16; +const MAX_CLIENT_IDENTITIES_PER_CALL: usize = 16; pub(crate) struct PluginModelPortSlot { port: OnceLock>, @@ -89,6 +91,64 @@ impl PluginModels { payload: Vec, ) -> Result { match method { + data::CLIENT_VERIFY => { + if params != serde_json::json!({}) { + return Err(invalid()); + } + let request: data::VerifyClientRequest = + serde_json::from_slice(&payload).map_err(|_| invalid())?; + let mut authentication = ClientAuthenticationRequest::new(request.authorization) + .map_err(|_| invalid())? + .with_invocation(call.child_origin(context)?); + if let Some(settings) = call.request_settings() { + authentication = authentication.with_settings(settings); + } + let port = self.slot.upgrade()?; + call.observe_dispatch()?; + let client = port + .verify_client(authentication) + .await + .map_err(super::error::gateway)?; + super::admin::encode(&data::VerifiedClient { + client_key_id: client.policy().key_id().as_str().to_owned(), + }) + } + data::ACCESS_CHECK => { + if params != serde_json::json!({}) { + return Err(invalid()); + } + let request: data::AccessQuery = + serde_json::from_slice(&payload).map_err(|_| invalid())?; + let port = self.slot.upgrade()?; + let client = self.identity(&port, call, request.client_key_id).await?; + let scope = client.policy().account_scope(); + let model = request + .upstream_model + .map(UpstreamModelId::new) + .transpose() + .map_err(|_| invalid())?; + let accounts = request + .account_ids + .into_iter() + .map(|account_id| { + let account = + ProviderAccountId::new(account_id.clone()).map_err(|_| invalid())?; + let allowed = model.as_ref().map_or_else( + || scope.allows(&account), + |model| scope.allows_model(&account, model.as_str()), + ); + Ok(data::AccountAccessFacts { + account_id, + allowed, + }) + }) + .collect::, PluginFault>>()?; + super::admin::encode(&data::AccessFacts { + schema_version: 1, + client_key_id: client.policy().key_id().as_str().to_owned(), + accounts, + }) + } "host.models.list" => { if !payload.is_empty() { return Err(invalid()); @@ -102,19 +162,14 @@ impl PluginModels { return Err(invalid()); } let port = self.slot.upgrade()?; - let bound = port - .bind(BoundModelExecutionBinding { - settings: call.request_settings(), - client_key_id: ClientApiKeyId::new(request.client_key_id) - .map_err(|_| invalid())?, - initiating_plugin_instance_id: context.instance_id.clone(), - cancellation: call.cancellation.child_token(), - extension_scope: call.scope.extension_scope.clone(), - }) - .await - .map_err(super::error::gateway)?; + let client = self.identity(&port, call, request.client_key_id).await?; let models = port - .models(bound, request.protocol, request.client_version) + .models( + client, + call.cancellation.clone(), + request.protocol, + request.client_version, + ) .await .map_err(super::error::gateway)?; encode_result(ModelListResult { @@ -141,9 +196,12 @@ impl PluginModels { let port = self.slot.upgrade()?; let stream = method == "host.model.execute_stream"; let started = if let Some(client_key_id) = request.client_key_id.take() { - let bound = self.binding(&port, context, call, client_key_id).await?; - check_size(bound.request_body_limit())?; - port.start_bound(self.bound_request(request, payload, stream, bound)?) + let client = self.identity(&port, call, client_key_id).await?; + check_size(client.snapshot().responses_max_decompressed_body_bytes())?; + let request = self.request_parts(request, payload, stream)?; + let request = Self::bound_request(context, call, request, client); + call.observe_dispatch()?; + port.start_bound(request) .await .map_err(super::error::gateway)? } else { @@ -237,42 +295,36 @@ impl PluginModels { } } - async fn binding( + async fn identity( &self, port: &Arc, - context: &CallContext, call: &CallResources, client_key_id: String, - ) -> Result { + ) -> Result { let key = ClientApiKeyId::new(client_key_id.clone()).map_err(|_| invalid())?; let settings = call.request_settings(); - let mut bindings = call.model_bindings.lock().await; - if let Some(bound) = bindings + let mut identities = call.client_identities.lock().await; + if let Some(client) = identities .get(&client_key_id) - .filter(|bound| bound.request_settings() == settings.as_ref()) + .filter(|client| client.request_settings() == settings.as_ref()) { - return Ok(bound.clone()); + return Ok(client.clone()); } - if !bindings.contains_key(&client_key_id) && bindings.len() >= MAX_MODEL_STREAMS_PER_CALL { + if !identities.contains_key(&client_key_id) + && identities.len() >= MAX_CLIENT_IDENTITIES_PER_CALL + { return Err(PluginFault::new( ErrorCode::Capacity, - "model identity capacity is exhausted", + "client identity capacity is exhausted", )); } - let bound = port - .bind(BoundModelExecutionBinding { - settings, - client_key_id: key, - initiating_plugin_instance_id: context.instance_id.clone(), - cancellation: call.cancellation.child_token(), - extension_scope: call.scope.extension_scope.clone(), - }) + let client = port + .resolve_client(key, settings) .await .map_err(super::error::gateway)?; - // 调用持有身份直到 RPC 结束，流式执行不会因单次 callback 返回而提前取消； - // 同一 Key 复用身份快照，每次执行由 Core 创建独立调用图 - bindings.insert(client_key_id, bound.clone()); - Ok(bound) + // 查询与执行共用调用内快照，执行链和子执行取消只在真正启动时建立 + identities.insert(client_key_id, client.clone()); + Ok(client) } fn request( @@ -310,21 +362,22 @@ impl PluginModels { } fn bound_request( - &self, - request: ModelExecuteRequest, - body: Vec, - stream: bool, - context: BoundModelExecutionContext, - ) -> Result { - let request = self.request_parts(request, body, stream)?; - Ok(BoundModelExecutionRequest { - context, + context: &CallContext, + call: &CallResources, + request: ModelRequestParts, + client: AuthenticatedClient, + ) -> BoundModelExecutionRequest { + BoundModelExecutionRequest { + client, + initiating_plugin_instance_id: context.instance_id.clone(), + cancellation: call.cancellation.clone(), + extension_scope: call.scope.extension_scope.clone(), public_model: request.public_model, operation: request.operation, metadata: request.metadata, provider: request.provider, account: request.account, - }) + } } fn request_parts( diff --git a/backend/crates/gateway-plugin/runtime/src/callback/scope.rs b/backend/crates/gateway-plugin/runtime/src/callback/scope.rs index 475847cae..ef240f2b8 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/scope.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/scope.rs @@ -7,14 +7,7 @@ use gateway_core::{ }; use gateway_plugin_sdk::{CallContext, Stage}; -/// 所有中间件共享的父调用身份；不混入协议操作或服务参数 -#[derive(Clone)] -pub(crate) struct InvocationContext { - pub request_id: String, - pub call_id: String, - pub cancellation: gateway_core::lifecycle::CancellationToken, - pub plan: gateway_core::engine::middleware::FrozenMiddlewarePlan, -} +pub(crate) use gateway_core::engine::extensions::InvocationOrigin as InvocationContext; /// 一次回调的调用链、账号与网络事实，由操作持有到完成 pub(crate) struct CallbackScope { @@ -25,7 +18,8 @@ pub(crate) struct CallbackScope { credential_revision: Option, pub(super) proxy: Option, pub(super) extension_scope: gateway_core::engine::extensions::ExtensionCallScope, - execution_effects: Option>, + pub(super) execution_effects: Option>, + pub(super) settings: Option, pub(super) upstream: Option>, } @@ -37,14 +31,13 @@ impl CallbackScope { } } - pub(super) fn child_extensions( + pub(super) fn check_child_depth( &self, instance_id: &str, - ) -> Result - { + ) -> Result<(), gateway_plugin_sdk::PluginFault> { // 上游适配器在进入回调前已登记当前实例；主动子调用沿用该集合 if self.extension_scope.contains(instance_id) { - return Ok(self.extension_scope.clone()); + return Ok(()); } if self.extension_scope.len() >= gateway_core::engine::extensions::ExtensionCallScope::MAXIMUM_DEPTH @@ -54,9 +47,7 @@ impl CallbackScope { "child call depth exceeded", )); } - self.extension_scope - .extending(instance_id.to_owned()) - .ok_or_else(super::invalid) + Ok(()) } pub(super) fn new(context: &CallContext, proxy: Option) -> Self { @@ -69,6 +60,7 @@ impl CallbackScope { proxy, extension_scope: Default::default(), execution_effects: None, + settings: None, upstream: None, } } diff --git a/backend/crates/gateway-plugin/runtime/src/callback/services.rs b/backend/crates/gateway-plugin/runtime/src/callback/services.rs index 5b4529307..7eda669df 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/services.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/services.rs @@ -63,27 +63,18 @@ impl ServicePorts { ) })?; let call = registry.call(&request.operation, request.input); - let scope = &resources.scope; - let origin = scope.origin.as_ref(); + resources.scope.check_child_depth(&context.instance_id)?; + let origin = resources.child_origin(context)?; let parent = gateway_admin::public_service::Origin { - request_id: origin.map_or_else( - || { - context - .request_id - .clone() - .unwrap_or_else(|| context.resource_scope_id.clone()) - }, - |origin| origin.request_id.clone(), - ), - call_id: origin.map_or_else( - || format!("{}:{}", context.resource_scope_id, context.call_id), - |origin| origin.call_id.clone(), - ), - cancellation: resources.cancellation.clone(), - extensions: scope.child_extensions(&context.instance_id)?, - plan: origin.map_or(gateway_admin::public_service::Plan::Current, |origin| { - gateway_admin::public_service::Plan::Frozen(Some(origin.plan.clone())) - }), + request_id: origin.request_id, + call_id: origin.call_id, + cancellation: origin.cancellation, + extensions: origin.extensions, + plan: origin + .plan + .map_or(gateway_admin::public_service::Plan::Current, |plan| { + gateway_admin::public_service::Plan::Frozen(Some(plan)) + }), }; // 即使终端只读，服务中间件也可提交写入；不存在的操作不会进入调用链 if registry diff --git a/backend/crates/gateway-plugin/runtime/src/callback/upstream/mod.rs b/backend/crates/gateway-plugin/runtime/src/callback/upstream/mod.rs index 8b0308466..472cc47d6 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/upstream/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/upstream/mod.rs @@ -197,8 +197,13 @@ pub(super) async fn dispatch( let request = serde_json::from_value(params).map_err(|_| invalid())?; let maximum_body_bytes = http.body_limit(&payload)?; let (request, purpose) = managed.request(request, payload, maximum_body_bytes)?; - http.send(request, method.ends_with("do_stream"), Some(purpose)) - .await + http.send( + request, + method.ends_with("do_stream"), + Some(purpose), + scope.proxy.as_ref(), + ) + .await } "host.upstream.http.stream_read" | "host.upstream.http.stream_close" => { http.dispatch( diff --git a/backend/crates/gateway-plugin/runtime/src/rpc/dispatch.rs b/backend/crates/gateway-plugin/runtime/src/rpc/dispatch.rs index 436060a8a..959403747 100644 --- a/backend/crates/gateway-plugin/runtime/src/rpc/dispatch.rs +++ b/backend/crates/gateway-plugin/runtime/src/rpc/dispatch.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use gateway_plugin_sdk::{ ErrorCode, Frame, Message, PluginFault, - call::middleware::NEXT_METHOD, client::{EncodedFrame, read_frame}, }; use tokio::{ @@ -100,6 +99,13 @@ async fn read_loop( } }; let result = match frame.message { + Message::Suspend { id, epoch } if frame.payload.is_empty() => shared.suspend(id, epoch), + Message::Resume { id, epoch } if frame.payload.is_empty() => { + shared.request_resume(id, epoch) + } + Message::CallbackConsumed { id, parent_id } if frame.payload.is_empty() => { + shared.consume_callback(parent_id, id) + } Message::Cancelled { id } if frame.payload.is_empty() => { shared.acknowledge_cancellation(id) } @@ -125,26 +131,17 @@ async fn read_loop( Err(RpcError::Protocol(None)) } else if let Some(context) = shared.context(parent_id) { last_callback = id; - if let Some(permit) = shared.try_callback_slot(&method) { + if let Some(permit) = shared.try_callback_slot() { let handler = Arc::clone(&callbacks); let response = Arc::clone(&shared); let (ready, started) = oneshot::channel(); - let pause = (method == NEXT_METHOD) - .then(|| shared.pause_next(parent_id)) - .flatten(); + let callback_method = method.clone(); let task = tokio::spawn(async move { // 先把任务归属登记到父调用，再允许执行宿主操作 if started.await.is_err() { return; } let result = handler.call(context, method, params, frame.payload).await; - if pause.is_some() - && let Err(error) = response.resume_next(parent_id).await - { - response.cancel(parent_id, error); - return; - } - drop(pause); let reply = match result { Ok(reply) => Frame { message: Message::Result { @@ -166,9 +163,14 @@ async fn read_loop( ), })), } - drop(permit); }); - if shared.track_callback(parent_id, task.abort_handle()) { + if shared.track_callback( + parent_id, + id, + &callback_method, + task.abort_handle(), + permit, + ) { let _ = ready.send(()); } } else { diff --git a/backend/crates/gateway-plugin/runtime/src/rpc/session.rs b/backend/crates/gateway-plugin/runtime/src/rpc/session.rs index d38f33d6e..8798c20da 100644 --- a/backend/crates/gateway-plugin/runtime/src/rpc/session.rs +++ b/backend/crates/gateway-plugin/runtime/src/rpc/session.rs @@ -46,7 +46,7 @@ impl Default for RpcLimits { maximum_executing_calls: 64, maximum_queued_calls: 64, maximum_queue_timeout: Duration::from_secs(30), - maximum_callbacks: 64, + maximum_callbacks: 320, handshake_timeout: Duration::from_secs(5), maximum_stderr_bytes: 64 * 1024, } @@ -225,7 +225,6 @@ pub(crate) struct Shared { pub process: ProcessControl, pub stopped: watch::Sender, callback_slots: Arc, - next_slots: Arc, execution: ExecutionCapacity, cancellations_changed: Notify, } @@ -245,14 +244,48 @@ struct PendingCall { context: CallContext, result: Option>>, stream: Option, - callbacks: Vec, + scheduling: Option, + callbacks: BTreeMap, +} + +// 调度状态只属于父任务；回调完成与恢复授权分别推进 +enum Scheduling { + Running { + epoch: u64, + }, + Suspended { + epoch: u64, + pause: Option, + }, + Resuming { + epoch: u64, + pause: Option, + task: tokio::task::AbortHandle, + }, +} + +struct CallbackTask { + task: tokio::task::AbortHandle, + pauses_budget: bool, + _permit: OwnedSemaphorePermit, +} + +impl Drop for CallbackTask { + fn drop(&mut self) { + self.task.abort(); + } } impl PendingCall { - fn finish(mut self, result: Result) { - for callback in self.callbacks.drain(..) { - callback.abort(); + fn stop_resumption(&mut self) { + if let Some(Scheduling::Resuming { task, .. }) = self.scheduling.take() { + task.abort(); } + } + + fn finish(mut self, result: Result) { + self.callbacks.clear(); + self.stop_resumption(); let stream = self.stream.take(); let sender = self.result.take(); // 先同步关闭调用资源，Core 才能在收到结果时看到全部发送副作用 @@ -318,12 +351,16 @@ impl Shared { { log_timeout(&call.context, "plugin", "initial_reply", "remote_timeout"); } + if result.is_ok() && !matches!(call.scheduling, Some(Scheduling::Running { .. })) { + return Err(RpcError::Protocol(None)); + } if call.stream.is_some() && result.is_ok() { if call.context.resource_stream { call.deadline.clear(); } let initial = call.result.take().ok_or(RpcError::Protocol(None))?; call.execution.take(); + call.scheduling.take(); call.handler.streaming(&call.context); let _ = initial.send(result); } else { @@ -404,43 +441,6 @@ impl Shared { .map(|call| call.context.clone()) } - pub fn pause_next(&self, id: u64) -> Option { - let mut state = self.state(); - let call = state - .calls - .get_mut(&id) - .filter(|call| call.cancellation.is_none())?; - let pause = call.deadline.pause()?; - call.execution.take(); - Some(pause) - } - - pub async fn resume_next(&self, id: u64) -> Result<(), RpcError> { - { - let state = self.state(); - let call = state - .calls - .get(&id) - .filter(|call| call.cancellation.is_none()) - .ok_or(RpcError::Cancelled)?; - // 已交付响应头的流只占在途资源，不重新占用首结果执行名额 - if call.result.is_none() { - return Ok(()); - } - } - let permit = self.execution.resume().await?; - let mut state = self.state(); - let call = state - .calls - .get_mut(&id) - .filter(|call| call.cancellation.is_none()) - .ok_or(RpcError::Cancelled)?; - if call.result.is_some() { - call.execution = Some(permit); - } - Ok(()) - } - pub fn is_cancelling(&self, id: u64) -> bool { self.state() .calls @@ -489,9 +489,8 @@ impl Shared { } call.cancellation = Some(error.clone()); call.handler.finish(&call.context); - for callback in call.callbacks.drain(..) { - callback.abort(); - } + call.callbacks.clear(); + call.stop_resumption(); if let Some(result) = call.result.take() { let _ = result.send(Err(error.clone())); } @@ -524,21 +523,134 @@ impl Shared { Ok(()) } - pub fn track_callback(&self, parent_id: u64, callback: tokio::task::AbortHandle) -> bool { - if let Some(call) = self - .state() + pub fn track_callback( + &self, + parent_id: u64, + id: u64, + method: &str, + task: tokio::task::AbortHandle, + permit: OwnedSemaphorePermit, + ) -> bool { + let mut state = self.state(); + let Some(call) = state .calls .get_mut(&parent_id) .filter(|call| call.cancellation.is_none()) + else { + task.abort(); + return false; + }; + let pauses_budget = method == gateway_plugin_sdk::call::middleware::NEXT_METHOD; + if !pauses_budget + && let Some(Scheduling::Suspended { pause, .. } | Scheduling::Resuming { pause, .. }) = + call.scheduling.as_mut() { - // 已完成回调的句柄不继续占用请求预算 - call.callbacks.retain(|handle| !handle.is_finished()); - call.callbacks.push(callback); - true - } else { - callback.abort(); - false + pause.take(); } + call.callbacks.insert( + id, + CallbackTask { + task, + pauses_budget, + _permit: permit, + }, + ); + true + } + + pub(super) fn consume_callback(&self, parent_id: u64, id: u64) -> Result<(), RpcError> { + if id == 0 || !id.is_multiple_of(2) || parent_id == 0 || parent_id.is_multiple_of(2) { + return Err(RpcError::Protocol(None)); + } + // 取消确认可以与消费确认交错；确认只能撤销资源，不能重新创建调用 + if let Some(call) = self.state().calls.get_mut(&parent_id) { + call.callbacks.remove(&id); + } + Ok(()) + } + + pub(super) fn suspend(&self, id: u64, epoch: u64) -> Result<(), RpcError> { + if id == 0 || id.is_multiple_of(2) || epoch == 0 { + return Err(RpcError::Protocol(None)); + } + let mut state = self.state(); + let Some(call) = state.calls.get_mut(&id) else { + return Ok(()); + }; + if call.cancellation.is_some() { + return Ok(()); + } + match call.scheduling.as_ref() { + Some(Scheduling::Running { epoch: previous }) + if previous.checked_add(1) == Some(epoch) => {} + _ => return Err(RpcError::Protocol(None)), + } + let pauses = !call.callbacks.is_empty() + && call + .callbacks + .values() + .all(|callback| callback.pauses_budget); + let pause = pauses.then(|| call.deadline.pause()).flatten(); + call.execution.take(); + call.scheduling = Some(Scheduling::Suspended { epoch, pause }); + Ok(()) + } + + pub(super) fn request_resume(self: &Arc, id: u64, epoch: u64) -> Result<(), RpcError> { + if id == 0 || id.is_multiple_of(2) || epoch == 0 { + return Err(RpcError::Protocol(None)); + } + let mut state = self.state(); + let Some(call) = state.calls.get_mut(&id) else { + return Ok(()); + }; + if call.cancellation.is_some() { + return Ok(()); + } + if !matches!(call.scheduling, Some(Scheduling::Suspended { epoch: current, .. }) if current == epoch) + { + return Err(RpcError::Protocol(None)); + } + let Some(Scheduling::Suspended { pause, .. }) = call.scheduling.take() else { + return Err(RpcError::Protocol(None)); + }; + let shared = Arc::clone(self); + let (ready, started) = oneshot::channel(); + let task = tokio::spawn(async move { + if started.await.is_err() { + return; + } + match shared.execution.resume().await { + Ok(permit) => shared.grant_resume(id, epoch, permit), + Err(error) => shared.cancel(id, error), + } + }); + call.scheduling = Some(Scheduling::Resuming { + epoch, + pause, + task: task.abort_handle(), + }); + let _ = ready.send(()); + Ok(()) + } + + fn grant_resume(&self, id: u64, epoch: u64, permit: OwnedSemaphorePermit) { + let mut state = self.state(); + let Some(call) = state + .calls + .get_mut(&id) + .filter(|call| call.cancellation.is_none()) + else { + return; + }; + if !matches!(call.scheduling, Some(Scheduling::Resuming { epoch: current, .. }) if current == epoch) + { + return; + } + call.execution = Some(permit); + call.scheduling = Some(Scheduling::Running { epoch }); + drop(state); + self.send_control(Frame::control(Message::Resumed { id, epoch })); } pub fn send_control(&self, frame: Frame) { @@ -555,23 +667,18 @@ impl Shared { } } - pub(super) fn try_callback_slot(&self, method: &str) -> Option { - let slots = if method == gateway_plugin_sdk::call::middleware::NEXT_METHOD { - &self.next_slots - } else { - &self.callback_slots - }; - slots.clone().try_acquire_owned().ok() + pub(super) fn try_callback_slot(&self) -> Option { + self.callback_slots.clone().try_acquire_owned().ok() } async fn wait_for_callbacks(&self, limits: RpcLimits) { - for (slots, maximum) in [ - (&self.callback_slots, limits.maximum_callbacks), - (&self.next_slots, limits.maximum_calls), - ] { - if let Ok(permits) = slots.clone().acquire_many_owned(maximum as u32).await { - drop(permits); - } + if let Ok(permits) = self + .callback_slots + .clone() + .acquire_many_owned(limits.maximum_callbacks as u32) + .await + { + drop(permits); } } } @@ -653,7 +760,7 @@ impl RpcSession { || limits.maximum_queued_calls > 256 || limits.maximum_queue_timeout.is_zero() || limits.maximum_callbacks == 0 - || limits.maximum_callbacks > 256 + || limits.maximum_callbacks > 512 || limits.maximum_buffered_body_bytes < 1024 || limits.maximum_buffered_body_bytes > 16 * 1024 * 1024 || limits.handshake_timeout.is_zero() @@ -771,7 +878,6 @@ impl RpcSession { process, stopped, callback_slots, - next_slots: Arc::new(Semaphore::new(limits.maximum_calls)), execution: ExecutionCapacity::new( limits.maximum_executing_calls, limits.maximum_queued_calls, @@ -902,7 +1008,8 @@ impl RpcSession { context, result: Some(result), stream, - callbacks: Vec::new(), + scheduling: Some(Scheduling::Running { epoch: 0 }), + callbacks: BTreeMap::new(), }, ); Ok(StartedCall { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/authentication.rs b/backend/crates/gateway-plugin/runtime/tests/generation/authentication.rs index e08c01510..f9660085a 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/authentication.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/authentication.rs @@ -273,3 +273,90 @@ async fn startup_failure_of_an_unrelated_plugin_preserves_native_authentication_ } close(environment, runtime, core).await; } + +#[tokio::test] +async fn callback_verification_reenters_the_authentication_instance_without_model_binding() { + use gateway_admin::{ + model::plugins::management::PluginManagementRequest, + ports::plugin_management::PluginManagement, + }; + use gateway_plugin_runtime::RpcLimits; + let Some(environment) = Environment::create().await else { + return; + }; + environment + .client_key("key_frontend", "sk-native-fixture") + .await; + let account = environment.account(None).await; + let marker = environment + .directory + .path() + .join("authenticated-facts.json"); + let facts = json!([ + {"method":"host.data.access.check", "query":{"client_key_id":"key_frontend", "account_ids":[account.as_str(), "acct_missing"]}}, + {"method":"host.models.list", "params":{"client_key_id":"key_frontend", "protocol":"codex", "client_version":"0.154.0"}} + ]); + let mut config = configuration( + json!({"outcome":"authenticated","principal":"external-user"}), + true, + "External controlled-fixture", + ); + config["sdk_session"] = json!(true); + config["frontend_authentication_marker"] = json!(marker); + config["auth_queries"] = facts.clone(); + config["management_registration"] = json!({"routes":[{"method":"GET","path":"facts","request_content_types":[],"response_content_types":["application/json"]}]}); + config["data_queries"] = json!([ + {"method":"host.client.verify", "query":{"authorization":"External controlled-fixture"}}, + facts[0].clone(), facts[1].clone() + ]); + environment.install_plugin(config).await; + let (runtime, core) = environment + .runtime_with_limits(RpcLimits { + maximum_executing_calls: 1, + maximum_queue_timeout: std::time::Duration::from_millis(500), + ..RpcLimits::default() + }) + .await; + let snapshot = core.snapshots().acquire().unwrap(); + let generation = snapshot.extensions().unwrap(); + let view = runtime.views(generation).await.unwrap().remove(0); + let response = runtime + .handle( + generation, + &view.target, + PluginManagementRequest { + request_context: None, + headers: Vec::new(), + method: "GET".into(), + path: "facts".into(), + query: String::new(), + content_type: None, + body: Vec::new(), + request_id: "verify-composition".into(), + }, + ) + .await + .unwrap(); + let results: Value = serde_json::from_slice(&response.body).unwrap(); + assert_eq!(results[0], json!({"client_key_id":"key_frontend"})); + let authenticated_facts: Value = + serde_json::from_slice(&std::fs::read(marker).unwrap()).unwrap(); + assert_eq!(authenticated_facts[0], results[1]); + assert_eq!( + authenticated_facts[0]["accounts"], + json!([ + {"account_id":account.as_str(),"allowed":true}, {"account_id":"acct_missing","allowed":false} + ]) + ); + assert_eq!(authenticated_facts[1], results[2]); + // 此替身只登记执行能力，没有原生目录；查询应成功返回空目录而不是执行绑定错误 + assert_eq!(results[2], json!({"models":[]})); + assert!( + environment + .bound_model_requests("key_frontend") + .await + .is_empty() + ); + drop(snapshot); + close(environment, runtime, core).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs index bc12eaff9..6c4c0b60b 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/data.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/data.rs @@ -22,7 +22,12 @@ use gateway_plugin_sdk::{Capability, Contributions, Stage}; use serde_json::{Value, json}; #[derive(Default)] -struct Facts(AtomicUsize, AtomicUsize, Option); +struct Facts { + reads: AtomicUsize, + refreshes: AtomicUsize, + email: Option, + runtime_account: Option, +} fn account() -> AccountRecord { let at = "2026-01-01T00:00:00Z".parse().unwrap(); @@ -71,7 +76,7 @@ impl PluginAccountAccess for Facts { &self, account: &ProviderAccountId, ) -> Result { - self.1.fetch_add(1, Ordering::SeqCst); + self.refreshes.fetch_add(1, Ordering::SeqCst); if account.as_str() == "acct_missing" { return Err(AdminError::not_found("missing account")); } @@ -88,19 +93,26 @@ impl PluginAccountAccess for Facts { } async fn list(&self, query: PluginAccountListQuery) -> Result { - self.0.fetch_add(1, Ordering::SeqCst); + self.reads.fetch_add(1, Ordering::SeqCst); assert_eq!(query.limit.get(), 1); assert_eq!(query.provider_kind.unwrap().as_str(), "openai"); assert_eq!(query.cursor.unwrap().as_str(), "acct_before"); let mut account = account(); - account.email.clone_from(&self.2); + account.email.clone_from(&self.email); Ok(PluginAccountPage { accounts: vec![account], next_cursor: Some(ProviderAccountId::new("acct_facts").unwrap()), }) } - async fn get_runtime(&self, _: &ProviderAccountId) -> Result { - panic!("事实回调只使用对应的窄查询") + async fn get_runtime(&self, id: &ProviderAccountId) -> Result { + let account = self + .runtime_account + .as_ref() + .expect("事实回调只使用对应的窄查询"); + if account.id != id.as_str() { + return Err(AdminError::not_found("fixture account is unavailable")); + } + Ok(account.clone()) } async fn get_credential( &self, @@ -116,10 +128,14 @@ impl PluginAccountAccess for Facts { panic!("不允许修改账号") } async fn get_quota(&self, account: &ProviderAccountId) -> Result { - self.0.fetch_add(1, Ordering::SeqCst); + self.reads.fetch_add(1, Ordering::SeqCst); assert_eq!(account.as_str(), "acct_facts"); Ok(ProviderQuota { - credits: None, + credits: Some(ProviderQuotaCredits { + has_credits: true, + unlimited: false, + balance: Some("123456789.123456789".into()), + }), plan_type: Some("private plan".into()), observed_at: None, refresh_token_expires_at: None, @@ -129,17 +145,24 @@ impl PluginAccountAccess for Facts { label: "private label".into(), limit_id: None, limit_name: None, - role: None, + role: Some(ProviderQuotaWindowRole::Primary), local_usage_attribution: QuotaLocalUsageAttribution::Unavailable, window_seconds: Some(604800), used_percent: None, reset_at: None, - limit_reached: false, + limit_reached: true, local_usage: None, provider_data: None, }], - limit_reached: false, - provider_data: None, + limit_reached: true, + provider_data: Some(ProviderDocument::new( + gateway_core::account::OpaqueProviderData::new( + json!({"private":"provider-internal-fixture"}) + .as_object() + .unwrap() + .clone(), + ), + )), }) } } @@ -157,11 +180,10 @@ async fn management_facts_validate_queries_without_permission_declarations() { Default::default(), ) .await; - let facts = Arc::new(Facts( - AtomicUsize::new(0), - AtomicUsize::new(0), - email.map(str::to_owned), - )); + let facts = Arc::new(Facts { + email: email.map(str::to_owned), + ..Default::default() + }); let port: Arc = facts.clone(); runtime.bind_account_ports(&port).unwrap(); { @@ -172,6 +194,7 @@ async fn management_facts_validate_queries_without_permission_declarations() { "data_queries":[ {"method":"host.data.accounts.list","query":{"provider_id":"openai","cursor":"acct_before","limit":1}}, {"method":"host.data.quota.get","query":{"account_id":"acct_facts"}}, + {"method":"host.data.quota.details","query":{"account_id":"acct_facts"}}, {"method":"host.data.accounts.list","query":{"limit":0}}, {"method":"host.data.accounts.list","query":{"limit":65535}}, {"method":"host.data.quota.get","query":{"account_id":"acct_facts","sql":"SELECT secret"}} @@ -209,7 +232,8 @@ async fn management_facts_validate_queries_without_permission_declarations() { .unwrap(); let result: Value = serde_json::from_slice(&response.body).unwrap(); - assert_eq!(facts.0.load(Ordering::SeqCst), 2, "{result}"); + assert_eq!(facts.reads.load(Ordering::SeqCst), 3, "{result}"); + assert_eq!(facts.refreshes.load(Ordering::SeqCst), 0); assert_eq!( result[0], json!({"schema_version":1,"accounts":[{ @@ -227,7 +251,19 @@ async fn management_facts_validate_queries_without_permission_declarations() { json!({"schema_version":1,"account_id":"acct_facts","observed_at_ms":null, "windows":[{"key":"weekly","window_seconds":604800,"used_percent":null,"reset_at_ms":null}]}) ); - for invalid in &result.as_array().unwrap()[2..] { + let details: gateway_plugin_sdk::call::data::QuotaDetails = + serde_json::from_value(result[2].clone()).unwrap(); + assert_eq!(details.observed_at_ms, None); + assert_eq!(details.windows[0].used_percent, None); + assert_eq!(details.windows[0].role.as_deref(), Some("primary")); + assert_eq!(details.windows[0].label, "private label"); + assert!(details.limit_reached && details.windows[0].limit_reached); + assert_eq!( + details.credits.unwrap().balance.as_deref(), + Some("123456789.123456789") + ); + assert!(!result[2].to_string().contains("provider-internal-fixture")); + for invalid in &result.as_array().unwrap()[3..] { assert_eq!(invalid["error"], "invalid_input"); } @@ -301,7 +337,7 @@ async fn quota_refresh_is_available_without_permission_declarations() { .unwrap(); let result: Value = serde_json::from_slice(&response.body).unwrap(); - assert_eq!(facts.1.load(Ordering::SeqCst), 4); + assert_eq!(facts.refreshes.load(Ordering::SeqCst), 4); assert_eq!(result[0]["observed_at_ms"], Value::Null); assert_eq!( result[1], @@ -323,3 +359,149 @@ async fn quota_refresh_is_available_without_permission_declarations() { super::wait_until_empty(cache.path()).await; } } + +async fn management_queries(accounts: Arc, queries: Value) -> Value { + let (cache, store, runtime) = + super::setup_with_contributions(Contributions::from([crate::support::contribution( + Capability::Management, + vec![Stage::Management], + vec![], + vec![], + )])) + .await; + runtime.bind_account_ports(&accounts).unwrap(); + store.snapshot.lock().unwrap().instances[0].configuration = json!({ + "management_registration":{"routes":[{"method":"GET","path":"facts","request_content_types":[],"response_content_types":["application/json"]}]}, + "data_queries":queries, + "sdk_session":true, + }); + let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) + .await + .unwrap(); + let view = runtime.views(&generation).await.unwrap().remove(0); + let response = runtime + .handle( + &generation, + &view.target, + PluginManagementRequest { + request_context: None, + headers: vec![], + method: "GET".into(), + path: "facts".into(), + query: String::new(), + content_type: None, + body: vec![], + request_id: "account-http-fixture".into(), + }, + ) + .await + .unwrap(); + let result = serde_json::from_slice(&response.body).unwrap(); + drop(generation); + runtime.shutdown().await; + super::wait_until_empty(cache.path()).await; + result +} + +async fn request_head(stream: &mut tokio::net::TcpStream) -> String { + use tokio::io::AsyncReadExt as _; + let mut bytes = Vec::new(); + while !bytes.ends_with(b"\r\n\r\n") { + assert!(bytes.len() < 16 * 1024); + bytes.push(stream.read_u8().await.unwrap()); + } + String::from_utf8(bytes).unwrap() +} + +#[tokio::test] +async fn account_http_uses_the_selected_proxy_and_rejects_unavailable_account_versions() { + use tokio::io::AsyncWriteExt as _; + for scenario in [ + "plain", + "explicit-auth", + "rotated", + "disabled", + "proxy-failed", + "missing", + ] { + let target = wiremock::MockServer::start().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = gateway_core::account::OutboundProxy::parse(&format!( + "http://{}", + listener.local_addr().unwrap() + )) + .unwrap(); + let successful = matches!(scenario, "plain" | "explicit-auth"); + let proxy_task = if successful { + Some(tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let request = request_head(&mut socket).await; + assert!(request.starts_with("GET http://127.0.0.1:")); + socket + .write_all( + b"HTTP/1.1 201 Created\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .unwrap(); + request + })) + } else { + drop(listener); + None + }; + let mut selected = account(); + selected.outbound_proxy = Some(proxy); + selected.enabled = scenario != "disabled"; + if scenario == "rotated" { + selected.credential_revision = Revision::new(2).unwrap(); + } + let accounts: Arc = Arc::new(Facts { + runtime_account: Some(selected), + ..Default::default() + }); + let headers = if scenario == "explicit-auth" { + vec![("authorization", b"Bearer plugin-fixture".to_vec())] + } else { + vec![] + }; + let results = management_queries(accounts, json!([{ + "method":gateway_plugin_sdk::call::host::ACCOUNT_HTTP, + "params":{ + "account_id":if scenario == "missing" {"acct_missing"} else {"acct_facts"}, + "credential_revision":1, + "request":{"method":"GET","url":format!("{}/facts", target.uri()),"headers":headers}, + }, + }])).await; + if let Some(task) = proxy_task { + assert_eq!(results[0]["status"], 201, "{scenario}: {results}"); + let request = tokio::time::timeout(std::time::Duration::from_secs(3), task) + .await + .unwrap() + .unwrap() + .to_ascii_lowercase(); + assert!( + request.starts_with("get http://127.0.0.1:") && request.contains("/facts http/1.1"), + "{request}" + ); + assert_eq!( + request.contains("authorization:"), + scenario == "explicit-auth" + ); + if scenario == "explicit-auth" { + assert!(request.contains("authorization: bearer plugin-fixture\r\n")); + } + } else { + let expected = match scenario { + "rotated" => "conflict", + "disabled" => "permission_denied", + "missing" => "rejected", + _ => "upstream", + }; + assert_eq!(results[0]["error"], expected, "{scenario}: {results}"); + } + assert!( + target.received_requests().await.unwrap().is_empty(), + "{scenario} must not connect directly" + ); + } +} diff --git a/backend/crates/gateway-plugin/runtime/tests/rpc/admission.rs b/backend/crates/gateway-plugin/runtime/tests/rpc/admission.rs index b83416e39..4b682a528 100644 --- a/backend/crates/gateway-plugin/runtime/tests/rpc/admission.rs +++ b/backend/crates/gateway-plugin/runtime/tests/rpc/admission.rs @@ -29,7 +29,7 @@ fn limits() -> RpcLimits { maximum_executing_calls: 1, maximum_queued_calls: 1, maximum_queue_timeout: Duration::from_millis(200), - maximum_callbacks: 1, + maximum_callbacks: 4, ..Default::default() } } @@ -146,7 +146,7 @@ async fn next_resumption_does_not_need_a_vacant_initial_queue_slot() { } }); started.recv().await.unwrap(); - // 等待 next 不挤占唯一的普通回调额度 + // 未消费的 next 与普通回调共享有界驻留额度 let mut occupying = Box::pin(call(&session, "callback")); assert!(futures::poll!(&mut occupying).is_pending()); started.recv().await.unwrap(); @@ -326,3 +326,246 @@ async fn resumption_queue_wait_preserves_the_remaining_plugin_processing_budget( assert!(session.is_ready()); session.shutdown(Duration::from_secs(1)).await; } + +async fn scheduled( + callbacks: Arc, + mut limits: RpcLimits, +) -> (tempfile::TempDir, Arc) { + let (cache, prepared, mut handshake, processes) = prepared_session(); + handshake.configuration = json!({"sdk_session": true}); + limits.maximum_queue_timeout = Duration::from_secs(2); + let session = RpcSession::start(prepared, handshake, limits, &processes, callbacks) + .await + .unwrap(); + (cache, Arc::new(session)) +} + +struct CrossCallbacks { + other: std::sync::Mutex>, + ready: Arc, +} + +impl CallbackHandler for CrossCallbacks { + fn call( + &self, + _context: CallContext, + _method: String, + _params: serde_json::Value, + _payload: Vec, + ) -> BoxFuture<'static, Result> { + let other = self.other.lock().unwrap().upgrade().unwrap(); + let ready = self.ready.clone(); + Box::pin(async move { + ready.wait().await; + other + .call( + "echo", + other.context(Stage::Request, Duration::from_secs(2)), + json!("child"), + Vec::new(), + ) + .await + .map_err(|_| PluginFault::new(gateway_plugin_sdk::ErrorCode::Fault, "child failed")) + }) + } +} + +#[tokio::test] +async fn independent_sessions_release_execution_for_every_callback_method() { + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let left = Arc::new(CrossCallbacks { + other: Default::default(), + ready: barrier.clone(), + }); + let right = Arc::new(CrossCallbacks { + other: Default::default(), + ready: barrier, + }); + let (_left_cache, a) = scheduled(left.clone(), limits()).await; + let (_right_cache, b) = scheduled(right.clone(), limits()).await; + *left.other.lock().unwrap() = Arc::downgrade(&b); + *right.other.lock().unwrap() = Arc::downgrade(&a); + for method in [ + "host.service.call", + "host.http.dispatch", + "host.model.execute", + "host.client.verify", + NEXT_METHOD, + "host.unclassified", + ] { + let invoke = |session: Arc| async move { + session + .call( + "callback", + session.context(Stage::Request, Duration::from_secs(2)), + json!({"method":method}), + Vec::new(), + ) + .await + }; + let (a_result, b_result) = tokio::time::timeout(Duration::from_secs(3), async { + tokio::join!(invoke(a.clone()), invoke(b.clone())) + }) + .await + .unwrap(); + assert_eq!(a_result.unwrap().result, json!("child"), "{method}"); + assert_eq!(b_result.unwrap().result, json!("child"), "{method}"); + } + a.shutdown(Duration::from_secs(1)).await; + b.shutdown(Duration::from_secs(1)).await; +} + +struct PartialCallbacks { + entered: mpsc::UnboundedSender, + fast: Arc, + slow: Arc, +} + +impl CallbackHandler for PartialCallbacks { + fn call( + &self, + _context: CallContext, + method: String, + _params: serde_json::Value, + _payload: Vec, + ) -> BoxFuture<'static, Result> { + self.entered.send(method.clone()).unwrap(); + let gate = match method.as_str() { + "fast" => Some(self.fast.clone()), + "slow" => Some(self.slow.clone()), + _ => None, + }; + Box::pin(async move { + if let Some(gate) = gate { + gate.acquire().await.unwrap().forget(); + } + Ok(RpcReply { + result: json!(method), + payload: Vec::new(), + }) + }) + } +} + +#[tokio::test] +async fn a_partial_callback_result_requires_a_permit_and_can_wait_again() { + let (entered, mut events) = mpsc::unbounded_channel(); + let fast = Arc::new(Semaphore::new(0)); + let slow = Arc::new(Semaphore::new(0)); + let (cache, session) = scheduled( + Arc::new(PartialCallbacks { + entered, + fast: fast.clone(), + slow: slow.clone(), + }), + limits(), + ) + .await; + let parent = tokio::spawn({ + let session = session.clone(); + async move { call(&session, "partial").await } + }); + for _ in 0..2 { + tokio::time::timeout(Duration::from_secs(2), events.recv()) + .await + .unwrap() + .unwrap(); + } + let marker = cache.path().join("processing"); + let busy = tokio::spawn({ + let session = session.clone(); + let marker = marker.clone(); + async move { + session + .call( + "busy", + session.context(Stage::Request, Duration::from_secs(2)), + json!({"marker":marker, "delay_ms":300}), + Vec::new(), + ) + .await + } + }); + tokio::time::timeout(Duration::from_secs(2), async { + while !marker.exists() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + fast.add_permits(1); + assert!( + tokio::time::timeout(Duration::from_millis(50), events.recv()) + .await + .is_err() + ); + busy.await.unwrap().unwrap(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(2), events.recv()) + .await + .unwrap() + .as_deref(), + Some("processed") + ); + assert!(!parent.is_finished()); + // 再次等待慢回调后，独立调用可取得同一个执行名额 + assert!(call(&session, "echo").await.is_ok()); + slow.add_permits(1); + assert_eq!(parent.await.unwrap().unwrap().result, json!("complete")); + session.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test] +async fn polling_wakeups_and_cancelled_suspensions_leave_the_session_usable() { + let (entered, mut events) = mpsc::unbounded_channel(); + let slow = Arc::new(Semaphore::new(0)); + let (_cache, session) = scheduled( + Arc::new(PartialCallbacks { + entered, + fast: Arc::new(Semaphore::new(0)), + slow, + }), + RpcLimits { + maximum_calls: 1, + ..limits() + }, + ) + .await; + assert_eq!( + call(&session, "yield").await.unwrap().result, + json!("yielded") + ); + let parent = tokio::spawn({ + let session = session.clone(); + async move { + session + .call( + "callback", + session.context(Stage::Request, Duration::from_secs(2)), + json!({"method":"slow"}), + Vec::new(), + ) + .await + } + }); + assert_eq!(events.recv().await.as_deref(), Some("slow")); + assert!(matches!( + call(&session, "echo").await, + Err(RpcError::Overloaded) + )); + parent.abort(); + assert!(matches!(parent.await, Err(error) if error.is_cancelled())); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + match call(&session, "echo").await { + Ok(_) => break, + Err(RpcError::Overloaded) => tokio::task::yield_now().await, + Err(error) => panic!("{error}"), + } + } + }) + .await + .unwrap(); + assert!(session.is_ready()); + session.shutdown(Duration::from_secs(1)).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs b/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs index 0d2767dd8..140a33007 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs @@ -23,6 +23,8 @@ use gateway_plugin_sdk::call::{ upstream_adapter::{UpstreamAdapterEvent, UpstreamAdapterRequest}, }; +mod scheduled; + type CallbackResult = Result<(Value, Vec), PluginFault>; struct PendingCallback { parent: u64, @@ -82,6 +84,7 @@ struct Peer { output: mpsc::Sender, callbacks: Mutex>, next_callback: AtomicU64, + resumptions: Mutex>>, streams: Mutex>>, } @@ -90,18 +93,21 @@ impl Peer { let queries = self.configuration["data_queries"].as_array()?; let mut results = Vec::new(); for query in queries { - let metadata_query = query["method"] == "host.keys.list"; + let metadata_query = + query["method"] == "host.keys.list" || query.get("params").is_some(); let result = self .callback_payload( id, query["method"].as_str().unwrap(), if metadata_query { - query["query"].clone() + query.get("params").unwrap_or(&query["query"]).clone() } else { json!({}) }, if metadata_query { - vec![] + query["body"] + .as_str() + .map_or_else(Vec::new, |body| body.as_bytes().to_vec()) } else { serde_json::to_vec(&query["query"]).unwrap() }, @@ -243,6 +249,36 @@ impl Peer { self.callback_payload(parent, method, params, vec![]).await } + async fn send_callback( + &self, + parent: u64, + method: &str, + params: Value, + payload: Vec, + ) -> (u64, oneshot::Receiver) { + let (response, received) = oneshot::channel(); + let mut callbacks = self.callbacks.lock().await; + let id = self.next_callback.fetch_add(2, Ordering::Relaxed); + callbacks.insert( + id, + PendingCallback { + parent, + response: Some(response), + }, + ); + self.send( + Message::Callback { + id, + parent_id: parent, + method: method.into(), + params, + }, + payload, + ) + .await; + (id, received) + } + async fn callback_payload( &self, parent: u64, @@ -250,29 +286,17 @@ impl Peer { params: Value, payload: Vec, ) -> CallbackResult { - let (response, received) = oneshot::channel(); - { - let mut callbacks = self.callbacks.lock().await; - let id = self.next_callback.fetch_add(2, Ordering::Relaxed); - callbacks.insert( + let (id, received) = self.send_callback(parent, method, params, payload).await; + let result = received.await.unwrap(); + self.send( + Message::CallbackConsumed { id, - PendingCallback { - parent, - response: Some(response), - }, - ); - self.send( - Message::Callback { - id, - parent_id: parent, - method: method.into(), - params, - }, - payload, - ) - .await; - } - received.await.unwrap() + parent_id: parent, + }, + vec![], + ) + .await; + result } async fn run_nested_model_fixture( @@ -1349,31 +1373,43 @@ impl Peer { "slow" => tokio::time::sleep(Duration::from_millis(400)).await, "hang" | "hang_uncancellable" => return, "crash" => std::process::exit(7), - "callback" | "callback_method" | "forged_callback" => { + "callback" | "callback_method" => { let callback_method = if method == "callback_method" { params["method"].as_str().unwrap().to_owned() } else { "host.log".into() }; - // 分配与入队保持同一顺序，避免并发回调制造不合法 ID 序列 - let mut callbacks = self.callbacks.lock().await; - let callback = self.next_callback.fetch_add(2, Ordering::Relaxed); - callbacks.insert( - callback, - PendingCallback { - parent: id, - response: None, + let (callback, received) = self + .send_callback(id, &callback_method, params, payload) + .await; + self.send(Message::Suspend { id, epoch: 1 }, vec![]).await; + let result = received.await.unwrap(); + let (resumed, granted) = oneshot::channel(); + self.resumptions.lock().await.insert(id, resumed); + self.send(Message::Resume { id, epoch: 1 }, vec![]).await; + granted.await.unwrap(); + self.send( + Message::CallbackConsumed { + id: callback, + parent_id: id, }, - ); + vec![], + ) + .await; + match result { + Ok((result, payload)) => { + self.send(Message::Result { id, result }, payload).await + } + Err(error) => self.send(Message::Error { id, error }, vec![]).await, + } + return; + } + "forged_callback" => { self.send( Message::Callback { - id: callback, - parent_id: if method == "forged_callback" { - 999999 - } else { - id - }, - method: callback_method, + id: self.next_callback.fetch_add(2, Ordering::Relaxed), + parent_id: 999999, + method: "host.log".into(), params, }, payload, @@ -1426,7 +1462,7 @@ impl Peer { } } -#[tokio::main(flavor = "current_thread")] +#[tokio::main(flavor = "multi_thread", worker_threads = 2)] async fn main() { let mut input = tokio::io::stdin(); let (output, mut frames) = mpsc::channel::(64); @@ -1437,6 +1473,11 @@ async fn main() { } }); let hello = read_frame(&mut input).await.unwrap(); + if matches!(&hello.message, Message::Hello { handshake } if handshake.configuration["sdk_session"] == true) + { + scheduled::run(hello, input).await; + return; + } let Message::Hello { mut handshake } = hello.message else { panic!("expected handshake") }; @@ -1491,6 +1532,7 @@ async fn main() { output, callbacks: Mutex::new(BTreeMap::new()), next_callback: AtomicU64::new(2), + resumptions: Mutex::new(BTreeMap::new()), streams: Mutex::new(BTreeMap::new()), }); let exit_after_ready = peer.configuration["exit_after_ready_delays_ms"] @@ -1603,6 +1645,11 @@ async fn main() { ), ); } + Message::Resumed { id, epoch: 1 } => { + if let Some(resumed) = peer.resumptions.lock().await.remove(&id) { + let _ = resumed.send(()); + } + } Message::Cancel { id } => { if uncancellable.contains(&id) { continue; diff --git a/backend/crates/gateway-plugin/runtime/tests/support/worker/scheduled.rs b/backend/crates/gateway-plugin/runtime/tests/support/worker/scheduled.rs new file mode 100644 index 000000000..6059be3e0 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/support/worker/scheduled.rs @@ -0,0 +1,173 @@ +//! 使用公开 SDK 驱动真实子进程，覆盖受管任务的并发与恢复 + +use gateway_plugin_sdk::{ + ErrorCode, Frame, PluginFault, + client::{ + CallFuture, CallReply, PluginCall, PluginHandler, PluginSession, SessionConfig, write_frame, + }, +}; +use serde_json::json; +use tokio::io::AsyncReadExt as _; + +pub(super) async fn run(hello: Frame, input: tokio::io::Stdin) { + let mut prefix = Vec::new(); + write_frame(&mut prefix, &hello).await.unwrap(); + let input = std::io::Cursor::new(prefix).chain(input); + let session = PluginSession::accept(input, tokio::io::stdout(), SessionConfig::default()) + .await + .unwrap(); + let handler = Handler { + configuration: session.handshake().configuration.clone(), + contributes: session.handshake().contributes.clone(), + }; + session.run(handler).await.unwrap(); +} + +struct Handler { + configuration: serde_json::Value, + contributes: gateway_plugin_sdk::Contributions, +} + +impl Handler { + async fn queries(&self, call: &PluginCall, key: &str) -> Vec { + let mut results = Vec::new(); + for query in self.configuration[key].as_array().unwrap() { + let metadata = query.get("params").is_some() || query["method"] == "host.keys.list"; + let params = if metadata { + query.get("params").unwrap_or(&query["query"]).clone() + } else { + json!({}) + }; + let payload = if metadata { + query["body"] + .as_str() + .map_or_else(Vec::new, |body| body.as_bytes().to_vec()) + } else { + serde_json::to_vec(&query["query"]).unwrap() + }; + let result = call + .host + .call(query["method"].as_str().unwrap(), params, payload) + .await; + results.push(match result { + Ok(reply) if metadata => { + assert!(reply.payload.is_empty()); + reply.result + } + Ok(reply) => { + assert_eq!(reply.result, json!({})); + serde_json::from_slice(&reply.payload).unwrap() + } + Err(error) => json!({"error":error.into_plugin_fault().code}), + }); + } + results + } +} + +impl PluginHandler for Handler { + fn call(&self, call: PluginCall) -> CallFuture<'_> { + Box::pin(async move { + match call.method.as_str() { + "plugin.register" => Ok(CallReply::unary( + serde_json::to_value(gateway_plugin_sdk::call::registration::Registration { + contributes: self.contributes.clone(), + }) + .unwrap(), + Vec::new(), + )), + "management.register" => Ok(CallReply::unary( + json!({}), + serde_json::to_vec(&self.configuration["management_registration"]).unwrap(), + )), + "frontend_auth.identifier" => Ok(CallReply::unary( + json!({"identifier":"fixture-auth"}), + Vec::new(), + )), + "frontend_auth.authenticate" => { + let request: gateway_plugin_sdk::call::frontend_authentication::FrontendAuthenticationRequest = serde_json::from_slice(&call.payload).unwrap(); + assert_eq!( + request.authorization, + self.configuration["expected_frontend_authorization"] + .as_str() + .unwrap() + ); + let facts = if self.configuration.get("auth_queries").is_some() { + self.queries(&call, "auth_queries").await + } else { + Vec::new() + }; + if let Some(path) = + self.configuration["frontend_authentication_marker"].as_str() + { + std::fs::write(path, serde_json::to_vec(&facts).unwrap()).unwrap(); + } + Ok(CallReply::unary( + json!({}), + serde_json::to_vec(&self.configuration["frontend_authentication_result"]) + .unwrap(), + )) + } + "management.handle" => { + let results = self.queries(&call, "data_queries").await; + Ok(CallReply::unary( + json!({"status":200,"content_type":"application/json"}), + serde_json::to_vec(&results).unwrap(), + )) + } + "callback" => { + let reply = call + .host + .call( + call.params["method"].as_str().unwrap(), + json!({}), + Vec::new(), + ) + .await + .map_err(|error| error.into_plugin_fault())?; + Ok(CallReply::unary(reply.result, reply.payload)) + } + "partial" => { + let fast = Box::pin(call.host.call("fast", json!({}), Vec::new())); + let slow = Box::pin(call.host.call("slow", json!({}), Vec::new())); + let (first, remaining) = match futures::future::select(fast, slow).await { + futures::future::Either::Left(pair) + | futures::future::Either::Right(pair) => pair, + }; + first.map_err(|error| error.into_plugin_fault())?; + call.host + .call("processed", json!({}), Vec::new()) + .await + .map_err(|error| error.into_plugin_fault())?; + remaining.await.map_err(|error| error.into_plugin_fault())?; + Ok(CallReply::unary(json!("complete"), Vec::new())) + } + "busy" => { + std::fs::write(call.params["marker"].as_str().unwrap(), "entered").unwrap(); + // 保持一次 poll 未返回，代表确实占用执行名额的工作 + std::thread::sleep(std::time::Duration::from_millis( + call.params["delay_ms"].as_u64().unwrap(), + )); + Ok(CallReply::unary(json!("busy complete"), Vec::new())) + } + "yield" => { + for _ in 0..16 { + tokio::task::yield_now().await; + } + Ok(CallReply::unary(json!("yielded"), Vec::new())) + } + "drop_callback" => { + let mut pending = Box::pin(call.host.call("slow", json!({}), Vec::new())); + assert!(futures::poll!(&mut pending).is_pending()); + drop(pending); + Ok(CallReply::unary(json!("dropped"), Vec::new())) + } + "echo" => Ok(CallReply::unary(call.params, call.payload)), + _ => Err(PluginFault::new( + ErrorCode::InvalidInput, + "unknown fixture method", + )), + } + }) + } +} diff --git a/backend/crates/gateway-plugin/sdk/README.md b/backend/crates/gateway-plugin/sdk/README.md index a81a3682c..e12cfe333 100644 --- a/backend/crates/gateway-plugin/sdk/README.md +++ b/backend/crates/gateway-plugin/sdk/README.md @@ -68,9 +68,13 @@ sequenceDiagram H->>P: 注册能力，核对清单 P-->>H: 注册结果 H->>P: Call（ID、上下文、载荷） - opt 需要宿主资源 + opt 父处理等待宿主回调 P->>H: 回调（关联父调用 ID） + P->>H: Suspend（父处理 Pending） H-->>P: 回调结果 + P->>H: Resume + H-->>P: Resumed（取得执行名额） + P->>H: CallbackConsumed end P-->>H: 结果或错误；流以 End 结束 ``` @@ -78,14 +82,17 @@ sequenceDiagram ### 并发与期限 每个宿主插件会话默认允许 64 个执行中的 RPC 调用、64 个初次排队调用和 256 个总在途上下文。 -初次排队和等待 `next` 后的恢复排队各限 30 秒，不扣插件自身的处理预算。 -恢复等待占总在途名额，不占初次排队名额。 -等待 `next` 时释放执行名额并暂停处理预算，返回后重新取得名额，继续同一次调用； -已返回响应头的流继续占用在途预算，不占首结果执行名额。具体组合行为见[洋葱中间件](docs/capabilities.md#洋葱中间件) +初次排队和恢复排队各限 30 秒；恢复等待占总在途名额，不占初次排队名额 -插件侧 `SessionConfig.maximum_calls` 默认 256，限制包含下游等待和响应流在内的在途调用； -`maximum_callbacks` 默认 64，限制普通宿主回调。宿主普通回调的默认上限同样为 64， -`next` 使用总在途上限对应的独立回调预算，不挤占普通回调容量 +SDK 只在父处理 future 返回 `Pending` 后报告挂起，Runtime 随即归还执行名额。 +收到唤醒后，SDK 请求恢复，取得授权才继续驱动父任务；并行回调只完成一个也遵守这一边界。 +等待普通回调及其恢复排队继续计入处理预算；只有父任务挂起且驻留回调全部为 `next` 时暂停预算。 +外部取消、请求期限和排队上限独立生效，调度不保证插件自行创建的线程、任务或锁能够推进 + +插件侧 `SessionConfig.maximum_calls` 默认 256，限制包含下游等待和响应流在内的在途调用。 +SDK 和宿主的 `maximum_callbacks` 默认 320、最大 512，统一限制全部宿主回调，包含已返回但尚未消费的结果。 +取消或丢弃回调 future 会释放该回调；宿主不保证已发出的业务操作可以撤销，不自动重试或补偿。 +已返回响应头的流继续占用总在途预算，不占首结果执行名额，由正文或连接的 owner 管理真实寿命 实际容量由宿主预算与 SDK 会话配置共同限制，较小的插件侧上限会提前拒绝新调用。 达到容量或等待上限时，中间件返回过载错误；HTTP 响应为 503（`middleware_capacity_unavailable`）， @@ -144,7 +151,7 @@ session.run(handler).await?; 首个流式结果成功返回后，正文、会话及关联回调由资源取消与终态回收，不再沿用初始期限。 普通调用继续使用原期限；资源模式仍保留并发容量、Credit、父调用关联及取消约束 -协议 v6 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成，随后依次传输 JSON 元数据和二进制载荷。 +协议 v7 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成，随后依次传输 JSON 元数据和二进制载荷。 元数据保留 64 KiB 上限；正文总量不由插件传输层额外限制。传输内部以 256 KiB 为单次 I/O 块， 不逐块发起 RPC，也不要求插件配置与宿主同步。读取按进度扩展缓冲区；截断正文必须报错，不能作为完整消息交付 @@ -165,6 +172,17 @@ SDK 帧读写器自动拆装这些字段，不拼接复制二进制正文。`Mes 底层帧读写方法不是取消安全的；自行编排会话时，不能在 `select!` 的其他分支完成后 重建半途中的帧读写。`PluginSession` 会在处理调用完成通知时保留同一个读取 future，直到完整帧到达或会话关闭 +RPC v7 的调度控制消息没有载荷：`Suspend { id, epoch }` 归还执行名额，`Resume { id, epoch }` +请求同一轮恢复，宿主 `Resumed { id, epoch }` 授权后才能继续首结果前的处理。`epoch` 从 1 递增， +每次挂起与恢复严格配对。插件发出的控制消息和业务帧使用同一有界 FIFO，保证回调登记先于挂起 + +`CallbackConsumed { id, parent_id }` 表示回调结果已消费或 future 已丢弃，宿主据此释放驻留额度并取消尚未完成的回调。 +结果接收本身不恢复父任务、不释放驻留额度；父调用终结统一回收剩余回调。 +取消、恢复授权和结果可能交错，终结后的消息不得重新打开调用 + +宿主只接受 RPC v7，不提供旧协议兼容或降级。使用旧 SDK 的插件须更新 SDK、重新构建并由 CLI 生成安装包后安装， +只修改清单中的版本号不能满足任务挂起、恢复和回调回收合同 + 插件须校验 `Hello` 的协议版本。一次 `Call` 的结果、错误和流帧按 ID 关联；实例、代次、阶段、账号或资源 ID 用于关联调用和管理资源生命周期 ## 本地验证 diff --git a/backend/crates/gateway-plugin/sdk/docs/capabilities.md b/backend/crates/gateway-plugin/sdk/docs/capabilities.md index 04b8d762a..e2af86b04 100644 --- a/backend/crates/gateway-plugin/sdk/docs/capabilities.md +++ b/backend/crates/gateway-plugin/sdk/docs/capabilities.md @@ -81,11 +81,14 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。宿主开放以下 11 类扩 | SDK 方法 | 回调 | 查询与结果 | | --- | --- | --- | | `call.host.account_facts(query)` | `host.data.accounts.list` | `AccountFactsQuery`：可选 `provider_id`、`cursor`，必填 `limit`（1～200）；按账号 ID 升序，`next_cursor=null` 表示本页已结束 | +| `call.host.quota_details(query)` | `host.data.quota.details` | `QuotaFactsQuery { account_id }`：读取完整额度事实，包含窗口归属、触顶和点数 | +| `call.host.check_access(query)` | `host.data.access.check` | `AccessQuery { client_key_id, account_ids, upstream_model }`：读取 Key 对指定账号及可选上游模型的冻结授权 | +| `call.host.verify_client(request)` | `host.client.verify` | `VerifyClientRequest { authorization }`：复用入口认证，返回 `VerifiedClient { client_key_id }` | | `call.host.quota_facts(query)` | `host.data.quota.get` | `QuotaFactsQuery { account_id }`：读取 Provider 现有观测，不访问上游刷新 | | `call.host.key_facts(query)` | `host.data.keys.get` | `ClientKeyFactsQuery { client_key_id }`：读取当前 Key 的启用状态及显式分组 ID，不返回密钥 | -类型在 `call::data`。控制参数为 `{}`，查询和结果使用二进制 JSON；结果固定 `schema_version=1`。 -SDK 解析响应时忽略未知字段，包含账号分页、账号及额度窗口；已知字段按声明校验类型与必填性，查询拒绝未知字段。 +类型在 `call::data`。控制参数为 `{}`，查询和结果使用二进制 JSON；事实结果固定 `schema_version=1`，认证结果只包含 `client_key_id`。 +SDK 解析响应时忽略未知字段，包含账号分页、账号及额度窗口；已知字段按声明校验类型与必填性，账号、Key 与额度查询拒绝未知字段，认证与权限查询忽略未知字段。 插件须通过 `engines.codex-proxy-rs` 限定支持所用接口和字段的宿主版本 `key_facts` 每次读取当前管理数据，返回 `client_key_id`、`enabled`、`group_ids`；不存在的 Key 沿用事实接口的 `rejected` 错误 @@ -93,10 +96,22 @@ SDK 解析响应时忽略未知字段，包含账号分页、账号及额度窗 `group_ids` 表示显式绑定，包含停用分组，不是最终可路由账号集合；空绑定也不代表单账号范围。通过 `account_facts` 关联账号时需完整遍历分页，结果包含停用账号。跨查询关联及同步策略由插件负责，这些调用不构成跨查询事务，管理员修改绑定后应重新检查。 账号仅返回 `account_id`、`provider_id`、`name`、`email`、`group_ids`、`enabled` 和 `updated_at_ms`，不附带令牌或代理信息。 `name` 和 `email` 来自宿主已保存的账号资料；没有邮箱时 `email=null`，读取不会请求上游个人信息。 -额度仅返回观测时间与窗口的 `key`、`window_seconds`、`used_percent`、`reset_at_ms`。 +`quota_facts` 仅返回观测时间与窗口的 `key`、`window_seconds`、`used_percent`、`reset_at_ms`。 时间均为 UTC Unix 毫秒，比例为百分数；未知值保留 `null`，不能解释为 0。`observed_at_ms=null` 表示没有可用观测时间， 不保证当前缓存新鲜，不提供历史快照或多个查询之间的原子一致性 +`quota_details` 保留 Provider 已解释的 `limit_reached`、窗口 `group`、`label`、`limit_id`、`limit_name`、`role` +及点数 `has_credits`、`unlimited`、`balance`；缺失值仍为 `null`，不返回上游原始文档，不发起刷新 + +`verify_client` 使用当前请求设置和原认证计划，包含外部 principal 到 Key 的映射；没有 HTTP 入口时使用当前发布配置。 +认证可以调用认证插件，继续扣减父预算并响应取消；验证成功不做推理准入、账号租约、使用记录或计费。 +授权原文只作本次认证输入，不进入返回值 + +`check_access` 与 `host.models.list`、显式 Key 模型执行共用父调用内的身份和冻结权限快照。 +账号或模型不在冻结范围内时返回 `allowed=false`，不存在或不可用的 Key 沿用 Core 的身份拒绝。 +它们不建立模型执行图，不受模型递归深度或当前实例已入链影响；健康、冷却、额度、租约与实时可用性由实际执行检查。 +`upstream_model` 使用上游模型名，省略时只判断账号范围；返回结果不承诺后续请求一定可以执行 + ```rust,ignore use gateway_plugin_sdk::call::data::{AccountFactsQuery, QuotaFactsQuery}; let page = call.host.account_facts(AccountFactsQuery { @@ -217,8 +232,8 @@ SDK 不自动重试。超时或断连不能证明写入未提交；重试上限 原有 8 MiB 编码上限、Runtime 缓冲预算及流控窗口限制 HTTP 和 request 中间件中的模型回调沿用有效请求设置；attempt、outbound、路由、调度、重试和上游适配器 -沿用该次模型请求冻结的正文预算。显式 Key 模型调用的预检与执行使用同一绑定快照， -同一父调用内复用 Key 时不受后续配置发布影响；父调用结束后重新绑定才读取新配置，无需重新启用插件 +沿用该次模型请求冻结的正文预算。显式 Key 模型调用的预检与执行使用同一身份快照， +同一父调用内复用 Key 时不受后续配置发布影响；父调用结束后重新解析身份才读取新配置，无需重新启用插件 `ModelExecuteRequest.client_key_id` 可显式选择执行 Key，不与触发插件的父请求 Key 求交。省略时继承现有模型父请求身份；没有模型父请求的调用必须提供 Key。所选 Key 的启用状态、模型规则、账号组、并发和预算继续由 Core 执行 嵌套调用与显式 Key 调用继承入口快照中的客户端 IP 和 User-Agent；没有 HTTP 入口的后台调用保持来源未知 @@ -238,12 +253,17 @@ HTTP 和 request 中间件中的模型回调沿用有效请求设置；attempt 父调用取消或结束时宿主回收未关闭的流，SDK 不后台预读。取消一次 `read()` 等待后，再次读取会继续同一次操作，不丢弃已经返回的分块。 账号上游使用同一正文对象，使用方式见[受管上游](upstream-adapters.md) +`HostClient::account_http(AccountHttpRequest { account_id, credential_revision, request }, body)` 使用指定账号的当前代理配置， +经 `host.http.account` 返回同一 `HostHttpBody`。每次发送核对账号存在、启用和凭据版本，版本变化须由插件重新读取后决定下一步。 +请求 headers 与正文完全由插件提供；宿主不注入账号认证、不申请推理租约，代理不可用时直接失败，不回退直连。 +这项能力与上游适配器的账号推理 HTTP 分开，适用于插件自行构造的账号网络请求 + 请求和响应 `headers` 使用 `(String, Vec)`，保留重复头与非 UTF-8 值；响应另有 `http_version`。 读取至 EOF 后通过 `body.trailers()` 取得完整尾部头，读取中只包含已收到的部分。 `collect` 只返回正文，需要尾部头时使用 `read`；底层非流式回复携带 `trailers`，流式回复在 `stream_read` 中交付尾部头 `host.http.do/do_stream` 和流读取／关闭方法提供受管 HTTP。正文使用独立二进制载荷，Host 统一执行代理、超时和背压，不按插件身份限制目标地址段。DNS 由宿主解析，连接保留目标 Host 与 TLS 域名校验；重定向不会自动跟随。 -受管 HTTP（包括 `host.upstream.http.do[_stream]`）的请求正文按有效设置 +受管 HTTP（包括 `host.http.account` 和 `host.upstream.http.do[_stream]`）的请求正文按有效设置 `responses_max_decompressed_body_bytes` 检查。HTTP/request 中间件使用有效请求设置； attempt、outbound、路由、调度、重试和上游适配器使用执行冻结预算，其余调用读取当前发布配置。 超限返回带 HTTP 413 的 `invalid_input` 和 `not_sent`，不向目标发送请求。 @@ -275,7 +295,7 @@ attempt、outbound、路由、调度、重试和上游适配器使用执行冻 `next.run(input)` 消费自身，最多执行一次；结果按相反顺序经过外层。重试或独立子请求通过宿主执行服务发起 调用上下文的 `timeout_ms` 限制插件自身处理时间，入站与返回处理合计使用同一预算。 -等待 `next` 时暂停该预算并释放执行名额；恢复后继续同一次调用，普通宿主回调不会暂停计时。 +父处理任务真正挂起时释放执行名额，恢复取得授权后继续同一次调用。仅等待 `next` 时暂停处理预算；并行包含普通回调时继续计时，具体规则见[并发与期限](../README.md#并发与期限)。 排队、在途容量与过载行为见 [SDK 并发与期限](../README.md#并发与期限)。 请求总期限、下游操作超时和取消信号独立生效，客户端取消或插件退出会结束等待并回收资源 diff --git a/backend/crates/gateway-plugin/sdk/docs/manifest.md b/backend/crates/gateway-plugin/sdk/docs/manifest.md index aa8903c0d..3c8f95b1a 100644 --- a/backend/crates/gateway-plugin/sdk/docs/manifest.md +++ b/backend/crates/gateway-plugin/sdk/docs/manifest.md @@ -26,7 +26,7 @@ | 版本字段 | 当前值 | 对应合同 | | --- | --- | --- | | `manifestVersion` / `manifest_schema_versions` | `2` | 插件清单格式 | -| `package.protocolVersion` / `protocol_versions` | `6` | 插件进程 RPC 封装，不是 OpenAI 或 xAI 的业务协议 | +| `package.protocolVersion` / `protocol_versions` | `7` | 插件进程 RPC 封装，不是 OpenAI 或 xAI 的业务协议 | | `contributes..version` / `capabilities[].versions` | 见[扩展项简写](#扩展项简写) | 指定扩展能力的行为合同 | | 宿主声明的 `schema_version` | `2` | 宿主兼容声明自身的格式 | @@ -34,7 +34,7 @@ 清单不接受 `permissions` 字段；插件按[完整信任](#完整信任)运行 -RPC v6 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建，并由 CLI 生成包含协议版本的安装包 +RPC v7 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建，并由 CLI 生成包含协议版本的安装包 安装包清单校验与已安装元数据读取是两个边界：清单拒绝未知字段，并检查必需字段、类型和版本； 数据库元数据的读取规则见[制品与发布](../../../../../docs/architecture.md#制品配置与发布)，不改变清单或 RPC 合同。 diff --git a/backend/crates/gateway-plugin/sdk/src/call/data.rs b/backend/crates/gateway-plugin/sdk/src/call/data.rs index 48927ec1b..df2cbc696 100644 --- a/backend/crates/gateway-plugin/sdk/src/call/data.rs +++ b/backend/crates/gateway-plugin/sdk/src/call/data.rs @@ -1,5 +1,5 @@ -//! 基础事实查询与主动额度刷新 -//! 响应忽略未知字段，以兼容宿主新增事实；查询仍严格校验字段 +//! 入口认证、权限与账号额度事实，以及主动额度刷新合同 +//! 响应忽略未知字段，以兼容宿主新增事实 use serde::{Deserialize, Serialize}; @@ -7,6 +7,44 @@ pub const ACCOUNTS_LIST: &str = "host.data.accounts.list"; pub const QUOTA_REFRESH: &str = "host.quota_observations.refresh"; pub const QUOTA_GET: &str = "host.data.quota.get"; pub const KEYS_GET: &str = "host.data.keys.get"; +pub const CLIENT_VERIFY: &str = "host.client.verify"; +pub const ACCESS_CHECK: &str = "host.data.access.check"; +pub const QUOTA_DETAILS: &str = "host.data.quota.details"; + +/// 交给宿主入口认证计划的凭据；只通过二进制载荷传输，不提供内容型 Debug +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct VerifyClientRequest { + pub authorization: String, +} + +/// 宿主认证与 principal 映射后的执行身份，不包含 Key 明文 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct VerifiedClient { + pub client_key_id: String, +} + +/// 候选账号的冻结权限查询；上游模型名不解析为公开模型别名 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AccessQuery { + pub client_key_id: String, + pub account_ids: Vec, + pub upstream_model: Option, +} + +/// 授权结果不代表账号健康、额度或执行准入结果 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AccessFacts { + pub schema_version: u32, + pub client_key_id: String, + /// 与查询候选保持相同顺序 + pub accounts: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AccountAccessFacts { + pub account_id: String, + pub allowed: bool, +} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -72,3 +110,37 @@ pub struct QuotaWindowFacts { pub used_percent: Option, pub reset_at_ms: Option, } + +/// Provider 已解析的窗口与点数事实，不包含原始 Provider 文档或宿主预测 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct QuotaDetails { + pub schema_version: u32, + pub account_id: String, + pub observed_at_ms: Option, + pub limit_reached: bool, + pub windows: Vec, + pub credits: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct QuotaWindowDetails { + pub key: String, + pub group: String, + pub label: String, + pub limit_id: Option, + pub limit_name: Option, + pub role: Option, + pub window_seconds: Option, + /// 百分比而非 0～1 比率；未知值为 null + pub used_percent: Option, + pub reset_at_ms: Option, + pub limit_reached: bool, +} + +/// 余额沿用上游的十进制文本，不转为浮点数或把未知值解释为零 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct QuotaCredits { + pub has_credits: bool, + pub unlimited: bool, + pub balance: Option, +} diff --git a/backend/crates/gateway-plugin/sdk/src/call/host.rs b/backend/crates/gateway-plugin/sdk/src/call/host.rs index 8b9316035..008910b47 100644 --- a/backend/crates/gateway-plugin/sdk/src/call/host.rs +++ b/backend/crates/gateway-plugin/sdk/src/call/host.rs @@ -66,6 +66,16 @@ pub struct HttpRequest { pub headers: Vec<(String, Vec)>, } +pub const ACCOUNT_HTTP: &str = "host.http.account"; + +/// 显式选择账号代理和凭据版本；认证头仍由插件提供，不取得推理租约 +#[derive(Clone, Serialize, Deserialize)] +pub struct AccountHttpRequest { + pub account_id: String, + pub credential_revision: u64, + pub request: HttpRequest, +} + #[derive(Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct HttpResponse { diff --git a/backend/crates/gateway-plugin/sdk/src/client/data.rs b/backend/crates/gateway-plugin/sdk/src/client/data.rs index fea28d541..b415352c3 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/data.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/data.rs @@ -5,6 +5,28 @@ use crate::{PluginFault, call::data}; use super::{HostClient, payload_call}; impl HostClient { + /// 通过宿主入口认证计划校验凭据并取得执行身份，不执行模型或记录使用事实 + /// + /// # Errors + /// 凭据无效、认证失败或父调用取消时返回错误 + pub async fn verify_client( + &self, + request: data::VerifyClientRequest, + ) -> Result { + payload_call(self, data::CLIENT_VERIFY, request).await + } + + /// 读取所选 Key 对候选账号及上游模型的有效权限，不判断健康、额度或租约 + /// + /// # Errors + /// 身份不可用、查询无效或父调用取消时返回错误 + pub async fn check_access( + &self, + query: data::AccessQuery, + ) -> Result { + payload_call(self, data::ACCESS_CHECK, query).await + } + /// 读取 Key 当前启用状态与显式分组绑定，不返回密钥 /// /// # Errors @@ -52,4 +74,15 @@ impl HostClient { ) -> Result { payload_call(self, data::QUOTA_GET, query).await } + + /// 读取已有额度窗口标签、触顶与点数事实，不触发刷新或补算 + /// + /// # Errors + /// 账号不存在或宿主读取失败时返回错误 + pub async fn quota_details( + &self, + query: data::QuotaFactsQuery, + ) -> Result { + payload_call(self, data::QUOTA_DETAILS, query).await + } } diff --git a/backend/crates/gateway-plugin/sdk/src/client/http.rs b/backend/crates/gateway-plugin/sdk/src/client/http.rs index 972ac81ed..d260709ce 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/http.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/http.rs @@ -39,6 +39,20 @@ impl HostClient { .await } + /// 使用指定账号的当前代理发起请求，认证头由调用方显式提供 + /// 账号必须启用且凭据版本匹配；代理失败不回退直连 + /// + /// # Errors + /// 账号不可用、版本冲突、期限或网络请求失败时返回错误 + pub async fn account_http( + &self, + request: wire::AccountHttpRequest, + body: Vec, + ) -> Result { + self.open_http(wire::ACCOUNT_HTTP, request, body, false) + .await + } + /// 由宿主附加已选账号的认证与代理，返回统一的惰性正文 /// /// # Errors diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/config.rs b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs index 7929b9ee5..c865d7d13 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/config.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs @@ -6,6 +6,7 @@ use super::SessionError; const MAXIMUM_STREAM_CHUNK_BYTES: usize = 16 * 1024 * 1024; const MAXIMUM_CONCURRENCY: usize = 256; +const MAXIMUM_CALLBACKS: usize = 512; const MAXIMUM_BUFFERED_STREAM_CHUNKS: usize = 65_536; /// 插件侧会话的内存、并发与期限边界 @@ -14,6 +15,7 @@ pub struct SessionConfig { /// 单个业务流分块的预算，不限制普通调用正文的总长度 pub maximum_stream_chunk_bytes: usize, pub maximum_calls: usize, + /// 包含等待下游与尚未消费结果的全部宿主回调 pub maximum_callbacks: usize, pub maximum_buffered_stream_chunks: usize, pub handshake_timeout: Duration, @@ -25,7 +27,7 @@ impl Default for SessionConfig { // 流分块还需匹配宿主授予的信用窗口 maximum_stream_chunk_bytes: 1024 * 1024, maximum_calls: 256, - maximum_callbacks: 64, + maximum_callbacks: 320, maximum_buffered_stream_chunks: 1_024, handshake_timeout: Duration::from_secs(5), } @@ -39,7 +41,7 @@ impl SessionConfig { || self.maximum_calls == 0 || self.maximum_calls > MAXIMUM_CONCURRENCY || self.maximum_callbacks == 0 - || self.maximum_callbacks > MAXIMUM_CONCURRENCY + || self.maximum_callbacks > MAXIMUM_CALLBACKS || self.maximum_buffered_stream_chunks == 0 || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS || self.handshake_timeout.is_zero() diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/deadline.rs b/backend/crates/gateway-plugin/sdk/src/client/session/deadline.rs index b7e6ea324..788d1d0b3 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/deadline.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/deadline.rs @@ -1,13 +1,21 @@ //! 宿主与插件共用的调用处理预算，受管续体等待期间暂停并保留剩余额度 -use std::time::Duration; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; -use tokio::{sync::watch, time::Instant}; +use tokio::{sync::Notify, time::Instant}; /// 同一次调用的所有等待点共享此期限，避免回调恢复后仍使用旧截止时间 #[derive(Clone)] pub struct CallDeadline { - state: watch::Sender, + shared: Arc, +} + +struct Shared { + state: Mutex, + changed: Notify, } #[derive(Clone, Copy)] @@ -25,13 +33,22 @@ pub struct CallDeadlinePause { impl CallDeadline { #[must_use] pub fn new(deadline: Instant) -> Self { - let (state, _) = watch::channel(State::Running(deadline)); - Self { state } + Self { + shared: Arc::new(Shared { + state: Mutex::new(State::Running(deadline)), + changed: Notify::new(), + }), + } } #[must_use] pub fn remaining(&self) -> Option { - match *self.state.borrow() { + match *self + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + { State::Running(at) => Some(at.saturating_duration_since(Instant::now())), State::Paused(remaining) => Some(remaining), State::Unlimited => None, @@ -47,32 +64,46 @@ impl CallDeadline { /// 只有尚未耗尽的运行预算可以暂停；重复 next 不能重置或复活预算 #[must_use] pub fn pause(&self) -> Option { - self.state - .send_if_modified(|state| { - let State::Running(at) = *state else { - return false; - }; - let remaining = at.saturating_duration_since(Instant::now()); - if remaining.is_zero() { - return false; - } - *state = State::Paused(remaining); - true - }) - .then(|| CallDeadlinePause { - deadline: self.clone(), - }) + let mut state = self + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let State::Running(at) = *state else { + return None; + }; + let remaining = at.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return None; + } + *state = State::Paused(remaining); + // 暂停只延后期限，无需唤醒处理任务；旧计时器到期会重新核对状态 + // 否则父任务会把自身停表通知误当成业务唤醒，反复请求恢复 + Some(CallDeadlinePause { + deadline: self.clone(), + }) } /// 初始资源流响应已交付，后续期限由资源 owner 持有 pub fn clear(&self) { - self.state.send_replace(State::Unlimited); + *self + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = State::Unlimited; + self.shared.changed.notify_waiters(); } pub async fn wait(&self) { - let mut state = self.state.subscribe(); loop { - let current = *state.borrow_and_update(); + let changed = self.shared.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let current = *self + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let elapsed = async { match current { State::Running(at) => tokio::time::sleep_until(at).await, @@ -81,7 +112,7 @@ impl CallDeadline { }; tokio::select! { biased; - _ = state.changed() => {} + () = changed => {} () = elapsed => { // 暂停可与旧计时器唤醒并发，返回前重新核对当前预算 if self.is_elapsed() { @@ -95,12 +126,16 @@ impl CallDeadline { impl Drop for CallDeadlinePause { fn drop(&mut self) { - self.deadline.state.send_if_modified(|state| { - let State::Paused(remaining) = *state else { - return false; - }; + let mut state = self + .deadline + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let State::Paused(remaining) = *state { *state = State::Running(Instant::now() + remaining); - true - }); + drop(state); + self.deadline.shared.changed.notify_waiters(); + } } } diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs index 3627b22ef..4424e7822 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs @@ -23,6 +23,7 @@ use super::{ lifecycle::{CancellationSource, LifecycleHooks}, output::{Output, writer_loop}, registry::CallbackRegistry, + scheduling::CallSchedule, stream::validate_stream_chunk, }; @@ -107,12 +108,13 @@ where .config .maximum_calls .checked_add(self.config.maximum_callbacks) - .and_then(|capacity| capacity.checked_add(self.config.maximum_calls)) .ok_or(SessionError::Configuration)?; let control_reserve = self .config .maximum_calls - .checked_add(8) + .checked_mul(2) + .and_then(|capacity| capacity.checked_add(self.config.maximum_callbacks)) + .and_then(|capacity| capacity.checked_add(8)) .ok_or(SessionError::Configuration)?; let queue_capacity = data_capacity .checked_add(control_reserve) @@ -121,17 +123,15 @@ where let output = Output { sender: outbound, data_slots: Arc::new(Semaphore::new(data_capacity)), + failed: tokio::sync::watch::channel(false).0, }; - let callbacks = Arc::new(CallbackRegistry::new( - self.config.maximum_callbacks, - self.config.maximum_calls, - )); + let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks)); let handler: Arc = Arc::new(handler); let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?; write_frame( &mut self.writer, &Frame::control(Message::Ready { - protocol_version: PROTOCOL_VERSION, + protocol_version: self.handshake.protocol_version, incarnation: self.handshake.incarnation.clone(), call_budgets: self.call_budgets, }), @@ -167,6 +167,7 @@ struct ActiveCall { credits: Arc, cancellation: CancellationSource, task: JoinHandle<()>, + schedule: Arc, } #[derive(Clone, Copy)] @@ -243,6 +244,7 @@ impl SessionDriver { let frame = loop { tokio::select! { biased; + () = self.output.failed() => return Err(SessionError::Capacity), writer = &mut self.writer => { return writer.map_err(|_| SessionError::Closed)?; } @@ -296,6 +298,16 @@ impl SessionDriver { } Ok(false) } + Message::Resumed { id, epoch } if frame.payload.is_empty() => { + if id == 0 || id.is_multiple_of(2) || id > self.last_call || epoch == 0 { + return Err(SessionError::Protocol); + } + // 取消和完成可能先于同一授权抵达，已退休调用不能被重新打开 + if let Some(active) = self.active.get(&id) { + active.schedule.resumed(epoch)?; + } + Ok(false) + } Message::Cancel { id } if frame.payload.is_empty() => { self.cancel_call(id)?; Ok(false) @@ -367,6 +379,7 @@ impl SessionDriver { let deadline = CallDeadline::new(deadline); let credits = Arc::new(CreditWindow::default()); let (cancellation, cancellation_signal) = CancellationSource::new(); + let schedule = Arc::new(CallSchedule::new(id, self.output.clone(), deadline.clone())); let host = HostClient { parent_id: id, deadline: deadline.clone(), @@ -375,6 +388,7 @@ impl SessionDriver { output: self.output.clone(), maximum_stream_chunk_bytes: self.config.maximum_stream_chunk_bytes, credits: credits.clone(), + schedule: schedule.clone(), }; let call = PluginCall { method, @@ -392,6 +406,8 @@ impl SessionDriver { let maximum_chunks = self.config.maximum_buffered_stream_chunks; let task_context = context.clone(); let lifecycle = self.lifecycle.clone(); + let task_schedule = schedule.clone(); + let callbacks = Arc::clone(&self.callbacks); let task = tokio::spawn(async move { let guard = CompletionGuard::new(id, completed); let result = run_call( @@ -404,6 +420,8 @@ impl SessionDriver { deadline, maximum_stream_chunk_bytes, maximum_chunks, + task_schedule, + callbacks, ) .await; guard.finish(if result.is_ok() { @@ -419,6 +437,7 @@ impl SessionDriver { credits, cancellation, task, + schedule, }, ); Ok(()) @@ -478,9 +497,12 @@ async fn run_call( deadline: CallDeadline, maximum_stream_chunk_bytes: usize, maximum_chunks: usize, + schedule: Arc, + callbacks: Arc, ) -> Result<(), SessionError> { let id = context.call_id; let cancellation = call.cancellation.clone(); + let response = schedule.drive(handler.call(call), || callbacks.pauses_budget(id)); let response = tokio::select! { biased; () = cancellation.cancelled() => return Ok(()), @@ -492,7 +514,7 @@ async fn run_call( ))?; return Ok(()); } - response = handler.call(call) => response, + response = response => response?, }; let mut reply = match response { Ok(reply) => reply, diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs index d400e0794..1c53d615d 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/mod.rs @@ -9,6 +9,7 @@ mod error; mod lifecycle; mod output; mod registry; +mod scheduling; mod stream; pub use call::{CallFuture, CallReply, PluginCall, PluginHandler}; diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/output.rs b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs index acaadeca8..163578d79 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/output.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/output.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use tokio::{ io::AsyncWrite, - sync::{OwnedSemaphorePermit, Semaphore, mpsc}, + sync::{OwnedSemaphorePermit, Semaphore, mpsc, watch}, }; use crate::Frame; @@ -15,6 +15,7 @@ use super::{CallCancellation, CallDeadline, SessionError}; pub(super) struct Output { pub(super) sender: mpsc::Sender, pub(super) data_slots: Arc, + pub(super) failed: watch::Sender, } impl Clone for Output { @@ -22,6 +23,7 @@ impl Clone for Output { Self { sender: self.sender.clone(), data_slots: Arc::clone(&self.data_slots), + failed: self.failed.clone(), } } } @@ -65,7 +67,16 @@ impl Output { frame, _data_slot: None, }) - .map_err(|_| SessionError::Capacity) + .map_err(|_| { + // drop 路径也会提交资源确认，队列耗尽不能静默丢失生命周期消息 + self.failed.send_replace(true); + SessionError::Capacity + }) + } + + pub(super) async fn failed(&self) { + let mut failed = self.failed.subscribe(); + let _ = failed.wait_for(|failed| *failed).await; } } diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs index 8261737f8..5809cc29f 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs @@ -12,8 +12,8 @@ use crate::{Frame, Message, call::middleware::NEXT_METHOD}; use super::super::frame::EncodedFrame; use super::{ - CallCancellation, CallDeadline, CallDeadlinePause, SessionError, credit::CreditWindow, - output::Output, + CallCancellation, CallDeadline, SessionError, credit::CreditWindow, output::Output, + scheduling::CallSchedule, }; /// 宿主回调的成功结果；不实现 `Debug`，避免载荷进入普通诊断 @@ -32,6 +32,7 @@ pub struct HostClient { pub(super) output: Output, pub(super) maximum_stream_chunk_bytes: usize, pub(super) credits: Arc, + pub(super) schedule: Arc, } impl HostClient { @@ -75,6 +76,16 @@ impl HostClient { }; let managed_next = method == NEXT_METHOD; let (id, received) = self.callbacks.begin(self.parent_id, managed_next)?; + let mut guard = CallbackGuard { + id, + parent_id: self.parent_id, + callbacks: &self.callbacks, + output: &self.output, + sent: false, + }; + if !managed_next { + self.schedule.charge_wait(); + } let frame = Frame { message: Message::Callback { id, @@ -84,32 +95,23 @@ impl HostClient { }, payload, }; - let frame = EncodedFrame::new(frame).map_err(|_| { - self.callbacks.retire(id); - SessionError::Protocol - })?; + let frame = EncodedFrame::new(frame).map_err(|_| SessionError::Protocol)?; let sent = tokio::select! { biased; () = self.cancellation.cancelled() => Err(SessionError::Cancelled), () = self.deadline.wait() => Err(SessionError::Timeout), result = self.output.send_data(frame) => result, }; - if let Err(error) = sent { - self.callbacks.retire(id); - return Err(error); - } - if managed_next { - self.callbacks.pause_next(id, &self.deadline); - } + sent?; + guard.sent = true; + self.callbacks.mark_sent(id); drop(ordered); tokio::select! { biased; () = self.cancellation.cancelled() => { - self.callbacks.retire(id); Err(SessionError::Cancelled) } () = self.deadline.wait() => { - self.callbacks.retire(id); Err(SessionError::Timeout) } result = received => result.map_err(|_| SessionError::Closed)?, @@ -121,7 +123,6 @@ pub(super) struct CallbackRegistry { state: Mutex, order: tokio::sync::Mutex<()>, maximum: usize, - maximum_next: usize, maximum_retired: usize, } @@ -135,12 +136,33 @@ struct CallbackState { struct PendingCallback { parent_id: u64, managed_next: bool, - pause: Option, - result: oneshot::Sender>, + sent: bool, + result: Option>>, +} + +struct CallbackGuard<'a> { + id: u64, + parent_id: u64, + callbacks: &'a CallbackRegistry, + output: &'a Output, + sent: bool, +} + +impl Drop for CallbackGuard<'_> { + fn drop(&mut self) { + if self.callbacks.retire(self.id) && self.sent { + let _ = self + .output + .send_control(Frame::control(Message::CallbackConsumed { + id: self.id, + parent_id: self.parent_id, + })); + } + } } impl CallbackRegistry { - pub(super) fn new(maximum: usize, maximum_next: usize) -> Self { + pub(super) fn new(maximum: usize) -> Self { Self { state: Mutex::new(CallbackState { next_id: 2, @@ -150,8 +172,7 @@ impl CallbackRegistry { }), order: tokio::sync::Mutex::new(()), maximum, - maximum_next, - maximum_retired: (maximum + maximum_next).saturating_mul(2), + maximum_retired: maximum.saturating_mul(2), } } @@ -164,18 +185,7 @@ impl CallbackRegistry { .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let maximum = if managed_next { - self.maximum_next - } else { - self.maximum - }; - if state - .pending - .values() - .filter(|pending| pending.managed_next == managed_next) - .count() - >= maximum - { + if state.pending.len() >= self.maximum { return Err(SessionError::Capacity); } let id = state.next_id; @@ -186,35 +196,54 @@ impl CallbackRegistry { PendingCallback { parent_id, managed_next, - result, - pause: None, + sent: false, + result: Some(result), }, ); Ok((id, received)) } - fn pause_next(&self, id: u64, deadline: &CallDeadline) { + fn mark_sent(&self, id: u64) { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(pending) = state.pending.get_mut(&id) { - pending.pause = deadline.pause(); + if let Some(callback) = state.pending.get_mut(&id) { + callback.sent = true; } } + pub(super) fn pauses_budget(&self, parent_id: u64) -> bool { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut has_next = false; + for callback in state + .pending + .values() + .filter(|call| call.parent_id == parent_id) + { + if !callback.managed_next || !callback.sent { + return false; + } + has_next = true; + } + has_next + } + pub(super) fn resolve(&self, frame: Frame) -> Result<(), SessionError> { let id = match &frame.message { Message::Result { id, .. } | Message::Error { id, .. } => *id, _ => return Err(SessionError::Protocol), }; - let pending = { + let sender = { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(pending) = state.pending.remove(&id) { - Some(pending) + if let Some(pending) = state.pending.get_mut(&id) { + Some(pending.result.take().ok_or(SessionError::Protocol)?) } else if state.retired.remove(&id) { state.retired_order.retain(|retired| *retired != id); None @@ -222,7 +251,7 @@ impl CallbackRegistry { return Err(SessionError::Protocol); } }; - let Some(mut pending) = pending else { + let Some(sender) = sender else { return Ok(()); }; let result = match frame.message { @@ -235,20 +264,22 @@ impl CallbackRegistry { } _ => return Err(SessionError::Protocol), }; - // 收到宿主结果就恢复预算，不依赖插件再次轮询 next future - pending.pause.take(); - let _ = pending.result.send(result); + let _ = sender.send(result); Ok(()) } - fn retire(&self, id: u64) { + fn retire(&self, id: u64) -> bool { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.pending.remove(&id).is_some() { - retire_callback(&mut state, id, self.maximum_retired); + if let Some(pending) = state.pending.remove(&id) { + if pending.result.is_some() { + retire_callback(&mut state, id, self.maximum_retired); + } + return true; } + false } pub(super) fn finish_parent(&self, parent_id: u64) { @@ -261,7 +292,9 @@ impl CallbackRegistry { let mut finished = Vec::new(); for (id, pending) in std::mem::take(&mut state.pending) { if pending.parent_id == parent_id { - retire_callback(&mut state, id, self.maximum_retired); + if pending.result.is_some() { + retire_callback(&mut state, id, self.maximum_retired); + } finished.push(pending); } else { remaining.insert(id, pending); @@ -271,7 +304,9 @@ impl CallbackRegistry { finished }; for pending in pending { - let _ = pending.result.send(Err(SessionError::Cancelled)); + if let Some(result) = pending.result { + let _ = result.send(Err(SessionError::Cancelled)); + } } } @@ -284,7 +319,9 @@ impl CallbackRegistry { std::mem::take(&mut state.pending) }; for pending in pending.into_values() { - let _ = pending.result.send(Err(SessionError::Closed)); + if let Some(result) = pending.result { + let _ = result.send(Err(SessionError::Closed)); + } } } } diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/scheduling.rs b/backend/crates/gateway-plugin/sdk/src/client/session/scheduling.rs new file mode 100644 index 000000000..9e2097cf7 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/client/session/scheduling.rs @@ -0,0 +1,156 @@ +//! 父处理任务的协作式挂起、唤醒与宿主恢复授权 + +use std::{ + future::{Future, poll_fn}, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll, Wake, Waker}, +}; + +use crate::{Frame, Message}; + +use super::{CallDeadline, CallDeadlinePause, SessionError, output::Output}; + +/// 这里只保存协议镜像；执行名额与准入决策始终由宿主持有 +pub(super) struct CallSchedule { + id: u64, + output: Output, + deadline: CallDeadline, + state: Mutex, + ready: AtomicBool, + waker: Mutex>, +} + +struct State { + epoch: u64, + phase: Phase, + pause: Option, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Phase { + Running, + Suspended, + Resuming, +} + +impl CallSchedule { + pub(super) fn new(id: u64, output: Output, deadline: CallDeadline) -> Self { + Self { + id, + output, + deadline, + state: Mutex::new(State { + epoch: 0, + phase: Phase::Running, + pause: None, + }), + ready: AtomicBool::new(false), + waker: Mutex::new(None), + } + } + + fn state(&self) -> MutexGuard<'_, State> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub(super) fn resumed(&self, epoch: u64) -> Result<(), SessionError> { + let waker = { + let mut state = self.state(); + if state.phase != Phase::Resuming || state.epoch != epoch { + return Err(SessionError::Protocol); + } + state.phase = Phase::Running; + state.pause.take(); + self.waker + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + }; + if let Some(waker) = waker { + waker.wake(); + } + Ok(()) + } + + /// 后台工作发起应计时的回调时，不能借父任务等待 next 停表 + pub(super) fn charge_wait(&self) { + self.state().pause.take(); + } + + pub(super) async fn drive( + self: &Arc, + future: F, + pauses_budget: impl Fn() -> bool, + ) -> Result { + let mut future = std::pin::pin!(future); + let waker = Waker::from(Arc::clone(self)); + poll_fn(|cx| { + { + let mut state = self.state(); + *self + .waker + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(cx.waker().clone()); + match state.phase { + Phase::Resuming => return Poll::Pending, + Phase::Suspended => { + if self.ready.load(Ordering::Acquire) { + self.output.send_control(Frame::control(Message::Resume { + id: self.id, + epoch: state.epoch, + }))?; + state.phase = Phase::Resuming; + } + return Poll::Pending; + } + Phase::Running => self.ready.store(false, Ordering::Release), + } + } + let mut task_context = Context::from_waker(&waker); + if let Poll::Ready(result) = future.as_mut().poll(&mut task_context) { + return Poll::Ready(Ok(result)); + } + let mut state = self.state(); + state.epoch = state.epoch.checked_add(1).ok_or(SessionError::Capacity)?; + // 与回调共享单一 FIFO，宿主看到挂起前一定已看到本次轮询提交的回调 + self.output.send_control(Frame::control(Message::Suspend { + id: self.id, + epoch: state.epoch, + }))?; + state.phase = Phase::Suspended; + if pauses_budget() { + state.pause = self.deadline.pause(); + } + if self.ready.load(Ordering::Acquire) { + // future 可以在 poll 返回 Pending 前唤醒自己，这次唤醒也必须经过授权 + cx.waker().wake_by_ref(); + } + Poll::Pending + }) + .await + } +} + +impl Wake for CallSchedule { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + // 期限恢复会同步唤醒回调等待者；唤醒路径不能再次取得调度状态锁 + self.ready.store(true, Ordering::Release); + let waker = self + .waker + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + if let Some(waker) = waker { + waker.wake(); + } + } +} diff --git a/backend/crates/gateway-plugin/sdk/src/message.rs b/backend/crates/gateway-plugin/sdk/src/message.rs index ad3531185..962cab0f6 100644 --- a/backend/crates/gateway-plugin/sdk/src/message.rs +++ b/backend/crates/gateway-plugin/sdk/src/message.rs @@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize}; use crate::{CallBudgets, CallContext, Handshake, PluginFault}; -pub const PROTOCOL_VERSION: u32 = 6; +pub const PROTOCOL_VERSION: u32 = 7; /// 元数据与二进制载荷分开，流分块不经过 JSON/base64 #[derive(Clone, PartialEq, Serialize, Deserialize)] @@ -42,6 +42,26 @@ pub enum Message { #[serde(skip)] params: serde_json::Value, }, + /// 父处理任务已交还执行权，epoch 在同一调用内严格递增 + Suspend { + id: u64, + epoch: u64, + }, + /// 已挂起任务被唤醒，请求重新取得执行名额 + Resume { + id: u64, + epoch: u64, + }, + /// 宿主已取得执行名额，插件才可再次驱动父处理任务 + Resumed { + id: u64, + epoch: u64, + }, + /// 回调 future 已消费结果或被释放，宿主可回收对应驻留额度 + CallbackConsumed { + id: u64, + parent_id: u64, + }, Cancel { id: u64, }, @@ -102,6 +122,10 @@ impl std::fmt::Debug for Message { Self::Result { id, .. } => ("Result", Some(id)), Self::Error { id, .. } => ("Error", Some(id)), Self::Callback { id, .. } => ("Callback", Some(id)), + Self::Suspend { id, .. } => ("Suspend", Some(id)), + Self::Resume { id, .. } => ("Resume", Some(id)), + Self::Resumed { id, .. } => ("Resumed", Some(id)), + Self::CallbackConsumed { id, .. } => ("CallbackConsumed", Some(id)), Self::Cancel { id } => ("Cancel", Some(id)), Self::Cancelled { id } => ("Cancelled", Some(id)), Self::Stream { id, .. } => ("Stream", Some(id)), diff --git a/backend/crates/gateway-plugin/sdk/tests/client/http.rs b/backend/crates/gateway-plugin/sdk/tests/client/http.rs index 7afd32dfc..bbfd210b2 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/http.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/http.rs @@ -5,7 +5,7 @@ use gateway_plugin_sdk::client::{CallFuture, CallReply, PluginCall, PluginHandle use gateway_plugin_sdk::{ErrorCode, Frame, Message}; use gateway_plugin_sdk::{ call::{ - host::HttpRequest, + host::{AccountHttpRequest, HttpRequest}, upstream_adapter::{UpstreamHttpRequest, UpstreamWebSocketRequest, WebSocketMessageKind}, }, client::UpstreamWebSocketUpgrade, @@ -91,6 +91,21 @@ impl PluginHandler for NetworkHandler { call.payload, ) .await? + } else if call.method == "account" { + call.host + .account_http( + AccountHttpRequest { + account_id: "acct_one".into(), + credential_revision: 7, + request: HttpRequest { + method: "POST".into(), + url: "https://example.test/responses".into(), + headers: vec![], + }, + }, + call.payload, + ) + .await? } else { call.host .upstream_http( @@ -157,7 +172,7 @@ async fn reply_callback( #[tokio::test] async fn http_response_objects_share_pull_semantics_and_retire_eof_handles() { - for method in ["http", "upstream"] { + for method in ["http", "account", "upstream"] { let (mut host, task) = start_session(NetworkHandler).await; send_call( &mut host, @@ -167,19 +182,28 @@ async fn http_response_objects_share_pull_semantics_and_retire_eof_handles() { b"request".to_vec(), ) .await; - let prefix = if method == "http" { + let prefix = if method != "upstream" { "host.http" } else { "host.upstream.http" }; + let open_method = if method == "account" { + gateway_plugin_sdk::call::host::ACCOUNT_HTTP.to_owned() + } else { + format!("{prefix}.do_stream") + }; let frame = reply_callback( &mut host, - &format!("{prefix}.do_stream"), + &open_method, json!({"status":200,"http_version":"HTTP/1.1","trailers":[],"headers":[["x-bytes",[255]],["x-bytes",b"other".to_vec()]],"stream":"private-handle"}), vec![], ) .await; assert_eq!(frame.payload, b"request"); + if method == "account" { + assert!(matches!(frame.message, Message::Callback { params, .. } + if params["account_id"] == "acct_one" && params["credential_revision"] == 7)); + } let frame = reply_callback( &mut host, &format!("{prefix}.stream_read"), diff --git a/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs index a4dd20f56..3528d8b23 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs @@ -519,10 +519,25 @@ async fn send_call( } async fn receive(host: &mut HostPeer) -> Frame { - tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) - .await - .expect("plugin response timed out") - .expect("plugin response frame must be valid") + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let frame = read_frame(&mut host.reader).await?; + match frame.message { + Message::Suspend { .. } | Message::CallbackConsumed { .. } => {} + Message::Resume { id, epoch } => { + write_frame( + &mut host.writer, + &Frame::control(Message::Resumed { id, epoch }), + ) + .await?; + } + _ => return Ok::<_, gateway_plugin_sdk::FrameError>(frame), + } + } + }) + .await + .expect("plugin response timed out") + .expect("plugin response frame must be valid") } fn unwrap_result(frame: Frame, expected_id: u64) -> (Value, Vec) { diff --git a/backend/crates/gateway-plugin/sdk/tests/client/read.rs b/backend/crates/gateway-plugin/sdk/tests/client/read.rs index bcbf3136e..e19da5f11 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/read.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/read.rs @@ -111,7 +111,7 @@ async fn verify_cancelled_read(method: &str) { .await; let first = callback_id(&mut host, read).await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; // 已取消的读取随后成功，下一次 read 应恢复同一读取结果 reply(&mut host, first, data.clone(), b"first").await; signals.resume.notify_one(); @@ -264,7 +264,7 @@ async fn middleware_http_and_websocket_reads_resume_the_pending_chunk() { }; let id = callback_id(&mut host, method).await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; reply(&mut host, id, json!({"eof":false}), b"first").await; signals.resume.notify_one(); let frame = session::receive(&mut host).await; @@ -296,7 +296,7 @@ async fn model_frame_survives_cancellation_while_fetching_its_facts() { .await; let id = callback_id(&mut host, BODY_FACTS_METHOD).await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; reply(&mut host, id, json!({"present":false}), b"").await; signals.resume.notify_one(); let frame = session::receive(&mut host).await; @@ -314,7 +314,7 @@ async fn http_forwarding_retains_a_pending_read_under_output_backpressure() { ) .await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; reply(&mut host, id, json!({"eof":false}), b"first").await; signals.resume.notify_one(); write_frame( @@ -359,7 +359,7 @@ async fn pending_http_read_is_split_when_credit_arrives_after_the_callback() { ) .await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; let expected = vec![b'a'; 8192]; reply(&mut host, id, json!({"eof":false}), &expected).await; signals.resume.notify_one(); @@ -417,7 +417,7 @@ async fn websocket_transfer_rejects_a_payload_with_a_pending_read() { let (mut host, task, signals) = start_middleware(Action::WebSocketTransfer).await; let _id = callback_id(&mut host, BODY_READ_METHOD).await; signals.cancel.notify_one(); - signals.dropped.notified().await; + session::wait_suspended(&mut host, &signals.dropped).await; signals.resume.notify_one(); assert!(matches!(session::receive(&mut host).await.message, Message::Error { id: 1, error } if error.code == ErrorCode::InvalidInput)); diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs index 8f1c5c040..7e500d3ab 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs @@ -268,12 +268,29 @@ async fn completed_call_does_not_discard_a_partially_read_next_frame() { tokio::time::timeout(Duration::from_secs(2), async { // 单字节缓冲保证前缀已经进入读取 future，而不是仍滞留在传输队列 let (mut host, task) = start_session_with_capacity(TestHandler::default(), 1).await; - send_call(&mut host, 1, "slow", json!("first"), Vec::new()).await; - host.writer.write_all(&encoded[..split]).await.unwrap(); + send_call(&mut host, 1, "delayed_stream", json!({}), Vec::new()).await; + send_control( + &mut host, + Message::Credit { + id: 1, + bytes: 1_024, + frames: 1, + }, + ) + .await; assert!(matches!( receive(&mut host).await.message, Message::Result { id: 1, .. } )); + host.writer.write_all(&encoded[..split]).await.unwrap(); + assert!(matches!( + receive(&mut host).await.message, + Message::Stream { id: 1, sequence: 0 } + )); + assert!(matches!( + receive(&mut host).await.message, + Message::End { id: 1, error: None } + )); host.writer.write_all(&encoded[split..]).await.unwrap(); assert_eq!( receive(&mut host).await, diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs index 571001cb7..e65d95e5a 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs @@ -5,6 +5,7 @@ mod deadline; mod driver; mod lifecycle; mod registry; +mod scheduling; use std::{ future::pending, @@ -153,6 +154,14 @@ impl PluginHandler for TestHandler { tokio::time::sleep(Duration::from_millis(80)).await; Ok(CallReply::unary(call.params, call.payload)) } + "delayed_stream" => { + let (sender, stream) = ResponseStream::channel(NonZeroUsize::new(1).unwrap()); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(80)).await; + sender.send(b"done".to_vec()).await.unwrap(); + }); + Ok(CallReply::stream(json!({}), Vec::new(), stream)) + } "callback" | "next_callback" => { let method = if call.method == "next_callback" { NEXT_METHOD @@ -466,10 +475,49 @@ async fn send_control(host: &mut HostPeer, message: Message) { } pub(super) async fn receive(host: &mut HostPeer) -> Frame { - tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) - .await - .expect("plugin response timed out") - .expect("plugin response frame must be valid") + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let frame = read_frame(&mut host.reader).await?; + match frame.message { + Message::Suspend { .. } | Message::CallbackConsumed { .. } => {} + Message::Resume { id, epoch } => { + send_control(host, Message::Resumed { id, epoch }).await; + } + _ => return Ok::<_, gateway_plugin_sdk::FrameError>(frame), + } + } + }) + .await + .expect("plugin response timed out") + .expect("plugin response frame must be valid") +} + +// 等待测试事件时仍推进恢复授权，且只在完整挂起帧边界返回，避免取消半帧读取 +pub(super) async fn wait_suspended(host: &mut HostPeer, event: &tokio::sync::Notify) { + let event = event.notified(); + tokio::pin!(event); + loop { + let frame = tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) + .await + .unwrap() + .unwrap(); + match frame.message { + Message::Suspend { .. } => { + if std::future::poll_fn(|cx| { + std::task::Poll::Ready(event.as_mut().poll(cx).is_ready()) + }) + .await + { + return; + } + } + Message::Resume { id, epoch } => { + send_control(host, Message::Resumed { id, epoch }).await + } + Message::CallbackConsumed { .. } => {} + message => panic!("unexpected frame while waiting for suspension: {message:?}"), + } + } } pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle>) { diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs index 473867923..cda12ef26 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/registry.rs @@ -297,57 +297,3 @@ async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() { } } } - -#[tokio::test] -async fn managed_next_does_not_consume_the_bounded_ordinary_callback_budget() { - let (mut host, task) = start_configured_session( - TestHandler::default(), - MAXIMUM_STREAM_CHUNK_BYTES * 2, - SessionConfig { - maximum_calls: 4, - maximum_callbacks: 1, - ..Default::default() - }, - ) - .await; - send_call(&mut host, 1, "next_callback", json!({}), Vec::new()).await; - let Message::Callback { - id: next, - parent_id: 1, - method, - .. - } = receive(&mut host).await.message - else { - panic!("expected next callback") - }; - assert_eq!(method, NEXT_METHOD); - send_call(&mut host, 3, "callback", json!({}), Vec::new()).await; - let Message::Callback { - id: ordinary, - parent_id: 3, - method, - .. - } = receive(&mut host).await.message - else { - panic!("next must not occupy ordinary callback capacity") - }; - assert_eq!(method, "host.log"); - send_call(&mut host, 5, "callback", json!({}), Vec::new()).await; - assert!( - matches!(receive(&mut host).await.message, Message::Error { id: 5, error } if error.code == ErrorCode::Capacity) - ); - for (callback, parent) in [(ordinary, 3), (next, 1)] { - send_control( - &mut host, - Message::Result { - id: callback, - result: json!({}), - }, - ) - .await; - assert!( - matches!(receive(&mut host).await.message, Message::Result { id, .. } if id == parent) - ); - } - shutdown(&mut host, task).await; -} diff --git a/backend/crates/gateway-plugin/sdk/tests/client/session/scheduling.rs b/backend/crates/gateway-plugin/sdk/tests/client/session/scheduling.rs new file mode 100644 index 000000000..9a29e88ea --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/scheduling.rs @@ -0,0 +1,332 @@ +//! 协作调度的部分完成、唤醒竞争、预算与回调回收合同测试 + +use std::{ + future::{Future, poll_fn}, + task::Poll, +}; + +use super::*; + +struct ScheduledHandler { + progressed: Arc, +} + +impl PluginHandler for ScheduledHandler { + fn call(&self, call: PluginCall) -> CallFuture<'_> { + Box::pin(async move { + match call.method.as_str() { + "parallel" | "mixed" => { + let first_method = if call.method == "mixed" { + NEXT_METHOD + } else { + "host.first" + }; + let first = async { + let reply = call.host.call(first_method, json!({}), Vec::new()).await; + self.progressed.fetch_add(1, Ordering::SeqCst); + reply + }; + let second = async { + let reply = call.host.call("host.second", json!({}), Vec::new()).await; + self.progressed.fetch_add(1, Ordering::SeqCst); + reply + }; + let (first, second) = tokio::join!(first, second); + first.map_err(SessionError::into_plugin_fault)?; + second.map_err(SessionError::into_plugin_fault)?; + } + "yield" => { + tokio::task::yield_now().await; + self.progressed.fetch_add(1, Ordering::SeqCst); + } + "drop_callback" => { + let mut callback = + Box::pin(call.host.call("host.detached", json!({}), Vec::new())); + poll_fn(|cx| { + assert!(callback.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(callback); + } + "saturate" => { + let mut callbacks = [ + Box::pin(call.host.call(NEXT_METHOD, json!({}), Vec::new())), + Box::pin(call.host.call(NEXT_METHOD, json!({}), Vec::new())), + ]; + for callback in &mut callbacks { + poll_fn(|cx| { + assert!(callback.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + } + let refused = call.host.call("host.third", json!({}), Vec::new()).await; + assert!(matches!(refused, Err(SessionError::Capacity))); + } + _ => { + return Err(PluginFault::new( + ErrorCode::Unsupported, + "unknown scheduling test call", + )); + } + } + Ok(CallReply::unary(json!({}), Vec::new())) + }) + } +} + +async fn wire(host: &mut HostPeer) -> Message { + let frame = tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader)) + .await + .expect("expected protocol progress") + .unwrap(); + assert!(frame.payload.is_empty()); + frame.message +} + +async fn callback(host: &mut HostPeer, method: &str) -> u64 { + let Message::Callback { + id, + parent_id: 1, + method: actual, + .. + } = wire(host).await + else { + panic!("expected callback before suspension"); + }; + assert_eq!(actual, method); + id +} + +async fn callback_result(host: &mut HostPeer, id: u64) { + send_control( + host, + Message::Result { + id, + result: json!({}), + }, + ) + .await; +} + +#[tokio::test] +async fn partial_callback_results_require_a_new_execution_grant_before_processing() { + let progressed = Arc::new(AtomicUsize::new(0)); + let (mut host, task) = start_session(ScheduledHandler { + progressed: Arc::clone(&progressed), + }) + .await; + send_call(&mut host, 1, "parallel", json!({}), Vec::new()).await; + let first = callback(&mut host, "host.first").await; + let second = callback(&mut host, "host.second").await; + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 1 } + )); + + callback_result(&mut host, first).await; + assert!(matches!( + wire(&mut host).await, + Message::Resume { id: 1, epoch: 1 } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 0); + send_control(&mut host, Message::Resumed { id: 1, epoch: 1 }).await; + assert!( + matches!(wire(&mut host).await, Message::CallbackConsumed { id, parent_id: 1 } if id == first) + ); + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 2 } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 1); + + callback_result(&mut host, second).await; + assert!(matches!( + wire(&mut host).await, + Message::Resume { id: 1, epoch: 2 } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 1); + send_control(&mut host, Message::Resumed { id: 1, epoch: 2 }).await; + assert!( + matches!(wire(&mut host).await, Message::CallbackConsumed { id, parent_id: 1 } if id == second) + ); + assert!(matches!( + wire(&mut host).await, + Message::Result { id: 1, .. } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 2); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn wake_during_poll_is_preserved_but_cannot_skip_host_authorization() { + let progressed = Arc::new(AtomicUsize::new(0)); + let (mut host, task) = start_session(ScheduledHandler { + progressed: Arc::clone(&progressed), + }) + .await; + send_call(&mut host, 1, "yield", json!({}), Vec::new()).await; + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 1 } + )); + assert!(matches!( + wire(&mut host).await, + Message::Resume { id: 1, epoch: 1 } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 0); + send_control(&mut host, Message::Resumed { id: 1, epoch: 1 }).await; + assert!(matches!( + wire(&mut host).await, + Message::Result { id: 1, .. } + )); + assert_eq!(progressed.load(Ordering::SeqCst), 1); + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn cancellation_during_resume_does_not_allow_late_grants_to_revive_the_call() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call(&mut host, 1, "callback", json!({}), Vec::new()).await; + let child = callback(&mut host, "host.log").await; + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 1 } + )); + callback_result(&mut host, child).await; + assert!(matches!( + wire(&mut host).await, + Message::Resume { id: 1, epoch: 1 } + )); + send_control(&mut host, Message::Cancel { id: 1 }).await; + send_control(&mut host, Message::Resumed { id: 1, epoch: 1 }).await; + assert!(matches!( + wire(&mut host).await, + Message::Cancelled { id: 1 } + )); + send_call(&mut host, 3, "echo", json!("alive"), Vec::new()).await; + assert!( + matches!(wire(&mut host).await, Message::Result { id: 3, result } if result == json!("alive")) + ); + shutdown(&mut host, task).await; +} + +#[tokio::test(start_paused = true)] +async fn next_pause_covers_resume_wait_but_cannot_hide_parallel_chargeable_work() { + let (mut host, task) = start_session(TestHandler::default()).await; + send_call_with_timeout( + &mut host, + 1, + "next_callback", + json!({}), + Vec::new(), + Duration::from_millis(100), + ) + .await; + let child = callback(&mut host, NEXT_METHOD).await; + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 1 } + )); + tokio::time::advance(Duration::from_secs(2)).await; + callback_result(&mut host, child).await; + assert!(matches!( + wire(&mut host).await, + Message::Resume { id: 1, epoch: 1 } + )); + tokio::time::advance(Duration::from_secs(2)).await; + send_control(&mut host, Message::Resumed { id: 1, epoch: 1 }).await; + assert!( + matches!(wire(&mut host).await, Message::CallbackConsumed { id, parent_id: 1 } if id == child) + ); + assert!(matches!( + wire(&mut host).await, + Message::Result { id: 1, .. } + )); + shutdown(&mut host, task).await; + + let (mut host, task) = start_session(ScheduledHandler { + progressed: Arc::new(AtomicUsize::new(0)), + }) + .await; + send_call_with_timeout( + &mut host, + 1, + "mixed", + json!({}), + Vec::new(), + Duration::from_millis(100), + ) + .await; + callback(&mut host, NEXT_METHOD).await; + callback(&mut host, "host.second").await; + assert!(matches!( + wire(&mut host).await, + Message::Suspend { id: 1, epoch: 1 } + )); + tokio::time::advance(Duration::from_millis(101)).await; + loop { + match wire(&mut host).await { + Message::CallbackConsumed { parent_id: 1, .. } => {} + Message::Error { id: 1, error } => { + assert_eq!(error.code, ErrorCode::Timeout); + break; + } + other => panic!("chargeable work must time out without a grant: {other:?}"), + } + } + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn dropping_a_callback_reports_consumption_before_the_parent_result() { + let (mut host, task) = start_session(ScheduledHandler { + progressed: Arc::new(AtomicUsize::new(0)), + }) + .await; + send_call(&mut host, 1, "drop_callback", json!({}), Vec::new()).await; + let child = callback(&mut host, "host.detached").await; + assert!( + matches!(wire(&mut host).await, Message::CallbackConsumed { id, parent_id: 1 } if id == child) + ); + assert!(matches!( + wire(&mut host).await, + Message::Result { id: 1, .. } + )); + callback_result(&mut host, child).await; + shutdown(&mut host, task).await; +} + +#[tokio::test] +async fn all_callback_methods_share_the_same_bounded_resident_capacity() { + let (mut host, task) = start_configured_session( + ScheduledHandler { + progressed: Arc::new(AtomicUsize::new(0)), + }, + MAXIMUM_STREAM_CHUNK_BYTES * 2, + SessionConfig { + maximum_calls: 1, + maximum_callbacks: 2, + ..Default::default() + }, + ) + .await; + send_call(&mut host, 1, "saturate", json!({}), Vec::new()).await; + let first = callback(&mut host, NEXT_METHOD).await; + let second = callback(&mut host, NEXT_METHOD).await; + let mut consumed = Vec::new(); + for _ in 0..2 { + let Message::CallbackConsumed { id, parent_id: 1 } = wire(&mut host).await else { + panic!("resident callbacks must be reclaimed"); + }; + consumed.push(id); + } + consumed.sort_unstable(); + assert_eq!(consumed, [first, second]); + assert!(matches!( + wire(&mut host).await, + Message::Result { id: 1, .. } + )); + shutdown(&mut host, task).await; +} diff --git a/docs/architecture.md b/docs/architecture.md index 495fbe13a..2db106f05 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -183,6 +183,15 @@ flowchart LR Provider 来源与执行事件快照可由插件读取，包含费用明细、上游观测、会话状态和原始 wire；读取快照不转移计量和资源所有权。 数据面插件通过受管 HTTP 已发送、无法证明未发送，或进入主动 HTTP／公开服务调用链时，Runtime 把该事实并入 Core 的请求副作用水位，后续不能按 Provider 的 `not_sent` 结果透明重放；RPC 返回或取消前同步封存未完成的发送，并禁止已关闭调用再分派 HTTP／公开服务 +- **身份查询与执行准入分离**：Core 的 `AuthenticatedClient` 持有身份和冻结权限，模型目录与账号授权直接读取这份事实。 + Runtime 只按父调用缓存身份快照，查询不建立模型执行图；真正发起模型执行时才检查深度、建立执行归属并申请原生资源。 + 认证复用原有认证计划与 principal 映射，通用 `InvocationOrigin` 传递父调用来源、取消、冻结计划和副作用归属，认证不依赖模型执行上下文 +- **挂起与恢复由共同边界管理**：SDK 在父处理 future 返回 `Pending` 后报告挂起，再次 poll 前等待 Runtime 的执行授权。 + Runtime 的调用表持有唯一准入状态，挂起保留总在途名额；结果接收独立于任务恢复，回调额度持续保留到消费、丢弃或父调用终结。 + 并行部分完成可以恢复处理，再次等待时重新归还名额；取消统一收尾，晚到结果与恢复授权不复活调用。 + 是否查询事实、是否调用下游插件、等待是否计时分别由能力合同确定，不用来决定执行名额的归还。 + 这是一项协作式协议，不分析插件私有线程、锁和业务依赖；资源不足返回明确过载，宿主不重放或补偿插件操作。 + 协议与容量见 [SDK 通信边界](../backend/crates/gateway-plugin/sdk/README.md#通信边界) - **主动 HTTP 请求复用总路由**：Runtime 通过 Core 的 `http::Dispatcher` 调用 API 已组装的路由，SDK 的 `dispatch_http` 与 `next` 共用正文资源池。 子调用持有父子关联、冻结计划和取消信号，跳过祖先发起实例；模型 request 与 attempt 延续该集合。内部管理请求使用宿主签发的插件身份，审计记录插件与调用来源，普通客户端不能用 headers 声明此身份。 正文转交延续真实交付寿命，EOF、显式关闭或父级取消释放资源；组合根保活 Dispatcher，Runtime 只保存 Weak，不增加逐接口注册清单 @@ -216,8 +225,9 @@ flowchart LR 注册在进程启动与握手完成后使用自己的预算；CLI 的执行与结果提交共用一次期限，重试委托链共享匹配实例中最长的预算并受请求剩余期限约束，单次调用仍受自身声明限制。 业务预算不延长进程启动与握手、排队、取消确认、关闭和外部请求期限；声明与默认值见 [SDK 调用预算](../backend/crates/gateway-plugin/sdk/README.md#声明调用预算) - **资源跟随调用**：账号、HTTP、模型、状态和日志回调绑定有效父调用，取消随调用传播。HTTP 与模型 request/attempt 中间件在首个流式响应头交付后，由正文 owner 管理取消与终态，不沿用初始处理期限；网络操作与模型执行各自计时。 - Runtime 与 SDK 的中间件处理预算在宿主 `next` 等待期间暂停；Runtime 同时归还执行名额，返回后取得名额再以剩余处理预算恢复原调用。 - 执行调度由 Runtime 统一负责，准入和恢复等待使用独立期限；初次排队、总在途上下文和普通回调各有上限；`next` 回调按在途上限独立计数。 + Runtime 与 SDK 以父处理任务的实际挂起和恢复授权管理执行名额，不按回调方法决定是否让出； + `next` 仅在父任务挂起且没有其他应计时回调时暂停处理预算，暂停覆盖恢复排队。 + 执行调度由 Runtime 统一负责，准入和恢复等待使用独立期限；初次排队、总在途上下文和驻留回调各有上限，`next` 与普通回调共享驻留额度。 流、取消确认与消费方共同持有同一在途租约，执行名额只覆盖首结果前的处理；调用等待与回调有效期共用各端的预算状态，请求总期限独立生效。 显式 Key 的身份快照随调用保留，每个模型执行拥有独立子调用图与结算；省略 Key 时继承现有模型父请求。页面 Responses 桥复核精确实例目标和 revision `frontend_authentication` 使用 principal 到 Key 的映射；公开登录回调的一次性票据用于关联流程，不限制插件回调能力。接口合同见 [SDK 回调](../backend/crates/gateway-plugin/sdk/docs/capabilities.md#宿主资源回调)