diff --git a/Cargo.lock b/Cargo.lock index 8f33b1f4..d06d1bc8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -576,6 +576,7 @@ dependencies = [ "thiserror", "tokio", "tracing", + "uuid", ] [[package]] diff --git a/_context/wiki/architecture.md b/_context/wiki/architecture.md index ae34d161..bf64ea41 100644 --- a/_context/wiki/architecture.md +++ b/_context/wiki/architecture.md @@ -16,6 +16,8 @@ TCP/TLS listener -> CORS layer -> mcp_header_limits_layer bounds MCP headers (431) -> virtual_host_id_layer inserts VirtualHostId from path (400) + -> HTTP pre hook global plugins, permitted request-header edits + -> mcp_header_limits_layer rechecks edited MCP headers (431) -> claims_layer verifies JWT, inserts AuthorizationClaims (401) -> PrincipalExtractorLayer inserts AuthorizedPrincipal (401) -> user_config_store_layer loads UserConfig (400 missing, 500 decode/error) @@ -26,7 +28,9 @@ TCP/TLS listener Origin is checked before authentication. The optional Host allowlist is checked at the RMCP boundary, so earlier middleware may return first. MCP header budgets -apply before JWT verification, configuration reads, and body parsing. See +apply before JWT verification, user-configuration reads, and body parsing. The +HTTP post hook runs as the response returns through the HTTP plugin layer, +before headers are sent; streaming MCP work may still be running. See [Security](security.md#mcp-origin-and-host-validation). Health is registered in every build outside the MCP auth/config layers, while @@ -89,7 +93,8 @@ into successful response hooks. See [Routing](routing.md) and | Principal, claims, virtual-host ID, config snapshot | HTTP request extensions | One request. | | Backend RMCP service | Routed operation | One request; explicitly closed after the call. | | Tool progress-token mapping | Request's backend client | While the tool call is in flight. | -| Active CPEX runtime | Registry | Reloadable; in-flight hook state pins its selected runtime. | +| Active CPEX policies | Registry | Reloadable; each HTTP request pins the full policy snapshot. | +| Plugin request context | Request-owned `PluginRequest` | Shared CPEX state and permitted extension updates across HTTP and MCP hooks, including streaming work. | | RMCP session manager | RMCP service (`LocalSessionManager`) | Transport implementation detail; no session is required by the supported modern request contract. | The library no longer has `BackendTransports`, `SessionId` middleware, or a @@ -114,8 +119,8 @@ Locks have specific scopes: shared state; do not assume all network I/O is globally lock-free. - Tool progress tracking holds a write guard while enqueuing the backend call so an early notification cannot race registration. -- Tool hook state uses a mutex to serialize progress and final-response plugin - context updates. Prompt/resource state belongs to one request. +- A request-scoped plugin mutex serializes HTTP/MCP hook state updates, including + tool progress and final responses. No registry lock is held during execution. There is no shared map of live backend transports to lock during routing. @@ -138,11 +143,15 @@ post-hooks can inspect stream events, and a denied notification is dropped. Prompt/resource operations should not be assumed to have the same explicit cancellation relay. -Pre-hooks select a runtime before backend I/O. Typed hook state retains both -that runtime and whether a post-hook was enabled. Reloads affect subsequent -requests; they cannot add a hook or change policy halfway through a call. -Invalid reloads mark the registry failed for new calls while already pinned -requests can finish. Details are in [Plugin Config](config.md#plugin-config-redis-key-contextforgegatewayruntimepluginconfig). +The HTTP boundary pins the full policy snapshot in one `PluginRequest`. HTTP +hooks use the global policy; tool hooks resolve published tool/team contexts, +and prompt/resource hooks use the global policy. Authentication and user-config +lookup attach verified identity; routing adds target metadata. Each operation +retains its selected runtime and matched pre/post hooks. Reloads affect new HTTP +requests; in-flight work keeps its policy and state, including streaming events. +Invalid reloads fail new requests while already pinned requests can finish. +See [HTTP hooks and shared request state](config.md#http-hooks-and-shared-request-state) +for ordering, failure, timeout, cancellation and streaming behavior. ## Startup diff --git a/_context/wiki/config.md b/_context/wiki/config.md index f9f9e222..42fa36b6 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -244,7 +244,7 @@ Tools select the resolved configuration using the tool's canonical `id`, gateway `name`, and owning `team_id`. Aliases select the same upstream tool context. With plugins enabled, a missing tool identity or policy context fails the call before backend I/O; it never falls back to global -policy. Prompts and resources use `global`, matching the built-in +policy. Prompts, resources and HTTP hooks use `global`, matching the built-in resolver, which only applies database bindings to team/tool context keys. All targets use the same policy resolver and execution engine. Only `enabled: false` explicitly bypasses the published policies. @@ -252,8 +252,8 @@ bypasses the published policies. The existing watcher checks Redis every 30 seconds, matching the built-in manager's default cache lifetime. Publication adds its own interval. A reload builds all policies before swapping the active set; invalid, expired or missing -configuration fails new requests. Each MCP operation pins its selected policy, so a reload cannot change it -halfway through. Operations retain their selected +configuration fails new requests. Each HTTP request pins the full policy snapshot, +so its MCP operations cannot pick up a newer policy halfway through. Operations retain their selected runtime, matched pre/post hooks, CPEX local/shared state, and permitted extension updates, including post-only hooks and tool progress/logging events. @@ -268,7 +268,7 @@ a policy or establish identity. Within each execution mode, hooks run in ascending numeric priority (1 before 90). Equal priorities retain configuration order. This order is compiled into -each policy snapshot and applies to all MCP pre/post hooks. +each policy snapshot and applies to both HTTP and MCP pre/post hooks. CPEX filters each plugin's view by its declared capabilities. The gateway guards every extension write-back before the next plugin executes: identity remains @@ -285,12 +285,51 @@ five seconds per plugin in both pre and post hooks. CPEX enforces this with `on_error` policy. This does not bound backend work or the total time for a chain of plugins. Like other asynchronous timeouts, handlers must yield to be cancelled. -Supported gateway hooks are tool, prompt and resource pre/post hooks. Unsupported +Supported gateway hooks are HTTP, tool, prompt and resource pre/post hooks. Unsupported hooks, route-based CPEX configuration, plugin directories and missing factories fail configuration loading. Redis policy cannot load Python modules or add new Rust factories. Authentication-resolution/permission hooks and the broader execution-pool and plugin-load-error parity work remain outside this adapter. +### HTTP hooks and shared request state + +`http_pre_request` runs after the existing Origin/header-budget checks and before +credential verification. `http_post_request` runs when response headers are +available, including authentication failures. Both use the published global +configuration. The native `HttpHook` payload provides method, original path, +client address when available, and response status on the post hook. Headers +are CPEX extensions, gated by `read_headers`, `write_headers`, and the published +hook's `headers` write policy. Use `GatewayPluginFactory` with `.with_http_hooks()` +and/or `.with_cmf_hooks()` to register handlers on the same plugin instance. + +Header changes are merged before the next plugin runs, preserving existing +request/response headers even when the writer lacks `read_headers`. Method, +path, host and scheme remain host-provided metadata. Invalid header updates +are ignored. Existing auth headers are protected unless +`plugins_can_override_auth_headers` is true. The +before hook may add a previously absent credential header, which normal +authentication must still verify. MCP header budgets are checked again after +plugin edits. HTTP hooks do not modify request bodies, response bodies or status. + +One `PluginRequest` carries the policy snapshot, correlation ID, CPEX state table +and permitted extension changes across all stages. Scoped/global managers give +plugins different runtime IDs; request-local state follows the configured plugin +kind and name when crossing that boundary. Verified identity comes only from the +host after authentication and user-configuration lookup, even if that lookup fails. +MCP routing adds target metadata to this existing state; it does not reconstruct +transport data or identity. A request-scoped mutex serializes hook state updates; +no registry lock is held while plugins run. + +As in the built-in HTTP middleware, HTTP hook errors, denials and timeouts are +logged and the request continues. The configured CPEX timeout bounds execution; +missing required configuration still fails the request. Dropping the request +future cancels foreground hook work. Invocation spans use `gateway_plugin_invoke` +with the hook name; logs omit header values. + +Streaming responses are not buffered. The HTTP post hook runs once before headers +are sent, so it cannot observe MCP work that completes later in the stream. +Those later MCP hooks retain the same request state and policy snapshot. + Compile bundled factories with `--features plugins` and enable execution with `--runtime-plugins-enabled true`. A valid document must exist before startup; a missing document fails initialization. The @@ -360,7 +399,7 @@ docker compose -f docker/docker-compose-local.yaml exec -T redis \ This example uses the `quickstart-counter` policy key from the quick-start routes. For other tools, publish a matching `contexts[context_id]` entry; -`global` alone applies to prompt and resource hooks. +`global` alone applies to prompt, resource and HTTP hooks. For local testing only, build and run with demo factories, `with_tools` helpers, and runtime execution enabled: diff --git a/_context/wiki/deployment.md b/_context/wiki/deployment.md index e577abb0..cd729f65 100644 --- a/_context/wiki/deployment.md +++ b/_context/wiki/deployment.md @@ -60,7 +60,7 @@ configuration-cache freshness and dependency availability. - A missing entry or Redis GET error currently produces HTTP `400`; undecodable configuration produces `500`. See [Failure Modes](failure-modes.md). - Enabled CPEX also needs its initial plugin document and checks for reloads - every 30 seconds. An invalid reload fails new plugin calls closed. + every 30 seconds. An invalid reload fails new HTTP requests closed. ## Builds and Images diff --git a/_context/wiki/failure-modes.md b/_context/wiki/failure-modes.md index 6711b59d..f5b14688 100644 --- a/_context/wiki/failure-modes.md +++ b/_context/wiki/failure-modes.md @@ -54,11 +54,12 @@ modern routing. A configured backend alone does not make its objects callable. | --- | --- | | MCP pre-hook denies | MCP error; no upstream call. | | MCP post-hook denies | MCP error; backend operation may already have completed. | +| HTTP hook error, denial or timeout | Logged; request/response processing continues. | | Tool identity or resolved policy context missing | MCP error before backend I/O; no fallback to global policy. | | Plugin supplies an error code | That code is used; a denial without one defaults to invalid request `-32600`. | | Soft plugin error | Logged; execution can continue under the runtime's soft-error behavior. | | Missing or invalid initial plugin config | Runtime initialization fails; gateway startup does not complete. | -| Invalid plugin reload | New plugin calls fail closed until valid configuration is loaded; already pinned requests keep their runtime. | +| Invalid plugin reload | New HTTP requests fail with `500` until valid configuration is loaded; already pinned requests keep their policies. | | Plugin edits cannot be represented faithfully as the operation's MCP result | MCP error; no fallback to the original unredacted response. | See [Configuration](config.md#plugin-config-redis-key-contextforgegatewayruntimepluginconfig) diff --git a/_context/wiki/routing.md b/_context/wiki/routing.md index e18b6966..44bbd9b3 100644 --- a/_context/wiki/routing.md +++ b/_context/wiki/routing.md @@ -66,14 +66,12 @@ validation. tool/prompt arguments or the resource URI. Resource URI edits must resolve through the caller's published routes. Post-hooks can rewrite or reject the response. -Tools select their published tool/team policy before backend I/O; prompts and -resources select the global policy. The host supplies verified identity and -canonical route metadata to each hook. It returns typed request state -whose `after_*` method runs the post-hook on that same runtime, even after a reload -or a reload failure. A request that started without post-hooks never gains one -mid-flight. Tool state is shared under a mutex so progress notifications and the -final response use the same correlation ID and serialize plugin-context updates. -Prompt and resource state is owned by a single request and needs no mutex. +The HTTP layer pins the policy snapshot before authentication. Tool hooks select +the routed backend's published tool policy context; prompt, resource and HTTP +hooks use the global policy. Typed operation state runs pre/post hooks on the +same runtime, even after a reload or reload failure. One request-scoped mutex +serializes shared plugin state across HTTP/MCP hooks and tool progress events. +A request that started without post-hooks never gains one mid-flight. The internal CPEX crate separates these responsibilities: diff --git a/crates/contextforge-data-plane-cpex/Cargo.toml b/crates/contextforge-data-plane-cpex/Cargo.toml index 4621dc2b..6ecf01f4 100644 --- a/crates/contextforge-data-plane-cpex/Cargo.toml +++ b/crates/contextforge-data-plane-cpex/Cargo.toml @@ -26,6 +26,7 @@ serde_json.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true +uuid.workspace = true [lints] workspace = true diff --git a/crates/contextforge-data-plane-cpex/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index aa39e33c..ee93b392 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -1,48 +1,13 @@ //! Common CMF envelopes and the response contract used by the hook runner. //! Conversion rules stay with each operation because their MCP semantics differ. +use crate::hooks::Operation; use cpex::cpex_core::{ cmf::{ContentPart, Message, MessagePayload, Role, constants::SCHEMA_VERSION}, executor::PipelineResult, - hooks::types::cmf_hook_names, }; use rmcp::{ErrorData, model::ErrorCode}; -#[derive(Clone, Debug, Copy)] -pub(crate) enum Operation { - Tool, - Prompt, - Resource, -} - -impl Operation { - pub(crate) const ALL: [Self; 3] = [Self::Tool, Self::Prompt, Self::Resource]; - - pub(crate) fn hooks(self) -> [&'static str; 2] { - match self { - Self::Tool => [cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE], - Self::Prompt => [cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], - Self::Resource => [cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], - } - } - - pub(crate) fn subject(self) -> &'static str { - match self { - Self::Tool => "tool call", - Self::Prompt => "prompt", - Self::Resource => "resource", - } - } - - pub(crate) fn id_prefix(self) -> &'static str { - match self { - Self::Tool => "gateway-tool-call", - Self::Prompt => "gateway-prompt-request", - Self::Resource => "gateway-resource-request", - } - } -} - /// Only the projection and application differ between response types. /// The runner owns invocation, unchanged payloads, context, and denial handling. pub(crate) trait CmfResponse: Sized { diff --git a/crates/contextforge-data-plane-cpex/src/extension_guard.rs b/crates/contextforge-data-plane-cpex/src/extension_guard.rs index d4c3cdcf..9cd608f3 100644 --- a/crates/contextforge-data-plane-cpex/src/extension_guard.rs +++ b/crates/contextforge-data-plane-cpex/src/extension_guard.rs @@ -13,18 +13,26 @@ use cpex::cpex_core::{ registry::{AnyHookHandler, HookEntry}, }; +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum HeaderWritePolicy { + ReadOnly, + PreserveCredentials, + AllowNewCredentials, + All, +} + struct ExtensionGuard { entry: HookEntry, canonical: Arc>, payload_writable: bool, - auth_headers_writable: bool, + header_policy: HeaderWritePolicy, } pub(crate) fn guarded_entries( entries: &[HookEntry], extensions: &Extensions, payload_writable: bool, - auth_headers_writable: bool, + header_policy: HeaderWritePolicy, ) -> Vec { let canonical = Arc::new(Mutex::new(extensions.clone())); entries @@ -35,7 +43,7 @@ pub(crate) fn guarded_entries( entry: entry.clone(), canonical: Arc::clone(&canonical), payload_writable, - auth_headers_writable, + header_policy, }), }) .collect() @@ -66,7 +74,8 @@ impl AnyHookHandler for ExtensionGuard { if canonical.validate_immutable(&modified) { let caps = &self.entry.plugin_ref.trusted_config().capabilities; let mut permitted = canonical.cow_copy(); - if caps.contains("write_headers") + if self.header_policy != HeaderWritePolicy::ReadOnly + && caps.contains("write_headers") && let Some(http) = modified.http { let mut http = http.into_inner(); @@ -78,8 +87,12 @@ impl AnyHookHandler for ExtensionGuard { &mut http.response_headers, canonical.http.as_ref().map(|http| &http.response_headers), )?; - if !self.auth_headers_writable { - preserve_auth_headers(&mut http.request_headers, canonical.http.as_deref()); + if self.header_policy != HeaderWritePolicy::All { + preserve_auth_headers( + &mut http.request_headers, + canonical.http.as_deref(), + self.header_policy == HeaderWritePolicy::AllowNewCredentials, + ); } // A write-only plugin cannot return headers or metadata it never saw. // Merge header edits while retaining the host's transport metadata. @@ -145,13 +158,16 @@ fn normalize_headers( fn preserve_auth_headers( headers: &mut std::collections::HashMap, original: Option<&cpex::cpex_core::extensions::HttpExtension>, + allow_new: bool, ) { fn protected(name: &str) -> bool { ["authorization", "proxy-authorization", "cookie", "x-api-key"] .iter() .any(|header| name.eq_ignore_ascii_case(header)) } - headers.retain(|name, _| !protected(name)); + headers.retain(|name, _| { + !protected(name) || (allow_new && !original.is_some_and(|http| http.has_request_header(name))) + }); if let Some(original) = original { headers.extend( original diff --git a/crates/contextforge-data-plane-cpex/src/factory.rs b/crates/contextforge-data-plane-cpex/src/factory.rs index 3034fbbd..7421ed3c 100644 --- a/crates/contextforge-data-plane-cpex/src/factory.rs +++ b/crates/contextforge-data-plane-cpex/src/factory.rs @@ -9,42 +9,71 @@ use cpex::cpex_core::{ registry::AnyHookHandler, }; -pub struct CmfPluginFactory

{ +use crate::{HttpHook, hooks::Operation}; + +type HandlerFactory

= fn(Arc

) -> Arc; + +/// Registers HTTP and/or MCP handlers on the same native plugin instance. +#[must_use] +pub struct GatewayPluginFactory

{ build: Box P + Send + Sync>, + handlers: Vec<(&'static str, HandlerFactory

)>, } -impl

CmfPluginFactory

{ +impl GatewayPluginFactory

{ pub fn new(build: impl Fn(PluginConfig) -> P + Send + Sync + 'static) -> Self { - Self { build: Box::new(build) } + Self { build: Box::new(build), handlers: Vec::new() } + } + + pub fn with_cmf_hooks(mut self) -> Self + where + P: HookHandler, + { + self.handlers.extend(Operation::MCP.into_iter().flat_map(Operation::hooks).map(|name| { + ( + name, + (|plugin| Arc::new(TypedHandlerAdapter::::new(plugin)) as Arc) + as HandlerFactory

, + ) + })); + self + } + + pub fn with_http_hooks(mut self) -> Self + where + P: HookHandler, + { + self.handlers.extend(Operation::Http.hooks().map(|name| { + ( + name, + (|plugin| Arc::new(TypedHandlerAdapter::::new(plugin)) as Arc) + as HandlerFactory

, + ) + })); + self } } -impl

