From 0a87ac7f767a5b8bb95c212b4987f906dfb2ca49 Mon Sep 17 00:00:00 2001 From: Zyy <0x@ainz.cc> Date: Sun, 25 Oct 2026 22:39:23 +0800 Subject: [PATCH] fix(plugins)!: honor declared stage call budgets (#513) Resolve scenario defaults and plugin declarations consistently across the SDK and runtime. Preserve maintenance across unrelated publications and cancel long preparation when the runtime shuts down. BREAKING CHANGE: Plugin RPC requires protocol v6. Rebuild plugins with the updated SDK; previous protocol versions are not supported. Co-authored-by: Zyy Co-authored-by: Codex --- .../runtime/plugin-host-compatibility.json | 2 +- .../runtime/src/adapter/catalog.rs | 4 +- .../src/adapter/command_line/execution.rs | 17 +- .../src/adapter/command_line/registration.rs | 4 +- .../src/adapter/frontend_authentication.rs | 12 +- .../src/adapter/management/callback.rs | 7 +- .../runtime/src/adapter/management/mod.rs | 2 +- .../src/adapter/management/registration.rs | 4 +- .../src/adapter/middleware/execution.rs | 32 ++-- .../runtime/src/adapter/middleware/http.rs | 5 +- .../runtime/src/adapter/middleware/mod.rs | 8 +- .../src/adapter/middleware/outbound.rs | 10 +- .../runtime/src/adapter/middleware/request.rs | 36 ++-- .../runtime/src/adapter/middleware/service.rs | 5 +- .../src/adapter/middleware/websocket.rs | 5 +- .../runtime/src/adapter/observer.rs | 17 +- .../runtime/src/adapter/policy/mod.rs | 11 +- .../runtime/src/adapter/policy/retry.rs | 45 +++-- .../src/adapter/policy/route_schedule.rs | 21 +-- .../src/adapter/upstream_adapter/execution.rs | 2 +- .../adapter/upstream_adapter/registration.rs | 4 +- .../src/generation/prepare/instance.rs | 104 ++++++++---- .../src/generation/prepare/maintenance.rs | 154 +++++++++++++---- .../runtime/src/generation/prepare/runtime.rs | 12 +- .../runtime/src/generation/prepare/set.rs | 14 +- .../src/generation/prepare/snapshot.rs | 55 +++--- .../runtime/src/generation/prepare/state.rs | 13 +- .../gateway-plugin/runtime/src/rpc/session.rs | 80 ++++++--- .../tests/adapter/command_line/execution.rs | 114 +++++++++++++ .../runtime/tests/adapter/command_line/mod.rs | 46 ++++- .../tests/generation/policy/http/mod.rs | 20 ++- .../generation/policy/http/websocket/mod.rs | 21 +-- .../tests/generation/policy/lifecycle.rs | 13 +- .../runtime/tests/generation/policy/mod.rs | 105 +++++++----- .../tests/generation/policy/outbound.rs | 17 +- .../runtime/tests/generation/policy/retry.rs | 6 +- .../tests/generation/prepare/instance.rs | 52 ++++++ .../tests/generation/prepare/maintenance.rs | 157 ++++++++++++++++++ .../runtime/tests/generation/prepare/mod.rs | 1 + .../tests/generation/upstream_adapter.rs | 10 +- .../runtime/tests/rpc/session.rs | 129 +++++++++++++- .../runtime/tests/support/middleware/mod.rs | 5 +- .../runtime/tests/support/worker/mod.rs | 47 +++++- backend/crates/gateway-plugin/sdk/README.md | 53 +++++- .../gateway-plugin/sdk/docs/capabilities.md | 4 +- .../gateway-plugin/sdk/docs/manifest.md | 4 +- .../crates/gateway-plugin/sdk/src/budget.rs | 119 +++++++++++++ .../sdk/src/client/session/config.rs | 3 - .../sdk/src/client/session/driver.rs | 23 ++- .../sdk/src/client/session/error.rs | 10 +- .../sdk/src/client/session/registry.rs | 7 + backend/crates/gateway-plugin/sdk/src/lib.rs | 2 + .../crates/gateway-plugin/sdk/src/message.rs | 5 +- .../crates/gateway-plugin/sdk/tests/budget.rs | 49 ++++++ .../sdk/tests/client/plugin/mod.rs | 2 +- .../sdk/tests/client/session/driver.rs | 80 +++++++++ .../sdk/tests/client/session/mod.rs | 15 +- .../crates/gateway-plugin/sdk/tests/main.rs | 1 + docs/architecture.md | 5 +- 59 files changed, 1432 insertions(+), 378 deletions(-) create mode 100644 backend/crates/gateway-plugin/runtime/tests/generation/prepare/maintenance.rs create mode 100644 backend/crates/gateway-plugin/sdk/src/budget.rs create mode 100644 backend/crates/gateway-plugin/sdk/tests/budget.rs diff --git a/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json b/backend/crates/gateway-plugin/runtime/plugin-host-compatibility.json index 86db0d164..6771f7d4a 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": [5], + "protocol_versions": [6], "capabilities": [ { "capability": "frontend_authentication", "versions": [2] }, { "capability": "scheduler", "versions": [3] }, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/catalog.rs b/backend/crates/gateway-plugin/runtime/src/adapter/catalog.rs index 926af5bcd..a5e0951ac 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/catalog.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/catalog.rs @@ -1,7 +1,5 @@ //! 编译插件模型目录贡献并校验能力绑定 -use std::time::Duration; - use gateway_admin::model::{ AdminError, plugins::instances::{PluginCapabilityBinding, PluginInstance}, @@ -23,7 +21,7 @@ pub(crate) async fn prepare( let reply = session .call( "model_catalog.register", - session.context(Stage::Registration, Duration::from_secs(5)), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), Vec::new(), ) diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/execution.rs b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/execution.rs index ff5f9b9c6..01ec04773 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/execution.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/execution.rs @@ -48,12 +48,13 @@ impl PluginCommandSession { let payload = serde_json::to_vec(&invocation) .map_err(|_| AdminError::invalid("插件命令参数无法编码"))?; let mut saved_accounts = 0; + let deadline = tokio::time::Instant::now() + command.session.budget(Stage::CommandLine); let result = tokio::select! { biased; () = cancellation.cancelled() => Ok(Err("调用已取消")), - result = tokio::time::timeout( - self.limits.maximum_call_timeout, - command.execute(payload, self.limits, &mut saved_accounts), + result = tokio::time::timeout_at( + deadline, + command.execute(payload, self.limits, deadline, &mut saved_accounts), ) => result, }; match result { @@ -76,11 +77,14 @@ impl PluginCommand { &self, payload: Vec, limits: RpcLimits, + deadline: tokio::time::Instant, saved_accounts: &mut usize, ) -> Result { - let mut context = self - .session - .context(Stage::CommandLine, limits.maximum_call_timeout); + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err("超过命令总期限"); + } + let mut context = self.session.context(Stage::CommandLine, remaining); context.request_id = Some(uuid::Uuid::new_v4().to_string()); let scope = self .callbacks @@ -135,6 +139,7 @@ fn rpc_failure(error: RpcError) -> &'static str { } RpcError::Capacity | RpcError::Overloaded => "插件调用容量不足", RpcError::Context(_) => "调用上下文无效", + RpcError::Budget(_) => "调用预算冲突", RpcError::Remote(_) => "插件拒绝命令", } } diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/registration.rs b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/registration.rs index 573cee3d3..177c961b5 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/command_line/registration.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/command_line/registration.rs @@ -4,7 +4,7 @@ use super::PluginCommand; use crate::{RpcSession, callback::PluginCallbacks}; use gateway_admin::model::{AdminError, plugins::instances::PluginInstance}; use gateway_plugin_sdk::{Capability, Manifest, Stage, call::management::CommandRegistration}; -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; const MAXIMUM_REGISTRATION_BYTES: usize = 64 * 1024; @@ -23,7 +23,7 @@ pub(crate) async fn prepare( let reply = session .call( "command_line.register", - session.context(Stage::Registration, Duration::from_secs(5)), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), Vec::new(), ) 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 fe9c60df8..4268dd024 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/frontend_authentication.rs @@ -1,6 +1,6 @@ //! 编译插件客户端认证绑定，并将认证结果映射为宿主身份 -use std::{collections::BTreeMap, sync::Arc, time::Duration}; +use std::{collections::BTreeMap, sync::Arc}; use futures::future::BoxFuture; use gateway_admin::model::{ @@ -26,8 +26,6 @@ use gateway_plugin_sdk::{ use crate::{RpcSession, callback::PluginCallbacks}; -const IDENTIFIER_TIMEOUT: Duration = Duration::from_secs(5); - #[derive(Clone)] pub(crate) struct FrontendAuthenticationEntry { identities: BTreeMap, @@ -47,7 +45,7 @@ pub(crate) async fn prepare_entry( let reply = session .call( "frontend_auth.identifier", - session.context(Stage::Registration, IDENTIFIER_TIMEOUT), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), Vec::new(), ) @@ -167,19 +165,16 @@ fn valid_identifier(value: &str) -> bool { pub(crate) struct PluginFrontendAuthenticationPlan { entry: FrontendAuthenticationEntry, - timeout: Duration, } impl PluginFrontendAuthenticationPlan { pub(crate) fn compile( mut entries: Vec, - maximum_call_timeout: Duration, ) -> Result>, AdminError> { match entries.len() { 0 => Ok(None), 1 => Ok(Some(Arc::new(Self { entry: entries.pop().expect("one entry was checked"), - timeout: maximum_call_timeout.min(Duration::from_secs(30)), }))), _ => Err(AdminError::invalid("同一发布集合只能启用一个入口认证插件")), } @@ -197,7 +192,8 @@ impl FrontendAuthenticationPlan for PluginFrontendAuthenticationPlan { .invocation .as_ref() .ok_or(FrontendAuthenticationError)?; - let mut context = session.context(Stage::Authentication, self.timeout); + let mut context = + session.context(Stage::Authentication, session.budget(Stage::Authentication)); context.request = crate::adapter::request_context(request.request_context()); let _scope = callbacks .prepare_frontend_authentication(&context, request.shared_request_context()) diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/management/callback.rs b/backend/crates/gateway-plugin/runtime/src/adapter/management/callback.rs index 06141929d..f880b4f3b 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/management/callback.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/management/callback.rs @@ -165,9 +165,10 @@ impl ManagementEntry { { return Err(invalid_state()); } - let mut context = self - .session - .context(Stage::PublicManagement, limits.maximum_call_timeout); + let mut context = self.session.context( + Stage::PublicManagement, + self.session.budget(Stage::PublicManagement), + ); context.request = crate::adapter::request_context(request.request_context.as_deref()); context.request_id = Some(request.request_id); let _scope = self diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/management/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/management/mod.rs index 14df27f5a..b33f1e505 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/management/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/management/mod.rs @@ -98,7 +98,7 @@ impl ManagementEntry { .map_err(|_| AdminError::invalid("插件管理请求无法编码"))?; let mut context = self .session - .context(Stage::Management, limits.maximum_call_timeout); + .context(Stage::Management, self.session.budget(Stage::Management)); context.request = crate::adapter::request_context(request.request_context.as_deref()); context.request_id = Some(request.request_id); let _scope = self diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/management/registration.rs b/backend/crates/gateway-plugin/runtime/src/adapter/management/registration.rs index 69dbbebd8..e177f3a2f 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/management/registration.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/management/registration.rs @@ -1,6 +1,6 @@ //! 将插件管理声明编译为页面、路由与资源入口 -use std::{collections::BTreeMap, sync::Arc, time::Duration}; +use std::{collections::BTreeMap, sync::Arc}; use gateway_admin::model::{ AdminError, @@ -34,7 +34,7 @@ pub(crate) async fn prepare( let reply = session .call( "management.register", - session.context(Stage::Registration, Duration::from_secs(5)), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), Vec::new(), ) diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/execution.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/execution.rs index 38777c730..b1060debd 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/execution.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/execution.rs @@ -7,7 +7,7 @@ use gateway_core::engine::middleware::{ MiddlewareRequest, MiddlewareResponse, }; use gateway_core::middleware::{Middleware, Next}; -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; // 协议适配只提供选择条件与单层调用；顺序、上下文捕获和续体拼接共用这一个入口 type Invoke = fn( @@ -15,7 +15,6 @@ type Invoke = fn( Context, Input, Next, - Duration, ) -> BoxFuture<'static, Result>; impl PluginMiddlewarePlan { @@ -42,8 +41,7 @@ impl PluginMiddlewarePlan { .map(|entry| -> Middleware { let entry = Arc::clone(entry); let context = context.clone(); - let timeout = self.middleware_timeout; - Box::new(move |input, next| invoke(entry, context, input, next, timeout)) + Box::new(move |input, next| invoke(entry, context, input, next)) }) .collect(); gateway_core::middleware::compose(layers, move |input| next.run(input)) @@ -70,10 +68,8 @@ impl MiddlewarePlan for PluginMiddlewarePlan { MiddlewareMount::Outbound, context, outbound_matches, - |entry, context, request, next, timeout| { - Box::pin(super::outbound::invoke( - entry, context, request, next, timeout, - )) + |entry, context, request, next| { + Box::pin(super::outbound::invoke(entry, context, request, next)) }, next, ) @@ -99,8 +95,8 @@ impl MiddlewarePlan for PluginMiddlewarePlan { MiddlewareMount::Service, context, |entry, context| !context.extensions.contains(&entry.instance_id), - |entry, context, input, next, timeout| { - Box::pin(super::service::invoke(entry, context, input, next, timeout)) + |entry, context, input, next| { + Box::pin(super::service::invoke(entry, context, input, next)) }, next, ) @@ -123,10 +119,8 @@ impl MiddlewarePlan for PluginMiddlewarePlan { MiddlewareMount::WebSocket, context, |_entry, _context| true, - |entry, context, message, next, timeout| { - Box::pin(super::websocket::invoke( - entry, context, message, next, timeout, - )) + |entry, context, message, next| { + Box::pin(super::websocket::invoke(entry, context, message, next)) }, next, ) @@ -147,8 +141,8 @@ impl MiddlewarePlan for PluginMiddlewarePlan { MiddlewareMount::Http, context, |entry, context| !context.extensions.contains(&entry.instance_id), - |entry, context, request, next, timeout| { - Box::pin(super::http::invoke(entry, context, request, next, timeout)) + |entry, context, request, next| { + Box::pin(super::http::invoke(entry, context, request, next)) }, next, ) @@ -173,10 +167,8 @@ impl MiddlewarePlan for PluginMiddlewarePlan { context.model(), ) }, - |entry, context, request, next, timeout| { - Box::pin(super::request::invoke( - entry, context, request, next, timeout, - )) + |entry, context, request, next| { + Box::pin(super::request::invoke(entry, context, request, next)) }, next, ) 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 e645e9c4b..80802116a 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/http.rs @@ -20,7 +20,6 @@ pub(super) async fn invoke( context: core::Context, request: core::Request, next: core::Next, - timeout: std::time::Duration, ) -> Result { let Some(ports) = &entry.invocation else { return if entry.failure_policy == PluginFailurePolicy::Delegate { @@ -35,7 +34,9 @@ pub(super) async fn invoke( .and_then(|settings| settings.runtime.as_ref()) .map_or(serde_json::Value::Null, |settings| settings.inspect()); let (invocation, request) = Invocation::new(request, next, entry.instance_id.clone())?; - let mut call = ports.session.context(Stage::Http, timeout); + let mut call = ports + .session + .context(Stage::Http, ports.session.budget(Stage::Http)); call.request = crate::adapter::request_context(context.extensions.request()); call.request_id = Some(context.request_id.clone()); call.resource_stream = true; diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/mod.rs index dc73bda18..3046f923a 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/mod.rs @@ -18,7 +18,6 @@ use std::{ collections::{BTreeMap, BTreeSet}, fmt, sync::Arc, - time::Duration, }; #[derive(Clone)] @@ -168,7 +167,6 @@ fn compile_entry( pub(crate) struct PluginMiddlewarePlan { middleware: BTreeMap]>>, - middleware_timeout: Duration, } impl fmt::Debug for PluginMiddlewarePlan { @@ -181,10 +179,7 @@ impl fmt::Debug for PluginMiddlewarePlan { } impl PluginMiddlewarePlan { - pub(crate) fn compile( - mut entries: Vec, - maximum_call_timeout: Duration, - ) -> Option> { + pub(crate) fn compile(mut entries: Vec) -> Option> { if entries.is_empty() { return None; } @@ -208,7 +203,6 @@ impl PluginMiddlewarePlan { .into_iter() .map(|(mount, entries)| (mount, entries.into_boxed_slice())) .collect(), - middleware_timeout: maximum_call_timeout, })) } } diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/outbound.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/outbound.rs index 089d6e452..d0fb5f0a5 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/outbound.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/outbound.rs @@ -1,6 +1,6 @@ //! 在 Provider 发送前组合完整出站请求，并沿用模型数据面的调用资源 -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; use gateway_admin::model::plugins::instances::PluginFailurePolicy; use gateway_core::engine::middleware::{MiddlewareContext, MiddlewareError, outbound as core}; @@ -17,7 +17,6 @@ pub(super) async fn invoke( context: MiddlewareContext, request: core::Request, next: core::Next, - timeout: Duration, ) -> Result { let Some(ports) = &entry.invocation else { return if entry.failure_policy == PluginFailurePolicy::Delegate { @@ -54,7 +53,12 @@ pub(super) async fn invoke( .get(), }, }; - let mut call = ports.session.context(Stage::Outbound, timeout); + let mut call = ports.session.context( + Stage::Outbound, + context + .deadline() + .bounded(ports.session.budget(Stage::Outbound)), + ); call.request = crate::adapter::request_context(context.extension_scope().request()); call.request_id = Some(input.metadata.request_id.clone()); call.attempt_id = Some(format!( diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/request.rs b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/request.rs index 82cc532ad..51ded81d0 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/request.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/request.rs @@ -1,9 +1,6 @@ //! 请求与尝试中间件的 RPC 调用、失败回退和响应正文生命周期 -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; +use std::sync::{Arc, Mutex}; use bytes::Bytes; use futures::future::BoxFuture; @@ -37,7 +34,6 @@ pub(super) async fn invoke( context: MiddlewareContext, request: MiddlewareRequest, next: MiddlewareNext, - maximum_timeout: Duration, ) -> Result { let Some(invocation_ports) = &entry.invocation else { return if entry.failure_policy @@ -48,10 +44,7 @@ pub(super) async fn invoke( Err(MiddlewareError::Fault) }; }; - if maximum_timeout.is_zero() - || context.deadline().is_elapsed() - || context.cancellation().is_cancelled() - { + if context.deadline().is_elapsed() || context.cancellation().is_cancelled() { return Err(MiddlewareError::Fault); } let invocation = MiddlewareInvocation::new(&context, request, next, entry.instance_id.clone()); @@ -59,18 +52,21 @@ pub(super) async fn invoke( Ok(projection) => projection, Err(_) => return recover_rpc(&entry, &invocation, RpcError::Context(None)).await, }; + let stage = match entry.mount { + MiddlewareMount::Http + | MiddlewareMount::WebSocket + | MiddlewareMount::Service + | MiddlewareMount::Outbound => { + return Err(MiddlewareError::InvalidState); + } + MiddlewareMount::Request => Stage::Request, + MiddlewareMount::Attempt => Stage::Attempt, + }; let mut call_context = invocation_ports.session.context( - match entry.mount { - MiddlewareMount::Http - | MiddlewareMount::WebSocket - | MiddlewareMount::Service - | MiddlewareMount::Outbound => { - return Err(MiddlewareError::InvalidState); - } - MiddlewareMount::Request => Stage::Request, - MiddlewareMount::Attempt => Stage::Attempt, - }, - maximum_timeout, + stage, + context + .deadline() + .bounded(invocation_ports.session.budget(stage)), ); call_context.resource_stream = true; call_context.request = crate::adapter::request_context(context.extension_scope().request()); 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 9a74b1a46..5543fb890 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/service.rs @@ -16,7 +16,6 @@ pub(super) async fn invoke( context: core::Context, input: core::Value, next: core::Next, - timeout: std::time::Duration, ) -> Result { let Some(ports) = &entry.invocation else { return if entry.failure_policy == PluginFailurePolicy::Delegate { @@ -26,7 +25,9 @@ pub(super) async fn invoke( }; }; let invocation = Invocation::new(next); - let mut call = ports.session.context(Stage::Service, timeout); + let mut call = ports + .session + .context(Stage::Service, ports.session.budget(Stage::Service)); call.request = crate::adapter::request_context(context.extensions.request()); call.request_id = Some(context.request_id.clone()); let _binding = ports 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 1a66d00f8..129d85ccf 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/middleware/websocket.rs @@ -16,7 +16,6 @@ pub(super) async fn invoke( context: core::Context, message: core::Message, next: core::Next, - timeout: std::time::Duration, ) -> Result, MiddlewareError> { let Some(ports) = &entry.invocation else { return if entry.failure_policy == PluginFailurePolicy::Delegate { @@ -44,7 +43,9 @@ pub(super) async fn invoke( .encode(message) .map_err(|_| MiddlewareError::Fault)?, }; - let mut call = ports.session.context(Stage::WebSocket, timeout); + let mut call = ports + .session + .context(Stage::WebSocket, ports.session.budget(Stage::WebSocket)); call.request = crate::adapter::request_context(context.extensions.request()); call.request_id = Some(context.connection_id.clone()); let _binding = ports diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs index 8de3437fa..64da3f571 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/observer.rs @@ -3,7 +3,7 @@ use std::{ collections::BTreeMap, sync::{Arc, Mutex}, - time::{Duration, SystemTime, UNIX_EPOCH}, + time::{SystemTime, UNIX_EPOCH}, }; use gateway_admin::model::{ @@ -31,7 +31,6 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use crate::{RpcSession, adapter::scope::BindingScope}; -const OBSERVATION_TIMEOUT: Duration = Duration::from_secs(2); const OBSERVATION_ENVELOPE_RESERVE: usize = 16 * 1024; #[derive(Clone)] @@ -57,10 +56,11 @@ impl ObserverEntry { event: Event, payload: Vec, extensions: gateway_core::engine::extensions::ExtensionCallScope, - timeout: Duration, ) { let request_id = event.request_id(); - let mut context = self.session.context(Stage::Observation, timeout); + let mut context = self + .session + .context(Stage::Observation, self.session.budget(Stage::Observation)); context.request = crate::adapter::request_context(extensions.request()); context.request_id = Some(request_id.to_owned()); let Ok(_scope) = self.callbacks.prepare_extensions(&context, extensions) else { @@ -166,7 +166,6 @@ pub(crate) struct PluginObserverPlan { websocket_queues: Arc>>>, websocket_queue_capacity: usize, maximum_websocket_payload_bytes: usize, - timeout: Duration, } struct QueuedWebSocketObservation { @@ -180,7 +179,6 @@ impl PluginObserverPlan { pub(crate) fn compile( mut entries: Vec, maximum_calls: usize, - maximum_call_timeout: Duration, maximum_buffered_body_bytes: usize, ) -> Option> { if entries.is_empty() { @@ -205,7 +203,6 @@ impl PluginObserverPlan { websocket_queues: Arc::new(Mutex::new(BTreeMap::new())), websocket_queue_capacity: maximum_calls.clamp(1, 256), maximum_websocket_payload_bytes, - timeout: maximum_call_timeout.min(OBSERVATION_TIMEOUT), })) } @@ -326,7 +323,6 @@ impl RequestObserverPlan for PluginObserverPlan { return; }; let entries = Arc::clone(&self.entries); - let timeout = self.timeout; handle.spawn(async move { let _generation = generation; let _slot = slot; @@ -342,7 +338,6 @@ impl RequestObserverPlan for PluginObserverPlan { Event::RequestCompleted(Box::new(wire_observation(&observation))), Vec::new(), observation.extension_scope().clone(), - timeout, ) .await; } @@ -372,10 +367,8 @@ impl PluginObserverPlan { let (sender, receiver) = mpsc::channel(self.websocket_queue_capacity); queues.insert(request_id.to_owned(), sender.clone()); let entries = Arc::clone(&self.entries); - let timeout = self.timeout; handle.spawn(dispatch_websocket_queue( entries, - timeout, generation, worker_slot, receiver, @@ -386,7 +379,6 @@ impl PluginObserverPlan { async fn dispatch_websocket_queue( entries: Arc<[ObserverEntry]>, - timeout: Duration, generation: ExtensionSetReference, worker_slot: OwnedSemaphorePermit, mut receiver: mpsc::Receiver, @@ -422,7 +414,6 @@ async fn dispatch_websocket_queue( Event::WebSocketResponse(Box::new(request)), queued.payload.to_vec(), queued.observation.extension_scope().clone(), - timeout, ) .await; } diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs index c106ff77f..66e3ecd77 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/mod.rs @@ -3,7 +3,7 @@ mod retry; mod route_schedule; -use std::{fmt, sync::Arc, time::Duration}; +use std::{fmt, sync::Arc}; use gateway_admin::model::{ AdminError, @@ -13,8 +13,6 @@ use gateway_plugin_sdk::{Capability, Manifest, Stage}; use crate::{RpcSession, adapter::scope::BindingScope, callback::PluginCallbacks}; -const POLICY_TIMEOUT: Duration = Duration::from_secs(2); - #[derive(Clone)] pub(crate) enum PolicyEntry { Router(ModelRouterEntry), @@ -219,7 +217,6 @@ pub(crate) struct PluginRequestPolicyPlan { routers: Arc<[ModelRouterEntry]>, schedulers: Arc<[AccountSchedulerEntry]>, retries: Arc<[RetryEntry]>, - policy_timeout: Duration, } impl fmt::Debug for PluginRequestPolicyPlan { @@ -234,10 +231,7 @@ impl fmt::Debug for PluginRequestPolicyPlan { } impl PluginRequestPolicyPlan { - pub(crate) fn compile( - entries: Vec, - maximum_call_timeout: Duration, - ) -> Result>, AdminError> { + pub(crate) fn compile(entries: Vec) -> Result>, AdminError> { let mut routers = Vec::new(); let mut schedulers = Vec::new(); let mut retries = Vec::new(); @@ -277,7 +271,6 @@ impl PluginRequestPolicyPlan { routers: routers.into(), schedulers: schedulers.into(), retries: retries.into(), - policy_timeout: maximum_call_timeout.min(POLICY_TIMEOUT), }))) } } diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/retry.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/retry.rs index e5117583f..aaa707dcb 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/retry.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/retry.rs @@ -17,11 +17,10 @@ use super::RetryEntry; pub(super) fn decide( entries: Arc<[RetryEntry]>, input: RetryInput, - timeout: Duration, ) -> BoxFuture<'static, Result> { Box::pin(async move { - let started = std::time::Instant::now(); - for entry in entries.iter().filter(|entry| { + let started = tokio::time::Instant::now(); + let matched = entries.iter().filter(|entry| { !input.extension_scope.contains(&entry.instance_id) && entry.scope.matches_processing( &input.client_key_id, @@ -29,23 +28,32 @@ pub(super) fn decide( Some(&input.facts.provider), input.facts.model.as_deref(), ) - }) { + }); + let timeout = matched + .clone() + .filter_map(|entry| entry.invocation.as_ref()) + .map(|invocation| invocation.session.budget(Stage::Retry)) + .max() + .unwrap_or_default() + .min(input.facts.remaining_deadline); + for entry in matched { // 整条委托链共享预算，不能每个插件重新取得完整等待时间 - let remaining = timeout - .min(input.facts.remaining_deadline) - .saturating_sub(started.elapsed()); + let remaining = timeout.saturating_sub(started.elapsed()); if remaining.is_zero() { break; } - match call(entry, &input, remaining).await { - Ok(RetryDecision::Delegate) => {} - Ok(decision) => return Ok(decision), - Err(()) => tracing::warn!( - plugin_id = entry.plugin_id, - instance_id = entry.instance_id, - request_id = input.request_id.as_str(), - "插件重试策略无效或调用失败，委托后续策略" - ), + match tokio::time::timeout(remaining, call(entry, &input, remaining)).await { + Err(_) => break, + Ok(result) => match result { + Ok(RetryDecision::Delegate) => {} + Ok(decision) => return Ok(decision), + Err(()) => tracing::warn!( + plugin_id = entry.plugin_id, + instance_id = entry.instance_id, + request_id = input.request_id.as_str(), + "插件重试策略无效或调用失败，委托后续策略" + ), + }, } } Ok(RetryDecision::Delegate) @@ -79,7 +87,10 @@ async fn call( .unwrap_or(u64::MAX), allowed_actions, }; - let mut context = invocation.session.context(Stage::Retry, timeout); + let mut context = invocation.session.context( + Stage::Retry, + timeout.min(invocation.session.budget(Stage::Retry)), + ); context.request = crate::adapter::request_context(input.extension_scope.request()); context.request_id = Some(request.request_id.clone()); context.attempt_id = Some(format!("{}:{}", request.request_id, request.attempt_index)); diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/policy/route_schedule.rs b/backend/crates/gateway-plugin/runtime/src/adapter/policy/route_schedule.rs index 8bcd264ed..865f94a21 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/policy/route_schedule.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/policy/route_schedule.rs @@ -1,6 +1,6 @@ //! 插件模型路由与账号调度调用，以及请求事实的线协议投影 -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::{SystemTime, UNIX_EPOCH}; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use futures::future::BoxFuture; @@ -38,7 +38,7 @@ impl RequestPolicyPlan for PluginRequestPolicyPlan { input: gateway_core::engine::policy::RetryInput, ) -> BoxFuture<'static, Result> { - super::retry::decide(self.retries.clone(), input, self.policy_timeout) + super::retry::decide(self.retries.clone(), input) } fn route_model( @@ -46,12 +46,11 @@ impl RequestPolicyPlan for PluginRequestPolicyPlan { input: ModelRouteInput, ) -> BoxFuture<'static, Result> { let routers = std::sync::Arc::clone(&self.routers); - let timeout = self.policy_timeout; Box::pin(async move { for entry in routers.iter().filter(|entry| { !input.suppresses_plugin(&entry.instance_id) && entry.scope.matches_route(&input) }) { - let result = route_with_entry(entry, &input, timeout).await; + let result = route_with_entry(entry, &input).await; match result { Ok(ModelRouteDecision::Unhandled) => continue, Ok(decision) => return Ok(decision), @@ -78,14 +77,13 @@ impl RequestPolicyPlan for PluginRequestPolicyPlan { input: AccountScheduleInput, ) -> BoxFuture<'static, Result> { let schedulers = std::sync::Arc::clone(&self.schedulers); - let timeout = self.policy_timeout; Box::pin(async move { let Some(entry) = schedulers.iter().find(|entry| { !input.suppresses_plugin(&entry.instance_id) && entry.scope.matches_schedule(&input) }) else { return Ok(AccountScheduleDecision::Delegate); }; - match schedule_with_entry(entry, &input, timeout).await { + match schedule_with_entry(entry, &input).await { Ok(decision) => Ok(decision), Err(()) if entry.failure_policy @@ -108,7 +106,6 @@ impl RequestPolicyPlan for PluginRequestPolicyPlan { async fn route_with_entry( entry: &ModelRouterEntry, input: &ModelRouteInput, - timeout: Duration, ) -> Result { let invocation = entry.invocation.as_ref().ok_or(())?; let projection = operation_projection(input.operation())?; @@ -126,7 +123,9 @@ async fn route_with_entry( .collect(), headers, }; - let mut context = invocation.session.context(Stage::Routing, timeout); + let mut context = invocation + .session + .context(Stage::Routing, invocation.session.budget(Stage::Routing)); context.request = crate::adapter::request_context(input.extension_scope().request()); context.request_id = Some(input.request_id().as_str().to_owned()); let _network_scope = invocation @@ -172,7 +171,6 @@ async fn route_with_entry( async fn schedule_with_entry( entry: &AccountSchedulerEntry, input: &AccountScheduleInput, - timeout: Duration, ) -> Result { let invocation = entry.invocation.as_ref().ok_or(())?; let candidates = input.candidates(); @@ -197,7 +195,10 @@ async fn schedule_with_entry( }) .collect(), }; - let mut context = invocation.session.context(Stage::Scheduling, timeout); + let mut context = invocation.session.context( + Stage::Scheduling, + invocation.session.budget(Stage::Scheduling), + ); context.request = crate::adapter::request_context(input.extension_scope().request()); context.request_id = Some(input.request_id().as_str().to_owned()); context.attempt_id = Some(format!( diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs index 9fc2a2b30..ee36e65f9 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/execution.rs @@ -167,7 +167,7 @@ impl Execution { let timeout = invocation .context .deadline() - .bounded(self.adapter.session.maximum_call_timeout()); + .bounded(self.adapter.session.budget(Stage::Upstream)); if timeout.is_zero() || invocation.context.cancellation().is_cancelled() { return Err(ProviderError::new( ProviderErrorKind::Cancelled, diff --git a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/registration.rs b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/registration.rs index 2c6c00fff..0e93b0141 100644 --- a/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/registration.rs +++ b/backend/crates/gateway-plugin/runtime/src/adapter/upstream_adapter/registration.rs @@ -1,6 +1,6 @@ //! 校验并准备插件上游适配声明与能力绑定 -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; use gateway_admin::model::{ AdminError, @@ -56,7 +56,7 @@ pub(crate) async fn prepare( let reply = session .call( REGISTER_METHOD, - session.context(Stage::Registration, Duration::from_secs(5)), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), vec![], ) diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs index 1d3b8c3ed..aaa0d0653 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/instance.rs @@ -16,18 +16,22 @@ use gateway_admin::{ use gateway_plugin_sdk::{Handshake, Stage, call::registration::Registration}; use std::{sync::Arc, time::Duration}; +const PROCESS_PREPARATION_TIMEOUT: Duration = Duration::from_secs(30); + +struct StartedInstance { + instance: PluginInstance, + package: Arc, + callbacks: Arc, + prepared: PreparedInstance, +} + impl PluginRuntime { - pub(super) async fn prepare_instance( + async fn start_instance( &self, instance: PluginInstance, target_revision: Revision, request_body_limit: Option, - ) -> Result, AdminError> { - let mut observer_entries = Vec::new(); - let mut commands = Vec::new(); - let mut management = Vec::new(); - let mut policy_entries = Vec::new(); - let mut authentication_entries = Vec::new(); + ) -> Result { let authorization_binding = instance_fingerprint(&instance)?; if self.restart_circuit_is_open(&RestartIdentity::new( instance.id.clone(), @@ -109,7 +113,6 @@ impl PluginRuntime { crate::compatibility::report(&manifest, &instance_id)?; let restart_identity = RestartIdentity::new(instance_id.clone(), authorization_binding.clone()); - let bindings = instance.bindings.clone(); let owner = self .state .load_owner( @@ -181,23 +184,69 @@ impl PluginRuntime { lifecycle, ) .await - .map_err(|_| AdminError::unavailable("插件进程握手失败"))?, + .map_err(|error| match error { + crate::RpcError::Budget(gateway_plugin_sdk::CallBudgetError::Invalid { + stage, + milliseconds, + }) => AdminError::invalid(format!( + "插件阶段 {stage:?} 的调用预算无效：{milliseconds} 毫秒" + )), + _ => AdminError::unavailable("插件进程握手失败"), + })?, ); - let mut prepared = PreparedInstance { + let prepared = PreparedInstance { instance_id: instance_id.clone(), artifact_sha256: instance.artifact_sha256.clone(), revision: instance.revision, - session: session.clone(), + session, private_state: private_state.clone(), maintenance: manifest .contributes .contains_key(&gateway_plugin_sdk::Capability::Maintenance), contributions: PreparedContributions::default(), }; + Ok(StartedInstance { + instance, + package, + callbacks, + prepared, + }) + } + + pub(super) async fn prepare_instance( + &self, + instance: PluginInstance, + target_revision: Revision, + request_body_limit: Option, + ) -> Result, AdminError> { + // 启动控制期限只包住制品准备与握手，Ready 后每次注册使用协商预算 + let StartedInstance { + instance, + package, + callbacks, + mut prepared, + } = tokio::time::timeout( + PROCESS_PREPARATION_TIMEOUT, + self.start_instance(instance, target_revision, request_body_limit), + ) + .await + .map_err(|_| AdminError::unavailable("插件启动超时"))??; + let session = &prepared.session; + let manifest = package.package().manifest(); + let plugin_id = manifest + .plugin_id() + .map_err(|_| AdminError::invalid("插件身份无效"))?; + let instance_id = &instance.id; + let bindings = &instance.bindings; + let mut observer_entries = Vec::new(); + let mut commands = Vec::new(); + let mut management = Vec::new(); + let mut policy_entries = Vec::new(); + let mut authentication_entries = Vec::new(); let registration = session .call( "plugin.register", - session.context(Stage::Registration, Duration::from_secs(5)), + session.context(Stage::Registration, session.budget(Stage::Registration)), serde_json::json!({}), vec![], ) @@ -211,7 +260,7 @@ impl PluginRuntime { commands.extend( crate::adapter::command_line::prepare( &instance, - &manifest, + manifest, session.clone(), callbacks.clone(), ) @@ -228,8 +277,8 @@ impl PluginRuntime { management.push(entry); } if let Some(entry) = crate::adapter::frontend_authentication::prepare_entry( - &manifest, - &bindings, + manifest, + bindings, session.clone(), callbacks.clone(), ) @@ -238,35 +287,34 @@ impl PluginRuntime { authentication_entries.push(entry); } if let Some(entry) = crate::adapter::observer::compile_entry( - &manifest, + manifest, &plugin_id, - &instance_id, - &bindings, + instance_id, + bindings, session.clone(), callbacks.clone(), )? { observer_entries.push(entry); } policy_entries.extend(crate::adapter::policy::compile_entries( - &manifest, - &instance_id, - &bindings, + manifest, + instance_id, + bindings, session.clone(), callbacks.clone(), )?); let middleware_entries = crate::adapter::middleware::compile_entries( - &manifest, - &instance_id, - &bindings, + manifest, + instance_id, + bindings, session.clone(), callbacks.clone(), )?; - let model_aliases = - crate::adapter::catalog::prepare(&manifest, &instance, &session).await?; + let model_aliases = crate::adapter::catalog::prepare(manifest, &instance, session).await?; let upstream_entries = crate::adapter::upstream_adapter::prepare( - &manifest, + manifest, &instance, - Arc::clone(&session), + Arc::clone(session), Arc::clone(&callbacks), ) .await?; diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs index 514c9a12e..5895092cf 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/maintenance.rs @@ -1,8 +1,11 @@ //! 对账只从已发布集合开始；通知合并，实例串行，失败和停用不阻塞其他实例 -use super::{PluginRuntime, set::PreparedSet}; +use super::{ + PluginRuntime, + set::{PreparedInstance, PreparedSet}, +}; use crate::RpcSession; -use futures::{StreamExt as _, future::BoxFuture}; +use futures::{StreamExt as _, future::BoxFuture, stream::FuturesUnordered}; use gateway_core::{ lifecycle::CancellationToken, runtime::RuntimeSnapshotHandle, @@ -13,12 +16,51 @@ use gateway_core::{ }, }; use gateway_plugin_sdk::Stage; -use std::{sync::Arc, time::Duration}; -use tokio::{sync::Notify, task::JoinSet}; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; +use tokio::{sync::Notify, task::JoinHandle}; const RECONCILE_INTERVAL: Duration = Duration::from_secs(30); const RETRY_DELAY: Duration = Duration::from_secs(5); -const CALL_TIMEOUT: Duration = Duration::from_secs(30); + +struct InstanceTask { + instance: Arc, + dirty: Arc, + cancellation: CancellationToken, + task: JoinHandle<()>, +} + +impl InstanceTask { + fn start(instance: Arc) -> Self { + let dirty = Arc::new(Notify::new()); + dirty.notify_one(); + let cancellation = CancellationToken::new(); + let task = tokio::spawn(reconcile( + Arc::clone(&instance.session), + Arc::clone(&dirty), + cancellation.clone(), + )); + Self { + instance, + dirty, + cancellation, + task, + } + } + + async fn stop(mut self) { + self.cancellation.cancel(); + let _ = (&mut self.task).await; + // 宿主资源在取消时同步撤销，替换任务还须等对端确认或进程退出 + self.instance.session.settle_cancellations().await; + } +} + +impl Drop for InstanceTask { + fn drop(&mut self) { + self.cancellation.cancel(); + self.task.abort(); + } +} struct MaintenanceTask { runtime: Arc, @@ -59,18 +101,12 @@ impl DaemonTask for MaintenanceTask { let mut publications = self.snapshots.subscribe_publications(); let mut active: Option> = None; let mut revision = None; - let mut stop = CancellationToken::new(); - let mut tasks = JoinSet::new(); - let mut notifications = Vec::new(); + let mut tasks = BTreeMap::::new(); loop { let snapshot = self.snapshots.acquire().ok(); let reference = snapshot.as_ref().and_then(|s| s.extensions()); let same_set = active.as_ref().map(|s| &s.id) == reference.map(|s| s.id()); if !same_set { - stop.cancel(); - while tasks.join_next().await.is_some() {} - notifications.clear(); - stop = CancellationToken::new(); active = match reference { Some(reference) => { Some(self.runtime.prepared_set(reference).await.map_err(|_| { @@ -79,17 +115,37 @@ impl DaemonTask for MaintenanceTask { } None => None, }; - if let Some(set) = &active { - for instance in set.sessions.iter().filter(|i| i.maintenance) { - let dirty = Arc::new(Notify::new()); - dirty.notify_one(); - notifications.push(dirty.clone()); - tasks.spawn(reconcile(instance.session.clone(), dirty, stop.clone())); - } + let desired = active + .iter() + .flat_map(|set| &set.sessions) + .filter(|instance| instance.maintenance) + .map(|instance| (instance.instance_id.clone(), Arc::clone(instance))) + .collect::>(); + let retired = tasks + .extract_if(.., |id, task| { + // 发布层只复用完全相同且就绪的实例，Arc 身份也包含会话 incarnation + !desired + .get(id) + .is_some_and(|instance| Arc::ptr_eq(instance, &task.instance)) + }) + .map(|(_, task)| task) + .collect::>(); + for task in &retired { + task.cancellation.cancel(); } - } else if revision != snapshot.as_ref().map(|s| s.revision()) { - for dirty in ¬ifications { - dirty.notify_one(); + futures::future::join_all(retired.into_iter().map(InstanceTask::stop)).await; + for task in tasks.values() { + task.dirty.notify_one(); + } + for (id, instance) in desired { + tasks + .entry(id) + .or_insert_with(|| InstanceTask::start(instance)); + } + } + if same_set && revision != snapshot.as_ref().map(|s| s.revision()) { + for task in tasks.values() { + task.dirty.notify_one(); } } revision = snapshot.as_ref().map(|s| s.revision()); @@ -101,22 +157,39 @@ impl DaemonTask for MaintenanceTask { biased; () = cancellation.cancelled() => break, changed = publications.next() => if changed.is_none() { break; }, - result = tasks.join_next(), if !tasks.is_empty() => { - if result.is_some() { - stop.cancel(); - while tasks.join_next().await.is_some() {} + result = next_finished(&mut tasks), if !tasks.is_empty() => { + if let Some(instance_id) = result { + tasks.remove(&instance_id); + stop_all(tasks).await; return Err(WorkerTaskError::safe("plugin maintenance task stopped unexpectedly")); } } } } - stop.cancel(); - while tasks.join_next().await.is_some() {} + stop_all(tasks).await; Ok(()) }) } } +async fn next_finished(tasks: &mut BTreeMap) -> Option { + let mut pending = tasks + .iter_mut() + .map(|(id, task)| async move { + let _ = (&mut task.task).await; + id.clone() + }) + .collect::>(); + pending.next().await +} + +async fn stop_all(tasks: BTreeMap) { + for task in tasks.values() { + task.cancellation.cancel(); + } + futures::future::join_all(tasks.into_values().map(InstanceTask::stop)).await; +} + async fn reconcile(session: Arc, dirty: Arc, cancellation: CancellationToken) { loop { tokio::select! { @@ -125,21 +198,38 @@ async fn reconcile(session: Arc, dirty: Arc, cancellation: C () = dirty.notified() => {}, () = tokio::time::sleep(RECONCILE_INTERVAL) => {}, } - let context = session.context(Stage::Maintenance, CALL_TIMEOUT); + let context = session.context(Stage::Maintenance, session.budget(Stage::Maintenance)); let instance_id = context.instance_id.clone(); + let resource_scope_id = context.resource_scope_id.clone(); + let budget_ms = context.timeout_ms; + let started = tokio::time::Instant::now(); let result = tokio::select! { biased; - () = cancellation.cancelled() => return, + () = cancellation.cancelled() => { + tracing::info!(target: "plugin.maintenance", %instance_id, %resource_scope_id, budget_ms, + budget_source = ?session.budget_source(Stage::Maintenance), + elapsed_ms = started.elapsed().as_millis() as u64, + reason = "cancelled", "plugin maintenance ended"); + return; + }, result = session.call("plugin.reconcile", context, serde_json::json!({}), Vec::new()) => result, }; - if result.is_err() { - tracing::warn!(%instance_id, "plugin maintenance failed; retrying after bounded delay"); + if let Err(error) = result { + tracing::warn!(target: "plugin.maintenance", %instance_id, %resource_scope_id, budget_ms, + budget_source = ?session.budget_source(Stage::Maintenance), + elapsed_ms = started.elapsed().as_millis() as u64, + reason = error.diagnostic().0, "plugin maintenance failed; retrying after bounded delay"); tokio::select! { biased; () = cancellation.cancelled() => return, () = tokio::time::sleep(RETRY_DELAY) => {}, } dirty.notify_one(); + } else { + tracing::info!(target: "plugin.maintenance", %instance_id, %resource_scope_id, budget_ms, + budget_source = ?session.budget_source(Stage::Maintenance), + elapsed_ms = started.elapsed().as_millis() as u64, + reason = "completed", "plugin maintenance ended"); } } } diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs index 3ded63376..b4e52b9bb 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/runtime.rs @@ -27,15 +27,13 @@ use gateway_core::{ observation::RequestObserverExtensionIndex, policy::RequestPolicyExtensionIndex, }, + lifecycle::CancellationToken, provider_ports::OAuthPendingFlowPort, }; use std::{ collections::BTreeMap, path::PathBuf, - sync::{ - Arc, Mutex as SyncMutex, Weak, - atomic::{AtomicBool, Ordering}, - }, + sync::{Arc, Mutex as SyncMutex, Weak}, time::Duration, }; use tokio::sync::{Mutex, Semaphore}; @@ -62,7 +60,7 @@ pub struct PluginRuntime { pub(super) prepared_instances: Mutex>>, pub(super) preparation_diagnostics: SyncMutex>, pub(super) restart_circuits: Arc, - pub(super) shutting_down: Arc, + pub(super) shutting_down: CancellationToken, shutdown_lock: Mutex<()>, pub(super) processes: Arc, pub(super) validators: Arc, @@ -105,7 +103,7 @@ impl PluginRuntime { prepared_instances: Mutex::new(BTreeMap::new()), preparation_diagnostics: SyncMutex::new(BTreeMap::new()), restart_circuits: RestartCircuits::new(restart_circuit), - shutting_down: Arc::new(AtomicBool::new(false)), + shutting_down: CancellationToken::new(), shutdown_lock: Mutex::new(()), processes, validators: Arc::new(Semaphore::new(MAXIMUM_PARALLEL_PREPARATIONS)), @@ -182,7 +180,7 @@ impl PluginRuntime { /// 调用方必须先停止接收 HTTP 与 Worker；这里拒绝新候选并等待全部插件 I/O 退出 pub async fn shutdown(&self) { - self.shutting_down.store(true, Ordering::Release); + self.shutting_down.cancel(); let _shutdown = self.shutdown_lock.lock().await; let _prepare = self.prepare_lock.lock().await; let sets = { diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs index 87abcce47..383df994b 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/set.rs @@ -10,18 +10,12 @@ use gateway_core::{ extensions::ExecutionExtensionPlans, observation::RequestObserverPlan, policy::RequestPolicyPlan, }, + lifecycle::CancellationToken, routing::extensions::{ExtensionSetId, ExtensionSetLease}, }; use secrecy::ExposeSecret as _; use sha2::{Digest as _, Sha256}; -use std::{ - collections::BTreeMap, - sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - }, - time::Duration, -}; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; pub(super) struct PreparedSet { pub(super) id: ExtensionSetId, @@ -32,7 +26,7 @@ pub(super) struct PreparedSet { Option>, pub(super) sessions: Vec>, pub(super) failures: BTreeMap, - pub(super) shutting_down: Arc, + pub(super) shutting_down: CancellationToken, pub(super) commands: Vec>, pub(super) management: Vec, pub(super) model_aliases: Vec, @@ -94,7 +88,7 @@ impl ExtensionSetLease for PreparedSet { fn can_serve(&self) -> bool { // 发布的是可用能力和故障绑定组成的完整计划；单个进程退出不能使原生转发失效 - !self.shutting_down.load(Ordering::Acquire) + !self.shutting_down.is_cancelled() } } diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs index 9391e2c0c..d1fb32285 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/snapshot.rs @@ -26,8 +26,7 @@ use gateway_core::{ }; use std::{ collections::{BTreeMap, BTreeSet}, - sync::{Arc, Weak, atomic::Ordering}, - time::Duration, + sync::{Arc, Weak}, }; impl PluginRuntime { @@ -39,9 +38,11 @@ impl PluginRuntime { ) -> Result { let target_revision = snapshot.config_revision.get(); self.record_preparation(target_revision, PreparationDiagnostic::Preparing); - let result = self - .prepare_inner(snapshot, required_revision, request_body_limit) - .await; + let result = tokio::select! { + biased; + () = self.shutting_down.cancelled() => Err(AdminError::unavailable("插件运行时正在关闭")), + result = self.prepare_inner(snapshot, required_revision, request_body_limit) => result, + }; match &result { Ok(reference) => self.record_preparation( target_revision, @@ -63,7 +64,7 @@ impl PluginRuntime { required_revision: Option, request_body_limit: Option, ) -> Result { - if self.shutting_down.load(Ordering::Acquire) { + if self.shutting_down.is_cancelled() { return Err(AdminError::unavailable("插件运行时正在关闭")); } if snapshot.instances.len() > 64 { @@ -72,7 +73,7 @@ impl PluginRuntime { let fingerprint = fingerprint(&snapshot)?; // 准备串行化与弱索引分锁；诊断和已发布请求不会等待外部 I/O let _prepare = self.prepare_lock.lock().await; - if self.shutting_down.load(Ordering::Acquire) { + if self.shutting_down.is_cancelled() { return Err(AdminError::unavailable("插件运行时正在关闭")); } let existing = { @@ -125,17 +126,12 @@ impl PluginRuntime { let result = if let Some(instance) = reusable { Ok(instance) } else { - tokio::time::timeout( - // 沿用准备阶段原有的 30 秒预算；打包校验与能力恢复也在预算内 - Duration::from_secs(30), - self.prepare_instance( - instance.clone(), - snapshot.config_revision, - request_body_limit, - ), + self.prepare_instance( + instance.clone(), + snapshot.config_revision, + request_body_limit, ) .await - .unwrap_or_else(|_| Err(AdminError::unavailable("插件启动超时"))) }; (instance, identity, result) }, @@ -208,7 +204,6 @@ impl PluginRuntime { .rpc_limits .maximum_calls .min(self.config.rpc_limits.maximum_executing_calls), - self.config.rpc_limits.maximum_call_timeout, self.config.rpc_limits.maximum_buffered_body_bytes, ) .map(|plan| { @@ -217,21 +212,16 @@ impl PluginRuntime { .map_err(|_| AdminError::invalid("插件观察计划注册冲突")) }) .transpose()?; - let policies = crate::adapter::policy::PluginRequestPolicyPlan::compile( - policy_entries, - self.config.rpc_limits.maximum_call_timeout, - )? - .map(|plan| { - self.policies - .register(id.clone(), plan) - .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) - }) - .transpose()?; - let middleware = crate::adapter::middleware::PluginMiddlewarePlan::compile( - middleware_entries, - self.config.rpc_limits.maximum_call_timeout, - ) - .map(|plan| plan as Arc); + let policies = crate::adapter::policy::PluginRequestPolicyPlan::compile(policy_entries)? + .map(|plan| { + self.policies + .register(id.clone(), plan) + .map_err(|_| AdminError::invalid("插件请求策略计划注册冲突")) + }) + .transpose()?; + let middleware = + crate::adapter::middleware::PluginMiddlewarePlan::compile(middleware_entries) + .map(|plan| plan as Arc); let upstream_adapters = crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile(upstream_entries); let execution = if middleware.is_some() || upstream_adapters.is_some() { @@ -252,7 +242,6 @@ impl PluginRuntime { let authentication = crate::adapter::frontend_authentication::PluginFrontendAuthenticationPlan::compile( authentication_entries, - self.config.rpc_limits.maximum_call_timeout, )? .map(|plan| { self.authentications diff --git a/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs index 2b83fa96a..6b0bab4f6 100644 --- a/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs +++ b/backend/crates/gateway-plugin/runtime/src/generation/prepare/state.rs @@ -90,9 +90,10 @@ impl PluginRuntime { .session .call( "plugin.state.migrate", - instance - .session - .context(Stage::Configuration, Duration::from_secs(30)), + instance.session.context( + Stage::Configuration, + instance.session.budget(Stage::Configuration), + ), serde_json::json!({}), serde_json::to_vec(&request) .map_err(|_| AdminError::internal("状态迁移批次无法编码"))?, @@ -207,7 +208,11 @@ impl gateway_admin::ports::plugins::PluginStateLifecycle for PluginRuntime { prepared: &ExtensionSetReference, transition: PluginStateTransition, ) -> Result<(), AdminError> { - self.run_state_migration(prepared, transition).await + tokio::select! { + biased; + () = self.shutting_down.cancelled() => Err(AdminError::unavailable("插件运行时正在关闭")), + result = self.run_state_migration(prepared, transition) => result, + } } } pub(super) fn map_state_admin_error( diff --git a/backend/crates/gateway-plugin/runtime/src/rpc/session.rs b/backend/crates/gateway-plugin/runtime/src/rpc/session.rs index 78c40ca3f..d38f33d6e 100644 --- a/backend/crates/gateway-plugin/runtime/src/rpc/session.rs +++ b/backend/crates/gateway-plugin/runtime/src/rpc/session.rs @@ -12,10 +12,13 @@ use std::{ use futures::future::BoxFuture; use gateway_host::process::{ProcessControl, ProcessSpec, ProcessStartError, ProcessSupervisor}; use gateway_plugin_sdk::{ - CallContext, Frame, Handshake, Message, PROTOCOL_VERSION, PluginFault, Stage, + CallBudgetError, CallBudgetSource, CallContext, Frame, Handshake, Message, PROTOCOL_VERSION, + PluginFault, ResolvedCallBudgets, Stage, client::{CallDeadline, CallDeadlinePause, EncodedFrame, read_frame, write_frame}, }; -use tokio::sync::{Mutex as AsyncMutex, OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}; +use tokio::sync::{ + Mutex as AsyncMutex, Notify, OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch, +}; use crate::{PreparedPackage, RpcStream, stream::StreamIngress}; @@ -32,7 +35,6 @@ pub struct RpcLimits { pub maximum_queue_timeout: Duration, pub maximum_callbacks: usize, pub handshake_timeout: Duration, - pub maximum_call_timeout: Duration, pub maximum_stderr_bytes: usize, } @@ -46,7 +48,6 @@ impl Default for RpcLimits { maximum_queue_timeout: Duration::from_secs(30), maximum_callbacks: 64, handshake_timeout: Duration::from_secs(5), - maximum_call_timeout: Duration::from_secs(120), maximum_stderr_bytes: 64 * 1024, } } @@ -58,6 +59,8 @@ pub enum RpcError { Start(#[source] ProcessStartError), #[error("plugin handshake does not match the prepared instance")] Handshake(#[source] Option), + #[error("plugin call budget negotiation failed: {0}")] + Budget(#[from] CallBudgetError), #[error("plugin protocol is invalid")] Protocol(#[source] Option), #[error("plugin returned an invalid response at {0:?}")] @@ -89,6 +92,7 @@ impl RpcError { match self { RpcError::Start(_) => ("process_start", "插件进程启动失败"), RpcError::Handshake(_) => ("handshake", "插件握手失败"), + RpcError::Budget(_) => ("call_budget", "插件调用预算协商失败"), RpcError::Protocol(_) => ("protocol", "插件协议错误"), RpcError::InvalidResponse(stage, _) => ( "invalid_response", @@ -191,6 +195,7 @@ pub trait CallbackHandler: Send + Sync { pub struct RpcSession { shared: Arc, handshake: Handshake, + budgets: ResolvedCallBudgets, data: mpsc::Sender, slots: Arc, next_call: AsyncMutex, @@ -222,6 +227,7 @@ pub(crate) struct Shared { callback_slots: Arc, next_slots: Arc, execution: ExecutionCapacity, + cancellations_changed: Notify, } struct State { @@ -295,6 +301,7 @@ impl Shared { } self.stopped.send_replace(true); self.process.stop(); + self.cancellations_changed.notify_waiters(); } fn expect_stop(&self) { @@ -512,6 +519,8 @@ impl Shared { return Err(RpcError::Protocol(None)); } state.calls.remove(&id); + drop(state); + self.cancellations_changed.notify_waiters(); Ok(()) } @@ -568,8 +577,36 @@ impl Shared { } impl RpcSession { - pub(crate) const fn maximum_call_timeout(&self) -> Duration { - self.limits.maximum_call_timeout + pub fn budget(&self, stage: Stage) -> Duration { + self.budgets.budget(stage) + } + + #[must_use] + pub(crate) fn budget_source(&self, stage: Stage) -> CallBudgetSource { + self.budgets.source(stage) + } + + /// 替换维护任务前等待已取消的调用确认；未确认仍沿用会话的有界故障回收 + pub(crate) async fn settle_cancellations(&self) { + let ids: Vec<_> = self + .shared + .state() + .calls + .iter() + .filter_map(|(id, call)| call.cancellation.is_some().then_some(*id)) + .collect(); + loop { + let changed = self.shared.cancellations_changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if !ids.iter().any(|id| self.shared.is_cancelling(*id)) { + break; + } + changed.await; + } + if !self.is_ready() { + self.shared.process.exited().await; + } } pub async fn start( @@ -620,7 +657,6 @@ impl RpcSession { || limits.maximum_buffered_body_bytes < 1024 || limits.maximum_buffered_body_bytes > 16 * 1024 * 1024 || limits.handshake_timeout.is_zero() - || limits.maximum_call_timeout.is_zero() { return Err(RpcError::Context(None)); } @@ -691,11 +727,12 @@ impl RpcSession { Message::Ready { protocol_version, incarnation, + call_budgets, } if protocol_version == handshake.protocol_version && incarnation == handshake.incarnation && reply.payload.is_empty() => { - Ok(()) + ResolvedCallBudgets::resolve(&call_budgets).map_err(RpcError::Budget) } _ => Err(RpcError::Handshake(None)), } @@ -703,15 +740,18 @@ impl RpcSession { .await .map_err(|_| RpcError::Timeout) .and_then(|result| result); - if let Err(error) = ready { - if let Some(lifecycle) = &lifecycle { - // 只有 Ready 之后的存活才可复位预算；长时间卡在握手不等于稳定运行 - lifecycle.fail(Duration::ZERO); + let budgets = match ready { + Ok(budgets) => budgets, + Err(error) => { + if let Some(lifecycle) = &lifecycle { + // 只有 Ready 之后的存活才可复位预算；长时间卡在握手不等于稳定运行 + lifecycle.fail(Duration::ZERO); + } + process.stop(); + process.exited().await; + return Err(error); } - process.stop(); - process.exited().await; - return Err(error); - } + }; let (data, data_receiver) = mpsc::channel(limits.maximum_calls); let (control, control_receiver) = mpsc::channel(limits.maximum_callbacks + 2 * limits.maximum_calls + 8); @@ -737,6 +777,7 @@ impl RpcSession { limits.maximum_queued_calls, limits.maximum_queue_timeout, ), + cancellations_changed: Notify::new(), }); dispatch::start( input, @@ -760,6 +801,7 @@ impl RpcSession { Ok(Self { shared, handshake, + budgets, data, slots: Arc::new(Semaphore::new(limits.maximum_calls)), next_call: AsyncMutex::new(1), @@ -799,8 +841,8 @@ impl RpcSession { #[must_use] pub fn context(&self, stage: gateway_plugin_sdk::Stage, timeout: Duration) -> CallContext { - // 宿主生成的上下文受本会话预算约束；固定注册期限不能使较小的总预算无法启动 - let timeout = timeout.min(self.limits.maximum_call_timeout); + // 调用方可以传入更短的外部剩余期限，不能扩大实例已协商的阶段预算 + let timeout = timeout.min(self.budget(stage)); CallContext { request: None, call_id: 0, @@ -823,7 +865,7 @@ impl RpcSession { || context.generation != self.handshake.generation || context.incarnation != self.handshake.incarnation || context.timeout_ms == 0 - || Duration::from_millis(context.timeout_ms) > self.limits.maximum_call_timeout + || Duration::from_millis(context.timeout_ms) > self.budget(context.stage) || method.is_empty() || method.len() > 128 { diff --git a/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/execution.rs b/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/execution.rs index 14a5f7d77..6828ea34b 100644 --- a/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/execution.rs +++ b/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/execution.rs @@ -6,6 +6,120 @@ use serde_json::json; use crate::support::environment::Environment; +#[derive(Default)] +struct BudgetAccountAccess { + calls: std::sync::atomic::AtomicUsize, + first_started: tokio::sync::Notify, + second_started: tokio::sync::Notify, +} + +#[async_trait::async_trait] +impl gateway_admin::ports::plugin_accounts::PluginAccountAccess for BudgetAccountAccess { + async fn list( + &self, + _: gateway_admin::model::provider_credentials::PluginAccountListQuery, + ) -> Result< + gateway_admin::model::provider_credentials::PluginAccountPage, + gateway_admin::model::AdminError, + > { + panic!("账号提交不应读取账号列表") + } + + async fn get_runtime( + &self, + _: &gateway_core::account::ProviderAccountId, + ) -> Result + { + panic!("创建账号不应读取已有账号") + } + + async fn get_credential( + &self, + _: &gateway_core::account::ProviderAccountId, + ) -> Result< + gateway_admin::model::provider_credentials::PluginAccountCredential, + gateway_admin::model::AdminError, + > { + panic!("账号提交不应读取凭据") + } + + async fn save( + &self, + _: gateway_admin::model::provider_credentials::PreparedPluginAccountSave, + _: &gateway_admin::model::MutationContext, + ) -> Result< + gateway_admin::model::provider_credentials::PluginAccountSaveResult, + gateway_admin::model::AdminError, + > { + if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + self.first_started.notify_one(); + tokio::time::sleep(std::time::Duration::from_secs(70)).await; + Ok( + gateway_admin::model::provider_credentials::PluginAccountSaveResult { + config_revision: gateway_admin::model::Revision::new(2).unwrap(), + account_id: gateway_core::account::ProviderAccountId::new( + "acct_command_budget", + ) + .unwrap(), + credential_revision: gateway_admin::model::Revision::new(1).unwrap(), + }, + ) + } else { + self.second_started.notify_one(); + std::future::pending().await + } + } +} + +#[tokio::test] +async fn account_commits_share_the_commands_declared_deadline_and_report_partial_completion() { + use gateway_admin::ports::plugin_accounts::PluginAccountAccess; + use std::{sync::Arc, time::Duration}; + + let account = json!({"action":"create","provider_id":"openai","facts":{ + "name":"预算测试","authentication_kind":"api_key","material":{"key":"fixture-only"} + }}); + let config = json!({ + "call_budgets_ms":{"command_line":180_000}, + "command_registration":{"commands":[{"name":"save","description":"提交账号"}]}, + "command_result":{"stdout":"","stderr":"","exit_code":0,"accounts":[account, account]}, + }); + let (_directory, _store, runtime) = super::setup(config).await; + let probe = Arc::new(BudgetAccountAccess::default()); + let accounts: Arc = probe.clone(); + runtime.bind_account_ports(&accounts).unwrap(); + let commands = runtime.prepare_command_line().await.unwrap(); + let task = tokio::spawn(async move { + let result = commands.execute("commands", "save", &[]).await; + (commands, result) + }); + tokio::time::timeout(Duration::from_secs(5), probe.first_started.notified()) + .await + .unwrap(); + + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(121)).await; + probe.second_started.notified().await; + assert!( + !task.is_finished(), + "账号提交沿用声明预算，不能被默认期限截断" + ); + tokio::time::advance(Duration::from_secs(60)).await; + let (commands, result) = task.await.unwrap(); + tokio::time::resume(); + assert!( + matches!( + result, + Err(PluginCommandError::Incomplete { + saved_accounts: 1, + reason: "超过命令总期限", + }) + ), + "后续账号提交不能重新获得完整预算" + ); + commands.shutdown().await; +} + #[tokio::test] async fn model_requests_have_a_body_limit_independent_of_the_response_buffer() { use gateway_plugin_sdk::call::model::request_payload_limit; diff --git a/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/mod.rs b/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/mod.rs index f9a15fb7c..46da41c7c 100644 --- a/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/adapter/command_line/mod.rs @@ -174,8 +174,7 @@ async fn failed_or_cancelled_commands_are_not_replayed_and_children_are_reaped() } else { // 正常完成注册并确认命令已进入子进程后，只推进执行期限 tokio::time::pause(); - tokio::time::advance(limits.maximum_call_timeout + Duration::from_millis(1)) - .await; + tokio::time::advance(Duration::from_secs(120) + Duration::from_millis(1)).await; } } }; @@ -209,3 +208,46 @@ async fn failed_or_cancelled_commands_are_not_replayed_and_children_are_reaped() } } } + +#[tokio::test] +async fn command_declared_budget_survives_the_default_and_expires_at_its_own_deadline() { + use std::time::Duration; + + let mut config = configuration(json!([])); + config["command_wait"] = json!(true); + config["call_budgets_ms"] = json!({"command_line":180_000}); + let (directory, store, runtime) = setup(config).await; + let marker = directory.path().join("command-budget.jsonl"); + store.snapshot.lock().unwrap().instances[0].configuration["command_marker"] = json!(marker); + let commands = runtime.prepare_command_line().await.unwrap(); + let task = tokio::spawn(async move { + let result = commands.execute("commands", "inspect", &[]).await; + (commands, result) + }); + tokio::time::timeout(Duration::from_secs(5), async { + while !marker.exists() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap(); + + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(121)).await; + tokio::task::yield_now().await; + assert!(!task.is_finished(), "声明的命令预算不受默认 120 秒截断"); + tokio::time::advance(Duration::from_secs(60)).await; + let (commands, result) = task.await.unwrap(); + tokio::time::resume(); + let error = result.err().unwrap(); + assert!(matches!( + error, + gateway_plugin_runtime::PluginCommandError::Incomplete { + saved_accounts: 0, + .. + } + )); + assert!(error.to_string().contains("超时") || error.to_string().contains("总期限")); + commands.shutdown().await; + assert_eq!(std::fs::read_to_string(marker).unwrap().lines().count(), 1); +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs index e9bcf020a..6964e8b1e 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/mod.rs @@ -41,6 +41,24 @@ async fn http_plugin_with_policy( PluginRuntime, ExtensionSetReference, gateway_core::engine::middleware::FrozenMiddlewarePlan, +) { + http_plugin_with_configuration( + serde_json::json!({"http":true,"mode":mode}), + limits, + failure_policy, + ) + .await +} + +async fn http_plugin_with_configuration( + configuration: serde_json::Value, + limits: RpcLimits, + failure_policy: PluginFailurePolicy, +) -> ( + tempfile::TempDir, + PluginRuntime, + ExtensionSetReference, + gateway_core::engine::middleware::FrozenMiddlewarePlan, ) { let worker = std::fs::read(env!("CARGO_BIN_EXE_gateway-plugin-test-middleware")).unwrap(); let package = crate::support::package_with_contributions( @@ -55,7 +73,7 @@ async fn http_plugin_with_policy( let (cache, runtime) = setup_package_with_limits( vec![InstanceFixture { id: "http-plugin", - configuration: serde_json::json!({"http":true,"mode":mode}), + configuration, bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, "http", 0, failure_policy)], }], package, diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs index 080b084ed..93f7b1be6 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/http/websocket/mod.rs @@ -64,14 +64,11 @@ async fn session_model_and_network_callbacks_outlive_handshake_and_each_executio let key = format!("key_{}", uuid::Uuid::new_v4().simple()); environment.client_key(&key, "sk-session-fixture").await; environment - .install_plugin(serde_json::json!({"http":true,"mode":"callbacks"})) - .await; - let (runtime, bundle) = environment - .runtime_with_limits(RpcLimits { - maximum_call_timeout: Duration::from_millis(500), - ..RpcLimits::default() - }) + .install_plugin( + serde_json::json!({"http":true,"mode":"callbacks", "call_budgets_ms":{"http":500}}), + ) .await; + let (runtime, bundle) = environment.runtime_with_limits(RpcLimits::default()).await; environment.store.start_command_line_writes().unwrap(); let snapshot = bundle.snapshots().snapshot_for_diagnostics().unwrap(); let plan = runtime @@ -185,12 +182,10 @@ async fn session_model_and_network_callbacks_outlive_handshake_and_each_executio #[tokio::test] async fn real_plugin_controls_upgrade_and_concurrent_duplex_with_large_binary_frames() { - let (_cache, runtime, generation, plan) = http_plugin_with_limits( - "upgrade", - RpcLimits { - maximum_call_timeout: Duration::from_millis(500), - ..RpcLimits::default() - }, + let (_cache, runtime, generation, plan) = http_plugin_with_configuration( + serde_json::json!({"http":true,"mode":"upgrade", "call_budgets_ms":{"http":500}}), + RpcLimits::default(), + PluginFailurePolicy::Reject, ) .await; let (send, incoming) = tokio::sync::mpsc::channel(1); diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/lifecycle.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/lifecycle.rs index c8f49d63f..1bc77b99f 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/lifecycle.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/lifecycle.rs @@ -147,7 +147,8 @@ async fn sdk_budget_runtime(stage: Stage, count: usize) -> (tempfile::TempDir, P .enumerate() .map(|(order, id)| InstanceFixture { id, - configuration: serde_json::json!({"mode":"map", "attempt":stage == Stage::Attempt}), + configuration: serde_json::json!({"mode":"map", "attempt":stage == Stage::Attempt, + "call_budgets_ms":{"request":300,"attempt":300}}), bindings: vec![binding( MIDDLEWARE_CONTRIBUTION, mount, @@ -156,15 +157,7 @@ async fn sdk_budget_runtime(stage: Stage, count: usize) -> (tempfile::TempDir, P )], }) .collect(); - setup_package_with_limits( - instances, - package, - RpcLimits { - maximum_call_timeout: Duration::from_millis(300), - ..RpcLimits::default() - }, - ) - .await + setup_package_with_limits(instances, package, RpcLimits::default()).await } fn budget_request() -> MiddlewareRequest { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs index 14e752949..14acd116e 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/mod.rs @@ -369,54 +369,85 @@ async fn retry_policy_preserves_decisions_and_leaves_replay_safety_to_core() { } #[tokio::test] -async fn retry_policy_timeout_exhausts_the_chain_budget_and_delegates() { +async fn retry_policy_combines_instance_budgets_with_one_request_bounded_chain() { use gateway_core::{ engine::policy::{RetryDecision, RetryFacts}, error::ProviderErrorKind, upstream::UpstreamSendState, }; - let records = tempfile::tempdir().unwrap(); - let second_marker = records.path().join("second.jsonl"); - let (_cache, runtime) = setup( - vec![ + for ( + first_budget, + first_delay, + second_budget, + second_delay, + request_budget, + expected, + reached_second, + ) in [ + (100, 500, 1000, 100, 2000, RetryDecision::Stop, true), + (1000, 100, 50, 300, 2000, RetryDecision::Delegate, true), + (1000, 300, 1000, 800, 2000, RetryDecision::Delegate, true), + (1000, 500, 1000, 0, 100, RetryDecision::Delegate, false), + (3000, 2100, 3000, 0, 5000, RetryDecision::Stop, true), + ] { + let records = tempfile::tempdir().unwrap(); + let second_marker = records.path().join("second.jsonl"); + let (_cache, runtime) = setup(vec![ InstanceFixture { id: "retry-delayed", - configuration: serde_json::json!({"retry_delay_ms":500,"retry_decision":{"decision":"stop"}}), - - bindings: vec![binding("test.example.retryPolicy", "retry", 0, PluginFailurePolicy::Delegate)], + configuration: serde_json::json!({ + "call_budgets_ms":{"retry":first_budget}, + "retry_delay_ms":first_delay, + "retry_decision":{"decision":"delegate"}, + }), + bindings: vec![binding( + "test.example.retryPolicy", + "retry", + 0, + PluginFailurePolicy::Delegate, + )], }, InstanceFixture { id: "retry-next", - configuration: serde_json::json!({"retry_marker":second_marker,"retry_decision":{"decision":"stop"}}), - - bindings: vec![binding("test.example.retryPolicy", "retry", 1, PluginFailurePolicy::Delegate)], + configuration: serde_json::json!({ + "call_budgets_ms":{"retry":second_budget}, + "retry_delay_ms":second_delay, + "retry_marker":second_marker, + "retry_decision":{"decision":"stop"}, + }), + bindings: vec![binding( + "test.example.retryPolicy", + "retry", + 1, + PluginFailurePolicy::Delegate, + )], }, - ], - - ).await; - let context = policy_context(&runtime, prepare(&runtime).await, "req_retry_timeout"); - let decision = tokio::time::timeout( - Duration::from_secs(2), - context.retry_decision(RetryFacts { - attempt_index: NonZeroU32::MIN, - provider: ProviderKind::new("openai").unwrap(), - model: Some("gpt-test".into()), - error_kind: ProviderErrorKind::RateLimited, - upstream_status: Some(429), - send_state: UpstreamSendState::NotSent, - remaining_routing_attempts: 1, - remaining_deadline: Duration::from_millis(100), - retry_allowed: true, - }), - ) - .await - .expect("the shared deadline bounds the whole retry chain"); - assert_eq!(decision, RetryDecision::Delegate); - assert!( - !second_marker.exists(), - "later policies cannot restart an exhausted budget" - ); - runtime.shutdown().await; + ]) + .await; + let context = policy_context(&runtime, prepare(&runtime).await, "req_retry_timeout"); + let decision = tokio::time::timeout( + Duration::from_secs(6), + context.retry_decision(RetryFacts { + attempt_index: NonZeroU32::MIN, + provider: ProviderKind::new("openai").unwrap(), + model: Some("gpt-test".into()), + error_kind: ProviderErrorKind::RateLimited, + upstream_status: Some(429), + send_state: UpstreamSendState::NotSent, + remaining_routing_attempts: 1, + remaining_deadline: Duration::from_millis(request_budget), + retry_allowed: true, + }), + ) + .await + .expect("重试链不能重置总预算"); + assert_eq!( + decision, expected, + "{first_budget}/{second_budget}/{request_budget}" + ); + assert_eq!(second_marker.exists(), reached_second); + runtime.shutdown().await; + } } async fn marker_lines(path: &std::path::Path, count: usize) -> Vec { diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/outbound.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/outbound.rs index 1d16aabe4..108def7fd 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/outbound.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/outbound.rs @@ -268,13 +268,18 @@ async fn outbound_next_wait_pauses_plugin_budget_and_cancel_releases_the_continu for cancel in [false, true] { let (cache, runtime) = runtime( vec![ - instance("outer", 0, serde_json::json!({})), - instance("inner", 1, serde_json::json!({})), + instance( + "outer", + 0, + serde_json::json!({"call_budgets_ms":{"outbound":300}}), + ), + instance( + "inner", + 1, + serde_json::json!({"call_budgets_ms":{"outbound":300}}), + ), ], - RpcLimits { - maximum_call_timeout: Duration::from_millis(300), - ..RpcLimits::default() - }, + RpcLimits::default(), ) .await; let generation = prepare(&runtime).await; diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/policy/retry.rs b/backend/crates/gateway-plugin/runtime/tests/generation/policy/retry.rs index 425ed4d93..2996100a1 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/policy/retry.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/policy/retry.rs @@ -136,6 +136,7 @@ async fn retry_callback_effects_survive_success_failure_and_cancellation() { "retry_http_return_marker":(outcome == "early_result").then_some(marker), "retry_service_request":is_service.then_some(service_request), "retry_service_marker":service_marker, + "call_budgets_ms":{"retry":500}, })) .await; let key = "sk-retry-effects-fixture"; @@ -146,10 +147,7 @@ async fn retry_callback_effects_survive_success_failure_and_cancellation() { }); let (runtime, core) = environment .runtime_with_providers( - RpcLimits { - maximum_call_timeout: Duration::from_millis(500), - ..RpcLimits::default() - }, + RpcLimits::default(), ProviderRegistry::new([provider.clone() as Arc]).unwrap(), ) .await; diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs index f9f729af2..6cebbd789 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/instance.rs @@ -552,3 +552,55 @@ async fn a_host_version_mismatch_still_restores_the_enabled_plugin() { runtime.shutdown().await; super::super::wait_until_empty(cache.path()).await; } + +#[tokio::test] +async fn declared_registration_budget_is_not_cut_short_by_process_preparation_timeout() { + let (cache, store, runtime) = super::super::setup().await; + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "registration_delay_ms":31_000, + "call_budgets_ms":{"registration":40_000}, + }); + let candidate = tokio::time::timeout( + Duration::from_secs(45), + PluginPreparation::prepare(&runtime, snapshot), + ) + .await + .expect("registration stays bounded by its declared budget") + .expect("ready process may register beyond the startup control deadline"); + assert!(candidate.is_ready()); + drop(candidate); + super::super::wait_until_empty(cache.path()).await; +} + +#[tokio::test] +async fn shutdown_cancels_long_registration_and_reclaims_the_unpublished_candidate() { + let (cache, store, runtime) = super::super::setup().await; + let runtime = std::sync::Arc::new(runtime); + let markers = tempfile::tempdir().unwrap(); + let started = markers.path().join("registration"); + let mut snapshot = store.snapshot.lock().unwrap().clone(); + snapshot.instances[0].configuration = serde_json::json!({ + "registration_started_marker":started, + "registration_delay_ms":600_000, + "call_budgets_ms":{"registration":600_000}, + }); + let preparing = tokio::spawn({ + let runtime = std::sync::Arc::clone(&runtime); + async move { PluginPreparation::prepare(runtime.as_ref(), snapshot).await } + }); + tokio::time::timeout(Duration::from_secs(5), async { + while !started.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(3), runtime.shutdown()) + .await + .expect( + "shutdown must cancel candidate preparation instead of waiting for its long budget", + ); + assert!(preparing.await.unwrap().is_err()); + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/maintenance.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/maintenance.rs new file mode 100644 index 000000000..ee7ba4b95 --- /dev/null +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/maintenance.rs @@ -0,0 +1,157 @@ +//! 验证维护任务跨发布保活、通知合并及实例替换后的取消顺序 + +use gateway_admin::{ + model::{Revision, plugins::instances::PluginInstanceSnapshot}, + ports::plugins::PluginPreparation, +}; +use gateway_core::{ + lifecycle::CancellationToken, + routing::{ConfigRevision, RuntimeSnapshot}, + runtime::RuntimeSnapshotHandle, + task::{WorkerContribution, WorkerDaemonContext, WorkerRunnable}, +}; +use gateway_plugin_runtime::PluginRuntime; +use gateway_plugin_sdk::{Capability, Contributions, Stage}; +use serde_json::{Value, json}; +use std::{path::Path, sync::Arc, time::Duration}; + +async fn publish( + runtime: &PluginRuntime, + snapshots: &RuntimeSnapshotHandle, + instances: PluginInstanceSnapshot, +) { + let revision = ConfigRevision::new(instances.config_revision.get()).unwrap(); + let prepared = PluginPreparation::prepare(runtime, instances) + .await + .unwrap(); + snapshots.publish( + RuntimeSnapshot::new( + revision, + gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None), + vec![], + vec![], + vec![], + ) + .unwrap() + .with_extensions(Some(prepared)), + ); +} + +fn records(path: &Path) -> Vec { + std::fs::read_to_string(path) + .unwrap_or_default() + .lines() + .filter_map(|line| serde_json::from_str(line).ok()) + .collect() +} + +async fn wait_records(path: &Path, count: usize) -> Vec { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let entries = records(path); + if entries.len() >= count { + return entries; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|error| panic!("missing maintenance records: {error}; {:?}", records(path))) +} + +#[tokio::test] +async fn unrelated_publications_preserve_inflight_maintenance_and_replacement_waits_for_cancel() { + let (cache, store, runtime) = super::super::setup_with_contributions(Contributions::from([ + crate::support::contribution( + Capability::Maintenance, + vec![Stage::Maintenance], + vec![], + vec![], + ), + ])) + .await; + let runtime = Arc::new(runtime); + let markers = tempfile::tempdir().unwrap(); + let mut instances = store.snapshot.lock().unwrap().clone(); + let template = instances.instances[0].clone(); + instances.instances = ["a", "b"] + .into_iter() + .map(|id| { + let mut instance = template.clone(); + instance.id = id.into(); + instance.configuration = json!({ + "maintenance_marker":markers.path().join(id), + "execution_cancelled_marker":markers.path().join(id), + "maintenance_wait_marker":markers.path().join(format!("release-{id}")), + "call_context_marker":markers.path().join(format!("context-{id}")), + "call_budgets_ms":{"maintenance":600_000}, + }); + instance + }) + .collect(); + let snapshots = RuntimeSnapshotHandle::default(); + publish(&runtime, &snapshots, instances.clone()).await; + let WorkerContribution::Registration(registration) = + runtime.maintenance_worker(snapshots.clone()).unwrap() + else { + panic!("maintenance worker registration") + }; + let WorkerRunnable::Daemon { task, .. } = registration.runnable else { + panic!("maintenance daemon") + }; + let stop = CancellationToken::new(); + let running = tokio::spawn({ + let stop = stop.clone(); + async move { task.run(WorkerDaemonContext::new(stop)).await.unwrap() } + }); + let a = markers.path().join("a"); + let b = markers.path().join("b"); + assert_eq!(wait_records(&a, 1).await[0]["phase"], "start"); + assert_eq!(wait_records(&b, 1).await[0]["phase"], "start"); + let contexts = records(&markers.path().join("context-a")); + let maintenance = contexts + .iter() + .find(|entry| entry["method"] == "plugin.reconcile") + .expect("maintenance call context"); + assert_eq!(maintenance["context"]["timeout_ms"], 600_000); + + instances.config_revision = Revision::new(2).unwrap(); + instances.instances[1].revision = instances.config_revision; + publish(&runtime, &snapshots, instances.clone()).await; + let changed = wait_records(&b, 3).await; + assert!( + changed[1].get("call_id").is_some(), + "旧维护必须先完成取消确认" + ); + assert_eq!(changed[2]["phase"], "start"); + assert_eq!(records(&a).len(), 1, "无关实例发布不得重启 A 的在途维护"); + + instances.config_revision = Revision::new(3).unwrap(); + instances.instances[1].revision = instances.config_revision; + instances.instances[1].enabled = false; + publish(&runtime, &snapshots, instances.clone()).await; + assert!(wait_records(&b, 4).await[3].get("call_id").is_some()); + instances.config_revision = Revision::new(4).unwrap(); + instances.instances[1].revision = instances.config_revision; + instances.instances[1].enabled = true; + publish(&runtime, &snapshots, instances.clone()).await; + assert_eq!(wait_records(&b, 5).await[4]["phase"], "start"); + assert_eq!(records(&a).len(), 1, "停用与恢复 B 不得取消 A"); + + // 多次发布通知合并为一次后续对账，首轮仍使用原调用 ID 完成 + std::fs::write(markers.path().join("release-a"), []).unwrap(); + let completed = wait_records(&a, 4).await; + assert_eq!(completed[1]["phase"], "done"); + assert_eq!(completed[1]["id"], completed[0]["id"]); + assert_eq!(completed[2]["phase"], "start"); + assert_eq!(completed[3]["phase"], "done"); + assert_ne!(completed[2]["id"], completed[0]["id"]); + assert_eq!(completed.len(), 4); + + stop.cancel(); + running.await.unwrap(); + assert!(wait_records(&b, 6).await[5].get("call_id").is_some()); + snapshots.suspend(); + runtime.shutdown().await; + super::super::wait_until_empty(cache.path()).await; +} diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs index e7531c5cc..008302184 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/prepare/mod.rs @@ -2,6 +2,7 @@ mod diagnostics; mod instance; +mod maintenance; mod management; mod runtime; mod set; diff --git a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs index 36553cbb7..a95244141 100644 --- a/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs +++ b/backend/crates/gateway-plugin/runtime/tests/generation/upstream_adapter.rs @@ -533,14 +533,8 @@ async fn adapter_sse_outlives_rpc_budget_and_cancellation_closes_its_resources() .as_object_mut() .unwrap() .remove("wire"); - let (cache, _, runtime) = setup_with_limits( - config, - gateway_plugin_runtime::RpcLimits { - maximum_call_timeout: Duration::from_millis(500), - ..Default::default() - }, - ) - .await; + config["call_budgets_ms"] = json!({"upstream":500}); + let (cache, _, runtime) = setup(config).await; let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap()) .await diff --git a/backend/crates/gateway-plugin/runtime/tests/rpc/session.rs b/backend/crates/gateway-plugin/runtime/tests/rpc/session.rs index 1be5f5e35..d438c0eab 100644 --- a/backend/crates/gateway-plugin/runtime/tests/rpc/session.rs +++ b/backend/crates/gateway-plugin/runtime/tests/rpc/session.rs @@ -36,6 +36,33 @@ pub(super) struct LifecycleCallbacks { struct OversizedReplyCallbacks; +#[derive(Default)] +struct BudgetCallbacks { + started: Arc, + release: Arc, +} + +impl CallbackHandler for BudgetCallbacks { + fn call( + &self, + context: CallContext, + _method: String, + _params: serde_json::Value, + payload: Vec, + ) -> BoxFuture<'static, Result> { + let started = Arc::clone(&self.started); + let release = Arc::clone(&self.release); + Box::pin(async move { + started.notify_one(); + release.notified().await; + Ok(RpcReply { + result: json!({"budget_ms": context.timeout_ms}), + payload, + }) + }) + } +} + impl CallbackHandler for OversizedReplyCallbacks { fn call( &self, @@ -336,7 +363,7 @@ async fn real_process_round_trips_raw_bytes_and_reentrant_callbacks() { session .context(Stage::Registration, Duration::MAX) .timeout_ms, - u64::try_from(RpcLimits::default().maximum_call_timeout.as_millis()).unwrap() + 5_000 ); *callbacks.nested.lock().unwrap() = Some(Arc::downgrade(&session)); let context = session.context(Stage::Management, Duration::from_secs(2)); @@ -355,6 +382,106 @@ async fn real_process_round_trips_raw_bytes_and_reentrant_callbacks() { session.shutdown(Duration::from_secs(1)).await; } +#[tokio::test] +async fn declared_stage_budget_covers_long_callback_without_affecting_other_calls() { + let (_cache, prepared, mut handshake, processes) = prepared_session(); + handshake.configuration = json!({"call_budgets_ms":{"maintenance":600_000,"management":40}}); + let callbacks = Arc::new(BudgetCallbacks::default()); + let session = Arc::new( + RpcSession::start( + prepared, + handshake, + RpcLimits::default(), + &processes, + callbacks.clone(), + ) + .await + .unwrap(), + ); + let calling = { + let session = Arc::clone(&session); + tokio::spawn(async move { + session + .call( + "callback", + session.context(Stage::Maintenance, session.budget(Stage::Maintenance)), + json!({}), + vec![], + ) + .await + }) + }; + callbacks.started.notified().await; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(125)).await; + assert!(!calling.is_finished()); + tokio::time::resume(); + let short = session + .call( + "hang", + session.context(Stage::Management, session.budget(Stage::Management)), + json!({}), + vec![], + ) + .await; + assert!(matches!(short, Err(RpcError::Timeout))); + callbacks.release.notify_one(); + let reply = calling.await.unwrap().unwrap(); + assert_eq!(reply.result, json!({"budget_ms":600_000})); + assert!(session.is_ready()); + session.shutdown(Duration::from_secs(1)).await; +} + +#[tokio::test] +async fn invalid_budget_rejects_candidate_before_business_and_releases_process_capacity() { + let (_cache, prepared, mut handshake, _) = prepared_session(); + let processes = + gateway_host::process::ProcessSupervisor::new(std::num::NonZeroUsize::new(1).unwrap()); + handshake.configuration = json!({"call_budgets_ms":{"maintenance":0}}); + let callbacks = Arc::new(LifecycleCallbacks::default()); + let result = RpcSession::start( + prepared.clone(), + handshake.clone(), + RpcLimits::default(), + &processes, + callbacks.clone(), + ) + .await; + assert!(matches!( + result, + Err(RpcError::Budget( + gateway_plugin_sdk::CallBudgetError::Invalid { + stage: Stage::Maintenance, + milliseconds: 0, + } + )) + )); + assert_eq!(callbacks.begun.load(Ordering::Relaxed), 0); + handshake.configuration = json!({}); + handshake.incarnation = uuid::Uuid::new_v4().to_string(); + let session = RpcSession::start( + prepared, + handshake, + RpcLimits::default(), + &processes, + callbacks, + ) + .await + .unwrap(); + assert!( + session + .call( + "echo", + session.context(Stage::Maintenance, session.budget(Stage::Maintenance)), + json!({}), + vec![] + ) + .await + .is_ok() + ); + session.shutdown(Duration::from_secs(1)).await; +} + #[tokio::test] async fn slow_management_call_does_not_block_an_independent_call() { let (_cache, session) = session(Arc::new(Callbacks::default())).await; diff --git a/backend/crates/gateway-plugin/runtime/tests/support/middleware/mod.rs b/backend/crates/gateway-plugin/runtime/tests/support/middleware/mod.rs index e035a7896..ca68b7de8 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/middleware/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/middleware/mod.rs @@ -11,7 +11,7 @@ use gateway_plugin_sdk::{ #[tokio::main(flavor = "current_thread")] async fn main() -> Result<(), Box> { - let session = PluginSession::accept( + let mut session = PluginSession::accept( tokio::io::stdin(), tokio::io::stdout(), SessionConfig::default(), @@ -22,6 +22,9 @@ async fn main() -> Result<(), Box> { .unwrap_or("map") .to_owned(); let configuration = session.handshake().configuration.clone(); + if let Some(budgets) = configuration.get("call_budgets_ms") { + session.set_call_budgets(serde_json::from_value(budgets.clone())?)?; + } if mode == "mixed_mounts" { let contributes = Contributions::from([( Capability::Middleware, 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 a1607c334..0d2767dd8 100644 --- a/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs +++ b/backend/crates/gateway-plugin/runtime/tests/support/worker/mod.rs @@ -528,15 +528,29 @@ impl Peer { return; } "plugin.reconcile" => { - let result = if let Some(results) = self.data_queries(id).await { - self.append_observation_marker( - "maintenance_marker", - &json!({"phase":"done", "results":results}), - ); - Ok(()) - } else { - self.reconcile_fixture(id).await - }; + let result = + if let Some(path) = self.configuration["maintenance_wait_marker"].as_str() { + self.append_observation_marker( + "maintenance_marker", + &json!({"phase":"start", "id":id}), + ); + while !std::path::Path::new(path).exists() { + tokio::time::sleep(Duration::from_millis(20)).await; + } + self.append_observation_marker( + "maintenance_marker", + &json!({"phase":"done", "id":id}), + ); + Ok(()) + } else if let Some(results) = self.data_queries(id).await { + self.append_observation_marker( + "maintenance_marker", + &json!({"phase":"done", "results":results}), + ); + Ok(()) + } else { + self.reconcile_fixture(id).await + }; match result { Ok(()) => { self.send( @@ -1374,6 +1388,10 @@ impl Peer { return; } "plugin.register" => { + self.append_observation_marker("registration_started_marker", &json!({"id":id})); + if let Some(delay) = self.configuration["registration_delay_ms"].as_u64() { + tokio::time::sleep(Duration::from_millis(delay)).await; + } if let Some(request) = self.configuration.get("registration_http") { let result = self .callback_payload( @@ -1489,6 +1507,13 @@ async fn main() { Message::Ready { protocol_version: gateway_plugin_sdk::PROTOCOL_VERSION, incarnation: handshake.incarnation, + call_budgets: serde_json::from_value( + peer.configuration + .get("call_budgets_ms") + .cloned() + .unwrap_or_else(|| json!({})), + ) + .unwrap(), }, vec![], ) @@ -1519,6 +1544,10 @@ async fn main() { } => { assert!(id > last_call && !id.is_multiple_of(2) && context.call_id == id); last_call = id; + peer.append_observation_marker( + "call_context_marker", + &json!({"method":method, "context":context}), + ); if matches!( method.as_str(), "malformed_truncated_frame" | "malformed_frame_length" | "malformed_metadata" diff --git a/backend/crates/gateway-plugin/sdk/README.md b/backend/crates/gateway-plugin/sdk/README.md index e13317b44..a81a3682c 100644 --- a/backend/crates/gateway-plugin/sdk/README.md +++ b/backend/crates/gateway-plugin/sdk/README.md @@ -64,7 +64,7 @@ sequenceDiagram participant H as 宿主 Runtime participant P as 插件会话 H->>P: Hello - P-->>H: Ready（相同 incarnation） + P-->>H: Ready（相同 incarnation、所需预算声明） H->>P: 注册能力，核对清单 P-->>H: 注册结果 H->>P: Call（ID、上下文、载荷） @@ -78,7 +78,7 @@ sequenceDiagram ### 并发与期限 每个宿主插件会话默认允许 64 个执行中的 RPC 调用、64 个初次排队调用和 256 个总在途上下文。 -初次排队和等待 `next` 后的恢复排队各限 30 秒，不扣插件自身默认上限 120 秒的处理预算。 +初次排队和等待 `next` 后的恢复排队各限 30 秒，不扣插件自身的处理预算。 恢复等待占总在途名额，不占初次排队名额。 等待 `next` 时释放执行名额并暂停处理预算，返回后重新取得名额，继续同一次调用； 已返回响应头的流继续占用在途预算，不占首结果执行名额。具体组合行为见[洋葱中间件](docs/capabilities.md#洋葱中间件) @@ -91,13 +91,60 @@ sequenceDiagram 达到容量或等待上限时，中间件返回过载错误；HTTP 响应为 503（`middleware_capacity_unavailable`）， 不跳过插件或重放请求。调整 SDK 会话配置需重新构建插件包，容量应结合插件耗时、正文和响应流的资源占用评估 +### 声明调用预算 + +插件在 `PluginSession::accept` 后、`run` 前通过 `set_call_budgets` 声明各阶段所需的毫秒预算，可读取 +`session.handshake().configuration`，由自身配置决定模型、批量任务等场景的时长 + +```rust,no_run +use gateway_plugin_sdk::{CallBudgets, Stage}; +use gateway_plugin_sdk::client::{PluginHandler, PluginSession, SessionConfig, SessionError}; + +# async fn configure(reader: R, writer: W, handler: H) -> Result<(), SessionError> +# where R: tokio::io::AsyncRead + Unpin + Send + 'static, +# W: tokio::io::AsyncWrite + Unpin + Send + 'static, H: PluginHandler { +let mut session = PluginSession::accept(reader, writer, SessionConfig::default()).await?; +let maintenance_ms = session.handshake().configuration["maintenanceBudgetMs"] + .as_u64() + .unwrap_or(600_000); +session.set_call_budgets(CallBudgets::from([ + (Stage::Maintenance, maintenance_ms), + (Stage::CommandLine, 1_800_000), +]))?; +session.run(handler).await?; +# Ok(()) +# } +``` + +未声明的阶段使用协议默认值，默认表由 SDK 集中维护，宿主与插件使用同一份合同。 +默认值不限制插件显式声明的最大时长 + +| 阶段 | 默认预算 | +| --- | ---: | +| `registration` | 5 秒 | +| `configuration`、`authentication` | 30 秒 | +| `routing`、`scheduling`、`retry`、`observation` | 2 秒 | +| `maintenance`、`management`、`public_management`、`command_line` | 120 秒 | +| `http`、`service`、`websocket`、`request`、`attempt`、`outbound`、`upstream` | 120 秒 | + +`Ready.call_budgets` 携带显式声明，双方按相同规则解析并冻结该会话预算。预算必须为正的有限毫秒值， +不接受零值、数值溢出和平台无法表示的期限。修改声明需要创建新会话，不支持调用中无限续期 + +每次 `CallContext.timeout_ms` 是该调用的初始处理预算，外部剩余期限较短时可以小于会话预算。 +`call.host.remaining_budget()` 读取共享期限中的剩余值，正确反映 `next` 暂停；资源流 owner 接管寿命后返回 `None`。 +普通模型和网络回调消耗父调用剩余预算，不会重置预算。请求总期限、网络操作期限、停用和取消仍独立生效， +声明长预算不保证模型必然完成，也不提供跨宿主重启恢复；检查点和业务幂等由插件维护 + +宿主关闭使用独立宽限并取消未完成调用，不等待插件处理预算耗尽。 +注册与配置迁移位于实例准备链路，其实际耗时仍会影响宿主启动和实例就绪 + ### 帧与会话 `CallContext.timeout_ms` 限定初始处理期限。宿主对由流或连接 owner 管理的调用设置 `resource_stream`： 首个流式结果成功返回后，正文、会话及关联回调由资源取消与终态回收，不再沿用初始期限。 普通调用继续使用原期限；资源模式仍保留并发容量、Credit、父调用关联及取消约束 -协议 v5 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成，随后依次传输 JSON 元数据和二进制载荷。 +协议 v6 的帧头由大端 `u32` 元数据长度和大端 `u64` 载荷长度组成，随后依次传输 JSON 元数据和二进制载荷。 元数据保留 64 KiB 上限；正文总量不由插件传输层额外限制。传输内部以 256 KiB 为单次 I/O 块， 不逐块发起 RPC，也不要求插件配置与宿主同步。读取按进度扩展缓冲区；截断正文必须报错，不能作为完整消息交付 diff --git a/backend/crates/gateway-plugin/sdk/docs/capabilities.md b/backend/crates/gateway-plugin/sdk/docs/capabilities.md index 0a518c9bc..04b8d762a 100644 --- a/backend/crates/gateway-plugin/sdk/docs/capabilities.md +++ b/backend/crates/gateway-plugin/sdk/docs/capabilities.md @@ -138,8 +138,8 @@ Provider 不支持刷新（如 OpenAI URL + Key 账号）返回 `invalid_input` ### 自有资源与维护 声明 `maintenance` 并注册 `methods::RECONCILE`，宿主会在实例启用发布、进程恢复及配置变更后调用 -`plugin.reconcile`，并每 30 秒补偿一次。处理器必须幂等；调用可能重复，配置通知会合并，不表示逐条账号事件。 -每个实例串行执行，不同实例相互独立；调用限时 30 秒，返回值不消费，调用失败后等待 5 秒重试。维护失败不回滚已经提交的资源， +`plugin.reconcile`，每轮完成后等待 30 秒补偿。处理器必须幂等；调用可能重复，配置通知会合并，不表示逐条账号事件。 +每个实例串行执行，不同实例相互独立；处理预算按[会话声明](../README.md#声明调用预算)生效，返回值不消费，调用失败后等待 5 秒重试。维护失败不回滚已经提交的资源， 下一次对账继续补齐。只读校验、准备候选与 CLI 帮助不会启动维护；停用、替换和宿主关闭时取消旧任务 维护与其他处理器共用以下资源方法： diff --git a/backend/crates/gateway-plugin/sdk/docs/manifest.md b/backend/crates/gateway-plugin/sdk/docs/manifest.md index 22d2c8e34..aa8903c0d 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` | `5` | 插件进程 RPC 封装，不是 OpenAI 或 xAI 的业务协议 | +| `package.protocolVersion` / `protocol_versions` | `6` | 插件进程 RPC 封装，不是 OpenAI 或 xAI 的业务协议 | | `contributes..version` / `capabilities[].versions` | 见[扩展项简写](#扩展项简写) | 指定扩展能力的行为合同 | | 宿主声明的 `schema_version` | `2` | 宿主兼容声明自身的格式 | @@ -34,7 +34,7 @@ 清单不接受 `permissions` 字段；插件按[完整信任](#完整信任)运行 -RPC v5 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建，并由 CLI 生成包含协议版本的安装包 +RPC v6 的帧布局、编码与会话生命周期见[通信边界](../README.md#通信边界)。插件使用 SDK 构建，并由 CLI 生成包含协议版本的安装包 安装包清单校验与已安装元数据读取是两个边界：清单拒绝未知字段，并检查必需字段、类型和版本； 数据库元数据的读取规则见[制品与发布](../../../../../docs/architecture.md#制品配置与发布)，不改变清单或 RPC 合同。 diff --git a/backend/crates/gateway-plugin/sdk/src/budget.rs b/backend/crates/gateway-plugin/sdk/src/budget.rs new file mode 100644 index 000000000..adb529139 --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/src/budget.rs @@ -0,0 +1,119 @@ +//! 按调用阶段解析有限处理预算，集中维护协议默认值与插件声明 + +use std::{ + collections::BTreeMap, + time::{Duration, Instant}, +}; + +use crate::Stage; + +/// 按阶段声明的毫秒预算；未声明的阶段采用协议默认值 +pub type CallBudgets = BTreeMap; + +const STAGES: [Stage; 18] = [ + Stage::Registration, + Stage::Configuration, + Stage::Authentication, + Stage::Routing, + Stage::Scheduling, + Stage::Retry, + Stage::Http, + Stage::Service, + Stage::WebSocket, + Stage::Request, + Stage::Attempt, + Stage::Outbound, + Stage::Upstream, + Stage::Observation, + Stage::Management, + Stage::CommandLine, + Stage::PublicManagement, + Stage::Maintenance, +]; + +/// 实际预算来自协议默认值或插件显式声明 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CallBudgetSource { + Default, + Plugin, +} + +/// 已校验且不可变的会话预算，不允许构造缺少阶段的解析结果 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedCallBudgets(BTreeMap); + +impl ResolvedCallBudgets { + /// 按协议默认值合并插件声明并冻结所有阶段的有效预算 + /// + /// # Errors + /// + /// 声明值为零或平台计时器无法表示时失败 + pub fn resolve(declarations: &CallBudgets) -> Result { + validate_values(declarations)?; + let resolved = STAGES + .into_iter() + .map(|stage| { + let (milliseconds, source) = match declarations.get(&stage) { + Some(&requested) => (requested, CallBudgetSource::Plugin), + None => (default_budget_ms(stage), CallBudgetSource::Default), + }; + (stage, (Duration::from_millis(milliseconds), source)) + }) + .collect(); + Ok(Self(resolved)) + } + + /// 读取当前会话在指定阶段的处理预算 + #[must_use] + pub fn budget(&self, stage: Stage) -> Duration { + self.0[&stage].0 + } + + /// 读取预算来源，用于调用诊断 + #[must_use] + pub fn source(&self, stage: Stage) -> CallBudgetSource { + self.0[&stage].1 + } +} + +/// 预算声明错误只暴露阶段和数值，不包含插件配置或业务数据 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum CallBudgetError { + #[error("plugin call budget for {stage:?} is invalid: {milliseconds} ms")] + Invalid { stage: Stage, milliseconds: u64 }, +} + +pub(crate) fn validate_values(values: &CallBudgets) -> Result<(), CallBudgetError> { + for (&stage, &milliseconds) in values { + if milliseconds == 0 + || Instant::now() + .checked_add(Duration::from_millis(milliseconds)) + .is_none() + { + return Err(CallBudgetError::Invalid { + stage, + milliseconds, + }); + } + } + Ok(()) +} + +const fn default_budget_ms(stage: Stage) -> u64 { + match stage { + Stage::Registration => 5_000, + Stage::Configuration | Stage::Authentication => 30_000, + Stage::Routing | Stage::Scheduling | Stage::Retry | Stage::Observation => 2_000, + Stage::Http + | Stage::Service + | Stage::WebSocket + | Stage::Request + | Stage::Attempt + | Stage::Outbound + | Stage::Upstream + | Stage::Management + | Stage::CommandLine + | Stage::PublicManagement + | Stage::Maintenance => 120_000, + } +} 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 29b1ba2d7..7929b9ee5 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/config.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/config.rs @@ -17,7 +17,6 @@ pub struct SessionConfig { pub maximum_callbacks: usize, pub maximum_buffered_stream_chunks: usize, pub handshake_timeout: Duration, - pub maximum_call_timeout: Duration, } impl Default for SessionConfig { @@ -29,7 +28,6 @@ impl Default for SessionConfig { maximum_callbacks: 64, maximum_buffered_stream_chunks: 1_024, handshake_timeout: Duration::from_secs(5), - maximum_call_timeout: Duration::from_secs(120), } } } @@ -45,7 +43,6 @@ impl SessionConfig { || self.maximum_buffered_stream_chunks == 0 || self.maximum_buffered_stream_chunks > MAXIMUM_BUFFERED_STREAM_CHUNKS || self.handshake_timeout.is_zero() - || self.maximum_call_timeout.is_zero() { return Err(SessionError::Configuration); } 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 fbe4408ba..3627b22ef 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/driver.rs @@ -11,7 +11,8 @@ use tokio::{ }; use crate::{ - CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, PluginFault, + CallBudgets, CallContext, ErrorCode, Frame, FrameError, Handshake, Message, PROTOCOL_VERSION, + PluginFault, ResolvedCallBudgets, }; use super::super::frame::{EncodedFrame, read_frame, write_frame}; @@ -31,6 +32,7 @@ pub struct PluginSession { writer: W, handshake: Handshake, config: SessionConfig, + call_budgets: CallBudgets, } impl PluginSession @@ -70,6 +72,7 @@ where writer, handshake, config, + call_budgets: CallBudgets::new(), }) } @@ -78,6 +81,18 @@ where &self.handshake } + /// 在 Ready 前声明所需预算，可根据握手中的实例配置计算 + /// 未声明的阶段采用协议默认值，同一会话的预算在 run 时冻结 + /// + /// # Errors + /// + /// 预算为零或平台计时器无法表示时失败，原声明保持不变 + pub fn set_call_budgets(&mut self, budgets: CallBudgets) -> Result<(), SessionError> { + crate::budget::validate_values(&budgets)?; + self.call_budgets = budgets; + Ok(()) + } + /// 发送 Ready 并持续处理调用、回调、流控、取消与关闭 /// /// # Errors @@ -87,6 +102,7 @@ where where H: PluginHandler, { + let call_budgets = ResolvedCallBudgets::resolve(&self.call_budgets)?; let data_capacity = self .config .maximum_calls @@ -117,6 +133,7 @@ where &Frame::control(Message::Ready { protocol_version: PROTOCOL_VERSION, incarnation: self.handshake.incarnation.clone(), + call_budgets: self.call_budgets, }), ) .await?; @@ -124,6 +141,7 @@ where let (completed, completions) = mpsc::channel(self.config.maximum_calls); let mut driver = SessionDriver { config: self.config, + call_budgets, instance_id: self.handshake.instance_id, generation: self.handshake.generation, incarnation: self.handshake.incarnation, @@ -200,6 +218,7 @@ impl Drop for CompletionGuard { struct SessionDriver { config: SessionConfig, + call_budgets: ResolvedCallBudgets, instance_id: String, generation: u64, incarnation: String, @@ -323,7 +342,7 @@ impl SessionDriver { || context.generation != self.generation || context.incarnation != self.incarnation || context.timeout_ms == 0 - || Duration::from_millis(context.timeout_ms) > self.config.maximum_call_timeout + || Duration::from_millis(context.timeout_ms) > self.call_budgets.budget(context.stage) || method.is_empty() || method.len() > 128 { diff --git a/backend/crates/gateway-plugin/sdk/src/client/session/error.rs b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs index 834027fc8..2308b3a0a 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/error.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/error.rs @@ -1,12 +1,14 @@ //! 插件侧传输和会话错误及稳定业务错误投影 -use crate::{ErrorCode, FrameError, PluginFault}; +use crate::{CallBudgetError, ErrorCode, FrameError, PluginFault}; /// 插件侧传输和会话错误 #[derive(Debug, thiserror::Error)] pub enum SessionError { #[error("plugin session configuration is invalid")] Configuration, + #[error(transparent)] + CallBudget(#[from] CallBudgetError), #[error("plugin handshake is invalid")] Handshake, #[error("plugin protocol is invalid")] @@ -36,7 +38,11 @@ impl SessionError { Self::Timeout => PluginFault::new(ErrorCode::Timeout, "host callback timed out"), Self::Cancelled => PluginFault::new(ErrorCode::Cancelled, "parent call was cancelled"), Self::Capacity => PluginFault::new(ErrorCode::Capacity, "plugin capacity is exhausted"), - Self::Configuration | Self::Protocol | Self::Handshake | Self::Frame(_) => { + Self::Configuration + | Self::CallBudget(_) + | Self::Protocol + | Self::Handshake + | Self::Frame(_) => { PluginFault::new(ErrorCode::InvalidInput, "plugin session input is invalid") } Self::Closed | Self::HandlerStopped => { 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 2e9fec9a8..8261737f8 100644 --- a/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs +++ b/backend/crates/gateway-plugin/sdk/src/client/session/registry.rs @@ -35,6 +35,13 @@ pub struct HostClient { } impl HostClient { + /// 读取父调用剩余处理预算，包含 next 暂停；资源流接管寿命后返回 None + /// 外部请求期限、网络操作期限和取消信号仍独立生效 + #[must_use] + pub fn remaining_budget(&self) -> Option { + self.deadline.remaining() + } + pub(crate) fn maximum_stream_chunk_bytes(&self) -> usize { let window = self.credits.maximum_bytes(); if window == 0 { diff --git a/backend/crates/gateway-plugin/sdk/src/lib.rs b/backend/crates/gateway-plugin/sdk/src/lib.rs index 3326b64e5..21853b35b 100644 --- a/backend/crates/gateway-plugin/sdk/src/lib.rs +++ b/backend/crates/gateway-plugin/sdk/src/lib.rs @@ -1,5 +1,6 @@ //! 实验性插件线协议；不依赖网关领域或基础设施类型 +mod budget; pub mod call; mod capability; #[cfg(feature = "io")] @@ -9,6 +10,7 @@ mod error; mod manifest; mod message; +pub use budget::{CallBudgetError, CallBudgetSource, CallBudgets, ResolvedCallBudgets}; pub use capability::{Capability, ContributionDeclaration, Contributions, FailurePolicy, Stage}; pub use context::{CallContext, Handshake, RequestContext}; pub use error::{ErrorCode, PluginFault, SendState}; diff --git a/backend/crates/gateway-plugin/sdk/src/message.rs b/backend/crates/gateway-plugin/sdk/src/message.rs index 223847c54..ad3531185 100644 --- a/backend/crates/gateway-plugin/sdk/src/message.rs +++ b/backend/crates/gateway-plugin/sdk/src/message.rs @@ -2,9 +2,9 @@ use serde::{Deserialize, Serialize}; -use crate::{CallContext, Handshake, PluginFault}; +use crate::{CallBudgets, CallContext, Handshake, PluginFault}; -pub const PROTOCOL_VERSION: u32 = 5; +pub const PROTOCOL_VERSION: u32 = 6; /// 元数据与二进制载荷分开，流分块不经过 JSON/base64 #[derive(Clone, PartialEq, Serialize, Deserialize)] @@ -16,6 +16,7 @@ pub enum Message { Ready { protocol_version: u32, incarnation: String, + call_budgets: CallBudgets, }, Call { id: u64, diff --git a/backend/crates/gateway-plugin/sdk/tests/budget.rs b/backend/crates/gateway-plugin/sdk/tests/budget.rs new file mode 100644 index 000000000..3cd3eaa8a --- /dev/null +++ b/backend/crates/gateway-plugin/sdk/tests/budget.rs @@ -0,0 +1,49 @@ +//! 验证场景默认预算、插件声明隔离与无效输入边界 + +use std::time::Duration; + +use gateway_plugin_sdk::{ + CallBudgetError, CallBudgetSource, CallBudgets, ResolvedCallBudgets, Stage, +}; + +#[test] +fn explicit_budgets_override_only_their_stage_without_a_global_ceiling() { + let resolved = ResolvedCallBudgets::resolve(&CallBudgets::from([ + (Stage::Maintenance, 1_800_000), + (Stage::CommandLine, 600_000), + (Stage::Management, 250), + ])) + .unwrap(); + for (stage, millis, source) in [ + (Stage::Maintenance, 1_800_000, CallBudgetSource::Plugin), + (Stage::CommandLine, 600_000, CallBudgetSource::Plugin), + (Stage::Management, 250, CallBudgetSource::Plugin), + (Stage::Registration, 5_000, CallBudgetSource::Default), + (Stage::Authentication, 30_000, CallBudgetSource::Default), + (Stage::Retry, 2_000, CallBudgetSource::Default), + (Stage::Request, 120_000, CallBudgetSource::Default), + ] { + assert_eq!(resolved.budget(stage), Duration::from_millis(millis)); + assert_eq!(resolved.source(stage), source); + } + assert_eq!( + ResolvedCallBudgets::resolve(&CallBudgets::new()) + .unwrap() + .budget(Stage::Maintenance), + Duration::from_secs(120) + ); +} + +#[test] +fn invalid_budget_values_fail_before_use() { + assert!(matches!( + ResolvedCallBudgets::resolve(&CallBudgets::from([(Stage::Maintenance, 0)])), + Err(CallBudgetError::Invalid { + stage: Stage::Maintenance, + milliseconds: 0 + }) + )); + assert!( + serde_json::from_str::(r#"{"maintenance":18446744073709551616}"#).is_err() + ); +} 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 18c51506a..a4dd20f56 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/plugin/mod.rs @@ -440,7 +440,6 @@ async fn start_session( maximum_callbacks: 4, maximum_buffered_stream_chunks: 16, handshake_timeout: Duration::from_secs(1), - maximum_call_timeout: Duration::from_secs(2), }, ) .await? @@ -473,6 +472,7 @@ async fn start_session( message: Message::Ready { protocol_version: PROTOCOL_VERSION, ref incarnation, + .. }, ref payload, } if incarnation == "test-incarnation" && payload.is_empty() 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 5627c648b..8f1c5c040 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/driver.rs @@ -1,6 +1,86 @@ //! 会话调用分派、并发帧读取与处理器响应边界测试 use super::*; +use gateway_plugin_sdk::CallBudgets; + +#[tokio::test(start_paused = true)] +async fn declared_long_budgets_survive_callbacks_and_expose_remaining_time() { + for budget_ms in [600_000, 1_800_000] { + let (host, plugin) = tokio::io::duplex(16_384); + let (plugin_reader, plugin_writer) = tokio::io::split(plugin); + let task = tokio::spawn(async move { + let mut session = + PluginSession::accept(plugin_reader, plugin_writer, SessionConfig::default()) + .await?; + let declared = session.handshake().configuration["maintenance_ms"] + .as_u64() + .unwrap(); + session.set_call_budgets(CallBudgets::from([(Stage::Maintenance, declared)]))?; + session.run(TestHandler::default()).await + }); + let (reader, writer) = tokio::io::split(host); + let mut host = HostPeer { reader, writer }; + let mut hello = handshake(); + hello.configuration = json!({"maintenance_ms": budget_ms}); + send_control(&mut host, Message::Hello { handshake: hello }).await; + let Message::Ready { call_budgets, .. } = receive(&mut host).await.message else { + panic!("expected Ready") + }; + assert_eq!(call_budgets[&Stage::Maintenance], budget_ms); + let mut call_context = context(1, Duration::from_millis(budget_ms)); + call_context.stage = Stage::Maintenance; + send_control( + &mut host, + Message::Call { + id: 1, + method: "budget".into(), + context: call_context, + params: json!({}), + }, + ) + .await; + let Message::Callback { + id, parent_id: 1, .. + } = receive(&mut host).await.message + else { + panic!("expected callback") + }; + tokio::time::advance(Duration::from_secs(125)).await; + send_control( + &mut host, + Message::Result { + id, + result: json!({}), + }, + ) + .await; + let Message::Result { id: 1, result } = receive(&mut host).await.message else { + panic!("long call must remain active") + }; + assert_eq!(result["initial_ms"], budget_ms); + let remaining = result["remaining_ms"].as_u64().unwrap(); + assert!(remaining <= budget_ms - 125_000 && remaining >= budget_ms - 126_000); + shutdown(&mut host, task).await; + } +} + +#[tokio::test] +async fn calls_cannot_exceed_the_budget_negotiated_for_their_stage() { + let (mut host, task) = start_session(TestHandler::default()).await; + let mut context = context(1, Duration::from_secs(3)); + context.stage = Stage::Retry; + send_control( + &mut host, + Message::Call { + id: 1, + method: "echo".into(), + context, + params: json!({}), + }, + ) + .await; + assert!(matches!(task.await.unwrap(), Err(SessionError::Protocol))); +} #[test] fn middleware_plugin_rejects_declarations_without_matching_handlers() { 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 f0292a39f..571001cb7 100644 --- a/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs +++ b/backend/crates/gateway-plugin/sdk/tests/client/session/mod.rs @@ -58,6 +58,19 @@ impl PluginHandler for TestHandler { Box::pin(async move { match call.method.as_str() { "echo" => Ok(CallReply::unary(call.params, call.payload)), + "budget" => { + call.host + .call("host.log", json!({}), vec![]) + .await + .map_err(SessionError::into_plugin_fault)?; + Ok(CallReply::unary( + json!({ + "initial_ms": call.context.timeout_ms, + "remaining_ms": call.host.remaining_budget().map(|value| value.as_millis()), + }), + vec![], + )) + } "key_facts" => { let result = call .host @@ -263,7 +276,6 @@ async fn start_session_with_capacity( maximum_callbacks: 4, maximum_buffered_stream_chunks: 16, handshake_timeout: Duration::from_secs(1), - maximum_call_timeout: Duration::from_secs(2), }, ) .await @@ -296,6 +308,7 @@ async fn start_configured_session( message: Message::Ready { protocol_version: PROTOCOL_VERSION, ref incarnation, + .. }, ref payload, } if incarnation == "test-incarnation" && payload.is_empty() diff --git a/backend/crates/gateway-plugin/sdk/tests/main.rs b/backend/crates/gateway-plugin/sdk/tests/main.rs index 1be68c88f..a7a33acd5 100644 --- a/backend/crates/gateway-plugin/sdk/tests/main.rs +++ b/backend/crates/gateway-plugin/sdk/tests/main.rs @@ -1,5 +1,6 @@ //! 插件 SDK 线协议、清单与客户端的测试入口 +mod budget; mod call; #[cfg(feature = "io")] mod client; diff --git a/docs/architecture.md b/docs/architecture.md index 447529a93..495fbe13a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -212,6 +212,9 @@ flowchart LR #### 调用资源与状态 +- **调用预算由阶段合同统一解释**：SDK 定义场景默认值，插件在 `Ready` 声明所需预算，Runtime 与 SDK 解析并冻结到会话；各入口按实例和阶段读取，宿主不提供部署侧业务时长上限。 + 注册在进程启动与握手完成后使用自己的预算；CLI 的执行与结果提交共用一次期限，重试委托链共享匹配实例中最长的预算并受请求剩余期限约束，单次调用仍受自身声明限制。 + 业务预算不延长进程启动与握手、排队、取消确认、关闭和外部请求期限；声明与默认值见 [SDK 调用预算](../backend/crates/gateway-plugin/sdk/README.md#声明调用预算) - **资源跟随调用**：账号、HTTP、模型、状态和日志回调绑定有效父调用，取消随调用传播。HTTP 与模型 request/attempt 中间件在首个流式响应头交付后，由正文 owner 管理取消与终态，不沿用初始处理期限；网络操作与模型执行各自计时。 Runtime 与 SDK 的中间件处理预算在宿主 `next` 等待期间暂停；Runtime 同时归还执行名额，返回后取得名额再以剩余处理预算恢复原调用。 执行调度由 Runtime 统一负责，准入和恢复等待使用独立期限；初次排队、总在途上下文和普通回调各有上限；`next` 回调按在途上限独立计数。 @@ -220,7 +223,7 @@ flowchart LR `frontend_authentication` 使用 principal 到 Key 的映射；公开登录回调的一次性票据用于关联流程，不限制插件回调能力。接口合同见 [SDK 回调](../backend/crates/gateway-plugin/sdk/docs/capabilities.md#宿主资源回调) - **观察合同统一**：`observer` 通过一个类型化 RPC 入口接收完成与实际上游 WebSocket 事件，发布时按事件类型编译独立作用域。完成事件直接携带终态、用量、费用和耗时，不按能力声明裁剪字段；异步投递的并发、队列和载荷有界，观察失败不改变业务结果 - **维护只作用于已发布实例**：`maintenance` 由 Host 监督的 Runtime worker 调用，启用、恢复、配置发布和周期补偿共用对账入口。 - 发布通知只用于唤醒，事实始终来自当前快照与数据库；同实例串行，通知合并，停用与版本替换取消旧任务。 + 发布通知只用于唤醒，事实始终来自当前快照与数据库；按发布层复用的实例身份保留在途任务，无关实例变更不重启维护。同实例串行，通知合并；自身停用、替换或恢复时取消旧任务，等取消确认或进程退出后启动替代任务。 自有分组和 Key 按实例、资源类型与稳定资源键登记；Admin 组织用例，Store 在同一事务复核当前启用 revision、 制品接受事实与归属，并复用原生资源写入、审计和配置发布。无变化不递增 revision；插件不能接管管理员资源， 成员变更不覆盖账号的其他分组。停用保留资源；删除实例只解除归属