diff --git a/.secrets.baseline b/.secrets.baseline index 38a1c0a6..2cd6744a 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -3,7 +3,7 @@ "files": "(?x)(Cargo\\.lock$|\\.lock$)|^\\.secrets\\.baseline$|^.secrets.baseline$", "lines": null }, - "generated_at": "2026-09-03T17:12:20Z", + "generated_at": "2026-09-04T15:25:27Z", "plugins_used": [ { "name": "AWSKeyDetector" @@ -210,7 +210,7 @@ "hashed_secret": "86de8c52637ec530fe39b0a8471da9b8764d5242", "is_secret": false, "is_verified": false, - "line_number": 610, + "line_number": 609, "type": "AWS Access Key", "verified_result": null } diff --git a/Cargo.lock b/Cargo.lock index 0a2824f2..1ec9dda9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -582,6 +582,7 @@ version = "0.1.0" dependencies = [ "arc-swap", "async-trait", + "base64 0.22.1", "contextforge-data-plane-apis", "cpex", "redis", diff --git a/Cargo.toml b/Cargo.toml index 6a211c8e..4e474b9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ clap = { version = "4.5.60", features = ["derive", "env"] } thiserror = "2.0.18" rmp-serde = "1.3.1" async-trait = "0.1.89" +base64 = "0.22.1" reqwest = "0.13" jsonwebtoken = { version = "11.0.0", features = ["aws_lc_rs"] } rustls = { version = "0.23", features = ["ring"] } diff --git a/README.md b/README.md index 814356db..8dad5f3d 100644 --- a/README.md +++ b/README.md @@ -60,8 +60,8 @@ Activation requires all three pieces: - Runtime flag: `--runtime-plugins-enabled true` - Redis config key: `ContextForgeGatewayRuntimePluginConfig` -The plugin kind is `validator/secrets-detection`. The data plane currently -wires only `cmf.tool_pre_invoke` and `cmf.tool_post_invoke`. +The plugin kind is `validator/secrets-detection`. The dataplane wires CMF hooks +for tool calls, prompt fetches, and resource reads. Example run command: diff --git a/_context/wiki/architecture.md b/_context/wiki/architecture.md index 2a196b90..e36a668d 100644 --- a/_context/wiki/architecture.md +++ b/_context/wiki/architecture.md @@ -159,9 +159,11 @@ The binary sets `tikv_jemallocator` as the global allocator. jemalloc holds up b - `initialize` opens one backend transport per configured backend concurrently (`futures::future::join_all`); a failed backend degrades that backend only. - List methods fan out to all connected backends concurrently and merge. - Targeted calls (except `call_tool`) resolve exactly one backend service handle from `BackendTransports`. -- `call_tool` creates a fresh per-request backend connection via `connect_backend_for_request`, runs pre/post plugin hooks, executes the call, then explicitly closes the connection before returning. +- Targeted tool, prompt, and resource calls run configured pre/post plugin hooks after backend routing. `call_tool` creates a fresh per-request backend connection via `connect_backend_for_request`, then explicitly closes it before returning. - `call_tool` watches the downstream cancellation token and forwards a cancel to the backend if the client gives up first; backend progress notifications are forwarded downstream while the call is in flight. +Resource reads carry a concrete, request-owned hook state across backend I/O. It pins the runtime selected before the read, or records that no post hook was configured. Post processing consumes that state without type erasure, downcasts, or a second registry lookup. + ## Startup And Response Flow Startup sequence (`main.rs` → `Gateway::run_gateway`): @@ -185,7 +187,7 @@ Response unwind order (Tower layers execute outside-in, so unwind is inside-out) ```text backend response - -> response plugin hooks (call_tool only) + -> response plugin hooks (tool, prompt, and resource calls) -> merge / namespace / pass through -> virtual_host_config_layer response side -> user_config_store_layer response side @@ -232,7 +234,7 @@ Do not bury transport security decisions inside MCP method handlers. They belong ## Plugin Hook Expansion Requirements -Current supported hooks are intentionally narrow (`cmf.tool_pre_invoke`, `cmf.tool_post_invoke`). Before adding any new hook point, define all of the following: +Current supported hooks cover tool, prompt, and resource pre/post lifecycles. Before adding any new hook point, define all of the following: | Requirement | Why | | --- | --- | diff --git a/_context/wiki/config.md b/_context/wiki/config.md index 4489abbb..10d57bf2 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -185,8 +185,8 @@ RuntimePluginConfigDocument cpex: CpexConfig ``` -Supported: `cmf.tool_pre_invoke`, `cmf.tool_post_invoke`, `cmf.prompt_pre_fetch`, `cmf.prompt_post_fetch` only. -Rejected: routing-based selection, plugin dirs, global policies, resource and LLM hooks, plugin conditions. +Supported: tool, prompt, and resource pre/post CMF hooks. +Rejected: routing-based selection, plugin dirs, global policies, LLM hooks, plugin conditions. 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. @@ -212,7 +212,15 @@ Writing plugin edits back follows three rules: MCP prompt results carry no error flag, so a plugin setting `is_error` on the CMF prompt result is rejecting the prompt rather than describing it. The gateway turns that into an MCP error carrying the plugin's `error_message`, and the rendered content never reaches the client. This differs from tools, where `is_error` is a field on `CallToolResult` and is forwarded as a successful response. -Binary resource blobs reach plugins by URI and MIME type but not by content: CMF stores decoded bytes while MCP sends base64. A plugin can deny such a message; editing one fails the write-back. +Binary resources embedded in prompts reach plugins by URI and MIME type but not by content. A plugin can deny such a message; editing one fails the write-back. Resource-read hooks below have their own binary conversion. + +### Resource Read Hook Behavior + +For `resources/read`, the pre hook receives the canonical backend-local URI and may allow, deny or rewrite it. A rewritten URI must resolve unambiguously through the caller's published virtual-host resources before a backend connection is opened. Aliases for the same backend target do not create ambiguity. + +The post hook may replace each returned resource's text or binary content, URI and MIME type, including converting text to a blob or a blob to text. Existing MCP `_meta` is preserved. CMF-only envelope and descriptive fields do not restrict these changes. Each resource still needs a valid MCP content representation; binary resource reads are decoded for CPEX and re-encoded after edits, while unchanged blob bytes retain their original wire value. This resource path does not add prompt-wide payload validation. + +The pre call returns an opaque, concrete `ResourceHookState` consumed by the post call. It captures both the runtime and the decision to run or skip post hooks before backend I/O. A reload only affects subsequent requests, including when it enables or disables resource hooks. Callers cannot construct missing or mismatched active state, and requests without a post hook allocate no correlation state. ### Demo Plugin Workflow diff --git a/_context/wiki/testing.md b/_context/wiki/testing.md index d3cc32bb..79706795 100644 --- a/_context/wiki/testing.md +++ b/_context/wiki/testing.md @@ -37,7 +37,7 @@ backends without recompiling a shared support tree for every feature file. | Area | Covers | | --- | --- | | `gateway/{tools,prompts,resources,subscriptions}.rs` | Active routed operations and exact routing failures. | -| `gateway/plugins.rs` | Gateway-owned CPEX ordering, mutation, denial, progress, and prompt seams using deterministic recording plugins. Concrete plugin behavior stays in each plugin crate. | +| `gateway/plugins.rs` | Gateway-owned CPEX ordering, mutation, denial, progress, and prompt seams using deterministic recording plugins. Resource coverage includes direct and aliased URIs, text/blob conversion, canonical pre-hook URIs, published-target rewrites, rejection of unpublished targets, metadata preservation, and pre/post denial. Concrete plugin behavior stays in each plugin crate. | | `gateway/harness/` | Authentication, modern and compatibility clients, in-memory configuration, concrete mock backends, and owned server fixtures. | | `gateway/future_contracts/` | Deferred fanout, pagination, TLS, completions, subscriptions, and cancellation contracts. | @@ -56,7 +56,7 @@ cargo nextest list --locked --workspace --all-features --run-ignored only ``` The two binary E2E tests and `tests/conformance/` remain separate infrastructure -boundaries. Active in-process tests run with no Docker or Redis dependency. +boundaries. Active in-process tests run with no Docker or Redis dependency. Resource policy coverage belongs in this active harness, not in ignored binary tests or new legacy-client cases. Runtime unit tests verify that enabling or disabling hooks during a resource read preserves its original policy decision. Parameter-header integration tests verify that calls without a published tool schema skip local `Mcp-Param-*` validation and still reach the backend. Unit and diff --git a/crates/contextforge-data-plane-cpex/Cargo.toml b/crates/contextforge-data-plane-cpex/Cargo.toml index d74d69bd..6df8e734 100644 --- a/crates/contextforge-data-plane-cpex/Cargo.toml +++ b/crates/contextforge-data-plane-cpex/Cargo.toml @@ -15,6 +15,7 @@ doctest = false [dependencies] arc-swap = "1.7" async-trait.workspace = true +base64.workspace = true contextforge-data-plane-apis.workspace = true cpex.workspace = true redis.workspace = true diff --git a/crates/contextforge-data-plane-cpex/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index eeb74c97..010606c8 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -1,12 +1,13 @@ use std::collections::HashMap; +use base64::{Engine as _, prelude::BASE64_STANDARD}; use cpex::cpex_core::cmf::{ AudioSource, ContentPart, ImageSource, Message, MessagePayload, PromptRequest, PromptResult, - Resource as CmfResource, ResourceReference, ResourceType, Role, ToolCall, ToolResult, + Resource as CmfResource, ResourceReference, ResourceType, Role, ToolCall, ToolResult, constants::SCHEMA_VERSION, }; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, GetPromptRequestParams, GetPromptResult, PromptMessage, - Resource as McpResource, ResourceContents, Role as McpRole, + ReadResourceResult, Resource as McpResource, ResourceContents, Role as McpRole, }; use serde_json::{Map, Value}; @@ -18,7 +19,7 @@ pub(crate) fn tool_call_payload( ) -> MessagePayload { MessagePayload { message: Message { - schema_version: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: Role::Assistant, content: vec![ContentPart::ToolCall { content: ToolCall { @@ -33,6 +34,108 @@ pub(crate) fn tool_call_payload( } } +pub(crate) fn resource_request_payload(resource_uri: &str, resource_request_id: &str) -> MessagePayload { + MessagePayload { + message: Message { + schema_version: SCHEMA_VERSION.to_owned(), + role: Role::User, + content: vec![ContentPart::ResourceRef { + content: ResourceReference { + resource_request_id: resource_request_id.to_owned(), + uri: resource_uri.to_owned(), + name: None, + resource_type: ResourceType::Uri, + range_start: None, + range_end: None, + selector: None, + }, + }], + channel: None, + }, + } +} + +pub(crate) fn resource_result_payload( + response: &ReadResourceResult, + resource_request_id: &str, +) -> Option { + let content = response + .contents + .iter() + .map(|content| { + cmf_resource_content(content, resource_request_id).map(|content| ContentPart::Resource { content }) + }) + .collect::>>()?; + Some(MessagePayload { + message: Message { schema_version: SCHEMA_VERSION.to_owned(), role: Role::Assistant, content, channel: None }, + }) +} + +fn cmf_resource_content(content: &ResourceContents, resource_request_id: &str) -> Option { + let (uri, mime_type, text, blob) = match content { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { + (uri.clone(), mime_type.clone(), Some(text.clone()), None) + }, + ResourceContents::BlobResourceContents { uri, mime_type, blob, .. } => { + (uri.clone(), mime_type.clone(), None, Some(BASE64_STANDARD.decode(blob).ok()?)) + }, + _ => return None, + }; + Some(CmfResource { + resource_request_id: resource_request_id.to_owned(), + uri, + resource_type: ResourceType::Uri, + content: text, + blob, + mime_type, + ..Default::default() + }) +} + +pub(crate) fn resource_result_response( + mut original: ReadResourceResult, + payload: &MessagePayload, +) -> Option { + // Resource post hooks replace each resource's content, not the read envelope. + if payload.message.content.len() != original.contents.len() { + return None; + } + for (original, modified) in original.contents.iter_mut().zip(&payload.message.content) { + let ContentPart::Resource { content } = modified else { return None }; + let meta = match original { + ResourceContents::TextResourceContents { meta, .. } + | ResourceContents::BlobResourceContents { meta, .. } => meta.clone(), + _ => return None, + }; + *original = match (&content.content, &content.blob) { + (Some(text), _) => ResourceContents::TextResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + text: text.clone(), + meta, + }, + (None, Some(bytes)) => { + let blob = match original { + ResourceContents::BlobResourceContents { blob, .. } + if BASE64_STANDARD.decode(blob.as_bytes()).ok().as_ref() == Some(bytes) => + { + blob.clone() + }, + _ => BASE64_STANDARD.encode(bytes), + }; + ResourceContents::BlobResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + blob, + meta, + } + }, + _ => return None, + }; + } + Some(original) +} + pub(crate) fn tool_result_payload(tool_name: &str, response: &CallToolResult, tool_call_id: &str) -> MessagePayload { tool_json_result_payload( tool_name, @@ -50,7 +153,7 @@ pub(crate) fn tool_json_result_payload( ) -> MessagePayload { MessagePayload { message: Message { - schema_version: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: Role::Tool, content: vec![ContentPart::ToolResult { content: ToolResult { @@ -126,7 +229,7 @@ pub(crate) fn prompt_request_payload( ) -> MessagePayload { MessagePayload { message: Message { - schema_version: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: Role::User, content: vec![ContentPart::PromptRequest { content: PromptRequest { @@ -347,6 +450,148 @@ fn mcp_prompt_message(message: &Message) -> Option { mod tests { use super::*; + #[test] + fn resource_result_response_applies_text_changes() { + let original = + ReadResourceResult::new(vec![ResourceContents::text("AWS_ACCESS_KEY_ID=secret", "file:///password.env")]); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource result is supported"); + let ContentPart::Resource { content } = &mut payload.message.content[0] else { + panic!("expected resource content"); + }; + content.content = Some("AWS_ACCESS_KEY_ID=[redacted]".to_owned()); + + let result = resource_result_response(original, &payload).expect("text edit applies"); + + let ResourceContents::TextResourceContents { text, uri, .. } = &result.contents[0] else { + panic!("expected text resource"); + }; + assert_eq!("AWS_ACCESS_KEY_ID=[redacted]", text); + assert_eq!("file:///password.env", uri); + } + + #[test] + fn resource_result_response_decodes_and_applies_blob_changes() { + let wire_blob = BASE64_STANDARD.encode(b"AWS_ACCESS_KEY_ID=secret"); + let original = ReadResourceResult::new(vec![ + ResourceContents::blob(wire_blob.clone(), "file:///password.bin") + .with_mime_type("application/octet-stream"), + ]); + let mut payload = resource_result_payload(&original, "resource-1").expect("valid blob is supported"); + let ContentPart::Resource { content } = &mut payload.message.content[0] else { + panic!("expected resource content"); + }; + assert_eq!(Some(b"AWS_ACCESS_KEY_ID=secret".as_slice()), content.blob.as_deref()); + content.blob = Some(b"AWS_ACCESS_KEY_ID=[redacted]".to_vec()); + + let result = resource_result_response(original, &payload).expect("blob edit applies"); + + let ResourceContents::BlobResourceContents { blob, uri, .. } = &result.contents[0] else { + panic!("expected blob resource"); + }; + assert_eq!(b"AWS_ACCESS_KEY_ID=[redacted]", BASE64_STANDARD.decode(blob).expect("valid base64").as_slice()); + assert_eq!("file:///password.bin", uri); + assert_ne!(&wire_blob, blob); + } + + #[test] + fn resource_result_response_preserves_unchanged_blob_wire_value() { + let wire_blob = BASE64_STANDARD.encode(b"unchanged"); + let original = ReadResourceResult::new(vec![ResourceContents::blob(&wire_blob, "file:///image.bin")]); + let payload = resource_result_payload(&original, "resource-1").expect("valid blob is supported"); + + let result = resource_result_response(original, &payload).expect("unchanged blob applies"); + + let ResourceContents::BlobResourceContents { blob, .. } = &result.contents[0] else { + panic!("expected blob resource"); + }; + assert_eq!(&wire_blob, blob); + } + + #[test] + fn resource_result_payload_rejects_invalid_base64_blob() { + let original = ReadResourceResult::new(vec![ResourceContents::blob("not base64!", "file:///image.bin")]); + + assert!(resource_result_payload(&original, "resource-1").is_none()); + } + + #[test] + fn resource_result_allows_mime_uri_and_content_type_changes() { + let original: ReadResourceResult = serde_json::from_value(serde_json::json!({ + "_meta": {"response": "preserved"}, + "contents": [ + {"uri": "file:///a", "mimeType": "text/plain", "text": "original", "_meta": {"item": 1}}, + {"uri": "file:///b", "mimeType": "application/octet-stream", "blob": "YmluYXJ5", "_meta": {"item": 2}} + ] + })) + .expect("valid resource response"); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource payload"); + for (index, part) in payload.message.content.iter_mut().enumerate() { + let ContentPart::Resource { content } = part else { panic!("resource content") }; + content.uri = format!("file:///changed-{index}"); + if index == 0 { + content.content = None; + content.blob = Some(b"binary edit".to_vec()); + content.mime_type = Some("application/octet-stream".to_owned()); + } else { + content.blob = None; + content.content = Some("text edit".to_owned()); + content.mime_type = Some("text/plain".to_owned()); + } + } + let actual = serde_json::to_value(resource_result_response(original, &payload).expect("valid changes apply")) + .expect("response serializes"); + assert_eq!(serde_json::json!({"response": "preserved"}), actual["_meta"]); + assert_eq!("file:///changed-0", actual["contents"][0]["uri"]); + assert_eq!("application/octet-stream", actual["contents"][0]["mimeType"]); + assert_eq!(BASE64_STANDARD.encode(b"binary edit"), actual["contents"][0]["blob"]); + assert_eq!(1, actual["contents"][0]["_meta"]["item"]); + assert_eq!("file:///changed-1", actual["contents"][1]["uri"]); + assert_eq!("text/plain", actual["contents"][1]["mimeType"]); + assert_eq!("text edit", actual["contents"][1]["text"]); + assert_eq!(2, actual["contents"][1]["_meta"]["item"]); + } + + #[test] + fn resource_result_ignores_cmf_fields_that_are_not_mcp_content() { + let original = ReadResourceResult::new(vec![ResourceContents::text("original", "file:///a")]); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource payload"); + payload.message.schema_version = "plugin value".to_owned(); + payload.message.role = Role::User; + payload.message.channel = Some(cpex::cpex_core::cmf::Channel::Analysis); + let ContentPart::Resource { content } = &mut payload.message.content[0] else { panic!("resource content") }; + content.resource_request_id = "plugin value".to_owned(); + content.name = Some("display name".to_owned()); + content.description = Some("description".to_owned()); + content.size_bytes = Some(8); + content.version = Some("v2".to_owned()); + content.annotations.insert("note".to_owned(), serde_json::json!("annotation")); + content.content = Some("redacted".to_owned()); + let result = resource_result_response(original, &payload).expect("MCP content remains usable"); + let ResourceContents::TextResourceContents { text, .. } = &result.contents[0] else { panic!("text content") }; + assert_eq!("redacted", text); + } + + #[test] + fn resource_result_prefers_text_when_both_content_fields_are_present() { + let original = ReadResourceResult::new(vec![ResourceContents::blob("YmluYXJ5", "file:///a")]); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource payload"); + let ContentPart::Resource { content } = &mut payload.message.content[0] else { panic!("resource content") }; + content.content = Some("text replacement".to_owned()); + let result = + resource_result_response(original, &payload).expect("text takes precedence, as in the built-in serializer"); + let ResourceContents::TextResourceContents { text, .. } = &result.contents[0] else { panic!("text resource") }; + assert_eq!("text replacement", text); + } + + #[test] + fn resource_result_rejects_content_without_a_valid_mcp_representation() { + let original = ReadResourceResult::new(vec![ResourceContents::text("original", "file:///a")]); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource payload"); + let ContentPart::Resource { content } = &mut payload.message.content[0] else { panic!("resource content") }; + content.content = None; + assert!(resource_result_response(original, &payload).is_none()); + } + fn text_prompt() -> GetPromptResult { GetPromptResult::new(vec![PromptMessage::new_text(McpRole::User, "review of weather")]) } diff --git a/crates/contextforge-data-plane-cpex/src/factory.rs b/crates/contextforge-data-plane-cpex/src/factory.rs index 598d8a2e..627abe1c 100644 --- a/crates/contextforge-data-plane-cpex/src/factory.rs +++ b/crates/contextforge-data-plane-cpex/src/factory.rs @@ -28,7 +28,7 @@ where let handlers = config .hooks .iter() - .filter_map(|hook| cmf_hook_name(hook)) + .filter_map(|hook| supported_cmf_hook_name(hook)) .map(|hook| { (hook, Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) as Arc) }) @@ -45,12 +45,14 @@ where } } -fn cmf_hook_name(hook: &str) -> Option<&'static str> { +pub(crate) fn supported_cmf_hook_name(hook: &str) -> Option<&'static str> { match hook { cmf_hook_names::TOOL_PRE_INVOKE => Some(cmf_hook_names::TOOL_PRE_INVOKE), cmf_hook_names::TOOL_POST_INVOKE => Some(cmf_hook_names::TOOL_POST_INVOKE), cmf_hook_names::PROMPT_PRE_FETCH => Some(cmf_hook_names::PROMPT_PRE_FETCH), cmf_hook_names::PROMPT_POST_FETCH => Some(cmf_hook_names::PROMPT_POST_FETCH), + cmf_hook_names::RESOURCE_PRE_FETCH => Some(cmf_hook_names::RESOURCE_PRE_FETCH), + cmf_hook_names::RESOURCE_POST_FETCH => Some(cmf_hook_names::RESOURCE_POST_FETCH), _ => None, } } diff --git a/crates/contextforge-data-plane-cpex/src/handle.rs b/crates/contextforge-data-plane-cpex/src/handle.rs index a2f504a7..f6e8e507 100644 --- a/crates/contextforge-data-plane-cpex/src/handle.rs +++ b/crates/contextforge-data-plane-cpex/src/handle.rs @@ -13,7 +13,9 @@ use cpex::cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult, ErrorCode, GetPromptRequestParams, GetPromptResult}, + model::{ + CallToolRequestParams, CallToolResult, ErrorCode, GetPromptRequestParams, GetPromptResult, ReadResourceResult, + }, serde::{Serialize, de::DeserializeOwned}, }; use tokio::task::JoinHandle; @@ -22,7 +24,7 @@ use crate::{ config::{LoadedRuntimePluginConfig, RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, error::GatewayPluginRuntimeError, hooks::{PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolPreCallResult}, - runtime::GatewayPluginRuntime, + runtime::{GatewayPluginRuntime, ResourceCallState}, }; const DEFAULT_CONFIG_WATCHER_INTERVAL: Duration = Duration::from_mins(10); @@ -45,6 +47,30 @@ struct RegistryCallState { state: Option, } +struct RegistryResourceCallState { + runtime: Arc, + state: ResourceCallState, +} + +/// Captures the resource post-hook decision and runtime for one request. +pub struct ResourceHookState { + rewritten_uri: Option, + call: Option, +} + +impl ResourceHookState { + pub fn rewritten_uri(&self) -> Option<&str> { + self.rewritten_uri.as_deref() + } + + pub async fn after_read_resource(self, response: ReadResourceResult) -> Result { + match self.call { + Some(call) => call.runtime.after_read_resource(response, call.state).await, + None => Ok(response), + } + } +} + enum RuntimeState { Active(Arc), Failed(String), @@ -280,6 +306,18 @@ impl GatewayPluginRuntimeHandle { Ok(result) } + pub async fn before_read_resource(&self, resource_uri: &str) -> Result { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + let (rewritten_uri, call) = runtime.before_read_resource(resource_uri).await?; + Ok(ResourceHookState { + rewritten_uri, + call: call.map(|state| RegistryResourceCallState { runtime: Arc::clone(runtime), state }), + }) + } + pub async fn after_get_prompt( &self, prompt_name: &str, @@ -324,7 +362,7 @@ impl GatewayPluginRuntimeHandle { fn runtime_failed_error(state: &RuntimeState) -> ErrorData { if let RuntimeState::Failed(error) = state { - tracing::warn!(%error, "rejecting tool call because CPEX runtime is failed"); + tracing::warn!(%error, "rejecting MCP call because CPEX runtime is failed"); } ErrorData { code: ErrorCode::INTERNAL_ERROR, message: "Runtime plugin reload failed".into(), data: None } } @@ -342,7 +380,7 @@ mod tests { use async_trait::async_trait; use cpex::cpex_core::{ - cmf::{CmfHook, ContentPart, MessagePayload, Role}, + cmf::{CmfHook, ContentPart, MessagePayload}, context::PluginContext, error::{PluginError, PluginViolation}, factory::{PluginFactory, PluginInstance}, @@ -352,6 +390,7 @@ mod tests { }; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, NumberOrString, ProgressNotificationParam, ProgressToken, + ReadResourceResult, ResourceContents, }; use serde_json::{Value, json}; use tokio::sync::Mutex as TokioMutex; @@ -517,7 +556,12 @@ mod tests { _extensions: &Extensions, ctx: &mut PluginContext, ) -> PluginResult { - let is_post = payload.message.role == Role::Tool; + let is_post = payload.message.content.iter().any(|part| { + matches!( + part, + ContentPart::ToolResult { .. } | ContentPart::PromptResult { .. } | ContentPart::Resource { .. } + ) + }); let mut observations = self.observations.lock().expect("observations lock poisoned"); if is_post { observations.post_calls += 1; @@ -640,6 +684,8 @@ mod tests { let hook = match hook.as_str() { cmf_hook_names::TOOL_PRE_INVOKE => cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE => cmf_hook_names::TOOL_POST_INVOKE, + cmf_hook_names::RESOURCE_PRE_FETCH => cmf_hook_names::RESOURCE_PRE_FETCH, + cmf_hook_names::RESOURCE_POST_FETCH => cmf_hook_names::RESOURCE_POST_FETCH, _ => return None, }; Some(( @@ -769,6 +815,60 @@ mod tests { runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn resource_pre_hook_runs_for_a_canonical_uri() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_PRE_FETCH])); + 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"); + + assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn resource_without_post_hook_keeps_its_decision_across_reload() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH])); + 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"); + assert!(state.call.is_none(), "no post-hook state allocation"); + runtime.apply_config(Some(plugin_config(&[plugin]).cpex)).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); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + 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"); + 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"); + assert_eq!("Plugin denied resource", error.message); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn resource_hooks_preserve_context_across_the_backend_call() { + let plugin = Arc::new( + TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH]) + .with_context_roundtrip(), + ); + 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 response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); + + pre.after_read_resource(response).await.expect("resource post hook receives pre context"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn runtime_config_loads_registered_factory_plugin() { let plugin = diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index 5d1be674..3d024272 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -9,7 +9,7 @@ mod runtime; pub use error::GatewayPluginRuntimeError; pub use factory::CmfPluginFactory; -pub use handle::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; +pub use handle::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle, ResourceHookState}; pub use hooks::{ PromptArgumentsUpdate, PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, diff --git a/crates/contextforge-data-plane-cpex/src/pipeline.rs b/crates/contextforge-data-plane-cpex/src/pipeline.rs index 9e4ecab0..93abb715 100644 --- a/crates/contextforge-data-plane-cpex/src/pipeline.rs +++ b/crates/contextforge-data-plane-cpex/src/pipeline.rs @@ -2,7 +2,7 @@ use cpex::cpex_core::cmf::MessagePayload; use cpex::cpex_core::executor::PipelineResult; use rmcp::{ ErrorData, - model::{CallToolResult, ErrorCode, GetPromptResult}, + model::{CallToolResult, ErrorCode, GetPromptResult, ReadResourceResult}, serde::de::DeserializeOwned, }; use tracing::warn; @@ -10,8 +10,8 @@ use tracing::warn; use crate::{ PromptArgumentsUpdate, ToolArgumentsUpdate, cmf::{ - prompt_request_arguments, prompt_result_rejection, prompt_result_response, tool_call_arguments, - tool_result_content, tool_result_response, + prompt_request_arguments, prompt_result_rejection, prompt_result_response, resource_result_response, + tool_call_arguments, tool_result_content, tool_result_response, }, }; @@ -97,6 +97,25 @@ pub(crate) fn effective_post_prompt_result( }) } +pub(crate) fn effective_pre_resource_uri(result: &PipelineResult) -> Result, ErrorData> { + let Some(payload) = modified_message_payload(result) else { return Ok(None) }; + let [cpex::cpex_core::cmf::ContentPart::ResourceRef { content }] = payload.message.content.as_slice() else { + return Err(ErrorData::internal_error("Plugin returned an invalid resource request", None)); + }; + Ok(Some(content.uri.clone())) +} + +pub(crate) fn effective_post_resource_result( + original: ReadResourceResult, + result: &PipelineResult, +) -> Result { + let Some(payload) = modified_message_payload(result) else { + return Ok(original); + }; + resource_result_response(original, payload) + .ok_or_else(|| ErrorData::internal_error("Plugin returned a resource result the gateway cannot apply", None)) +} + pub(crate) fn effective_post_json(original: T, result: &PipelineResult) -> Result where T: DeserializeOwned, diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index 36c2efb7..2d4b4914 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -14,20 +14,23 @@ use cpex::cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult, GetPromptRequestParams, GetPromptResult}, + model::{CallToolRequestParams, CallToolResult, GetPromptRequestParams, GetPromptResult, ReadResourceResult}, serde::{Serialize, de::DeserializeOwned}, }; use tokio::sync::Mutex; use crate::{ cmf::{ - prompt_request_payload, prompt_result_payload, tool_call_payload, tool_json_result_payload, tool_result_payload, + prompt_request_payload, prompt_result_payload, resource_request_payload, resource_result_payload, + tool_call_payload, tool_json_result_payload, tool_result_payload, }, error::GatewayPluginRuntimeError, + factory::supported_cmf_hook_name, hooks::{PromptPreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, pipeline::{ - effective_post_json, effective_post_prompt_result, effective_post_result, effective_pre_args, - effective_pre_prompt_args, log_pipeline_errors, plugin_denied_error, + effective_post_json, effective_post_prompt_result, effective_post_resource_result, effective_post_result, + effective_pre_args, effective_pre_prompt_args, effective_pre_resource_uri, log_pipeline_errors, + plugin_denied_error, }, }; @@ -41,6 +44,7 @@ struct HookPair { struct HookPresence { tool: HookPair, prompt: HookPair, + resource: HookPair, } #[derive(Default)] @@ -82,6 +86,15 @@ fn new_prompt_call_state(context_table: PluginContextTable, prompt_request_id: S Arc::new(PromptCallState { context_table, prompt_request_id }) } +pub(crate) struct ResourceCallState { + context_table: PluginContextTable, + resource_request_id: String, +} + +fn next_resource_request_id() -> String { + format!("gateway-resource-request-{}", CORRELATION_ID.fetch_add(1, Ordering::Relaxed)) +} + impl GatewayPluginRuntime { pub(crate) fn has_post_hook(&self) -> bool { self.hooks.tool.post @@ -106,6 +119,10 @@ impl GatewayPluginRuntime { pre: declares(&config, cmf_hook_names::PROMPT_PRE_FETCH), post: declares(&config, cmf_hook_names::PROMPT_POST_FETCH), }, + resource: HookPair { + pre: declares(&config, cmf_hook_names::RESOURCE_PRE_FETCH), + post: declares(&config, cmf_hook_names::RESOURCE_POST_FETCH), + }, }; let manager = PluginManager::from_config(config, factories) .map_err(|source| GatewayPluginRuntimeError::Configuration { hook: "config", source })?; @@ -128,13 +145,6 @@ impl Drop for GatewayPluginRuntime { } } -const SUPPORTED_HOOKS: [&str; 4] = [ - cmf_hook_names::TOOL_PRE_INVOKE, - cmf_hook_names::TOOL_POST_INVOKE, - cmf_hook_names::PROMPT_PRE_FETCH, - cmf_hook_names::PROMPT_POST_FETCH, -]; - fn declares(config: &CpexConfig, hook_name: &str) -> bool { config.plugins.iter().any(|plugin| plugin.hooks.iter().any(|hook| hook == hook_name)) } @@ -155,7 +165,7 @@ fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayP return Err(GatewayPluginRuntimeError::ConfigUnsupported); } - if plugin.hooks.iter().any(|hook| !SUPPORTED_HOOKS.contains(&hook.as_str())) { + if plugin.hooks.iter().any(|hook| supported_cmf_hook_name(hook).is_none()) { return Err(GatewayPluginRuntimeError::ConfigUnsupported); } } @@ -164,26 +174,15 @@ fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayP } impl GatewayPluginRuntime { - async fn invoke_tool_pre(&self, payload: MessagePayload) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::TOOL_PRE_INVOKE, payload, Extensions::default(), None) - .await; - log_pipeline_errors(cmf_hook_names::TOOL_PRE_INVOKE, &result); - drop(background_tasks); - result - } - - async fn invoke_tool_post( + async fn invoke_cmf_hook( &self, + hook_name: &'static str, payload: MessagePayload, context_table: Option, ) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::TOOL_POST_INVOKE, payload, Extensions::default(), context_table) - .await; - log_pipeline_errors(cmf_hook_names::TOOL_POST_INVOKE, &result); + let (result, background_tasks) = + self.manager.invoke_named::(hook_name, payload, Extensions::default(), context_table).await; + log_pipeline_errors(hook_name, &result); drop(background_tasks); result } @@ -201,7 +200,7 @@ impl GatewayPluginRuntime { let tool_call_id = next_tool_call_id(); let original_payload = tool_call_payload(request, tool_name, backend_name, &tool_call_id); - let pre_result = self.invoke_tool_pre(original_payload).await; + let pre_result = self.invoke_cmf_hook(cmf_hook_names::TOOL_PRE_INVOKE, original_payload, None).await; if pre_result.is_denied() { return Err(plugin_denied_error("tool call", pre_result)); } @@ -211,30 +210,6 @@ impl GatewayPluginRuntime { Ok(ToolPreCallResult { arguments, state: Some(Arc::new(state)) }) } - async fn invoke_prompt_pre(&self, payload: MessagePayload) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::PROMPT_PRE_FETCH, payload, Extensions::default(), None) - .await; - log_pipeline_errors(cmf_hook_names::PROMPT_PRE_FETCH, &result); - drop(background_tasks); - result - } - - async fn invoke_prompt_post( - &self, - payload: MessagePayload, - context_table: Option, - ) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::PROMPT_POST_FETCH, payload, Extensions::default(), context_table) - .await; - log_pipeline_errors(cmf_hook_names::PROMPT_POST_FETCH, &result); - drop(background_tasks); - result - } - pub(crate) async fn before_get_prompt( &self, request: &GetPromptRequestParams, @@ -253,7 +228,7 @@ impl GatewayPluginRuntime { let prompt_request_id = next_prompt_request_id(); let payload = prompt_request_payload(request, prompt_name, backend_name, &prompt_request_id); - let pre_result = self.invoke_prompt_pre(payload).await; + let pre_result = self.invoke_cmf_hook(cmf_hook_names::PROMPT_PRE_FETCH, payload, None).await; if pre_result.is_denied() { return Err(plugin_denied_error("prompt", pre_result)); } @@ -270,6 +245,37 @@ impl GatewayPluginRuntime { Ok(PromptPreFetchResult { arguments, state }) } + pub(crate) async fn before_read_resource( + &self, + resource_uri: &str, + ) -> Result<(Option, Option), ErrorData> { + if !self.hooks.resource.pre && !self.hooks.resource.post { + return Ok((None, None)); + } + + let resource_request_id = next_resource_request_id(); + if !self.hooks.resource.pre { + return Ok(( + None, + Some(ResourceCallState { context_table: PluginContextTable::default(), resource_request_id }), + )); + } + + let payload = resource_request_payload(resource_uri, &resource_request_id); + let pre_result = self.invoke_cmf_hook(cmf_hook_names::RESOURCE_PRE_FETCH, payload, None).await; + if pre_result.is_denied() { + return Err(plugin_denied_error("resource", pre_result)); + } + let uri = effective_pre_resource_uri(&pre_result)?; + Ok(( + uri, + self.hooks + .resource + .post + .then_some(ResourceCallState { context_table: pre_result.context_table, resource_request_id }), + )) + } + pub(crate) async fn after_get_prompt( &self, prompt_name: &str, @@ -284,7 +290,8 @@ impl GatewayPluginRuntime { let Some(state) = state else { return Ok(response) }; let payload = prompt_result_payload(&response, prompt_name, &state.prompt_request_id); - let post_result = self.invoke_prompt_post(payload, Some(state.context_table.clone())).await; + let post_result = + self.invoke_cmf_hook(cmf_hook_names::PROMPT_POST_FETCH, payload, Some(state.context_table.clone())).await; if post_result.is_denied() { return Err(plugin_denied_error("prompt", post_result)); } @@ -292,6 +299,21 @@ impl GatewayPluginRuntime { effective_post_prompt_result(response, &post_result, prompt_name, &state.prompt_request_id) } + pub(crate) async fn after_read_resource( + &self, + response: ReadResourceResult, + state: ResourceCallState, + ) -> Result { + let payload = resource_result_payload(&response, &state.resource_request_id) + .ok_or_else(|| ErrorData::internal_error("Resource response contains an unsupported content type", None))?; + let post_result = + self.invoke_cmf_hook(cmf_hook_names::RESOURCE_POST_FETCH, payload, Some(state.context_table)).await; + if post_result.is_denied() { + return Err(plugin_denied_error("resource", post_result)); + } + effective_post_resource_result(response, &post_result) + } + pub(crate) async fn after_tool_call( &self, tool_name: &str, @@ -307,7 +329,8 @@ impl GatewayPluginRuntime { let mut state = state.lock().await; let post_result = self - .invoke_tool_post( + .invoke_cmf_hook( + cmf_hook_names::TOOL_POST_INVOKE, tool_result_payload(tool_name, &response, &state.tool_call_id), Some(state.context_table.clone()), ) @@ -339,7 +362,8 @@ impl GatewayPluginRuntime { let content = serde_json::to_value(&event).unwrap_or(serde_json::Value::Null); let mut state = state.lock().await; let post_result = self - .invoke_tool_post( + .invoke_cmf_hook( + cmf_hook_names::TOOL_POST_INVOKE, tool_json_result_payload(tool_name, content, false, &state.tool_call_id), Some(state.context_table.clone()), ) diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 7c2df933..191050f9 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -35,7 +35,7 @@ clap.workspace = true thiserror.workspace = true rmp-serde.workspace = true async-trait.workspace = true -base64 = "0.22.1" +base64.workspace = true reqwest.workspace = true uuid.workspace = true lru_time_cache = "0.11.11" 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 3633da67..36ab5497 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 @@ -22,7 +22,7 @@ pub(super) async fn read_resource( let (virtual_host, _claims) = mcp_call_validator.validate_stateless()?; let downstream_name = request.uri.clone(); - let Some(route) = virtual_host.resources.get(&downstream_name) else { + let Some(mut route) = virtual_host.resources.get(&downstream_name) else { return Err(ErrorData { code: ErrorCode::INVALID_PARAMS, message: "Routing problem... resource not found".into(), @@ -30,9 +30,25 @@ 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?) + } else { + None + }; + if let Some(uri) = resource_hook.as_ref().and_then(|hook| hook.rewritten_uri()) + && uri != route.upstream_name + { + let mut candidates = virtual_host.resources.values().filter(|candidate| candidate.upstream_name == uri); + let rewritten = candidates.next().ok_or_else(|| { + ErrorData::invalid_params("Plugin resource target is not available in this virtual host", None) + })?; + if candidates.any(|candidate| candidate.backend_name != rewritten.backend_name) { + return Err(ErrorData::invalid_params("Plugin resource target is ambiguous", None)); + } + route = rewritten; + } let backend_name = route.backend_name.clone(); let resource_uri = route.upstream_name.clone(); - let backend = virtual_host.backends.get(&backend_name).ok_or_else(|| ErrorData { code: ErrorCode::INVALID_PARAMS, message: "Routing problem... backend not found".into(), @@ -48,6 +64,11 @@ pub(super) async fn read_resource( tracing::warn!("read_resource: backend cleanup failed backend_name = {backend_name} error = {error:?}"); } let response = response.map_err(|error| backend_forward_error("read_resource", &backend_name, &error))?; + let response = if let Some(resource_hook) = resource_hook { + resource_hook.after_read_resource(response).await? + } else { + response + }; info!("read_resource: backend {backend_name} returned {} contents", response.contents.len()); diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs index a44603b1..5b9b5c1c 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs @@ -32,6 +32,7 @@ pub(crate) struct Observations { pub(crate) pre_payload_namespace: Option, pub(crate) pre_payload_role: Option, pub(crate) pre_tool_call_id: Option, + pub(crate) pre_resource_uri: Option, pub(crate) post_payload_name: Option, pub(crate) post_tool_call_ids: Vec, pub(crate) post_result_text: Option, @@ -39,6 +40,7 @@ pub(crate) struct Observations { #[derive(Clone, Copy, Default)] pub(crate) enum PreBehavior { + ResourceUri(&'static str), #[default] Allow, Rewrite, @@ -49,6 +51,7 @@ pub(crate) enum PreBehavior { #[derive(Clone, Copy, Default)] pub(crate) enum PostBehavior { + ResourceText, #[default] Allow, Rewrite, @@ -81,6 +84,16 @@ impl TestPlugin { } } + pub(crate) fn with_resource_uri(mut self, uri: &'static str) -> Self { + self.pre_behavior = PreBehavior::ResourceUri(uri); + self + } + + pub(crate) fn with_resource_text(mut self) -> Self { + self.post_behavior = PostBehavior::ResourceText; + self + } + pub(crate) fn with_pre_rewrite(mut self) -> Self { self.pre_behavior = PreBehavior::Rewrite; self @@ -152,7 +165,8 @@ impl HookHandler for TestPlugin { _extensions: &Extensions, ctx: &mut PluginContext, ) -> PluginResult { - let is_post = payload.message.role == Role::Tool; + let is_post = payload.message.role == Role::Tool + || payload.message.content.iter().any(|part| matches!(part, ContentPart::Resource { .. })); let mut observations = self.observations.lock().expect("observations lock poisoned"); if is_post { observations.post_calls += 1; @@ -163,6 +177,9 @@ impl HookHandler for TestPlugin { observations.post_result_text = Some(cmf_result_text(payload)); } else { observations.pre_calls += 1; + if let Some(ContentPart::ResourceRef { content }) = payload.message.content.first() { + observations.pre_resource_uri = Some(content.uri.clone()); + } if let Some(call) = payload.message.get_tool_calls().first() { observations.pre_payload_name = Some(call.name.clone()); observations.pre_payload_namespace.clone_from(&call.namespace); @@ -175,8 +192,25 @@ impl HookHandler for TestPlugin { if is_post { match self.post_behavior { PostBehavior::Allow => PluginResult::allow(), - PostBehavior::Rewrite => { + PostBehavior::Rewrite | PostBehavior::ResourceText => { let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part { + if matches!(self.post_behavior, PostBehavior::ResourceText) { + content.content = Some("converted".to_owned()); + content.blob = None; + content.mime_type = Some("text/plain".to_owned()); + "file:///converted.txt".clone_into(&mut content.uri); + continue; + } + if let Some(text) = &mut content.content { + "post:[redacted]".clone_into(text); + } + if let Some(blob) = &mut content.blob { + *blob = b"post:[redacted]".to_vec(); + } + } + } let result_text = cmf_result_text(payload); if let Some(ContentPart::ToolResult { content }) = modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::ToolResult { .. })) @@ -245,6 +279,13 @@ impl HookHandler for TestPlugin { } } else { match self.pre_behavior { + PreBehavior::ResourceUri(uri) => { + let mut modified = payload.clone(); + if let Some(ContentPart::ResourceRef { content }) = modified.message.content.first_mut() { + uri.clone_into(&mut content.uri); + } + PluginResult::modify_payload(modified) + }, PreBehavior::Allow => PluginResult::allow(), PreBehavior::Rewrite => { let mut modified = payload.clone(); @@ -507,21 +548,22 @@ impl PluginFactory for TestPluginFactory { pre_behavior: self.pre_behavior, post_behavior: self.post_behavior, }); - let mut handlers = Vec::new(); - if config.hooks.iter().any(|hook| hook == cmf_hook_names::TOOL_PRE_INVOKE) { - handlers.push(( - cmf_hook_names::TOOL_PRE_INVOKE, + let handlers = [ + cmf_hook_names::TOOL_PRE_INVOKE, + cmf_hook_names::TOOL_POST_INVOKE, + cmf_hook_names::RESOURCE_PRE_FETCH, + cmf_hook_names::RESOURCE_POST_FETCH, + ] + .into_iter() + .filter(|hook| config.hooks.iter().any(|configured| configured == hook)) + .map(|hook| { + ( + hook, Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) as Arc, - )); - } - if config.hooks.iter().any(|hook| hook == cmf_hook_names::TOOL_POST_INVOKE) { - handlers.push(( - cmf_hook_names::TOOL_POST_INVOKE, - Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) - as Arc, - )); - } + ) + }) + .collect(); Ok(PluginInstance { plugin: Arc::::clone(&plugin), handlers }) } } diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs index cd60e3b3..7b324834 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin_gateway.rs @@ -20,7 +20,8 @@ use rmcp::{ model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ErrorCode, GetPromptRequestParams, GetPromptResponse, GetPromptResult, Implementation, InitializeRequestParams, InitializeResult, NumberOrString, - ProgressNotificationParam, ProgressToken, PromptMessage, ResourceContents, Role, ServerCapabilities, + ProgressNotificationParam, ProgressToken, PromptMessage, ReadResourceRequestParams, ReadResourceResponse, + ReadResourceResult, ResourceContents, Role, ServerCapabilities, }, service::{RequestContext, Service}, transport::{ @@ -48,6 +49,7 @@ pub(crate) struct BackendState { pub(crate) calls: Arc>>, pub(crate) request_headers: Arc>>, pub(crate) prompts: Arc>>, + pub(crate) resources: Arc>>, pub(crate) cancellations: Arc>>, pub(crate) events: Arc>>, parameter_headers: bool, @@ -90,8 +92,24 @@ impl ServerHandler for TestBackend { _request: InitializeRequestParams, _cx: RequestContext, ) -> Result { - Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().enable_prompts().build()) - .with_server_info(Implementation::new("test-backend", "0.1.0"))) + Ok(InitializeResult::new( + ServerCapabilities::builder().enable_tools().enable_prompts().enable_resources().build(), + ) + .with_server_info(Implementation::new("test-backend", "0.1.0"))) + } + + async fn read_resource( + &self, + request: ReadResourceRequestParams, + _cx: RequestContext, + ) -> Result { + self.state.resources.lock().expect("resource calls lock poisoned").push(request.uri.clone()); + let content = if request.uri == "file:///password.bin" { + ResourceContents::blob("c2VjcmV0", request.uri) + } else { + ResourceContents::text("secret", request.uri) + }; + Ok(ReadResourceResult::new(vec![content]).into()) } async fn get_prompt( @@ -235,10 +253,11 @@ pub const TOOL_NAMES: &[&str] = &[ "reflect_text", "wait_for_cancellation", ]; -pub const RESOURCE_URIS: &[&str] = &[""]; +pub const RESOURCE_URIS: &[&str] = &["file:///password.env", "file:///password.bin"]; pub const PROMPT_NAMES: &[&str] = &["review_bundle", "review"]; pub(crate) struct RunningGateway { + pub(crate) user_store: MemoryUserConfigStore, pub(crate) backend_state: BackendState, pub(crate) backend_name: String, gateway_url: String, @@ -396,6 +415,13 @@ async fn start_gateway_with_state( format!("{backend_name}-sum"), ServiceRoute { backend_name: backend_name.clone(), upstream_name: "sum".to_owned() }, ); + let mut resources = construct_services(&backend_name, RESOURCE_URIS); + for uri in RESOURCE_URIS { + resources.insert( + format!("{backend_name}-{uri}"), + ServiceRoute { backend_name: backend_name.clone(), upstream_name: (*uri).to_owned() }, + ); + } let user_store = MemoryUserConfigStore::default(); user_store .set_config( @@ -418,7 +444,7 @@ async fn start_gateway_with_state( }, )]), tools, - resources: construct_services(&backend_name, RESOURCE_URIS), + resources, resource_templates: HashMap::new(), prompts: construct_services(&backend_name, PROMPT_NAMES), }, @@ -434,7 +460,7 @@ async fn start_gateway_with_state( runtime_plugins_enabled: Some(runtime_plugins_enabled), ..create_default_config() }, - user_store, + user_store: user_store.clone(), user_id: user.to_owned(), virtual_host_id: virtual_host_id.to_owned(), backends: vec![backend], @@ -444,5 +470,5 @@ async fn start_gateway_with_state( .expect("gateway starts"); let gateway_url = fixture.gateway_url(); - RunningGateway { backend_state, backend_name, gateway_url, fixture } + RunningGateway { user_store, backend_state, backend_name, gateway_url, fixture } } diff --git a/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs b/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs index 0fd159cc..c97cefc6 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs @@ -8,7 +8,7 @@ use rmcp::{ model::{ CallToolRequestParams, CallToolResult, ClientCapabilities, ClientRequest, ContentBlock, ErrorCode, GetPromptRequestParams, GetPromptResult, Implementation, InitializeRequestParams, ProgressNotificationParam, - Request, ResourceContents, Role as McpRole, ServerResult, + ReadResourceRequestParams, Request, ResourceContents, Role as McpRole, ServerResult, }, service::{NotificationContext, PeerRequestOptions, RequestHandle, RoleClient, RunningService}, }; @@ -862,3 +862,149 @@ async fn prompt_pre_and_post_hooks_share_gateway_call_context() { assert_eq!(1, observations.pre_calls); assert_eq!(1, observations.post_calls); } + +#[tokio::test] +async fn resource_hooks_inspect_canonical_uris_and_redact_direct_and_aliased_reads() { + let plugin = Arc::new( + TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH]) + .with_post_rewrite(), + ); + let observations = plugin.observations(); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_pre(plugin).await).await; + let service = gateway.connect(TEST_USER_ID).await; + for uri in ["file:///password.env", "file:///password.bin"] { + for requested in [uri.to_owned(), format!("{}-{uri}", gateway.backend_name)] { + let result = + service.read_resource(ReadResourceRequestParams::new(requested)).await.expect("resource is returned"); + match &result.contents[0] { + ResourceContents::TextResourceContents { text, uri: returned, .. } => { + assert_eq!("post:[redacted]", text); + assert_eq!(uri, returned); + }, + ResourceContents::BlobResourceContents { blob, uri: returned, .. } => { + use base64::{Engine as _, prelude::BASE64_STANDARD}; + assert_eq!(b"post:[redacted]", BASE64_STANDARD.decode(blob).expect("valid base64").as_slice()); + assert_eq!(uri, returned); + }, + _ => panic!("expected text or blob resource"), + } + assert_eq!(Some(uri), observations.lock().expect("observations lock").pre_resource_uri.as_deref()); + } + } + let observed = observations.lock().expect("observations lock"); + assert_eq!(4, observed.pre_calls); + assert_eq!(4, observed.post_calls); + assert_eq!( + ["file:///password.env", "file:///password.env", "file:///password.bin", "file:///password.bin"], + gateway.backend_state.resources.lock().expect("resource calls lock").as_slice() + ); +} + +#[tokio::test] +async fn resource_pre_hook_denies_before_the_backend_read() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_PRE_FETCH]).with_pre_deny()); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_pre(plugin).await).await; + let error = gateway + .connect(TEST_USER_ID) + .await + .read_resource(ReadResourceRequestParams::new("file:///password.env")) + .await + .expect_err("resource policy denies"); + assert_eq!(ErrorCode(PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.resources.lock().expect("resource calls lock").is_empty()); +} + +#[tokio::test] +async fn resource_post_hook_denies_the_backend_response() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH]).with_post_deny()); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_post(plugin).await).await; + let error = gateway + .connect(TEST_USER_ID) + .await + .read_resource(ReadResourceRequestParams::new("file:///password.env")) + .await + .expect_err("resource policy denies"); + let (_, message) = error_parts(error); + assert!(message.contains("Plugin denied resource"), "{message}"); + assert_eq!(1, gateway.backend_state.resources.lock().expect("resource calls lock").len()); +} + +#[tokio::test] +async fn resource_plugins_can_rewrite_a_published_target_and_convert_the_result() { + let plugin = Arc::new( + TestPlugin::new( + "resource-rewrite", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ) + .with_resource_uri("file:///password.bin") + .with_resource_text(), + ); + let observations = plugin.observations(); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_pre(plugin).await).await; + let service = gateway.connect(TEST_USER_ID).await; + for uri in ["file:///password.env".to_owned(), format!("{}-file:///password.env", gateway.backend_name)] { + let result = + service.read_resource(ReadResourceRequestParams::new(uri)).await.expect("published target rewrite applies"); + let ResourceContents::TextResourceContents { text, uri, mime_type, .. } = &result.contents[0] else { + panic!("blob was converted to text"); + }; + assert_eq!("converted", text); + assert_eq!("file:///converted.txt", uri); + assert_eq!(Some("text/plain"), mime_type.as_deref()); + } + assert_eq!( + ["file:///password.bin", "file:///password.bin"], + gateway.backend_state.resources.lock().expect("resource calls lock").as_slice() + ); + assert_eq!(2, observations.lock().expect("observations lock").post_calls); +} + +#[tokio::test] +async fn resource_plugin_cannot_route_to_an_unpublished_target() { + let plugin = Arc::new( + TestPlugin::new("resource-rewrite", vec![cmf_hook_names::RESOURCE_PRE_FETCH]) + .with_resource_uri("file:///unpublished"), + ); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_pre(plugin).await).await; + let error = gateway + .connect(TEST_USER_ID) + .await + .read_resource(ReadResourceRequestParams::new("file:///password.env")) + .await + .expect_err("target is outside the published resource routes"); + assert_eq!(ErrorCode::INVALID_PARAMS, error_code(error)); + assert!(gateway.backend_state.resources.lock().expect("resource calls lock").is_empty()); +} + +#[tokio::test] +async fn resource_plugin_rejects_a_target_shared_by_different_backends() { + use contextforge_data_plane_apis::{User, user_store::ServiceRoute}; + use contextforge_data_plane_lib::UserConfigStore; + + let plugin = Arc::new( + TestPlugin::new("resource-rewrite", vec![cmf_hook_names::RESOURCE_PRE_FETCH]) + .with_resource_uri("file:///password.bin"), + ); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_pre(plugin).await).await; + let user = User::new(TEST_USER_ID); + let mut config = gateway.user_store.get_config(&user).await.expect("published config"); + let host = config.virtual_hosts.values_mut().next().expect("virtual host"); + let backend = host.backends[&gateway.backend_name].clone(); + host.backends.insert("second-backend".to_owned(), backend); + host.resources.insert( + "second-file".to_owned(), + ServiceRoute { backend_name: "second-backend".to_owned(), upstream_name: "file:///password.bin".to_owned() }, + ); + gateway.user_store.set_config(&user, &config).await.expect("updated config"); + + let error = gateway + .connect(TEST_USER_ID) + .await + .read_resource(ReadResourceRequestParams::new("file:///password.env")) + .await + .expect_err("rewritten URI is ambiguous"); + let (code, message) = error_parts(error); + assert_eq!(ErrorCode::INVALID_PARAMS, code); + assert_eq!("Plugin resource target is ambiguous", message); + assert!(gateway.backend_state.resources.lock().expect("resource calls lock").is_empty()); +} diff --git a/crates/plugins/cpex-secrets-detection/README.md b/crates/plugins/cpex-secrets-detection/README.md index d77b96f2..d30c709d 100644 --- a/crates/plugins/cpex-secrets-detection/README.md +++ b/crates/plugins/cpex-secrets-detection/README.md @@ -21,7 +21,7 @@ Example config: { "name": "secrets-detection", "kind": "validator/secrets-detection", - "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], + "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke", "cmf.resource_post_fetch"], "config": { "redact": true, "redaction_text": "[redacted]", @@ -33,14 +33,12 @@ Example config: } ``` -The dataplane integration currently wires the tool-call path: +The dataplane integration supports: - `cmf.tool_pre_invoke`: scans tool arguments before the backend receives them. - `cmf.tool_post_invoke`: scans tool results before the client receives them. - -The crate also keeps prompt/resource stage handling for CPEX parity and future -hosts, but the current dataplane runtime config only uses the tool pre/post -hooks. +- `cmf.prompt_pre_fetch`: scans prompt arguments before rendering. +- `cmf.resource_post_fetch`: scans text resource contents before returning them. ## Behavior diff --git a/crates/plugins/cpex-secrets-detection/src/lib.rs b/crates/plugins/cpex-secrets-detection/src/lib.rs index 6d3259be..32abf830 100644 --- a/crates/plugins/cpex-secrets-detection/src/lib.rs +++ b/crates/plugins/cpex-secrets-detection/src/lib.rs @@ -54,8 +54,7 @@ impl Plugin for SecretsDetectionCore { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Stage { - // The dataplane runtime config currently uses the tool pre/post stages. - // Prompt/resource stages are kept for CPEX parity and future hosts. + // Runtime configuration registers each stage independently. PromptPreFetch, ToolPreInvoke, ToolPostInvoke,