PluginFactory for CmfPluginFactory

-where - P: Plugin + HookHandler + 'static, -{ +impl PluginFactory for GatewayPluginFactory

{ fn create(&self, config: &PluginConfig) -> Result> { let plugin = Arc::new((self.build)(config.clone())); let handlers = config .hooks .iter() - .filter_map(|hook| supported_cmf_hook_name(hook)) .map(|hook| { - (hook, Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) as Arc) + let (name, build) = self.handlers.iter().find(|(name, _)| *name == hook).ok_or_else(|| { + Box::new(PluginError::Config { + message: format!("plugin '{}' has no handler for '{hook}'", config.name), + }) + })?; + Ok((*name, build(Arc::clone(&plugin)))) }) - .collect::>(); - - if handlers.is_empty() { - return Err(Box::new(PluginError::Config { - message: format!("plugin '{}' does not declare supported CMF hooks", config.name), - })); - } - + .collect::, Box>>()?; let plugin: Arc = plugin; Ok(PluginInstance { plugin, handlers }) } } -pub(crate) fn supported_cmf_hook_name(hook: &str) -> Option<&'static str> { - crate::cmf::Operation::ALL.into_iter().flat_map(crate::cmf::Operation::hooks).find(|name| *name == hook) +pub(crate) fn supported_hook_name(hook: &str) -> Option<&'static str> { + Operation::ALL.into_iter().flat_map(Operation::hooks).find(|name| *name == hook) } diff --git a/crates/contextforge-data-plane-cpex/src/hooks.rs b/crates/contextforge-data-plane-cpex/src/hooks.rs index 46d5d5a0..20407ba3 100644 --- a/crates/contextforge-data-plane-cpex/src/hooks.rs +++ b/crates/contextforge-data-plane-cpex/src/hooks.rs @@ -1,12 +1,15 @@ use contextforge_data_plane_apis::user_store::ToolPolicyContext; -use cpex::cpex_core::hooks::Extensions; +use cpex::cpex_core::hooks::{Extensions, types::cmf_hook_names}; use serde_json::{Map, Value}; /// Host-authored context from the authorized route and verified request. /// Client metadata and plugin state never select a policy or establish identity. #[derive(Default)] pub struct PluginRequestContext { + pub request: Option, pub tool: Option, + /// With an existing request, only routed `meta` and `mcp` fields are applied. + /// Transport data and verified identity belong to the shared request. pub extensions: Extensions, } @@ -46,3 +49,43 @@ impl Default for PreHookResult { Self { arguments: ArgumentsUpdate::Unchanged, state: None } } } + +#[derive(Clone, Debug, Copy)] +pub(crate) enum Operation { + Tool, + Prompt, + Resource, + Http, +} + +impl Operation { + pub(crate) const MCP: [Self; 3] = [Self::Tool, Self::Prompt, Self::Resource]; + pub(crate) const ALL: [Self; 4] = [Self::Tool, Self::Prompt, Self::Resource, Self::Http]; + + pub(crate) fn hooks(self) -> [&'static str; 2] { + match self { + Self::Tool => [cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE], + Self::Prompt => [cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + Self::Resource => [cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + Self::Http => ["http_pre_request", "http_post_request"], + } + } + + pub(crate) fn subject(self) -> &'static str { + match self { + Self::Tool => "tool call", + Self::Prompt => "prompt", + Self::Resource => "resource", + Self::Http => "HTTP request", + } + } + + pub(crate) fn id_prefix(self) -> &'static str { + match self { + Self::Tool => "gateway-tool-call", + Self::Prompt => "gateway-prompt-request", + Self::Resource => "gateway-resource-request", + Self::Http => "gateway-http-request", + } + } +} diff --git a/crates/contextforge-data-plane-cpex/src/http.rs b/crates/contextforge-data-plane-cpex/src/http.rs new file mode 100644 index 00000000..aaf5018c --- /dev/null +++ b/crates/contextforge-data-plane-cpex/src/http.rs @@ -0,0 +1,79 @@ +//! HTTP metadata is typed; header access and changes use capability-gated CPEX extensions. +use std::{net::SocketAddr, sync::Arc}; + +use cpex::cpex_core::{ + extensions::{Extensions, HttpExtension}, + hooks::PluginResult, +}; +use rmcp::ErrorData; + +use crate::{GatewayPluginRuntimeHandle, PluginRequest, PluginRequestContext, hooks::Operation}; + +#[derive(Clone, Debug)] +pub struct HttpHookPayload { + pub method: String, + pub path: String, + pub client_addr: Option, + /// Present only on the after hook. The hook runs before response headers are sent. + pub status_code: Option, +} + +cpex::cpex_core::impl_plugin_payload!(HttpHookPayload); +cpex::cpex_core::define_hook! { + HttpHook, "http" => { payload: HttpHookPayload, result: PluginResult } +} + +pub struct HttpHookState { + pub request: PluginRequest, + payload: HttpHookPayload, + runtime: Arc, +} + +impl GatewayPluginRuntimeHandle { + pub async fn before_http_request( + &self, + payload: HttpHookPayload, + extensions: Extensions, + ) -> Result { + let context = PluginRequestContext { extensions, ..Default::default() }; + let (request, runtime) = self.resolve(Operation::Http, &context)?; + let state = HttpHookState { request, payload, runtime }; + state.invoke(0).await; + Ok(state) + } +} + +impl HttpHookState { + pub async fn after_http_request( + mut self, + status_code: u16, + response_headers: std::collections::HashMap, + ) -> Extensions { + self.payload.status_code = Some(status_code); + { + let mut state = self.request.state.lock().await; + let http = state.extensions.http.get_or_insert_with(|| Arc::new(HttpExtension::default())); + Arc::make_mut(http).response_headers = response_headers; + } + self.invoke(1).await; + self.request.extensions().await + } + + async fn invoke(&self, phase: usize) { + // The same resolver and invocation engine serve all target kinds. + let runtime = &self.runtime; + let mut host = self.request.extensions().await; + // The HTTP boundary has no routed MCP target, including during streaming. + host.meta = None; + host.mcp = None; + let entries = runtime.entries(Operation::Http, phase, &host); + let hook = Operation::Http.hooks()[phase]; + let result = runtime + .invoke::((Operation::Http, phase, &entries), self.payload.clone(), &self.request, &host) + .await; + if result.is_denied() { + // Match the built-in HTTP middleware: hook failures do not reject HTTP requests. + tracing::warn!(hook, "HTTP plugin hook did not complete; continuing request"); + } + } +} diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index 190f7a71..d9b80d32 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -4,16 +4,20 @@ mod error; mod extension_guard; mod factory; mod hooks; +mod http; mod prompts; mod registry; +mod request; mod resources; mod runtime; mod tools; pub use error::GatewayPluginRuntimeError; -pub use factory::CmfPluginFactory; +pub use factory::GatewayPluginFactory; pub use hooks::{ArgumentsUpdate, PluginRequestContext, PreHookResult, RuntimeHookError}; +pub use http::{HttpHook, HttpHookPayload, HttpHookState}; pub use prompts::PromptHookState; pub use registry::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; +pub use request::PluginRequest; pub use resources::ResourceHookState; pub use tools::ToolHookState; diff --git a/crates/contextforge-data-plane-cpex/src/prompts/mod.rs b/crates/contextforge-data-plane-cpex/src/prompts/mod.rs index ae46b042..a8368b27 100644 --- a/crates/contextforge-data-plane-cpex/src/prompts/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/prompts/mod.rs @@ -15,7 +15,8 @@ use serde_json::{Map, Value}; use crate::{ ArgumentsUpdate, GatewayPluginRuntimeHandle, PluginRequestContext, PreHookResult, - cmf::{CmfResponse, Operation, message_payload}, + cmf::{CmfResponse, message_payload}, + hooks::Operation, runtime::CallState, }; @@ -243,12 +244,11 @@ impl GatewayPluginRuntimeHandle { backend_name: &str, context: PluginRequestContext, ) -> Result, ErrorData> { - let (arguments, state) = self - .current()? - .global + let (request_state, runtime) = self.resolve(Operation::Prompt, &context)?; + let (arguments, state) = runtime .before( (Operation::Prompt, prompt_name), - context.extensions, + (request_state, context.extensions), |id| prompt_request_payload(request, prompt_name, backend_name, id), |payload, id| { let arguments = @@ -264,7 +264,7 @@ impl GatewayPluginRuntimeHandle { } impl PromptHookState { - pub async fn after_get_prompt(mut self, response: GetPromptResult) -> Result { + pub async fn after_get_prompt(self, response: GetPromptResult) -> Result { self.0.after(response).await } } diff --git a/crates/contextforge-data-plane-cpex/src/registry/mod.rs b/crates/contextforge-data-plane-cpex/src/registry/mod.rs index f47fb981..faee23cb 100644 --- a/crates/contextforge-data-plane-cpex/src/registry/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/registry/mod.rs @@ -17,9 +17,10 @@ use rmcp::{ErrorData, model::ErrorCode}; use tokio::task::JoinHandle; use crate::{ + PluginRequest, PluginRequestContext, config::{RedisRuntimePluginConfigStore, RuntimePluginConfigStore}, error::GatewayPluginRuntimeError, - hooks::RuntimeHookError, + hooks::{Operation, RuntimeHookError}, runtime::GatewayPluginRuntime, }; @@ -51,8 +52,12 @@ pub(crate) struct RuntimePolicies { } impl RuntimePolicies { - pub(crate) fn tool(&self, tool: Option<&ToolPolicyContext>) -> Result, ErrorData> { - if !self.scoped { + pub(crate) fn resolve( + &self, + operation: Operation, + tool: Option<&ToolPolicyContext>, + ) -> Result, ErrorData> { + if !matches!(operation, Operation::Tool) || !self.scoped { return Ok(Arc::clone(&self.global)); } let tool = tool @@ -211,6 +216,19 @@ impl CpexRuntimeRegistry { } impl GatewayPluginRuntimeHandle { + pub(crate) fn resolve( + &self, + operation: Operation, + context: &PluginRequestContext, + ) -> Result<(PluginRequest, Arc), ErrorData> { + let request = match &context.request { + Some(request) => request.clone(), + None => PluginRequest::new(self.current()?, context.extensions.clone()), + }; + let runtime = request.policies.resolve(operation, context.tool.as_ref())?; + Ok((request, runtime)) + } + pub(crate) fn current(&self) -> Result, ErrorData> { match self.runtime.load().as_ref() { RuntimeState::Active(runtime) => Ok(Arc::clone(runtime)), diff --git a/crates/contextforge-data-plane-cpex/src/registry/tests.rs b/crates/contextforge-data-plane-cpex/src/registry/tests.rs index 431b0804..47b767bd 100644 --- a/crates/contextforge-data-plane-cpex/src/registry/tests.rs +++ b/crates/contextforge-data-plane-cpex/src/registry/tests.rs @@ -8,6 +8,7 @@ use std::{ }; use async_trait::async_trait; +use contextforge_data_plane_apis::runtime_plugin_config::{RuntimePluginConfigDocument, RuntimePluginSettings}; use cpex::cpex_core::{ cmf::{CmfHook, ContentPart, MessagePayload}, context::PluginContext, @@ -18,19 +19,14 @@ use cpex::cpex_core::{ registry::AnyHookHandler, }; use rmcp::model::{ - CallToolRequestParams, CallToolResult, ContentBlock, NumberOrString, ProgressNotificationParam, ProgressToken, - ReadResourceResult, ResourceContents, + CallToolRequestParams, CallToolResult, ContentBlock, GetPromptRequestParams, NumberOrString, + ProgressNotificationParam, ProgressToken, ReadResourceResult, ResourceContents, }; use serde_json::{Value, json}; use tokio::sync::Mutex as TokioMutex; -use contextforge_data_plane_apis::runtime_plugin_config::{RuntimePluginConfigDocument, RuntimePluginSettings}; - -use crate::config::LoadedRuntimePluginConfig; -use crate::{ArgumentsUpdate, CmfPluginFactory, PreHookResult, ToolHookState}; -use rmcp::model::GetPromptRequestParams; - use super::*; +use crate::{ArgumentsUpdate, GatewayPluginFactory, PreHookResult, ToolHookState, config::LoadedRuntimePluginConfig}; const TEST_MISSING_CONTEXT_ERROR_CODE: i64 = -32003; const TEST_REWRITTEN_SUM_A: i64 = 10; @@ -311,7 +307,7 @@ impl PluginFactory for TestPluginFactory { .hooks .iter() .filter_map(|hook| { - let hook = crate::factory::supported_cmf_hook_name(hook)?; + let hook = crate::factory::supported_hook_name(hook)?; Some(( hook, Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) as Arc, @@ -541,7 +537,10 @@ async fn runtime_config_loads_generic_cmf_factory_plugin() { })); let mut runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); runtime - .register_factory("generic", Box::new(CmfPluginFactory::new(TestPlugin::rewrite_from_config))) + .register_factory( + "generic", + Box::new(GatewayPluginFactory::new(TestPlugin::rewrite_from_config).with_cmf_hooks()), + ) .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); @@ -564,7 +563,10 @@ async fn generic_cmf_factory_registers_prompt_only_plugin() { })); let mut runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); runtime - .register_factory("generic", Box::new(CmfPluginFactory::new(TestPlugin::rewrite_from_config))) + .register_factory( + "generic", + Box::new(GatewayPluginFactory::new(TestPlugin::rewrite_from_config).with_cmf_hooks()), + ) .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); @@ -591,7 +593,10 @@ async fn generic_cmf_factory_registers_mixed_tool_and_prompt_plugin() { })); let mut runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); runtime - .register_factory("generic", Box::new(CmfPluginFactory::new(TestPlugin::rewrite_from_config))) + .register_factory( + "generic", + Box::new(GatewayPluginFactory::new(TestPlugin::rewrite_from_config).with_cmf_hooks()), + ) .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); @@ -930,10 +935,11 @@ fn request_id(payload: &MessagePayload) -> Option { #[tokio::test] async fn hook_combinations_preserve_correlation_for_each_operation() { - use crate::cmf::Operation; use rmcp::model::{GetPromptResult, PromptMessage, Role}; - for operation in Operation::ALL { + use crate::hooks::Operation; + + for operation in Operation::MCP { for (pre_enabled, post_enabled) in [(false, false), (true, false), (false, true), (true, true)] { let [pre, post] = operation.hooks(); let hooks = [(pre, pre_enabled), (post, post_enabled)] @@ -946,6 +952,7 @@ async fn hook_combinations_preserve_correlation_for_each_operation() { let handle = runtime.handle(); match operation { + Operation::Http => unreachable!("HTTP lifecycle is tested separately"), Operation::Tool => { let pre = handle .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) @@ -1073,6 +1080,7 @@ async fn dropping_the_last_in_flight_state_releases_the_replaced_runtime() { fn test_request_context() -> crate::PluginRequestContext { crate::PluginRequestContext { + request: None, tool: Some(ToolPolicyContext { id: "test".to_owned(), name: "sum".to_owned(), diff --git a/crates/contextforge-data-plane-cpex/src/request.rs b/crates/contextforge-data-plane-cpex/src/request.rs new file mode 100644 index 00000000..357371a7 --- /dev/null +++ b/crates/contextforge-data-plane-cpex/src/request.rs @@ -0,0 +1,64 @@ +//! One policy snapshot and CPEX state table for the whole HTTP/MCP request. +use std::{collections::HashMap, sync::Arc}; + +use cpex::cpex_core::{context::PluginContextTable, extensions::Extensions, registry::HookEntry}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::registry::RuntimePolicies; + +#[derive(Clone)] +pub struct PluginRequest { + pub(crate) policies: Arc, + pub(crate) state: Arc>, +} + +impl PluginRequest { + pub(crate) fn new(policies: Arc, extensions: Extensions) -> Self { + Self { policies, state: Arc::new(Mutex::new(RequestState { extensions, ..Default::default() })) } + } + + pub async fn set_verified_subject(&self, subject: cpex::cpex_core::extensions::SubjectExtension) { + let mut state = self.state.lock().await; + let security = state + .extensions + .security + .get_or_insert_with(|| Arc::new(cpex::cpex_core::extensions::SecurityExtension::default())); + Arc::make_mut(security).subject = Some(subject); + } + + pub async fn sync_request_headers(&self, headers: HashMap) { + let mut state = self.state.lock().await; + if let Some(http) = &mut state.extensions.http { + Arc::make_mut(http).request_headers = headers; + } + } + + pub async fn extensions(&self) -> Extensions { + self.state.lock().await.extensions.clone() + } +} + +#[derive(Default)] +pub(crate) struct RequestState { + pub(crate) extensions: Extensions, + pub(crate) contexts: PluginContextTable, + // Global and scoped managers assign different IDs to the same configured plugin. + // Move its CPEX local state when crossing that boundary; never share by name alone. + plugin_ids: HashMap<(String, String), Uuid>, +} + +impl RequestState { + pub(crate) fn prepare(&mut self, entries: &[HookEntry]) { + for entry in entries { + let config = entry.plugin_ref.trusted_config(); + let id = entry.plugin_ref.id(); + if let Some(previous) = self.plugin_ids.insert((config.kind.clone(), config.name.clone()), id) + && previous != id + && let Some(local) = self.contexts.local_states.remove(&previous) + { + self.contexts.local_states.insert(id, local); + } + } + } +} diff --git a/crates/contextforge-data-plane-cpex/src/resources/mod.rs b/crates/contextforge-data-plane-cpex/src/resources/mod.rs index e0bee2b8..530dd5d8 100644 --- a/crates/contextforge-data-plane-cpex/src/resources/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/resources/mod.rs @@ -9,7 +9,8 @@ use rmcp::{ use crate::{ GatewayPluginRuntimeHandle, PluginRequestContext, - cmf::{CmfResponse, Operation, message_payload}, + cmf::{CmfResponse, message_payload}, + hooks::Operation, runtime::CallState, }; @@ -116,7 +117,7 @@ impl ResourceHookState { pub async fn after_read_resource(self, response: ReadResourceResult) -> Result { match self.call { - Some(mut call) => call.after(response).await, + Some(call) => call.after(response).await, None => Ok(response), } } @@ -128,12 +129,11 @@ impl GatewayPluginRuntimeHandle { resource_uri: &str, context: PluginRequestContext, ) -> Result { - let (rewritten_uri, call) = self - .current()? - .global + let (request_state, runtime) = self.resolve(Operation::Resource, &context)?; + let (rewritten_uri, call) = runtime .before( (Operation::Resource, resource_uri), - context.extensions, + (request_state, context.extensions), |id| resource_request_payload(resource_uri, id), |payload, _| { let [ContentPart::ResourceRef { content }] = payload.message.content.as_slice() else { diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index f0c7b09e..8667847a 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -8,10 +8,9 @@ use contextforge_data_plane_apis::runtime_plugin_config::RuntimePluginSettings; use cpex::cpex_core::{ cmf::{CmfHook, MessagePayload}, config::CpexConfig, - context::PluginContextTable, executor::PipelineResult, factory::PluginFactoryRegistry, - hooks::payload::Extensions, + hooks::{HookTypeDef, payload::Extensions}, manager::PluginManager, plugin::{MatchContext, PluginMode}, registry::HookEntry, @@ -20,31 +19,33 @@ use rmcp::ErrorData; use tracing::instrument; use crate::{ - cmf::{CmfResponse, Operation, modified_message_payload, plugin_denied_error}, + PluginRequest, + cmf::{CmfResponse, modified_message_payload, plugin_denied_error}, error::GatewayPluginRuntimeError, - factory::supported_cmf_hook_name, + extension_guard::HeaderWritePolicy, + factory::supported_hook_name, + hooks::Operation, }; #[derive(Default)] struct HookPair { pre: Vec, post: Vec, + payload_writes: [bool; 2], } #[derive(Default)] pub(crate) struct GatewayPluginRuntime { manager: PluginManager, - hooks: [HookPair; 3], - payload_writes: [[bool; 2]; 3], + hooks: [HookPair; 4], auth_headers_writable: bool, } -/// Pins the selected runtime and correlation context until the request finishes. -/// Only tool calls wrap this in a mutex, because events also update their context. +/// Pins operation-specific payload metadata. Mutable CPEX state belongs to the request. pub(crate) struct CallState { runtime: Arc, - context_table: PluginContextTable, - extensions: Extensions, + request: PluginRequest, + target: Extensions, post: Vec, name: String, id: String, @@ -66,28 +67,29 @@ impl GatewayPluginRuntime { for plugin in &mut config.plugins { for hook in &mut plugin.hooks { // Python publishes native hook names; registered Rust handlers use CMF. - if !hook.starts_with("cmf.") { + if !hook.starts_with("cmf.") && !Operation::Http.hooks().contains(&hook.as_str()) { *hook = format!("cmf.{hook}"); } } } let mut runtime = Self::from_config(config, factories).await?; runtime.auth_headers_writable = settings.plugins_can_override_auth_headers; - runtime.payload_writes = Operation::ALL.map(|operation| { - operation.hooks().map(|hook| { + for (hooks, operation) in runtime.hooks.iter_mut().zip(Operation::ALL) { + hooks.payload_writes = operation.hooks().map(|hook| { let name = hook.strip_prefix("cmf.").unwrap_or(hook); let field = match name { "tool_pre_invoke" | "prompt_pre_fetch" => "args", "tool_post_invoke" | "prompt_post_fetch" => "result", "resource_pre_fetch" => "uri", + "http_pre_request" | "http_post_request" => "headers", _ => "content", }; settings .hook_policies .get(name) .map_or(settings.default_hook_policy == "allow", |policy| policy.writable_fields.contains(field)) - }) - }); + }); + } Ok(runtime) } @@ -108,19 +110,23 @@ impl GatewayPluginRuntime { let [pre, post] = operation .hooks() .map(|name| entries.iter().filter(|(hook, _)| hook == name).map(|(_, entry)| entry.clone()).collect()); - HookPair { pre, post } + HookPair { pre, post, payload_writes: [true; 2] } }); - Ok(Self { manager, hooks, payload_writes: [[true; 2]; 3], auth_headers_writable: false }) + Ok(Self { manager, hooks, auth_headers_writable: false }) } - #[instrument(name = "cmf_plugin_before", level = "info", skip(self, target, payload, update, extensions))] + #[instrument(name = "cmf_plugin_before", level = "info", skip_all)] pub(crate) async fn before( self: &Arc, target: (Operation, &str), - mut extensions: Extensions, + context: (PluginRequest, Extensions), payload: impl FnOnce(&str) -> MessagePayload, update: impl FnOnce(&MessagePayload, &str) -> Result, ) -> Result<(U, Option), ErrorData> { + let (request, target_extensions) = context; + let mut extensions = request.extensions().await; + extensions.meta.clone_from(&target_extensions.meta); + extensions.mcp.clone_from(&target_extensions.mcp); let (operation, name) = target; let hooks = &self.hooks[operation as usize]; let pre = matching_entries(&hooks.pre, &extensions); @@ -130,26 +136,22 @@ impl GatewayPluginRuntime { } let id = format!("{}-{}", operation.id_prefix(), CORRELATION_ID.fetch_add(1, Ordering::Relaxed)); - let (update, context_table) = if pre.is_empty() { - (U::default(), PluginContextTable::default()) + let update = if pre.is_empty() { + U::default() } else { - let result = self.invoke((operation.hooks()[0], &pre), payload(&id), extensions.clone(), None).await; + let result = self.invoke::((operation, 0, &pre), payload(&id), &request, &target_extensions).await; if result.is_denied() { return Err(plugin_denied_error(operation.subject(), result)); } - let update = match modified_message_payload(&result) { + match modified_message_payload(&result) { Some(payload) => update(payload, &id)?, None => U::default(), - }; - if let Some(modified) = result.modified_extensions { - extensions = modified; } - (update, result.context_table) }; let state = (!post.is_empty()).then(|| CallState { runtime: Arc::clone(self), - context_table, - extensions, + request, + target: target_extensions, post, name: name.to_owned(), id, @@ -157,26 +159,43 @@ impl GatewayPluginRuntime { Ok((update, state)) } - #[instrument(name = "cmf_plugin_invoke", level = "info", skip_all)] - async fn invoke( + pub(crate) fn entries(&self, operation: Operation, phase: usize, extensions: &Extensions) -> Vec { + let hooks = &self.hooks[operation as usize]; + matching_entries(if phase == 0 { &hooks.pre } else { &hooks.post }, extensions) + } + + #[instrument(name = "gateway_plugin_invoke", level = "info", skip_all, fields(hook = invocation.0.hooks()[invocation.1]))] + pub(crate) async fn invoke( &self, - invocation: (&'static str, &[HookEntry]), - payload: MessagePayload, - extensions: Extensions, - context_table: Option, + invocation: (Operation, usize, &[HookEntry]), + payload: H::Payload, + request: &PluginRequest, + target: &Extensions, ) -> PipelineResult { - let (hook, entries) = invocation; - let payload_writable = Operation::ALL - .iter() - .enumerate() - .find_map(|(i, operation)| { - operation.hooks().iter().position(|name| *name == hook).map(|j| self.payload_writes[i][j]) - }) - .unwrap_or(false); - let entries = - crate::extension_guard::guarded_entries(entries, &extensions, payload_writable, self.auth_headers_writable); - let (result, background_tasks) = - self.manager.invoke_entries::(&entries, payload, extensions, context_table).await; + let (operation, phase, entries) = invocation; + let hook = operation.hooks()[phase]; + // Only this request is serialized. No registry or gateway lock crosses plugin I/O. + let mut state = request.state.lock().await; + state.extensions.meta.clone_from(&target.meta); + state.extensions.mcp.clone_from(&target.mcp); + state.prepare(entries); + let writable = self.hooks[operation as usize].payload_writes[phase]; + let http = matches!(operation, Operation::Http); + let header_policy = if http && !writable { + HeaderWritePolicy::ReadOnly + } else if self.auth_headers_writable { + HeaderWritePolicy::All + } else if http && phase == 0 { + HeaderWritePolicy::AllowNewCredentials + } else { + HeaderWritePolicy::PreserveCredentials + }; + let guarded = + crate::extension_guard::guarded_entries(entries, &state.extensions, !http && writable, header_policy); + let (result, background_tasks) = self + .manager + .invoke_entries::(&guarded, payload, state.extensions.clone(), Some(state.contexts.clone())) + .await; for error in &result.errors { tracing::warn!( hook, @@ -186,13 +205,19 @@ impl GatewayPluginRuntime { "CPEX plugin soft error" ); } + if !result.is_denied() { + state.contexts = result.context_table.clone(); + if let Some(extensions) = &result.modified_extensions { + state.extensions = extensions.clone(); + } + } drop(background_tasks); result } } impl CallState { - pub(crate) async fn after(&mut self, response: T) -> Result { + pub(crate) async fn after(&self, response: T) -> Result { let result = self.invoke(&response).await?; if result.is_denied() { return Err(plugin_denied_error(T::OPERATION.subject(), result)); @@ -200,23 +225,10 @@ impl CallState { self.apply(response, &result) } - pub(crate) async fn invoke(&mut self, response: &T) -> Result { + pub(crate) async fn invoke(&self, response: &T) -> Result { let payload = response.to_payload(&self.name, &self.id)?; - let result = self - .runtime - .invoke( - (T::OPERATION.hooks()[1], &self.post), - payload, - self.extensions.clone(), - Some(self.context_table.clone()), - ) - .await; - if !result.is_denied() { - self.context_table = result.context_table.clone(); - if let Some(extensions) = &result.modified_extensions { - self.extensions = extensions.clone(); - } - } + let result = + self.runtime.invoke::((T::OPERATION, 1, &self.post), payload, &self.request, &self.target).await; Ok(result) } @@ -289,7 +301,7 @@ fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayP } for plugin in &config.plugins { - if plugin.hooks.iter().any(|hook| supported_cmf_hook_name(hook).is_none()) { + if plugin.hooks.iter().any(|hook| supported_hook_name(hook).is_none()) { return Err(GatewayPluginRuntimeError::ConfigUnsupported); } } diff --git a/crates/contextforge-data-plane-cpex/src/tools/mod.rs b/crates/contextforge-data-plane-cpex/src/tools/mod.rs index e85a75b4..1d4da449 100644 --- a/crates/contextforge-data-plane-cpex/src/tools/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/tools/mod.rs @@ -7,11 +7,11 @@ use rmcp::{ }; use serde::{Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; -use tokio::sync::Mutex; use crate::{ ArgumentsUpdate, GatewayPluginRuntimeHandle, PluginRequestContext, PreHookResult, - cmf::{CmfResponse, Operation, message_payload}, + cmf::{CmfResponse, message_payload}, + hooks::Operation, runtime::CallState, }; @@ -99,7 +99,7 @@ fn raw_tool_result(value: Value, is_error: bool) -> CallToolResult { /// Shared by the final tool response and its progress notifications. #[derive(Clone)] -pub struct ToolHookState(Arc>); +pub struct ToolHookState(Arc); impl std::fmt::Debug for ToolHookState { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -115,12 +115,11 @@ impl GatewayPluginRuntimeHandle { backend_name: &str, context: PluginRequestContext, ) -> Result, ErrorData> { - let (arguments, state) = self - .current()? - .tool(context.tool.as_ref())? + let (request_state, runtime) = self.resolve(Operation::Tool, &context)?; + let (arguments, state) = runtime .before( (Operation::Tool, tool_name), - context.extensions, + (request_state, context.extensions), |id| tool_call_payload(request, tool_name, backend_name, id), |payload, _| { let arguments = tool_call_arguments(payload).ok_or_else(|| { @@ -130,13 +129,13 @@ impl GatewayPluginRuntimeHandle { }, ) .await?; - Ok(PreHookResult { arguments, state: state.map(|state| ToolHookState(Arc::new(Mutex::new(state)))) }) + Ok(PreHookResult { arguments, state: state.map(|state| ToolHookState(Arc::new(state))) }) } } impl ToolHookState { pub async fn after_tool_call(self, response: CallToolResult) -> Result { - self.0.lock().await.after(response).await + self.0.after(response).await } /// Returns `None` when a plugin denies a progress or logging notification. @@ -144,7 +143,7 @@ impl ToolHookState { where T: Serialize + DeserializeOwned, { - let mut state = self.0.lock().await; + let state = &self.0; let event = ToolEvent(event); let result = state.invoke(&event).await?; if result.is_denied() { diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index f13fe28f..9817c8eb 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -23,7 +23,7 @@ tracing-opentelemetry.workspace = true opentelemetry.workspace = true tokio.workspace = true tokio-util = "0.7" -axum.workspace = true +axum = { workspace = true, features = ["original-uri"] } axum-otel-metrics = "0.14" tower-http = { version = "0.7.1", features = ["cors", "trace"] } tower = "0.5.3" diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs index fe7259d6..aa24cc1d 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs @@ -1,24 +1,15 @@ -use std::{collections::HashSet, sync::Arc}; +use std::sync::Arc; -use contextforge_data_plane_apis::{ - User, - user_store::{BackendMCPGateway, UserConfig}, -}; -use contextforge_data_plane_cpex::PluginRequestContext; +use contextforge_data_plane_apis::user_store::BackendMCPGateway; +use contextforge_data_plane_cpex::{PluginRequest, PluginRequestContext}; use cpex::cpex_core::extensions::{ - Extensions, HttpExtension, MCPExtension, MetaExtension, PromptMetadata, RequestExtension, ResourceMetadata, - SecurityExtension, SubjectExtension, SubjectType, ToolMetadata, + Extensions, MCPExtension, MetaExtension, PromptMetadata, ResourceMetadata, ToolMetadata, }; use http::request::Parts; -use opentelemetry::trace::TraceContextExt; use rmcp::{ErrorData, RoleServer, service::RequestContext}; use serde_json::Value; -use tracing_opentelemetry::OpenTelemetrySpanExt; -use crate::{ - authorization::{AuthorizationClaims, AuthorizedPrincipal}, - layers::virtual_host_id::VirtualHostId, -}; +use crate::{authorization::AuthorizedPrincipal, layers::virtual_host_id::VirtualHostId}; /// Construct plugin input only after authorization and route resolution. Nothing /// in MCP arguments or client metadata participates in identity or policy lookup. @@ -37,32 +28,13 @@ pub(super) fn request_context( .extensions .get::() .ok_or_else(|| ErrorData::internal_error("Plugin verified principal is missing", None))?; - let claims = parts - .extensions - .get::() - .ok_or_else(|| ErrorData::internal_error("Plugin verified claims are missing", None))?; - let claims = Value::from(claims); - let user_config = parts + let request = parts .extensions - .get::() - .ok_or_else(|| ErrorData::internal_error("Plugin user configuration is missing", None))?; + .get::() + .ok_or_else(|| ErrorData::internal_error("Plugin request state is missing", None))?; let server_id = parts.extensions.get::().map(|id| id.value().clone()); let tool = (kind == "tool").then(|| backend.tool_policy_contexts.get(name)).flatten().cloned(); let canonical_name = tool.as_ref().map_or(name, |tool| tool.name.as_str()); - let user = User::from(principal); - let subject = SubjectExtension { - id: Some(user_config.user_email.as_deref().unwrap_or(user.key()).to_owned()), - subject_type: Some(SubjectType::User), - roles: string_set(claims.get("roles")), - teams: string_set(claims.get("teams")), - permissions: string_set(claims.get("scopes").and_then(|scopes| scopes.get("permissions"))), - claims: claims - .as_object() - .into_iter() - .flatten() - .map(|(key, value)| (key.clone(), value.as_str().map_or_else(|| value.to_string(), str::to_owned))) - .collect(), - }; let tenant = tool.as_ref().and_then(|tool| tool.team_id.as_deref()).unwrap_or(principal.tenant_id()); let mut meta = MetaExtension { entity_type: Some(kind.to_owned()), @@ -98,38 +70,9 @@ pub(super) fn request_context( ..Default::default() }, }; - let trace = tracing::Span::current().context(); - let span = trace.span(); - let span_context = span.span_context(); Ok(PluginRequestContext { + request: Some(request.clone()), tool, - extensions: Extensions { - request: Some(Arc::new(RequestExtension { - request_id: Some(uuid::Uuid::new_v4().to_string()), - trace_id: span_context.is_valid().then(|| span_context.trace_id().to_string()), - span_id: span_context.is_valid().then(|| span_context.span_id().to_string()), - ..Default::default() - })), - http: Some(Arc::new(HttpExtension { - request_headers: parts - .headers - .iter() - .filter_map(|(key, value)| value.to_str().ok().map(|value| (key.to_string(), value.to_owned()))) - .collect(), - method: Some(parts.method.to_string()), - path: Some(parts.uri.path().to_owned()), - host: parts.headers.get(http::header::HOST).and_then(|value| value.to_str().ok()).map(str::to_owned), - scheme: parts.uri.scheme_str().map(str::to_owned), - ..Default::default() - })), - security: Some(Arc::new(SecurityExtension { subject: Some(subject), ..Default::default() })), - meta: Some(Arc::new(meta)), - mcp: Some(Arc::new(mcp)), - ..Default::default() - }, + extensions: Extensions { meta: Some(Arc::new(meta)), mcp: Some(Arc::new(mcp)), ..Default::default() }, }) } - -fn string_set(value: Option<&Value>) -> HashSet { - value.and_then(Value::as_array).into_iter().flatten().filter_map(Value::as_str).map(str::to_owned).collect() -} diff --git a/crates/contextforge-data-plane-lib/src/layers/http_plugins.rs b/crates/contextforge-data-plane-lib/src/layers/http_plugins.rs new file mode 100644 index 00000000..d4482280 --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/layers/http_plugins.rs @@ -0,0 +1,121 @@ +use std::{collections::HashMap, net::SocketAddr, sync::Arc}; + +use axum::{ + extract::{ConnectInfo, OriginalUri, Request, State}, + middleware::Next, + response::Response, +}; +use contextforge_data_plane_cpex::{GatewayPluginRuntimeHandle, HttpHookPayload}; +use cpex::cpex_core::extensions::{Extensions, HttpExtension, RequestExtension, SecurityExtension}; +use http::{HeaderMap, HeaderName, HeaderValue, StatusCode}; +use opentelemetry::trace::TraceContextExt; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use crate::errors::custom_error; + +pub(crate) async fn http_plugin_layer( + State(runtime): State, + mut request: Request, + next: Next, +) -> Response { + let trace = tracing::Span::current().context(); + let span = trace.span(); + let context = span.span_context(); + let uri = request.extensions().get::().map_or(request.uri(), |uri| &uri.0); + let payload = HttpHookPayload { + method: request.method().to_string(), + path: uri.path().to_owned(), + client_addr: request.extensions().get::>().map(|info| info.0), + status_code: None, + }; + let extensions = Extensions { + request: Some(Arc::new(RequestExtension { + request_id: Some(uuid::Uuid::new_v4().to_string()), + trace_id: context.is_valid().then(|| context.trace_id().to_string()), + span_id: context.is_valid().then(|| context.span_id().to_string()), + ..Default::default() + })), + http: Some(Arc::new(HttpExtension { + method: Some(payload.method.clone()), + path: Some(payload.path.clone()), + host: request.headers().get(http::header::HOST).and_then(|value| value.to_str().ok()).map(str::to_owned), + scheme: uri.scheme_str().map(str::to_owned), + request_headers: header_values(request.headers()), + ..Default::default() + })), + security: Some(Arc::new(SecurityExtension::default())), + ..Default::default() + }; + let state = match runtime.before_http_request(payload, extensions).await { + Ok(state) => state, + Err(error) => { + tracing::warn!(%error, "HTTP plugin configuration is unavailable"); + return custom_error(StatusCode::INTERNAL_SERVER_ERROR, "Runtime plugin configuration is unavailable"); + }, + }; + let updated = state.request.extensions().await; + if let Some(http) = updated.http { + apply_headers(request.headers_mut(), &http.request_headers); + } + state.request.sync_request_headers(header_values(request.headers())).await; + request.extensions_mut().insert(state.request.clone()); + let mut response = next.run(request).await; + // Headers are still mutable here. Do not collect or wrap the response body: + // streaming MCP operations may continue after this hook has run. + let updated = state.after_http_request(response.status().as_u16(), header_values(response.headers())).await; + if let Some(http) = updated.http { + apply_headers(response.headers_mut(), &http.response_headers); + } + response +} + +fn header_values(headers: &HeaderMap) -> HashMap { + headers + .iter() + .filter_map(|(name, value)| value.to_str().ok().map(|value| (name.to_string(), value.to_owned()))) + .collect() +} + +fn apply_headers(headers: &mut HeaderMap, updates: &HashMap) { + // Validate the entire update first. Preserve repeated and non-UTF8 headers + // unless a plugin actually changes their value; the built-in hooks merge edits. + let parsed = updates + .iter() + .map(|(name, value)| { + Some((HeaderName::try_from(name.as_str()).ok()?, HeaderValue::try_from(value.as_str()).ok()?)) + }) + .collect::>>(); + let Some(parsed) = parsed else { + tracing::warn!("HTTP plugin returned invalid headers; ignoring header changes"); + return; + }; + for (name, value) in parsed { + if headers.get_all(&name).iter().next_back() != Some(&value) { + headers.insert(name, value); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn header_updates_preserve_repeated_and_binary_headers_and_validate_atomically() { + let mut headers = HeaderMap::new(); + headers.append("set-cookie", HeaderValue::from_static("first=1")); + headers.append("set-cookie", HeaderValue::from_static("second=2")); + headers.insert("x-binary", HeaderValue::from_bytes(&[0xff]).expect("opaque header")); + let mut updates = header_values(&headers); + updates.insert("x-plugin".to_owned(), "added".to_owned()); + apply_headers(&mut headers, &updates); + assert_eq!(headers.get_all("set-cookie").iter().count(), 2); + assert_eq!(headers["x-binary"].as_bytes(), &[0xff]); + assert_eq!(headers["x-plugin"], "added"); + let before = headers.clone(); + updates.insert("x-plugin".to_owned(), "changed".to_owned()); + updates.insert("x-invalid".to_owned(), "invalid\r\nvalue".to_owned()); + apply_headers(&mut headers, &updates); + assert_eq!(headers, before); + } +} diff --git a/crates/contextforge-data-plane-lib/src/layers/mod.rs b/crates/contextforge-data-plane-lib/src/layers/mod.rs index 9b96028c..a7a02c42 100644 --- a/crates/contextforge-data-plane-lib/src/layers/mod.rs +++ b/crates/contextforge-data-plane-lib/src/layers/mod.rs @@ -1,4 +1,5 @@ pub mod claims_id; +pub(crate) mod http_plugins; pub mod mcp_header_limits; pub mod mcp_origin; pub mod principal_extractor; diff --git a/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs b/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs index abab56d0..56484e9e 100644 --- a/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs +++ b/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs @@ -1,5 +1,9 @@ +use std::collections::HashSet; + use axum::{extract::State, middleware::Next, response::Response}; use contextforge_data_plane_apis::User; +use cpex::cpex_core::extensions::{SubjectExtension, SubjectType}; +use serde_json::Value; use tracing::{debug, info, warn}; @@ -23,7 +27,14 @@ pub async fn user_config_store_layer( "user_config_store_layer - getting user config for principal {principal:?} method = {method} path = {path}" ); let user = User::from(principal); - match state.config_store.get_config(&user).await { + let config = state.config_store.get_config(&user).await; + if let Some(plugin_request) = request.extensions().get::() + && let Some(claims) = request.extensions().get::() + { + let email = config.as_ref().ok().and_then(|config| config.user_email.as_deref()); + plugin_request.set_verified_subject(verified_subject(principal, claims, email)).await; + } + match config { Ok(user_config) => { let virtual_hosts = user_config.virtual_hosts.len(); info!( @@ -52,3 +63,29 @@ pub async fn user_config_store_layer( bad_request("No claims in the token") } } + +fn verified_subject( + principal: &crate::authorization::AuthorizedPrincipal, + claims: &crate::AuthorizationClaims, + email: Option<&str>, +) -> SubjectExtension { + let claims = Value::from(claims); + let user = User::from(principal); + SubjectExtension { + id: Some(email.unwrap_or(user.key()).to_owned()), + subject_type: Some(SubjectType::User), + roles: string_set(claims.get("roles")), + teams: string_set(claims.get("teams")), + permissions: string_set(claims.get("scopes").and_then(|scopes| scopes.get("permissions"))), + claims: claims + .as_object() + .into_iter() + .flatten() + .map(|(key, value)| (key.clone(), value.as_str().map_or_else(|| value.to_string(), str::to_owned))) + .collect(), + } +} + +fn string_set(value: Option<&Value>) -> HashSet { + value.and_then(Value::as_array).into_iter().flatten().filter_map(Value::as_str).map(str::to_owned).collect() +} diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index a36fcd59..13e05c82 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -125,11 +125,12 @@ impl Gateway { let reqwest_backend_client = reqwest::Client::try_from(&config)?; // Create streamable HTTP service + let mcp_plugin_runtime = plugin_runtime.clone(); let mcp_service: StreamableHttpService = StreamableHttpService::new( move || { Ok(McpService::builder() .with_http_client(reqwest_backend_client.clone()) - .with_plugin_runtime(plugin_runtime.clone()) + .with_plugin_runtime(mcp_plugin_runtime.clone()) .build()) }, session_manager, @@ -158,8 +159,14 @@ impl Gateway { app.layer(layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {})) }; + let app = app.layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), claims_layer)); + let app = if let Some(runtime) = plugin_runtime { + app.layer(middleware::from_fn_with_state(mcp_standard_header_limits.clone(), mcp_header_limits_layer)) + .layer(middleware::from_fn_with_state(runtime, layers::http_plugins::http_plugin_layer)) + } else { + app + }; let app = app - .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), claims_layer)) .layer(middleware::from_fn(virtual_host_id_layer)) .layer(middleware::from_fn_with_state(mcp_standard_header_limits, mcp_header_limits_layer)) .layer(cors_layer) diff --git a/crates/contextforge-data-plane-lib/src/transports/tcp.rs b/crates/contextforge-data-plane-lib/src/transports/tcp.rs index 84620e35..7593c532 100644 --- a/crates/contextforge-data-plane-lib/src/transports/tcp.rs +++ b/crates/contextforge-data-plane-lib/src/transports/tcp.rs @@ -19,7 +19,7 @@ impl Tcp { info!("Starting TCP listener at {}", self.address); let tcp_listener: TcpListener = self.try_into()?; - Ok(axum::serve(tcp_listener, service) + Ok(axum::serve(tcp_listener, service.into_make_service_with_connect_info::()) .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.ok(); info!("Shutting down..."); diff --git a/crates/contextforge-data-plane-lib/src/transports/tls.rs b/crates/contextforge-data-plane-lib/src/transports/tls.rs index 75b99c42..1f4f65c2 100644 --- a/crates/contextforge-data-plane-lib/src/transports/tls.rs +++ b/crates/contextforge-data-plane-lib/src/transports/tls.rs @@ -62,7 +62,8 @@ impl DownstreamTls { let stream = TokioIo::new(stream); - let hyper_service = hyper::service::service_fn(move |request: Request| { + let hyper_service = hyper::service::service_fn(move |mut request: Request| { + request.extensions_mut().insert(axum::extract::ConnectInfo(addr)); tower_service.clone().call(request) }); diff --git a/crates/contextforge-data-plane-lib/tests/gateway.rs b/crates/contextforge-data-plane-lib/tests/gateway.rs index 92aa36bd..95c9138d 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway.rs @@ -7,6 +7,8 @@ mod compatibility; mod completions; #[path = "gateway/future_contracts/mod.rs"] mod future_contracts; +#[path = "gateway/http_plugins.rs"] +mod http_plugins; #[path = "gateway/plugin_context.rs"] mod plugin_context; #[path = "gateway/plugins.rs"] diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs index c6e84cf4..3bbe60e0 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs @@ -226,6 +226,12 @@ impl ServerHandler for TestBackend { Ok(CallToolResult::success(vec![ContentBlock::text(text.to_owned())])) }, "wait_for_cancellation" => { + if let Some(token) = cx.meta.get_progress_token() { + cx.peer + .notify_progress(ProgressNotificationParam::new(token.clone(), 0.0)) + .await + .map_err(|_| ErrorData::internal_error("progress notification failed", None))?; + } cx.ct.cancelled().await; self.state .cancellations diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/runtime.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/runtime.rs index fc4c6c0b..806c4da2 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/runtime.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/runtime.rs @@ -1,18 +1,19 @@ use std::sync::Arc; -use contextforge_data_plane_cpex::CpexRuntimeRegistry; +use contextforge_data_plane_cpex::{CpexRuntimeRegistry, GatewayPluginFactory}; use cpex::cpex_core::config::CpexConfig; use serde_json::json; -use contextforge_data_plane_cpex::CmfPluginFactory; - use super::{PromptTestPlugin, TestPlugin, TestPluginFactory}; pub(crate) async fn runtime_with_prompt_plugin(plugin: Arc) -> Arc { let mut runtime = CpexRuntimeRegistry::default(); let template = Arc::clone(&plugin); runtime - .register_factory("prompt-test", Box::new(CmfPluginFactory::new(move |config| template.rebuild(config)))) + .register_factory( + "prompt-test", + Box::new(GatewayPluginFactory::new(move |config| template.rebuild(config)).with_cmf_hooks()), + ) .expect("prompt test factory registers"); let config = serde_json::from_value(json!({ "plugins": [{ diff --git a/crates/contextforge-data-plane-lib/tests/gateway/http_plugins.rs b/crates/contextforge-data-plane-lib/tests/gateway/http_plugins.rs new file mode 100644 index 00000000..60ebf039 --- /dev/null +++ b/crates/contextforge-data-plane-lib/tests/gateway/http_plugins.rs @@ -0,0 +1,533 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use async_trait::async_trait; +use contextforge_data_plane_apis::{runtime_plugin_config::RuntimePluginConfigDocument, user_store::ToolPolicyContext}; +use contextforge_data_plane_cpex::{ + CpexRuntimeRegistry, GatewayPluginFactory, HttpHook, HttpHookPayload, PluginRequestContext, +}; +use cpex::cpex_core::{ + cmf::{CmfHook, MessagePayload}, + context::PluginContext, + error::PluginViolation, + extensions::{Extensions, HttpExtension, RequestExtension, SecurityExtension, SubjectExtension}, + hooks::{HookHandler, PluginResult}, + plugin::{Plugin, PluginConfig}, +}; +use rmcp::model::{CallToolRequestParams, CallToolResult}; +use serde_json::{Value, json}; + +use crate::harness::{MemoryUserConfigStore, TEST_USER_ID, TestServer, create_default_config, token}; +use contextforge_data_plane_lib::{ + AuthorizationClaims, AuthorizationService, Gateway, UserConfigStore, UserConfigStoreType, +}; + +#[derive(Debug)] +struct VerifyJwt; +#[async_trait] +impl AuthorizationService for VerifyJwt { + async fn authorize(&self, header: &http::HeaderValue) -> Option { + let token = header.to_str().ok()?.strip_prefix("Bearer ")?; + let key = jsonwebtoken::DecodingKey::from_rsa_pem(&std::fs::read("../../assets/jwt.key.pub").ok()?).ok()?; + let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256); + validation.set_audience(&["mcpgateway-api"]); + validation.set_issuer(&["mcpgateway"]); + jsonwebtoken::decode::(token, &key, &validation).ok().map(|data| AuthorizationClaims::from(data.claims)) + } +} + +async fn start_verified_gateway(runtime: Arc) -> TestServer { + let store = MemoryUserConfigStore::default(); + store + .set_config( + &contextforge_data_plane_apis::User::new(TEST_USER_ID), + &serde_json::from_value(json!({ + "virtual_hosts":{"test":{"backends":{}}} + })) + .expect("user config"), + ) + .await + .expect("store config"); + let router = Gateway::builder() + .with_config(create_default_config()) + .with_session_manager(Arc::new( + rmcp::transport::streamable_http_server::session::local::LocalSessionManager::default(), + )) + .with_user_config_store_type(UserConfigStoreType::Test(Arc::new(store))) + .with_plugin_runtime(Some(runtime.handle())) + .with_authorization_service(Arc::new(VerifyJwt)) + .build() + .into_router() + .await + .expect("router"); + TestServer::start_http(router).await.expect("server") +} + +type Seen = Arc>>; +struct Probe { + config: PluginConfig, + seen: Seen, +} + +#[async_trait] +impl Plugin for Probe { + fn config(&self) -> &PluginConfig { + &self.config + } +} + +impl Probe { + fn observe(&self, stage: &str, extensions: &Extensions, ctx: &mut PluginContext) { + let count = ctx.get_local("count").and_then(Value::as_u64).unwrap_or(0); + self.seen.lock().expect("observations").push(json!({ + "plugin": self.config.name, + "stage": stage, "count": count, "global": ctx.get_global("count"), + "subject": extensions.security.as_ref().and_then(|security| security.subject.as_ref()).and_then(|subject| subject.id.as_deref()), + "request": extensions.request.as_ref().and_then(|request| request.request_id.as_deref()), + "headers_visible": extensions.http.is_some(), + "request_content_type": extensions.http.as_ref().and_then(|http| http.get_request_header("content-type")), + "response_content_type": extensions.http.as_ref().and_then(|http| http.get_response_header("content-type")), + })); + ctx.set_local("count", json!(count + 1)); + ctx.set_global("count", json!(count + 1)); + } +} + +#[allow(clippy::unused_async_trait_impl)] +impl HookHandler for Probe { + async fn handle( + &self, + _: &MessagePayload, + extensions: &Extensions, + ctx: &mut PluginContext, + ) -> PluginResult { + self.observe("mcp", extensions, ctx); + PluginResult::allow() + } +} + +impl HookHandler for Probe { + async fn handle( + &self, + payload: &HttpHookPayload, + extensions: &Extensions, + ctx: &mut PluginContext, + ) -> PluginResult { + let post = payload.status_code.is_some(); + self.observe(if post { "http_post" } else { "http_pre" }, extensions, ctx); + let config = self.config.config.clone().unwrap_or_default(); + if config.get("timeout") == Some(&json!(true)) { + tokio::time::sleep(Duration::from_secs(60)).await; + } + if config.get("deny") == Some(&json!(true)) { + return PluginResult::deny(PluginViolation::new("test", "HTTP hook failure")); + } + let mut updated = extensions.cow_copy(); + // Returning hidden fields is intentional: the host must enforce capabilities. + let mut http = extensions.http.as_deref().cloned().unwrap_or_default(); + if post { + http.set_response_header("x-plugin-status", payload.status_code.expect("post status").to_string()); + http.set_response_header("x-plugin-count", ctx.get_local("count").expect("count").to_string()); + } else { + http.set_request_header("x-from-plugin", "present"); + if let Some(auth) = config.get("authorization").and_then(Value::as_str) { + http.set_request_header("Authorization", auth); + } + if config.get("invalid") == Some(&json!(true)) { + http.set_request_header("x-invalid", "bad\r\nheader"); + } + if config.get("oversized") == Some(&json!(true)) { + http.set_request_header("Mcp-Name", "x".repeat(100_000)); + } + } + updated.http = Some(cpex::cpex_core::extensions::Guarded::new(http)); + PluginResult::modify_extensions(updated) + } +} + +fn configured(hooks: &[&str], config: Value) -> Value { + let mut plugin = json!({"name":"probe", "kind":"http-probe", "hooks":hooks, + "capabilities":["read_headers","write_headers","read_subject"]}); + plugin["config"] = config; + plugin +} + +fn document(global: Value, scoped: Value) -> RuntimePluginConfigDocument { + let mut document = json!({"enabled":true, "global":{}, "contexts":{"team::sum":{}}}); + document["global"]["plugins"] = global; + document["contexts"]["team::sum"]["plugins"] = scoped; + serde_json::from_value(document).expect("document") +} + +async fn runtime(document: RuntimePluginConfigDocument) -> (Arc, Seen) { + let seen = Seen::default(); + let captured = Arc::clone(&seen); + let mut registry = CpexRuntimeRegistry::default(); + registry + .register_factory( + "http-probe", + Box::new( + GatewayPluginFactory::new(move |config| Probe { config, seen: Arc::clone(&captured) }) + .with_cmf_hooks() + .with_http_hooks(), + ), + ) + .expect("factory"); + registry.apply_document(document).await.expect("configure"); + (Arc::new(registry), seen) +} + +fn http_extensions() -> Extensions { + Extensions { + request: Some(Arc::new(RequestExtension { + request_id: Some("shared-request".to_owned()), + ..Default::default() + })), + http: Some(Arc::new(HttpExtension { + method: Some("POST".to_owned()), + path: Some("/mcp".to_owned()), + ..Default::default() + })), + security: Some(Arc::new(SecurityExtension::default())), + ..Default::default() + } +} + +#[tokio::test] +async fn http_and_scoped_mcp_hooks_preserve_priority_and_state_across_reload() { + let plugins = |hooks: &[&str]| { + [("later", 90), ("earlier", 1), ("equal-priority", 1)].map(|(name, priority)| { + let mut plugin = configured(hooks, json!({})); + plugin["name"] = json!(name); + plugin["priority"] = json!(priority); + plugin + }) + }; + let global = plugins(&["http_pre_request", "http_post_request"]); + let scoped = plugins(&["tool_pre_invoke", "tool_post_invoke"]); + let (registry, seen) = runtime(document(json!(global), json!(scoped))).await; + let http = registry + .handle() + .before_http_request( + HttpHookPayload { + method: "POST".to_owned(), + path: "/mcp".to_owned(), + client_addr: None, + status_code: None, + }, + http_extensions(), + ) + .await + .expect("HTTP before"); + registry.apply_config(None).await.expect("reload between HTTP and MCP"); + http.request.set_verified_subject(SubjectExtension { id: Some("verified".to_owned()), ..Default::default() }).await; + let state = registry + .handle() + .before_tool_call( + &CallToolRequestParams::new("sum"), + "sum", + "backend", + PluginRequestContext { + request: Some(http.request.clone()), + tool: Some(ToolPolicyContext { + id: "tool-id".to_owned(), + name: "sum".to_owned(), + team_id: Some("team".to_owned()), + context_id: "team::sum".to_owned(), + }), + extensions: Extensions { + security: Some(Arc::new(SecurityExtension { + subject: Some(SubjectExtension { + id: Some("stale-route-copy".to_owned()), + ..Default::default() + }), + ..Default::default() + })), + ..Default::default() + }, + }, + ) + .await + .expect("MCP before uses pinned scoped config") + .state + .expect("post state"); + state.after_tool_call(CallToolResult::success(vec![])).await.expect("MCP after"); + let extensions = http.after_http_request(200, HashMap::new()).await; + assert_eq!(extensions.http.as_ref().expect("headers").get_response_header("x-plugin-count"), Some("4")); + let observations = seen.lock().expect("observations"); + assert_eq!(observations.len(), 12); + for (index, observation) in observations.iter().enumerate() { + let count = index / 3; + assert_eq!(observation["plugin"], ["earlier", "equal-priority", "later"][index % 3]); + assert_eq!(observation["count"], count); + assert_eq!(observation["request"], "shared-request"); + assert_eq!(observation["subject"], if count == 0 { Value::Null } else { json!("verified") }); + } +} + +#[tokio::test] +async fn write_only_http_hooks_preserve_headers_and_transport_metadata() { + for auth_override in [false, true] { + let mut writer = configured(&["http_pre_request", "http_post_request"], json!({"authorization":"changed"})); + writer["name"] = json!("writer"); + writer["capabilities"] = json!(["write_headers"]); + writer["priority"] = json!(1); + let reader = configured(&["http_pre_request", "http_post_request"], json!({})); + let mut config = document(json!([writer, reader]), json!([])); + config.settings.plugins_can_override_auth_headers = auth_override; + let (registry, seen) = runtime(config).await; + let mut extensions = http_extensions(); + let http = Arc::make_mut(extensions.http.as_mut().expect("HTTP metadata")); + http.host = Some("gateway.example".to_owned()); + http.scheme = Some("https".to_owned()); + http.set_request_header("content-type", "application/json"); + http.set_request_header("authorization", "original"); + http.set_request_header("x-from-plugin", "old"); + let state = registry + .handle() + .before_http_request( + HttpHookPayload { + method: "POST".to_owned(), + path: "/mcp".to_owned(), + client_addr: None, + status_code: None, + }, + extensions, + ) + .await + .expect("HTTP before"); + let extensions = state + .after_http_request( + 200, + HashMap::from([ + ("content-type".to_owned(), "application/json".to_owned()), + ("x-plugin-status".to_owned(), "old".to_owned()), + ]), + ) + .await; + let http = extensions.http.expect("HTTP metadata survives blind writes"); + assert_eq!(http.method.as_deref(), Some("POST")); + assert_eq!(http.path.as_deref(), Some("/mcp")); + assert_eq!(http.host.as_deref(), Some("gateway.example")); + assert_eq!(http.scheme.as_deref(), Some("https")); + assert_eq!(http.get_request_header("content-type"), Some("application/json")); + assert_eq!(http.get_response_header("content-type"), Some("application/json")); + assert_eq!(http.get_request_header("x-from-plugin"), Some("present")); + assert_eq!(http.get_response_header("x-plugin-status"), Some("200")); + assert_eq!(http.get_request_header("authorization"), Some(if auth_override { "changed" } else { "original" })); + let observations = seen.lock().expect("observations"); + assert_eq!(observations.len(), 4); + for index in [0, 2] { + assert_eq!(observations[index]["headers_visible"], false, "write access must not grant read access"); + assert_eq!(observations[index + 1]["request_content_type"], "application/json"); + } + assert_eq!(observations[3]["response_content_type"], "application/json"); + } +} + +async fn send(gateway: &TestServer, authorization: Option<&str>) -> reqwest::Response { + let client = reqwest::Client::new(); + let request = client + .post(gateway.url("/contextforge-rs/servers/test/mcp")) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "server/discover") + .json(&json!({"jsonrpc":"2.0","id":1,"method":"server/discover","params":{"_meta":{"io.modelcontextprotocol/protocolVersion":"2026-07-28","io.modelcontextprotocol/clientCapabilities":{}}}})); + let request = if let Some(auth) = authorization { request.header("Authorization", auth) } else { request }; + request.send().await.expect("HTTP response") +} + +#[tokio::test] +async fn request_header_edits_run_before_auth_and_response_hooks_cover_auth_failures() { + let valid = format!("Bearer {}", token(TEST_USER_ID)); + let probe = configured(&["http_pre_request", "http_post_request"], json!({"authorization":valid})); + let (registry, seen) = runtime(document(json!([probe]), json!([]))).await; + let gateway = start_verified_gateway(Arc::clone(®istry)).await; + let response = send(&gateway, None).await; + assert_eq!(response.status(), 200, "plugin supplies a credential which normal authentication verifies"); + assert_eq!(response.headers()["x-plugin-status"], "200"); + { + let events = seen.lock().expect("observations"); + assert!(events[0]["subject"].is_null()); + assert_eq!(events[1]["subject"], TEST_USER_ID, "HTTP after sees verified identity without any MCP plugin"); + } + let response = send(&gateway, Some("Bearer invalid")).await; + assert_eq!(response.status(), 401, "plugin cannot replace an existing credential by default"); + assert_eq!(response.headers()["x-plugin-status"], "401"); + let missing_user = format!("Bearer {}", token("missing-config")); + assert_eq!(send(&gateway, Some(&missing_user)).await.status(), 400); + assert_eq!( + seen.lock().expect("observations").last().expect("HTTP after")["subject"], + "missing-config", + "verified identity survives user configuration lookup failure" + ); + let mut override_config = document( + json!([configured(&["http_pre_request", "http_post_request"], json!({"authorization":valid}))]), + json!([]), + ); + override_config.settings.plugins_can_override_auth_headers = true; + registry.apply_document(override_config).await.expect("allow credential replacement"); + assert_eq!(send(&gateway, Some("Bearer invalid")).await.status(), 200); + gateway.shutdown().await.expect("shutdown"); +} + +#[tokio::test] +async fn capabilities_and_published_header_policy_are_enforced() { + for (capabilities, policy_allows) in [(false, true), (true, false)] { + let mut probe = configured(&["http_pre_request", "http_post_request"], json!({})); + if !capabilities { + probe["capabilities"] = json!([]); + } + let mut config = document(json!([probe]), json!([])); + if !policy_allows { + config.settings.default_hook_policy = "deny".to_owned(); + } + let (registry, seen) = runtime(config).await; + let gateway = start_verified_gateway(registry).await; + let response = send(&gateway, None).await; + assert_eq!(response.status(), 401); + assert!(response.headers().get("x-plugin-status").is_none()); + assert!( + seen.lock().expect("observations").iter().all(|observation| observation["headers_visible"] == capabilities) + ); + gateway.shutdown().await.expect("shutdown"); + } +} + +#[tokio::test] +async fn post_only_hooks_run_and_invalid_header_edits_are_ignored() { + for (hooks, config) in [ + (vec!["http_post_request"], json!({})), + ( + vec!["http_pre_request", "http_post_request"], + json!({"invalid":true,"authorization":format!("Bearer {}", token(TEST_USER_ID))}), + ), + ] { + let (registry, _) = runtime(document(json!([configured(&hooks, config)]), json!([]))).await; + let gateway = start_verified_gateway(registry).await; + let response = send(&gateway, None).await; + assert_eq!(response.status(), 401); + assert_eq!(response.headers()["x-plugin-status"], "401"); + gateway.shutdown().await.expect("shutdown"); + } +} + +#[tokio::test] +async fn denied_or_timed_out_http_hooks_continue_but_header_limits_still_apply() { + for config in [json!({"deny":true}), json!({"timeout":true}), json!({"oversized":true})] { + let oversized = config.get("oversized").is_some(); + let mut config = document(json!([configured(&["http_pre_request", "http_post_request"], config)]), json!([])); + config.settings.plugin_timeout = 1; + let (registry, _) = runtime(config).await; + let gateway = start_verified_gateway(registry).await; + let response = + tokio::time::timeout(Duration::from_secs(5), send(&gateway, None)).await.expect("bounded hook execution"); + assert_eq!(response.status().as_u16(), if oversized { 431 } else { 401 }); + gateway.shutdown().await.expect("shutdown"); + } +} + +#[tokio::test] +async fn streaming_http_after_runs_before_a_long_running_tool_finishes() { + let probe = configured( + &["http_pre_request", "http_post_request", "cmf.tool_pre_invoke", "cmf.tool_post_invoke"], + json!({}), + ); + let (registry, seen) = runtime(document(json!([probe.clone()]), json!([]))).await; + // This fixture intentionally uses one static policy; the scoped case is covered above. + registry + .apply_config(Some(serde_json::from_value(json!({"plugins":[probe]})).expect("config"))) + .await + .expect("static policy"); + let gateway = crate::harness::start_gateway(TEST_USER_ID, true, registry).await; + let request = reqwest::Client::new() + .post(gateway.gateway_url()) + .bearer_auth(token(TEST_USER_ID)) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2026-07-28") + .header("Mcp-Method", "tools/call") + .header("Mcp-Name", "wait_for_cancellation") + .json(&json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{ + "name":"wait_for_cancellation","arguments":{},"_meta":{ + "io.modelcontextprotocol/protocolVersion":"2026-07-28", + "io.modelcontextprotocol/clientCapabilities":{}, "progressToken":"http-streaming" + } + }})); + let response = tokio::time::timeout(Duration::from_secs(2), request.send()) + .await + .expect("headers must not wait for tool completion") + .expect("HTTP response"); + assert_eq!(response.status(), 200); + assert!(response.headers()["content-type"].to_str().expect("content type").starts_with("text/event-stream")); + assert_eq!(response.headers()["x-plugin-status"], "200"); + { + let events = seen.lock().expect("observations"); + // The MCP post hook also processes the first progress notification. + assert_eq!( + events.iter().map(|event| event["stage"].as_str().expect("stage")).collect::>(), + ["http_pre", "mcp", "mcp", "http_post"] + ); + for (count, event) in events.iter().enumerate() { + assert_eq!(event["count"], count, "real HTTP and MCP hooks share local state"); + assert_eq!(event["request"], events[0]["request"]); + } + } + drop(response); + gateway.shutdown().await.expect("shutdown"); +} + +#[tokio::test] +async fn prompt_and_resource_hooks_use_the_same_request_lifecycle_with_global_policy() { + use rmcp::model::{GetPromptRequestParams, GetPromptResult, ReadResourceResult}; + let probe = + configured(&["http_pre_request", "http_post_request", "prompt_post_fetch", "resource_post_fetch"], json!({})); + let (registry, seen) = runtime(document(json!([probe]), json!([]))).await; + for prompt in [true, false] { + let http = registry + .handle() + .before_http_request( + HttpHookPayload { + method: "POST".to_owned(), + path: "/mcp".to_owned(), + client_addr: None, + status_code: None, + }, + http_extensions(), + ) + .await + .expect("HTTP before"); + let context = PluginRequestContext { request: Some(http.request.clone()), ..Default::default() }; + if prompt { + registry + .handle() + .before_get_prompt(&GetPromptRequestParams::new("review"), "review", "backend", context) + .await + .expect("prompt before") + .state + .expect("post-only state") + .after_get_prompt(GetPromptResult::new(vec![])) + .await + .expect("prompt after"); + } else { + registry + .handle() + .before_read_resource("file:///test", context) + .await + .expect("resource before") + .after_read_resource(ReadResourceResult::new(vec![])) + .await + .expect("resource after"); + } + http.after_http_request(200, HashMap::new()).await; + } + let observations = seen.lock().expect("observations"); + assert_eq!(observations.len(), 6); + for (index, event) in observations.iter().enumerate() { + assert_eq!(event["count"], index % 3, "state is shared within each request and isolated between requests"); + } +} diff --git a/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs b/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs index 9c27c3ce..4504c22d 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs @@ -8,7 +8,7 @@ use async_trait::async_trait; use contextforge_data_plane_apis::{ User, runtime_plugin_config::RuntimePluginConfigDocument, user_store::ToolPolicyContext, }; -use contextforge_data_plane_cpex::{CmfPluginFactory, CpexRuntimeRegistry}; +use contextforge_data_plane_cpex::{CpexRuntimeRegistry, GatewayPluginFactory}; use contextforge_data_plane_lib::UserConfigStore; use cpex::cpex_core::{ cmf::{CmfHook, MessagePayload}, @@ -136,7 +136,10 @@ async fn runtime(document: RuntimePluginConfigDocument) -> (Arc Result<(), Box> { runtime.register_factory( "contextforge/payload-marker", - Box::new(CmfPluginFactory::new(cpex_payload_marker::PayloadMarkerPlugin::new)), + Box::new(GatewayPluginFactory::new(cpex_payload_marker::PayloadMarkerPlugin::new).with_cmf_hooks()), )?; runtime.register_factory( "contextforge/text-prefixer", - Box::new(CmfPluginFactory::new(cpex_text_prefixer::TextPrefixerPlugin::new)), + Box::new(GatewayPluginFactory::new(cpex_text_prefixer::TextPrefixerPlugin::new).with_cmf_hooks()), )?; runtime.register_factory( "contextforge/tool-namespace", - Box::new(CmfPluginFactory::new(cpex_tool_namespace::ToolNamespacePlugin::new)), + Box::new(GatewayPluginFactory::new(cpex_tool_namespace::ToolNamespacePlugin::new).with_cmf_hooks()), )?; Ok(()) }