feat(plugin): add managed upstream adapters for built-in providers

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-28 23:23:41 +08:00
co-authored by Codex
parent a5a844af16
commit faee3b2642
75 changed files with 5732 additions and 299 deletions
+2
View File
@@ -1776,6 +1776,7 @@ dependencies = [
"thiserror",
"tokio",
"tokio-rustls",
"tokio-tungstenite 0.29.0",
"tower-service",
"tracing",
"tracing-appender",
@@ -1813,6 +1814,7 @@ dependencies = [
"tempfile",
"thiserror",
"tokio",
"tokio-tungstenite 0.29.0",
"tracing",
"url",
"uuid",
@@ -61,7 +61,7 @@ pub struct NativeContinuationPin {
provider: ProviderKind,
account: ProviderAccountId,
scope: NativeContinuationScope,
session_state: Option<ProviderSessionState>,
session_state: Option<Box<ProviderSessionState>>,
}
impl NativeContinuationPin {
@@ -94,7 +94,7 @@ impl NativeContinuationPin {
/// 附着仅由对应 Provider 解释的不透明会话状态。
#[must_use]
pub fn with_session_state(mut self, state: ProviderSessionState) -> Self {
self.session_state = Some(state);
self.session_state = Some(Box::new(state));
self
}
@@ -136,8 +136,8 @@ impl NativeContinuationPin {
/// 返回与本 pin 同账号绑定的 Provider 私有会话状态。
#[must_use]
pub const fn session_state(&self) -> Option<&ProviderSessionState> {
self.session_state.as_ref()
pub fn session_state(&self) -> Option<&ProviderSessionState> {
self.session_state.as_deref()
}
/// 校验本次 route/account 选择没有破坏 native pin。
@@ -59,6 +59,7 @@ pub(super) struct CoordinationExtensions {
execution_effects: Option<Arc<ExecutionEffects>>,
execution_effects_baseline: usize,
middleware: Option<super::middleware::FrozenMiddlewarePlan>,
upstream_adapters: Option<super::upstream_adapter::FrozenUpstreamAdapterPlan>,
account_group_ids: Arc<[crate::account::scope::AccountGroupId]>,
endpoint: String,
client_transport: super::execution::ClientTransport,
@@ -66,6 +67,14 @@ pub(super) struct CoordinationExtensions {
}
impl CoordinationExtensions {
pub(super) fn with_upstream_adapters(
mut self,
plan: Option<super::upstream_adapter::FrozenUpstreamAdapterPlan>,
) -> Self {
self.upstream_adapters = plan;
self
}
pub(super) fn with_response_control(
mut self,
control: Option<super::response_control::ResponseControl>,
@@ -86,6 +95,7 @@ impl CoordinationExtensions {
execution_effects: None,
execution_effects_baseline: 0,
middleware: None,
upstream_adapters: None,
account_group_ids: Arc::from([]),
endpoint: String::new(),
client_transport: super::execution::ClientTransport::InternalProbe,
@@ -247,6 +257,7 @@ where
execution_effects,
execution_effects_baseline,
middleware,
upstream_adapters,
account_group_ids,
endpoint,
client_transport,
@@ -304,6 +315,7 @@ where
.unwrap_or(Duration::ZERO),
)
.fuse(),
requested_model: request.requested_model.clone(),
pending_request: Some(request),
request_persisted: false,
response_control,
@@ -313,6 +325,7 @@ where
execution_effects,
execution_effects_baseline,
middleware,
upstream_adapters,
account_group_ids,
endpoint,
client_transport,
@@ -420,6 +433,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
deadline_timer: Fuse<Delay>,
pending_request: Option<NewModelRequest>,
requested_model: Option<crate::routing::PublicModelId>,
request_persisted: bool,
response_control: Option<super::response_control::ResponseControl>,
operation: Operation,
@@ -428,6 +442,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
execution_effects: Option<Arc<ExecutionEffects>>,
execution_effects_baseline: usize,
middleware: Option<super::middleware::FrozenMiddlewarePlan>,
upstream_adapters: Option<super::upstream_adapter::FrozenUpstreamAdapterPlan>,
account_group_ids: Arc<[crate::account::scope::AccountGroupId]>,
endpoint: String,
client_transport: super::execution::ClientTransport,
@@ -756,7 +771,16 @@ where
provider,
account,
)
.with_scope(NativeContinuationScope::Persisted)
.with_scope(
if state
.extension_owner()
.is_some_and(|owner| owner.connection_local)
{
NativeContinuationScope::ConnectionLocal
} else {
NativeContinuationScope::Persisted
},
)
.with_session_state(state.clone()),
)
}
@@ -1019,6 +1043,8 @@ where
.with_timing_started_at(self.observation.timing_started_at)
.with_request_policy(self.request_policy.clone())
.with_execution_effects(self.execution_effects.as_ref().map(Arc::clone))
.with_upstream_adapters(self.upstream_adapters.clone())
.with_requested_model(self.requested_model.clone())
.with_middleware(
self.middleware.clone(),
Arc::clone(&self.account_group_ids),
@@ -973,6 +973,11 @@ impl DefaultExecutionService {
.iter()
.map(|group| group.id().clone())
.collect::<Vec<_>>();
let upstream_adapters = request
.client
.snapshot
.extensions()
.and_then(|set| set.upstream_adapters());
let middleware = self.middlewares.as_ref().and_then(|middlewares| {
let generation = request.client.snapshot.extensions()?.clone();
middlewares.resolve(&generation)
@@ -1167,6 +1172,7 @@ impl DefaultExecutionService {
CoordinationExtensions::new(continuation, request_observation.clone())
.with_response_control(authorization.response_control.clone())
.with_request_policy(request_policy)
.with_upstream_adapters(upstream_adapters)
.with_execution_effects(
execution_effects,
authorization.execution_effects_baseline,
@@ -15,6 +15,7 @@ pub mod policy;
pub mod probe;
pub mod provider;
pub mod response_control;
pub mod upstream_adapter;
pub use coordinator::{AttemptCoordinator, ResponseExecutionSession};
@@ -338,6 +339,8 @@ pub struct RequestAttemptContext {
request_policy: Option<policy::RequestPolicyContext>,
execution_effects: Option<Arc<nested::ExecutionEffects>>,
middleware: Option<middleware::FrozenMiddlewarePlan>,
upstream_adapters: Option<upstream_adapter::FrozenUpstreamAdapterPlan>,
requested_model: Option<PublicModelId>,
account_group_ids: Arc<[crate::account::scope::AccountGroupId]>,
endpoint: String,
client_transport: execution::ClientTransport,
@@ -345,6 +348,21 @@ pub struct RequestAttemptContext {
}
impl RequestAttemptContext {
#[must_use]
pub fn with_requested_model(mut self, model: Option<PublicModelId>) -> Self {
self.requested_model = model;
self
}
#[must_use]
pub fn with_upstream_adapters(
mut self,
plan: Option<upstream_adapter::FrozenUpstreamAdapterPlan>,
) -> Self {
self.upstream_adapters = plan;
self
}
#[must_use]
pub fn with_response_control(
mut self,
@@ -401,6 +419,8 @@ impl RequestAttemptContext {
request_policy: None,
execution_effects: None,
middleware: None,
upstream_adapters: None,
requested_model: None,
account_group_ids: Arc::from([]),
endpoint: String::new(),
client_transport: execution::ClientTransport::InternalProbe,
@@ -539,6 +559,38 @@ pub struct AttemptContext {
}
impl AttemptContext {
#[must_use]
pub fn execution_effects(&self) -> Option<Arc<nested::ExecutionEffects>> {
self.request.execution_effects.as_ref().map(Arc::clone)
}
#[must_use]
pub const fn requested_model(&self) -> Option<&PublicModelId> {
self.request.requested_model.as_ref()
}
/// 从本次请求冻结的发布代次选择上游适配器;不建立连接、不重新选号。
pub fn upstream_adapter(
&self,
provider: &ProviderKind,
model: &UpstreamModelId,
) -> Result<Option<Arc<dyn upstream_adapter::UpstreamAdapter>>, ProviderError> {
self.request
.upstream_adapters
.as_ref()
.map_or(Ok(None), |plan| plan.select(self, provider, model))
}
#[must_use]
pub fn account_group_ids(&self) -> &[crate::account::scope::AccountGroupId] {
&self.request.account_group_ids
}
#[must_use]
pub const fn client_transport(&self) -> execution::ClientTransport {
self.request.client_transport
}
#[must_use]
pub fn response_control(&self) -> Option<&response_control::ResponseControl> {
self.request.response_control.as_ref()
@@ -22,7 +22,9 @@ use crate::engine::middleware::{
MiddlewareHeader, MiddlewareNext, MiddlewareRequest, MiddlewareResponse,
};
use crate::error::{PreDeliveryRetry, ProviderError, ProviderErrorKind};
use crate::event::{EventSequenceValidator, GatewayEvent, ProtocolWireEvent, ProviderEvent};
use crate::event::{
EventSequenceError, EventSequenceValidator, GatewayEvent, ProtocolWireEvent, ProviderEvent,
};
use crate::identity::ProviderKind;
use crate::operation::Operation;
use crate::policy::ClientApiKeyId;
@@ -216,7 +218,7 @@ pub struct ProviderStream {
_lease: Box<dyn ResourceLease>,
native_response_translator: Option<Box<dyn NativeResponseTranslator>>,
account_feedback: Option<ProviderStreamAccountFeedback>,
validator: EventSequenceValidator,
validator: Result<EventSequenceValidator, EventSequenceError>,
strict_canonical_seen: bool,
terminated: bool,
}
@@ -323,7 +325,7 @@ impl ProviderStream {
_lease: Box::new(lease),
native_response_translator: None,
account_feedback: None,
validator: EventSequenceValidator::new(),
validator: Ok(EventSequenceValidator::new()),
strict_canonical_seen: false,
terminated: false,
}
@@ -857,24 +859,26 @@ impl Stream for ProviderStream {
match this.events.as_mut().poll_next(context) {
Poll::Pending => Poll::Pending,
Poll::Ready(Some(Ok(event))) => {
// 带 wire 的 canonical facts 只是旁路观测:wire 才是客户端协议的
// 权威表达。只有 canonical-only Provider 输出需要以状态机作为交付条件。
if event.wire_event().is_none() && !event.canonical_facts().is_empty() {
this.strict_canonical_seen = true;
for fact in event.canonical_facts() {
if this.validator.observe(fact).is_err() {
this.terminated = true;
let error = ProviderError::new(
ProviderErrorKind::Protocol,
UpstreamSendState::Sent,
);
if let Some(feedback) = this.account_feedback.as_mut() {
feedback.report_failure(&error);
}
return Poll::Ready(Some(Err(error)));
}
// 纯 wire 流的 facts 仍只是旁路观测;一旦交付 canonical-only
// 事件,就必须校验完整事实序列,不能漏掉带 wire 的开始或终态。
this.strict_canonical_seen |=
event.wire_event().is_none() && !event.canonical_facts().is_empty();
for fact in event.canonical_facts() {
if let Ok(validator) = &mut this.validator
&& let Err(error) = validator.observe(fact)
{
this.validator = Err(error);
}
}
if this.strict_canonical_seen && this.validator.is_err() {
this.terminated = true;
let error =
ProviderError::new(ProviderErrorKind::Protocol, UpstreamSendState::Sent);
if let Some(feedback) = this.account_feedback.as_mut() {
feedback.report_failure(&error);
}
return Poll::Ready(Some(Err(error)));
}
if let Some(feedback) = this.account_feedback.as_mut() {
feedback.observe(&event);
}
@@ -890,7 +894,10 @@ impl Stream for ProviderStream {
Poll::Ready(None) => {
this.terminated = true;
let validation = if this.strict_canonical_seen {
this.validator.finish()
this.validator
.as_ref()
.map_err(Clone::clone)
.and_then(EventSequenceValidator::finish)
} else {
Ok(())
};
@@ -0,0 +1,99 @@
//! 内置 Provider 已选账号上的受管上游适配端口。
//!
//! 适配器不注册 Provider 或选择账号。Provider 持有租约、解释凭据并计算费用;
//! Runtime 只实现适配器与受管网络之间的连接,返回值继续经过原生事件和结算链。
use std::{fmt, sync::Arc};
use futures::future::BoxFuture;
use crate::{
account::{CredentialRevision, OutboundProxy, ProviderAccountId},
error::ProviderError,
identity::ProviderKind,
metering::{CalculatedCost, Usage},
operation::Operation,
routing::UpstreamModelId,
runtime::extensions::ExtensionSetReference,
};
use super::{
AttemptContext,
middleware::MiddlewareHeader,
provider::{EventStream, ProviderCallMetadata},
};
/// 发布代次内冻结的适配器选择计划;匹配只读宿主身份,不调用插件进程。
pub trait UpstreamAdapterPlan: fmt::Debug + Send + Sync {
fn select(
&self,
context: &AttemptContext,
provider: &ProviderKind,
model: &UpstreamModelId,
) -> Result<Option<Arc<dyn UpstreamAdapter>>, ProviderError>;
}
/// 请求同时持有计划及发布集合,防止旧集合先排空进程而中断在途调用。
#[derive(Debug, Clone)]
pub struct FrozenUpstreamAdapterPlan {
plan: Arc<dyn UpstreamAdapterPlan>,
_generation: ExtensionSetReference,
}
impl FrozenUpstreamAdapterPlan {
#[must_use]
pub const fn new(
plan: Arc<dyn UpstreamAdapterPlan>,
generation: ExtensionSetReference,
) -> Self {
Self {
plan,
_generation: generation,
}
}
pub(super) fn select(
&self,
context: &AttemptContext,
provider: &ProviderKind,
model: &UpstreamModelId,
) -> Result<Option<Arc<dyn UpstreamAdapter>>, ProviderError> {
self.plan.select(context, provider, model)
}
}
/// 为本次调用选定的适配器,持有对应代次的运行资源。
pub trait UpstreamAdapter: fmt::Debug + Send + Sync {
/// 返回冷流。首次消费前不得调用插件、建立业务连接或发送请求。
fn execute(self: Arc<Self>, invocation: UpstreamAdapterInvocation) -> EventStream;
/// 注册时声明的上游传输,用于宿主登记 attempt。
fn transport(&self) -> &str;
}
/// 一次执行的宿主事实;账号及 Provider 归属不可由插件返回值覆盖。
pub struct UpstreamAdapterInvocation {
pub operation: Operation,
pub headers: Vec<MiddlewareHeader>,
pub context: AttemptContext,
pub metadata: ProviderCallMetadata,
pub account: Arc<dyn UpstreamAccountConnection>,
}
/// Provider 针对已选账号租约提供的凭据与反馈端口。
///
/// 实现必须持有原账号租约。鉴权头只交给宿主网络边界,不能投影到插件 RPC。
/// 目标授权由适配器注册与受管网络检查,凭据格式及刷新协议仍归 Provider。
pub trait UpstreamAccountConnection: Send + Sync {
fn account_id(&self) -> &ProviderAccountId;
fn credential_revision(&self) -> CredentialRevision;
fn authentication_kind(&self) -> &str;
fn outbound_proxy(&self) -> Option<&OutboundProxy>;
fn authorization(&self) -> Result<Vec<MiddlewareHeader>, ProviderError>;
/// 复用 Provider 的价格来源;未知用量、模型或档位不构造零费用。
fn calculate_cost(&self, service_tier: Option<&str>, usage: &Usage) -> Option<CalculatedCost>;
/// 失败的刷新、账号状态和冷却反馈仍由 Provider 处理,不由 Runtime 写账号。
fn record_failure(&self, error: ProviderError) -> BoxFuture<'_, ProviderError>;
}
@@ -120,6 +120,20 @@ impl CapabilityRequirements {
pub struct ProviderSessionState {
provider: String,
payload: Map<String, Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
extension_owner: Option<Box<ExtensionSessionOwner>>,
}
/// 受管扩展会话的宿主归属,不能把插件私有状态交给原生 Provider 解码。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ExtensionSessionOwner {
pub instance_id: String,
pub contribution_id: String,
pub adapter_id: String,
pub generation: u64,
pub incarnation: String,
pub connection_local: bool,
}
impl ProviderSessionState {
@@ -136,7 +150,11 @@ impl ProviderSessionState {
validate_text(&provider, 64, true, None).map_err(|_| OperationError::EmptyField {
field: "provider_session_state provider",
})?;
Ok(Self { provider, payload })
Ok(Self {
provider,
payload,
extension_owner: None,
})
}
#[must_use]
@@ -148,6 +166,17 @@ impl ProviderSessionState {
pub const fn payload(&self) -> &Map<String, Value> {
&self.payload
}
#[must_use]
pub fn with_extension_owner(mut self, owner: ExtensionSessionOwner) -> Self {
self.extension_owner = Some(Box::new(owner));
self
}
#[must_use]
pub fn extension_owner(&self) -> Option<&ExtensionSessionOwner> {
self.extension_owner.as_deref()
}
}
impl fmt::Debug for ProviderSessionState {
@@ -32,6 +32,13 @@ impl ExtensionSetId {
pub trait ExtensionSetLease: Send + Sync {
fn is_ready(&self) -> bool;
/// 与发布集合一起冻结的上游适配计划,不能在请求中解析为新代次。
fn upstream_adapters(
&self,
) -> Option<Arc<dyn crate::engine::upstream_adapter::UpstreamAdapterPlan>> {
None
}
/// 与集合一起冻结的目录事实,读取不调用插件进程。
fn model_aliases(&self) -> &[ContributedModelAlias] {
&[]
@@ -50,6 +57,15 @@ pub struct ExtensionSetReference {
}
impl ExtensionSetReference {
#[must_use]
pub fn upstream_adapters(
&self,
) -> Option<crate::engine::upstream_adapter::FrozenUpstreamAdapterPlan> {
self.lease.upstream_adapters().map(|plan| {
crate::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(plan, self.clone())
})
}
#[must_use]
pub fn new(id: ExtensionSetId, lease: Arc<dyn ExtensionSetLease>) -> Self {
Self { id, lease }
@@ -276,6 +276,80 @@ fn provider_stream_should_not_sequence_gate_canonical_observation_attached_to_wi
assert!(futures::executor::block_on(provider_stream.next()).is_none());
}
#[test]
fn provider_stream_should_validate_complete_facts_across_mixed_wire_events() {
for wire_started in [false, true] {
let metadata = ProviderCallMetadata::new(
ProviderKind::new("openai").unwrap(),
UpstreamModelId::new("gpt-5").unwrap(),
ProviderAccountId::new("acct_mixed").unwrap(),
UpstreamTransport::new("http_json").unwrap(),
);
let response = ResponseMeta::new("response-mixed", "gpt-5");
let wire = || ProtocolWireEvent::json("openai", None, json!({})).unwrap();
let started = GatewayEvent::Started(response.clone());
let started = if wire_started {
ProviderEvent::canonical_with_wire(vec![started], wire())
} else {
ProviderEvent::canonical(started)
};
let events = vec![
Ok(started),
Ok(ProviderEvent::canonical(GatewayEvent::Usage(
gateway_core::metering::Usage {
input_tokens: Some(7),
..Default::default()
},
))),
Ok(ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Completed(response)],
wire(),
)),
];
let mut provider = ProviderStream::new(metadata, Box::pin(stream::iter(events)), ());
futures::executor::block_on(async {
for _ in 0..3 {
assert!(provider.next().await.unwrap().is_ok());
}
assert!(provider.next().await.is_none());
});
}
}
#[test]
fn provider_stream_should_reject_invalid_wire_facts_when_canonical_output_is_used() {
for invalid_wire_first in [false, true] {
let metadata = ProviderCallMetadata::new(
ProviderKind::new("openai").unwrap(),
UpstreamModelId::new("gpt-5").unwrap(),
ProviderAccountId::new("acct_mixed").unwrap(),
UpstreamTransport::new("http_json").unwrap(),
);
let started = GatewayEvent::Started(ResponseMeta::new("response-mixed", "gpt-5"));
let canonical = ProviderEvent::canonical(started.clone());
let wire = ProviderEvent::canonical_with_wire(
vec![if invalid_wire_first {
GatewayEvent::Usage(Default::default())
} else {
started
}],
ProtocolWireEvent::json("openai", None, json!({})).unwrap(),
);
let events = if invalid_wire_first {
vec![Ok(wire), Ok(canonical)]
} else {
vec![Ok(canonical), Ok(wire)]
};
let mut provider = ProviderStream::new(metadata, Box::pin(stream::iter(events)), ());
futures::executor::block_on(async {
assert!(provider.next().await.unwrap().is_ok());
let error = provider.next().await.unwrap().unwrap_err();
assert_eq!(error.kind(), ProviderErrorKind::Protocol);
assert!(provider.next().await.is_none());
});
}
}
#[test]
fn provider_stream_should_report_confirmed_sent_failure_but_ignore_unconfirmed_failure() {
let feedback = Arc::new(AccountFeedbackStats::default());
+1
View File
@@ -40,6 +40,7 @@ tar = "0.4.44"
thiserror = "2.0.18"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "signal", "time", "sync", "net", "process", "io-util"] }
tower-service = "0.3.3"
tokio-tungstenite = { version = "0.29.0", default-features = false, features = ["handshake"] }
tracing = "0.1.44"
tracing-appender = "0.2.5"
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
@@ -30,6 +30,7 @@ struct ClientKey {
addresses: Vec<std::net::SocketAddr>,
proxy: Option<OutboundProxy>,
scope: Option<String>,
websocket: bool,
}
pub struct HttpRequest {
@@ -94,7 +95,7 @@ impl HttpError {
self
}
fn sent(reason: &'static str) -> Self {
pub(super) fn sent(reason: &'static str) -> Self {
Self {
reason,
kind: HttpErrorKind::Response,
@@ -139,7 +140,7 @@ impl HttpClient {
return client.clone();
}
let connector = Connector::new(key.addresses.clone(), key.proxy.clone(), self.tls.clone());
let connector = connector::https(connector, self.tls.as_ref().clone(), true);
let connector = connector::https(connector, self.tls.as_ref().clone(), !key.websocket);
let client = Client::builder(TokioExecutor::new())
.pool_max_idle_per_host(4)
.http1_max_buf_size(64 * 1024)
@@ -173,6 +174,21 @@ impl HttpClient {
timeout: Duration,
scope: Option<&str>,
) -> Result<HttpResponse, HttpError> {
let response = self
.send(request, proxy, network, timeout, scope, None)
.await?;
Ok(response.into_http())
}
pub(super) async fn send(
&self,
request: HttpRequest,
proxy: Option<&OutboundProxy>,
network: &NetworkPolicy,
timeout: Duration,
scope: Option<&str>,
websocket_key: Option<&str>,
) -> Result<PendingResponse, HttpError> {
if request.body.len() > 16 * 1024 * 1024
|| request.headers.len() > 128
|| timeout.is_zero()
@@ -241,6 +257,7 @@ impl HttpClient {
addresses: addresses.clone(),
proxy: proxy.cloned(),
scope: scope.map(str::to_owned),
websocket: websocket_key.is_some(),
});
let mut uri: http::Uri = url
.as_str()
@@ -278,6 +295,18 @@ impl HttpClient {
}
}
let mut outgoing = http::Request::new(Full::new(Bytes::from(request.body)));
if let Some(key) = websocket_key {
headers.insert(
http::header::CONNECTION,
HeaderValue::from_static("Upgrade"),
);
headers.insert(http::header::UPGRADE, HeaderValue::from_static("websocket"));
headers.insert("sec-websocket-version", HeaderValue::from_static("13"));
headers.insert(
"sec-websocket-key",
HeaderValue::from_str(key).map_err(|_| HttpError::invalid("WebSocket key"))?,
);
}
*outgoing.method_mut() = method;
*outgoing.uri_mut() = uri;
*outgoing.headers_mut() = headers;
@@ -311,17 +340,37 @@ impl HttpClient {
{
return Err(HttpError::sent("response headers"));
}
Ok(HttpResponse {
Ok(PendingResponse {
status,
headers,
response,
permit,
deadline,
})
}
}
pub(super) struct PendingResponse {
pub(super) status: u16,
pub(super) headers: Vec<(String, Vec<u8>)>,
pub(super) response: http::Response<hyper::body::Incoming>,
pub(super) permit: OwnedSemaphorePermit,
pub(super) deadline: Instant,
}
impl PendingResponse {
pub(super) fn into_http(self) -> HttpResponse {
HttpResponse {
status: self.status,
headers: self.headers,
body: HttpBody {
response: response.into_body(),
response: self.response.into_body(),
pending: Bytes::new(),
received: 0,
deadline,
_permit: permit,
deadline: self.deadline,
_permit: self.permit,
},
})
}
}
}
@@ -3,6 +3,8 @@
mod connector;
mod http;
mod network;
mod websocket;
pub use http::{HttpBody, HttpClient, HttpError, HttpErrorKind, HttpRequest, HttpResponse};
pub use network::{DnsResolver, NetworkPolicy};
pub use websocket::{ManagedWebSocket, WebSocketMessage, WebSocketResponse};
@@ -0,0 +1,166 @@
use std::time::Duration;
use bytes::Bytes;
use futures::{SinkExt as _, StreamExt as _};
use gateway_core::account::OutboundProxy;
use hyper_util::rt::TokioIo;
use tokio::sync::OwnedSemaphorePermit;
use tokio_tungstenite::{
WebSocketStream,
tungstenite::{
Message,
handshake::{client::generate_key, derive_accept_key},
protocol::{Role, WebSocketConfig},
},
};
use super::{HttpBody, HttpClient, HttpError, HttpErrorKind, HttpRequest, NetworkPolicy};
const MAX_MESSAGE_BYTES: usize = 8 * 1024 * 1024;
/// 无内容型 Debug,连接由持有者独占;丢弃时关闭底层连接并释放容量。
pub struct ManagedWebSocket {
socket: WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>,
_permit: OwnedSemaphorePermit,
}
pub enum WebSocketMessage {
Text(String),
Binary(Bytes),
}
pub struct WebSocketResponse {
pub status: u16,
pub headers: Vec<(String, Vec<u8>)>,
pub connection: Option<ManagedWebSocket>,
pub body: Option<HttpBody>,
}
impl HttpClient {
/// 使用与 HTTP 相同的 DNS、代理、证书和容量边界执行 HTTP/1.1 Upgrade。
pub async fn open_websocket(
&self,
request: HttpRequest,
proxy: Option<&OutboundProxy>,
network: &NetworkPolicy,
timeout: Duration,
) -> Result<WebSocketResponse, HttpError> {
if request.method != "GET"
|| !request.body.is_empty()
|| request
.headers
.iter()
.any(|(name, _)| name.to_ascii_lowercase().starts_with("sec-websocket-"))
{
return Err(HttpError::invalid("WebSocket request"));
}
let key = generate_key();
let mut response = self
.send(request, proxy, network, timeout, None, Some(&key))
.await?;
if response.status != 101 {
let response = response.into_http();
return Ok(WebSocketResponse {
status: response.status,
headers: response.headers,
connection: None,
body: Some(response.body),
});
}
let headers = response.response.headers();
if headers
.get("sec-websocket-accept")
.and_then(|value| value.to_str().ok())
!= Some(derive_accept_key(key.as_bytes()).as_str())
|| !headers
.get("upgrade")
.is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket"))
|| !headers
.get("connection")
.and_then(|value| value.to_str().ok())
.is_some_and(|value| {
value
.split(',')
.any(|part| part.trim().eq_ignore_ascii_case("upgrade"))
})
|| headers.contains_key("sec-websocket-extensions")
|| headers.contains_key("sec-websocket-protocol")
{
return Err(HttpError::sent("WebSocket handshake"));
}
let upgraded = tokio::time::timeout_at(
response.deadline,
hyper::upgrade::on(&mut response.response),
)
.await
.map_err(|_| timeout_error())?
.map_err(|_| HttpError::sent("WebSocket upgrade"))?;
let config = WebSocketConfig::default()
.max_message_size(Some(MAX_MESSAGE_BYTES))
.max_frame_size(Some(MAX_MESSAGE_BYTES));
let socket =
WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Client, Some(config))
.await;
Ok(WebSocketResponse {
status: response.status,
headers: response.headers,
connection: Some(ManagedWebSocket {
socket,
_permit: response.permit,
}),
body: None,
})
}
}
impl ManagedWebSocket {
pub async fn send(
&mut self,
message: WebSocketMessage,
timeout: Duration,
) -> Result<(), HttpError> {
let message = match message {
WebSocketMessage::Text(text) if text.len() <= MAX_MESSAGE_BYTES => {
Message::Text(text.into())
}
WebSocketMessage::Binary(bytes) if bytes.len() <= MAX_MESSAGE_BYTES => {
Message::Binary(bytes)
}
_ => return Err(HttpError::invalid("WebSocket message limits")),
};
tokio::time::timeout(timeout, self.socket.send(message))
.await
.map_err(|_| timeout_error())?
.map_err(|_| HttpError::sent("WebSocket send"))
}
pub async fn read(&mut self, timeout: Duration) -> Result<Option<WebSocketMessage>, HttpError> {
tokio::time::timeout(timeout, async {
loop {
match self.socket.next().await {
Some(Ok(Message::Text(text))) => {
return Ok(Some(WebSocketMessage::Text(text.to_string())));
}
Some(Ok(Message::Binary(bytes))) => {
return Ok(Some(WebSocketMessage::Binary(bytes)));
}
Some(Ok(Message::Ping(_))) => self
.socket
.flush()
.await
.map_err(|_| HttpError::sent("WebSocket heartbeat"))?,
Some(Ok(Message::Pong(_))) => {}
Some(Ok(Message::Close(_))) | None => return Ok(None),
Some(Ok(Message::Frame(_))) => return Err(HttpError::sent("WebSocket frame")),
Some(Err(_)) => return Err(HttpError::sent("WebSocket receive")),
}
}
})
.await
.map_err(|_| timeout_error())?
}
}
fn timeout_error() -> HttpError {
HttpError::sent("WebSocket deadline").with_kind(HttpErrorKind::Timeout)
}
@@ -1,3 +1,4 @@
mod connector;
mod http;
mod network;
mod websocket;
@@ -0,0 +1,155 @@
use std::time::Duration;
use futures::{SinkExt as _, StreamExt as _};
use gateway_core::{account::OutboundProxy, upstream::UpstreamSendState};
use gateway_host::outbound::{
HttpClient, HttpErrorKind, HttpRequest, NetworkPolicy, WebSocketMessage,
};
use tokio::net::TcpListener;
use tokio_tungstenite::tungstenite::Message;
fn request(url: String) -> HttpRequest {
HttpRequest {
method: "GET".into(),
url,
headers: vec![],
body: vec![],
}
}
fn network() -> NetworkPolicy {
NetworkPolicy::new(&["127.0.0.0/8".into()]).unwrap()
}
#[tokio::test]
async fn websocket_preserves_frames_handles_ping_and_releases_connection_on_drop() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}/responses", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
let (socket, _) = listener.accept().await.unwrap();
let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap();
assert_eq!(
websocket.next().await.unwrap().unwrap(),
Message::Text("first".into())
);
websocket
.send(Message::Ping(vec![1, 2].into()))
.await
.unwrap();
websocket
.send(Message::Binary(vec![0, 255, 7].into()))
.await
.unwrap();
assert_eq!(
websocket.next().await.unwrap().unwrap(),
Message::Pong(vec![1, 2].into())
);
assert!(
websocket
.next()
.await
.is_none_or(|message| message.is_err())
);
});
let mut response = HttpClient::new()
.unwrap()
.open_websocket(request(url), None, &network(), Duration::from_secs(2))
.await
.unwrap();
assert_eq!(response.status, 101);
let mut websocket = response.connection.take().unwrap();
websocket
.send(
WebSocketMessage::Text("first".into()),
Duration::from_secs(2),
)
.await
.unwrap();
let Some(WebSocketMessage::Binary(bytes)) =
websocket.read(Duration::from_secs(2)).await.unwrap()
else {
panic!("binary frame");
};
assert_eq!(&bytes[..], &[0, 255, 7]);
drop(websocket);
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
async fn websocket_rejections_return_http_body_and_never_fall_back_from_proxy() {
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(401).set_body_string("denied"))
.expect(1)
.mount(&server)
.await;
let client = HttpClient::new().unwrap();
let response = client
.open_websocket(
request(server.uri()),
None,
&network(),
Duration::from_secs(2),
)
.await
.unwrap();
assert_eq!(response.status, 401);
assert!(response.connection.is_none());
assert_eq!(
&response.body.unwrap().read(64).await.unwrap().unwrap()[..],
b"denied"
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let proxy =
OutboundProxy::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap();
drop(listener);
let error = client
.open_websocket(
request(server.uri()),
Some(&proxy),
&network(),
Duration::from_secs(2),
)
.await
.err()
.unwrap();
assert_eq!(error.send_state, UpstreamSendState::NotSent);
assert_eq!(error.kind(), HttpErrorKind::Transport);
}
#[tokio::test]
async fn websocket_refuses_invalid_upgrade_and_header_injection() {
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}/", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut bytes = Vec::new();
while !bytes.ends_with(b"\r\n\r\n") {
assert!(bytes.len() < 4096);
bytes.push(socket.read_u8().await.unwrap());
}
socket.write_all(b"HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: invalid\r\n\r\n").await.unwrap();
});
let client = HttpClient::new().unwrap();
let mut invalid = request(url.clone());
invalid
.headers
.push(("sec-websocket-key".into(), b"injected".to_vec()));
let error = client
.open_websocket(invalid, None, &network(), Duration::from_secs(2))
.await
.err()
.unwrap();
assert_eq!(error.send_state, UpstreamSendState::NotSent);
let error = client
.open_websocket(request(url), None, &network(), Duration::from_secs(2))
.await
.err()
.unwrap();
assert_eq!(error.send_state, UpstreamSendState::Sent);
server.await.unwrap();
}
@@ -37,6 +37,7 @@ uuid = { version = "1.23.3", features = ["v4"] }
url = "2.5.8"
[dev-dependencies]
tokio-tungstenite = { version = "0.29.0", default-features = false, features = ["handshake"] }
wiremock = "0.6.5"
gateway-store = { path = "../../gateway-store" }
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres"] }
@@ -9,6 +9,7 @@
{ "capability": "model_catalog", "versions": [1] },
{ "capability": "retry_policy", "versions": [1] },
{ "capability": "middleware", "versions": [1, 2] },
{ "capability": "upstream_adapter", "versions": [1] },
{ "capability": "request_lifecycle", "versions": [1] },
{ "capability": "web_socket_observer", "versions": [1] },
{ "capability": "usage", "versions": [1] },
@@ -18,6 +19,7 @@
],
"permissions": [
"network",
"upstream_connections",
"models",
"accounts",
"data",
@@ -5,3 +5,4 @@ pub(super) mod management;
pub(super) mod observer;
pub(super) mod policy;
pub(super) mod scope;
pub(super) mod upstream_adapter;
@@ -0,0 +1,364 @@
use std::time::Duration;
use bytes::Bytes;
use gateway_core::{
engine::upstream_adapter::UpstreamAccountConnection,
error::{ClientVisibleUpstreamError, ProviderError, ProviderErrorKind},
event::{
ContentItem, ContentKind, EventSequenceValidator, FinishReason, GatewayEvent,
ProtocolWireEvent, ProviderEvent, ProviderResponseObservation, ReasoningDelta,
ResponseMeta, TextDelta, ToolCallDelta,
},
metering::Usage,
upstream::{OpaqueUpstreamValue, UpstreamSendState, UpstreamTransport},
};
use gateway_plugin_sdk::call::{
model::{CanonicalEvent, WireEvent, WirePayload},
upstream_adapter::{
UpstreamAdapterEvent, UpstreamContinuation, UpstreamFailure, UpstreamFailureKind,
},
};
use gateway_protocol::openai::sse::SseEventDecoder;
#[derive(Default)]
pub(super) struct EventDecoder {
sequence: EventSequenceValidator,
response_id: Option<String>,
service_tier: Option<String>,
usage: Usage,
contents: usize,
pub(super) completed: bool,
}
pub(super) struct DecodedEvent {
pub(super) event: ProviderEvent,
pub(super) continuation: Option<UpstreamContinuation>,
}
impl EventDecoder {
pub(super) fn decode(
&mut self,
payload: &[u8],
protocol: &str,
transport: &str,
account: &dyn UpstreamAccountConnection,
sent: UpstreamSendState,
) -> Result<DecodedEvent, ProviderError> {
let message = UpstreamAdapterEvent::decode(payload).map_err(|_| invalid(sent))?;
if self.completed || message.event.facts.len() > 64 {
return Err(invalid(sent));
}
let wire = message
.event
.wire
.map(|wire| decode_wire(wire, protocol, sent))
.transpose()?;
if let Some(failure) = message.failure {
if !message.event.facts.is_empty() || message.continuation.is_some() {
return Err(invalid(sent));
}
let mut error = failure_error(failure, sent)?;
if let Some(wire) = wire {
error = error.with_atomic_client_events(vec![ProviderEvent::wire(wire)]);
}
return Err(error);
}
if let Some(tier) = message.service_tier {
if tier.is_empty() || tier.len() > 64 || tier.chars().any(char::is_control) {
return Err(invalid(sent));
}
// 上游可能在完成时才把 auto 解析为实际档位;与原生观测一样保留最新值。
self.service_tier = Some(tier);
}
let mut facts = Vec::new();
for fact in message.event.facts {
let fact = match fact {
CanonicalEvent::Started { id, model } => {
validate_id(&id, sent)?;
validate_model(model.as_deref(), sent)?;
self.response_id = Some(id.clone());
GatewayEvent::Started(meta(id, model))
}
CanonicalEvent::ContentAdded { index, kind } => {
self.contents += 1;
if self.contents > 4096 {
return Err(invalid(sent));
}
GatewayEvent::ContentAdded(ContentItem::new(
index,
match kind {
gateway_plugin_sdk::call::model::ContentKind::Text => ContentKind::Text,
gateway_plugin_sdk::call::model::ContentKind::Reasoning => {
ContentKind::Reasoning
}
gateway_plugin_sdk::call::model::ContentKind::ToolCall => {
ContentKind::ToolCall
}
gateway_plugin_sdk::call::model::ContentKind::Image => {
ContentKind::Image
}
gateway_plugin_sdk::call::model::ContentKind::Audio => {
ContentKind::Audio
}
},
))
}
CanonicalEvent::TextDelta { index, text } => GatewayEvent::TextDelta(TextDelta {
content_index: index,
text,
}),
CanonicalEvent::ReasoningDelta { index, text } => {
GatewayEvent::ReasoningDelta(ReasoningDelta {
content_index: index,
text,
})
}
CanonicalEvent::ToolCallDelta {
index,
id,
name,
arguments,
} => {
validate_id(&id, sent)?;
if name
.as_ref()
.is_some_and(|name| name.len() > 512 || name.chars().any(char::is_control))
{
return Err(invalid(sent));
}
GatewayEvent::ToolCallDelta(ToolCallDelta {
content_index: index,
call_id: id,
name,
arguments_delta: arguments,
})
}
CanonicalEvent::Usage { usage } => {
if [
usage.input_tokens,
usage.output_tokens,
usage.cached_tokens,
usage.cache_write_tokens,
usage.reasoning_tokens,
usage.image_input_tokens,
usage.image_output_tokens,
usage.total_tokens,
]
.into_iter()
.flatten()
.any(|count| count > i64::MAX as u64)
{
return Err(invalid(sent));
}
let usage = Usage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cached_tokens: usage.cached_tokens,
cache_write_tokens: usage.cache_write_tokens,
reasoning_tokens: usage.reasoning_tokens,
image_input_tokens: usage.image_input_tokens,
image_output_tokens: usage.image_output_tokens,
total_tokens: usage.total_tokens,
};
self.usage.merge(&usage);
GatewayEvent::Usage(usage)
}
CanonicalEvent::Completed { id, model, reason } => {
validate_model(model.as_deref(), sent)?;
if self.response_id.as_deref() != Some(id.as_str()) {
return Err(invalid(sent));
}
if let Some(cost) =
account.calculate_cost(self.service_tier.as_deref(), &self.usage)
{
let cost = GatewayEvent::CalculatedCost(cost);
self.sequence.observe(&cost).map_err(|_| invalid(sent))?;
facts.push(cost);
}
self.completed = true;
GatewayEvent::Completed(meta(id, model).with_finish_reason(match reason {
gateway_plugin_sdk::call::model::FinishReason::Stop => FinishReason::Stop,
gateway_plugin_sdk::call::model::FinishReason::Length => {
FinishReason::Length
}
gateway_plugin_sdk::call::model::FinishReason::ToolCall => {
FinishReason::ToolCall
}
gateway_plugin_sdk::call::model::FinishReason::ContentFilter => {
FinishReason::ContentFilter
}
gateway_plugin_sdk::call::model::FinishReason::Other => FinishReason::Other,
}))
}
};
self.sequence.observe(&fact).map_err(|_| invalid(sent))?;
facts.push(fact);
}
if message.continuation.is_some() && !self.completed {
return Err(invalid(sent));
}
let mut event = if let Some(wire) = wire {
ProviderEvent::canonical_with_wire(facts, wire)
} else {
let mut facts = facts.into_iter();
let first = facts.next().ok_or_else(|| invalid(sent))?;
facts.fold(ProviderEvent::canonical(first), ProviderEvent::with_fact)
};
if let Some(tier) = &self.service_tier {
let observation = ProviderResponseObservation::new(
UpstreamTransport::new(transport).map_err(|_| invalid(sent))?,
)
.try_with_service_tier(tier.clone())
.map_err(|_| invalid(sent))?;
event.attach_observation(observation);
}
Ok(DecodedEvent {
event,
continuation: message.continuation,
})
}
pub(super) fn finish(&self, sent: UpstreamSendState) -> Result<(), ProviderError> {
self.sequence.finish().map_err(|_| invalid(sent))
}
}
fn meta(id: String, model: Option<String>) -> ResponseMeta {
model.map_or_else(
|| ResponseMeta::for_provider_endpoint(id.clone()),
|model| ResponseMeta::new(id.clone(), model),
)
}
fn validate_id(id: &str, sent: UpstreamSendState) -> Result<(), ProviderError> {
if id.is_empty() || id.len() > 512 || id.chars().any(char::is_control) {
return Err(invalid(sent));
}
Ok(())
}
fn validate_model(model: Option<&str>, sent: UpstreamSendState) -> Result<(), ProviderError> {
if model.is_some_and(|model| {
model.is_empty() || model.len() > 256 || model.chars().any(char::is_control)
}) {
return Err(invalid(sent));
}
Ok(())
}
fn decode_wire(
wire: WireEvent,
protocol: &str,
sent: UpstreamSendState,
) -> Result<ProtocolWireEvent, ProviderError> {
if wire.protocol != protocol {
return Err(invalid(sent));
}
let result = match wire.payload {
WirePayload::Json {
event,
data,
id,
retry,
raw_sse,
} => {
if let Some(raw) = raw_sse {
let frame = single_sse(&raw, sent)?;
let [parsed] = frame.events() else {
return Err(invalid(sent));
};
if parsed.event != event
|| parsed.id != id
|| parsed.retry != retry
|| serde_json::from_str::<serde_json::Value>(&parsed.data)
.ok()
.as_ref()
!= Some(&data)
{
return Err(invalid(sent));
}
ProtocolWireEvent::json_with_raw_sse_metadata(
protocol,
event,
data,
Bytes::from(raw),
id,
retry,
)
} else {
ProtocolWireEvent::json_with_sse_metadata(protocol, event, data, id, retry)
}
}
WirePayload::RawSse { frame } => {
single_sse(&frame, sent)?;
ProtocolWireEvent::raw_sse(protocol, Bytes::from(frame))
}
WirePayload::RawJson { body } => {
serde_json::from_slice::<serde_json::Value>(&body).map_err(|_| invalid(sent))?;
ProtocolWireEvent::raw_json(protocol, Bytes::from(body))
}
WirePayload::RawBody { body } => {
ProtocolWireEvent::raw_http_body(protocol, Bytes::from(body))
}
};
result.map_err(|_| invalid(sent))
}
fn single_sse(
bytes: &[u8],
sent: UpstreamSendState,
) -> Result<gateway_protocol::openai::sse::SseFrame, ProviderError> {
let mut frames = SseEventDecoder::default().push_frames(bytes);
if frames.len() != 1 || frames[0].raw() != bytes || std::str::from_utf8(bytes).is_err() {
return Err(invalid(sent));
}
frames.pop().ok_or_else(|| invalid(sent))
}
fn failure_error(
failure: UpstreamFailure,
sent: UpstreamSendState,
) -> Result<ProviderError, ProviderError> {
if failure.message.len() > 64 * 1024
|| failure.code.as_ref().is_some_and(|code| code.len() > 256)
|| failure
.status
.is_some_and(|status| !(400..=599).contains(&status))
|| failure
.retry_after_ms
.is_some_and(|delay| delay > 24 * 60 * 60 * 1000)
{
return Err(invalid(sent));
}
let kind = match failure.kind {
UpstreamFailureKind::InvalidRequest => ProviderErrorKind::InvalidRequest,
UpstreamFailureKind::Unsupported => ProviderErrorKind::Unsupported,
UpstreamFailureKind::Unauthorized => ProviderErrorKind::Unauthorized,
UpstreamFailureKind::PermissionDenied => ProviderErrorKind::PermissionDenied,
UpstreamFailureKind::RateLimited => ProviderErrorKind::RateLimited,
UpstreamFailureKind::QuotaExhausted => ProviderErrorKind::QuotaExhausted,
UpstreamFailureKind::Timeout => ProviderErrorKind::Timeout,
UpstreamFailureKind::Unavailable => ProviderErrorKind::Unavailable,
UpstreamFailureKind::Protocol => ProviderErrorKind::Protocol,
};
let mut error = ProviderError::new(kind, sent);
if let Some(code) = &failure.code {
error = error.with_upstream_code(OpaqueUpstreamValue::new(code.clone()));
}
error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
failure.message,
failure.code,
None,
));
if let Some(status) = failure.status {
error = error.with_status(status);
}
if let Some(delay) = failure.retry_after_ms {
error = error.with_retry_after(Duration::from_millis(delay));
}
Ok(error)
}
pub(super) fn invalid(sent: UpstreamSendState) -> ProviderError {
ProviderError::new(ProviderErrorKind::Protocol, sent)
}
@@ -0,0 +1,386 @@
use std::{sync::Arc, time::SystemTime};
use futures::stream;
use gateway_core::{
engine::{provider::EventStream, upstream_adapter::UpstreamAdapterInvocation},
error::{ProviderError, ProviderErrorKind},
operation::{ExtensionSessionOwner, ProviderSessionState},
upstream::UpstreamSendState,
};
use gateway_plugin_sdk::{
CallContext, ErrorCode, PluginFault, Stage,
call::upstream_adapter::{
ContinuationScope, EXECUTE_METHOD, UpstreamAdapterRequest, UpstreamContinuation,
UpstreamTransport,
},
};
use serde::{Deserialize, Serialize};
use super::{
PluginUpstreamAdapter,
event::{DecodedEvent, EventDecoder, invalid},
};
use crate::{
RpcError, RpcStream,
callback::{
NetworkScope,
upstream::{ConnectionOwner, ManagedUpstream},
},
};
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct StoredContinuation {
continuation: UpstreamContinuation,
client_key_id: String,
account_id: String,
credential_revision: u64,
connection: Option<String>,
}
struct Execution {
adapter: Arc<PluginUpstreamAdapter>,
invocation: UpstreamAdapterInvocation,
active: Option<ActiveCall>,
decoder: EventDecoder,
ended: bool,
}
struct ActiveCall {
context: CallContext,
stream: RpcStream,
managed: Arc<ManagedUpstream>,
_scope: Arc<NetworkScope>,
}
pub(super) fn execute(
adapter: Arc<PluginUpstreamAdapter>,
invocation: UpstreamAdapterInvocation,
) -> EventStream {
let state = Execution {
adapter,
invocation,
active: None,
decoder: EventDecoder::default(),
ended: false,
};
// try_unfold 首次 poll 才运行闭包;Core 在此前登记 attempt 并持有账号租约。
Box::pin(stream::try_unfold(state, |mut state| async move {
if state.ended {
return Ok(None);
}
let result = state.next().await;
match result {
Ok(event) => Ok(Some((event, state))),
Err(error) => Err(state.invocation.account.record_failure(error).await),
}
}))
}
impl Execution {
async fn next(&mut self) -> Result<gateway_core::event::ProviderEvent, ProviderError> {
if self.active.is_none() {
self.active = Some(self.begin().await?);
}
let active = self
.active
.as_mut()
.ok_or_else(|| invalid(UpstreamSendState::NotSent))?;
let cancellation = self.invocation.context.cancellation();
let chunk = tokio::select! {
biased;
() = cancellation.cancelled() => return Err(ProviderError::new(ProviderErrorKind::Cancelled, active.managed.send_state.get())),
result = active.stream.next() => result.map_err(|error| rpc_error(error, active.managed.send_state.get()))?,
}.ok_or_else(|| invalid(active.managed.send_state.get()))?;
let mut decoded = self.decoder.decode(
&chunk,
&self.adapter.declaration.protocol,
self.adapter.declaration.transport.as_str(),
self.invocation.account.as_ref(),
active.managed.send_state.get(),
)?;
if self.decoder.completed {
// 先确认 RPC 正常终结,再发布唯一 Completed;否则 Core 不会继续 poll 尾部错误。
let terminal = async {
let tail = tokio::select! {
biased;
() = cancellation.cancelled() => return Err(ProviderError::new(ProviderErrorKind::Cancelled, active.managed.send_state.get())),
result = active.stream.next() => result.map_err(|error| rpc_error(error, active.managed.send_state.get()))?,
};
if tail.is_some() {
return Err(invalid(active.managed.send_state.get()));
}
self.decoder.finish(active.managed.send_state.get())?;
attach_continuation(&self.adapter, &self.invocation, active, &mut decoded).await
}.await;
if let Err(error) = terminal {
// 尾部故障不得交付成功 wire,但已经解析的计量事实仍归 Core 结算。
let metering = decoded
.event
.canonical_facts()
.iter()
.filter(|fact| {
matches!(
fact,
gateway_core::event::GatewayEvent::Usage(_)
| gateway_core::event::GatewayEvent::CalculatedCost(_)
)
})
.cloned()
.map(gateway_core::event::ProviderEvent::canonical)
.collect();
return Err(error.with_atomic_client_events(metering));
}
self.ended = true;
}
Ok(decoded.event)
}
async fn begin(&self) -> Result<ActiveCall, ProviderError> {
let invalid = || invalid(UpstreamSendState::NotSent);
let declaration = &self.adapter.declaration;
let invocation = &self.invocation;
if invocation.operation.protocol() != declaration.protocol
|| invocation.metadata.provider().as_str() != declaration.provider.as_str()
|| invocation.metadata.provider_account_id() != invocation.account.account_id()
|| !declaration
.authentication_kinds
.iter()
.any(|kind| kind == invocation.account.authentication_kind())
{
return Err(invalid());
}
let timeout = invocation
.context
.deadline()
.duration_since(SystemTime::now())
.map_err(|_| {
ProviderError::new(ProviderErrorKind::Timeout, UpstreamSendState::NotSent)
})?
.min(self.adapter.session.maximum_call_timeout());
if timeout.is_zero() || invocation.context.cancellation().is_cancelled() {
return Err(ProviderError::new(
ProviderErrorKind::Cancelled,
UpstreamSendState::NotSent,
));
}
let mut context = self.adapter.session.context(Stage::Upstream, timeout);
context.request_id = Some(invocation.context.request_id().as_str().to_owned());
context.attempt_id = Some(format!(
"{}:{}",
invocation.context.request_id().as_str(),
invocation.context.attempt_index()
));
context.account_id = Some(invocation.account.account_id().as_str().to_owned());
context.credential_revision = Some(invocation.account.credential_revision().get());
let owner = ConnectionOwner {
client_key_id: invocation.context.client_api_key_ref().as_str().to_owned(),
account_id: invocation.account.account_id().as_str().to_owned(),
credential_revision: invocation.account.credential_revision().get(),
};
let previous = previous_continuation(&self.adapter, invocation, &context, &owner)?;
let managed = ManagedUpstream::new(
Arc::clone(&invocation.account),
Arc::clone(&self.adapter.target),
invocation.context.execution_effects(),
Arc::clone(&self.adapter.connections),
owner,
previous
.as_ref()
.and_then(|state| state.connection.as_deref()),
declaration.transport == UpstreamTransport::WebSocket,
)
.map_err(|error| fault_error(error, UpstreamSendState::NotSent))?;
let extension_scope = invocation
.context
.extension_scope()
.extending(self.adapter.instance_id.clone())
.ok_or_else(invalid)?;
let scope = self
.adapter
.callbacks
.prepare_upstream(&context, Arc::clone(&managed), extension_scope)
.map_err(|error| fault_error(error, UpstreamSendState::NotSent))?;
let request = UpstreamAdapterRequest {
adapter_id: declaration.id.clone(),
provider: declaration.provider,
upstream_model: invocation
.metadata
.upstream_model()
.ok_or_else(invalid)?
.as_str()
.to_owned(),
client_key_id: invocation.context.client_api_key_ref().as_str().to_owned(),
account_id: invocation.account.account_id().as_str().to_owned(),
credential_revision: invocation.account.credential_revision().get(),
protocol: invocation.operation.protocol().to_owned(),
client_transport: invocation.context.client_transport().as_str().to_owned(),
disable_fast: invocation.context.disable_fast(),
headers: invocation
.headers
.iter()
.map(|header| (header.name().to_owned(), header.value().to_vec()))
.collect(),
continuation: previous.map(|state| state.continuation),
};
let payload = request
.encode(
invocation
.operation
.middleware_body()
.map_err(|_| invalid())?
.to_vec(),
)
.map_err(|_| invalid())?;
invocation.context.trace().record("plugin.upstream", serde_json::json!({ "instanceId": self.adapter.instance_id, "adapterId": declaration.id, "generation": context.generation, "transport": declaration.transport.as_str() }));
let stream = tokio::select! {
biased;
() = invocation.context.cancellation().cancelled() => return Err(ProviderError::new(ProviderErrorKind::Cancelled, managed.send_state.get())),
result = self.adapter.session.call_stream(EXECUTE_METHOD, context.clone(), serde_json::json!({}), payload) => result.map_err(|error| rpc_error(error, managed.send_state.get()))?,
};
if stream.initial.result != serde_json::json!({}) || !stream.initial.payload.is_empty() {
return Err(super::event::invalid(managed.send_state.get()));
}
Ok(ActiveCall {
context,
stream,
managed,
_scope: scope,
})
}
}
fn previous_continuation(
adapter: &PluginUpstreamAdapter,
invocation: &UpstreamAdapterInvocation,
context: &CallContext,
owner: &ConnectionOwner,
) -> Result<Option<StoredContinuation>, ProviderError> {
let invalid = || invalid(UpstreamSendState::NotSent);
let Some(state) = invocation
.operation
.provider_session_state(adapter.declaration.provider.as_str())
else {
return if invocation.context.continuation().is_some() {
Err(invalid())
} else {
Ok(None)
};
};
let expected = extension_owner(
adapter,
context,
state
.extension_owner()
.is_some_and(|owner| owner.connection_local),
);
if state.extension_owner() != Some(&expected) {
return Err(invalid());
}
let stored: StoredContinuation =
serde_json::from_value(serde_json::Value::Object(state.payload().clone()))
.map_err(|_| invalid())?;
if stored.client_key_id != owner.client_key_id
|| stored.account_id != owner.account_id
|| stored.credential_revision != owner.credential_revision
|| (stored.continuation.scope == ContinuationScope::ConnectionLocal)
!= stored.connection.is_some()
|| expected.connection_local != stored.connection.is_some()
{
return Err(invalid());
}
Ok(Some(stored))
}
async fn attach_continuation(
adapter: &PluginUpstreamAdapter,
invocation: &UpstreamAdapterInvocation,
active: &ActiveCall,
decoded: &mut DecodedEvent,
) -> Result<(), ProviderError> {
let invalid = || invalid(active.managed.send_state.get());
let Some(continuation) = decoded.continuation.take() else {
if invocation
.operation
.provider_session_state(adapter.declaration.provider.as_str())
.is_some()
{
return Err(invalid());
}
return Ok(());
};
if continuation.upstream_response_id.is_empty()
|| continuation.upstream_response_id.len() > 512
|| serde_json::to_vec(&continuation)
.map_err(|_| invalid())?
.len()
> 32 * 1024
{
return Err(invalid());
}
let local = continuation.scope == ContinuationScope::ConnectionLocal;
let connection = if local {
Some(
active
.managed
.retain_connection()
.await
.map_err(|error| fault_error(error, active.managed.send_state.get()))?,
)
} else {
None
};
let state = StoredContinuation {
continuation,
connection,
client_key_id: invocation.context.client_api_key_ref().as_str().to_owned(),
account_id: invocation.account.account_id().as_str().to_owned(),
credential_revision: invocation.account.credential_revision().get(),
};
let serde_json::Value::Object(payload) = serde_json::to_value(state).map_err(|_| invalid())?
else {
return Err(invalid());
};
let state = ProviderSessionState::new(adapter.declaration.provider.as_str(), payload)
.map_err(|_| invalid())?
.with_extension_owner(extension_owner(adapter, &active.context, local));
decoded.event.attach_session_update(state);
Ok(())
}
fn extension_owner(
adapter: &PluginUpstreamAdapter,
context: &CallContext,
connection_local: bool,
) -> ExtensionSessionOwner {
ExtensionSessionOwner {
instance_id: adapter.instance_id.clone(),
contribution_id: adapter.contribution_id.clone(),
adapter_id: adapter.declaration.id.clone(),
generation: context.generation,
incarnation: context.incarnation.clone(),
connection_local,
}
}
fn rpc_error(error: RpcError, sent: UpstreamSendState) -> ProviderError {
match error {
RpcError::Remote(fault) => fault_error(fault, sent),
RpcError::Timeout => ProviderError::new(ProviderErrorKind::Timeout, sent),
RpcError::Cancelled => ProviderError::new(ProviderErrorKind::Cancelled, sent),
_ => ProviderError::new(ProviderErrorKind::Unavailable, sent),
}
}
fn fault_error(fault: PluginFault, sent: UpstreamSendState) -> ProviderError {
let kind = match fault.code {
ErrorCode::Timeout => ProviderErrorKind::Timeout,
ErrorCode::Cancelled => ProviderErrorKind::Cancelled,
ErrorCode::Unsupported => ProviderErrorKind::Unsupported,
ErrorCode::InvalidInput | ErrorCode::Rejected | ErrorCode::PermissionDenied => {
ProviderErrorKind::InvalidRequest
}
_ => ProviderErrorKind::Unavailable,
};
ProviderError::new(kind, sent)
}
@@ -0,0 +1,157 @@
mod event;
mod execution;
mod registration;
pub(crate) mod target;
use std::{fmt, sync::Arc};
use gateway_admin::model::AdminError;
use gateway_core::{
engine::{
AttemptContext,
upstream_adapter::{UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan},
},
error::{ProviderError, ProviderErrorKind},
identity::ProviderKind,
routing::UpstreamModelId,
upstream::UpstreamSendState,
};
use gateway_plugin_sdk::call::upstream_adapter::UpstreamAdapterDeclaration;
use super::scope::BindingScope;
use crate::{RpcSession, callback::PluginCallbacks};
pub(crate) use registration::{prepare, unavailable_entries, validate_bindings};
pub(crate) struct AdapterEntry {
instance_id: String,
scope: BindingScope,
adapter: Option<Arc<PluginUpstreamAdapter>>,
}
pub(crate) struct PluginUpstreamAdapter {
instance_id: String,
contribution_id: String,
declaration: UpstreamAdapterDeclaration,
target: Arc<target::UpstreamTarget>,
session: Arc<RpcSession>,
callbacks: Arc<PluginCallbacks>,
connections: Arc<crate::callback::upstream::ConnectionPool>,
}
impl fmt::Debug for PluginUpstreamAdapter {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("PluginUpstreamAdapter")
.field("instance_id", &self.instance_id)
.field("adapter_id", &self.declaration.id)
.finish_non_exhaustive()
}
}
impl UpstreamAdapter for PluginUpstreamAdapter {
fn transport(&self) -> &str {
self.declaration.transport.as_str()
}
fn execute(
self: Arc<Self>,
invocation: UpstreamAdapterInvocation,
) -> gateway_core::engine::provider::EventStream {
execution::execute(self, invocation)
}
}
pub(crate) struct PluginUpstreamAdapterPlan {
entries: Vec<AdapterEntry>,
}
impl PluginUpstreamAdapterPlan {
pub(crate) fn compile(
entries: Vec<AdapterEntry>,
) -> Result<Option<Arc<dyn UpstreamAdapterPlan>>, AdminError> {
for (index, left) in entries.iter().enumerate() {
for right in &entries[index + 1..] {
if left.scope.overlaps(&right.scope)
&& adapters_overlap(left.adapter.as_deref(), right.adapter.as_deref())
{
return Err(AdminError::invalid(
"上游适配器的 Provider 与模型绑定范围重叠",
));
}
}
}
Ok((!entries.is_empty())
.then(|| Arc::new(Self { entries }) as Arc<dyn UpstreamAdapterPlan>))
}
}
impl fmt::Debug for PluginUpstreamAdapterPlan {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("PluginUpstreamAdapterPlan")
.field("adapters", &self.entries.len())
.finish()
}
}
impl UpstreamAdapterPlan for PluginUpstreamAdapterPlan {
fn select(
&self,
context: &AttemptContext,
provider: &ProviderKind,
model: &UpstreamModelId,
) -> Result<Option<Arc<dyn UpstreamAdapter>>, ProviderError> {
let model = context
.requested_model()
.map_or(model.as_str(), |model| model.as_str());
for entry in &self.entries {
if context.extension_scope().contains(&entry.instance_id)
|| !entry.scope.matches_processing(
context.client_api_key_ref(),
context.account_group_ids(),
Some(provider),
Some(model),
)
{
continue;
}
let adapter = entry.adapter.as_ref().ok_or_else(unavailable)?;
if adapter.declaration.provider.as_str() == provider.as_str()
&& (adapter.declaration.models.is_empty()
|| adapter
.declaration
.models
.iter()
.any(|candidate| candidate == model))
{
if !adapter.session.is_ready() {
return Err(unavailable());
}
return Ok(Some(Arc::clone(adapter) as Arc<dyn UpstreamAdapter>));
}
}
Ok(None)
}
}
fn adapters_overlap(
left: Option<&PluginUpstreamAdapter>,
right: Option<&PluginUpstreamAdapter>,
) -> bool {
let (Some(left), Some(right)) = (left, right) else {
return true;
};
left.declaration.provider == right.declaration.provider
&& (left.declaration.models.is_empty()
|| right.declaration.models.is_empty()
|| left
.declaration
.models
.iter()
.any(|model| right.declaration.models.contains(model)))
}
fn unavailable() -> ProviderError {
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
}
@@ -0,0 +1,154 @@
use std::{collections::BTreeSet, sync::Arc, time::Duration};
use gateway_admin::model::{
AdminError,
plugins::instances::{PluginCapabilityBinding, PluginFailurePolicy, PluginInstance},
};
use gateway_plugin_sdk::{
Capability, Manifest, Permission, Stage,
call::upstream_adapter::{REGISTER_METHOD, UpstreamAdapterRegistration},
};
use super::{AdapterEntry, PluginUpstreamAdapter, target::UpstreamTarget};
use crate::{RpcSession, adapter::scope::BindingScope, callback::PluginCallbacks};
pub(crate) fn validate_bindings(
manifest: &Manifest,
bindings: &[PluginCapabilityBinding],
) -> Result<(), AdminError> {
let Some(declaration) = manifest.contributes.get(&Capability::UpstreamAdapter) else {
return Ok(());
};
let mut found = false;
for binding in bindings
.iter()
.filter(|binding| binding.contribution == declaration.id)
{
if std::mem::replace(&mut found, true)
|| binding.stage != "upstream"
|| binding.failure_policy != PluginFailurePolicy::Reject
|| binding
.provider_ids
.iter()
.any(|provider| !matches!(provider.as_str(), "openai" | "xai"))
{
return Err(AdminError::invalid(
"上游适配器只允许一次 upstream 绑定,使用 reject 策略及内置 Provider 范围",
));
}
BindingScope::compile(binding)?;
}
Ok(())
}
pub(crate) async fn prepare(
manifest: &Manifest,
instance: &PluginInstance,
session: Arc<RpcSession>,
callbacks: Arc<PluginCallbacks>,
) -> Result<Vec<AdapterEntry>, AdminError> {
let Some(contribution) = manifest.contributes.get(&Capability::UpstreamAdapter) else {
return Ok(Vec::new());
};
validate_bindings(manifest, &instance.bindings)?;
for permission in [Permission::Requests, Permission::UpstreamConnections] {
if !instance
.grants
.iter()
.any(|grant| grant.permission == permission.as_str())
{
return Err(AdminError::invalid(
"上游适配器需要 requests 和 upstream_connections 授权",
));
}
}
let reply = session
.call(
REGISTER_METHOD,
session.context(Stage::Registration, Duration::from_secs(5)),
serde_json::json!({}),
vec![],
)
.await
.map_err(|_| AdminError::invalid("上游适配器注册失败"))?;
if reply.result != serde_json::json!({}) || reply.payload.len() > 64 * 1024 {
return Err(AdminError::invalid("上游适配器注册信封无效或过大"));
}
let registration: UpstreamAdapterRegistration = serde_json::from_slice(&reply.payload)
.map_err(|_| AdminError::invalid("上游适配器声明无效"))?;
if registration.adapters.is_empty() || registration.adapters.len() > 16 {
return Err(AdminError::invalid("单个插件须声明 1 至 16 个上游适配器"));
}
let mut ids = BTreeSet::new();
let binding = instance
.bindings
.iter()
.find(|binding| binding.contribution == contribution.id);
let mut entries = Vec::new();
for declaration in registration.adapters {
if declaration.id.is_empty()
|| declaration.id.len() > 64
|| !declaration
.id
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))
|| !ids.insert(declaration.id.clone())
|| declaration.authentication_kinds.is_empty()
|| declaration.authentication_kinds.len() > 8
|| declaration.authentication_kinds.iter().any(|kind| {
kind.is_empty()
|| kind.len() > 64
|| !kind
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
})
|| !contribution.input_formats.contains(&declaration.protocol)
|| !contribution.output_formats.contains(&declaration.protocol)
|| declaration.models.len() > 256
{
return Err(AdminError::invalid("上游适配器身份、认证或协议范围无效"));
}
let mut models = BTreeSet::new();
for model in &declaration.models {
gateway_core::routing::PublicModelId::new(model.clone())
.map_err(|_| AdminError::invalid("上游适配器模型无效"))?;
if !models.insert(model) {
return Err(AdminError::invalid("上游适配器模型重复"));
}
}
let target = Arc::new(UpstreamTarget::compile(&declaration)?);
if let Some(binding) = binding {
entries.push(AdapterEntry {
instance_id: instance.id.clone(),
scope: BindingScope::compile(binding)?,
adapter: Some(Arc::new(PluginUpstreamAdapter {
instance_id: instance.id.clone(),
contribution_id: contribution.id.clone(),
declaration,
target,
session: Arc::clone(&session),
callbacks: Arc::clone(&callbacks),
connections: Arc::default(),
})),
});
}
}
Ok(entries)
}
pub(crate) fn unavailable_entries(
instance: &PluginInstance,
) -> Result<Vec<AdapterEntry>, AdminError> {
instance
.bindings
.iter()
.filter(|binding| binding.stage == "upstream")
.map(|binding| {
Ok(AdapterEntry {
instance_id: instance.id.clone(),
scope: BindingScope::compile(binding)?,
adapter: None,
})
})
.collect()
}
@@ -0,0 +1,108 @@
use std::collections::BTreeMap;
use gateway_admin::model::AdminError;
use gateway_plugin_sdk::{
ErrorCode, PluginFault,
call::upstream_adapter::{UpstreamAdapterDeclaration, UpstreamPathPurpose},
};
use url::Url;
pub(crate) struct UpstreamTarget {
base: Url,
paths: BTreeMap<String, UpstreamPathPurpose>,
}
impl UpstreamTarget {
pub(super) fn compile(declaration: &UpstreamAdapterDeclaration) -> Result<Self, AdminError> {
let base = Url::parse(&declaration.base_url)
.map_err(|_| AdminError::invalid("上游 Base URL 无效"))?;
if !matches!(base.scheme(), "https" | "http")
|| base.host_str().is_none()
|| !base.username().is_empty()
|| base.password().is_some()
|| base.query().is_some()
|| base.fragment().is_some()
|| !base.path().ends_with('/')
|| declaration.base_url.len() > 2048
|| declaration.paths.is_empty()
|| declaration.paths.len() > 16
{
return Err(AdminError::invalid(
"上游 Base URL 须为无凭据、查询串或片段且以 / 结尾的 HTTP(S) 地址",
));
}
let mut paths = BTreeMap::new();
for path in &declaration.paths {
if !valid_path(&path.path) || paths.insert(path.path.clone(), path.purpose).is_some() {
return Err(AdminError::invalid("上游业务路径无效或重复"));
}
}
if !paths
.values()
.any(|purpose| *purpose == UpstreamPathPurpose::Inference)
{
return Err(AdminError::invalid("上游适配器须声明推理路径"));
}
Ok(Self { base, paths })
}
pub(crate) fn resolve(
&self,
path: &str,
query: &[(String, String)],
) -> Result<(Url, UpstreamPathPurpose), PluginFault> {
let purpose = *self.paths.get(path).ok_or_else(invalid)?;
if query.len() > 64
|| query
.iter()
.any(|(key, value)| key.len() > 256 || value.len() > 4096)
{
return Err(invalid());
}
let mut target = self.base.join(path).map_err(|_| invalid())?;
if target.origin() != self.base.origin() || !target.path().starts_with(self.base.path()) {
return Err(invalid());
}
if !query.is_empty() {
target
.query_pairs_mut()
.extend_pairs(query.iter().map(|(key, value)| (key, value)));
}
Ok((target, purpose))
}
}
fn valid_path(path: &str) -> bool {
!path.is_empty()
&& path.len() <= 2048
&& path.is_ascii()
&& !path.starts_with('/')
&& path.bytes().all(|byte| {
byte.is_ascii_alphanumeric() || matches!(byte, b'/' | b'-' | b'_' | b'.' | b'~')
})
&& path
.split('/')
.all(|segment| !matches!(segment, "" | "." | ".."))
}
pub(crate) fn protected_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
matches!(
name.as_str(),
"authorization"
| "proxy-authorization"
| "cookie"
| "host"
| "connection"
| "upgrade"
| "transfer-encoding"
| "content-length"
| "proxy-connection"
| "te"
| "trailer"
) || name.starts_with("sec-websocket-")
}
fn invalid() -> PluginFault {
PluginFault::new(ErrorCode::InvalidInput, "upstream target is not authorized")
}
@@ -8,30 +8,112 @@ use serde::de::DeserializeOwned;
use super::{CallResources, HttpAuthorization, HttpStream, NetworkScope, denied, invalid};
use crate::RpcReply;
#[expect(
clippy::too_many_arguments,
reason = "回调只接受宿主已解析的调用资源与权限,不重建第二份上下文"
)]
pub(super) async fn dispatch(
client: &HttpClient,
authorization: &HttpAuthorization,
scope: &Arc<NetworkScope>,
call: &CallResources,
method: &str,
params: serde_json::Value,
payload: Vec<u8>,
maximum_payload: usize,
) -> Result<RpcReply, PluginFault> {
match method {
"host.http.do" | "host.http.do_stream" => {
let request: wire::HttpRequest = decode(params)?;
let timeout = call
.deadline
.checked_duration_since(tokio::time::Instant::now())
.ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed"))?;
let attempt = scope.start_http();
let response = client
.open(
/// 回调的公共网络外层;账号鉴权只装饰请求,发送、读流和回收共用同一终端。
pub(super) struct HttpCallbacks<'a> {
pub(super) client: &'a HttpClient,
pub(super) authorization: &'a HttpAuthorization,
pub(super) scope: &'a Arc<NetworkScope>,
pub(super) call: &'a CallResources,
pub(super) maximum_payload: usize,
}
impl HttpCallbacks<'_> {
pub(super) async fn send(
&self,
request: HttpRequest,
streaming: bool,
upstream_purpose: Option<gateway_plugin_sdk::call::upstream_adapter::UpstreamPathPurpose>,
) -> Result<RpcReply, PluginFault> {
let Self {
client,
authorization,
scope,
call,
maximum_payload,
} = *self;
let timeout = call
.deadline
.checked_duration_since(tokio::time::Instant::now())
.ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed"))?;
let attempt = scope.start_upstream(upstream_purpose);
let response = client
.open(
request,
scope.proxy.as_ref(),
&authorization.network,
timeout,
)
.await;
attempt.finish(
response
.as_ref()
.map_or_else(|error| error.send_state, |_| UpstreamSendState::Sent),
);
let response = response.map_err(http_error)?;
let mut body = response.body;
let headers = response
.headers
.into_iter()
.map(|(name, value)| String::from_utf8(value).map(|value| (name, value)))
.collect::<Result<Vec<_>, _>>()
.map_err(|_| sent_fault(ErrorCode::Fault, "HTTP header cannot be represented"))?;
let mut result = wire::HttpResponse {
status: response.status,
headers,
stream: None,
};
let mut payload = Vec::new();
if streaming {
let id = uuid::Uuid::new_v4().to_string();
let mut state = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.closed {
return Err(denied());
}
if state.streams.len() >= 16 {
return Err(sent_fault(
ErrorCode::Capacity,
"HTTP stream capacity exhausted",
));
}
state.streams.insert(
id.clone(),
Arc::new(HttpStream {
body: tokio::sync::Mutex::new(Some(body)),
closed: tokio::sync::watch::channel(false).0,
}),
);
result.stream = Some(id);
} else {
while let Some(chunk) = body.read(64 * 1024).await.map_err(http_error)? {
if payload.len() + chunk.len() > maximum_payload {
return Err(sent_fault(
ErrorCode::Capacity,
"HTTP body requires streaming",
));
}
payload.extend_from_slice(&chunk);
}
}
Ok(RpcReply {
result: serde_json::to_value(result).map_err(|_| invalid())?,
payload,
})
}
pub(super) async fn dispatch(
&self,
method: &str,
params: serde_json::Value,
payload: Vec<u8>,
) -> Result<RpcReply, PluginFault> {
let call = self.call;
match method {
"host.http.do" | "host.http.do_stream" => {
let request: wire::HttpRequest = decode(params)?;
self.send(
HttpRequest {
method: request.method,
url: request.url,
@@ -42,128 +124,71 @@ pub(super) async fn dispatch(
.collect(),
body: payload,
},
scope.proxy.as_ref(),
&authorization.network,
timeout,
method == "host.http.do_stream",
None,
)
.await;
attempt.finish(
response
.as_ref()
.map_or_else(|error| error.send_state, |_| UpstreamSendState::Sent),
);
let response = response.map_err(http_error)?;
let mut body = response.body;
let headers = response
.headers
.into_iter()
.map(|(name, value)| String::from_utf8(value).map(|value| (name, value)))
.collect::<Result<Vec<_>, _>>()
.map_err(|_| sent_fault(ErrorCode::Fault, "HTTP header cannot be represented"))?;
let mut result = wire::HttpResponse {
status: response.status,
headers,
stream: None,
};
let mut payload = Vec::new();
if method == "host.http.do_stream" {
let id = uuid::Uuid::new_v4().to_string();
let mut state = call
.await
}
"host.http.stream_read" => {
if !payload.is_empty() {
return Err(invalid());
}
let request: wire::StreamRead = decode(params)?;
if request.maximum_bytes == 0 || request.maximum_bytes > 64 * 1024 {
return Err(invalid());
}
let stream = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.closed {
return Err(denied());
}
if state.streams.len() >= 16 {
return Err(sent_fault(
ErrorCode::Capacity,
"HTTP stream capacity exhausted",
));
}
state.streams.insert(
id.clone(),
Arc::new(HttpStream {
body: tokio::sync::Mutex::new(Some(body)),
closed: tokio::sync::watch::channel(false).0,
}),
);
result.stream = Some(id);
} else {
while let Some(chunk) = body.read(64 * 1024).await.map_err(http_error)? {
if payload.len() + chunk.len() > maximum_payload {
return Err(sent_fault(
ErrorCode::Capacity,
"HTTP body requires streaming",
));
}
payload.extend_from_slice(&chunk);
}
}
Ok(RpcReply {
result: serde_json::to_value(result).map_err(|_| invalid())?,
payload,
})
}
"host.http.stream_read" => {
if !payload.is_empty() {
return Err(invalid());
}
let request: wire::StreamRead = decode(params)?;
if request.maximum_bytes == 0 || request.maximum_bytes > 64 * 1024 {
return Err(invalid());
}
let stream = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.streams
.get(&request.stream)
.cloned()
.ok_or_else(denied)?;
let mut body = stream.body.try_lock().map_err(|_| {
PluginFault::new(ErrorCode::Capacity, "HTTP stream already has a reader")
})?;
let mut closed = stream.closed.subscribe();
let result = tokio::select! {
biased;
_ = closed.wait_for(|closed| *closed) => Err(denied()),
result = body.as_mut().ok_or_else(denied)?.read(request.maximum_bytes as usize) => result.map_err(http_error),
};
if !matches!(&result, Ok(Some(_))) {
body.take();
stream.closed.send_replace(true);
call.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.streams
.remove(&request.stream);
.get(&request.stream)
.cloned()
.ok_or_else(denied)?;
let mut body = stream.body.try_lock().map_err(|_| {
PluginFault::new(ErrorCode::Capacity, "HTTP stream already has a reader")
})?;
let mut closed = stream.closed.subscribe();
let result = tokio::select! {
biased;
_ = closed.wait_for(|closed| *closed) => Err(denied()),
result = body.as_mut().ok_or_else(denied)?.read(request.maximum_bytes as usize) => result.map_err(http_error),
};
if !matches!(&result, Ok(Some(_))) {
body.take();
stream.closed.send_replace(true);
call.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.streams
.remove(&request.stream);
}
let chunk = result?;
Ok(RpcReply {
result: serde_json::json!({"eof":chunk.is_none()}),
payload: chunk.map_or_else(Vec::new, |chunk| chunk.to_vec()),
})
}
let chunk = result?;
Ok(RpcReply {
result: serde_json::json!({"eof":chunk.is_none()}),
payload: chunk.map_or_else(Vec::new, |chunk| chunk.to_vec()),
})
}
"host.http.stream_close" => {
if !payload.is_empty() {
return Err(invalid());
"host.http.stream_close" => {
if !payload.is_empty() {
return Err(invalid());
}
let request: wire::StreamClose = decode(params)?;
let stream = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.streams
.remove(&request.stream)
.ok_or_else(denied)?;
stream.close();
Ok(RpcReply {
result: serde_json::json!({}),
payload: vec![],
})
}
let request: wire::StreamClose = decode(params)?;
let stream = call
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.streams
.remove(&request.stream)
.ok_or_else(denied)?;
stream.close();
Ok(RpcReply {
result: serde_json::json!({}),
payload: vec![],
})
_ => Err(denied()),
}
_ => Err(denied()),
}
}
@@ -177,7 +202,7 @@ fn sent_fault(code: ErrorCode, message: &'static str) -> PluginFault {
error
}
fn http_error(error: HttpError) -> PluginFault {
pub(super) fn http_error(error: HttpError) -> PluginFault {
tracing::warn!(
reason = error.reason(),
kind = ?error.kind(),
@@ -10,6 +10,7 @@ mod model;
pub(crate) mod private_state;
mod resources;
mod scope;
pub(crate) mod upstream;
use std::{
collections::BTreeMap,
@@ -87,6 +88,7 @@ pub(crate) struct PluginCallbacks {
struct HttpAuthorization {
network: NetworkPolicy,
authorized: bool,
upstream_authorized: bool,
}
struct CallResources {
@@ -155,6 +157,9 @@ impl PluginCallbacks {
http_authorization: Arc::new(HttpAuthorization {
network: ports.network,
authorized: grants.iter().any(|grant| grant.permission == "network"),
upstream_authorized: grants
.iter()
.any(|grant| grant.permission == "upstream_connections"),
}),
scopes: Mutex::new(BTreeMap::new()),
pending_middleware: Mutex::new(BTreeMap::new()),
@@ -218,6 +223,25 @@ impl PluginCallbacks {
)
}
pub(crate) fn prepare_upstream(
&self,
context: &CallContext,
managed: Arc<upstream::ManagedUpstream>,
extension_scope: gateway_core::engine::extensions::ExtensionCallScope,
) -> Result<Arc<NetworkScope>, PluginFault> {
if context.stage != gateway_plugin_sdk::Stage::Upstream
|| !self.http_authorization.upstream_authorized
|| context.account_id.as_deref() != Some(managed.account.account_id().as_str())
|| context.credential_revision != Some(managed.account.credential_revision().get())
{
return Err(denied());
}
let mut scope = NetworkScope::new(context, managed.account.outbound_proxy().cloned())
.with_upstream(managed);
scope.extension_scope = extension_scope;
self.register_scope(context, scope)
}
pub(crate) fn prepare_management(
&self,
context: &CallContext,
@@ -480,20 +504,25 @@ impl CallbackHandler for PluginCallbacks {
.call(&context, scope, &method, params, &payload)
.await;
}
let http = http::HttpCallbacks {
client: &http,
authorization: &authorization,
scope,
call: &call,
maximum_payload,
};
if method.starts_with("host.upstream.") {
if !authorization.upstream_authorized
|| context.stage != gateway_plugin_sdk::Stage::Upstream
{
return Err(denied());
}
return upstream::dispatch(&http, &method, params, payload).await;
}
if !authorization.authorized {
return Err(denied());
}
http::dispatch(
&http,
&authorization,
scope,
&call,
&method,
params,
payload,
maximum_payload,
)
.await
http.dispatch(&method, params, payload).await
})
}
}
@@ -13,6 +13,7 @@ pub(crate) struct NetworkScope {
pub(super) proxy: Option<OutboundProxy>,
pub(super) extension_scope: gateway_core::engine::extensions::ExtensionCallScope,
execution_effects: Option<Arc<ExecutionEffects>>,
pub(super) upstream: Option<Arc<super::upstream::ManagedUpstream>>,
}
impl NetworkScope {
@@ -24,6 +25,7 @@ impl NetworkScope {
proxy,
extension_scope: Default::default(),
execution_effects: None,
upstream: None,
}
}
@@ -36,6 +38,12 @@ impl NetworkScope {
self
}
pub(super) fn with_upstream(mut self, managed: Arc<super::upstream::ManagedUpstream>) -> Self {
self.execution_effects = managed.effects.clone();
self.upstream = Some(managed);
self
}
pub(super) fn account_id(&self) -> Option<&str> {
self.account_id.as_deref()
}
@@ -50,9 +58,22 @@ impl NetworkScope {
&& context.credential_revision == self.credential_revision
}
pub(super) fn start_http(&self) -> HttpAttempt {
pub(super) fn start_upstream(
&self,
purpose: Option<gateway_plugin_sdk::call::upstream_adapter::UpstreamPathPurpose>,
) -> HttpAttempt {
HttpAttempt {
effects: self.execution_effects.clone(),
effects: if purpose
== Some(gateway_plugin_sdk::call::upstream_adapter::UpstreamPathPurpose::Inference)
{
None
} else {
self.execution_effects.clone()
},
upstream: self
.upstream
.as_ref()
.map(|managed| Arc::clone(&managed.send_state)),
completed: false,
}
}
@@ -61,6 +82,7 @@ impl NetworkScope {
pub(super) struct HttpAttempt {
effects: Option<Arc<ExecutionEffects>>,
completed: bool,
upstream: Option<Arc<super::upstream::SendWatermark>>,
}
impl HttpAttempt {
@@ -70,6 +92,9 @@ impl HttpAttempt {
}
fn observe(&self, observed: UpstreamSendState) {
if let Some(upstream) = &self.upstream {
upstream.observe(observed);
}
if observed != UpstreamSendState::NotSent
&& let Some(effects) = &self.effects
{
@@ -0,0 +1,331 @@
mod sessions;
use std::{
sync::{
Arc,
atomic::{AtomicU8, Ordering},
},
time::Duration,
};
use gateway_core::{
engine::{nested::ExecutionEffects, upstream_adapter::UpstreamAccountConnection},
upstream::UpstreamSendState,
};
use gateway_host::outbound::{HttpRequest, ManagedWebSocket, WebSocketMessage};
use gateway_plugin_sdk::{
ErrorCode, PluginFault,
call::{
host as wire,
upstream_adapter::{self as adapter, UpstreamPathPurpose},
},
};
use super::{denied, http::HttpCallbacks, invalid};
use crate::{
RpcReply,
adapter::upstream_adapter::target::{UpstreamTarget, protected_header},
};
pub(crate) use sessions::{ConnectionOwner, ConnectionPool};
#[derive(Default)]
pub(crate) struct SendWatermark(AtomicU8);
impl SendWatermark {
pub(crate) fn observe(&self, state: UpstreamSendState) {
self.0.fetch_max(
match state {
UpstreamSendState::NotSent => 0,
UpstreamSendState::Sent => 1,
UpstreamSendState::Ambiguous => 2,
},
Ordering::AcqRel,
);
}
pub(crate) fn get(&self) -> UpstreamSendState {
match self.0.load(Ordering::Acquire) {
0 => UpstreamSendState::NotSent,
1 => UpstreamSendState::Sent,
_ => UpstreamSendState::Ambiguous,
}
}
}
pub(crate) struct ManagedUpstream {
pub(crate) account: Arc<dyn UpstreamAccountConnection>,
pub(crate) send_state: Arc<SendWatermark>,
pub(crate) effects: Option<Arc<ExecutionEffects>>,
target: Arc<UpstreamTarget>,
socket: tokio::sync::Mutex<Option<(String, ManagedWebSocket)>>,
pool: Arc<ConnectionPool>,
owner: ConnectionOwner,
websocket_allowed: bool,
}
impl ManagedUpstream {
pub(crate) fn new(
account: Arc<dyn UpstreamAccountConnection>,
target: Arc<UpstreamTarget>,
effects: Option<Arc<ExecutionEffects>>,
pool: Arc<ConnectionPool>,
owner: ConnectionOwner,
connection: Option<&str>,
websocket_allowed: bool,
) -> Result<Arc<Self>, PluginFault> {
let socket = connection.map(|id| pool.take(id, &owner)).transpose()?;
Ok(Arc::new(Self {
account,
target,
effects,
pool,
owner,
websocket_allowed,
send_state: Arc::default(),
socket: tokio::sync::Mutex::new(socket),
}))
}
pub(crate) async fn retain_connection(&self) -> Result<String, PluginFault> {
let socket = self.socket.lock().await.take().ok_or_else(denied)?;
self.pool.put(self.owner.clone(), socket)
}
fn request(
&self,
request: adapter::UpstreamHttpRequest,
body: Vec<u8>,
) -> Result<(HttpRequest, UpstreamPathPurpose), PluginFault> {
let (target, purpose) = self.target.resolve(&request.path, &request.query)?;
if request.headers.len() > 128
|| request
.headers
.iter()
.any(|(name, _)| protected_header(name))
{
return Err(denied());
}
let authorization = self.account.authorization().map_err(|_| {
PluginFault::new(
ErrorCode::Rejected,
"selected account authentication is unavailable",
)
})?;
if request.headers.iter().any(|(name, _)| {
authorization
.iter()
.any(|header| name.eq_ignore_ascii_case(header.name()))
}) {
return Err(denied());
}
let headers = request
.headers
.into_iter()
.map(|(name, value)| (name, value.into_bytes()))
.chain(authorization.into_iter().map(|header| {
let (name, value) = header.into_parts();
(name, value.to_vec())
}))
.collect();
Ok((
HttpRequest {
method: request.method,
url: target.into(),
headers,
body,
},
purpose,
))
}
}
pub(super) async fn dispatch(
http: &HttpCallbacks<'_>,
method: &str,
params: serde_json::Value,
payload: Vec<u8>,
) -> Result<RpcReply, PluginFault> {
let HttpCallbacks {
client,
authorization,
scope,
call,
maximum_payload,
} = *http;
let managed = scope.upstream.as_ref().ok_or_else(denied)?;
let timeout = || {
call.deadline
.checked_duration_since(tokio::time::Instant::now())
.filter(|remaining| !remaining.is_zero())
.map(|remaining| remaining.min(Duration::from_secs(120)))
.ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "upstream deadline elapsed"))
};
match method {
"host.upstream.http.do" | "host.upstream.http.do_stream" => {
let request = serde_json::from_value(params).map_err(|_| invalid())?;
let (request, purpose) = managed.request(request, payload)?;
http.send(request, method.ends_with("do_stream"), Some(purpose))
.await
}
"host.upstream.http.stream_read" | "host.upstream.http.stream_close" => {
http.dispatch(
if method.ends_with("stream_read") {
"host.http.stream_read"
} else {
"host.http.stream_close"
},
params,
payload,
)
.await
}
"host.upstream.websocket.open" => {
if !payload.is_empty() || !managed.websocket_allowed {
return Err(denied());
}
let request: adapter::UpstreamWebSocketRequest =
serde_json::from_value(params).map_err(|_| invalid())?;
let (request, purpose) = managed.request(
adapter::UpstreamHttpRequest {
method: "GET".into(),
path: request.path,
query: request.query,
headers: request.headers,
},
vec![],
)?;
if purpose != UpstreamPathPurpose::Inference {
return Err(denied());
}
let mut socket = managed.socket.try_lock().map_err(|_| denied())?;
if let Some((url, _)) = socket.as_ref() {
if *url != request.url {
return Err(denied());
}
return Ok(RpcReply {
result: serde_json::json!({"status":101,"headers":[],"stream":null}),
payload: vec![],
});
}
let url = request.url.clone();
let attempt = scope.start_upstream(Some(UpstreamPathPurpose::Inference));
let response = client
.open_websocket(
request,
managed.account.outbound_proxy(),
&authorization.network,
timeout()?,
)
.await;
// 握手也已向上游发送账号身份;失败不能伪装成未出站或跳过账号反馈。
attempt.finish(
response
.as_ref()
.map_or_else(|error| error.send_state, |_| UpstreamSendState::Sent),
);
let mut response = response.map_err(super::http::http_error)?;
let headers = response
.headers
.into_iter()
.map(|(name, value)| String::from_utf8(value).map(|value| (name, value)))
.collect::<Result<Vec<_>, _>>()
.map_err(|_| invalid())?;
let mut payload = Vec::new();
if let Some(body) = response.body.as_mut() {
while let Some(chunk) = body
.read(64 * 1024)
.await
.map_err(super::http::http_error)?
{
if payload.len() + chunk.len() > maximum_payload {
return Err(invalid());
}
payload.extend_from_slice(&chunk);
}
}
if let Some(connection) = response.connection {
*socket = Some((url, connection));
}
Ok(RpcReply {
result: serde_json::to_value(wire::HttpResponse {
status: response.status,
headers,
stream: None,
})
.map_err(|_| invalid())?,
payload,
})
}
"host.upstream.websocket.send" => {
let message: adapter::UpstreamWebSocketMessage =
serde_json::from_value(params).map_err(|_| invalid())?;
let message = match message.kind {
adapter::WebSocketMessageKind::Text => {
WebSocketMessage::Text(String::from_utf8(payload).map_err(|_| invalid())?)
}
adapter::WebSocketMessageKind::Binary => WebSocketMessage::Binary(payload.into()),
};
let mut socket = managed.socket.try_lock().map_err(|_| denied())?;
let (_, socket) = socket.as_mut().ok_or_else(denied)?;
let attempt = scope.start_upstream(Some(UpstreamPathPurpose::Inference));
let result = socket.send(message, timeout()?).await;
attempt.finish(
result
.as_ref()
.map_or_else(|error| error.send_state, |()| UpstreamSendState::Sent),
);
result.map_err(super::http::http_error)?;
Ok(RpcReply {
result: serde_json::json!({}),
payload: vec![],
})
}
"host.upstream.websocket.read" => {
if params != serde_json::json!({}) || !payload.is_empty() {
return Err(invalid());
}
let mut socket = managed.socket.try_lock().map_err(|_| denied())?;
let (_, connection) = socket.as_mut().ok_or_else(denied)?;
let result = connection
.read(timeout()?)
.await
.map_err(super::http::http_error)?;
let (kind, payload) = match result {
Some(WebSocketMessage::Text(text)) => {
(Some(adapter::WebSocketMessageKind::Text), text.into_bytes())
}
Some(WebSocketMessage::Binary(bytes)) => {
(Some(adapter::WebSocketMessageKind::Binary), bytes.to_vec())
}
None => {
socket.take();
(None, vec![])
}
};
if payload.len() > maximum_payload {
socket.take();
return Err(invalid());
}
Ok(RpcReply {
result: serde_json::to_value(adapter::UpstreamWebSocketRead {
eof: kind.is_none(),
kind,
})
.map_err(|_| invalid())?,
payload,
})
}
"host.upstream.websocket.close" => {
if params != serde_json::json!({}) || !payload.is_empty() {
return Err(invalid());
}
managed.socket.try_lock().map_err(|_| denied())?.take();
Ok(RpcReply {
result: serde_json::json!({}),
payload: vec![],
})
}
_ => Err(denied()),
}
}
@@ -0,0 +1,101 @@
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::Duration,
};
use gateway_host::outbound::ManagedWebSocket;
use gateway_plugin_sdk::{ErrorCode, PluginFault};
#[derive(Clone, PartialEq, Eq)]
pub(crate) struct ConnectionOwner {
pub client_key_id: String,
pub account_id: String,
pub credential_revision: u64,
}
/// 每个适配器代次独立持有连接;取出即独占,失败或取消不会归还半完成会话。
#[derive(Default)]
pub(crate) struct ConnectionPool {
idle: Mutex<BTreeMap<String, IdleConnection>>,
}
struct IdleConnection {
owner: ConnectionOwner,
socket: Option<(String, ManagedWebSocket)>,
expires: tokio::task::JoinHandle<()>,
}
impl Drop for IdleConnection {
fn drop(&mut self) {
self.expires.abort();
}
}
impl ConnectionPool {
pub(crate) fn take(
&self,
id: &str,
owner: &ConnectionOwner,
) -> Result<(String, ManagedWebSocket), PluginFault> {
let mut idle = self
.idle
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if !idle
.get(id)
.is_some_and(|connection| connection.owner == *owner)
{
return Err(PluginFault::new(
ErrorCode::Rejected,
"upstream continuation connection is unavailable",
));
}
idle.remove(id)
.and_then(|mut connection| connection.socket.take())
.ok_or_else(|| {
PluginFault::new(
ErrorCode::Rejected,
"upstream continuation connection is unavailable",
)
})
}
pub(crate) fn put(
self: &Arc<Self>,
owner: ConnectionOwner,
socket: (String, ManagedWebSocket),
) -> Result<String, PluginFault> {
let mut idle = self
.idle
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if idle.len() >= 64 {
return Err(PluginFault::new(
ErrorCode::Capacity,
"upstream continuation capacity exhausted",
));
}
let id = uuid::Uuid::new_v4().to_string();
let weak = Arc::downgrade(self);
let expired_id = id.clone();
let expires = tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(30 * 60)).await;
if let Some(pool) = weak.upgrade() {
pool.idle
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(&expired_id);
}
});
idle.insert(
id.clone(),
IdleConnection {
owner,
socket: Some(socket),
expires,
},
);
Ok(id)
}
}
@@ -57,6 +57,10 @@ impl PluginRuntime {
&instance.bindings,
)?;
crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?;
crate::adapter::upstream_adapter::validate_bindings(
package.manifest(),
&instance.bindings,
)?;
crate::adapter::frontend_authentication::validate_bindings(
package.manifest(),
&instance.bindings,
@@ -235,6 +239,13 @@ impl PluginRuntime {
)?);
let model_aliases =
crate::adapter::catalog::prepare(&manifest, &instance, &session).await?;
let upstream_entries = crate::adapter::upstream_adapter::prepare(
&manifest,
&instance,
Arc::clone(&session),
Arc::clone(&callbacks),
)
.await?;
sessions.push(PreparedInstance {
instance_id,
artifact_sha256: instance.artifact_sha256,
@@ -251,6 +262,7 @@ impl PluginRuntime {
commands,
management,
policy_entries,
upstream_entries,
authentication_entries,
model_aliases,
})
@@ -120,6 +120,7 @@ pub(super) struct PreparedSet {
_observers: Option<Arc<dyn RequestObserverPlan>>,
_policies: Option<Arc<dyn RequestPolicyPlan>>,
_middleware: Option<Arc<dyn MiddlewarePlan>>,
upstream_adapters: Option<Arc<dyn gateway_core::engine::upstream_adapter::UpstreamAdapterPlan>>,
_authentication:
Option<Arc<dyn gateway_core::engine::authentication::FrontendAuthenticationPlan>>,
sessions: Vec<PreparedInstance>,
@@ -137,6 +138,7 @@ struct PreparedContributions {
commands: Vec<Arc<crate::adapter::command_line::PluginCommand>>,
management: Vec<crate::adapter::management::ManagementEntry>,
policy_entries: Vec<crate::adapter::policy::PolicyEntry>,
upstream_entries: Vec<crate::adapter::upstream_adapter::AdapterEntry>,
authentication_entries:
Vec<crate::adapter::frontend_authentication::FrontendAuthenticationEntry>,
model_aliases: Vec<gateway_core::routing::ContributedModelAlias>,
@@ -150,6 +152,7 @@ impl PreparedContributions {
self.management.extend(other.management);
self.model_aliases.extend(other.model_aliases);
self.policy_entries.extend(other.policy_entries);
self.upstream_entries.extend(other.upstream_entries);
self.authentication_entries
.extend(other.authentication_entries);
}
@@ -173,6 +176,12 @@ impl PreparedSet {
}
impl ExtensionSetLease for PreparedSet {
fn upstream_adapters(
&self,
) -> Option<Arc<dyn gateway_core::engine::upstream_adapter::UpstreamAdapterPlan>> {
self.upstream_adapters.clone()
}
fn is_ready(&self) -> bool {
self.sessions_ready() && self.can_serve()
}
@@ -456,6 +465,9 @@ impl PluginRuntime {
contributions
.policy_entries
.extend(crate::adapter::policy::unavailable_entries(instance)?);
contributions.upstream_entries.extend(
crate::adapter::upstream_adapter::unavailable_entries(instance)?,
);
if let Some(entry) =
crate::adapter::frontend_authentication::unavailable_entry(instance)
{
@@ -471,6 +483,7 @@ impl PluginRuntime {
commands,
management,
policy_entries,
upstream_entries,
authentication_entries,
model_aliases,
} = contributions;
@@ -535,6 +548,10 @@ impl PluginRuntime {
_observers: observers,
_policies: policies,
_middleware: middleware,
upstream_adapters:
crate::adapter::upstream_adapter::PluginUpstreamAdapterPlan::compile(
upstream_entries,
)?,
_authentication: authentication,
sessions,
commands,
@@ -865,6 +882,10 @@ impl PluginPreparation for PluginRuntime {
crate::adapter::observer::validate_bindings(package.manifest(), &instance.bindings)?;
crate::adapter::policy::validate_bindings(package.manifest(), &instance.bindings)?;
crate::adapter::catalog::validate_bindings(package.manifest(), &instance.bindings)?;
crate::adapter::upstream_adapter::validate_bindings(
package.manifest(),
&instance.bindings,
)?;
crate::adapter::frontend_authentication::validate_bindings(
package.manifest(),
&instance.bindings,
@@ -4,7 +4,7 @@ use gateway_admin::model::{
AdminError,
plugins::{PluginCompatibilityRequirements, PluginHostCompatibility},
};
use gateway_plugin_sdk::Manifest;
use gateway_plugin_sdk::{Capability, MANIFEST_VERSION, Manifest, PROTOCOL_VERSION, Permission};
const HOST_COMPATIBILITY_JSON: &str = include_str!("../../plugin-host-compatibility.json");
@@ -15,7 +15,31 @@ pub(crate) fn host_compatibility() -> Result<&'static PluginHostCompatibility, A
let compatibility =
serde_json::from_str::<PluginHostCompatibility>(HOST_COMPATIBILITY_JSON)
.map_err(|_| ())?;
compatibility.is_valid().then_some(compatibility).ok_or(())
// 发行声明可以是 SDK 合同的子集,但不能声称支持当前二进制无法解释的合同。
let known_contracts = compatibility
.manifest_schema_versions
.iter()
.all(|version| *version == MANIFEST_VERSION)
&& compatibility
.protocol_versions
.iter()
.all(|version| *version == PROTOCOL_VERSION)
&& compatibility.capabilities.iter().all(|entry| {
serde_json::from_value::<Capability>(entry.capability.clone().into()).is_ok_and(
|capability| {
entry
.versions
.iter()
.all(|version| capability.contract_versions().contains(version))
},
)
})
&& compatibility.permissions.iter().all(|permission| {
serde_json::from_value::<Permission>(permission.clone().into()).is_ok()
});
(compatibility.is_valid() && known_contracts)
.then_some(compatibility)
.ok_or(())
})
.as_ref()
.map_err(|()| AdminError::internal("宿主插件兼容声明不合法"))
@@ -35,13 +59,15 @@ pub(crate) fn requirements(
capabilities: manifest
.contributes
.iter()
.map(|(capability, declaration)| Ok((identifier(*capability)?, declaration.version)))
.collect::<Result<_, AdminError>>()?,
.map(|(capability, declaration)| {
(capability.identifier().to_owned(), declaration.version)
})
.collect(),
permissions: manifest
.permissions
.iter()
.map(|permission| identifier(*permission))
.collect::<Result<_, _>>()?,
.map(|permission| permission.as_str().to_owned())
.collect(),
})
}
@@ -63,10 +89,3 @@ pub(crate) fn supports(manifest: &Manifest) -> Result<bool, AdminError> {
.iter()
.all(|permission| compatibility.supports_permission(permission)))
}
fn identifier(value: impl serde::Serialize) -> Result<String, AdminError> {
serde_json::to_value(value)
.ok()
.and_then(|value| value.as_str().map(str::to_owned))
.ok_or_else(|| AdminError::internal("插件描述转换失败"))
}
@@ -253,6 +253,18 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
) && permissions.contains(&Permission::KeyBudgets));
}
let permission = match method {
"host.upstream.http.do"
| "host.upstream.http.do_stream"
| "host.upstream.http.stream_read"
| "host.upstream.http.stream_close"
| "host.upstream.websocket.open"
| "host.upstream.websocket.send"
| "host.upstream.websocket.read"
| "host.upstream.websocket.close"
if stage == Stage::Upstream =>
{
Permission::UpstreamConnections
}
"host.http.do"
| "host.http.do_stream"
| "host.http.stream_read"
@@ -417,6 +417,10 @@ impl Shared {
}
impl RpcSession {
pub(crate) const fn maximum_call_timeout(&self) -> Duration {
self.limits.maximum_call_timeout
}
pub async fn start(
package: Arc<PreparedPackage>,
handshake: Handshake,
@@ -6,6 +6,7 @@ mod observer;
mod policy;
mod prepare;
mod private_state;
mod upstream_adapter;
use gateway_admin::{
model::{
@@ -0,0 +1,792 @@
use std::{
num::NonZeroU32,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::{Duration, SystemTime},
};
use bytes::Bytes;
use futures::{StreamExt as _, future::BoxFuture};
use gateway_admin::model::plugins::instances::{
PluginCapabilityBinding, PluginFailurePolicy, PluginPermissionGrant,
};
use gateway_core::{
account::{AccountSelectionPolicy, CredentialRevision, OutboundProxy, ProviderAccountId},
engine::{
AccountAttemptContext, AttemptContext, ModelRequestId, RequestAttemptContext,
middleware::MiddlewareHeader,
provider::{EventStream, ProviderCallMetadata},
upstream_adapter::{UpstreamAccountConnection, UpstreamAdapterInvocation},
},
error::{ProviderError, ProviderErrorKind},
event::GatewayEvent,
identity::ProviderKind,
lifecycle::CancellationToken,
metering::{CalculatedCost, Usage},
operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState},
policy::ClientApiKeyId,
routing::{ConfigRevision, UpstreamModelId},
runtime::extensions::{ExtensionPreparationPort, ExtensionSetReference},
upstream::{UpstreamSendState, UpstreamTransport},
};
use gateway_plugin_sdk::{Capability, Contributions, Permission, Stage};
use serde_json::{Value, json};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{header, method, path},
};
struct Account {
id: ProviderAccountId,
revision: u64,
failures: AtomicUsize,
}
impl Account {
fn new() -> Arc<Self> {
Self::with_identity("acct_one", 1)
}
fn with_identity(id: &str, revision: u64) -> Arc<Self> {
Arc::new(Self {
id: ProviderAccountId::new(id).unwrap(),
revision,
failures: AtomicUsize::new(0),
})
}
}
impl UpstreamAccountConnection for Account {
fn account_id(&self) -> &ProviderAccountId {
&self.id
}
fn credential_revision(&self) -> CredentialRevision {
CredentialRevision::new(self.revision).unwrap()
}
fn authentication_kind(&self) -> &str {
"oauth"
}
fn outbound_proxy(&self) -> Option<&OutboundProxy> {
None
}
fn authorization(&self) -> Result<Vec<MiddlewareHeader>, ProviderError> {
Ok(vec![MiddlewareHeader::new(
"authorization",
Bytes::from_static(b"Bearer fixture-native-token"),
)])
}
fn calculate_cost(&self, _: Option<&str>, usage: &Usage) -> Option<CalculatedCost> {
assert_eq!(usage.input_tokens, Some(7));
CalculatedCost::from_usd_ticks(123).ok()
}
fn record_failure(&self, error: ProviderError) -> BoxFuture<'_, ProviderError> {
self.failures.fetch_add(1, Ordering::SeqCst);
Box::pin(async move { error })
}
}
fn configuration(base: &str, streaming: bool) -> Value {
json!({
"upstream_registration":{"adapters":[{"id":"example","provider":"openai","base_url":format!("{base}/adapter/"),"paths":[{"path":"responses","purpose":"inference"}],"authentication_kinds":["oauth"],"transport":if streaming {"http_sse"} else {"http_json"},"protocol":"openai"}]},
"upstream_callbacks":[{"method":if streaming {"host.upstream.http.do_stream"} else {"host.upstream.http.do"},"params":{"method":"POST","path":"responses","headers":[["content-type","application/json"]]},"body":"request"}],
"upstream_events":[
{"event":{"facts":[{"type":"started","id":"response-one","model":"native-model"}]}},
{"event":{"facts":[{"type":"usage","usage":{"input_tokens":7,"output_tokens":2}},{"type":"completed","id":"response-one","model":"native-model","reason":"stop"}]},"wire":if streaming {"sse"} else {"http"},"continuation":{"scope":"persisted","upstream_response_id":"private-upstream-id","state":{"cursor":"opaque"}}}
]
})
}
async fn setup(
config: Value,
) -> (
tempfile::TempDir,
Arc<super::Store>,
gateway_plugin_runtime::PluginRuntime,
) {
let permissions = vec![Permission::Requests, Permission::UpstreamConnections];
let contributes = Contributions::from([crate::support::contribution(
Capability::UpstreamAdapter,
vec![Stage::Upstream],
vec!["openai".into()],
vec!["openai".into()],
)]);
let (cache, store, runtime) =
super::setup_with_permissions(contributes, Default::default(), permissions.clone()).await;
{
let mut snapshot = store.snapshot.lock().unwrap();
let instance = &mut snapshot.instances[0];
instance.configuration = config;
instance.grants = permissions
.into_iter()
.map(|permission| PluginPermissionGrant {
permission: permission.as_str().into(),
})
.collect();
instance.bindings = vec![PluginCapabilityBinding {
contribution: "test.example.upstreamAdapter".into(),
stage: "upstream".into(),
order: 0,
failure_policy: PluginFailurePolicy::Reject,
client_key_ids: vec![],
account_group_ids: vec![],
provider_ids: vec!["openai".into()],
models: vec![],
identity_bindings: vec![],
}];
}
(
cache,
store,
runtime.with_network_policy(
gateway_host::outbound::NetworkPolicy::new(&["127.0.0.0/8".into()]).unwrap(),
),
)
}
fn context(generation: &ExtensionSetReference, key: &str) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_one").unwrap(),
ClientApiKeyId::new(key).unwrap(),
)
.with_upstream_adapters(generation.upstream_adapters()),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
AccountSelectionPolicy::new(
gateway_core::account::RotationStrategy::RoundRobin,
std::num::NonZeroU32::new(1).unwrap(),
Duration::ZERO,
),
AccountAttemptContext::new(Default::default(), None, None),
None,
CancellationToken::new(),
)
}
fn execute(
generation: &ExtensionSetReference,
account: Arc<Account>,
key: &str,
previous: Option<ProviderSessionState>,
) -> EventStream {
let context = context(generation, key);
let provider = ProviderKind::new("openai").unwrap();
let model = UpstreamModelId::new("native-model").unwrap();
let adapter = context
.upstream_adapter(&provider, &model)
.unwrap()
.unwrap();
let metadata = ProviderCallMetadata::new(
provider,
model,
account.id.clone(),
UpstreamTransport::new(adapter.transport()).unwrap(),
);
let mut operation = Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
json!({"model":"native-model","input":"hello"})
.as_object()
.unwrap()
.clone(),
)
.unwrap(),
));
if let Some(previous) = previous {
operation = operation.with_provider_session_state(previous);
}
adapter.execute(UpstreamAdapterInvocation {
operation,
headers: vec![],
context,
metadata,
account,
})
}
#[tokio::test]
async fn adapter_http_and_sse_are_cold_authenticated_and_settled_once() {
for streaming in [false, true] {
let server = MockServer::start().await;
let body = if streaming {
"data: {\"id\":\"response-one\"}\n\n"
} else {
"{\"id\":\"response-one\"}"
};
Mock::given(method("POST"))
.and(path("/adapter/responses"))
.and(header("authorization", "Bearer fixture-native-token"))
.respond_with(ResponseTemplate::new(200).set_body_string(body))
.expect(1)
.mount(&server)
.await;
let (cache, _, runtime) = setup(configuration(&server.uri(), streaming)).await;
let generation =
ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let mut stream = execute(&generation, Arc::clone(&account), "key-one", None);
assert!(
server.received_requests().await.unwrap().is_empty(),
"构造冷流不能先发送业务请求"
);
let mut costs = 0;
let mut completed = 0;
let mut state = None;
while let Some(event) = stream.next().await {
let mut event = event.unwrap();
for fact in event.canonical_facts() {
costs += usize::from(matches!(fact, GatewayEvent::CalculatedCost(_)));
completed += usize::from(matches!(fact, GatewayEvent::Completed(_)));
}
if event.session_update().is_some() {
state = event.take_session_update();
}
}
assert_eq!((costs, completed), (1, 1));
let state = state.unwrap();
assert_eq!(state.extension_owner().unwrap().adapter_id, "example");
assert_eq!(account.failures.load(Ordering::SeqCst), 0);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn adapter_accepts_service_tier_resolved_by_the_terminal_response() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.expect(1)
.mount(&server)
.await;
let mut config = configuration(&server.uri(), false);
config["upstream_events"][0]["service_tier"] = json!("auto");
config["upstream_events"][1]["service_tier"] = json!("default");
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&generation, Account::new(), "key-one", None);
let mut tiers = Vec::new();
let mut completed = false;
while let Some(event) = stream.next().await {
let event = event.unwrap();
tiers.push(
event
.response_observation()
.unwrap()
.service_tier()
.unwrap()
.to_owned(),
);
completed |= event
.canonical_facts()
.iter()
.any(|fact| matches!(fact, GatewayEvent::Completed(_)));
}
assert_eq!(tiers, ["auto", "default"]);
assert!(completed);
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn adapter_failure_preserves_upstream_diagnostics_and_native_feedback() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(429).set_body_string("{}"))
.expect(1)
.mount(&server)
.await;
let mut config = configuration(&server.uri(), false);
config["upstream_events"] = json!([{
"event": {"facts": []},
"failure": {
"kind": "rate_limited", "status": 429, "retry_after_ms": 1500,
"message": "fixture rate limit", "code": "rate_limit_exceeded"
}
}]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let mut stream = execute(&generation, Arc::clone(&account), "key-one", None);
let error = stream.next().await.unwrap().unwrap_err();
assert_eq!(error.kind(), ProviderErrorKind::RateLimited);
assert_eq!(error.send_state(), UpstreamSendState::Sent);
assert_eq!(error.upstream_status(), Some(429));
assert_eq!(
error.upstream_code().unwrap().as_str(),
"rate_limit_exceeded"
);
assert_eq!(error.retry_after(), Some(Duration::from_millis(1500)));
assert_eq!(account.failures.load(Ordering::SeqCst), 1);
drop(stream);
assert_eq!(Arc::strong_count(&account), 1);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn adapter_rejects_forged_auth_headers_and_cross_key_continuation_before_send() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.expect(1)
.mount(&server)
.await;
let (cache, _, runtime) = setup(configuration(&server.uri(), false)).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&generation, Account::new(), "key-one", None);
let mut state = None;
while let Some(event) = stream.next().await {
state = event.unwrap().take_session_update().or(state);
}
drop(stream);
for (key, account) in [
("key-two", Account::new()),
("key-one", Account::with_identity("acct_other", 1)),
("key-one", Account::with_identity("acct_one", 2)),
] {
let error = execute(&generation, account, key, state.clone())
.next()
.await
.unwrap()
.err()
.unwrap();
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
let mut config = configuration(&server.uri(), false);
config["upstream_callbacks"][0]["params"]["headers"] = json!([["Authorization", "injected"]]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let error = execute(&generation, Account::new(), "key-one", None)
.next()
.await
.unwrap()
.err()
.unwrap();
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn adapter_does_not_publish_completed_when_rpc_fails_after_terminal_frame() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.mount(&server)
.await;
let mut config = configuration(&server.uri(), false);
config["upstream_tail_error"] = json!(true);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&generation, Account::new(), "key-one", None);
stream.next().await.unwrap().unwrap();
let error = stream.next().await.unwrap().err().unwrap();
assert_eq!(error.kind(), ProviderErrorKind::Unavailable);
assert_eq!(error.send_state(), UpstreamSendState::Sent);
let mut error = error;
let facts = error
.take_atomic_client_events()
.into_iter()
.flat_map(|event| event.into_parts().0)
.collect::<Vec<_>>();
assert_eq!(facts.len(), 2);
assert!(matches!(facts[0], GatewayEvent::Usage(_)));
assert!(matches!(facts[1], GatewayEvent::CalculatedCost(_)));
drop(stream);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn adapter_registration_rejects_target_escape_and_custom_providers() {
for invalid in [
"../responses",
"https://elsewhere.invalid/responses",
"responses?other=1",
"%2e%2e/responses",
"/responses",
] {
let mut config = configuration("https://example.invalid", false);
config["upstream_registration"]["adapters"][0]["paths"][0]["path"] = json!(invalid);
let (_, store, runtime) = setup(config).await;
let snapshot = store.snapshot.lock().unwrap().clone();
assert!(
gateway_admin::ports::plugins::PluginPreparation::prepare(&runtime, snapshot)
.await
.is_err()
);
runtime.shutdown().await;
}
let mut config = configuration("https://example.invalid", false);
config["upstream_registration"]["adapters"][0]["provider"] = json!("custom");
let (_, store, runtime) = setup(config).await;
let snapshot = store.snapshot.lock().unwrap().clone();
assert!(
gateway_admin::ports::plugins::PluginPreparation::prepare(&runtime, snapshot)
.await
.is_err()
);
runtime.shutdown().await;
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "握手回调的错误类型由 tungstenite 固定"
)]
async fn websocket_continuation_reuses_exact_connection_and_consumes_each_handle_once() {
use futures::SinkExt as _;
use tokio_tungstenite::tungstenite::Message;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
let (socket, _) = listener.accept().await.unwrap();
let mut ws = tokio_tungstenite::accept_hdr_async(
socket,
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
assert_eq!(
request.headers()["authorization"],
"Bearer fixture-native-token"
);
Ok(response)
},
)
.await
.unwrap();
for _ in 0..2 {
assert!(matches!(
ws.next().await.unwrap().unwrap(),
Message::Text(_)
));
ws.send(Message::Text("{}".into())).await.unwrap();
}
assert!(ws.next().await.is_none_or(|result| result.is_err()));
});
let mut config = configuration(&base, false);
config["upstream_registration"]["adapters"][0]["transport"] = json!("websocket");
config["upstream_events"][1]["continuation"]["scope"] = json!("connection_local");
config["upstream_callbacks"] = json!([
{"method":"host.upstream.websocket.open","params":{"path":"responses","headers":[]}},
{"method":"host.upstream.websocket.send","params":{"kind":"text"},"body":"request"},
{"method":"host.upstream.websocket.read","params":{}}
]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut first = execute(&generation, Account::new(), "key-one", None);
let mut state = None;
while let Some(event) = first.next().await {
state = event.unwrap().take_session_update().or(state);
}
drop(first);
let state = state.unwrap();
assert!(state.extension_owner().unwrap().connection_local);
let error = execute(&generation, Account::new(), "key-two", Some(state.clone()))
.next()
.await
.unwrap()
.err()
.unwrap();
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
let mut second = execute(&generation, Account::new(), "key-one", Some(state.clone()));
while let Some(event) = second.next().await {
event.unwrap();
}
drop(second);
let error = execute(&generation, Account::new(), "key-one", Some(state))
.next()
.await
.unwrap()
.err()
.unwrap();
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap();
}
#[tokio::test]
async fn cancelling_adapter_body_closes_pending_websocket_and_releases_selected_account() {
use tokio_tungstenite::tungstenite::Message;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let (sent, received) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (socket, _) = listener.accept().await.unwrap();
let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
assert!(matches!(
socket.next().await.unwrap().unwrap(),
Message::Text(_)
));
sent.send(()).unwrap();
assert!(socket.next().await.is_none_or(|result| result.is_err()));
});
let mut config = configuration(&base, false);
config["upstream_registration"]["adapters"][0]["transport"] = json!("websocket");
config["upstream_callbacks"] = json!([
{"method":"host.upstream.websocket.open","params":{"path":"responses","headers":[]}},
{"method":"host.upstream.websocket.send","params":{"kind":"text"},"body":"request"},
{"method":"host.upstream.websocket.read","params":{}}
]);
let (cache, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let mut stream = execute(&generation, account.clone(), "key-one", None);
let mut reading = Box::pin(stream.next());
tokio::select! {
result = &mut reading => panic!("上游等待期间不应提前产生结果:{}", result.is_some()),
result = tokio::time::timeout(Duration::from_secs(3), received) => result.unwrap().unwrap(),
}
drop(reading);
drop(stream);
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap();
tokio::time::timeout(Duration::from_secs(2), async {
while Arc::strong_count(&account) != 1 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
#[tokio::test]
async fn usage_without_response_model_still_uses_native_account_pricing() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.mount(&server)
.await;
let mut config = configuration(&server.uri(), false);
config["upstream_events"][0]["event"]["facts"][0]["model"] = Value::Null;
config["upstream_events"][1]["event"]["facts"][1]["model"] = Value::Null;
let (_, _, runtime) = setup(config).await;
let generation = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let mut stream = execute(&generation, Account::new(), "key-one", None);
let mut costs = 0;
while let Some(event) = stream.next().await {
costs += event
.unwrap()
.canonical_facts()
.iter()
.filter(|fact| matches!(fact, GatewayEvent::CalculatedCost(_)))
.count();
}
assert_eq!(costs, 1);
drop(stream);
drop(generation);
runtime.shutdown().await;
}
#[tokio::test]
async fn core_published_adapter_uses_normal_attempt_usage_and_cost_ledger() {
use crate::support::environment::{Environment, account_grant};
let Some(mut environment) = Environment::create_command().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(header("authorization", "Bearer fixture-native-token"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.expect(1)
.mount(&server)
.await;
let account = environment.account(None).await;
let group = environment.account_group_with_account(&account).await;
let key_id = "key_adapter_ledger";
environment
.client_key_with_limits_and_groups(
key_id,
"sk-adapter-ledger-fixture",
gateway_core::policy::RateLimits::unlimited(),
vec![group],
)
.await;
environment
.install_plugin(
configuration(&server.uri(), false),
vec![
account_grant("requests"),
account_grant("upstream_connections"),
],
)
.await;
environment.install_plugin(json!({
"plugin_id":"test.adapter-ledger-caller",
"command_registration":{"commands":[{"name":"run","description":"执行模型"}]},
"command_nested_model_fixture":{
"request":{"client_key_id":key_id,"model":crate::support::native::MODEL,"protocol":"openai","operation":"generate","provider":"openai","account_id":account.as_str()},
"body":{"model":crate::support::native::MODEL,"input":"adapter ledger"}
},
"command_result":{"stdout":"ok","stderr":"","exit_code":0}
}), vec![account_grant("models"), account_grant("accounts")]).await;
let (runtime, core) = environment.command_plane().await;
let snapshot = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap();
let caller = snapshot
.instances
.iter()
.find(|instance| instance.configuration.get("command_registration").is_some())
.unwrap();
let commands = runtime.prepare_command_line().await.unwrap();
environment.store.start_command_line_writes().unwrap();
let result = commands.execute(&caller.id, "run", &[]).await.unwrap();
assert_eq!(result.stdout, "ok");
commands.shutdown().await;
environment
.store
.shutdown_command_line_writes()
.await
.unwrap();
environment
.assert_upstream_accounting(key_id, account.as_str())
.await;
drop(core);
drop(runtime);
environment.close().await;
}
#[tokio::test]
async fn adapter_reconfiguration_disable_and_rollback_keep_inflight_generation_and_reap_workers() {
let original = MockServer::start().await;
let replacement = MockServer::start().await;
for (server, expected) in [(&original, 2), (&replacement, 1)] {
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.expect(expected)
.mount(server)
.await;
}
let (cache, store, runtime) = setup(configuration(&original.uri(), false)).await;
let old = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
let account = Account::new();
let mut old_stream = execute(&old, account.clone(), "key-one", None);
old_stream.next().await.unwrap().unwrap();
{
let mut snapshot = store.snapshot.lock().unwrap();
snapshot.config_revision = gateway_admin::model::Revision::new(2).unwrap();
snapshot.instances[0].revision = snapshot.config_revision;
snapshot.instances[0].configuration = configuration(&replacement.uri(), false);
}
let new = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(2).unwrap())
.await
.unwrap();
assert_ne!(old.id(), new.id());
let mut new_stream = execute(&new, account.clone(), "key-one", None);
new_stream.next().await.unwrap().unwrap();
drop(old);
let old_state = old_stream
.next()
.await
.unwrap()
.unwrap()
.take_session_update()
.unwrap();
assert!(old_stream.next().await.is_none());
drop(old_stream);
{
let mut snapshot = store.snapshot.lock().unwrap();
snapshot.config_revision = gateway_admin::model::Revision::new(3).unwrap();
snapshot.instances[0].revision = snapshot.config_revision;
snapshot.instances[0].enabled = false;
}
let disabled = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(3).unwrap())
.await
.unwrap();
assert!(
context(&disabled, "key-one")
.upstream_adapter(
&ProviderKind::new("openai").unwrap(),
&UpstreamModelId::new("native-model").unwrap()
)
.unwrap()
.is_none()
);
drop(new);
assert!(
new_stream
.next()
.await
.unwrap()
.unwrap()
.canonical_facts()
.iter()
.any(|fact| matches!(fact, GatewayEvent::Completed(_)))
);
assert!(new_stream.next().await.is_none());
drop(new_stream);
drop(disabled);
super::wait_until_empty(cache.path()).await;
assert_eq!(Arc::strong_count(&account), 1);
{
let mut snapshot = store.snapshot.lock().unwrap();
snapshot.config_revision = gateway_admin::model::Revision::new(4).unwrap();
snapshot.instances[0].revision = snapshot.config_revision;
snapshot.instances[0].enabled = true;
snapshot.instances[0].configuration = configuration(&original.uri(), false);
}
let restored = ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(4).unwrap())
.await
.unwrap();
// 恢复同一配置仍是新代次,不能复活旧进程中的续接身份。
let error = execute(&restored, account.clone(), "key-one", Some(old_state))
.next()
.await
.unwrap()
.unwrap_err();
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
let mut fresh = execute(&restored, account.clone(), "key-one", None);
while let Some(event) = fresh.next().await {
event.unwrap();
}
drop(fresh);
drop(restored);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
assert_eq!(Arc::strong_count(&account), 1);
}
@@ -6,30 +6,51 @@ use gateway_plugin_sdk::{Capability, Contributions, Permission, Stage};
use sha2::{Digest as _, Sha256};
#[test]
fn host_compatibility_covers_every_sdk_capability_and_permission() {
fn host_compatibility_only_advertises_known_contracts() {
let compatibility: PluginHostCompatibility =
serde_json::from_str(include_str!("../../plugin-host-compatibility.json"))
.expect("compatibility JSON");
assert!(compatibility.is_valid());
for capability in all_capabilities() {
assert_exhaustive_capability(capability);
let identifier = identifier(capability);
assert!(
compatibility.supports_capability(&identifier, 1),
"missing capability {identifier}"
);
assert_eq!(
compatibility.manifest_schema_versions,
[gateway_plugin_sdk::MANIFEST_VERSION]
);
assert_eq!(
compatibility.protocol_versions,
[gateway_plugin_sdk::PROTOCOL_VERSION]
);
for entry in &compatibility.capabilities {
let capability: Capability = serde_json::from_value(entry.capability.clone().into())
.expect("host capability is understood by the SDK");
assert_eq!(entry.capability, capability.identifier());
for version in &entry.versions {
assert!(
capability.contract_versions().contains(version),
"unknown contract {} v{version}",
entry.capability
);
}
}
for permission in Permission::ALL {
assert_exhaustive_permission(permission);
let identifier = identifier(permission);
assert!(
compatibility.supports_permission(&identifier),
"missing permission {identifier}"
);
for identifier in &compatibility.permissions {
let permission: Permission = serde_json::from_value(identifier.clone().into())
.expect("host permission is understood by the SDK");
assert_eq!(identifier, permission.as_str());
}
}
#[test]
fn host_compatibility_preserves_both_middleware_contracts() {
let compatibility: PluginHostCompatibility =
serde_json::from_str(include_str!("../../plugin-host-compatibility.json"))
.expect("compatibility JSON");
assert!(compatibility.supports_capability("middleware", 1));
assert!(compatibility.supports_capability("middleware", 2));
assert!(!compatibility.supports_capability("middleware", 3));
assert!(!compatibility.supports_capability("openai", 1));
assert!(!compatibility.supports_permission("xai"));
}
#[tokio::test]
async fn package_inspector_returns_static_requirements_without_starting_the_plugin() {
let archive = crate::support::package_with_contributions(
@@ -61,60 +82,3 @@ async fn package_inspector_returns_static_requirements_without_starting_the_plug
assert_eq!(requirements.capabilities, vec![("middleware".into(), 1)]);
assert!(requirements.permissions.is_empty());
}
fn identifier(value: impl serde::Serialize) -> String {
serde_json::to_value(value)
.expect("serialize identifier")
.as_str()
.expect("string identifier")
.to_owned()
}
fn assert_exhaustive_capability(capability: Capability) {
match capability {
Capability::FrontendAuthentication
| Capability::Scheduler
| Capability::ModelRouter
| Capability::ModelCatalog
| Capability::RetryPolicy
| Capability::Middleware
| Capability::RequestLifecycle
| Capability::WebSocketObserver
| Capability::Usage
| Capability::CommandLine
| Capability::Management
| Capability::Maintenance => {}
}
}
fn assert_exhaustive_permission(permission: Permission) {
match permission {
Permission::Network
| Permission::Requests
| Permission::Models
| Permission::Accounts
| Permission::Data
| Permission::PublicEndpoints
| Permission::Groups
| Permission::Keys
| Permission::KeyBudgets
| Permission::QuotaObservations => {}
}
}
fn all_capabilities() -> [Capability; 12] {
[
Capability::FrontendAuthentication,
Capability::Scheduler,
Capability::ModelRouter,
Capability::ModelCatalog,
Capability::RetryPolicy,
Capability::Middleware,
Capability::RequestLifecycle,
Capability::WebSocketObserver,
Capability::Usage,
Capability::CommandLine,
Capability::Management,
Capability::Maintenance,
]
}
@@ -147,6 +147,49 @@ async fn registration_and_configuration_reject_network_and_model_callbacks() {
session.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn upstream_connections_require_their_own_grant_and_exact_execution_stage() {
let methods = [
"host.upstream.http.do",
"host.upstream.http.do_stream",
"host.upstream.http.stream_read",
"host.upstream.http.stream_close",
"host.upstream.websocket.open",
"host.upstream.websocket.send",
"host.upstream.websocket.read",
"host.upstream.websocket.close",
];
let callbacks = Arc::new(Callbacks::default());
let (_cache, session) = session_with_permissions(
Arc::clone(&callbacks),
vec![Permission::UpstreamConnections],
)
.await;
for method in methods {
assert!(
invoke_callback(&session, Stage::Upstream, method)
.await
.is_ok()
);
for stage in [
Stage::Request,
Stage::Attempt,
Stage::Management,
Stage::Observation,
Stage::Registration,
] {
assert_permission_denied(invoke_callback(&session, stage, method).await);
}
}
session.shutdown(Duration::from_secs(1)).await;
let (_cache, session) =
session_with_permissions(callbacks, vec![Permission::Network, Permission::Accounts]).await;
for method in methods {
assert_permission_denied(invoke_callback(&session, Stage::Upstream, method).await);
}
session.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn active_stages_use_resource_domains_without_operation_whitelists() {
let callbacks = Arc::new(Callbacks::default());
@@ -214,6 +214,16 @@ impl Environment {
vec![],
)]);
}
let has_upstream_adapter = configuration.get("upstream_registration").is_some();
if has_upstream_adapter {
contributes.extend([super::contribution_for_id(
&plugin_id,
Capability::UpstreamAdapter,
vec![Stage::Upstream],
vec!["openai".into()],
vec!["openai".into()],
)]);
}
let has_middleware = configuration.get("middleware_marker").is_some();
if has_middleware {
contributes.extend([super::contribution_for_id(
@@ -322,6 +332,17 @@ impl Environment {
secrets: BTreeMap::new(),
grants,
bindings: [
has_upstream_adapter.then(|| PluginCapabilityBinding {
contribution: format!("{plugin_id}.upstreamAdapter"),
stage: "upstream".into(),
order: 0,
failure_policy: PluginFailurePolicy::Reject,
client_key_ids: vec![],
account_group_ids: vec![],
provider_ids: vec![],
models: vec![],
identity_bindings: vec![],
}),
has_model_router.then(|| PluginCapabilityBinding {
contribution: format!("{plugin_id}.modelRouter"),
stage: "routing".into(),
@@ -667,6 +688,16 @@ impl Environment {
.unwrap()
}
pub async fn assert_upstream_accounting(&self, key: &str, account: &str) {
let rows: Vec<(String, String, i32, i64, i64, bool)> = sqlx::query_as(sqlx::AssertSqlSafe(format!(
"select outcome, provider_account_ref, attempt_count, input_tokens, output_tokens, cost_amount > 0 from {}.model_requests where client_api_key_ref=$1", self.schema,
))).bind(key).fetch_all(&self.admin).await.unwrap();
assert_eq!(
rows,
vec![("succeeded".into(), account.into(), 1, 7, 2, true)]
);
}
pub async fn bound_model_requests(&self, client_key_id: &str) -> Vec<(String, String)> {
sqlx::query_as(sqlx::AssertSqlSafe(format!(
"select request_kind, subagent_kind from {}.model_requests
@@ -45,6 +45,7 @@ pub fn contribution_for_id(
Capability::ModelCatalog => "modelCatalog",
Capability::RetryPolicy => "retryPolicy",
Capability::Middleware => "middleware",
Capability::UpstreamAdapter => "upstreamAdapter",
Capability::RequestLifecycle => "requestLifecycle",
Capability::WebSocketObserver => "webSocketObserver",
Capability::Usage => "usage",
@@ -73,12 +73,31 @@ impl Provider for NativeProvider {
let model = candidate.upstream_model().cloned().ok_or_else(|| {
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
let adapter = context.upstream_adapter(candidate.provider(), &model)?;
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
model,
account,
UpstreamTransport::new("http_sse").unwrap(),
UpstreamTransport::new(
adapter
.as_ref()
.map_or("http_sse", |adapter| adapter.transport()),
)
.unwrap(),
);
if let Some(adapter) = adapter {
let account = Arc::new(FixtureConnection(metadata.provider_account_id().clone()));
let events = adapter.execute(
gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation {
operation: request.operation().clone(),
headers: vec![],
context,
metadata: metadata.clone(),
account,
},
);
return Ok(ProviderStream::new(metadata, events, ()));
}
let response = ResponseMeta::new("fixture-response", MODEL);
let event = |name: &str, fact| {
ProviderEvent::canonical_with_wire(
@@ -108,6 +127,47 @@ impl Provider for NativeProvider {
}
}
// 仅测试 Core / Runtime / Store 组合;真实 Provider 的凭据解释由各自合同测试验证。
struct FixtureConnection(ProviderAccountId);
impl gateway_core::engine::upstream_adapter::UpstreamAccountConnection for FixtureConnection {
fn account_id(&self) -> &ProviderAccountId {
&self.0
}
fn credential_revision(&self) -> gateway_core::account::CredentialRevision {
gateway_core::account::CredentialRevision::new(1).unwrap()
}
fn authentication_kind(&self) -> &str {
"oauth"
}
fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
None
}
fn authorization(
&self,
) -> Result<Vec<gateway_core::engine::middleware::MiddlewareHeader>, ProviderError> {
Ok(vec![
gateway_core::engine::middleware::MiddlewareHeader::new(
"authorization",
bytes::Bytes::from_static(b"Bearer fixture-native-token"),
),
])
}
fn calculate_cost(
&self,
_: Option<&str>,
_: &gateway_core::metering::Usage,
) -> Option<gateway_core::metering::CalculatedCost> {
gateway_core::metering::CalculatedCost::from_usd_ticks(123).ok()
}
fn record_failure(
&self,
error: ProviderError,
) -> futures::future::BoxFuture<'_, ProviderError> {
Box::pin(async move { error })
}
}
pub fn provider_registry() -> ProviderRegistry {
ProviderRegistry::new([
Arc::new(NativeProvider("openai")) as Arc<dyn Provider>,
@@ -18,6 +18,11 @@ use serde_json::{Value, json};
use tokio::io::AsyncWriteExt as _;
use tokio::sync::{Mutex, Notify, mpsc, oneshot};
use gateway_plugin_sdk::call::{
model::{ExecutionEvent, WireEvent, WirePayload},
upstream_adapter::{UpstreamAdapterEvent, UpstreamAdapterRequest},
};
type CallbackResult = Result<(Value, Vec<u8>), PluginFault>;
struct PendingCallback {
parent: u64,
@@ -396,6 +401,23 @@ impl Peer {
return;
}
match method.as_str() {
"upstream_adapter.register" => {
self.send(
Message::Result {
id,
result: json!({}),
},
serde_json::to_vec(&self.configuration["upstream_registration"]).unwrap(),
)
.await;
return;
}
"upstream_adapter.execute" => {
if let Err(error) = self.upstream_adapter(id, payload).await {
self.send(Message::Error { id, error }, vec![]).await;
}
return;
}
"plugin.reconcile" => {
let result = if let Some(results) = self.data_queries(id).await {
self.append_observation_marker(
@@ -1308,7 +1330,12 @@ async fn main() {
if method == "hang_uncancellable" {
uncancellable.insert(id);
}
if method.starts_with("stream") || matches!(method.as_str(), "middleware.handle") {
if method.starts_with("stream")
|| matches!(
method.as_str(),
"middleware.handle" | "upstream_adapter.execute"
)
{
peer.streams
.lock()
.await
@@ -1385,3 +1412,101 @@ async fn main() {
}
}
}
impl Peer {
async fn upstream_adapter(&self, id: u64, payload: Vec<u8>) -> Result<(), PluginFault> {
let (request, body) = UpstreamAdapterRequest::decode(&payload).unwrap();
// 假凭据用于检测宿主是否把已选账号令牌放入了插件输入。
assert!(!String::from_utf8_lossy(&payload).contains("fixture-native-token"));
self.append_observation_marker("upstream_marker", &json!({"key":request.client_key_id,"account":request.account_id,"continuation":request.continuation}));
self.send(
Message::Result {
id,
result: json!({}),
},
vec![],
)
.await;
let mut response_body = vec![];
for callback in self.configuration["upstream_callbacks"]
.as_array()
.into_iter()
.flatten()
{
let payload = if callback["body"] == "request" {
body.clone()
} else {
vec![]
};
let (result, payload) = self
.callback_payload(
id,
callback["method"].as_str().unwrap(),
callback["params"].clone(),
payload,
)
.await?;
if let Some(stream) = result["stream"].as_str() {
loop {
let (result, chunk) = self
.callback_payload(
id,
"host.upstream.http.stream_read",
json!({"stream":stream,"maximum_bytes":65536}),
vec![],
)
.await?;
response_body.extend(chunk);
if result["eof"] == true {
break;
}
}
} else if !payload.is_empty() {
response_body = payload;
}
}
let events = self.configuration["upstream_events"].as_array().unwrap();
let credits = self.streams.lock().await.get(&id).unwrap().clone();
for (sequence, event) in events.iter().enumerate() {
let mut message = UpstreamAdapterEvent::new(
serde_json::from_value::<ExecutionEvent>(event["event"].clone()).unwrap(),
);
message.service_tier = event["service_tier"].as_str().map(str::to_owned);
message.continuation = event
.get("continuation")
.map(|value| serde_json::from_value(value.clone()).unwrap());
message.failure = event
.get("failure")
.map(|value| serde_json::from_value(value.clone()).unwrap());
if event["wire"] == "http" {
message.event.wire = Some(WireEvent {
protocol: request.protocol.clone(),
payload: WirePayload::RawJson {
body: response_body.clone(),
},
});
} else if event["wire"] == "sse" {
message.event.wire = Some(WireEvent {
protocol: request.protocol.clone(),
payload: WirePayload::RawSse {
frame: response_body.clone(),
},
});
}
let payload = message.encode().unwrap();
credits.take(payload.len() as u64).await;
self.send(
Message::Stream {
id,
sequence: sequence as u64,
},
payload,
)
.await;
}
let error = (self.configuration["upstream_tail_error"] == true)
.then(|| PluginFault::new(ErrorCode::Fault, "fixture terminal fault"));
self.send(Message::End { id, error }, vec![]).await;
Ok(())
}
}
+4 -2
View File
@@ -25,6 +25,7 @@ SDK 仍处于实验阶段;能力是否可用取决于宿主支持、清单声
| `call::frontend_authentication` | 数据面认证信封、认证器标识与外部 principal 结果 |
| `call::model` | 宿主模型调用的 canonical 事实、原生 wire 与有界事件编解码 |
| `call::middleware` | 洋葱中间件请求/响应 head、single-use next 与惰性正文 frame 合同 |
| `call::upstream_adapter` | 内置 OpenAI / xAI 的[受管上游适配器](docs/upstream-adapters.md)注册、执行、出站与续接合同 |
| `call::policy` | 模型路由、账号调度、受约束重试及终态用量观察合同 |
| `call::catalog` | 固定 Provider 的模型别名注册合同 |
| `call::data` | 最小账号事实与已有额度观测的只读合同 |
@@ -33,8 +34,9 @@ SDK 仍处于实验阶段;能力是否可用取决于宿主支持、清单声
| `call::management` | 管理 API、页面与资源声明、公开回调,以及 CLI 命令与待保存账号 |
| `client` | 通过 `io` feature 开启的 `PluginBuilder`、会话、中间件与异步帧收发 |
`call/` 按业务能力组织数据,`message.rs` 定义消息与帧。请求/响应改写、协议转换和思考参数映射统一使用
`call::middleware`;WebSocket 观察接口不能改写连接
`call/` 按业务能力组织数据,`message.rs` 定义消息与帧。已有请求链的请求/响应改写、协议转换和思考参数映射使用
`call::middleware`;需要指定业务上游并解析结果时使用 `call::upstream_adapter`,作为洋葱链终端复用宿主账号与结算。
WebSocket 观察接口不能改写连接
插件使用的导入路径例如:
@@ -9,7 +9,7 @@
| --- | --- |
| 组合处理器并声明权限 | [作者入口](#类型化作者入口) · [能力与方法](#能力与方法) · [访问域](#访问域) |
| 使用宿主资源 | [基础事实](#基础事实) · [额度刷新](#额度观测刷新) · [账号](#账号与凭据) · [自有资源](#自有资源与维护) · [Key 预算](#client-key-预算) · [模型](#key模型与模型调用) · [网络](#网络) |
| 扩展请求链 | [中间件](#洋葱中间件) · [模型目录](#模型目录) · [重试](#重试决策) · [路由、调度与观察](#路由调度与观察) · [入口认证](#数据面入口认证) |
| 扩展请求链 | [中间件](#洋葱中间件) · [上游适配器](upstream-adapters.md) · [模型目录](#模型目录) · [重试](#重试决策) · [路由、调度与观察](#路由调度与观察) · [入口认证](#数据面入口认证) |
| 管理与交互 | [状态、日志与迁移](#状态日志与迁移) · [管理页面](#管理页面公开入口与-cli) · [页面宿主桥](#页面宿主桥-v2) · [命令行](#命令行) |
## 类型化作者入口
@@ -52,6 +52,7 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
| 能力 | 类型化方法 |
| --- | --- |
| `middleware` | `PluginBuilder::middleware` / `middleware.handle` |
| `upstream_adapter` | `methods::UPSTREAM_ADAPTER_REGISTER`、`methods::UPSTREAM_ADAPTER_EXECUTE`,见[受管上游](upstream-adapters.md) |
| `model_router` | `methods::ROUTE_MODEL` |
| `model_catalog` | `PluginBuilder::model_catalog` / `methods::MODEL_CATALOG_REGISTER` |
| `retry_policy` | `methods::RETRY_DECISION` |
@@ -71,6 +72,7 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
| 权限 | 开放的宿主能力 |
| --- | --- |
| `network` | 通过 `host.http.*` 使用宿主受管出站网络;仍受统一代理、超时、大小和流控规则约束 |
| `upstream_connections` | 仅在 `upstream` 阶段使用已选账号的受管 HTTP / WebSocket;宿主注入凭据和账号代理,不能读取原始凭据或重新选号 |
| `models` | 列出非秘密 Key、按所选 Key 查询模型,以及通过 `host.model.*` 调用模型;调用可能产生消耗 |
| `accounts` | 查询账号、读取原始凭据及创建或替换账号;写入仍经过 revision CAS、审计和发布事务 |
| `data` | 在管理/命令/维护阶段只读账号、Key 的最小基础信息及已有额度观测,不包含凭据、写入或预测 |
@@ -244,6 +246,10 @@ SDK 不自动重试。超时或断连不能证明写入未提交;重试上限
### 网络
`HostClient::http(request, body)` 返回响应头与 `HostHttpBody`。正文用 `read()` 按需读取,
用 `collect(maximum_bytes)` 有界收集,提前结束用 `close()`;EOF 与重复关闭不再调用宿主。
父调用取消或结束时宿主回收未关闭的流,SDK 不后台预读。账号上游使用同一正文对象,权限与目标见[受管上游](upstream-adapters.md)
`network` 域通过 `host.http.do/do_stream` 和流读取/关闭方法提供受管 HTTP。URL、header 和正文可能含敏感数据,
正文使用独立二进制载荷。宿主统一施加代理、超时、帧与流限制;插件不能绕过宿主端口取得额外调用身份。
地址被拒绝返回 `permission_denied`,期限耗尽返回 `timeout`,解析、连接或响应读取失败返回 `upstream`。
@@ -291,6 +297,21 @@ restore_and_validate_response(response).await
仅实现中间件时也可使用轻量的 `MiddlewarePlugin`;需要与管理或其他方法组合时使用
`PluginBuilder`,两者复用相同的 `MiddlewareCall` 和 `MiddlewareResponse`
### 扩展的组合边界
请求处理沿 `request middleware → 原生选路与选号 → attempt middleware → 原生上游或 upstream_adapter` 推进,
响应按相反顺序返回;适配器是现有洋葱链的终端,不能再次调用 `next` 或组织换号重试
| 工作 | 组合方式 |
| --- | --- |
| 请求头/正文改写、提前拒绝、响应头处理 | `middleware` 包裹一次 `next.run(request)` |
| 响应逐帧转换、观察 | 返回惰性 `map_frames` / `inspect_frames`,不能在 `next` 返回时就把流视为结束 |
| 新上游路径、协议编码、错误及用量解析 | `upstream_adapter` 终端,复用已选账号和原生结算 |
| 目标限制、凭据注入、网络发送状态 | 宿主网络外层,普通 HTTP 与账号 HTTP 共用发送和正文资源实现 |
| 路由、调度、重试决策 | 各自的决策端口,Core 保持最终裁决和预算 |
| RPC Credit、取消、流关闭、WebSocket 续接 | 会话及资源生命周期,不作为可重排的中间件 |
| 终态观察、管理与维护 | 原有观察/管理/维护入口,不嵌入请求 `next` 链 |
### 模型目录
`model_catalog` v1 在 registration 阶段一次提交 `ModelCatalogRegistration { models: Vec<ModelAlias> }`,不需要请求级 binding。
@@ -22,6 +22,18 @@
`package_for(host, os, architecture)` 检查版本与平台。运行模式仅为 `trustedProcess`:插件拥有与宿主
相同的系统身份,**不是进程沙箱**
宿主实际开放的合同由 [plugin-host-compatibility.json](../../runtime/plugin-host-compatibility.json)
声明,SDK 的 `Capability::contract_versions()` 只表示 SDK 能描述的行为版本,不能代替宿主支持检查
| 版本字段 | 对应合同 |
| --- | --- |
| `manifestVersion` / `manifest_schema_versions` | 插件清单格式 |
| `package.protocolVersion` / `protocol_versions` | 插件进程 RPC 封装,不是 OpenAI 或 xAI 的业务协议 |
| `contributes.<capability>.version` / `capabilities[].versions` | 指定扩展能力的行为合同 |
| 宿主声明的 `schema_version` | 宿主兼容声明自身的格式 |
`capabilities` 使用扩展能力标识,`permissions` 使用资源访问域标识;Provider、模型和账号 ID 不属于这两个集合
## 扩展项简写
普通作者声明可以省略 `id`、`version` 和固定阶段:
@@ -0,0 +1,86 @@
# 受管上游适配器
`upstream_adapter` v1 是现有生成请求洋葱链的终端,只挂到内置 `openai` / `xai`,不替换独立图像或压缩端点。
插件定义 Base URL、业务路径、请求编码和响应解析;宿主继续持有身份、ACL、选号、租约、凭据、代理、重试和账本。
普通请求或响应加工使用 [middleware](capabilities.md#洋葱中间件)
## 声明与绑定
清单声明 `upstream_adapter`、版本 `1`、阶段 `upstream`,以及匹配的 `input_formats` / `output_formats`;
需要 `requests` 与 `upstream_connections`。使用 `PluginBuilder::on` 注册
`methods::UPSTREAM_ADAPTER_REGISTER` 与 `methods::UPSTREAM_ADAPTER_EXECUTE`
注册返回 `UpstreamAdapterRegistration`,包含 1 至 16 个适配器:
| 字段 | 含义 |
| --- | --- |
| `id` | 实例内唯一标识,不是 Provider 或账号类型 |
| `provider` | `openai` 或 `xai` |
| `base_url` | 以 `/` 结尾的 HTTP(S) 基址,无用户信息、查询串或 fragment |
| `paths` | 不以 `/` 开头的基址相对路径,含 `inference` 或 `auxiliary` 用途,无通配符 |
| `authentication_kinds` | 可使用的现有认证类型,宿主复核已选账号 |
| `transport` | `http_json`、`http_sse` 或 `websocket` |
| `protocol` | 请求和客户端 wire 使用的协议,须在清单输入、输出格式内 |
| `models` | 公开请求模型的精确匹配,空集合沿用绑定范围 |
实例绑定沿用 Key、账号组、Provider、模型范围,只允许一次 `upstream` 绑定和 `reject` 故障策略。
同一请求不能同时命中两个适配器;配置准备时拒绝重叠,执行时再次复核。适配器未命中时仍使用原生上游,
已命中但不可用时拒绝请求。安装、授权、配置、启停和升级使用已有插件管理
能力版本、清单版本、进程 RPC 版本和宿主兼容声明格式版本分别判断,见[清单](manifest.md)。
Provider、业务协议和传输名称不能写入宿主兼容清单的 capability / permission / RPC 版本位置
## 执行与受管出站
`TypedCall<UpstreamAdapterRequest>` 包含只读的 Key、账号、凭据版本、模型、协议、Header 和宿主续接投影;
原始请求正文位于 `call.payload`,原始账号凭据不进入插件输入。宿主首次消费冷流才调用插件
`disable_fast` 是宿主冻结的请求策略。OpenAI 请求正文在进入适配器前复用原生 Fast 策略处理;
适配器转换为其他上游协议时也须遵守该字段,不能重新启用 Fast / priority 档位
HTTP 请求使用 `call.host.upstream_http(request, body).await?`,返回 `HostHttpResponse`:
```rust,ignore
let response = call.host.upstream_http(UpstreamHttpRequest {
method: "POST".into(), path: "responses".into(), query: vec![], headers: vec![],
}, encode_request(&call.payload)?).await?;
let status = response.status;
let mut body = response.body;
while let Some(bytes) = body.read().await? {
// 分段不等于 SSE 事件;插件按自己的上游协议增量解码。
decode_upstream(status, bytes)?;
}
```
正文与普通 `HostClient::http` 共用 `read`、`collect(maximum_bytes)` 和 `close`;每次读取至多 64 KiB,
不后台预读,不向作者暴露流 ID。提前关闭及收集超限会释放正文,父调用结束或取消时宿主兜底回收。
JSON 可有界收集,SSE 应增量解析,不能把读取块当成完整事件
WebSocket 使用 `call.host.upstream_websocket(request).await?`,返回
`UpstreamWebSocketUpgrade::Connected { headers, connection }` 或 `Rejected { status, headers, body }`。
连接提供 `send(kind, body)`、`read()`、`close()`,支持完整文本与二进制消息,Ping/Pong 由宿主处理。
同一连接一次只执行一个操作,握手拒绝保留 HTTP 状态和正文。底层 `host.upstream.*` 是 SDK 使用的资源协议,
不是额外的业务中间件链
宿主只对声明的基址和精确路径注入已选账号认证,禁止覆盖 Host、认证、Cookie 与握手保护头;
重定向不能携带凭据逃逸。出站复用该账号的代理,代理失败不能回退直连。
`auxiliary` 路径的出站计入请求副作用水位,防止附件上传等操作在失败后被透明重复执行
## 事件、结算和续接
执行返回 `TypedReply<Empty>` 的有界 `ResponseStream`。每个 chunk 使用 `UpstreamAdapterEvent::encode()`,
包含标准事实与可选原始 wire;事件需符合开始、内容、工具、用量和唯一结束的顺序。
原始 JSON / SSE 字节与标准事实分开传递,错误使用 `UpstreamFailure`;宿主网络观测决定发送状态,插件不能将其降为未发送
用量提交标准 `Usage`,费用由内置 Provider 按已冻结的上游模型与价格计算,响应回显模型不切换计价。
`service_tier` 记录上游回显的最新档位,允许上游在完成时确定实际值;OpenAI 计价沿用请求策略,回显档位仅作观测。
未知用量或价格保留未知。宿主确认 RPC 正常结束后才发布完成事实,插件不能独立结算或组织换号重试
成功终态可附带 `UpstreamContinuation`,包含上游响应 ID、不超过 32 KiB 的私有续接材料和作用域:
- `persisted`:续接材料随宿主会话保存,仍受实例、代次、Key、账号和凭据版本限制
- `connection_local`:还必须保有本次已建立的 WebSocket;宿主在成功终态后保存同一连接,续接时独占取回
连接内状态不能跨账号、Key、凭据版本、插件进程或代次复用,失效时拒绝续接,不重新建立连接假装原生续接成功。
每个适配器代次最多保留 64 条闲置连接,闲置期限 30 分钟;保留连接不占用上一轮账号租约。
显式 `close()` 放弃连接内续接。普通成功、错误、取消与退出释放未保留连接
@@ -12,3 +12,4 @@ pub mod observation;
pub mod policy;
pub mod registration;
pub mod resources;
pub mod upstream_adapter;
@@ -20,7 +20,7 @@ pub(super) fn decode_event(bytes: &[u8]) -> Result<ExecutionEvent, ExecutionEnco
Ok(event)
}
fn pack<T: Serialize, const N: usize>(
pub(in crate::call) fn pack<T: Serialize, const N: usize>(
prefix: [u8; 4],
metadata: &T,
payloads: [Vec<u8>; N],
@@ -55,7 +55,7 @@ fn pack<T: Serialize, const N: usize>(
Ok(result)
}
fn unpack<T: DeserializeOwned, const N: usize>(
pub(in crate::call) fn unpack<T: DeserializeOwned, const N: usize>(
prefix: [u8; 4],
bytes: &[u8],
) -> Result<(T, [&[u8]; N]), ExecutionEncodingError> {
@@ -2,7 +2,7 @@
use serde::{Deserialize, Serialize};
mod codec;
pub(in crate::call) mod codec;
pub use codec::{ExecutionEncodingError, MAX_EXECUTION_PAYLOAD_BYTES};
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
@@ -0,0 +1,48 @@
use serde::{Deserialize, Serialize};
use super::{UpstreamAdapterEvent, UpstreamAdapterRequest, UpstreamContinuation, UpstreamFailure};
use crate::call::model::{
ExecutionEncodingError, ExecutionEvent,
codec::{pack, unpack},
};
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Metadata {
service_tier: Option<String>,
continuation: Option<UpstreamContinuation>,
failure: Option<UpstreamFailure>,
}
pub(super) fn encode(value: UpstreamAdapterEvent) -> Result<Vec<u8>, ExecutionEncodingError> {
let metadata = Metadata {
service_tier: value.service_tier,
continuation: value.continuation,
failure: value.failure,
};
pack(*b"GPA1", &metadata, [value.event.encode()?])
}
pub(super) fn decode(bytes: &[u8]) -> Result<UpstreamAdapterEvent, ExecutionEncodingError> {
let (metadata, [event]): (Metadata, _) = unpack(*b"GPA1", bytes)?;
Ok(UpstreamAdapterEvent {
event: ExecutionEvent::decode(event)?,
service_tier: metadata.service_tier,
continuation: metadata.continuation,
failure: metadata.failure,
})
}
pub(super) fn encode_request(
request: UpstreamAdapterRequest,
body: Vec<u8>,
) -> Result<Vec<u8>, ExecutionEncodingError> {
pack(*b"GPAQ", &request, [body])
}
pub(super) fn decode_request(
bytes: &[u8],
) -> Result<(UpstreamAdapterRequest, Vec<u8>), ExecutionEncodingError> {
let (request, [body]) = unpack(*b"GPAQ", bytes)?;
Ok((request, body.to_vec()))
}
@@ -0,0 +1,227 @@
//! 内置 OpenAI / xAI 账号上的上游转换合同,凭据不进入此线协议。
use serde::{Deserialize, Serialize};
use super::model::ExecutionEvent;
mod codec;
pub const REGISTER_METHOD: &str = "upstream_adapter.register";
pub const EXECUTE_METHOD: &str = "upstream_adapter.execute";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum BuiltinProvider {
OpenAi,
Xai,
}
impl BuiltinProvider {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::OpenAi => "openai",
Self::Xai => "xai",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UpstreamTransport {
HttpJson,
HttpSse,
#[serde(rename = "websocket")]
WebSocket,
}
impl UpstreamTransport {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::HttpJson => "http_json",
Self::HttpSse => "http_sse",
Self::WebSocket => "websocket",
}
}
}
/// 注册结果随实例配置冻结;同一实例可为两个内置 Provider 分别声明适配器。
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamAdapterRegistration {
pub adapters: Vec<UpstreamAdapterDeclaration>,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamAdapterDeclaration {
pub id: String,
pub provider: BuiltinProvider,
pub base_url: String,
/// 相对 Base URL 的精确业务路径,不接受绝对 URL、查询串或通配符。
pub paths: Vec<UpstreamPath>,
pub authentication_kinds: Vec<String>,
pub transport: UpstreamTransport,
/// 客户端请求及适配完成后的 wire 协议;标准事实独立于 wire 传递。
pub protocol: String,
/// 公开请求模型的匹配范围;空集合沿用功能绑定范围。
#[serde(default)]
pub models: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UpstreamPathPurpose {
Inference,
Auxiliary,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamPath {
pub path: String,
pub purpose: UpstreamPathPurpose,
}
/// 元数据为宿主只读投影;请求正文独立置于调用二进制载荷。
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamAdapterRequest {
pub adapter_id: String,
pub provider: BuiltinProvider,
pub upstream_model: String,
pub client_key_id: String,
pub account_id: String,
pub credential_revision: u64,
pub protocol: String,
pub client_transport: String,
/// 宿主冻结的强制策略;自定义上游编码同样不得请求 Fast / priority 档位。
pub disable_fast: bool,
pub headers: Vec<(String, Vec<u8>)>,
/// 只能来自同账号、同插件代次的宿主续接记录,不从客户端正文信任此值。
pub continuation: Option<UpstreamContinuation>,
}
impl UpstreamAdapterRequest {
/// 请求身份、业务头与私有续接材料同原始正文一起放入二进制载荷。
pub fn encode(self, body: Vec<u8>) -> Result<Vec<u8>, super::model::ExecutionEncodingError> {
codec::encode_request(self, body)
}
pub fn decode(bytes: &[u8]) -> Result<(Self, Vec<u8>), super::model::ExecutionEncodingError> {
codec::decode_request(bytes)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContinuationScope {
Persisted,
ConnectionLocal,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamContinuation {
pub scope: ContinuationScope,
pub upstream_response_id: String,
/// 适配器私有的续接材料,不包含账号凭据。
pub state: serde_json::Map<String, serde_json::Value>,
}
/// 一个完整 GPE1 事件与可选的结算/续接事实;见 [`Self::encode`]。
pub struct UpstreamAdapterEvent {
pub event: ExecutionEvent,
/// 上游回显的最新服务档位,可在终态确定;是否参与计价由内置 Provider 决定。
pub service_tier: Option<String>,
pub continuation: Option<UpstreamContinuation>,
pub failure: Option<UpstreamFailure>,
}
impl UpstreamAdapterEvent {
#[must_use]
pub const fn new(event: ExecutionEvent) -> Self {
Self {
event,
service_tier: None,
continuation: None,
failure: None,
}
}
/// 编码有界事件,wire 原字节仍使用 GPE1 二进制段。
pub fn encode(self) -> Result<Vec<u8>, super::model::ExecutionEncodingError> {
codec::encode(self)
}
/// 校验版本、长度及元数据,再解码事件载荷。
pub fn decode(bytes: &[u8]) -> Result<Self, super::model::ExecutionEncodingError> {
codec::decode(bytes)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UpstreamFailureKind {
InvalidRequest,
Unsupported,
Unauthorized,
PermissionDenied,
RateLimited,
QuotaExhausted,
Timeout,
Unavailable,
Protocol,
}
/// 解析后的上游拒绝事实;发送状态始终由宿主网络观测,不接受插件声明。
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamFailure {
pub kind: UpstreamFailureKind,
pub status: Option<u16>,
pub retry_after_ms: Option<u64>,
pub message: String,
pub code: Option<String>,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamHttpRequest {
pub method: String,
pub path: String,
pub headers: Vec<(String, String)>,
#[serde(default)]
pub query: Vec<(String, String)>,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamWebSocketRequest {
pub path: String,
pub headers: Vec<(String, String)>,
#[serde(default)]
pub query: Vec<(String, String)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WebSocketMessageKind {
Text,
Binary,
}
/// 消息数据独立放入二进制载荷,宿主处理 Ping/Pong 与关闭握手。
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamWebSocketMessage {
pub kind: WebSocketMessageKind,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct UpstreamWebSocketRead {
pub kind: Option<WebSocketMessageKind>,
pub eof: bool,
}
@@ -12,6 +12,7 @@ pub enum Capability {
ModelCatalog,
RetryPolicy,
Middleware,
UpstreamAdapter,
RequestLifecycle,
WebSocketObserver,
Usage,
@@ -21,6 +22,26 @@ pub enum Capability {
}
impl Capability {
/// SDK 能描述的行为合同版本;具体宿主可以只开放其中一部分。
#[must_use]
pub const fn contract_versions(self) -> &'static [u32] {
match self {
Self::Middleware => &[1, 2],
Self::FrontendAuthentication
| Self::Scheduler
| Self::ModelRouter
| Self::ModelCatalog
| Self::RetryPolicy
| Self::RequestLifecycle
| Self::WebSocketObserver
| Self::Usage
| Self::CommandLine
| Self::Management
| Self::Maintenance => &[1],
Self::UpstreamAdapter => &[1],
}
}
/// 稳定能力标识;默认扩展项 ID 由它派生,不受显示名称影响。
#[must_use]
pub const fn identifier(self) -> &'static str {
@@ -31,6 +52,7 @@ impl Capability {
Self::ModelCatalog => "model_catalog",
Self::RetryPolicy => "retry_policy",
Self::Middleware => "middleware",
Self::UpstreamAdapter => "upstream_adapter",
Self::RequestLifecycle => "request_lifecycle",
Self::WebSocketObserver => "web_socket_observer",
Self::Usage => "usage",
@@ -50,6 +72,7 @@ impl Capability {
Self::ModelCatalog => &[Stage::Registration],
Self::RetryPolicy => &[Stage::Retry],
Self::Middleware => &[],
Self::UpstreamAdapter => &[Stage::Upstream],
Self::RequestLifecycle | Self::WebSocketObserver | Self::Usage => &[Stage::Observation],
Self::CommandLine => &[Stage::CommandLine],
Self::Management => &[Stage::Management],
@@ -70,6 +93,8 @@ pub enum Stage {
Retry,
Request,
Attempt,
/// Core 登记已选账号的 attempt 后才启动的受管上游执行。
Upstream,
Observation,
Management,
CommandLine,
@@ -153,11 +178,12 @@ pub enum Permission {
Keys,
KeyBudgets,
QuotaObservations,
UpstreamConnections,
}
impl Permission {
/// 当前公开访问域,安装摘要与授权校验复用同一集合。
pub const ALL: [Self; 10] = [
pub const ALL: [Self; 11] = [
Self::Network,
Self::Models,
Self::Accounts,
@@ -168,6 +194,7 @@ impl Permission {
Self::Keys,
Self::KeyBudgets,
Self::QuotaObservations,
Self::UpstreamConnections,
];
#[must_use]
@@ -183,6 +210,7 @@ impl Permission {
Self::Keys => "keys",
Self::KeyBudgets => "key_budgets",
Self::QuotaObservations => "quota_observations",
Self::UpstreamConnections => "upstream_connections",
}
}
@@ -200,6 +228,7 @@ impl Permission {
Self::Keys => "专用 API Key",
Self::KeyBudgets => "API Key 预算",
Self::QuotaObservations => "账号额度观测",
Self::UpstreamConnections => "账号上游连接",
}
}
@@ -208,6 +237,9 @@ impl Permission {
pub const fn description(self) -> &'static str {
match self {
Self::Network => "访问网络",
Self::UpstreamConnections => {
"使用当前请求所选账号的身份与代理连接插件声明的上游,不提供原始凭据读取或账号修改"
}
Self::Models => "查询模型与 API Key 信息并调用模型,可能产生消耗",
Self::Accounts => "读取和修改账号,包括访问原始凭据",
Self::Data => {
@@ -0,0 +1,204 @@
//! HTTP 正文句柄封装;普通出站和已选账号出站共用惰性读取与提前关闭语义。
use serde::{Serialize, de::DeserializeOwned};
use super::{HostClient, SessionError};
use crate::{
ErrorCode, PluginFault,
call::{host as wire, upstream_adapter::UpstreamHttpRequest},
};
/// 响应头与惰性正文;不包含宿主的内部流句柄。
pub struct HostHttpResponse {
pub status: u16,
pub headers: Vec<(String, String)>,
pub body: HostHttpBody,
}
/// 父调用内的单消费者正文;不后台预读,父调用结束或取消时宿主兜底回收。
pub struct HostHttpBody {
host: HostClient,
stream: Option<String>,
upstream: bool,
}
impl HostClient {
/// 使用普通网络授权发送请求;账号上游使用 [`Self::upstream_http`]。
///
/// # Errors
/// 权限、期限、请求或网络无效时失败。
pub async fn http(
&self,
request: wire::HttpRequest,
body: Vec<u8>,
) -> Result<HostHttpResponse, PluginFault> {
self.open_http("host.http.do_stream", request, body, false)
.await
}
/// 由宿主附加已选账号的认证与代理,返回统一的惰性正文。
///
/// # Errors
/// 路径、权限、期限、网络或宿主响应无效时失败。
pub async fn upstream_http(
&self,
request: UpstreamHttpRequest,
body: Vec<u8>,
) -> Result<HostHttpResponse, PluginFault> {
self.open_http("host.upstream.http.do_stream", request, body, true)
.await
}
async fn open_http<T: Serialize>(
&self,
method: &str,
request: T,
body: Vec<u8>,
upstream: bool,
) -> Result<HostHttpResponse, PluginFault> {
let (response, payload): (wire::HttpResponse, _) =
callback(self, method, request, body).await?;
if !payload.is_empty() || response.stream.as_ref().is_none_or(String::is_empty) {
return Err(invalid());
}
Ok(HostHttpResponse {
status: response.status,
headers: response.headers,
body: HostHttpBody {
host: self.clone(),
stream: response.stream,
upstream,
},
})
}
}
impl HostHttpBody {
/// 按需读取至多 64 KiB,EOF 后再次读取不再调用宿主。
///
/// # Errors
/// 父调用结束、流无效或网络失败时返回错误。
pub async fn read(&mut self) -> Result<Option<Vec<u8>>, PluginFault> {
let Some(stream) = &self.stream else {
return Ok(None);
};
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct Read {
eof: bool,
}
let maximum_bytes = u32::try_from(self.host.maximum_stream_chunk_bytes().min(64 * 1024))
.map_err(|_| invalid())?;
let (result, payload): (Read, _) = callback(
&self.host,
if self.upstream {
"host.upstream.http.stream_read"
} else {
"host.http.stream_read"
},
wire::StreamRead {
stream: stream.clone(),
maximum_bytes,
},
vec![],
)
.await?;
if payload.len() > maximum_bytes as usize || (result.eof && !payload.is_empty()) {
let _ = self.close().await;
return Err(invalid());
}
if result.eof {
self.stream = None;
Ok(None)
} else {
Ok(Some(payload))
}
}
/// 收集有限正文;超出调用方预算或读取失败时提前关闭,不把流无限缓冲到内存。
///
/// # Errors
/// 读取失败或超出 `maximum_bytes` 时返回错误。
pub async fn collect(mut self, maximum_bytes: usize) -> Result<Vec<u8>, PluginFault> {
let mut body = Vec::new();
loop {
match self.read().await {
Ok(Some(chunk)) if chunk.len() <= maximum_bytes.saturating_sub(body.len()) => {
body.extend(chunk)
}
Ok(None) => return Ok(body),
result => {
let _ = self.close().await;
return Err(result.err().unwrap_or_else(|| {
PluginFault::new(
ErrorCode::Capacity,
"HTTP body exceeds the collection limit",
)
}));
}
}
}
}
/// 提前释放正文,重复关闭或 EOF 后关闭不再调用宿主。
///
/// # Errors
/// 父调用已结束或宿主拒绝释放时失败。
pub async fn close(&mut self) -> Result<(), PluginFault> {
let Some(stream) = &self.stream else {
return Ok(());
};
empty_callback(
&self.host,
if self.upstream {
"host.upstream.http.stream_close"
} else {
"host.http.stream_close"
},
wire::StreamClose {
stream: stream.clone(),
},
vec![],
)
.await?;
self.stream = None;
Ok(())
}
}
pub(super) async fn callback<T: Serialize, R: DeserializeOwned>(
host: &HostClient,
method: &str,
request: T,
body: Vec<u8>,
) -> Result<(R, Vec<u8>), PluginFault> {
let reply = host
.call(
method,
serde_json::to_value(request).map_err(|_| invalid())?,
body,
)
.await
.map_err(SessionError::into_plugin_fault)?;
Ok((
serde_json::from_value(reply.result).map_err(|_| invalid())?,
reply.payload,
))
}
pub(super) async fn empty_callback<T: Serialize>(
host: &HostClient,
method: &str,
request: T,
body: Vec<u8>,
) -> Result<(), PluginFault> {
let (metadata, payload): (serde_json::Value, _) = callback(host, method, request, body).await?;
if metadata != serde_json::json!({}) || !payload.is_empty() {
return Err(invalid());
}
Ok(())
}
pub(super) fn invalid() -> PluginFault {
PluginFault::new(ErrorCode::InvalidInput, "invalid network callback payload")
}
@@ -2,16 +2,19 @@
mod data;
mod frame;
mod http;
mod keys;
mod middleware;
mod plugin;
mod resources;
mod session;
mod upstream_adapter;
use crate::{ErrorCode, PluginFault};
use serde::{Serialize, de::DeserializeOwned};
pub use frame::{read_frame, validate_frame, write_frame};
pub use http::{HostHttpBody, HostHttpResponse};
pub use middleware::{
MiddlewareBody, MiddlewareBodySender, MiddlewareCall, MiddlewareNext, MiddlewarePlugin,
MiddlewareRequest, MiddlewareResponse,
@@ -23,6 +26,7 @@ pub use session::{
CallCancellation, CallFuture, CallReply, HostClient, HostReply, PluginCall, PluginHandler,
PluginSession, ResponseStream, SessionConfig, SessionError, StreamSender,
};
pub use upstream_adapter::{UpstreamWebSocket, UpstreamWebSocketUpgrade};
async fn payload_call<T: Serialize, R: DeserializeOwned>(
host: &HostClient,
@@ -2,9 +2,30 @@
use crate::{
Capability as C, Stage as S,
call::{catalog, frontend_authentication as frontend, host, management, observation, policy},
call::{
catalog, frontend_authentication as frontend, host, management, observation, policy,
upstream_adapter,
},
};
pub const UPSTREAM_ADAPTER_REGISTER: Method<Empty, upstream_adapter::UpstreamAdapterRegistration> =
Method::new(
upstream_adapter::REGISTER_METHOD,
&[C::UpstreamAdapter],
&[S::Registration],
decode_metadata_without_payload,
encode_payload,
);
pub const UPSTREAM_ADAPTER_EXECUTE: Method<upstream_adapter::UpstreamAdapterRequest, Empty> =
Method::new(
upstream_adapter::EXECUTE_METHOD,
&[C::UpstreamAdapter],
&[S::Upstream],
super::typed::decode_upstream,
super::typed::encode_stream,
);
use super::{
Empty, Method,
typed::{
@@ -126,6 +126,21 @@ pub(super) fn decode_payload<P: DeserializeOwned>(
Ok(typed_call(call, request, false))
}
pub(super) fn decode_upstream(
mut call: PluginCall,
stages: &'static [Stage],
) -> Result<TypedCall<crate::call::upstream_adapter::UpstreamAdapterRequest>, PluginFault> {
validate_stage(&call, stages)?;
if !empty_object(&call.params) {
return Err(invalid_input());
}
let (request, body) =
crate::call::upstream_adapter::UpstreamAdapterRequest::decode(&call.payload)
.map_err(|_| invalid_input())?;
call.payload = body;
Ok(typed_call(call, request, true))
}
pub(super) fn encode_metadata<R: Serialize>(
reply: TypedReply<R>,
) -> Result<CallReply, PluginFault> {
@@ -141,6 +156,22 @@ pub(super) fn encode_metadata<R: Serialize>(
Ok(CallReply::unary(result, Vec::new()))
}
pub(super) fn encode_stream<R: Serialize>(reply: TypedReply<R>) -> Result<CallReply, PluginFault> {
let TypedReply {
result,
payload,
stream,
} = reply;
if !payload.is_empty() {
return Err(invalid_input());
}
Ok(CallReply::stream(
serde_json::to_value(result).map_err(|_| invalid_input())?,
payload,
stream.ok_or_else(invalid_input)?,
))
}
pub(super) fn encode_metadata_with_payload<R: Serialize>(
reply: TypedReply<R>,
) -> Result<CallReply, PluginFault> {
@@ -18,6 +18,10 @@ pub(super) fn validate(
match capability {
C::Maintenance => required(methods::RECONCILE.name)?,
C::Middleware => required(crate::call::middleware::HANDLE_METHOD)?,
C::UpstreamAdapter => {
required(methods::UPSTREAM_ADAPTER_REGISTER.name)?;
required(methods::UPSTREAM_ADAPTER_EXECUTE.name)?;
}
C::ModelRouter => required(methods::ROUTE_MODEL.name)?,
C::ModelCatalog => required(methods::MODEL_CATALOG_REGISTER.name)?,
C::RetryPolicy => required(methods::RETRY_DECISION.name)?,
@@ -0,0 +1,140 @@
//! 已选账号的受管 WebSocket;连接归属与跨轮续接由宿主管理。
use super::{
HostClient,
http::{callback, empty_callback, invalid},
};
use crate::{
ErrorCode, PluginFault,
call::{
host::HttpResponse,
upstream_adapter::{
UpstreamWebSocketMessage, UpstreamWebSocketRead, UpstreamWebSocketRequest,
WebSocketMessageKind,
},
},
};
/// 握手拒绝仍保留 HTTP 状态及正文,供适配器解析上游错误。
pub enum UpstreamWebSocketUpgrade {
Connected {
headers: Vec<(String, String)>,
connection: UpstreamWebSocket,
},
Rejected {
status: u16,
headers: Vec<(String, String)>,
body: Vec<u8>,
},
}
/// 本次调用的单消费者连接入口;不暴露凭据或可跨账号复用的连接句柄。
///
/// 结束调用时宿主关闭连接;只有成功终态携带连接内续接状态时才保留连接。
/// 因此丢弃此对象不会提前关闭宿主尚待确认的续接连接。
pub struct UpstreamWebSocket {
host: HostClient,
closed: bool,
}
impl HostClient {
/// 建立连接,续接时复用宿主恢复的同一连接。
///
/// # Errors
/// 目标、权限、握手或续接归属无效时失败。
pub async fn upstream_websocket(
&self,
request: UpstreamWebSocketRequest,
) -> Result<UpstreamWebSocketUpgrade, PluginFault> {
let (response, body): (HttpResponse, _) =
callback(self, "host.upstream.websocket.open", request, vec![]).await?;
if response.stream.is_some() || (response.status == 101 && !body.is_empty()) {
return Err(invalid());
}
Ok(if response.status == 101 {
UpstreamWebSocketUpgrade::Connected {
headers: response.headers,
connection: UpstreamWebSocket {
host: self.clone(),
closed: false,
},
}
} else {
UpstreamWebSocketUpgrade::Rejected {
status: response.status,
headers: response.headers,
body,
}
})
}
}
impl UpstreamWebSocket {
/// 发送完整文本或二进制消息。
///
/// # Errors
/// 连接关闭、消息无效、网络失败或期限到达时失败。
pub async fn send(
&mut self,
kind: WebSocketMessageKind,
body: Vec<u8>,
) -> Result<(), PluginFault> {
if self.closed {
return Err(PluginFault::new(
ErrorCode::InvalidInput,
"WebSocket is closed",
));
}
empty_callback(
&self.host,
"host.upstream.websocket.send",
UpstreamWebSocketMessage { kind },
body,
)
.await
}
/// 按需读取完整消息,EOF 后再次读取返回 None;Ping/Pong 由宿主处理。
///
/// # Errors
/// 父调用结束、网络失败或期限到达时失败。
pub async fn read(&mut self) -> Result<Option<(WebSocketMessageKind, Vec<u8>)>, PluginFault> {
if self.closed {
return Ok(None);
}
let (read, body): (UpstreamWebSocketRead, _) = callback(
&self.host,
"host.upstream.websocket.read",
serde_json::json!({}),
vec![],
)
.await?;
match (read.eof, read.kind) {
(true, None) if body.is_empty() => {
self.closed = true;
Ok(None)
}
(false, Some(kind)) => Ok(Some((kind, body))),
_ => Err(invalid()),
}
}
/// 提前关闭并放弃连接内续接,重复关闭不调用宿主。
///
/// # Errors
/// 父调用已结束或宿主拒绝释放时失败。
pub async fn close(&mut self) -> Result<(), PluginFault> {
if self.closed {
return Ok(());
}
empty_callback(
&self.host,
"host.upstream.websocket.close",
serde_json::json!({}),
vec![],
)
.await?;
self.closed = true;
Ok(())
}
}
@@ -261,8 +261,9 @@ impl Manifest {
let stages = declaration.stages.iter().copied().collect::<BTreeSet<_>>();
let input_formats = declaration.input_formats.iter().collect::<BTreeSet<_>>();
let output_formats = declaration.output_formats.iter().collect::<BTreeSet<_>>();
if !(declaration.version == 1
|| (*capability == Capability::Middleware && declaration.version == 2))
if !capability
.contract_versions()
.contains(&declaration.version)
|| declaration.id.len() > 128
|| local_id.is_empty()
|| !local_id.is_ascii()
@@ -0,0 +1,223 @@
use super::session::{HostPeer, receive, send_call, shutdown, start_session};
use gateway_plugin_sdk::client::{CallFuture, CallReply, PluginCall, PluginHandler, write_frame};
use gateway_plugin_sdk::{ErrorCode, Frame, Message};
use gateway_plugin_sdk::{
call::{
host::HttpRequest,
upstream_adapter::{UpstreamHttpRequest, UpstreamWebSocketRequest, WebSocketMessageKind},
},
client::UpstreamWebSocketUpgrade,
};
use serde_json::{Value, json};
struct NetworkHandler;
impl PluginHandler for NetworkHandler {
fn call(&self, call: PluginCall) -> CallFuture<'_> {
Box::pin(async move {
if call.method == "websocket" {
let upgrade = call
.host
.upstream_websocket(UpstreamWebSocketRequest {
path: "responses".into(),
query: vec![],
headers: vec![],
})
.await?;
let UpstreamWebSocketUpgrade::Connected { mut connection, .. } = upgrade else {
let UpstreamWebSocketUpgrade::Rejected { status, body, .. } = upgrade else {
unreachable!()
};
return Ok(CallReply::unary(json!({"status":status}), body));
};
connection
.send(WebSocketMessageKind::Text, b"request".to_vec())
.await?;
let (kind, body) = connection.read().await?.unwrap();
assert_eq!(kind, WebSocketMessageKind::Text);
assert!(connection.read().await?.is_none());
assert!(connection.read().await?.is_none());
connection.close().await?;
assert!(connection.send(kind, vec![]).await.is_err());
return Ok(CallReply::unary(json!({}), body));
}
let response = if call.method == "http" {
call.host
.http(
HttpRequest {
method: "POST".into(),
url: "https://example.test/responses".into(),
headers: vec![],
},
call.payload,
)
.await?
} else {
call.host
.upstream_http(
UpstreamHttpRequest {
method: "POST".into(),
path: "responses".into(),
query: vec![],
headers: vec![],
},
call.payload,
)
.await?
};
assert_eq!(response.status, 200);
if call.params["collect"] == true {
return Ok(CallReply::unary(json!({}), response.body.collect(3).await?));
}
let mut body = response.body;
let bytes = body.read().await?.unwrap();
assert!(body.read().await?.is_none());
assert!(body.read().await?.is_none());
body.close().await?;
body.close().await?;
Ok(CallReply::unary(json!({}), bytes))
})
}
}
async fn reply_callback(
host: &mut HostPeer,
expected: &str,
result: Value,
payload: Vec<u8>,
) -> Frame {
let frame = receive(host).await;
let Message::Callback { id, method, .. } = &frame.message else {
panic!("expected callback")
};
assert_eq!(method, expected);
write_frame(
&mut host.writer,
&Frame {
message: Message::Result { id: *id, result },
payload,
},
)
.await
.unwrap();
frame
}
#[tokio::test]
async fn http_response_objects_share_pull_semantics_and_retire_eof_handles() {
for method in ["http", "upstream"] {
let (mut host, task) = start_session(NetworkHandler).await;
send_call(&mut host, 1, method, json!({}), b"request".to_vec()).await;
let prefix = if method == "http" {
"host.http"
} else {
"host.upstream.http"
};
let frame = reply_callback(
&mut host,
&format!("{prefix}.do_stream"),
json!({"status":200,"headers":[],"stream":"private-handle"}),
vec![],
)
.await;
assert_eq!(frame.payload, b"request");
let frame = reply_callback(
&mut host,
&format!("{prefix}.stream_read"),
json!({"eof":false}),
b"abc".to_vec(),
)
.await;
assert!(
matches!(frame.message, Message::Callback { params, .. } if params == json!({"stream":"private-handle","maximum_bytes":65536}))
);
reply_callback(
&mut host,
&format!("{prefix}.stream_read"),
json!({"eof":true}),
vec![],
)
.await;
let result = receive(&mut host).await;
assert!(matches!(result.message, Message::Result { id: 1, .. }));
assert_eq!(result.payload, b"abc");
shutdown(&mut host, task).await;
}
}
#[tokio::test]
async fn bounded_collection_closes_the_host_stream_before_reporting_overflow() {
let (mut host, task) = start_session(NetworkHandler).await;
send_call(&mut host, 1, "upstream", json!({"collect":true}), vec![]).await;
reply_callback(
&mut host,
"host.upstream.http.do_stream",
json!({"status":200,"headers":[],"stream":"private-handle"}),
vec![],
)
.await;
reply_callback(
&mut host,
"host.upstream.http.stream_read",
json!({"eof":false}),
b"abcd".to_vec(),
)
.await;
reply_callback(
&mut host,
"host.upstream.http.stream_close",
json!({}),
vec![],
)
.await;
assert!(
matches!(receive(&mut host).await.message, Message::Error { id: 1, error } if error.code == ErrorCode::Capacity)
);
shutdown(&mut host, task).await;
}
#[tokio::test]
async fn websocket_session_handles_messages_eof_and_http_rejection() {
let (mut host, task) = start_session(NetworkHandler).await;
send_call(&mut host, 1, "websocket", json!({}), vec![]).await;
reply_callback(
&mut host,
"host.upstream.websocket.open",
json!({"status":101,"headers":[],"stream":null}),
vec![],
)
.await;
let sent = reply_callback(&mut host, "host.upstream.websocket.send", json!({}), vec![]).await;
assert_eq!(sent.payload, b"request");
reply_callback(
&mut host,
"host.upstream.websocket.read",
json!({"eof":false,"kind":"text"}),
b"response".to_vec(),
)
.await;
reply_callback(
&mut host,
"host.upstream.websocket.read",
json!({"eof":true,"kind":null}),
vec![],
)
.await;
let response = receive(&mut host).await;
assert!(matches!(response.message, Message::Result { id: 1, .. }));
assert_eq!(response.payload, b"response");
send_call(&mut host, 3, "websocket", json!({}), vec![]).await;
reply_callback(
&mut host,
"host.upstream.websocket.open",
json!({"status":401,"headers":[],"stream":null}),
b"unauthorized".to_vec(),
)
.await;
let response = receive(&mut host).await;
assert!(
matches!(response.message, Message::Result { id: 3, result } if result == json!({"status":401}))
);
assert_eq!(response.payload, b"unauthorized");
shutdown(&mut host, task).await;
}
@@ -1,2 +1,3 @@
mod http;
mod plugin;
mod session;
@@ -301,12 +301,12 @@ impl PluginHandler for TestHandler {
}
}
struct HostPeer {
reader: ReadHalf<DuplexStream>,
writer: WriteHalf<DuplexStream>,
pub(super) struct HostPeer {
pub(super) reader: ReadHalf<DuplexStream>,
pub(super) writer: WriteHalf<DuplexStream>,
}
async fn start_session<H: PluginHandler>(
pub(super) async fn start_session<H: PluginHandler>(
handler: H,
) -> (HostPeer, JoinHandle<Result<(), SessionError>>) {
start_session_with_capacity(handler, MAXIMUM_STREAM_CHUNK_BYTES * 2).await
@@ -603,7 +603,13 @@ fn middleware_request() -> MiddlewareRequestHead {
}
}
async fn send_call(host: &mut HostPeer, id: u64, method: &str, params: Value, payload: Vec<u8>) {
pub(super) async fn send_call(
host: &mut HostPeer,
id: u64,
method: &str,
params: Value,
payload: Vec<u8>,
) {
send_call_with_timeout(host, id, method, params, payload, Duration::from_secs(1)).await;
}
@@ -637,14 +643,14 @@ async fn send_control(host: &mut HostPeer, message: Message) {
.unwrap();
}
async fn receive(host: &mut HostPeer) -> Frame {
pub(super) async fn receive(host: &mut HostPeer) -> Frame {
tokio::time::timeout(Duration::from_secs(1), read_frame(&mut host.reader))
.await
.expect("plugin response timed out")
.expect("plugin response frame must be valid")
}
async fn shutdown(host: &mut HostPeer, task: JoinHandle<Result<(), SessionError>>) {
pub(super) async fn shutdown(host: &mut HostPeer, task: JoinHandle<Result<(), SessionError>>) {
send_control(host, Message::Shutdown).await;
tokio::time::timeout(Duration::from_secs(1), task)
.await
@@ -126,6 +126,40 @@ fn installation_checks_protocol_engine_and_target() {
assert_eq!(manifest.validate(), Err(ManifestError::Invalid));
}
#[test]
fn capability_contract_versions_are_checked_independently_of_rpc_version() {
let mut manifest = packaged_manifest();
for version in [1, 2] {
manifest
.contributes
.get_mut(&Capability::Middleware)
.unwrap()
.version = version;
assert!(manifest.validate().is_ok());
}
for version in [0, 3, u32::MAX] {
manifest
.contributes
.get_mut(&Capability::Middleware)
.unwrap()
.version = version;
assert_eq!(manifest.validate(), Err(ManifestError::Invalid));
}
manifest.contributes.clear();
manifest.contributes.insert(
Capability::Management,
ContributionDeclaration {
id: "9acme.request-tags.management".into(),
version: 2,
stages: vec![Stage::Management],
input_formats: vec![],
output_formats: vec![],
},
);
assert_eq!(manifest.validate(), Err(ManifestError::Invalid));
}
#[test]
fn package_rejects_file_directory_conflicts_and_missing_declared_files() {
let mut manifest = packaged_manifest();
@@ -103,6 +103,7 @@ use crate::transport::{
mod execution;
mod failure;
mod observation;
mod upstream_adapter;
mod workers;
pub(crate) use workers::ClientReleaseServices;
@@ -472,9 +473,18 @@ impl Provider for CodexProvider {
UpstreamSendState::NotSent,
));
};
let adapter = context.upstream_adapter(candidate.provider(), upstream_model)?;
if adapter.is_none()
&& generate
.provider_session_state(PROVIDER_NAME)
.is_some_and(|state| state.extension_owner().is_some())
{
return Err(continuation_replay_required_error("scope_unavailable"));
}
// 其他协议必须先取得真实账号,再按固定 attempt 阶段调用转换器;选号前不能
// 把未知正文当成 OpenAI wire 解释会话、亲和或传输字段。
let preselection = (generate.protocol_payload().protocol() == PROVIDER_NAME)
let preselection = (adapter.is_none()
&& generate.protocol_payload().protocol() == PROVIDER_NAME)
.then(|| self.prepare_generate_request(generate, upstream_model, &context))
.transpose()?;
let (selection_session_affinity, selection_cyber_policy_key, requires_websocket) =
@@ -539,6 +549,18 @@ impl Provider for CodexProvider {
account_id,
Box::new(move |operation, middleware_headers| {
Box::pin(async move {
if let Some(adapter) = adapter {
return provider.execute_upstream_adapter(
operation,
middleware_headers,
terminal_model,
terminal_context,
lease,
account_selection_wait_ms,
frozen_requirements,
adapter,
);
}
provider
.execute_selected_generate(
operation,
@@ -0,0 +1,206 @@
//! 插件只替换业务 wire;账号租约、凭据解释与费用来源仍属于 OpenAI。
use super::*;
use crate::transport::usage::{OpenAiBillingUsage, openai_billing_breakdown_with_override};
use gateway_core::{
account::{CredentialRevision, OutboundProxy, ProviderAccountId},
engine::upstream_adapter::{
UpstreamAccountConnection, UpstreamAdapter, UpstreamAdapterInvocation,
},
metering::{CalculatedCost, Usage},
};
struct SelectedConnection {
provider: Arc<CodexProvider>,
lease: Arc<CodexCredentialLease>,
context: AttemptContext,
model: UpstreamModelId,
requested_service_tier: Option<String>,
}
impl CodexProvider {
// 与原生 terminal 共用一次选号结果,不能为适配器重新获取账号或租约。
#[expect(clippy::too_many_arguments)]
pub(super) fn execute_upstream_adapter(
self: Arc<Self>,
mut operation: Operation,
headers: Vec<MiddlewareHeader>,
model: UpstreamModelId,
context: AttemptContext,
lease: CodexCredentialLease,
account_selection_wait_ms: u64,
requirements: CapabilityRequirements,
adapter: Arc<dyn UpstreamAdapter>,
) -> Result<ProviderStream, ProviderError> {
if operation.capability_requirements() != requirements {
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
));
}
if context.disable_fast()
&& let Operation::Generate(generate) = &operation
&& generate.protocol_payload().protocol() == PROVIDER_NAME
{
// 适配器终端同样在 attempt 加工后应用原生强制策略,不另写一套档位规则。
let mut request =
CodexResponsesRequest::from_body(generate.protocol_payload().body().clone());
request.apply_fast_policy(true);
let body = serde_json::to_vec(request.body()).map_err(|_| {
provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
)
})?;
operation = operation
.replace_middleware_wire(PROVIDER_NAME, body.into())
.map_err(|_| {
provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
)
})?;
}
let requested_service_tier = match &operation {
Operation::Generate(generate)
if generate.protocol_payload().protocol() == PROVIDER_NAME =>
{
let request =
CodexResponsesRequest::from_body(generate.protocol_payload().body().clone());
normalize_service_tier(request.service_tier())
}
_ => None,
};
let lease = Arc::new(lease);
let metadata = ProviderCallMetadata::new(
ProviderKind::new(PROVIDER_NAME).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
})?,
model.clone(),
lease.account_id().clone(),
UpstreamTransport::new(adapter.transport()).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
})?,
)
.with_selection_observation(ProviderSelectionObservation::new(
account_selection_wait_ms,
lease.capacity_snapshot(),
));
let account = Arc::new(SelectedConnection {
provider: Arc::clone(&self),
lease: Arc::clone(&lease),
context: context.clone(),
model,
requested_service_tier,
});
let events = adapter.execute(UpstreamAdapterInvocation {
operation,
headers,
context,
metadata: metadata.clone(),
account,
});
let mutation = lease.allows_account_state_mutation();
let stream = ProviderStream::new(metadata, events, lease);
Ok(if mutation {
stream.with_filtered_account_feedback(
Arc::clone(&self.account_feedback),
openai_failure_affects_account_score,
)
} else {
stream
})
}
}
impl UpstreamAccountConnection for SelectedConnection {
fn account_id(&self) -> &ProviderAccountId {
self.lease.account_id()
}
fn credential_revision(&self) -> CredentialRevision {
self.lease.account().revision()
}
fn authentication_kind(&self) -> &str {
self.lease.account().authentication_kind()
}
fn outbound_proxy(&self) -> Option<&OutboundProxy> {
self.lease.account().outbound_proxy()
}
fn authorization(&self) -> Result<Vec<MiddlewareHeader>, ProviderError> {
let authorization = self
.lease
.authentication()
.authorization_header()
.map_err(|_| {
provider_error(ProviderErrorKind::Unauthorized, UpstreamSendState::NotSent)
})?;
let mut headers = vec![MiddlewareHeader::new(
"authorization",
bytes::Bytes::copy_from_slice(authorization.expose_secret().as_bytes()),
)];
if let Some(account_id) = self.lease.account().upstream_account_id() {
headers.push(MiddlewareHeader::new(
"chatgpt-account-id",
bytes::Bytes::copy_from_slice(account_id.as_bytes()),
));
}
// 业务目标可能与原生 Responses 不同,不能把原域名的 Cookie 带到新目标。
Ok(headers)
}
fn calculate_cost(&self, _: Option<&str>, usage: &Usage) -> Option<CalculatedCost> {
// 与原生路径一样按请求策略计价;响应回显只作观测,不能切换模型或档位。
let usage = gateway_protocol::openai::events::TokenUsage {
input_tokens: usage.input_tokens?,
output_tokens: usage.output_tokens?,
cached_tokens: usage.cached_tokens.unwrap_or(0),
cache_write_tokens: usage.cache_write_tokens.unwrap_or(0),
image_input_tokens: usage.image_input_tokens.unwrap_or(0),
image_output_tokens: usage.image_output_tokens.unwrap_or(0),
..Default::default()
};
openai_billing_breakdown_with_override(
self.model.as_str(),
OpenAiBillingUsage::from(usage),
self.requested_service_tier.as_deref(),
self.context
.pricing()
.get(PROVIDER_NAME)
.and_then(|models| models.get(self.model.as_str())),
)
.map(|cost| cost.calculated_cost())
}
fn record_failure(&self, error: ProviderError) -> BoxFuture<'_, ProviderError> {
Box::pin(async move {
if !self.lease.allows_account_state_mutation()
|| error.send_state() == UpstreamSendState::NotSent
{
return error;
}
let failure = match error.kind() {
ProviderErrorKind::Unauthorized => Some(CodexAccountFailure::CredentialExpired),
ProviderErrorKind::RateLimited => Some(CodexAccountFailure::RateLimited {
retry_after: error.retry_after(),
}),
ProviderErrorKind::QuotaExhausted => Some(CodexAccountFailure::QuotaExhausted),
_ => None,
};
if let Some(failure) = failure {
let _ = self
.provider
.selector
.record_failure(self.lease.account(), failure, None)
.await;
if error.kind() == ProviderErrorKind::QuotaExhausted {
schedule_authoritative_quota_refresh_after_failure(
&self.provider.quota,
self.lease.account(),
);
}
}
error
})
}
}
@@ -2,6 +2,7 @@ mod account_isolation;
mod capacity;
mod precommit;
mod response_interrupt;
mod upstream_adapter;
use std::collections::{BTreeMap, BTreeSet};
use std::io::{self, Write};
@@ -0,0 +1,346 @@
use super::*;
use gateway_core::engine::upstream_adapter::{
UpstreamAccountConnection, UpstreamAdapter, UpstreamAdapterInvocation, UpstreamAdapterPlan,
};
#[derive(Debug)]
struct AdapterProbe {
polls: Arc<std::sync::atomic::AtomicUsize>,
}
impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AdapterProbe {
fn select(
&self,
_: &AttemptContext,
provider: &ProviderKind,
_: &UpstreamModelId,
) -> Result<
Option<Arc<dyn gateway_core::engine::upstream_adapter::UpstreamAdapter>>,
gateway_core::error::ProviderError,
> {
assert_eq!(provider.as_str(), "openai");
Ok(Some(Arc::new(Self {
polls: self.polls.clone(),
})))
}
}
impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe {
fn transport(&self) -> &str {
"http_sse"
}
fn execute(
self: Arc<Self>,
invocation: gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation,
) -> gateway_core::engine::provider::EventStream {
Box::pin(futures::stream::once(async move {
self.polls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
assert_eq!(
invocation.account.account_id().as_str(),
"acct_provider_contract"
);
assert_eq!(
invocation.metadata.provider_account_id(),
invocation.account.account_id()
);
assert_eq!(invocation.account.authentication_kind(), "oauth");
let headers = invocation.account.authorization().unwrap();
let header = |name: &str| {
headers
.iter()
.find(|h| h.name() == name)
.map(|h| h.value().to_vec())
};
assert_eq!(
header("authorization"),
Some("Bearer at-acct_provider_contract".as_bytes().to_vec())
);
assert_eq!(
header("chatgpt-account-id"),
Some("chatgpt-acct_provider_contract".as_bytes().to_vec())
);
assert!(header("cookie").is_none());
assert!(invocation.context.disable_fast());
let Operation::Generate(generate) = &invocation.operation else {
panic!("adapter must receive the original operation kind");
};
assert_eq!(
generate.protocol_payload().body()["service_tier"],
"default"
);
assert!(
invocation
.headers
.iter()
.any(|h| h.name() == "x-adapter-onion")
);
assert!(
invocation
.account
.calculate_cost(
None,
&gateway_core::metering::Usage {
input_tokens: Some(7),
output_tokens: Some(2),
..Default::default()
}
)
.is_some()
);
let usage = gateway_core::metering::Usage {
input_tokens: Some(7),
output_tokens: Some(2),
..Default::default()
};
assert_eq!(
invocation.account.calculate_cost(Some("priority"), &usage),
invocation.account.calculate_cost(Some("default"), &usage),
"响应回显不得覆盖已冻结的禁用 Fast 策略"
);
Err(gateway_core::error::ProviderError::new(
ProviderErrorKind::Cancelled,
UpstreamSendState::NotSent,
))
}))
}
}
#[tokio::test]
async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_and_stays_cold() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let provider = provider_with_base_url(&store, server.uri());
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let lease = ExtensionSetReference::new(
ExtensionSetId::new("adapter-native-test".to_owned()).unwrap(),
Arc::new(TestExtensionLease),
);
let middleware = Arc::new(RecordingMiddleware {
observed: Default::default(),
replacement: ("service_tier".into(), json!("priority")),
request_headers: vec![MiddlewareHeader::new(
"x-adapter-onion",
Bytes::from_static(b"present"),
)],
});
let context = AttemptContext::new(
gateway_core::engine::RequestAttemptContext::new(
ModelRequestId::new("req_adapter_native").unwrap(),
ClientApiKeyId::new("key_openai_contract").unwrap(),
)
.with_disable_fast(true)
.with_upstream_adapters(Some(
gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(
Arc::new(AdapterProbe {
polls: polls.clone(),
}),
lease.clone(),
),
))
.with_middleware(
Some(FrozenMiddlewarePlan::new(middleware, lease)),
Arc::from([]),
"/v1/responses".to_owned(),
ClientTransport::HttpSse,
),
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
);
let mut stream = provider
.execute(planned_request("openai", generate_operation()), context)
.await
.unwrap();
assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0);
let error = stream.next().await.unwrap().unwrap_err();
assert_eq!(error.kind(), ProviderErrorKind::Cancelled);
assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 1);
assert!(server.received_requests().await.unwrap().is_empty());
}
type SelectedConnection = Arc<dyn UpstreamAccountConnection>;
struct ConnectionProbe {
selected: Arc<Mutex<Option<oneshot::Sender<SelectedConnection>>>>,
}
impl std::fmt::Debug for ConnectionProbe {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ConnectionProbe")
.finish_non_exhaustive()
}
}
impl UpstreamAdapterPlan for ConnectionProbe {
fn select(
&self,
_: &AttemptContext,
_: &ProviderKind,
_: &UpstreamModelId,
) -> Result<Option<Arc<dyn UpstreamAdapter>>, gateway_core::error::ProviderError> {
Ok(Some(Arc::new(Self {
selected: Arc::clone(&self.selected),
})))
}
}
impl UpstreamAdapter for ConnectionProbe {
fn transport(&self) -> &str {
"http_sse"
}
fn execute(
self: Arc<Self>,
invocation: UpstreamAdapterInvocation,
) -> gateway_core::engine::provider::EventStream {
Box::pin(futures::stream::once(async move {
assert!(
self.selected
.lock()
.unwrap()
.take()
.unwrap()
.send(invocation.account)
.is_ok()
);
Err(gateway_core::error::ProviderError::new(
ProviderErrorKind::Cancelled,
UpstreamSendState::NotSent,
))
}))
}
}
async fn selected_adapter_connection(
provider: Arc<CodexProvider>,
request_id: &str,
) -> SelectedConnection {
let (selected, receiver) = oneshot::channel();
let plan = gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(
Arc::new(ConnectionProbe {
selected: Arc::new(Mutex::new(Some(selected))),
}),
ExtensionSetReference::new(
ExtensionSetId::new("adapter-credential-test".to_owned()).unwrap(),
Arc::new(TestExtensionLease),
),
);
let context = AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new(request_id).unwrap(),
ClientApiKeyId::new("key_openai_contract").unwrap(),
)
.with_upstream_adapters(Some(plan)),
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
);
let mut stream = provider
.execute(planned_request("openai", generate_operation()), context)
.await
.unwrap();
assert_eq!(
stream.next().await.unwrap().unwrap_err().kind(),
ProviderErrorKind::Cancelled
);
receiver.await.unwrap()
}
#[tokio::test]
async fn upstream_adapter_reloads_rotated_credentials_and_proxy_without_accepting_stale_failure() {
use gateway_core::account::OutboundProxy;
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_provider_contract";
create_account(&store, account_id).await;
let old_proxy = OutboundProxy::parse("http://127.0.0.1:18001").unwrap();
let new_proxy = OutboundProxy::parse("http://127.0.0.1:18002").unwrap();
store.set_egress(account_id, Some(old_proxy.clone()), None);
let upstream = MockServer::start().await;
let provider = provider_with_base_url(&store, upstream.uri());
let old =
selected_adapter_connection(Arc::clone(&provider), "req_adapter_before_refresh").await;
let before = store.account(account_id).unwrap();
assert_eq!(old.outbound_proxy(), Some(&old_proxy));
// 使用原生刷新成功后的 CAS 入口;适配器不能复制令牌或缓存第二份账号代理。
let revision = store
.repository()
.rotate_refreshed_oauth_secret(
&before,
secret("rotated-adapter-access-token"),
Some(SystemTime::now() + Duration::from_secs(3600)),
None,
)
.await
.unwrap();
store.set_egress(account_id, Some(new_proxy.clone()), None);
let current = selected_adapter_connection(provider, "req_adapter_after_refresh").await;
assert_eq!(current.account_id(), old.account_id());
assert_eq!(current.credential_revision(), revision);
assert_ne!(current.credential_revision(), old.credential_revision());
assert_eq!(current.outbound_proxy(), Some(&new_proxy));
assert_eq!(old.outbound_proxy(), Some(&old_proxy));
for (connection, expected) in [
(&old, b"Bearer at-acct_provider_contract".as_slice()),
(&current, b"Bearer rotated-adapter-access-token".as_slice()),
] {
let headers = connection.authorization().unwrap();
assert_eq!(
headers
.iter()
.find(|header| header.name() == "authorization")
.unwrap()
.value(),
expected
);
assert_eq!(
headers
.iter()
.find(|header| header.name() == "chatgpt-account-id")
.unwrap()
.value(),
b"chatgpt-acct_provider_contract".as_slice()
);
}
// 旧请求的 401 不能使已轮换凭据失效,未发送失败也不能污染账号状态。
for (connection, send_state) in [
(&old, UpstreamSendState::Sent),
(&current, UpstreamSendState::NotSent),
] {
let error = connection
.record_failure(gateway_core::error::ProviderError::new(
ProviderErrorKind::Unauthorized,
send_state,
))
.await;
assert_eq!(error.kind(), ProviderErrorKind::Unauthorized);
assert_eq!(error.send_state(), send_state);
let after = store.account(account_id).unwrap();
assert_eq!(after.revision(), revision);
assert_eq!(after.credential_state(), CredentialState::Ready);
}
current
.record_failure(gateway_core::error::ProviderError::new(
ProviderErrorKind::Unauthorized,
UpstreamSendState::Sent,
))
.await;
assert_eq!(
store.account(account_id).unwrap().credential_state(),
CredentialState::Expired
);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
@@ -72,6 +72,7 @@ use crate::{GrokCatalogCapabilityEvidence, GrokCatalogModel};
mod continuation;
mod failure;
mod stream;
mod upstream_adapter;
mod workers;
use continuation::*;
@@ -266,8 +267,17 @@ impl GrokBuildProvider {
candidate: &ProviderCandidate,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let adapter =
context.upstream_adapter(candidate.provider(), candidate_upstream_model(candidate)?)?;
if adapter.is_none()
&& generate
.provider_session_state(XAI_PROVIDER_NAME)
.is_some_and(|state| state.extension_owner().is_some())
{
return Err(invalid_continuation());
}
let selected = self
.select_grok_attempt(generate, candidate, context)
.select_grok_attempt(generate, candidate, context, adapter.is_some())
.await?;
let provider_kind =
ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?;
@@ -283,6 +293,14 @@ impl GrokBuildProvider {
account_id,
Box::new(move |operation, middleware_headers| {
Box::pin(async move {
if let Some(adapter) = adapter {
return provider.execute_upstream_adapter(
operation,
middleware_headers,
selected,
adapter,
);
}
provider
.execute_selected_grok(operation, middleware_headers, selected)
.await
@@ -297,10 +315,15 @@ impl GrokBuildProvider {
generate: &GenerateRequest,
candidate: &ProviderCandidate,
context: AttemptContext,
adapted: bool,
) -> Result<SelectedGrokAttempt, ProviderError> {
let upstream_model = candidate_upstream_model(candidate)?;
let previous_session = decode_xai_session_state(generate)?;
let native_source = generate.protocol_payload().protocol() == "openai";
let previous_session = if adapted {
None
} else {
decode_xai_session_state(generate)?
};
let native_source = !adapted && generate.protocol_payload().protocol() == "openai";
let native_compaction = native_source
&& crate::transport::compaction::has_terminal_compaction_trigger(generate);
let (selection_model, operation_account, affinity) = if native_compaction {
@@ -0,0 +1,137 @@
//! xAI 已选 OAuth 会话的受管连接;插件不读取令牌或重新选号。
use super::*;
use gateway_core::{
account::{CredentialRevision, OutboundProxy},
engine::upstream_adapter::{
UpstreamAccountConnection, UpstreamAdapter, UpstreamAdapterInvocation,
},
metering::{CalculatedCost, Usage},
};
struct SelectedConnection {
provider: Arc<GrokBuildProvider>,
session: Arc<SelectedGrokSession>,
context: AttemptContext,
model: UpstreamModelId,
}
impl GrokBuildProvider {
pub(super) fn execute_upstream_adapter(
self: Arc<Self>,
operation: Operation,
headers: Vec<MiddlewareHeader>,
selected: SelectedGrokAttempt,
adapter: Arc<dyn UpstreamAdapter>,
) -> Result<ProviderStream, ProviderError> {
if operation.capability_requirements() != selected.frozen_requirements {
return Err(provider_error(
ProviderErrorKind::InvalidRequest,
UpstreamSendState::NotSent,
));
}
let session = Arc::new(selected.selected);
let metadata = ProviderCallMetadata::new(
ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| protocol_not_sent())?,
selected.upstream_model.clone(),
session.account_id().clone(),
UpstreamTransport::new(adapter.transport()).map_err(|_| protocol_not_sent())?,
)
.with_selection_observation(ProviderSelectionObservation::new(
selected.account_selection_wait_ms,
session.capacity_snapshot(),
));
let account = Arc::new(SelectedConnection {
provider: Arc::clone(&self),
session: Arc::clone(&session),
context: selected.context.clone(),
model: selected.upstream_model,
});
let events = adapter.execute(UpstreamAdapterInvocation {
operation,
headers,
context: selected.context,
metadata: metadata.clone(),
account,
});
let mutation = session.allows_account_state_mutation();
let stream = ProviderStream::new(metadata, events, session);
Ok(if mutation {
// 插件协议或 RPC 故障不能降低原生账号评分;只接纳已解析的上游拒绝事实。
stream.with_filtered_account_feedback(Arc::clone(&self.account_feedback), |error| {
error.client_visible_upstream_error().is_some()
})
} else {
stream
})
}
}
impl UpstreamAccountConnection for SelectedConnection {
fn account_id(&self) -> &ProviderAccountId {
self.session.account_id()
}
fn credential_revision(&self) -> CredentialRevision {
self.session.credential_revision()
}
fn authentication_kind(&self) -> &str {
crate::credential::XAI_AUTHENTICATION_KIND_OAUTH
}
fn outbound_proxy(&self) -> Option<&OutboundProxy> {
self.session.binding().outbound_proxy()
}
fn authorization(&self) -> Result<Vec<MiddlewareHeader>, ProviderError> {
Ok(vec![
MiddlewareHeader::new(
"authorization",
bytes::Bytes::from(format!("Bearer {}", self.session.access_token().expose())),
),
MiddlewareHeader::new(
"x-grok-user-id",
bytes::Bytes::copy_from_slice(self.session.user_id().expose().as_bytes()),
),
MiddlewareHeader::new(
"x-xai-token-auth",
bytes::Bytes::from_static(b"xai-grok-cli"),
),
])
}
fn calculate_cost(&self, tier: Option<&str>, usage: &Usage) -> Option<CalculatedCost> {
crate::transport::canonical::grok_billing_breakdown_with_override(
self.model.as_str(),
usage.input_tokens?,
usage.output_tokens?,
usage.cached_tokens.unwrap_or(0),
usage.cache_write_tokens.unwrap_or(0),
tier,
self.context
.pricing()
.get(XAI_PROVIDER_NAME)
.and_then(|models| models.get(self.model.as_str())),
)
.map(|cost| cost.calculated_cost())
}
fn record_failure(&self, mut error: ProviderError) -> BoxFuture<'_, ProviderError> {
Box::pin(async move {
if error.send_state() == UpstreamSendState::NotSent {
return error;
}
if error.kind() == ProviderErrorKind::Unauthorized {
error = error.with_credential_recovery();
}
let failure = stream_credential_failure(&error, &self.model);
recover_or_record_failure(
self.provider.selector.as_ref(),
self.provider.credential_recovery.as_ref(),
&self.session,
error,
failure,
self.context.credential_recovery_attempted(),
)
.await
})
}
}
@@ -4150,3 +4150,141 @@ async fn configured_request_profile_reaches_inference_headers() {
assert_eq!(header("x-grok-client-mode"), "headless");
assert_eq!(header("user-agent"), "grok-shell/9.8.7 (windows; aarch64)");
}
#[derive(Debug)]
struct AdapterProbe {
polls: Arc<std::sync::atomic::AtomicUsize>,
}
impl gateway_core::engine::upstream_adapter::UpstreamAdapterPlan for AdapterProbe {
fn select(
&self,
_: &AttemptContext,
provider: &ProviderKind,
_: &UpstreamModelId,
) -> Result<
Option<Arc<dyn gateway_core::engine::upstream_adapter::UpstreamAdapter>>,
gateway_core::error::ProviderError,
> {
assert_eq!(provider.as_str(), "xai");
Ok(Some(Arc::new(Self {
polls: self.polls.clone(),
})))
}
}
impl gateway_core::engine::upstream_adapter::UpstreamAdapter for AdapterProbe {
fn transport(&self) -> &str {
"http_sse"
}
fn execute(
self: Arc<Self>,
invocation: gateway_core::engine::upstream_adapter::UpstreamAdapterInvocation,
) -> gateway_core::engine::provider::EventStream {
Box::pin(futures::stream::once(async move {
self.polls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
assert_eq!(invocation.account.account_id().as_str(), "acct_provider");
assert_eq!(
invocation.metadata.provider_account_id(),
invocation.account.account_id()
);
assert_eq!(invocation.account.authentication_kind(), "oauth");
let headers = invocation.account.authorization().unwrap();
let header = |name: &str| {
headers
.iter()
.find(|h| h.name() == name)
.map(|h| h.value().to_vec())
};
assert_eq!(
header("authorization"),
Some("Bearer oauth-access".as_bytes().to_vec())
);
assert_eq!(
header("x-grok-user-id"),
Some("verified-user".as_bytes().to_vec())
);
assert!(header("cookie").is_none());
assert!(
invocation
.headers
.iter()
.any(|h| h.name() == "x-adapter-onion")
);
assert!(
invocation
.account
.calculate_cost(
None,
&gateway_core::metering::Usage {
input_tokens: Some(7),
output_tokens: Some(2),
..Default::default()
}
)
.is_some()
);
Err(gateway_core::error::ProviderError::new(
ProviderErrorKind::Cancelled,
UpstreamSendState::NotSent,
))
}))
}
}
#[tokio::test]
async fn upstream_adapter_reuses_selected_native_account_inside_attempt_onion_and_stays_cold() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::success();
let provider = provider(selector.clone(), transport.clone()).await;
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let lease = ExtensionSetReference::new(
ExtensionSetId::new("adapter-native-test".to_owned()).unwrap(),
Arc::new(TestExtensionLease),
);
let middleware = Arc::new(RecordingMiddleware {
observed: Default::default(),
replacement: Some(("input".into(), json!("rewritten"))),
request_headers: vec![MiddlewareHeader::new(
"x-adapter-onion",
Bytes::from_static(b"present"),
)],
});
let context = AttemptContext::new(
gateway_core::engine::RequestAttemptContext::new(
ModelRequestId::new("req_adapter_native").unwrap(),
ClientApiKeyId::new("key_xai_contract").unwrap(),
)
.with_upstream_adapters(Some(
gateway_core::engine::upstream_adapter::FrozenUpstreamAdapterPlan::new(
Arc::new(AdapterProbe {
polls: polls.clone(),
}),
lease.clone(),
),
))
.with_middleware(
Some(FrozenMiddlewarePlan::new(middleware, lease)),
Arc::from([]),
"/v1/responses".to_owned(),
ClientTransport::HttpSse,
),
NonZeroU32::MIN,
SystemTime::now() + Duration::from_secs(5),
selection_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None),
None,
CancellationToken::new(),
);
let mut stream = provider
.execute(provider_request("xai"), context)
.await
.unwrap();
assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 0);
let error = stream.next().await.unwrap().unwrap_err();
assert_eq!(error.kind(), ProviderErrorKind::Cancelled);
assert_eq!(polls.load(std::sync::atomic::Ordering::SeqCst), 1);
assert_eq!(selector.calls.load(Ordering::SeqCst), 1);
assert_eq!(transport.calls.load(Ordering::SeqCst), 0);
}
+5
View File
@@ -141,6 +141,11 @@ flowchart LR
正文按需读取,透传无需额外读取权限;转换和改写不能绕过 Core 的交付、取消与结算边界。
数据面插件通过受管 HTTP 已发送或无法证明未发送时,Runtime 把该事实并入 Core 的请求副作用水位,后续不能按
Provider 的 `not_sent` 结果透明重放
- **上游适配器是洋葱链终端**:`upstream_adapter` 只扩展内置 OpenAI / xAI,沿原 Provider 的选号、租约及 attempt 中间件进入冷流。
Provider 解释凭据、恢复账号状态和计算价格;Runtime 限定目标、绑定账号版本并解析标准事件,Host 负责真实出站。
普通 HTTP 与账号 HTTP 共用发送、背压、取消和正文回收,账号鉴权只装饰已授权请求。
WebSocket 连接续接绑定 Key、账号、凭据版本、实例和代次,不能当作可重排的中间件或全局共享句柄。
路由、调度、重试和终态观察保留各自端口,适配器不能另建执行引擎或提交最终账单
- **访问域不跨调用**:账号、HTTP、模型、私有状态和日志回调绑定有效父调用与阶段。管理页和 CLI 不预绑定 Client Key,模型调用按次
选择当前 Key;页面 Responses 桥还复核精确实例目标和 `models` 域,并在等待与交付期间持续撤销检查。
`frontend_authentication` 必须显式配置 principal 到 Key 的映射;公开登录回调使用一次性票据且不继承宿主回调权限。
@@ -12,6 +12,7 @@ const emit = defineEmits<{ validityChange: [valid: boolean] }>()
const bindings = defineModel<PluginCapabilityBinding[]>({ required: true })
const stages: Record<string, string[]> = {
middleware: ['request', 'attempt'],
upstream_adapter: ['upstream'],
model_router: ['routing'],
scheduler: ['scheduling'],
retry_policy: ['retry'],
@@ -19,7 +20,7 @@ const stages: Record<string, string[]> = {
usage: ['observation'],
web_socket_observer: ['observation'],
}
const stageLabels: Record<string, string> = { request: '请求开始', attempt: '每次尝试', routing: '模型路由', scheduling: '账号调度', retry: '重试决策', observation: '请求结束' }
const stageLabels: Record<string, string> = { request: '请求开始', attempt: '每次尝试', upstream: '上游调用', routing: '模型路由', scheduling: '账号调度', retry: '重试决策', observation: '请求结束' }
const entries = computed(() => Object.entries(props.metadata.contributes).flatMap(([capability, contribution]) =>
contribution.stages.filter(stage => stages[capability]?.includes(stage))
.map(stage => ({ capability, contribution: contribution.id, stage, key: `${contribution.id}:${stage}` })),
@@ -98,8 +99,8 @@ watch(valid, value => emit('validityChange', value), { immediate: true })
<PluginScopeInput :model-value="bindingFor(entry.key)!.models" label="模型范围" placeholder="每行一个模型名称,精确匹配" :class="['routing', 'request'].includes(entry.stage) ? 'sm:col-span-2' : undefined" :disabled="disabled" @update:model-value="patch(entry.key, { models: $event })" />
<PluginScopeInput v-if="!['routing', 'request'].includes(entry.stage)" :model-value="bindingFor(entry.key)!.providerIds" label="Provider 范围" placeholder="每行一个 Provider 标识" :disabled="disabled" @update:model-value="patch(entry.key, { providerIds: $event })" />
</div>
<div v-if="entry.stage !== 'observation'" class="grid items-end gap-4 sm:grid-cols-2">
<BaseFormItem label="执行顺序">
<div v-if="entry.stage !== 'observation'" class="grid items-end gap-4" :class="entry.stage !== 'upstream' ? 'sm:grid-cols-2' : undefined">
<BaseFormItem v-if="entry.stage !== 'upstream'" label="执行顺序">
<template #label-extra>
<PluginHelpPopover label="执行顺序说明">
数值较小的先执行
@@ -108,7 +109,7 @@ watch(valid, value => emit('validityChange', value), { immediate: true })
<BaseNumberInput :model-value="bindingFor(entry.key)!.order" label="执行顺序" size="md" :min="-2147483648" :max="2147483647" :disabled="disabled" class="w-full" @update:model-value="patch(entry.key, { order: $event })" />
</BaseFormItem>
<BaseFormItem label="插件失败时">
<BaseSelect :model-value="bindingFor(entry.key)!.failurePolicy" :options="entry.stage === 'retry' ? [{ label: '交给后续处理', value: 'delegate' }] : [{ label: '拒绝请求', value: 'reject' }, { label: '交给后续处理', value: 'delegate' }]" :disabled="disabled || entry.stage === 'retry'" class="w-full" @update:model-value="patch(entry.key, { failurePolicy: $event as 'reject' | 'delegate' })" />
<BaseSelect :model-value="bindingFor(entry.key)!.failurePolicy" :options="entry.stage === 'retry' ? [{ label: '交给后续处理', value: 'delegate' }] : entry.stage === 'upstream' ? [{ label: '拒绝请求', value: 'reject' }] : [{ label: '拒绝请求', value: 'reject' }, { label: '交给后续处理', value: 'delegate' }]" :disabled="disabled || ['retry', 'upstream'].includes(entry.stage)" class="w-full" @update:model-value="patch(entry.key, { failurePolicy: $event as 'reject' | 'delegate' })" />
</BaseFormItem>
</div>
</template>
@@ -37,7 +37,7 @@ const authenticationBinding = computed({
bindings.value = binding ? [...others, cloneJsonValue(binding)] : others
},
})
const hasRequestBindings = computed(() => Object.values(props.artifact?.metadata.contributes ?? {}).some(contribution => contribution.stages.some(stage => ['request', 'attempt', 'routing', 'scheduling', 'retry', 'observation'].includes(stage))))
const hasRequestBindings = computed(() => Object.values(props.artifact?.metadata.contributes ?? {}).some(contribution => contribution.stages.some(stage => ['request', 'attempt', 'upstream', 'routing', 'scheduling', 'retry', 'observation'].includes(stage))))
const sections = computed(() => [
{ label: '插件参数', value: 'general', icon: Settings2 },
...(hasRequestBindings.value ? [{ label: '高级设置', value: 'requests', icon: Blocks }] : []),
@@ -14,6 +14,7 @@ const props = defineProps<{
const stageLabels: Record<string, string> = {
request: '请求开始',
attempt: '每次尝试',
upstream: '上游调用',
routing: '模型路由',
scheduling: '账号调度',
retry: '重试决策',
@@ -23,7 +24,7 @@ const stageLabels: Record<string, string> = {
maintenance: '后台维护',
authentication: '客户端认证',
}
const requestStages = new Set(['request', 'attempt', 'routing', 'scheduling', 'retry', 'observation'])
const requestStages = new Set(['request', 'attempt', 'upstream', 'routing', 'scheduling', 'retry', 'observation'])
const groups = computed(() => {
const result = new Map<string, { label: string, bindings: PluginCapabilityBinding[] }>()
for (const binding of props.instance.bindings) {
@@ -41,6 +41,7 @@ export const PLUGIN_CAPABILITY_LABELS: Record<string, string> = {
retry_policy: '重试策略',
executor: 'Provider 执行',
middleware: '请求中间件',
upstream_adapter: '上游适配器',
request_lifecycle: '请求生命周期',
web_socket_observer: 'WebSocket 观察',
usage: '用量处理',