From ebb1820e1ded88cf15a284d695c2eeec97bca7bd Mon Sep 17 00:00:00 2001 From: lucarlig Date: Wed, 9 Sep 2026 16:30:15 +0100 Subject: [PATCH] feat: apply published plugin policies and context to MCP hooks Signed-off-by: lucarlig --- README.md | 2 +- _context/wiki/config.md | 102 ++++- _context/wiki/deployment.md | 4 +- _context/wiki/failure-modes.md | 5 +- _context/wiki/getting-started.md | 38 +- _context/wiki/routing.md | 4 +- .../src/runtime_plugin_config.rs | 67 ++- .../src/user_store.rs | 13 + .../src/config.rs | 30 +- .../src/extension_guard.rs | 164 ++++++++ .../contextforge-data-plane-cpex/src/hooks.rs | 10 + .../contextforge-data-plane-cpex/src/lib.rs | 3 +- .../src/prompts/mod.rs | 8 +- .../src/registry/mod.rs | 73 +++- .../src/registry/tests.rs | 242 ++++++++--- .../src/resources/mod.rs | 13 +- .../src/runtime.rs | 166 ++++++-- .../src/tools/mod.rs | 8 +- crates/contextforge-data-plane-lib/Cargo.toml | 2 +- .../authorization/principal_extractor/mod.rs | 6 + .../src/gateway/mcp_service/initialization.rs | 1 + .../src/gateway/mcp_service/mod.rs | 1 + .../src/gateway/mcp_service/plugin_context.rs | 135 ++++++ .../src/gateway/mcp_service/prompts.rs | 9 +- .../src/gateway/mcp_service/resources.rs | 13 +- .../src/gateway/mcp_service/tools.rs | 9 +- .../src/layers/virtual_host_config.rs | 1 + .../user_config_store/redis_config_store.rs | 2 +- .../tests/gateway.rs | 2 + .../gateway/future_contracts/pagination.rs | 2 + .../tests/gateway/harness/plugin_gateway.rs | 2 + .../tests/gateway/harness/test_gateways.rs | 2 + .../tests/gateway/plugin_context.rs | 398 ++++++++++++++++++ .../tests/secrets_detection_e2e.rs | 27 +- .../plugins/cpex-secrets-detection/README.md | 14 +- .../plugins/cpex-secrets-detection/src/lib.rs | 4 +- .../tests/plugin_manager.rs | 2 +- schemas/user_config.json | 85 +++- 38 files changed, 1493 insertions(+), 176 deletions(-) create mode 100644 crates/contextforge-data-plane-cpex/src/extension_guard.rs create mode 100644 crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs create mode 100644 crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs diff --git a/README.md b/README.md index 2f4afbfb..2109390c 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ Activation requires all three pieces: - Runtime flag: `--runtime-plugins-enabled true` - Redis config key: `ContextForgeGatewayRuntimePluginConfig` -The plugin kind is `validator/secrets-detection`. The dataplane wires CMF hooks +The plugin kind is `cpex_secrets_detection.SecretsDetectionPlugin`. The dataplane wires CMF hooks for tool calls, prompt fetches, and resource reads. Example run command: diff --git a/_context/wiki/config.md b/_context/wiki/config.md index 0b1bd36c..f9f9e222 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -139,6 +139,7 @@ the dataplane after removing a key if that cache must be cleared immediately. ```text UserConfig virtual_hosts: HashMap + user_email: String | null ← optional control-plane user name VirtualHost backends: HashMap ← backend key, not a parsed prefix @@ -160,6 +161,13 @@ BackendMCPGateway remove_headers: Vec ← defaults to [] completion: HashMap ← defaults to {}; completion is not implemented tool_schemas: HashMap ← defaults to {}; upstream tool name → schema + tool_policy_contexts: HashMap ← upstream tool name → identity/policy + +ToolPolicyContext + id: String ← canonical tool ID + name: String ← canonical gateway tool name + team_id: String | null ← owning team + context_id: String ← key in plugin document contexts ``` The virtual-host object maps default to empty. A backend must be referenced by @@ -219,16 +227,69 @@ cargo run -p contextforge-data-plane-apis ```text RuntimePluginConfigDocument - version: 1 - cpex: CpexConfig + enabled: bool + global: CpexConfig | null + contexts: HashMap + settings: runtime execution settings ``` -Supported: tool, prompt, and resource pre/post CMF hooks. -Rejected: routing-based selection, routes, plugin directories, global policies/defaults, -`plugin_settings.fail_on_plugin_error`, plugin conditions, and unsupported hooks -(including LLM hooks). -Config validation and `CmfPluginFactory` registration must agree on that list: a hook accepted by validation but not registered leaves the plugin loaded and silently inert. -Reload watcher: 10-minute interval. Invalid reload → runtime marked failed. +The unversioned document is published by the control plane in MessagePack; +JSON is also accepted. Python `plugins: null` means an empty plugin list. +Native hook names such as `tool_pre_invoke` are mapped to CPEX CMF handlers. +The secrets-detection factory uses the publisher's +`cpex_secrets_detection.SecretsDetectionPlugin` kind. + +Tools select the resolved configuration using +`backend.tool_policy_contexts[upstream_name].context_id`. The entry also carries +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 +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. + +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 +runtime, matched pre/post hooks, CPEX local/shared state, and permitted extension +updates, including post-only hooks and tool progress/logging events. + +Plugin extensions expose the server-generated request ID, active trace/span, +canonical target, tool schema/identity, gateway and virtual-server IDs, HTTP +request data, and verified subject claims. `UserConfig.user_email` supplies the +control-plane user name when present; otherwise the verified principal is used. +Target policy scope is stored in `meta.scope`, separately from subject claims. +CPEX's condition matcher uses that scope, canonical target, authenticated user, +server and request content type. Client arguments and plugin state cannot select +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. + +CPEX filters each plugin's view by its declared capabilities. The gateway guards +every extension write-back before the next plugin executes: identity remains +host-owned, hidden fields remain intact, labels are append-only with +`append_labels`, and HTTP writes require `write_headers`. Authentication headers +remain unchanged unless the published setting permits them. Permitted custom +state and extension updates survive pre/post execution; upstream actions remain +with the separate actions integration. Published hook payload policies control +whether the existing gateway projection accepts argument, result, URI or content +edits. `settings.plugin_timeout` sets the maximum duration of each plugin hook invocation +in seconds (default: 30). For example, `"settings": {"plugin_timeout": 5}` allows +five seconds per plugin in both pre and post hooks. CPEX enforces this with +`tokio::time::timeout`; expiry drops the handler future and follows the plugin's +`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 +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. Compile bundled factories with `--features plugins` and enable execution with `--runtime-plugins-enabled true`. A valid document must exist before startup; @@ -281,19 +342,26 @@ configuration; run it before starting the ContextForge external dataplane: ```bash docker compose -f docker/docker-compose-local.yaml exec -T redis \ redis-cli SET ContextForgeGatewayRuntimePluginConfig '{ - "version": 1, - "cpex": { - "plugins": [ - { - "name": "payload-marker", - "kind": "contextforge/payload-marker", - "hooks": ["cmf.tool_post_invoke"] - } - ] + "enabled": true, + "global": {"plugins": []}, + "contexts": { + "quickstart-counter": { + "plugins": [ + { + "name": "payload-marker", + "kind": "contextforge/payload-marker", + "hooks": ["cmf.tool_post_invoke"] + } + ] + } } }' ``` +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. + 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 b0cef3b2..e577abb0 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 ten minutes. An invalid reload fails new plugin calls closed. + every 30 seconds. An invalid reload fails new plugin calls closed. ## Builds and Images @@ -113,7 +113,7 @@ publisher interval + user-config cache expiry + publication/read latency The Rust cache defaults to 60 seconds; check the deployed publisher's actual interval. For functional tests, shorten publication and use cache expiry `0`. For benchmarks, report both values and keep them consistent between runs. -CPEX reloads use a separate ten-minute interval. +CPEX reloads use a separate 30-second interval. ## Security Posture diff --git a/_context/wiki/failure-modes.md b/_context/wiki/failure-modes.md index 7c7c42cf..6711b59d 100644 --- a/_context/wiki/failure-modes.md +++ b/_context/wiki/failure-modes.md @@ -52,8 +52,9 @@ modern routing. A configured backend alone does not make its objects callable. | Failure | Behavior | | --- | --- | -| Pre-hook denies | MCP error; no upstream call. | -| Post-hook denies | MCP error; backend operation may already have completed. | +| MCP pre-hook denies | MCP error; no upstream call. | +| MCP post-hook denies | MCP error; backend operation may already have completed. | +| 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. | diff --git a/_context/wiki/getting-started.md b/_context/wiki/getting-started.md index 2f4a5e6d..17937852 100644 --- a/_context/wiki/getting-started.md +++ b/_context/wiki/getting-started.md @@ -40,20 +40,27 @@ secrets detection before and after tool calls, blocking detected secrets: ```bash docker compose -f docker/docker-compose-local.yaml exec -T redis \ redis-cli SET ContextForgeGatewayRuntimePluginConfig '{ - "version": 1, - "cpex": { - "plugins": [{ - "name": "secrets-detection", - "kind": "validator/secrets-detection", - "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], - "config": {"block_on_detection": true} - }] + "enabled": true, + "global": {"plugins": []}, + "contexts": { + "quickstart-counter": { + "plugins": [{ + "name": "secrets-detection", + "kind": "cpex_secrets_detection.SecretsDetectionPlugin", + "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], + "config": {"block_on_detection": true} + }] + } } }' NX ``` `NX` preserves an existing plugin document. `OK` means the example was inserted; -an empty reply means an existing document remains in use. Its plugin kinds must +an empty reply means an existing document remains in use. Existing documents must +use the unversioned `enabled`/`global`/`contexts` format and include the tool policy +selected by the routes below; the former `version`/`cpex` format is no longer +accepted. The `quickstart-counter` key is a local example, paired with the +backend's `tool_policy_contexts.get_value.context_id` in step 5. Plugin kinds must be compiled into the binary. See [Plugin Config](config.md#plugin-config-redis-key-contextforgegatewayruntimepluginconfig) for configuration and the optional [demo plugins](config.md#demo-plugin-workflow). @@ -127,7 +134,15 @@ curl --fail --silent --show-error --request POST \ "name": "gateway-one", "url": "http://127.0.0.1:5555/mcp", "mcp_protocol_version": "2026-07-28", - "passthrough_headers": [] + "passthrough_headers": [], + "tool_policy_contexts": { + "get_value": { + "id": "counter-get-value", + "name": "counter-get_value", + "team_id": "team_awesome", + "context_id": "quickstart-counter" + } + } } }, "tools": { @@ -144,6 +159,8 @@ curl --fail --silent --show-error --request POST \ Expect `Added` (HTTP `202`). The `USER_ID` must match the token's `sub`. The backend protocol version and explicit tool route are required for the tool call below. The public tool name maps to the backend's original `get_value`. +Its policy context selects the secrets-detection configuration seeded in step 2; +tool calls do not fall back to `global` when that context is missing. ### 6. Discover the server and call the counter @@ -219,6 +236,7 @@ tool name directly instead of expecting `tools/list` here. | `404 {"detail":"Server not found"}` | The URL's virtual-host ID must exist in that user's config. | | `400` mentioning request metadata | Include the matching MCP protocol header, method/name headers, and per-request `_meta`. | | MCP error for an unpublished tool | Add the tool's explicit route to the virtual host before calling it. | +| `Runtime plugin tool context is missing` / `Runtime plugin policy context is missing` | Publish the backend's tool identity and a matching entry in the plugin document's `contexts`; `global` alone does not configure tool policies. | | Backend unavailable | Check `gateway-one` logs, port `5555`, and `--upstream-connection-mode plain-text-or-tls`. | For JWT diagnostics, restart with `RUST_LOG=debug` and look for diff --git a/_context/wiki/routing.md b/_context/wiki/routing.md index fc8393dc..e18b6966 100644 --- a/_context/wiki/routing.md +++ b/_context/wiki/routing.md @@ -66,7 +66,9 @@ 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. -The handle selects a runtime before backend I/O. It returns typed request state +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 diff --git a/crates/contextforge-data-plane-apis/src/runtime_plugin_config.rs b/crates/contextforge-data-plane-apis/src/runtime_plugin_config.rs index d5e09305..a128daa6 100644 --- a/crates/contextforge-data-plane-apis/src/runtime_plugin_config.rs +++ b/crates/contextforge-data-plane-apis/src/runtime_plugin_config.rs @@ -1,11 +1,68 @@ +use std::collections::{HashMap, HashSet}; + use cpex::cpex_core::config::CpexConfig; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; pub const RUNTIME_PLUGIN_CONFIG_KEY: &str = "ContextForgeGatewayRuntimePluginConfig"; -pub const RUNTIME_PLUGIN_CONFIG_VERSION: u8 = 1; - #[derive(Clone, Debug, Serialize, Deserialize)] pub struct RuntimePluginConfigDocument { - pub version: u8, - pub cpex: CpexConfig, + pub enabled: bool, + #[serde(deserialize_with = "optional_config")] + pub global: Option, + #[serde(deserialize_with = "context_configs")] + pub contexts: HashMap, + #[serde(default)] + pub settings: RuntimePluginSettings, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default)] +pub struct RuntimePluginSettings { + /// Maximum seconds per plugin hook invocation; defaults to 30. + pub plugin_timeout: u64, + pub fail_on_plugin_error: bool, + pub execution_pool: usize, + pub default_hook_policy: String, + pub hook_policies: HashMap, + pub plugins_can_override_rbac: bool, + pub plugins_can_override_auth_headers: bool, +} + +impl Default for RuntimePluginSettings { + fn default() -> Self { + Self { + plugin_timeout: 30, + fail_on_plugin_error: false, + execution_pool: 10, + default_hook_policy: "allow".to_owned(), + hook_policies: HashMap::new(), + plugins_can_override_rbac: false, + plugins_can_override_auth_headers: false, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct HookPayloadPolicy { + pub writable_fields: HashSet, +} + +// The publisher serializes Python Config.plugins=None as JSON/MessagePack null. +// CPEX Rust represents the same empty configuration with an empty vector. +fn published_config(mut value: serde_json::Value) -> Result { + if value.get("plugins").is_some_and(serde_json::Value::is_null) { + value["plugins"] = serde_json::json!([]); + } + serde_json::from_value(value).map_err(E::custom) +} + +fn optional_config<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + Option::::deserialize(deserializer)?.map(published_config).transpose() +} + +fn context_configs<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + HashMap::::deserialize(deserializer)? + .into_iter() + .map(|(key, value)| published_config(value).map(|config| (key, config))) + .collect() } diff --git a/crates/contextforge-data-plane-apis/src/user_store.rs b/crates/contextforge-data-plane-apis/src/user_store.rs index 6a6eb669..88be293b 100644 --- a/crates/contextforge-data-plane-apis/src/user_store.rs +++ b/crates/contextforge-data-plane-apis/src/user_store.rs @@ -38,6 +38,17 @@ pub struct BackendMCPGateway { /// Input schemas keyed by the original upstream tool name. #[serde(default)] pub tool_schemas: HashMap>, + /// Canonical tool identity and resolved policy key, indexed by upstream name. + #[serde(default)] + pub tool_policy_contexts: HashMap, +} + +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)] +pub struct ToolPolicyContext { + pub id: String, + pub name: String, + pub team_id: Option, + pub context_id: String, } #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)] @@ -62,4 +73,6 @@ pub struct VirtualHost { #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)] pub struct UserConfig { pub virtual_hosts: HashMap, + #[serde(default)] + pub user_email: Option, } diff --git a/crates/contextforge-data-plane-cpex/src/config.rs b/crates/contextforge-data-plane-cpex/src/config.rs index 4dc67006..08a245f8 100644 --- a/crates/contextforge-data-plane-cpex/src/config.rs +++ b/crates/contextforge-data-plane-cpex/src/config.rs @@ -7,10 +7,7 @@ use redis::{ use tokio::sync::Mutex; use crate::error::GatewayPluginRuntimeError; -use contextforge_data_plane_apis::runtime_plugin_config::{ - RUNTIME_PLUGIN_CONFIG_KEY, RUNTIME_PLUGIN_CONFIG_VERSION, RuntimePluginConfigDocument, -}; -use cpex::cpex_core::config::CpexConfig; +use contextforge_data_plane_apis::runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_KEY, RuntimePluginConfigDocument}; #[async_trait] pub(crate) trait RuntimePluginConfigStore: Send + Sync { @@ -69,13 +66,6 @@ impl RuntimePluginConfigStore for RedisRuntimePluginConfigStore { } } -pub(crate) fn cpex_config(document: &RuntimePluginConfigDocument) -> Result { - if document.version != RUNTIME_PLUGIN_CONFIG_VERSION { - return Err(GatewayPluginRuntimeError::ConfigWrongFormat); - } - Ok(document.cpex.clone()) -} - pub(crate) fn decode_config_document(config: &[u8]) -> Result { serde_json::from_slice::(config) .or_else(|_| rmp_serde::decode::from_slice::(config)) @@ -84,32 +74,28 @@ pub(crate) fn decode_config_document(config: &[u8]) -> Result>, + payload_writable: bool, + auth_headers_writable: bool, +} + +pub(crate) fn guarded_entries( + entries: &[HookEntry], + extensions: &Extensions, + payload_writable: bool, + auth_headers_writable: bool, +) -> Vec { + let canonical = Arc::new(Mutex::new(extensions.clone())); + entries + .iter() + .map(|entry| HookEntry { + plugin_ref: Arc::clone(&entry.plugin_ref), + handler: Arc::new(ExtensionGuard { + entry: entry.clone(), + canonical: Arc::clone(&canonical), + payload_writable, + auth_headers_writable, + }), + }) + .collect() +} + +#[async_trait] +impl AnyHookHandler for ExtensionGuard { + async fn invoke( + &self, + payload: &dyn PluginPayload, + extensions: &Extensions, + ctx: &mut PluginContext, + ) -> Result, Box> { + let result = self.entry.handler.invoke(payload, extensions, ctx).await?; + let mut result = result.downcast::().map_err(|_| { + Box::new(PluginError::Config { message: "Plugin returned an invalid hook result".to_owned() }) + })?; + if !self.payload_writable { + result.modified_payload = None; + } + if let Some(modified) = result.modified_extensions.take() + && self.entry.plugin_ref.mode().can_modify() + && result.continue_processing + { + let mut canonical = self.canonical.lock().map_err(|_| { + Box::new(PluginError::Config { message: "Plugin extension state is unavailable".to_owned() }) + })?; + 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") + && let Some(http) = modified.http + { + let mut http = http.into_inner(); + normalize_headers( + &mut http.request_headers, + canonical.http.as_ref().map(|http| &http.request_headers), + )?; + normalize_headers( + &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()); + } + // A write-only plugin cannot return headers or metadata it never saw. + // Merge header edits while retaining the host's transport metadata. + let mut merged = permitted.http.take().unwrap_or_default().into_inner(); + merged.request_headers.extend(http.request_headers); + merged.response_headers.extend(http.response_headers); + permitted.http = Some(cpex::cpex_core::extensions::Guarded::new(merged)); + } + if let Some(updated) = modified.security + && let Some(security) = &mut permitted.security + { + if caps.contains("append_labels") { + for label in updated.labels.iter() { + security.add_label(label.clone()); + } + } + security.objects = updated.objects; + security.data = updated.data; + security.classification = updated.classification; + } + permitted.custom = modified.custom; + canonical.merge_owned(permitted); + result.modified_extensions = Some(canonical.cow_copy()); + } + } + Ok(result) + } + + fn hook_type_name(&self) -> &'static str { + self.entry.handler.hook_type_name() + } +} + +fn normalize_headers( + headers: &mut std::collections::HashMap, + original: Option<&std::collections::HashMap>, +) -> Result<(), Box> { + let mut normalized = std::collections::HashMap::new(); + for (name, value) in std::mem::take(headers) { + let name = name.to_ascii_lowercase(); + if let Some(previous) = normalized.get(&name) { + if previous == &value { + continue; + } + let original = original.and_then(|headers| { + headers.iter().find(|(key, _)| key.eq_ignore_ascii_case(&name)).map(|(_, value)| value) + }); + if original == Some(&value) { + continue; + } + if original != Some(previous) { + return Err(Box::new(PluginError::Config { + message: "Plugin returned conflicting header values".to_owned(), + })); + } + } + normalized.insert(name, value); + } + *headers = normalized; + Ok(()) +} + +fn preserve_auth_headers( + headers: &mut std::collections::HashMap, + original: Option<&cpex::cpex_core::extensions::HttpExtension>, +) { + 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)); + if let Some(original) = original { + headers.extend( + original + .request_headers + .iter() + .filter(|(name, _)| protected(name)) + .map(|(name, value)| (name.clone(), value.clone())), + ); + } +} diff --git a/crates/contextforge-data-plane-cpex/src/hooks.rs b/crates/contextforge-data-plane-cpex/src/hooks.rs index 9291c391..46d5d5a0 100644 --- a/crates/contextforge-data-plane-cpex/src/hooks.rs +++ b/crates/contextforge-data-plane-cpex/src/hooks.rs @@ -1,5 +1,15 @@ +use contextforge_data_plane_apis::user_store::ToolPolicyContext; +use cpex::cpex_core::hooks::Extensions; 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 tool: Option, + pub extensions: Extensions, +} + pub type RuntimeHookError = Box; #[derive(Debug, Default)] diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index 5613e4cf..190f7a71 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -1,6 +1,7 @@ mod cmf; mod config; mod error; +mod extension_guard; mod factory; mod hooks; mod prompts; @@ -11,7 +12,7 @@ mod tools; pub use error::GatewayPluginRuntimeError; pub use factory::CmfPluginFactory; -pub use hooks::{ArgumentsUpdate, PreHookResult, RuntimeHookError}; +pub use hooks::{ArgumentsUpdate, PluginRequestContext, PreHookResult, RuntimeHookError}; pub use prompts::PromptHookState; pub use registry::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; pub use resources::ResourceHookState; diff --git a/crates/contextforge-data-plane-cpex/src/prompts/mod.rs b/crates/contextforge-data-plane-cpex/src/prompts/mod.rs index 74fbe51e..ae46b042 100644 --- a/crates/contextforge-data-plane-cpex/src/prompts/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/prompts/mod.rs @@ -14,7 +14,7 @@ use rmcp::{ use serde_json::{Map, Value}; use crate::{ - ArgumentsUpdate, GatewayPluginRuntimeHandle, PreHookResult, + ArgumentsUpdate, GatewayPluginRuntimeHandle, PluginRequestContext, PreHookResult, cmf::{CmfResponse, Operation, message_payload}, runtime::CallState, }; @@ -241,12 +241,14 @@ impl GatewayPluginRuntimeHandle { request: &GetPromptRequestParams, prompt_name: &str, backend_name: &str, + context: PluginRequestContext, ) -> Result, ErrorData> { let (arguments, state) = self .current()? + .global .before( - Operation::Prompt, - prompt_name, + (Operation::Prompt, prompt_name), + context.extensions, |id| prompt_request_payload(request, prompt_name, backend_name, id), |payload, id| { let arguments = diff --git a/crates/contextforge-data-plane-cpex/src/registry/mod.rs b/crates/contextforge-data-plane-cpex/src/registry/mod.rs index 180653de..f47fb981 100644 --- a/crates/contextforge-data-plane-cpex/src/registry/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/registry/mod.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -7,6 +8,7 @@ use std::{ }; use arc_swap::ArcSwap; +use contextforge_data_plane_apis::{runtime_plugin_config::RuntimePluginConfigDocument, user_store::ToolPolicyContext}; use cpex::cpex_core::{ config::CpexConfig, factory::{PluginFactory, PluginFactoryRegistry}, @@ -15,13 +17,13 @@ use rmcp::{ErrorData, model::ErrorCode}; use tokio::task::JoinHandle; use crate::{ - config::{RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, + config::{RedisRuntimePluginConfigStore, RuntimePluginConfigStore}, error::GatewayPluginRuntimeError, hooks::RuntimeHookError, runtime::GatewayPluginRuntime, }; -const DEFAULT_CONFIG_WATCHER_INTERVAL: Duration = Duration::from_mins(10); +const DEFAULT_CONFIG_WATCHER_INTERVAL: Duration = Duration::from_secs(30); pub struct CpexRuntimeRegistry { runtime: Arc>, @@ -37,14 +39,36 @@ pub struct GatewayPluginRuntimeHandle { } enum RuntimeState { - Active(Arc), + Active(Arc), Failed(String), } +#[derive(Default)] +pub(crate) struct RuntimePolicies { + pub(crate) global: Arc, + contexts: HashMap>, + scoped: bool, +} + +impl RuntimePolicies { + pub(crate) fn tool(&self, tool: Option<&ToolPolicyContext>) -> Result, ErrorData> { + if !self.scoped { + return Ok(Arc::clone(&self.global)); + } + let tool = tool + .filter(|tool| !tool.id.is_empty() && !tool.name.is_empty() && !tool.context_id.is_empty()) + .ok_or_else(|| ErrorData::internal_error("Runtime plugin tool context is missing", None))?; + self.contexts + .get(&tool.context_id) + .cloned() + .ok_or_else(|| ErrorData::internal_error("Runtime plugin policy context is missing", None)) + } +} + impl Default for CpexRuntimeRegistry { fn default() -> Self { Self { - runtime: Arc::new(ArcSwap::from_pointee(RuntimeState::Active(Arc::new(GatewayPluginRuntime::default())))), + runtime: Arc::new(ArcSwap::from_pointee(RuntimeState::Active(Arc::new(RuntimePolicies::default())))), config_store: None, factories: Arc::new(PluginFactoryRegistry::new()), watcher_started: AtomicBool::new(false), @@ -76,6 +100,10 @@ impl CpexRuntimeRegistry { apply_runtime_config(&self.runtime, &self.factories, config).await } + pub async fn apply_document(&self, document: RuntimePluginConfigDocument) -> Result<(), GatewayPluginRuntimeError> { + apply_document(&self.runtime, &self.factories, document).await + } + pub fn handle(&self) -> GatewayPluginRuntimeHandle { GatewayPluginRuntimeHandle { runtime: Arc::clone(&self.runtime) } } @@ -123,7 +151,7 @@ async fn reload_runtime( if last_applied_config == Some(config.fingerprint.as_slice()) { return Ok(None); } - apply_runtime_config(runtime, factories, Some(cpex_config(&config.document)?)).await?; + apply_document(runtime, factories, config.document).await?; Ok(Some(config.fingerprint)) } .await; @@ -139,14 +167,35 @@ async fn apply_runtime_config( config: Option, ) -> Result<(), GatewayPluginRuntimeError> { let Some(config) = config else { - drop(runtime.swap(Arc::new(RuntimeState::Active(Arc::new(GatewayPluginRuntime::default()))))); + drop(runtime.swap(Arc::new(RuntimeState::Active(Arc::new(RuntimePolicies::default()))))); return Ok(()); }; - drop( - runtime.swap(Arc::new(RuntimeState::Active(Arc::new( - GatewayPluginRuntime::from_config(config, factories).await?, - )))), - ); + drop(runtime.swap(Arc::new(RuntimeState::Active(Arc::new(RuntimePolicies { + global: Arc::new(GatewayPluginRuntime::from_config(config, factories).await?), + ..Default::default() + }))))); + Ok(()) +} + +async fn apply_document( + runtime: &ArcSwap, + factories: &PluginFactoryRegistry, + document: RuntimePluginConfigDocument, +) -> Result<(), GatewayPluginRuntimeError> { + if !document.enabled { + return apply_runtime_config(runtime, factories, None).await; + } + let global = document.global.ok_or(GatewayPluginRuntimeError::ConfigMissing)?; + let global = Arc::new(GatewayPluginRuntime::from_published_config(global, &document.settings, factories).await?); + let mut contexts = HashMap::new(); + for (key, config) in document.contexts { + if key.is_empty() { + return Err(GatewayPluginRuntimeError::ConfigWrongFormat); + } + let policy = GatewayPluginRuntime::from_published_config(config, &document.settings, factories).await?; + contexts.insert(key, Arc::new(policy)); + } + drop(runtime.swap(Arc::new(RuntimeState::Active(Arc::new(RuntimePolicies { global, contexts, scoped: true }))))); Ok(()) } @@ -162,7 +211,7 @@ impl CpexRuntimeRegistry { } impl GatewayPluginRuntimeHandle { - pub(crate) fn current(&self) -> Result, ErrorData> { + pub(crate) fn current(&self) -> Result, ErrorData> { match self.runtime.load().as_ref() { RuntimeState::Active(runtime) => Ok(Arc::clone(runtime)), RuntimeState::Failed(error) => { diff --git a/crates/contextforge-data-plane-cpex/src/registry/tests.rs b/crates/contextforge-data-plane-cpex/src/registry/tests.rs index ee2320bb..431b0804 100644 --- a/crates/contextforge-data-plane-cpex/src/registry/tests.rs +++ b/crates/contextforge-data-plane-cpex/src/registry/tests.rs @@ -24,7 +24,7 @@ use rmcp::model::{ use serde_json::{Value, json}; use tokio::sync::Mutex as TokioMutex; -use contextforge_data_plane_apis::runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_VERSION, RuntimePluginConfigDocument}; +use contextforge_data_plane_apis::runtime_plugin_config::{RuntimePluginConfigDocument, RuntimePluginSettings}; use crate::config::LoadedRuntimePluginConfig; use crate::{ArgumentsUpdate, CmfPluginFactory, PreHookResult, ToolHookState}; @@ -339,9 +339,22 @@ fn progress_event() -> ProgressNotificationParam { } fn config_document(cpex: Value) -> RuntimePluginConfigDocument { + let config: CpexConfig = serde_json::from_value(cpex).expect("test CPEX config parses"); + let mut global = config.clone(); + let mut scoped = config; + for plugin in &mut global.plugins { + plugin.hooks.retain(|hook| !hook.contains("tool_")); + } + global.plugins.retain(|plugin| !plugin.hooks.is_empty()); + for plugin in &mut scoped.plugins { + plugin.hooks.retain(|hook| hook.contains("tool_")); + } + scoped.plugins.retain(|plugin| !plugin.hooks.is_empty()); RuntimePluginConfigDocument { - version: RUNTIME_PLUGIN_CONFIG_VERSION, - cpex: serde_json::from_value(cpex).expect("test CPEX config parses"), + enabled: true, + global: Some(global), + contexts: HashMap::from([("test".to_owned(), scoped)]), + settings: RuntimePluginSettings::default(), } } @@ -398,11 +411,16 @@ async fn missing_runtime_plugin_config_is_rejected_on_initialize() { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn invalid_runtime_plugin_config_documents_are_rejected() { - for config in [RuntimePluginConfigDocument { version: 2, cpex: CpexConfig::default() }] { + for config in [RuntimePluginConfigDocument { + enabled: true, + global: None, + contexts: HashMap::new(), + settings: RuntimePluginSettings::default(), + }] { let runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); let error = runtime.initialize().await.expect_err("invalid config is rejected"); - assert_eq!("runtime plugin config is in wrong format", error.to_string()); + assert_eq!("runtime plugin config is missing", error.to_string()); } } @@ -411,14 +429,6 @@ async fn unsupported_runtime_plugin_config_is_rejected() { for cpex in [ json!({ "plugin_settings": { "routing_enabled": true }, "plugins": [] }), json!({ "plugin_settings": { "fail_on_plugin_error": true }, "plugins": [] }), - json!({ - "plugins": [{ - "name": "scoped", - "kind": "test", - "hooks": [cmf_hook_names::TOOL_PRE_INVOKE], - "conditions": [{ "tools": ["sum"] }] - }] - }), json!({ "plugins": [{ "name": "llm", "kind": "test", "hooks": [cmf_hook_names::LLM_INPUT] }] }), ] { let runtime = @@ -442,7 +452,11 @@ async fn resource_pre_hook_runs_for_a_canonical_uri() { let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - runtime.handle().before_read_resource("file:///password.env").await.expect("resource pre hook runs"); + runtime + .handle() + .before_read_resource("file:///password.env", test_request_context()) + .await + .expect("resource pre hook runs"); assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); } @@ -453,8 +467,12 @@ async fn resource_without_post_hook_keeps_its_decision_across_reload() { let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; runtime.apply_config(None).await.expect("disable hooks"); - let state = runtime.handle().before_read_resource("file:///password.env").await.expect("request starts"); - runtime.apply_config(Some(plugin_config(&[plugin]).cpex)).await.expect("enable hooks"); + let state = runtime + .handle() + .before_read_resource("file:///password.env", test_request_context()) + .await + .expect("request starts"); + runtime.apply_config(Some(plugin_config(&[plugin]).global.expect("global policy"))).await.expect("enable hooks"); let response = ReadResourceResult::new(vec![ResourceContents::text("original", "file:///password.env")]); state.after_read_resource(response).await.expect("in-flight decision survives reload"); assert_eq!(0, observations.lock().expect("observations lock poisoned").post_calls); @@ -464,7 +482,11 @@ async fn resource_without_post_hook_keeps_its_decision_across_reload() { async fn resource_post_hook_keeps_its_runtime_across_reload() { let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH]).with_post_deny()); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let state = runtime.handle().before_read_resource("file:///password.env").await.expect("request starts"); + let state = runtime + .handle() + .before_read_resource("file:///password.env", test_request_context()) + .await + .expect("request starts"); runtime.apply_config(None).await.expect("disable hooks"); let response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); let error = state.after_read_resource(response).await.expect_err("captured policy still denies"); @@ -479,7 +501,11 @@ async fn resource_hooks_preserve_context_across_the_backend_call() { ); let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.handle().before_read_resource("file:///password.env").await.expect("resource pre hook runs"); + let pre = runtime + .handle() + .before_read_resource("file:///password.env", test_request_context()) + .await + .expect("resource pre hook runs"); let response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); pre.after_read_resource(response).await.expect("resource post hook receives pre context"); @@ -495,7 +521,10 @@ async fn runtime_config_loads_registered_factory_plugin() { let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); assert!(matches!(result.arguments, ArgumentsUpdate::Replace(Some(_)))); assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); @@ -516,7 +545,10 @@ async fn runtime_config_loads_generic_cmf_factory_plugin() { .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); assert!(matches!(result.arguments, ArgumentsUpdate::Replace(Some(_)))); } @@ -538,7 +570,7 @@ async fn generic_cmf_factory_registers_prompt_only_plugin() { let result = runtime .handle() - .before_get_prompt(&review_request("weather"), "review", "backend") + .before_get_prompt(&review_request("weather"), "review", "backend", test_request_context()) .await .expect("prompt pre hook runs"); @@ -563,10 +595,13 @@ async fn generic_cmf_factory_registers_mixed_tool_and_prompt_plugin() { .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - let tool = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("tool pre hook runs"); + let tool = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("tool pre hook runs"); let prompt = runtime .handle() - .before_get_prompt(&review_request("weather"), "review", "backend") + .before_get_prompt(&review_request("weather"), "review", "backend", test_request_context()) .await .expect("prompt pre hook runs"); @@ -585,17 +620,26 @@ async fn runtime_reload_replaces_and_clears_current_runtime() { .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook skips"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook skips"); assert!(matches!(result.arguments, ArgumentsUpdate::Unchanged)); config_store.set_config(plugin_config(&[Arc::clone(&plugin)])).await; runtime.reload().await.expect("runtime reloads"); - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); assert!(matches!(result.arguments, ArgumentsUpdate::Replace(Some(_)))); config_store.set_config(config_document(json!({ "plugins": [] }))).await; runtime.reload().await.expect("runtime reloads"); - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook skips"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook skips"); assert!(matches!(result.arguments, ArgumentsUpdate::Unchanged)); assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); } @@ -611,21 +655,35 @@ async fn failed_runtime_reload_rejects_new_calls_until_valid_reload() { .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - config_store.set_config(RuntimePluginConfigDocument { version: 2, cpex: CpexConfig::default() }).await; + config_store + .set_config(RuntimePluginConfigDocument { + enabled: true, + global: None, + contexts: HashMap::new(), + settings: RuntimePluginSettings::default(), + }) + .await; runtime.reload().await.expect_err("invalid reload fails"); - let error = expect_runtime_failed(runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await); + let error = expect_runtime_failed( + runtime.before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()).await, + ); assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); assert_eq!("Runtime plugin reload failed", error.message); config_store.clear_config().await; runtime.reload().await.expect_err("missing reload fails"); - let error = expect_runtime_failed(runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await); + let error = expect_runtime_failed( + runtime.before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()).await, + ); assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); assert_eq!("Runtime plugin reload failed", error.message); config_store.set_config(plugin_config(&[Arc::clone(&plugin)])).await; runtime.reload().await.expect("runtime recovers"); - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); assert!(matches!(result.arguments, ArgumentsUpdate::Replace(Some(_)))); assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); } @@ -644,7 +702,10 @@ async fn combined_plugin_preserves_context_from_pre_to_post_across_replacement() .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); config_store.set_config(config_document(json!({ "plugins": [] }))).await; runtime.reload().await.expect("runtime reloads"); let response = CallToolResult::success(vec![ContentBlock::text("3")]); @@ -666,7 +727,10 @@ async fn post_only_runtime_does_not_apply_new_post_hook_to_in_flight_call() { .expect("test factory registers"); runtime.initialize().await.expect("runtime initializes"); - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook skips"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook skips"); config_store.set_config(plugin_config(&[Arc::clone(&plugin)])).await; runtime.reload().await.expect("runtime reloads"); let response = CallToolResult::success(vec![ContentBlock::text("3")]); @@ -682,7 +746,10 @@ async fn disabled_runtime_does_not_create_tool_post_state() { let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; runtime.apply_config(None).await.expect("disable hooks"); - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("request starts"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("request starts"); assert!(pre.state.is_none()); assert_eq!(0, observations.lock().expect("observations lock poisoned").post_calls); @@ -694,7 +761,10 @@ async fn stream_event_is_rewritten_by_post_hook() { let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre state is created"); let event = pre.state.expect("post state").after_stream_event(progress_event()).await.expect("event passes"); assert_eq!(Some("plugin:step 1/2"), event.expect("event is kept").message.as_deref()); @@ -706,7 +776,10 @@ async fn denied_stream_event_is_dropped() { let plugin = Arc::new(TestPlugin::new("post", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_post_deny()); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre state is created"); let event = pre.state.expect("post state").after_stream_event(progress_event()).await.expect("deny drops the event"); @@ -719,7 +792,10 @@ async fn invalid_stream_event_rewrite_is_rejected() { Arc::new(TestPlugin::new("post", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_invalid_stream_rewrite()); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre state is created"); let error = pre .state .expect("post state") @@ -768,16 +844,22 @@ async fn watcher_applies_config_changes() { config_store.set_config(plugin_config(&[Arc::clone(&plugin)])).await; for _ in 0..TEST_WATCHER_RETRY_COUNT { - let result = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre hook runs"); + let result = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("pre hook runs"); if matches!(result.arguments, ArgumentsUpdate::Replace(Some(_))) { config_store.clear_config().await; tokio::time::sleep(TEST_WATCHER_INTERVAL + TEST_WATCHER_RETRY_INTERVAL).await; - let error = expect_runtime_failed(runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await); + let error = expect_runtime_failed( + runtime.before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()).await, + ); assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); config_store.set_config(plugin_config(&[Arc::clone(&plugin)])).await; for _ in 0..TEST_WATCHER_RETRY_COUNT { - if let Ok(result) = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await + if let Ok(result) = + runtime.before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()).await && matches!(result.arguments, ArgumentsUpdate::Replace(Some(_))) { assert_eq!(2, observations.lock().expect("observations lock poisoned").pre_calls); @@ -816,8 +898,9 @@ impl CpexRuntimeRegistry { request: &CallToolRequestParams, tool_name: &str, backend_name: &str, + context: crate::PluginRequestContext, ) -> Result, ErrorData> { - self.handle().before_tool_call(request, tool_name, backend_name).await + self.handle().before_tool_call(request, tool_name, backend_name, context).await } async fn after_tool_call( @@ -864,7 +947,10 @@ async fn hook_combinations_preserve_correlation_for_each_operation() { match operation { Operation::Tool => { - let pre = handle.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("tool starts"); + let pre = handle + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("tool starts"); assert_eq!(post_enabled, pre.state.is_some()); if let Some(state) = pre.state { state.after_tool_call(CallToolResult::success(vec![])).await.expect("tool finishes"); @@ -872,7 +958,7 @@ async fn hook_combinations_preserve_correlation_for_each_operation() { }, Operation::Prompt => { let pre = handle - .before_get_prompt(&review_request("weather"), "review", "backend") + .before_get_prompt(&review_request("weather"), "review", "backend", test_request_context()) .await .expect("prompt starts"); assert_eq!(post_enabled, pre.state.is_some()); @@ -887,7 +973,10 @@ async fn hook_combinations_preserve_correlation_for_each_operation() { } }, Operation::Resource => { - let state = handle.before_read_resource("file:///test").await.expect("resource starts"); + let state = handle + .before_read_resource("file:///test", test_request_context()) + .await + .expect("resource starts"); state .after_read_resource(ReadResourceResult::new(vec![ResourceContents::text( "weather", @@ -922,11 +1011,12 @@ async fn prompt_context_and_policy_survive_a_failed_reload() { runtime.initialize().await.expect("runtime initializes"); let handle = runtime.handle(); let request = review_request("weather"); - let pre = handle.before_get_prompt(&request, "review", "backend").await.expect("prompt starts"); + let pre = + handle.before_get_prompt(&request, "review", "backend", test_request_context()).await.expect("prompt starts"); config_store.clear_config().await; runtime.reload().await.expect_err("missing config fails reload"); - assert!(handle.before_get_prompt(&request, "review", "backend").await.is_err()); - assert!(handle.before_read_resource("file:///test").await.is_err()); + assert!(handle.before_get_prompt(&request, "review", "backend", test_request_context()).await.is_err()); + assert!(handle.before_read_resource("file:///test", test_request_context()).await.is_err()); pre.state .expect("prompt state") .after_get_prompt(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, "weather")])) @@ -943,7 +1033,10 @@ async fn concurrent_tool_events_share_context_with_the_final_response_after_relo plugin.post_behavior = PostBehavior::CountEvents; let plugin = Arc::new(plugin); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("tool starts"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("tool starts"); let state = pre.state.expect("tool state"); runtime.apply_config(None).await.expect("disable hooks for new calls"); let (first, second) = @@ -958,7 +1051,10 @@ async fn concurrent_tool_events_share_context_with_the_final_response_after_relo async fn dropping_the_last_in_flight_state_releases_the_replaced_runtime() { let plugin = Arc::new(TestPlugin::new("pending", vec![cmf_hook_names::TOOL_POST_INVOKE])); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("tool starts"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("tool starts"); let state = pre.state.expect("tool state"); let event_state = state.clone(); runtime.apply_config(None).await.expect("replace runtime"); @@ -974,3 +1070,53 @@ async fn dropping_the_last_in_flight_state_releases_the_replaced_runtime() { } panic!("abandoned request retained its runtime"); } + +fn test_request_context() -> crate::PluginRequestContext { + crate::PluginRequestContext { + tool: Some(ToolPolicyContext { + id: "test".to_owned(), + name: "sum".to_owned(), + team_id: None, + context_id: "test".to_owned(), + }), + ..Default::default() + } +} + +#[tokio::test] +async fn published_hook_policy_controls_payload_writes_without_skipping_execution() { + use contextforge_data_plane_apis::runtime_plugin_config::HookPayloadPolicy; + let plugin = Arc::new(TestPlugin::new("policy", vec![cmf_hook_names::TOOL_PRE_INVOKE]).with_pre_rewrite()); + let mut config = plugin_config(&[Arc::clone(&plugin)]); + config.settings.default_hook_policy = "deny".to_owned(); + let runtime = runtime_with_plugin(&plugin, config.clone()).await; + let request = sum_request(1, 2); + let pre = + runtime.before_tool_call(&request, "sum", "backend", test_request_context()).await.expect("hook executes"); + assert!(matches!(pre.arguments, ArgumentsUpdate::Unchanged)); + assert_eq!(plugin.observations.lock().expect("observations").pre_calls, 1); + + config + .settings + .hook_policies + .insert("tool_pre_invoke".to_owned(), HookPayloadPolicy { writable_fields: ["args".to_owned()].into() }); + runtime.apply_document(config).await.expect("allow argument edits"); + let pre = + runtime.before_tool_call(&request, "sum", "backend", test_request_context()).await.expect("hook executes"); + assert!(matches!(pre.arguments, ArgumentsUpdate::Replace(Some(_)))); +} + +#[tokio::test] +async fn explicitly_disabled_plugins_do_not_require_a_factory() { + let config = config_document(json!({"plugins": [{ + "name": "disabled", "kind": "unavailable", "mode": "disabled", "hooks": ["tool_pre_invoke"], + }]})); + let runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); + runtime.initialize().await.expect("disabled plugin does not load"); + let pre = runtime + .before_tool_call(&sum_request(1, 2), "sum", "backend", test_request_context()) + .await + .expect("call allowed"); + assert!(pre.state.is_none()); + assert!(matches!(pre.arguments, ArgumentsUpdate::Unchanged)); +} diff --git a/crates/contextforge-data-plane-cpex/src/resources/mod.rs b/crates/contextforge-data-plane-cpex/src/resources/mod.rs index 069adc08..e0bee2b8 100644 --- a/crates/contextforge-data-plane-cpex/src/resources/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/resources/mod.rs @@ -8,7 +8,7 @@ use rmcp::{ }; use crate::{ - GatewayPluginRuntimeHandle, + GatewayPluginRuntimeHandle, PluginRequestContext, cmf::{CmfResponse, Operation, message_payload}, runtime::CallState, }; @@ -123,12 +123,17 @@ impl ResourceHookState { } impl GatewayPluginRuntimeHandle { - pub async fn before_read_resource(&self, resource_uri: &str) -> Result { + pub async fn before_read_resource( + &self, + resource_uri: &str, + context: PluginRequestContext, + ) -> Result { let (rewritten_uri, call) = self .current()? + .global .before( - Operation::Resource, - resource_uri, + (Operation::Resource, resource_uri), + 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 b9f8d543..f0c7b09e 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -4,6 +4,7 @@ use std::sync::{ atomic::{AtomicU64, Ordering}, }; +use contextforge_data_plane_apis::runtime_plugin_config::RuntimePluginSettings; use cpex::cpex_core::{ cmf::{CmfHook, MessagePayload}, config::CpexConfig, @@ -12,6 +13,8 @@ use cpex::cpex_core::{ factory::PluginFactoryRegistry, hooks::payload::Extensions, manager::PluginManager, + plugin::{MatchContext, PluginMode}, + registry::HookEntry, }; use rmcp::ErrorData; use tracing::instrument; @@ -24,14 +27,16 @@ use crate::{ #[derive(Default)] struct HookPair { - pre: bool, - post: bool, + pre: Vec, + post: Vec, } #[derive(Default)] pub(crate) struct GatewayPluginRuntime { manager: PluginManager, hooks: [HookPair; 3], + payload_writes: [[bool; 2]; 3], + auth_headers_writable: bool, } /// Pins the selected runtime and correlation context until the request finishes. @@ -39,6 +44,8 @@ pub(crate) struct GatewayPluginRuntime { pub(crate) struct CallState { runtime: Arc, context_table: PluginContextTable, + extensions: Extensions, + post: Vec, name: String, id: String, } @@ -46,38 +53,87 @@ pub(crate) struct CallState { static CORRELATION_ID: AtomicU64 = AtomicU64::new(1); impl GatewayPluginRuntime { + pub(crate) async fn from_published_config( + mut config: CpexConfig, + settings: &RuntimePluginSettings, + factories: &PluginFactoryRegistry, + ) -> Result { + if !matches!(settings.default_hook_policy.as_str(), "allow" | "deny") { + return Err(GatewayPluginRuntimeError::ConfigWrongFormat); + } + config.plugin_settings.plugin_timeout = settings.plugin_timeout; + config.plugins.retain(|plugin| plugin.mode != PluginMode::Disabled); + 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.") { + *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| { + 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", + _ => "content", + }; + settings + .hook_policies + .get(name) + .map_or(settings.default_hook_policy == "allow", |policy| policy.writable_fields.contains(field)) + }) + }); + Ok(runtime) + } + pub(crate) async fn from_config( config: CpexConfig, factories: &PluginFactoryRegistry, ) -> Result { validate_gateway_supported_config(&config)?; - let hooks = Operation::ALL.map(|operation| { - let [pre, post] = operation.hooks().map(|name| declares(&config, name)); - HookPair { pre, post } - }); + let names = config.plugins.iter().map(|plugin| plugin.name.clone()).collect::>(); let manager = PluginManager::from_config(config, factories) .map_err(|source| GatewayPluginRuntimeError::Configuration { hook: "config", source })?; manager.initialize().await.map_err(|source| GatewayPluginRuntimeError::Initialization { source })?; - Ok(Self { manager, hooks }) + let mut entries = names.iter().flat_map(|name| manager.find_plugin_entries(name)).collect::>(); + // invoke_entries expects priority order; preserve config order for ties. + entries.sort_by_key(|(_, entry)| entry.plugin_ref.priority()); + let hooks = Operation::ALL.map(|operation| { + let [pre, post] = operation + .hooks() + .map(|name| entries.iter().filter(|(hook, _)| hook == name).map(|(_, entry)| entry.clone()).collect()); + HookPair { pre, post } + }); + Ok(Self { manager, hooks, payload_writes: [[true; 2]; 3], auth_headers_writable: false }) } - #[instrument(name = "cmf_plugin_before", level = "info", skip(self, payload, update))] + #[instrument(name = "cmf_plugin_before", level = "info", skip(self, target, payload, update, extensions))] pub(crate) async fn before( self: &Arc, - operation: Operation, - name: &str, + target: (Operation, &str), + mut extensions: Extensions, payload: impl FnOnce(&str) -> MessagePayload, update: impl FnOnce(&MessagePayload, &str) -> Result, ) -> Result<(U, Option), ErrorData> { + let (operation, name) = target; let hooks = &self.hooks[operation as usize]; - if !hooks.pre && !hooks.post { + let pre = matching_entries(&hooks.pre, &extensions); + let post = matching_entries(&hooks.post, &extensions); + if pre.is_empty() && post.is_empty() { return Ok((U::default(), None)); } let id = format!("{}-{}", operation.id_prefix(), CORRELATION_ID.fetch_add(1, Ordering::Relaxed)); - let (update, context_table) = if hooks.pre { - let result = self.invoke(operation.hooks()[0], payload(&id), None).await; + let (update, context_table) = if pre.is_empty() { + (U::default(), PluginContextTable::default()) + } else { + let result = self.invoke((operation.hooks()[0], &pre), payload(&id), extensions.clone(), None).await; if result.is_denied() { return Err(plugin_denied_error(operation.subject(), result)); } @@ -85,24 +141,42 @@ impl GatewayPluginRuntime { Some(payload) => update(payload, &id)?, None => U::default(), }; + if let Some(modified) = result.modified_extensions { + extensions = modified; + } (update, result.context_table) - } else { - (U::default(), PluginContextTable::default()) }; - let state = - hooks.post.then(|| CallState { runtime: Arc::clone(self), context_table, name: name.to_owned(), id }); + let state = (!post.is_empty()).then(|| CallState { + runtime: Arc::clone(self), + context_table, + extensions, + post, + name: name.to_owned(), + id, + }); Ok((update, state)) } - #[instrument(name = "cmf_plugin_invoke", level = "info", skip(self, payload, context_table))] + #[instrument(name = "cmf_plugin_invoke", level = "info", skip_all)] async fn invoke( &self, - hook: &'static str, + invocation: (&'static str, &[HookEntry]), payload: MessagePayload, + extensions: Extensions, context_table: Option, ) -> 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_named::(hook, payload, Extensions::default(), context_table).await; + self.manager.invoke_entries::(&entries, payload, extensions, context_table).await; for error in &result.errors { tracing::warn!( hook, @@ -128,9 +202,20 @@ impl CallState { pub(crate) async fn invoke(&mut self, response: &T) -> Result { let payload = response.to_payload(&self.name, &self.id)?; - let result = self.runtime.invoke(T::OPERATION.hooks()[1], payload, Some(self.context_table.clone())).await; + 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(); + } } Ok(result) } @@ -157,8 +242,39 @@ impl Drop for GatewayPluginRuntime { } } -fn declares(config: &CpexConfig, hook_name: &str) -> bool { - config.plugins.iter().any(|plugin| plugin.hooks.iter().any(|hook| hook == hook_name)) +fn matching_entries(entries: &[HookEntry], extensions: &Extensions) -> Vec { + let meta = extensions.meta.as_deref(); + let name = meta.and_then(|meta| meta.entity_name.as_deref()); + let kind = meta.and_then(|meta| meta.entity_type.as_deref()); + let mcp = extensions.mcp.as_deref(); + let context = MatchContext { + server_id: mcp.and_then(|mcp| { + mcp.tool + .as_ref() + .and_then(|tool| tool.server_id.as_deref()) + .or_else(|| mcp.prompt.as_ref().and_then(|prompt| prompt.server_id.as_deref())) + .or_else(|| mcp.resource.as_ref().and_then(|resource| resource.server_id.as_deref())) + }), + tenant_id: meta.and_then(|meta| meta.scope.as_deref()), + tool: (kind == Some("tool")).then_some(name).flatten(), + prompt: (kind == Some("prompt")).then_some(name).flatten(), + resource: (kind == Some("resource")).then_some(name).flatten(), + user: extensions + .security + .as_ref() + .and_then(|security| security.subject.as_ref()) + .and_then(|subject| subject.id.as_deref()), + content_type: extensions.http.as_ref().and_then(|http| http.get_request_header("content-type")), + agent: extensions.agent.as_ref().and_then(|agent| agent.agent_id.as_deref()), + }; + entries + .iter() + .filter(|entry| { + let config = entry.plugin_ref.trusted_config(); + config.conditions.is_empty() || config.conditions.iter().any(|condition| condition.matches(&context)) + }) + .cloned() + .collect() } fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayPluginRuntimeError> { @@ -173,10 +289,6 @@ fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayP } for plugin in &config.plugins { - if !plugin.conditions.is_empty() { - return Err(GatewayPluginRuntimeError::ConfigUnsupported); - } - if plugin.hooks.iter().any(|hook| supported_cmf_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 7a8e13a4..e85a75b4 100644 --- a/crates/contextforge-data-plane-cpex/src/tools/mod.rs +++ b/crates/contextforge-data-plane-cpex/src/tools/mod.rs @@ -10,7 +10,7 @@ use serde_json::{Map, Value}; use tokio::sync::Mutex; use crate::{ - ArgumentsUpdate, GatewayPluginRuntimeHandle, PreHookResult, + ArgumentsUpdate, GatewayPluginRuntimeHandle, PluginRequestContext, PreHookResult, cmf::{CmfResponse, Operation, message_payload}, runtime::CallState, }; @@ -113,12 +113,14 @@ impl GatewayPluginRuntimeHandle { request: &CallToolRequestParams, tool_name: &str, backend_name: &str, + context: PluginRequestContext, ) -> Result, ErrorData> { let (arguments, state) = self .current()? + .tool(context.tool.as_ref())? .before( - Operation::Tool, - tool_name, + (Operation::Tool, tool_name), + context.extensions, |id| tool_call_payload(request, tool_name, backend_name, id), |payload, _| { let arguments = tool_call_arguments(payload).ok_or_else(|| { diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 5b151020..f13fe28f 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -14,6 +14,7 @@ repository.workspace = true [dependencies] contextforge-data-plane-apis.workspace = true contextforge-data-plane-cpex.workspace = true +cpex.workspace = true rmcp.workspace = true serde.workspace = true serde_json.workspace = true @@ -59,7 +60,6 @@ with_tools = ["axum/json", "axum/query"] [dev-dependencies] axum = { workspace = true, features = ["json"] } opentelemetry_sdk.workspace = true -cpex.workspace = true rmcp = { workspace = true, features = ["macros"] } test-log = "0.2.21" axum-server = { version = "0.8.0", features = ["tls-rustls"] } diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs index c38b376e..7fa6efb9 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs @@ -14,6 +14,12 @@ pub struct AuthorizedPrincipal { scopes: Vec, } +impl AuthorizedPrincipal { + pub fn tenant_id(&self) -> &str { + &self.tenant_id + } +} + impl<'a> From<&'a AuthorizedPrincipal> for User<'a> { fn from(value: &'a AuthorizedPrincipal) -> Self { Self::new(&value.user_id) diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/initialization.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/initialization.rs index f2d3fec2..ad3a33a0 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/initialization.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/initialization.rs @@ -171,6 +171,7 @@ mod tests { fn backend(passthrough: &[&str], add: &[(&str, &str)], remove: &[&str]) -> BackendMCPGateway { BackendMCPGateway { + tool_policy_contexts: std::collections::HashMap::new(), name: "b".into(), url: "https://upstream.example/mcp".parse().unwrap(), mcp_protocol_version: rmcp::model::ProtocolVersion::V_2026_07_28, diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/mod.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/mod.rs index 9f4f6675..88d1c93e 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/mod.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/mod.rs @@ -1,5 +1,6 @@ mod completion; mod initialization; +mod plugin_context; mod prompts; mod resources; mod tools; 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 new file mode 100644 index 00000000..fe7259d6 --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/plugin_context.rs @@ -0,0 +1,135 @@ +use std::{collections::HashSet, sync::Arc}; + +use contextforge_data_plane_apis::{ + User, + user_store::{BackendMCPGateway, UserConfig}, +}; +use contextforge_data_plane_cpex::PluginRequestContext; +use cpex::cpex_core::extensions::{ + Extensions, HttpExtension, MCPExtension, MetaExtension, PromptMetadata, RequestExtension, ResourceMetadata, + SecurityExtension, SubjectExtension, SubjectType, 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, +}; + +/// Construct plugin input only after authorization and route resolution. Nothing +/// in MCP arguments or client metadata participates in identity or policy lookup. +pub(super) fn request_context( + cx: &RequestContext, + kind: &str, + name: &str, + backend_name: &str, + backend: &BackendMCPGateway, +) -> Result { + let parts = cx + .extensions + .get::() + .ok_or_else(|| ErrorData::internal_error("Plugin request context is missing", None))?; + let principal = parts + .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 + .extensions + .get::() + .ok_or_else(|| ErrorData::internal_error("Plugin user configuration 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()), + entity_name: Some(canonical_name.to_owned()), + scope: Some(tenant.to_owned()), + ..Default::default() + }; + if let Some(tool) = &tool { + meta.properties.insert("tool_id".to_owned(), tool.id.clone()); + meta.properties.insert("context_id".to_owned(), tool.context_id.clone()); + } + meta.properties.insert("backend_id".to_owned(), backend_name.to_owned()); + if let Some(server_id) = &server_id { + meta.properties.insert("virtual_server_id".to_owned(), server_id.clone()); + } + let mcp = match kind { + "tool" => MCPExtension { + tool: Some(ToolMetadata { + name: canonical_name.to_owned(), + server_id: Some(backend_name.to_owned()), + namespace: Some(backend_name.to_owned()), + input_schema: backend.tool_schemas.get(name).cloned().map(Value::Object), + ..Default::default() + }), + ..Default::default() + }, + "prompt" => MCPExtension { + prompt: Some(PromptMetadata { name: name.to_owned(), server_id, ..Default::default() }), + ..Default::default() + }, + _ => MCPExtension { + resource: Some(ResourceMetadata { uri: name.to_owned(), server_id, ..Default::default() }), + ..Default::default() + }, + }; + let trace = tracing::Span::current().context(); + let span = trace.span(); + let span_context = span.span_context(); + Ok(PluginRequestContext { + 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() + }, + }) +} + +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/gateway/mcp_service/prompts.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/prompts.rs index 306139f5..b28a7254 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/prompts.rs @@ -37,7 +37,14 @@ pub(super) async fn get_prompt( data: None, })?; let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { - plugin_runtime.before_get_prompt(&request, &prompt_name, &backend_name).await? + plugin_runtime + .before_get_prompt( + &request, + &prompt_name, + &backend_name, + super::plugin_context::request_context(&cx, "prompt", &prompt_name, &backend_name, backend)?, + ) + .await? } else { PreHookResult::default() }; diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/resources.rs index 36ab5497..853bf5c3 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/resources.rs @@ -31,7 +31,18 @@ pub(super) async fn read_resource( }; let resource_hook = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { - Some(plugin_runtime.before_read_resource(&route.upstream_name).await?) + let backend = virtual_host + .backends + .get(&route.backend_name) + .ok_or_else(|| ErrorData::invalid_params("Routing problem... backend not found", None))?; + let context = super::plugin_context::request_context( + &cx, + "resource", + &route.upstream_name, + &route.backend_name, + backend, + )?; + Some(plugin_runtime.before_read_resource(&route.upstream_name, context).await?) } else { None }; diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/tools.rs index b3471ba4..18c5a72b 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_service/tools.rs @@ -54,7 +54,14 @@ pub(super) async fn call_tool( } let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { - plugin_runtime.before_tool_call(&request, &tool_name, &backend_name).await? + plugin_runtime + .before_tool_call( + &request, + &tool_name, + &backend_name, + super::plugin_context::request_context(&cx, "tool", &tool_name, &backend_name, backend)?, + ) + .await? } else { PreHookResult::default() }; diff --git a/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs b/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs index a0e0262d..18da60a4 100644 --- a/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs +++ b/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs @@ -54,6 +54,7 @@ mod tests { fn user_config_with_virtual_host(virtual_host_id: &str) -> UserConfig { UserConfig { + user_email: None, virtual_hosts: HashMap::from([( virtual_host_id.to_owned(), VirtualHost { diff --git a/crates/contextforge-data-plane-lib/src/user_config_store/redis_config_store.rs b/crates/contextforge-data-plane-lib/src/user_config_store/redis_config_store.rs index 7c07e0cb..880378f6 100644 --- a/crates/contextforge-data-plane-lib/src/user_config_store/redis_config_store.rs +++ b/crates/contextforge-data-plane-lib/src/user_config_store/redis_config_store.rs @@ -185,7 +185,7 @@ mod tests { use super::*; fn empty_config() -> UserConfig { - UserConfig { virtual_hosts: HashMap::new() } + UserConfig { user_email: None, virtual_hosts: HashMap::new() } } fn instant_ago(duration: Duration) -> Instant { diff --git a/crates/contextforge-data-plane-lib/tests/gateway.rs b/crates/contextforge-data-plane-lib/tests/gateway.rs index 0abd8d15..92aa36bd 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/plugin_context.rs"] +mod plugin_context; #[path = "gateway/plugins.rs"] mod plugins; #[path = "gateway/prompts.rs"] diff --git a/crates/contextforge-data-plane-lib/tests/gateway/future_contracts/pagination.rs b/crates/contextforge-data-plane-lib/tests/gateway/future_contracts/pagination.rs index bab5244c..364ad18e 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/future_contracts/pagination.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/future_contracts/pagination.rs @@ -34,6 +34,7 @@ async fn start_paginating_gateway(backend_count: usize) -> Result Result BackendMCPGateway { BackendMCPGateway { + tool_policy_contexts: std::collections::HashMap::new(), name: backend_id.to_owned(), url, mcp_protocol_version: protocol_version, diff --git a/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs b/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs new file mode 100644 index 00000000..9c27c3ce --- /dev/null +++ b/crates/contextforge-data-plane-lib/tests/gateway/plugin_context.rs @@ -0,0 +1,398 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +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_lib::UserConfigStore; +use cpex::cpex_core::{ + cmf::{CmfHook, MessagePayload}, + context::PluginContext, + hooks::{Extensions, HookHandler, PluginResult}, + plugin::{Plugin, PluginConfig}, +}; +use rmcp::model::{CallToolRequestParams, GetPromptRequestParams, ReadResourceRequestParams}; +use serde_json::{Value, json}; + +use crate::harness::{RunningGateway, TEST_USER_ID, start_gateway, sum_request}; + +type Seen = Arc>>; + +struct ContextPlugin { + config: PluginConfig, + seen: Seen, +} + +#[async_trait] +impl Plugin for ContextPlugin { + fn config(&self) -> &PluginConfig { + &self.config + } +} + +impl HookHandler for ContextPlugin { + async fn handle( + &self, + _: &MessagePayload, + extensions: &Extensions, + ctx: &mut PluginContext, + ) -> PluginResult { + if let Some(delay) = + self.config.config.as_ref().and_then(|config| config.get("delay_ms")).and_then(Value::as_u64) + { + tokio::time::sleep(Duration::from_millis(delay)).await; + } + let subject = extensions.security.as_ref().and_then(|security| security.subject.as_ref()); + let count = ctx.get_local("calls").and_then(Value::as_u64).unwrap_or(0); + self.seen.lock().expect("observations").push(json!({ + "plugin": self.config.name, + "subject": subject.and_then(|subject| subject.id.as_deref()), + "tenant": subject.and_then(|subject| subject.claims.get("tenant_id")), + "headers": extensions.http.is_some(), + "method": extensions.http.as_ref().and_then(|http| http.method.as_deref()), + "content_type": extensions.http.as_ref().and_then(|http| http.get_request_header("content-type")), + "policy_header": extensions.http.as_ref().is_some_and(|http| http.get_request_header("x-policy") == Some("checked")), + "auth_header_changed": extensions.http.as_ref().is_some_and(|http| http.get_request_header("authorization") == Some("changed")), + "request": extensions.request.as_ref().and_then(|request| request.request_id.as_ref()), + "meta": extensions.meta, + "mcp": extensions.mcp, + "count": count, + "label": extensions.security.as_ref().is_some_and(|security| security.has_label("pre-checked")), + })); + ctx.set_local("calls", json!(count + 1)); + ctx.set_global("user", json!("attacker")); + ctx.set_global("context_id", json!("other-team")); + let mut updated = extensions.cow_copy(); + if self.config.config.as_ref().and_then(|config| config.get("headers")) == Some(&json!(true)) { + let mut http = extensions.http.as_deref().cloned().unwrap_or_default(); + http.set_request_header("x-policy", "checked"); + http.set_request_header("authorization", "changed"); + updated.http = Some(cpex::cpex_core::extensions::Guarded::new(http)); + } + if self.config.config.as_ref().and_then(|config| config.get("spoof")) == Some(&json!(true)) { + if let Some(subject) = updated.security.as_mut().and_then(|security| security.subject.as_mut()) { + subject.id = Some("attacker".to_owned()); + } + } else if let Some(security) = &mut updated.security { + security.add_label("pre-checked"); + } + PluginResult::modify_extensions(updated) + } +} + +#[tokio::test] +async fn extension_header_writes_require_capability_and_auth_override_permission() { + let (runtime, seen) = runtime(document(&[], &[])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + for (writable, auth_override) in [(false, false), (true, false), (true, true)] { + let mut configured = plugin("headers", &["tool_pre_invoke", "tool_post_invoke"], true); + configured["config"] = json!({"headers": true}); + if writable { + configured["capabilities"].as_array_mut().expect("caps").push(json!("write_headers")); + } + let mut config = document(&[], &[configured]); + config.settings.plugins_can_override_auth_headers = auth_override; + runtime.apply_document(config).await.expect("publish header policy"); + seen.lock().expect("observations").clear(); + client.call_tool(sum_request("sum", 1, 2)).await.expect("call"); + let observations = seen.lock().expect("observations").clone(); + assert_eq!(observations.len(), 2); + assert_eq!(observations[0]["policy_header"], false); + assert_eq!(observations[1]["policy_header"], writable); + assert_eq!(observations[1]["auth_header_changed"], auth_override); + assert_eq!(observations[1]["subject"], "operator@example.com"); + } + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +fn plugin(name: &str, hooks: &[&str], trusted: bool) -> Value { + json!({ + "name": name, "kind": "context-test", "hooks": hooks, + "capabilities": if trusted { + vec!["read_subject", "read_claims", "read_teams", "read_permissions", "read_roles", "read_headers", "read_labels", "append_labels"] + } else { vec![] }, + }) +} + +fn document(global: &[Value], scoped: &[Value]) -> RuntimePluginConfigDocument { + serde_json::from_value(json!({ + "enabled": true, "global": {"plugins": global}, + "contexts": {"team1::gateway_sum": {"plugins": scoped}} + })) + .expect("published config") +} + +async fn runtime(document: RuntimePluginConfigDocument) -> (Arc, Seen) { + let seen = Seen::default(); + let captured = Arc::clone(&seen); + let mut runtime = CpexRuntimeRegistry::default(); + runtime + .register_factory( + "context-test", + Box::new(CmfPluginFactory::new(move |config| ContextPlugin { config, seen: Arc::clone(&captured) })), + ) + .expect("factory"); + runtime.apply_document(document).await.expect("published policy applies"); + (Arc::new(runtime), seen) +} + +async fn gateway(runtime: &Arc) -> RunningGateway { + let gateway = start_gateway(TEST_USER_ID, true, Arc::clone(runtime)).await; + let key = User::new(TEST_USER_ID); + let mut config = gateway.user_store.get_config(&key).await.expect("user config"); + config.user_email = Some("operator@example.com".to_owned()); + for host in config.virtual_hosts.values_mut() { + host.prompts.insert( + format!("{}-review", gateway.backend_name), + host.prompts.get("review").expect("prompt route").clone(), + ); + let backend = host.backends.get_mut(&gateway.backend_name).expect("backend"); + backend.tool_policy_contexts = HashMap::from([( + "sum".to_owned(), + ToolPolicyContext { + id: "tool-id".to_owned(), + name: "gateway_sum".to_owned(), + team_id: Some("team1".to_owned()), + context_id: "team1::gateway_sum".to_owned(), + }, + )]); + } + gateway.user_store.set_config(&key, &config).await.expect("publish context"); + gateway +} + +#[tokio::test] +async fn direct_and_aliased_tools_select_published_policy_and_preserve_context() { + let hooks = ["tool_pre_invoke", "tool_post_invoke"]; + let mut scoped = plugin("scoped", &hooks, true); + scoped["conditions"] = + json!([{"tools": ["gateway_sum"], "tenant_ids": ["team1"], "user_patterns": ["*@example.com"]}]); + let (runtime, seen) = runtime(document(&[plugin("wrong-global", &hooks, true)], &[scoped])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + for name in ["sum".to_owned(), format!("{}-sum", gateway.backend_name)] { + let mut request = sum_request(&name, 1, 2); + request.arguments.as_mut().expect("args").insert("user".to_owned(), json!("attacker")); + request.arguments.as_mut().expect("args").insert("context_id".to_owned(), json!("other-team")); + client.call_tool(request).await.expect("authorized call"); + } + let observations = seen.lock().expect("observations").clone(); + assert_eq!(observations.len(), 4); + for pair in observations.as_chunks::<2>().0 { + assert_eq!(pair[0]["plugin"], "scoped"); + assert_eq!(pair[1]["subject"], "operator@example.com"); + assert_eq!(pair[1]["tenant"], "test_tenant"); + assert_eq!(pair[1]["meta"]["properties"]["tool_id"], "tool-id"); + assert_eq!(pair[1]["mcp"]["tool"]["name"], "gateway_sum"); + assert_eq!(pair[0]["count"], 0); + assert_eq!(pair[1]["count"], 1); + assert_eq!(pair[0]["label"], false); + assert_eq!(pair[1]["label"], true); + assert_eq!(pair[0]["request"], pair[1]["request"]); + } + assert_ne!(observations[0]["request"], observations[2]["request"]); + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +#[tokio::test] +async fn post_only_stream_context_stays_pinned_across_reload() { + use contextforge_data_plane_cpex::PluginRequestContext; + use cpex::cpex_core::extensions::{RequestExtension, SecurityExtension, SubjectExtension}; + use rmcp::model::{CallToolResult, NumberOrString, ProgressNotificationParam, ProgressToken}; + + let (runtime, seen) = runtime(document(&[], &[plugin("stream", &["tool_post_invoke"], true)])).await; + let context = PluginRequestContext { + tool: Some(ToolPolicyContext { + id: "id".to_owned(), + name: "gateway_sum".to_owned(), + team_id: Some("team1".to_owned()), + context_id: "team1::gateway_sum".to_owned(), + }), + extensions: Extensions { + request: Some(Arc::new(RequestExtension { request_id: Some("request".to_owned()), ..Default::default() })), + security: Some(Arc::new(SecurityExtension { + subject: Some(SubjectExtension { id: Some("verified".to_owned()), ..Default::default() }), + ..Default::default() + })), + ..Default::default() + }, + }; + let state = runtime + .handle() + .before_tool_call(&CallToolRequestParams::new("sum"), "sum", "backend", context) + .await + .expect("post-only request starts") + .state + .expect("post state"); + runtime.apply_config(None).await.expect("disable new requests"); + for count in 1..=2 { + state + .after_stream_event(ProgressNotificationParam::new( + ProgressToken(NumberOrString::String("progress".into())), + f64::from(count), + )) + .await + .expect("progress") + .expect("event allowed"); + } + state.after_tool_call(CallToolResult::success(vec![])).await.expect("final response"); + let observations = seen.lock().expect("observations"); + assert_eq!(observations.len(), 3); + for (count, observation) in observations.iter().enumerate() { + assert_eq!(observation["plugin"], "stream"); + assert_eq!(observation["subject"], "verified"); + assert_eq!(observation["count"], count); + assert_eq!(observation["label"], count > 0); + } +} + +#[tokio::test] +async fn capabilities_gate_identity_and_headers_and_plugin_cannot_replace_identity() { + let hooks = ["tool_pre_invoke", "tool_post_invoke"]; + let mut spoof = plugin("spoof", &hooks, true); + spoof["config"] = json!({"spoof": true}); + spoof["priority"] = json!(1); + let mut unprivileged = plugin("unprivileged", &hooks, false); + unprivileged["priority"] = json!(1); + let mut observer = plugin("observer", &hooks, true); + observer["priority"] = json!(90); + let (runtime, seen) = runtime(document(&[], &[observer, spoof, unprivileged])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + client.call_tool(sum_request("sum", 1, 2)).await.expect("call"); + let observations = seen.lock().expect("observations").clone(); + assert_eq!(observations.len(), 6); + assert_eq!( + observations.iter().map(|event| event["plugin"].as_str().expect("plugin")).collect::>(), + ["spoof", "unprivileged", "observer", "spoof", "unprivileged", "observer"], + "lower priority runs first and equal priorities preserve configuration order in both phases" + ); + for observation in observations { + if observation["plugin"] == "unprivileged" { + assert!(observation["subject"].is_null()); + assert_eq!(observation["headers"], false); + } else { + assert_eq!(observation["subject"], "operator@example.com"); + assert_eq!(observation["headers"], true); + } + } + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +#[tokio::test] +async fn missing_tool_or_policy_context_fails_before_backend_and_disabled_policy_allows() { + let (runtime, _) = runtime(document(&[], &[])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + let error = client.call_tool(CallToolRequestParams::new("reflect_text")).await.expect_err("missing tool context"); + assert!(error.to_string().contains("tool context is missing")); + let mut missing = document(&[], &[]); + missing.contexts.clear(); + runtime.apply_document(missing).await.expect("publish missing scope"); + let error = client.call_tool(sum_request("sum", 1, 2)).await.expect_err("missing policy context"); + assert!(error.to_string().contains("policy context is missing")); + assert!(gateway.backend_state.calls.lock().expect("calls").is_empty()); + runtime + .apply_document( + serde_json::from_value(json!({"enabled":false,"global":null,"contexts":{}})).expect("disabled config"), + ) + .await + .expect("disable"); + client.call_tool(sum_request("sum", 1, 2)).await.expect("explicitly disabled plugins"); + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +#[tokio::test] +async fn prompt_and_resource_post_only_hooks_use_global_context_for_direct_and_aliased_requests() { + let global = plugin("global", &["prompt_post_fetch", "resource_post_fetch"], true); + let (runtime, seen) = runtime(document(&[global], &[plugin("wrong-scope", &["prompt_post_fetch"], true)])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + for name in ["review".to_owned(), format!("{}-review", gateway.backend_name)] { + client + .get_prompt( + GetPromptRequestParams::new(name) + .with_arguments(serde_json::Map::from_iter([("topic".to_owned(), json!("weather"))])), + ) + .await + .expect("prompt"); + } + for uri in ["file:///password.env".to_owned(), format!("{}-file:///password.env", gateway.backend_name)] { + client.read_resource(ReadResourceRequestParams::new(uri)).await.expect("resource"); + } + let observations = seen.lock().expect("observations").clone(); + assert_eq!(observations.len(), 4); + for observation in observations { + assert_eq!(observation["plugin"], "global"); + assert_eq!(observation["subject"], "operator@example.com"); + assert_eq!(observation["count"], 0); + assert!(observation["request"].is_string()); + } + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +#[tokio::test] +async fn write_only_mcp_hooks_preserve_hidden_headers_and_transport_metadata() { + let hooks = ["tool_pre_invoke", "tool_post_invoke"]; + let mut writer = plugin("writer", &hooks, false); + writer["capabilities"] = json!(["write_headers"]); + writer["config"] = json!({"headers":true}); + let (runtime, seen) = runtime(document(&[], &[writer, plugin("reader", &hooks, true)])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + client.call_tool(sum_request("sum", 1, 2)).await.expect("call"); + { + let observations = seen.lock().expect("observations"); + assert_eq!(observations.len(), 4); + for pair in observations.as_chunks::<2>().0 { + assert_eq!(pair[0]["headers"], false); + assert_eq!(pair[1]["policy_header"], true); + assert_eq!(pair[1]["method"], "POST"); + assert_eq!(pair[1]["content_type"], "application/json"); + assert_eq!(pair[1]["auth_header_changed"], false); + } + } + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} + +#[tokio::test] +async fn published_timeout_controls_each_mcp_plugin_invocation() { + let (runtime, seen) = runtime(document(&[], &[])).await; + let gateway = gateway(&runtime).await; + let client = gateway.connect(TEST_USER_ID).await; + for hook in ["tool_pre_invoke", "tool_post_invoke"] { + for seconds in [1, 3] { + let mut slow = plugin("slow", &[hook], false); + slow["config"] = json!({"delay_ms":1500}); + slow["on_error"] = json!("fail"); + let mut config = document(&[], &[slow]); + config.settings.plugin_timeout = seconds; + runtime.apply_document(config).await.expect("publish timeout"); + seen.lock().expect("observations").clear(); + let result = tokio::time::timeout(Duration::from_secs(5), client.call_tool(sum_request("sum", 1, 2))) + .await + .expect("hook execution is bounded"); + if seconds == 1 { + assert!(result.expect_err("slow plugin must time out").to_string().contains("Plugin denied tool call")); + assert!(seen.lock().expect("observations").is_empty(), "timed out handler did not complete"); + } else { + result.expect("same plugin completes with the longer published timeout"); + assert_eq!(seen.lock().expect("observations").len(), 1); + } + } + } + client.cancel().await.expect("client closes"); + gateway.shutdown().await.expect("gateway stops"); +} diff --git a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs index 7ed9868b..2dc0389e 100644 --- a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs +++ b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs @@ -15,7 +15,7 @@ use std::{ use contextforge_data_plane_apis::{ User, - runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_KEY, RUNTIME_PLUGIN_CONFIG_VERSION}, + runtime_plugin_config::RUNTIME_PLUGIN_CONFIG_KEY, user_store::{BackendMCPGateway, UserConfig, VirtualHost}, }; use http::{HeaderMap, HeaderValue}; @@ -373,12 +373,27 @@ async fn write_redis_config(redis_port: u16, backend: &RunningBackend) { let key = rmp_serde::encode::to_vec(&User::new(TEST_USER_ID)).expect("user key encodes"); let config = UserConfig { + user_email: None, virtual_hosts: HashMap::from([( TEST_VIRTUAL_HOST_ID.to_owned(), VirtualHost { backends: HashMap::from([( "backend".to_owned(), BackendMCPGateway { + tool_policy_contexts: ["sum", "reflect_text"] + .into_iter() + .map(|name| { + ( + name.to_owned(), + contextforge_data_plane_apis::user_store::ToolPolicyContext { + id: name.to_owned(), + name: name.to_owned(), + team_id: None, + context_id: TEST_VIRTUAL_HOST_ID.to_owned(), + }, + ) + }) + .collect(), name: "backend".to_owned(), url: backend.url.parse().expect("backend URL parses"), mcp_protocol_version: rmcp::model::ProtocolVersion::V_2026_07_28, @@ -414,17 +429,19 @@ async fn write_runtime_plugin_config(redis_port: u16, plugin_config: Value) { .get_connection_manager_with_config(ConnectionManagerConfig::default()) .await .expect("redis connection opens"); - let document = json!({ - "version": RUNTIME_PLUGIN_CONFIG_VERSION, - "cpex": { + let mut document = json!({ + "enabled": true, + "contexts": {}, + "global": { "plugins": [{ "name": "secrets-detection", - "kind": "validator/secrets-detection", + "kind": "cpex_secrets_detection.SecretsDetectionPlugin", "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], "config": plugin_config, }] } }); + document["contexts"][TEST_VIRTUAL_HOST_ID] = document["global"].clone(); redis::cmd("SET") .arg(RUNTIME_PLUGIN_CONFIG_KEY) .arg(serde_json::to_vec(&document).expect("runtime plugin config serializes")) diff --git a/crates/plugins/cpex-secrets-detection/README.md b/crates/plugins/cpex-secrets-detection/README.md index d30c709d..a5265e2c 100644 --- a/crates/plugins/cpex-secrets-detection/README.md +++ b/crates/plugins/cpex-secrets-detection/README.md @@ -15,12 +15,13 @@ Example config: ```json { - "version": 1, - "cpex": { + "enabled": true, + "contexts": {}, + "global": { "plugins": [ { "name": "secrets-detection", - "kind": "validator/secrets-detection", + "kind": "cpex_secrets_detection.SecretsDetectionPlugin", "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke", "cmf.resource_post_fetch"], "config": { "redact": true, @@ -33,6 +34,11 @@ Example config: } ``` +For tool calls, the publisher also places the resolved plugin configuration in +`contexts[context_id]` and publishes the corresponding `tool_policy_contexts` +entry on the routed backend. An enabled document without the tool's policy +context fails the call. Prompts and resources select `global`. + The dataplane integration supports: - `cmf.tool_pre_invoke`: scans tool arguments before the backend receives them. @@ -53,7 +59,7 @@ content. Depending on config, it can: The CPEX plugin kind is: ```text -validator/secrets-detection +cpex_secrets_detection.SecretsDetectionPlugin ``` ## Known CPEX 0.2.2 Gaps diff --git a/crates/plugins/cpex-secrets-detection/src/lib.rs b/crates/plugins/cpex-secrets-detection/src/lib.rs index 32abf830..9d40f989 100644 --- a/crates/plugins/cpex-secrets-detection/src/lib.rs +++ b/crates/plugins/cpex-secrets-detection/src/lib.rs @@ -21,7 +21,7 @@ pub mod config; pub mod patterns; pub mod scanner; -pub const KIND: &str = "validator/secrets-detection"; +pub const KIND: &str = "cpex_secrets_detection.SecretsDetectionPlugin"; const VIOLATION_CODE: &str = "SECRETS_DETECTED"; const MAX_SECRET_TYPES: usize = 32; @@ -658,7 +658,7 @@ mod tests { format!( r#"plugins: - name: secrets-detection - kind: validator/secrets-detection + kind: cpex_secrets_detection.SecretsDetectionPlugin hooks: ["{hook}"] mode: sequential {config}"# diff --git a/crates/plugins/cpex-secrets-detection/tests/plugin_manager.rs b/crates/plugins/cpex-secrets-detection/tests/plugin_manager.rs index bda206c0..1ee7c0ab 100644 --- a/crates/plugins/cpex-secrets-detection/tests/plugin_manager.rs +++ b/crates/plugins/cpex-secrets-detection/tests/plugin_manager.rs @@ -321,7 +321,7 @@ fn plugin_yaml(hook: &str, config_block: &str) -> String { format!( r#"plugins: - name: secrets-detection - kind: validator/secrets-detection + kind: cpex_secrets_detection.SecretsDetectionPlugin hooks: ["{hook}"] mode: sequential {config}"# diff --git a/schemas/user_config.json b/schemas/user_config.json index 7df9c60a..27f4f1a0 100644 --- a/schemas/user_config.json +++ b/schemas/user_config.json @@ -8,6 +8,13 @@ "additionalProperties": { "$ref": "#/$defs/VirtualHost" } + }, + "user_email": { + "type": [ + "string", + "null" + ], + "default": null } }, "required": [ @@ -22,6 +29,34 @@ "additionalProperties": { "$ref": "#/$defs/BackendMCPGateway" } + }, + "tools": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ServiceRoute" + }, + "default": {} + }, + "resources": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ServiceRoute" + }, + "default": {} + }, + "resource_templates": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ServiceRoute" + }, + "default": {} + }, + "prompts": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ServiceRoute" + }, + "default": {} } }, "required": [ @@ -79,6 +114,14 @@ "additionalProperties": true }, "default": {} + }, + "tool_policy_contexts": { + "description": "Canonical tool identity and resolved policy key, indexed by upstream name.", + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ToolPolicyContext" + }, + "default": {} } }, "required": [ @@ -91,6 +134,46 @@ "ProtocolVersion": { "description": "Represents the MCP protocol version used for communication.\n\nThis ensures compatibility between clients and servers by specifying\nwhich version of the Model Context Protocol is being used.", "type": "string" - } + }, + "ToolPolicyContext": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "team_id": { + "type": [ + "string", + "null" + ] + }, + "context_id": { + "type": "string" + } + }, + "required": [ + "id", + "name", + "context_id" + ] + }, + "ServiceRoute": { + "type": "object", + "properties": { + "backend_name": { + "type": "string" + }, + "upstream_name": { + "type": "string" + } + }, + "required": [ + "backend_name", + "upstream_name" + ] + } } }