From 433a5611f2265b6bc78f40585738b1b62ce8af45 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Thu, 3 Sep 2026 10:50:40 +0100 Subject: [PATCH 1/5] feat: run CPEX hooks for resource reads Signed-off-by: lucarlig --- .secrets.baseline | 4 +- README.md | 4 +- _context/wiki/architecture.md | 6 +- _context/wiki/config.md | 8 +- .../contextforge-data-plane-cpex/src/cmf.rs | 157 +++++++++++++++++- .../src/factory.rs | 2 + .../src/handle.rs | 46 ++++- .../contextforge-data-plane-cpex/src/hooks.rs | 10 ++ .../contextforge-data-plane-cpex/src/lib.rs | 4 +- .../src/pipeline.rs | 33 +++- .../src/runtime.rs | 103 +++++++++++- .../src/gateway/mcp_service/resources.rs | 11 ++ .../tests/secrets_detection_e2e.rs | 39 ++++- .../plugins/cpex-secrets-detection/README.md | 10 +- .../plugins/cpex-secrets-detection/src/lib.rs | 3 +- 15 files changed, 402 insertions(+), 38 deletions(-) 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/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..a07ceab4 100644 --- a/_context/wiki/architecture.md +++ b/_context/wiki/architecture.md @@ -159,7 +159,7 @@ 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. ## Startup And Response Flow @@ -185,7 +185,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 +232,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..c57b687c 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. @@ -214,6 +214,10 @@ MCP prompt results carry no error flag, so a plugin setting `is_error` on the CM 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. +### Resource Read Hook Behavior + +For `resources/read`, the pre hook runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources or deny the response. URI, MIME type, item count, and blob identity must remain stable; unsupported or lossy edits fail closed. + ### Demo Plugin Workflow The optional `test-plugins` feature compiles demo factories from the `cpex-plugins-rs` repository. Redis configuration activates factories already present in the binary; it never loads new Rust code into a running process. diff --git a/crates/contextforge-data-plane-cpex/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index eeb74c97..5f1a66a8 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -6,7 +6,7 @@ use cpex::cpex_core::cmf::{ }; 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}; @@ -33,6 +33,119 @@ 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: "2.0".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_request_matches( + payload: &MessagePayload, + resource_uri: &str, + resource_request_id: &str, +) -> bool { + let [ContentPart::ResourceRef { content }] = payload.message.content.as_slice() else { + return false; + }; + payload.message.role == Role::User + && content.resource_request_id == resource_request_id + && content.uri == resource_uri + && matches!(content.resource_type, ResourceType::Uri) + && content.name.is_none() + && content.range_start.is_none() + && content.range_end.is_none() + && content.selector.is_none() +} + +pub(crate) fn resource_result_payload( + response: &ReadResourceResult, + resource_request_id: &str, +) -> Option { + let content = response + .contents + .iter() + .map(|content| { + let (uri, mime_type, text) = match content { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { + (uri.clone(), mime_type.clone(), Some(text.clone())) + }, + ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), + _ => return None, + }; + Some(ContentPart::Resource { + content: CmfResource { + resource_request_id: resource_request_id.to_owned(), + uri, + resource_type: ResourceType::Uri, + content: text, + mime_type, + ..Default::default() + }, + }) + }) + .collect::>>()?; + Some(MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + }) +} + +pub(crate) fn resource_result_response( + mut original: ReadResourceResult, + payload: &MessagePayload, + resource_request_id: &str, +) -> Option { + if payload.message.role != Role::Assistant || 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; + }; + if content.resource_request_id != resource_request_id + || !matches!(content.resource_type, ResourceType::Uri) + || content.name.is_some() + || content.description.is_some() + || content.blob.is_some() + || content.size_bytes.is_some() + || !content.annotations.is_empty() + || content.version.is_some() + { + return None; + } + match original { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { + if content.uri != *uri || content.mime_type != *mime_type { + return None; + } + *text = content.content.clone()?; + }, + ResourceContents::BlobResourceContents { uri, mime_type, .. } => { + if content.uri != *uri || content.mime_type != *mime_type || content.content.is_some() { + return None; + } + }, + _ => 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, @@ -347,6 +460,48 @@ fn mcp_prompt_message(message: &Message) -> Option { mod tests { use super::*; + #[test] + fn resource_request_rejects_uri_mutation() { + let mut payload = resource_request_payload("file:///password.env", "resource-1"); + let ContentPart::ResourceRef { content } = &mut payload.message.content[0] else { + panic!("expected resource reference"); + }; + content.uri = "file:///other.env".to_owned(); + + assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); + } + + #[test] + fn resource_result_response_applies_only_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, "resource-1").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_rejects_uri_mutation() { + let original = ReadResourceResult::new(vec![ResourceContents::text("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.uri = "file:///other.env".to_owned(); + + assert!(resource_result_response(original, &payload, "resource-1").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..a2c8c6e3 100644 --- a/crates/contextforge-data-plane-cpex/src/factory.rs +++ b/crates/contextforge-data-plane-cpex/src/factory.rs @@ -51,6 +51,8 @@ fn cmf_hook_name(hook: &str) -> Option<&'static str> { 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..6dcd0dd8 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; @@ -21,7 +23,7 @@ use tokio::task::JoinHandle; use crate::{ config::{LoadedRuntimePluginConfig, RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, error::GatewayPluginRuntimeError, - hooks::{PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolPreCallResult}, + hooks::{PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookError, RuntimeHookState, ToolPreCallResult}, runtime::GatewayPluginRuntime, }; @@ -280,6 +282,21 @@ 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 mut result = runtime.before_read_resource(resource_uri).await?; + if runtime.has_resource_post_hook() { + let state = result.state.take(); + result.state = Some(Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state })); + } else { + result.state = None; + } + Ok(result) + } + pub async fn after_get_prompt( &self, prompt_name: &str, @@ -292,6 +309,17 @@ impl GatewayPluginRuntimeHandle { } } + pub async fn after_read_resource( + &self, + response: ReadResourceResult, + state: Option, + ) -> Result { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_read_resource(response, state.state.clone()).await, + None => Ok(response), + } + } + pub async fn after_tool_call( &self, tool_name: &str, @@ -324,7 +352,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 } } @@ -640,6 +668,7 @@ 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, _ => return None, }; Some(( @@ -769,6 +798,17 @@ 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 runtime_config_loads_registered_factory_plugin() { let plugin = diff --git a/crates/contextforge-data-plane-cpex/src/hooks.rs b/crates/contextforge-data-plane-cpex/src/hooks.rs index b6a9a63a..0b481601 100644 --- a/crates/contextforge-data-plane-cpex/src/hooks.rs +++ b/crates/contextforge-data-plane-cpex/src/hooks.rs @@ -52,6 +52,16 @@ pub struct PromptPreFetchResult { pub state: Option, } +pub struct ResourcePreFetchResult { + pub state: Option, +} + +impl ResourcePreFetchResult { + pub fn unchanged() -> Self { + Self { state: None } + } +} + impl PromptPreFetchResult { pub fn unchanged() -> Self { Self { arguments: PromptArgumentsUpdate::Unchanged, state: None } diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index 5d1be674..944bfb22 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -11,6 +11,6 @@ pub use error::GatewayPluginRuntimeError; pub use factory::CmfPluginFactory; pub use handle::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; pub use hooks::{ - PromptArgumentsUpdate, PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolArgumentsUpdate, - ToolPreCallResult, + PromptArgumentsUpdate, PromptPreFetchResult, ResourcePreFetchResult, 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..a03081ca 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_request_matches, + resource_result_response, tool_call_arguments, tool_result_content, tool_result_response, }, }; @@ -97,6 +97,33 @@ pub(crate) fn effective_post_prompt_result( }) } +pub(crate) fn validate_pre_resource_result( + result: &PipelineResult, + resource_uri: &str, + resource_request_id: &str, +) -> Result<(), ErrorData> { + let Some(payload) = modified_message_payload(result) else { + return Ok(()); + }; + if resource_request_matches(payload, resource_uri, resource_request_id) { + Ok(()) + } else { + Err(ErrorData::internal_error("Plugin attempted to modify the canonical resource route", None)) + } +} + +pub(crate) fn effective_post_resource_result( + original: ReadResourceResult, + result: &PipelineResult, + resource_request_id: &str, +) -> Result { + let Some(payload) = modified_message_payload(result) else { + return Ok(original); + }; + resource_result_response(original, payload, resource_request_id) + .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..fd9d0ce6 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -14,20 +14,22 @@ 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, - hooks::{PromptPreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, + hooks::{PromptPreFetchResult, ResourcePreFetchResult, 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, log_pipeline_errors, plugin_denied_error, + validate_pre_resource_result, }, }; @@ -41,6 +43,7 @@ struct HookPair { struct HookPresence { tool: HookPair, prompt: HookPair, + resource: HookPair, } #[derive(Default)] @@ -82,6 +85,19 @@ fn new_prompt_call_state(context_table: PluginContextTable, prompt_request_id: S Arc::new(PromptCallState { context_table, prompt_request_id }) } +struct ResourceCallState { + context_table: Option, + resource_request_id: String, +} + +fn next_resource_request_id() -> String { + format!("gateway-resource-request-{}", CORRELATION_ID.fetch_add(1, Ordering::Relaxed)) +} + +fn new_resource_call_state(context_table: Option, resource_request_id: String) -> RuntimeHookState { + Arc::new(ResourceCallState { context_table, resource_request_id }) +} + impl GatewayPluginRuntime { pub(crate) fn has_post_hook(&self) -> bool { self.hooks.tool.post @@ -91,6 +107,10 @@ impl GatewayPluginRuntime { self.hooks.prompt.post } + pub(crate) fn has_resource_post_hook(&self) -> bool { + self.hooks.resource.post + } + pub(crate) async fn from_config( config: CpexConfig, factories: &PluginFactoryRegistry, @@ -106,6 +126,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,11 +152,13 @@ impl Drop for GatewayPluginRuntime { } } -const SUPPORTED_HOOKS: [&str; 4] = [ +const SUPPORTED_HOOKS: [&str; 6] = [ cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE, cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH, + cmf_hook_names::RESOURCE_PRE_FETCH, + cmf_hook_names::RESOURCE_POST_FETCH, ]; fn declares(config: &CpexConfig, hook_name: &str) -> bool { @@ -270,6 +296,51 @@ impl GatewayPluginRuntime { Ok(PromptPreFetchResult { arguments, state }) } + async fn invoke_resource_pre(&self, payload: MessagePayload) -> PipelineResult { + let (result, background_tasks) = self + .manager + .invoke_named::(cmf_hook_names::RESOURCE_PRE_FETCH, payload, Extensions::default(), None) + .await; + log_pipeline_errors(cmf_hook_names::RESOURCE_PRE_FETCH, &result); + drop(background_tasks); + result + } + + async fn invoke_resource_post( + &self, + payload: MessagePayload, + context_table: Option, + ) -> PipelineResult { + let (result, background_tasks) = self + .manager + .invoke_named::(cmf_hook_names::RESOURCE_POST_FETCH, payload, Extensions::default(), context_table) + .await; + log_pipeline_errors(cmf_hook_names::RESOURCE_POST_FETCH, &result); + drop(background_tasks); + result + } + + pub(crate) async fn before_read_resource(&self, resource_uri: &str) -> Result { + let resource_request_id = next_resource_request_id(); + if !self.hooks.resource.pre { + let state = self.hooks.resource.post.then(|| new_resource_call_state(None, resource_request_id)); + return Ok(ResourcePreFetchResult { state }); + } + + let payload = resource_request_payload(resource_uri, &resource_request_id); + let pre_result = self.invoke_resource_pre(payload).await; + if pre_result.is_denied() { + return Err(plugin_denied_error("resource", pre_result)); + } + validate_pre_resource_result(&pre_result, resource_uri, &resource_request_id)?; + let state = self + .hooks + .resource + .post + .then(|| new_resource_call_state(Some(pre_result.context_table), resource_request_id)); + Ok(ResourcePreFetchResult { state }) + } + pub(crate) async fn after_get_prompt( &self, prompt_name: &str, @@ -292,6 +363,26 @@ 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: Option, + ) -> Result { + if !self.hooks.resource.post { + return Ok(response); + } + + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(response) }; + 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_resource_post(payload, state.context_table.clone()).await; + if post_result.is_denied() { + return Err(plugin_denied_error("resource", post_result)); + } + effective_post_resource_result(response, &post_result, &state.resource_request_id) + } + pub(crate) async fn after_tool_call( &self, tool_name: &str, 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..4cd7d66c 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 @@ -1,3 +1,4 @@ +use contextforge_data_plane_cpex::ResourcePreFetchResult; use rmcp::{ ErrorData, RoleServer, model::{ @@ -39,6 +40,11 @@ pub(super) async fn read_resource( data: None, })?; + let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { + plugin_runtime.before_read_resource(&resource_uri).await? + } else { + ResourcePreFetchResult::unchanged() + }; let mut backend_service = connect_backend_for_request(mcp_service, &backend_name, backend, &cx).await?; let mut routed_request = request; @@ -48,6 +54,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(plugin_runtime) = &mcp_service.plugin_runtime { + plugin_runtime.after_read_resource(response, pre_result.state).await? + } else { + response + }; info!("read_resource: backend {backend_name} returned {} contents", response.contents.len()); diff --git a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs index b19bf231..12475d5c 100644 --- a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs +++ b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs @@ -16,7 +16,7 @@ use std::{ use contextforge_data_plane_apis::{ User, runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_KEY, RUNTIME_PLUGIN_CONFIG_VERSION}, - user_store::{BackendMCPGateway, UserConfig, VirtualHost}, + user_store::{BackendMCPGateway, ServiceRoute, UserConfig, VirtualHost}, }; use http::{HeaderMap, HeaderValue}; use jsonwebtoken::{Algorithm, EncodingKey, Header, encode}; @@ -25,7 +25,8 @@ use rmcp::{ ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ClientCapabilities, ContentBlock, ErrorCode, - Implementation, InitializeRequestParams, InitializeResult, ServerCapabilities, + Implementation, InitializeRequestParams, InitializeResult, ReadResourceRequestParams, ReadResourceResponse, + ReadResourceResult, ResourceContents, ServerCapabilities, }, service::{RequestContext, ServiceError}, transport::{ @@ -43,6 +44,7 @@ const TEST_USER_EMAIL: &str = "admin@example.com"; const TEST_VIRTUAL_HOST_ID: &str = "vh-secrets-e2e"; const TEST_TOKEN_TTL_SECS: u64 = 60 * 60; const REDACTED: &str = "[redacted]"; +const BACKEND_RESOURCE_SECRET: &str = "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"; // pragma: allowlist secret #[derive(Clone, Debug)] struct BackendObservation { @@ -66,10 +68,18 @@ impl ServerHandler for TestBackend { _request: InitializeRequestParams, _cx: RequestContext, ) -> Result { - Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().enable_resources().build()) .with_server_info(Implementation::new("secrets-e2e-backend", "0.1.0"))) } + async fn read_resource( + &self, + request: ReadResourceRequestParams, + _cx: RequestContext, + ) -> Result { + Ok(ReadResourceResult::new(vec![ResourceContents::text(BACKEND_RESOURCE_SECRET, request.uri)]).into()) + } + async fn call_tool( &self, request: CallToolRequestParams, @@ -177,7 +187,7 @@ struct E2eEnvironment { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ignore = "spawns redis-server and the contextforge-data-plane binary"] -async fn binary_e2e_redacts_tool_arguments_and_results() { +async fn binary_e2e_redacts_tool_arguments_results_and_resources() { let backend = start_backend().await; let env = start_environment( backend, @@ -211,6 +221,17 @@ async fn binary_e2e_redacts_tool_arguments_and_results() { .expect("secret result is redacted and call succeeds"); assert_eq!(REDACTED, tool_text(&result)); + + let result = service + .read_resource(ReadResourceRequestParams::new("file:///password.env")) + .await + .expect("resource is returned"); + let Some(ResourceContents::TextResourceContents { text, uri, .. }) = result.contents.first() else { + panic!("expected text resource contents"); + }; + assert_eq!("file:///password.env", uri); + assert_ne!(BACKEND_RESOURCE_SECRET, text); + assert!(text.contains(REDACTED)); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -389,7 +410,13 @@ async fn write_redis_config(redis_port: u16, backend: &RunningBackend) { }, )]), tools: HashMap::new(), - resources: HashMap::new(), + resources: HashMap::from([( + "file:///password.env".to_owned(), + ServiceRoute { + backend_name: "backend".to_owned(), + upstream_name: "file:///password.env".to_owned(), + }, + )]), resource_templates: HashMap::new(), prompts: HashMap::new(), }, @@ -416,7 +443,7 @@ async fn write_runtime_plugin_config(redis_port: u16, plugin_config: Value) { "plugins": [{ "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": plugin_config, }] } 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, From 3ad8afb5b86a7c91863adcc946529bf1fda9f339 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Thu, 3 Sep 2026 11:20:03 +0100 Subject: [PATCH 2/5] refactor: centralize CPEX hook plumbing Signed-off-by: lucarlig --- .../contextforge-data-plane-cpex/src/cmf.rs | 59 ++++------- .../src/factory.rs | 4 +- .../src/runtime.rs | 99 ++++--------------- 3 files changed, 40 insertions(+), 122 deletions(-) diff --git a/crates/contextforge-data-plane-cpex/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index 5f1a66a8..06a5229c 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -80,23 +80,7 @@ pub(crate) fn resource_result_payload( .contents .iter() .map(|content| { - let (uri, mime_type, text) = match content { - ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { - (uri.clone(), mime_type.clone(), Some(text.clone())) - }, - ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), - _ => return None, - }; - Some(ContentPart::Resource { - content: CmfResource { - resource_request_id: resource_request_id.to_owned(), - uri, - resource_type: ResourceType::Uri, - content: text, - mime_type, - ..Default::default() - }, - }) + cmf_resource_content(content, resource_request_id).map(|content| ContentPart::Resource { content }) }) .collect::>>()?; Some(MessagePayload { @@ -104,6 +88,24 @@ pub(crate) fn resource_result_payload( }) } +fn cmf_resource_content(content: &ResourceContents, resource_request_id: &str) -> Option { + let (uri, mime_type, text) = match content { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { + (uri.clone(), mime_type.clone(), Some(text.clone())) + }, + ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), + _ => return None, + }; + Some(CmfResource { + resource_request_id: resource_request_id.to_owned(), + uri, + resource_type: ResourceType::Uri, + content: text, + mime_type, + ..Default::default() + }) +} + pub(crate) fn resource_result_response( mut original: ReadResourceResult, payload: &MessagePayload, @@ -379,28 +381,7 @@ fn cmf_content_part(block: &ContentBlock, prompt_request_id: &str) -> Option { - let (uri, mime_type, content) = match &resource.resource { - ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { - (uri.clone(), mime_type.clone(), Some(text.clone())) - }, - ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), - _ => return None, - }; - ContentPart::Resource { - content: CmfResource { - resource_request_id: prompt_request_id.to_owned(), - uri, - name: None, - description: None, - resource_type: ResourceType::Uri, - content, - blob: None, - mime_type, - size_bytes: None, - annotations: HashMap::new(), - version: None, - }, - } + ContentPart::Resource { content: cmf_resource_content(&resource.resource, prompt_request_id)? } }, ContentBlock::ResourceLink(link) => ContentPart::ResourceRef { content: ResourceReference { diff --git a/crates/contextforge-data-plane-cpex/src/factory.rs b/crates/contextforge-data-plane-cpex/src/factory.rs index a2c8c6e3..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,7 +45,7 @@ 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), diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index fd9d0ce6..11243e10 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -25,6 +25,7 @@ use crate::{ tool_call_payload, tool_json_result_payload, tool_result_payload, }, error::GatewayPluginRuntimeError, + factory::supported_cmf_hook_name, hooks::{PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, pipeline::{ effective_post_json, effective_post_prompt_result, effective_post_resource_result, effective_post_result, @@ -152,15 +153,6 @@ impl Drop for GatewayPluginRuntime { } } -const SUPPORTED_HOOKS: [&str; 6] = [ - cmf_hook_names::TOOL_PRE_INVOKE, - cmf_hook_names::TOOL_POST_INVOKE, - cmf_hook_names::PROMPT_PRE_FETCH, - cmf_hook_names::PROMPT_POST_FETCH, - cmf_hook_names::RESOURCE_PRE_FETCH, - cmf_hook_names::RESOURCE_POST_FETCH, -]; - fn declares(config: &CpexConfig, hook_name: &str) -> bool { config.plugins.iter().any(|plugin| plugin.hooks.iter().any(|hook| hook == hook_name)) } @@ -181,7 +173,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); } } @@ -190,26 +182,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 } @@ -227,7 +208,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)); } @@ -237,30 +218,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, @@ -279,7 +236,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)); } @@ -296,30 +253,6 @@ impl GatewayPluginRuntime { Ok(PromptPreFetchResult { arguments, state }) } - async fn invoke_resource_pre(&self, payload: MessagePayload) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::RESOURCE_PRE_FETCH, payload, Extensions::default(), None) - .await; - log_pipeline_errors(cmf_hook_names::RESOURCE_PRE_FETCH, &result); - drop(background_tasks); - result - } - - async fn invoke_resource_post( - &self, - payload: MessagePayload, - context_table: Option, - ) -> PipelineResult { - let (result, background_tasks) = self - .manager - .invoke_named::(cmf_hook_names::RESOURCE_POST_FETCH, payload, Extensions::default(), context_table) - .await; - log_pipeline_errors(cmf_hook_names::RESOURCE_POST_FETCH, &result); - drop(background_tasks); - result - } - pub(crate) async fn before_read_resource(&self, resource_uri: &str) -> Result { let resource_request_id = next_resource_request_id(); if !self.hooks.resource.pre { @@ -328,7 +261,7 @@ impl GatewayPluginRuntime { } let payload = resource_request_payload(resource_uri, &resource_request_id); - let pre_result = self.invoke_resource_pre(payload).await; + 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)); } @@ -355,7 +288,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)); } @@ -376,7 +310,8 @@ impl GatewayPluginRuntime { let Some(state) = state else { return Ok(response) }; 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_resource_post(payload, state.context_table.clone()).await; + let post_result = + self.invoke_cmf_hook(cmf_hook_names::RESOURCE_POST_FETCH, payload, state.context_table.clone()).await; if post_result.is_denied() { return Err(plugin_denied_error("resource", post_result)); } @@ -398,7 +333,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()), ) @@ -430,7 +366,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()), ) From 4aff72d206dbd8fd4a8fc8cbbb4b2ef413a0a256 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Thu, 3 Sep 2026 11:59:04 +0100 Subject: [PATCH 3/5] fix: harden CPEX resource hook lifecycle Signed-off-by: lucarlig --- Cargo.lock | 1 + Cargo.toml | 1 + _context/wiki/config.md | 4 +- .../contextforge-data-plane-cpex/Cargo.toml | 1 + .../contextforge-data-plane-cpex/src/cmf.rs | 171 +++++++++++++++--- .../src/handle.rs | 104 +++++++++-- .../contextforge-data-plane-cpex/src/hooks.rs | 30 ++- .../contextforge-data-plane-cpex/src/lib.rs | 4 +- .../src/runtime.rs | 44 +++-- crates/contextforge-data-plane-lib/Cargo.toml | 2 +- 10 files changed, 299 insertions(+), 63 deletions(-) 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/_context/wiki/config.md b/_context/wiki/config.md index c57b687c..87eb6f9c 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -212,11 +212,11 @@ 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 resource blobs reach plugins as decoded CMF bytes and are encoded back to MCP base64 after an edit. Unchanged blobs retain the backend's exact wire representation. ### Resource Read Hook Behavior -For `resources/read`, the pre hook runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources or deny the response. URI, MIME type, item count, and blob identity must remain stable; unsupported or lossy edits fail closed. +For `resources/read`, the pre hook runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources, transform binary resources, or deny the response. URI, MIME type, item count, CMF schema version, and channel must remain stable; unsupported or lossy edits and invalid hook lifecycle state fail closed. ### Demo Plugin Workflow 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 06a5229c..eab86bfc 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -1,8 +1,9 @@ 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, @@ -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 { @@ -36,7 +37,7 @@ 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: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: Role::User, content: vec![ContentPart::ResourceRef { content: ResourceReference { @@ -62,7 +63,7 @@ pub(crate) fn resource_request_matches( let [ContentPart::ResourceRef { content }] = payload.message.content.as_slice() else { return false; }; - payload.message.role == Role::User + canonical_message_envelope(payload, Role::User) && content.resource_request_id == resource_request_id && content.uri == resource_uri && matches!(content.resource_type, ResourceType::Uri) @@ -84,16 +85,18 @@ pub(crate) fn resource_result_payload( }) .collect::>>()?; Some(MessagePayload { - message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + 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) = match content { + let (uri, mime_type, text, blob) = match content { ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { - (uri.clone(), mime_type.clone(), Some(text.clone())) + (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()?)) }, - ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), _ => return None, }; Some(CmfResource { @@ -101,6 +104,7 @@ fn cmf_resource_content(content: &ResourceContents, resource_request_id: &str) - uri, resource_type: ResourceType::Uri, content: text, + blob, mime_type, ..Default::default() }) @@ -111,7 +115,8 @@ pub(crate) fn resource_result_response( payload: &MessagePayload, resource_request_id: &str, ) -> Option { - if payload.message.role != Role::Assistant || payload.message.content.len() != original.contents.len() { + if !canonical_message_envelope(payload, Role::Assistant) || payload.message.content.len() != original.contents.len() + { return None; } @@ -123,7 +128,6 @@ pub(crate) fn resource_result_response( || !matches!(content.resource_type, ResourceType::Uri) || content.name.is_some() || content.description.is_some() - || content.blob.is_some() || content.size_bytes.is_some() || !content.annotations.is_empty() || content.version.is_some() @@ -132,15 +136,20 @@ pub(crate) fn resource_result_response( } match original { ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { - if content.uri != *uri || content.mime_type != *mime_type { + if content.uri != *uri || content.mime_type != *mime_type || content.blob.is_some() { return None; } *text = content.content.clone()?; }, - ResourceContents::BlobResourceContents { uri, mime_type, .. } => { + ResourceContents::BlobResourceContents { uri, mime_type, blob, .. } => { if content.uri != *uri || content.mime_type != *mime_type || content.content.is_some() { return None; } + let modified_blob = content.blob.as_ref()?; + let original_blob = BASE64_STANDARD.decode(blob.as_bytes()).ok()?; + if modified_blob != &original_blob { + *blob = BASE64_STANDARD.encode(modified_blob); + } }, _ => return None, } @@ -148,6 +157,12 @@ pub(crate) fn resource_result_response( Some(original) } +fn canonical_message_envelope(payload: &MessagePayload, role: Role) -> bool { + payload.message.schema_version == SCHEMA_VERSION + && payload.message.role == role + && payload.message.channel.is_none() +} + pub(crate) fn tool_result_payload(tool_name: &str, response: &CallToolResult, tool_call_id: &str) -> MessagePayload { tool_json_result_payload( tool_name, @@ -165,7 +180,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 { @@ -241,7 +256,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 { @@ -283,7 +298,7 @@ pub(crate) fn prompt_result_payload( MessagePayload { message: Message { - schema_version: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: Role::Assistant, content: vec![ContentPart::PromptResult { content: PromptResult { @@ -352,7 +367,7 @@ pub(crate) fn prompt_result_response( fn cmf_prompt_message(message: &PromptMessage, prompt_request_id: &str) -> Message { Message { - schema_version: "2.0".to_owned(), + schema_version: SCHEMA_VERSION.to_owned(), role: match message.role { McpRole::Assistant => Role::Assistant, McpRole::User => Role::User, @@ -422,12 +437,7 @@ fn mcp_prompt_message(message: &Message) -> Option { ContentPart::Audio { content } => { ContentBlock::audio(inline_media_data(&content.source_type, &content.data)?, content.media_type.clone()?) }, - ContentPart::Resource { content } => ContentBlock::resource(ResourceContents::TextResourceContents { - uri: content.uri.clone(), - mime_type: content.mime_type.clone(), - text: content.content.clone()?, - meta: None, - }), + ContentPart::Resource { content } => ContentBlock::resource(mcp_resource_content(content)?), ContentPart::ResourceRef { content } => { ContentBlock::ResourceLink(McpResource::new(content.uri.clone(), content.name.clone()?)) }, @@ -437,8 +447,28 @@ fn mcp_prompt_message(message: &Message) -> Option { Some(PromptMessage::new(role, content)) } +fn mcp_resource_content(content: &CmfResource) -> Option { + match (&content.content, &content.blob) { + (Some(text), None) => Some(ResourceContents::TextResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + text: text.clone(), + meta: None, + }), + (None, Some(blob)) => Some(ResourceContents::BlobResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + blob: BASE64_STANDARD.encode(blob), + meta: None, + }), + _ => None, + } +} + #[cfg(test)] mod tests { + use cpex::cpex_core::cmf::Channel; + use super::*; #[test] @@ -452,6 +482,19 @@ mod tests { assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); } + #[test] + fn resource_request_rejects_modified_envelope() { + let mut payload = resource_request_payload("file:///password.env", "resource-1"); + payload.message.schema_version = "3.0".to_owned(); + + assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); + + let mut payload = resource_request_payload("file:///password.env", "resource-1"); + payload.message.channel = Some(Channel::Analysis); + + assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); + } + #[test] fn resource_result_response_applies_only_text_changes() { let original = @@ -483,6 +526,65 @@ mod tests { assert!(resource_result_response(original, &payload, "resource-1").is_none()); } + #[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, "resource-1").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, "resource-1").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_response_rejects_modified_envelope() { + let original = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); + let mut payload = resource_result_payload(&original, "resource-1").expect("resource result is supported"); + payload.message.schema_version = "3.0".to_owned(); + + assert!(resource_result_response(original.clone(), &payload, "resource-1").is_none()); + + let mut payload = resource_result_payload(&original, "resource-1").expect("resource result is supported"); + payload.message.channel = Some(Channel::Final); + + assert!(resource_result_response(original, &payload, "resource-1").is_none()); + } + fn text_prompt() -> GetPromptResult { GetPromptResult::new(vec![PromptMessage::new_text(McpRole::User, "review of weather")]) } @@ -918,6 +1020,31 @@ mod tests { assert_eq!("file:///app.env", uri); } + #[test] + fn prompt_result_response_round_trips_embedded_blob_resource() { + let original = GetPromptResult::new(vec![PromptMessage::new( + McpRole::User, + ContentBlock::resource(ResourceContents::blob(BASE64_STANDARD.encode(b"token=secret"), "file:///app.bin")), + )]); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { + panic!("embedded resource reaches the plugin as a CMF resource part"); + }; + assert_eq!(Some(b"token=secret".as_slice()), content.blob.as_deref()); + content.blob = Some(b"token=[REDACTED]".to_vec()); + + let result = prompt_result_response(original, &payload, "review", "prompt-1").expect("resource edit applies"); + + let ContentBlock::Resource(resource) = &result.messages[0].content else { + panic!("expected an embedded resource"); + }; + let ResourceContents::BlobResourceContents { blob, uri, .. } = &resource.resource else { + panic!("expected blob resource contents"); + }; + assert_eq!(b"token=[REDACTED]", BASE64_STANDARD.decode(blob).expect("valid base64").as_slice()); + assert_eq!("file:///app.bin", uri); + } + #[test] fn tool_result_response_uses_cmf_error_flag_for_nested_mcp_result() { let original = CallToolResult::success(vec![ContentBlock::text("original")]); diff --git a/crates/contextforge-data-plane-cpex/src/handle.rs b/crates/contextforge-data-plane-cpex/src/handle.rs index 6dcd0dd8..5fd30c31 100644 --- a/crates/contextforge-data-plane-cpex/src/handle.rs +++ b/crates/contextforge-data-plane-cpex/src/handle.rs @@ -23,7 +23,10 @@ use tokio::task::JoinHandle; use crate::{ config::{LoadedRuntimePluginConfig, RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, error::GatewayPluginRuntimeError, - hooks::{PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookError, RuntimeHookState, ToolPreCallResult}, + hooks::{ + PromptPreFetchResult, ResourceHookState, ResourcePreFetchResult, RuntimeHookError, RuntimeHookState, + ToolPreCallResult, invalid_resource_hook_state_error, + }, runtime::GatewayPluginRuntime, }; @@ -47,6 +50,11 @@ struct RegistryCallState { state: Option, } +struct RegistryResourceCallState { + runtime: Arc, + state: RuntimeHookState, +} + enum RuntimeState { Active(Arc), Failed(String), @@ -287,14 +295,13 @@ impl GatewayPluginRuntimeHandle { let RuntimeState::Active(runtime) = state.as_ref() else { return Err(runtime_failed_error(state.as_ref())); }; - let mut result = runtime.before_read_resource(resource_uri).await?; - if runtime.has_resource_post_hook() { - let state = result.state.take(); - result.state = Some(Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state })); - } else { - result.state = None; + match runtime.before_read_resource(resource_uri).await?.state.into_inner() { + Some(state) => Ok(ResourcePreFetchResult::with_post_state(Arc::new(RegistryResourceCallState { + runtime: Arc::clone(runtime), + state, + }))), + None => Ok(ResourcePreFetchResult::unchanged()), } - Ok(result) } pub async fn after_get_prompt( @@ -312,11 +319,19 @@ impl GatewayPluginRuntimeHandle { pub async fn after_read_resource( &self, response: ReadResourceResult, - state: Option, + state: ResourceHookState, ) -> Result { - match state.and_then(|state| state.downcast::().ok()) { - Some(state) => state.runtime.after_read_resource(response, state.state.clone()).await, - None => Ok(response), + let Some(state) = state.into_inner() else { + let current = self.current(); + return match current.as_ref() { + RuntimeState::Active(runtime) if !runtime.has_resource_post_hook() => Ok(response), + RuntimeState::Active(_) => Err(invalid_resource_hook_state_error()), + failed @ RuntimeState::Failed(_) => Err(runtime_failed_error(failed)), + }; + }; + match state.downcast::() { + Ok(state) => state.runtime.after_read_resource(response, Arc::clone(&state.state)).await, + Err(_) => Err(invalid_resource_hook_state_error()), } } @@ -370,7 +385,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}, @@ -380,6 +395,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; @@ -545,7 +561,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; @@ -669,6 +690,7 @@ mod tests { 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(( @@ -809,6 +831,60 @@ mod tests { assert_eq!(1, observations.lock().expect("observations lock poisoned").pre_calls); } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn resource_post_hook_rejects_missing_lifecycle_state() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH])); + let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; + let response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); + + let error = runtime + .handle() + .after_read_resource(response, ResourcePreFetchResult::unchanged().state) + .await + .expect_err("missing state fails closed"); + + assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); + assert_eq!("Resource post-hook state is missing or invalid", error.message); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn resource_post_hook_rejects_invalid_lifecycle_state() { + let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH])); + let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; + let response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); + + let error = runtime + .handle() + .after_read_resource(response, ResourceHookState::active(Arc::new(()))) + .await + .expect_err("invalid state fails closed"); + + assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); + assert_eq!("Resource post-hook state is missing or invalid", 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")]); + + runtime + .handle() + .after_read_resource(response, pre.state) + .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/hooks.rs b/crates/contextforge-data-plane-cpex/src/hooks.rs index 0b481601..392d0fd0 100644 --- a/crates/contextforge-data-plane-cpex/src/hooks.rs +++ b/crates/contextforge-data-plane-cpex/src/hooks.rs @@ -1,6 +1,9 @@ use std::{any::Any, sync::Arc}; -use rmcp::model::{CallToolRequestParams, GetPromptRequestParams}; +use rmcp::{ + ErrorData, + model::{CallToolRequestParams, GetPromptRequestParams}, +}; use serde_json::{Map, Value}; pub type RuntimeHookError = Box; @@ -53,15 +56,36 @@ pub struct PromptPreFetchResult { } pub struct ResourcePreFetchResult { - pub state: Option, + pub state: ResourceHookState, +} + +/// Opaque state connecting a resource pre-fetch hook to its post-fetch hook. +pub struct ResourceHookState(Option); + +impl ResourceHookState { + pub(crate) fn active(state: RuntimeHookState) -> Self { + Self(Some(state)) + } + + pub(crate) fn into_inner(self) -> Option { + self.0 + } } impl ResourcePreFetchResult { pub fn unchanged() -> Self { - Self { state: None } + Self { state: ResourceHookState(None) } + } + + pub(crate) fn with_post_state(state: RuntimeHookState) -> Self { + Self { state: ResourceHookState::active(state) } } } +pub(crate) fn invalid_resource_hook_state_error() -> ErrorData { + ErrorData::internal_error("Resource post-hook state is missing or invalid", None) +} + impl PromptPreFetchResult { pub fn unchanged() -> Self { Self { arguments: PromptArgumentsUpdate::Unchanged, state: None } diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index 944bfb22..e51bd818 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -11,6 +11,6 @@ pub use error::GatewayPluginRuntimeError; pub use factory::CmfPluginFactory; pub use handle::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; pub use hooks::{ - PromptArgumentsUpdate, PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookError, RuntimeHookState, - ToolArgumentsUpdate, ToolPreCallResult, + PromptArgumentsUpdate, PromptPreFetchResult, ResourceHookState, ResourcePreFetchResult, RuntimeHookError, + RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, }; diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index 11243e10..28c4a712 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -26,7 +26,10 @@ use crate::{ }, error::GatewayPluginRuntimeError, factory::supported_cmf_hook_name, - hooks::{PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, + hooks::{ + PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, + invalid_resource_hook_state_error, + }, pipeline::{ effective_post_json, effective_post_prompt_result, effective_post_resource_result, effective_post_result, effective_pre_args, effective_pre_prompt_args, log_pipeline_errors, plugin_denied_error, @@ -87,7 +90,7 @@ fn new_prompt_call_state(context_table: PluginContextTable, prompt_request_id: S } struct ResourceCallState { - context_table: Option, + context_table: PluginContextTable, resource_request_id: String, } @@ -95,7 +98,7 @@ fn next_resource_request_id() -> String { format!("gateway-resource-request-{}", CORRELATION_ID.fetch_add(1, Ordering::Relaxed)) } -fn new_resource_call_state(context_table: Option, resource_request_id: String) -> RuntimeHookState { +fn new_resource_call_state(context_table: PluginContextTable, resource_request_id: String) -> RuntimeHookState { Arc::new(ResourceCallState { context_table, resource_request_id }) } @@ -254,10 +257,16 @@ impl GatewayPluginRuntime { } pub(crate) async fn before_read_resource(&self, resource_uri: &str) -> Result { + if !self.hooks.resource.pre && !self.hooks.resource.post { + return Ok(ResourcePreFetchResult::unchanged()); + } + let resource_request_id = next_resource_request_id(); if !self.hooks.resource.pre { - let state = self.hooks.resource.post.then(|| new_resource_call_state(None, resource_request_id)); - return Ok(ResourcePreFetchResult { state }); + return Ok(ResourcePreFetchResult::with_post_state(new_resource_call_state( + PluginContextTable::default(), + resource_request_id, + ))); } let payload = resource_request_payload(resource_uri, &resource_request_id); @@ -266,12 +275,14 @@ impl GatewayPluginRuntime { return Err(plugin_denied_error("resource", pre_result)); } validate_pre_resource_result(&pre_result, resource_uri, &resource_request_id)?; - let state = self - .hooks - .resource - .post - .then(|| new_resource_call_state(Some(pre_result.context_table), resource_request_id)); - Ok(ResourcePreFetchResult { state }) + if self.hooks.resource.post { + Ok(ResourcePreFetchResult::with_post_state(new_resource_call_state( + pre_result.context_table, + resource_request_id, + ))) + } else { + Ok(ResourcePreFetchResult::unchanged()) + } } pub(crate) async fn after_get_prompt( @@ -300,18 +311,13 @@ impl GatewayPluginRuntime { pub(crate) async fn after_read_resource( &self, response: ReadResourceResult, - state: Option, + state: RuntimeHookState, ) -> Result { - if !self.hooks.resource.post { - return Ok(response); - } - - let state = state.and_then(|state| state.downcast::().ok()); - let Some(state) = state else { return Ok(response) }; + let state = state.downcast::().map_err(|_| invalid_resource_hook_state_error())?; 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, state.context_table.clone()).await; + self.invoke_cmf_hook(cmf_hook_names::RESOURCE_POST_FETCH, payload, Some(state.context_table.clone())).await; if post_result.is_denied() { return Err(plugin_denied_error("resource", post_result)); } 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" From ae9fe68ada18970f6d956e15dba5c4468926df5f Mon Sep 17 00:00:00 2001 From: lucarlig Date: Fri, 4 Sep 2026 21:18:25 +0100 Subject: [PATCH 4/5] fix: preserve CPEX resource policy across reloads Signed-off-by: lucarlig --- _context/wiki/architecture.md | 2 + _context/wiki/config.md | 6 +- _context/wiki/testing.md | 4 +- .../contextforge-data-plane-cpex/src/cmf.rs | 95 ++++++++++-------- .../src/handle.rs | 98 +++++++------------ .../contextforge-data-plane-cpex/src/hooks.rs | 36 +------ .../contextforge-data-plane-cpex/src/lib.rs | 6 +- .../src/runtime.rs | 49 ++++------ .../src/gateway/mcp_service/resources.rs | 11 +-- .../tests/gateway/harness/plugin.rs | 44 ++++++--- .../tests/gateway/harness/plugin_gateway.rs | 44 +++++++-- .../tests/gateway/plugins.rs | 85 +++++++++++++++- .../tests/secrets_detection_e2e.rs | 39 ++------ 13 files changed, 286 insertions(+), 233 deletions(-) diff --git a/_context/wiki/architecture.md b/_context/wiki/architecture.md index a07ceab4..e36a668d 100644 --- a/_context/wiki/architecture.md +++ b/_context/wiki/architecture.md @@ -162,6 +162,8 @@ The binary sets `tikv_jemallocator` as the global allocator. jemalloc holds up b - 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`): diff --git a/_context/wiki/config.md b/_context/wiki/config.md index 87eb6f9c..e3b686e9 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -212,11 +212,13 @@ 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 as decoded CMF bytes and are encoded back to MCP base64 after an edit. Unchanged blobs retain the backend's exact wire representation. +Binary resource blobs reach plugins as decoded CMF bytes and are encoded back to MCP base64 after an edit. Unchanged blobs retain the backend's exact wire representation. Invalid embedded blobs or unsupported prompt content fail the entire prompt before post hooks; they must never disappear from the policy payload while remaining in the client response. ### Resource Read Hook Behavior -For `resources/read`, the pre hook runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources, transform binary resources, or deny the response. URI, MIME type, item count, CMF schema version, and channel must remain stable; unsupported or lossy edits and invalid hook lifecycle state fail closed. +For `resources/read`, the pre hook runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources, transform binary resources, or deny the response. URI, MIME type, item count, CMF schema version, and channel must remain stable; unsupported or lossy edits fail closed. + +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..38b2b9b0 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 redaction, canonical pre-hook URIs, pre/post denial, and invalid embedded prompt blobs. 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/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index eab86bfc..0f0308ff 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -292,11 +292,14 @@ pub(crate) fn prompt_result_payload( response: &GetPromptResult, prompt_name: &str, prompt_request_id: &str, -) -> MessagePayload { - let messages = - response.messages.iter().map(|message| cmf_prompt_message(message, prompt_request_id)).collect::>(); +) -> Option { + let messages = response + .messages + .iter() + .map(|message| cmf_prompt_message(message, prompt_request_id)) + .collect::>>()?; - MessagePayload { + Some(MessagePayload { message: Message { schema_version: SCHEMA_VERSION.to_owned(), role: Role::Assistant, @@ -312,7 +315,7 @@ pub(crate) fn prompt_result_payload( }], channel: None, }, - } + }) } fn prompt_result(payload: &MessagePayload) -> Option<&PromptResult> { @@ -348,13 +351,13 @@ pub(crate) fn prompt_result_response( } for (message, edited) in original.messages.iter_mut().zip(&result.messages) { - let projected = cmf_prompt_message(message, prompt_request_id); + let projected = cmf_prompt_message(message, prompt_request_id)?; if serde_json::to_value(&projected).ok()? == serde_json::to_value(edited).ok()? { continue; } let rebuilt = mcp_prompt_message(edited)?; - if serde_json::to_value(cmf_prompt_message(&rebuilt, prompt_request_id)).ok()? + if serde_json::to_value(cmf_prompt_message(&rebuilt, prompt_request_id)?).ok()? != serde_json::to_value(edited).ok()? { return None; @@ -365,16 +368,16 @@ pub(crate) fn prompt_result_response( Some(original) } -fn cmf_prompt_message(message: &PromptMessage, prompt_request_id: &str) -> Message { - Message { +fn cmf_prompt_message(message: &PromptMessage, prompt_request_id: &str) -> Option { + Some(Message { schema_version: SCHEMA_VERSION.to_owned(), role: match message.role { McpRole::Assistant => Role::Assistant, McpRole::User => Role::User, }, - content: cmf_content_part(&message.content, prompt_request_id).into_iter().collect(), + content: vec![cmf_content_part(&message.content, prompt_request_id)?], channel: None, - } + }) } fn cmf_content_part(block: &ContentBlock, prompt_request_id: &str) -> Option { @@ -608,7 +611,7 @@ mod tests { #[test] fn prompt_result_response_rejects_added_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let extra = edited_messages(&mut payload).first().cloned().expect("one message"); edited_messages(&mut payload).push(extra); @@ -618,7 +621,7 @@ mod tests { #[test] fn prompt_result_response_rejects_extra_prompt_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let duplicate = payload.message.content[0].clone(); payload.message.content.push(duplicate); @@ -628,7 +631,7 @@ mod tests { #[test] fn prompt_result_rejection_reports_the_plugin_error_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let result = prompt_result_mut(&mut payload); result.is_error = true; result.error_message = Some("blocked by policy".to_owned()); @@ -639,7 +642,7 @@ mod tests { #[test] fn prompt_result_rejection_falls_back_when_the_plugin_gives_no_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); prompt_result_mut(&mut payload).is_error = true; assert_eq!(Some("Plugin rejected the rendered prompt".to_owned()), prompt_result_rejection(&payload)); @@ -648,7 +651,7 @@ mod tests { #[test] fn prompt_result_rejection_is_absent_for_a_normal_result() { let original = text_prompt(); - let payload = prompt_result_payload(&original, "review", "prompt-1"); + let payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); assert_eq!(None, prompt_result_rejection(&payload)); } @@ -717,7 +720,7 @@ mod tests { #[test] fn prompt_result_response_rejects_envelope_content_edit() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); prompt_result_mut(&mut payload).content = Some("[REDACTED]".to_owned()); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -726,7 +729,7 @@ mod tests { #[test] fn prompt_result_response_rejects_renamed_prompt() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); prompt_result_mut(&mut payload).prompt_name = "other".to_owned(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -735,7 +738,7 @@ mod tests { #[test] fn prompt_result_response_rejects_recorrelated_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); prompt_result_mut(&mut payload).prompt_request_id = "prompt-2".to_owned(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -744,7 +747,7 @@ mod tests { #[test] fn prompt_result_response_rejects_error_message_without_error_flag() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); prompt_result_mut(&mut payload).error_message = Some("blocked".to_owned()); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -760,7 +763,7 @@ mod tests { #[test] fn prompt_result_response_rejects_resource_type_edit() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -772,7 +775,7 @@ mod tests { #[test] fn prompt_result_response_rejects_dropped_resource_metadata() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -788,7 +791,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_an_image_edit() { let original = media_prompt(ContentBlock::image("aW1hZ2U=", "image/png")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -804,7 +807,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_an_audio_edit() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("audio reaches the plugin as a CMF audio part"); }; @@ -820,7 +823,7 @@ mod tests { #[test] fn prompt_result_response_rejects_url_sourced_audio() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an audio part"); }; @@ -833,7 +836,7 @@ mod tests { #[test] fn prompt_result_response_rejects_audio_without_media_type() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an audio part"); }; @@ -845,7 +848,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_a_resource_link_edit() { let original = media_prompt(ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env"))); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -861,7 +864,7 @@ mod tests { #[test] fn prompt_result_response_rejects_resource_with_removed_text() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -873,7 +876,7 @@ mod tests { #[test] fn prompt_result_response_rejects_multiple_content_parts() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); edited_messages(&mut payload)[0].content.push(ContentPart::Text { text: "extra".to_owned() }); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -882,7 +885,7 @@ mod tests { #[test] fn prompt_result_response_rejects_a_cmf_only_content_part() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); edited_messages(&mut payload)[0].content = vec![ContentPart::Thinking { text: "reasoning".to_owned() }]; assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -891,7 +894,7 @@ mod tests { #[test] fn prompt_result_response_rejects_a_payload_without_a_prompt_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); payload.message.content.clear(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -909,7 +912,7 @@ mod tests { fn prompt_result_response_rejects_url_sourced_image() { let original = GetPromptResult::new(vec![PromptMessage::new(McpRole::User, ContentBlock::image("aW1hZ2U=", "image/png"))]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -923,7 +926,7 @@ mod tests { fn prompt_result_response_rejects_image_without_media_type() { let original = GetPromptResult::new(vec![PromptMessage::new(McpRole::User, ContentBlock::image("aW1hZ2U=", "image/png"))]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -938,7 +941,7 @@ mod tests { McpRole::User, ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -953,7 +956,7 @@ mod tests { McpRole::User, ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -965,7 +968,7 @@ mod tests { #[test] fn prompt_result_response_rejects_removed_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); edited_messages(&mut payload).clear(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -974,16 +977,28 @@ mod tests { #[test] fn prompt_result_response_rejects_unmappable_role() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); edited_messages(&mut payload)[0].role = Role::System; assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); } + #[test] + fn prompt_result_payload_rejects_invalid_embedded_blob() { + let original = GetPromptResult::new(vec![ + PromptMessage::new_text(McpRole::User, "ordinary text"), + PromptMessage::new( + McpRole::User, + ContentBlock::resource(ResourceContents::blob("not base64!", "file:///blocked.env")), + ), + ]); + assert!(prompt_result_payload(&original, "review", "prompt-1").is_none()); + } + #[test] fn prompt_result_response_preserves_unmodified_messages() { let original = text_prompt(); - let payload = prompt_result_payload(&original, "review", "prompt-1"); + let payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let result = prompt_result_response(original.clone(), &payload, "review", "prompt-1") .expect("unmodified payload applies"); @@ -1000,7 +1015,7 @@ mod tests { McpRole::User, ContentBlock::resource(ResourceContents::text("token=secret", "file:///app.env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("embedded resource reaches the plugin as a CMF resource part"); @@ -1026,7 +1041,7 @@ mod tests { McpRole::User, ContentBlock::resource(ResourceContents::blob(BASE64_STANDARD.encode(b"token=secret"), "file:///app.bin")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("embedded resource reaches the plugin as a CMF resource part"); }; diff --git a/crates/contextforge-data-plane-cpex/src/handle.rs b/crates/contextforge-data-plane-cpex/src/handle.rs index 5fd30c31..5ebebb1a 100644 --- a/crates/contextforge-data-plane-cpex/src/handle.rs +++ b/crates/contextforge-data-plane-cpex/src/handle.rs @@ -23,11 +23,8 @@ use tokio::task::JoinHandle; use crate::{ config::{LoadedRuntimePluginConfig, RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, error::GatewayPluginRuntimeError, - hooks::{ - PromptPreFetchResult, ResourceHookState, ResourcePreFetchResult, RuntimeHookError, RuntimeHookState, - ToolPreCallResult, invalid_resource_hook_state_error, - }, - runtime::GatewayPluginRuntime, + hooks::{PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolPreCallResult}, + runtime::{GatewayPluginRuntime, ResourceCallState}, }; const DEFAULT_CONFIG_WATCHER_INTERVAL: Duration = Duration::from_mins(10); @@ -52,7 +49,19 @@ struct RegistryCallState { struct RegistryResourceCallState { runtime: Arc, - state: RuntimeHookState, + state: ResourceCallState, +} + +/// Captures the resource post-hook decision and runtime for one request. +pub struct ResourceHookState(Option); + +impl ResourceHookState { + pub async fn after_read_resource(self, response: ReadResourceResult) -> Result { + match self.0 { + Some(call) => call.runtime.after_read_resource(response, call.state).await, + None => Ok(response), + } + } } enum RuntimeState { @@ -290,18 +299,17 @@ impl GatewayPluginRuntimeHandle { Ok(result) } - pub async fn before_read_resource(&self, resource_uri: &str) -> 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())); }; - match runtime.before_read_resource(resource_uri).await?.state.into_inner() { - Some(state) => Ok(ResourcePreFetchResult::with_post_state(Arc::new(RegistryResourceCallState { - runtime: Arc::clone(runtime), - state, - }))), - None => Ok(ResourcePreFetchResult::unchanged()), - } + Ok(ResourceHookState( + runtime + .before_read_resource(resource_uri) + .await? + .map(|state| RegistryResourceCallState { runtime: Arc::clone(runtime), state }), + )) } pub async fn after_get_prompt( @@ -316,25 +324,6 @@ impl GatewayPluginRuntimeHandle { } } - pub async fn after_read_resource( - &self, - response: ReadResourceResult, - state: ResourceHookState, - ) -> Result { - let Some(state) = state.into_inner() else { - let current = self.current(); - return match current.as_ref() { - RuntimeState::Active(runtime) if !runtime.has_resource_post_hook() => Ok(response), - RuntimeState::Active(_) => Err(invalid_resource_hook_state_error()), - failed @ RuntimeState::Failed(_) => Err(runtime_failed_error(failed)), - }; - }; - match state.downcast::() { - Ok(state) => state.runtime.after_read_resource(response, Arc::clone(&state.state)).await, - Err(_) => Err(invalid_resource_hook_state_error()), - } - } - pub async fn after_tool_call( &self, tool_name: &str, @@ -832,35 +821,28 @@ mod tests { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] - async fn resource_post_hook_rejects_missing_lifecycle_state() { + 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; - let response = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); - - let error = runtime - .handle() - .after_read_resource(response, ResourcePreFetchResult::unchanged().state) - .await - .expect_err("missing state fails closed"); - - assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); - assert_eq!("Resource post-hook state is missing or invalid", error.message); + runtime.apply_config(None).await.expect("disable hooks"); + let state = runtime.handle().before_read_resource("file:///password.env").await.expect("request starts"); + assert!(state.0.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_rejects_invalid_lifecycle_state() { - let plugin = Arc::new(TestPlugin::new("resource", vec![cmf_hook_names::RESOURCE_POST_FETCH])); + 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 = runtime - .handle() - .after_read_resource(response, ResourceHookState::active(Arc::new(()))) - .await - .expect_err("invalid state fails closed"); - - assert_eq!(ErrorCode::INTERNAL_ERROR, error.code); - assert_eq!("Resource post-hook state is missing or invalid", error.message); + 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)] @@ -874,11 +856,7 @@ mod tests { 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")]); - runtime - .handle() - .after_read_resource(response, pre.state) - .await - .expect("resource post hook receives pre context"); + 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); diff --git a/crates/contextforge-data-plane-cpex/src/hooks.rs b/crates/contextforge-data-plane-cpex/src/hooks.rs index 392d0fd0..b6a9a63a 100644 --- a/crates/contextforge-data-plane-cpex/src/hooks.rs +++ b/crates/contextforge-data-plane-cpex/src/hooks.rs @@ -1,9 +1,6 @@ use std::{any::Any, sync::Arc}; -use rmcp::{ - ErrorData, - model::{CallToolRequestParams, GetPromptRequestParams}, -}; +use rmcp::model::{CallToolRequestParams, GetPromptRequestParams}; use serde_json::{Map, Value}; pub type RuntimeHookError = Box; @@ -55,37 +52,6 @@ pub struct PromptPreFetchResult { pub state: Option, } -pub struct ResourcePreFetchResult { - pub state: ResourceHookState, -} - -/// Opaque state connecting a resource pre-fetch hook to its post-fetch hook. -pub struct ResourceHookState(Option); - -impl ResourceHookState { - pub(crate) fn active(state: RuntimeHookState) -> Self { - Self(Some(state)) - } - - pub(crate) fn into_inner(self) -> Option { - self.0 - } -} - -impl ResourcePreFetchResult { - pub fn unchanged() -> Self { - Self { state: ResourceHookState(None) } - } - - pub(crate) fn with_post_state(state: RuntimeHookState) -> Self { - Self { state: ResourceHookState::active(state) } - } -} - -pub(crate) fn invalid_resource_hook_state_error() -> ErrorData { - ErrorData::internal_error("Resource post-hook state is missing or invalid", None) -} - impl PromptPreFetchResult { pub fn unchanged() -> Self { Self { arguments: PromptArgumentsUpdate::Unchanged, state: None } diff --git a/crates/contextforge-data-plane-cpex/src/lib.rs b/crates/contextforge-data-plane-cpex/src/lib.rs index e51bd818..3d024272 100644 --- a/crates/contextforge-data-plane-cpex/src/lib.rs +++ b/crates/contextforge-data-plane-cpex/src/lib.rs @@ -9,8 +9,8 @@ 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, ResourceHookState, ResourcePreFetchResult, RuntimeHookError, - RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, + PromptArgumentsUpdate, PromptPreFetchResult, RuntimeHookError, RuntimeHookState, ToolArgumentsUpdate, + ToolPreCallResult, }; diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index 28c4a712..65ef1e4e 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -26,10 +26,7 @@ use crate::{ }, error::GatewayPluginRuntimeError, factory::supported_cmf_hook_name, - hooks::{ - PromptPreFetchResult, ResourcePreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, - invalid_resource_hook_state_error, - }, + hooks::{PromptPreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, pipeline::{ effective_post_json, effective_post_prompt_result, effective_post_resource_result, effective_post_result, effective_pre_args, effective_pre_prompt_args, log_pipeline_errors, plugin_denied_error, @@ -89,7 +86,7 @@ fn new_prompt_call_state(context_table: PluginContextTable, prompt_request_id: S Arc::new(PromptCallState { context_table, prompt_request_id }) } -struct ResourceCallState { +pub(crate) struct ResourceCallState { context_table: PluginContextTable, resource_request_id: String, } @@ -98,10 +95,6 @@ fn next_resource_request_id() -> String { format!("gateway-resource-request-{}", CORRELATION_ID.fetch_add(1, Ordering::Relaxed)) } -fn new_resource_call_state(context_table: PluginContextTable, resource_request_id: String) -> RuntimeHookState { - Arc::new(ResourceCallState { context_table, resource_request_id }) -} - impl GatewayPluginRuntime { pub(crate) fn has_post_hook(&self) -> bool { self.hooks.tool.post @@ -111,10 +104,6 @@ impl GatewayPluginRuntime { self.hooks.prompt.post } - pub(crate) fn has_resource_post_hook(&self) -> bool { - self.hooks.resource.post - } - pub(crate) async fn from_config( config: CpexConfig, factories: &PluginFactoryRegistry, @@ -256,17 +245,17 @@ impl GatewayPluginRuntime { Ok(PromptPreFetchResult { arguments, state }) } - pub(crate) async fn before_read_resource(&self, resource_uri: &str) -> Result { + pub(crate) async fn before_read_resource( + &self, + resource_uri: &str, + ) -> Result, ErrorData> { if !self.hooks.resource.pre && !self.hooks.resource.post { - return Ok(ResourcePreFetchResult::unchanged()); + return Ok(None); } let resource_request_id = next_resource_request_id(); if !self.hooks.resource.pre { - return Ok(ResourcePreFetchResult::with_post_state(new_resource_call_state( - PluginContextTable::default(), - resource_request_id, - ))); + return Ok(Some(ResourceCallState { context_table: PluginContextTable::default(), resource_request_id })); } let payload = resource_request_payload(resource_uri, &resource_request_id); @@ -275,14 +264,11 @@ impl GatewayPluginRuntime { return Err(plugin_denied_error("resource", pre_result)); } validate_pre_resource_result(&pre_result, resource_uri, &resource_request_id)?; - if self.hooks.resource.post { - Ok(ResourcePreFetchResult::with_post_state(new_resource_call_state( - pre_result.context_table, - resource_request_id, - ))) - } else { - Ok(ResourcePreFetchResult::unchanged()) - } + Ok(self + .hooks + .resource + .post + .then_some(ResourceCallState { context_table: pre_result.context_table, resource_request_id })) } pub(crate) async fn after_get_prompt( @@ -298,7 +284,9 @@ impl GatewayPluginRuntime { let state = state.and_then(|state| state.downcast::().ok()); let Some(state) = state else { return Ok(response) }; - let payload = prompt_result_payload(&response, prompt_name, &state.prompt_request_id); + let payload = prompt_result_payload(&response, prompt_name, &state.prompt_request_id).ok_or_else(|| { + ErrorData::internal_error("Prompt response contains unsupported or invalid content", None) + })?; 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() { @@ -311,13 +299,12 @@ impl GatewayPluginRuntime { pub(crate) async fn after_read_resource( &self, response: ReadResourceResult, - state: RuntimeHookState, + state: ResourceCallState, ) -> Result { - let state = state.downcast::().map_err(|_| invalid_resource_hook_state_error())?; 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.clone())).await; + 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)); } 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 4cd7d66c..88db4948 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 @@ -1,4 +1,3 @@ -use contextforge_data_plane_cpex::ResourcePreFetchResult; use rmcp::{ ErrorData, RoleServer, model::{ @@ -40,10 +39,10 @@ pub(super) async fn read_resource( data: None, })?; - let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { - plugin_runtime.before_read_resource(&resource_uri).await? + let resource_hook = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { + Some(plugin_runtime.before_read_resource(&resource_uri).await?) } else { - ResourcePreFetchResult::unchanged() + None }; let mut backend_service = connect_backend_for_request(mcp_service, &backend_name, backend, &cx).await?; let mut routed_request = request; @@ -54,8 +53,8 @@ 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(plugin_runtime) = &mcp_service.plugin_runtime { - plugin_runtime.after_read_resource(response, pre_result.state).await? + let response = if let Some(resource_hook) = resource_hook { + resource_hook.after_read_resource(response).await? } else { response }; 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..0eb9ee80 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, @@ -152,7 +153,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 +165,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); @@ -177,6 +182,16 @@ impl HookHandler for TestPlugin { PostBehavior::Allow => PluginResult::allow(), PostBehavior::Rewrite => { let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part { + 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 { .. })) @@ -507,21 +522,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..a329ae5e 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( @@ -106,6 +124,13 @@ impl ServerHandler for TestBackend { .push(BackendObservation { tool_name: request.name.clone(), args: request.arguments.clone() }); self.state.events.lock().expect("backend events lock poisoned").push("backend"); + if request.name == "review_invalid_blob" { + return Ok(GetPromptResult::new(vec![PromptMessage::new( + Role::User, + ContentBlock::resource(ResourceContents::blob("not base64!", "file:///blocked.env")), + )]) + .into()); + } let topic = request .arguments .as_ref() @@ -235,8 +260,8 @@ pub const TOOL_NAMES: &[&str] = &[ "reflect_text", "wait_for_cancellation", ]; -pub const RESOURCE_URIS: &[&str] = &[""]; -pub const PROMPT_NAMES: &[&str] = &["review_bundle", "review"]; +pub const RESOURCE_URIS: &[&str] = &["file:///password.env", "file:///password.bin"]; +pub const PROMPT_NAMES: &[&str] = &["review_bundle", "review", "review_invalid_blob"]; pub(crate) struct RunningGateway { pub(crate) backend_state: BackendState, @@ -396,6 +421,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 +450,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), }, diff --git a/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs b/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs index 0fd159cc..22f1ad91 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,86 @@ 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 invalid_prompt_blob_is_rejected_before_post_policy() { + let plugin = Arc::new(PromptTestPlugin::new("prompt", vec![cmf_hook_names::PROMPT_POST_FETCH])); + let observations = plugin.observations(); + let gateway = start_gateway(TEST_USER_ID, true, runtime_with_prompt_plugin(plugin).await).await; + let error = gateway + .connect(TEST_USER_ID) + .await + .get_prompt(GetPromptRequestParams::new("review_invalid_blob")) + .await + .expect_err("invalid content fails closed"); + let (code, message) = error_parts(error); + assert_eq!(ErrorCode::INTERNAL_ERROR, code); + assert_eq!("Prompt response contains unsupported or invalid content", message); + assert_eq!(0, observations.lock().expect("observations lock").post_calls); +} diff --git a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs index 12475d5c..b19bf231 100644 --- a/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs +++ b/crates/contextforge-data-plane/tests/secrets_detection_e2e.rs @@ -16,7 +16,7 @@ use std::{ use contextforge_data_plane_apis::{ User, runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_KEY, RUNTIME_PLUGIN_CONFIG_VERSION}, - user_store::{BackendMCPGateway, ServiceRoute, UserConfig, VirtualHost}, + user_store::{BackendMCPGateway, UserConfig, VirtualHost}, }; use http::{HeaderMap, HeaderValue}; use jsonwebtoken::{Algorithm, EncodingKey, Header, encode}; @@ -25,8 +25,7 @@ use rmcp::{ ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ClientCapabilities, ContentBlock, ErrorCode, - Implementation, InitializeRequestParams, InitializeResult, ReadResourceRequestParams, ReadResourceResponse, - ReadResourceResult, ResourceContents, ServerCapabilities, + Implementation, InitializeRequestParams, InitializeResult, ServerCapabilities, }, service::{RequestContext, ServiceError}, transport::{ @@ -44,7 +43,6 @@ const TEST_USER_EMAIL: &str = "admin@example.com"; const TEST_VIRTUAL_HOST_ID: &str = "vh-secrets-e2e"; const TEST_TOKEN_TTL_SECS: u64 = 60 * 60; const REDACTED: &str = "[redacted]"; -const BACKEND_RESOURCE_SECRET: &str = "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"; // pragma: allowlist secret #[derive(Clone, Debug)] struct BackendObservation { @@ -68,18 +66,10 @@ impl ServerHandler for TestBackend { _request: InitializeRequestParams, _cx: RequestContext, ) -> Result { - Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().enable_resources().build()) + Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) .with_server_info(Implementation::new("secrets-e2e-backend", "0.1.0"))) } - async fn read_resource( - &self, - request: ReadResourceRequestParams, - _cx: RequestContext, - ) -> Result { - Ok(ReadResourceResult::new(vec![ResourceContents::text(BACKEND_RESOURCE_SECRET, request.uri)]).into()) - } - async fn call_tool( &self, request: CallToolRequestParams, @@ -187,7 +177,7 @@ struct E2eEnvironment { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ignore = "spawns redis-server and the contextforge-data-plane binary"] -async fn binary_e2e_redacts_tool_arguments_results_and_resources() { +async fn binary_e2e_redacts_tool_arguments_and_results() { let backend = start_backend().await; let env = start_environment( backend, @@ -221,17 +211,6 @@ async fn binary_e2e_redacts_tool_arguments_results_and_resources() { .expect("secret result is redacted and call succeeds"); assert_eq!(REDACTED, tool_text(&result)); - - let result = service - .read_resource(ReadResourceRequestParams::new("file:///password.env")) - .await - .expect("resource is returned"); - let Some(ResourceContents::TextResourceContents { text, uri, .. }) = result.contents.first() else { - panic!("expected text resource contents"); - }; - assert_eq!("file:///password.env", uri); - assert_ne!(BACKEND_RESOURCE_SECRET, text); - assert!(text.contains(REDACTED)); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -410,13 +389,7 @@ async fn write_redis_config(redis_port: u16, backend: &RunningBackend) { }, )]), tools: HashMap::new(), - resources: HashMap::from([( - "file:///password.env".to_owned(), - ServiceRoute { - backend_name: "backend".to_owned(), - upstream_name: "file:///password.env".to_owned(), - }, - )]), + resources: HashMap::new(), resource_templates: HashMap::new(), prompts: HashMap::new(), }, @@ -443,7 +416,7 @@ async fn write_runtime_plugin_config(redis_port: u16, plugin_config: Value) { "plugins": [{ "name": "secrets-detection", "kind": "validator/secrets-detection", - "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke", "cmf.resource_post_fetch"], + "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], "config": plugin_config, }] } From 05c86581d17e9e4303c0f431b90769665d01b84a Mon Sep 17 00:00:00 2001 From: lucarlig Date: Fri, 4 Sep 2026 22:11:42 +0100 Subject: [PATCH 5/5] fix: align resource hooks with built-in payload behavior Signed-off-by: lucarlig --- _context/wiki/config.md | 6 +- _context/wiki/testing.md | 2 +- .../contextforge-data-plane-cpex/src/cmf.rs | 379 ++++++++---------- .../src/handle.rs | 24 +- .../src/pipeline.rs | 24 +- .../src/runtime.rs | 33 +- .../src/gateway/mcp_service/resources.rs | 25 +- .../tests/gateway/harness/plugin.rs | 28 +- .../tests/gateway/harness/plugin_gateway.rs | 14 +- .../tests/gateway/plugins.rs | 79 +++- 10 files changed, 339 insertions(+), 275 deletions(-) diff --git a/_context/wiki/config.md b/_context/wiki/config.md index e3b686e9..10d57bf2 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -212,11 +212,13 @@ 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 as decoded CMF bytes and are encoded back to MCP base64 after an edit. Unchanged blobs retain the backend's exact wire representation. Invalid embedded blobs or unsupported prompt content fail the entire prompt before post hooks; they must never disappear from the policy payload while remaining in the client response. +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 runs after routing and may allow or deny the canonical backend-local URI, but cannot change it. The post hook can redact text resources, transform binary resources, or deny the response. URI, MIME type, item count, CMF schema version, and channel must remain stable; unsupported or lossy edits fail closed. +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. diff --git a/_context/wiki/testing.md b/_context/wiki/testing.md index 38b2b9b0..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. Resource coverage includes direct and aliased URIs, text/blob redaction, canonical pre-hook URIs, pre/post denial, and invalid embedded prompt blobs. 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. | diff --git a/crates/contextforge-data-plane-cpex/src/cmf.rs b/crates/contextforge-data-plane-cpex/src/cmf.rs index 0f0308ff..010606c8 100644 --- a/crates/contextforge-data-plane-cpex/src/cmf.rs +++ b/crates/contextforge-data-plane-cpex/src/cmf.rs @@ -55,24 +55,6 @@ pub(crate) fn resource_request_payload(resource_uri: &str, resource_request_id: } } -pub(crate) fn resource_request_matches( - payload: &MessagePayload, - resource_uri: &str, - resource_request_id: &str, -) -> bool { - let [ContentPart::ResourceRef { content }] = payload.message.content.as_slice() else { - return false; - }; - canonical_message_envelope(payload, Role::User) - && content.resource_request_id == resource_request_id - && content.uri == resource_uri - && matches!(content.resource_type, ResourceType::Uri) - && content.name.is_none() - && content.range_start.is_none() - && content.range_end.is_none() - && content.selector.is_none() -} - pub(crate) fn resource_result_payload( response: &ReadResourceResult, resource_request_id: &str, @@ -113,56 +95,47 @@ fn cmf_resource_content(content: &ResourceContents, resource_request_id: &str) - pub(crate) fn resource_result_response( mut original: ReadResourceResult, payload: &MessagePayload, - resource_request_id: &str, ) -> Option { - if !canonical_message_envelope(payload, Role::Assistant) || payload.message.content.len() != original.contents.len() - { + // 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 ContentPart::Resource { content } = modified else { return None }; + let meta = match original { + ResourceContents::TextResourceContents { meta, .. } + | ResourceContents::BlobResourceContents { meta, .. } => meta.clone(), + _ => return None, }; - if content.resource_request_id != resource_request_id - || !matches!(content.resource_type, ResourceType::Uri) - || content.name.is_some() - || content.description.is_some() - || content.size_bytes.is_some() - || !content.annotations.is_empty() - || content.version.is_some() - { - return None; - } - match original { - ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { - if content.uri != *uri || content.mime_type != *mime_type || content.blob.is_some() { - return None; - } - *text = content.content.clone()?; + *original = match (&content.content, &content.blob) { + (Some(text), _) => ResourceContents::TextResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + text: text.clone(), + meta, }, - ResourceContents::BlobResourceContents { uri, mime_type, blob, .. } => { - if content.uri != *uri || content.mime_type != *mime_type || content.content.is_some() { - return None; - } - let modified_blob = content.blob.as_ref()?; - let original_blob = BASE64_STANDARD.decode(blob.as_bytes()).ok()?; - if modified_blob != &original_blob { - *blob = BASE64_STANDARD.encode(modified_blob); + (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) } -fn canonical_message_envelope(payload: &MessagePayload, role: Role) -> bool { - payload.message.schema_version == SCHEMA_VERSION - && payload.message.role == role - && payload.message.channel.is_none() -} - pub(crate) fn tool_result_payload(tool_name: &str, response: &CallToolResult, tool_call_id: &str) -> MessagePayload { tool_json_result_payload( tool_name, @@ -292,16 +265,13 @@ pub(crate) fn prompt_result_payload( response: &GetPromptResult, prompt_name: &str, prompt_request_id: &str, -) -> Option { - let messages = response - .messages - .iter() - .map(|message| cmf_prompt_message(message, prompt_request_id)) - .collect::>>()?; +) -> MessagePayload { + let messages = + response.messages.iter().map(|message| cmf_prompt_message(message, prompt_request_id)).collect::>(); - Some(MessagePayload { + MessagePayload { message: Message { - schema_version: SCHEMA_VERSION.to_owned(), + schema_version: "2.0".to_owned(), role: Role::Assistant, content: vec![ContentPart::PromptResult { content: PromptResult { @@ -315,7 +285,7 @@ pub(crate) fn prompt_result_payload( }], channel: None, }, - }) + } } fn prompt_result(payload: &MessagePayload) -> Option<&PromptResult> { @@ -351,13 +321,13 @@ pub(crate) fn prompt_result_response( } for (message, edited) in original.messages.iter_mut().zip(&result.messages) { - let projected = cmf_prompt_message(message, prompt_request_id)?; + let projected = cmf_prompt_message(message, prompt_request_id); if serde_json::to_value(&projected).ok()? == serde_json::to_value(edited).ok()? { continue; } let rebuilt = mcp_prompt_message(edited)?; - if serde_json::to_value(cmf_prompt_message(&rebuilt, prompt_request_id)?).ok()? + if serde_json::to_value(cmf_prompt_message(&rebuilt, prompt_request_id)).ok()? != serde_json::to_value(edited).ok()? { return None; @@ -368,16 +338,16 @@ pub(crate) fn prompt_result_response( Some(original) } -fn cmf_prompt_message(message: &PromptMessage, prompt_request_id: &str) -> Option { - Some(Message { - schema_version: SCHEMA_VERSION.to_owned(), +fn cmf_prompt_message(message: &PromptMessage, prompt_request_id: &str) -> Message { + Message { + schema_version: "2.0".to_owned(), role: match message.role { McpRole::Assistant => Role::Assistant, McpRole::User => Role::User, }, - content: vec![cmf_content_part(&message.content, prompt_request_id)?], + content: cmf_content_part(&message.content, prompt_request_id).into_iter().collect(), channel: None, - }) + } } fn cmf_content_part(block: &ContentBlock, prompt_request_id: &str) -> Option { @@ -399,7 +369,28 @@ fn cmf_content_part(block: &ContentBlock, prompt_request_id: &str) -> Option { - ContentPart::Resource { content: cmf_resource_content(&resource.resource, prompt_request_id)? } + let (uri, mime_type, content) = match &resource.resource { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => { + (uri.clone(), mime_type.clone(), Some(text.clone())) + }, + ResourceContents::BlobResourceContents { uri, mime_type, .. } => (uri.clone(), mime_type.clone(), None), + _ => return None, + }; + ContentPart::Resource { + content: CmfResource { + resource_request_id: prompt_request_id.to_owned(), + uri, + name: None, + description: None, + resource_type: ResourceType::Uri, + content, + blob: None, + mime_type, + size_bytes: None, + annotations: HashMap::new(), + version: None, + }, + } }, ContentBlock::ResourceLink(link) => ContentPart::ResourceRef { content: ResourceReference { @@ -440,7 +431,12 @@ fn mcp_prompt_message(message: &Message) -> Option { ContentPart::Audio { content } => { ContentBlock::audio(inline_media_data(&content.source_type, &content.data)?, content.media_type.clone()?) }, - ContentPart::Resource { content } => ContentBlock::resource(mcp_resource_content(content)?), + ContentPart::Resource { content } => ContentBlock::resource(ResourceContents::TextResourceContents { + uri: content.uri.clone(), + mime_type: content.mime_type.clone(), + text: content.content.clone()?, + meta: None, + }), ContentPart::ResourceRef { content } => { ContentBlock::ResourceLink(McpResource::new(content.uri.clone(), content.name.clone()?)) }, @@ -450,56 +446,12 @@ fn mcp_prompt_message(message: &Message) -> Option { Some(PromptMessage::new(role, content)) } -fn mcp_resource_content(content: &CmfResource) -> Option { - match (&content.content, &content.blob) { - (Some(text), None) => Some(ResourceContents::TextResourceContents { - uri: content.uri.clone(), - mime_type: content.mime_type.clone(), - text: text.clone(), - meta: None, - }), - (None, Some(blob)) => Some(ResourceContents::BlobResourceContents { - uri: content.uri.clone(), - mime_type: content.mime_type.clone(), - blob: BASE64_STANDARD.encode(blob), - meta: None, - }), - _ => None, - } -} - #[cfg(test)] mod tests { - use cpex::cpex_core::cmf::Channel; - use super::*; #[test] - fn resource_request_rejects_uri_mutation() { - let mut payload = resource_request_payload("file:///password.env", "resource-1"); - let ContentPart::ResourceRef { content } = &mut payload.message.content[0] else { - panic!("expected resource reference"); - }; - content.uri = "file:///other.env".to_owned(); - - assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); - } - - #[test] - fn resource_request_rejects_modified_envelope() { - let mut payload = resource_request_payload("file:///password.env", "resource-1"); - payload.message.schema_version = "3.0".to_owned(); - - assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); - - let mut payload = resource_request_payload("file:///password.env", "resource-1"); - payload.message.channel = Some(Channel::Analysis); - - assert!(!resource_request_matches(&payload, "file:///password.env", "resource-1")); - } - - #[test] - fn resource_result_response_applies_only_text_changes() { + 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"); @@ -508,7 +460,7 @@ mod tests { }; content.content = Some("AWS_ACCESS_KEY_ID=[redacted]".to_owned()); - let result = resource_result_response(original, &payload, "resource-1").expect("text edit applies"); + let result = resource_result_response(original, &payload).expect("text edit applies"); let ResourceContents::TextResourceContents { text, uri, .. } = &result.contents[0] else { panic!("expected text resource"); @@ -517,18 +469,6 @@ mod tests { assert_eq!("file:///password.env", uri); } - #[test] - fn resource_result_response_rejects_uri_mutation() { - let original = ReadResourceResult::new(vec![ResourceContents::text("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.uri = "file:///other.env".to_owned(); - - assert!(resource_result_response(original, &payload, "resource-1").is_none()); - } - #[test] fn resource_result_response_decodes_and_applies_blob_changes() { let wire_blob = BASE64_STANDARD.encode(b"AWS_ACCESS_KEY_ID=secret"); @@ -543,7 +483,7 @@ mod tests { 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, "resource-1").expect("blob edit applies"); + let result = resource_result_response(original, &payload).expect("blob edit applies"); let ResourceContents::BlobResourceContents { blob, uri, .. } = &result.contents[0] else { panic!("expected blob resource"); @@ -559,7 +499,7 @@ mod tests { 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, "resource-1").expect("unchanged blob applies"); + let result = resource_result_response(original, &payload).expect("unchanged blob applies"); let ResourceContents::BlobResourceContents { blob, .. } = &result.contents[0] else { panic!("expected blob resource"); @@ -575,17 +515,81 @@ mod tests { } #[test] - fn resource_result_response_rejects_modified_envelope() { - let original = ReadResourceResult::new(vec![ResourceContents::text("secret", "file:///password.env")]); - let mut payload = resource_result_payload(&original, "resource-1").expect("resource result is supported"); - payload.message.schema_version = "3.0".to_owned(); + 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"]); + } - assert!(resource_result_response(original.clone(), &payload, "resource-1").is_none()); + #[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); + } - let mut payload = resource_result_payload(&original, "resource-1").expect("resource result is supported"); - payload.message.channel = Some(Channel::Final); + #[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); + } - assert!(resource_result_response(original, &payload, "resource-1").is_none()); + #[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 { @@ -611,7 +615,7 @@ mod tests { #[test] fn prompt_result_response_rejects_added_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let extra = edited_messages(&mut payload).first().cloned().expect("one message"); edited_messages(&mut payload).push(extra); @@ -621,7 +625,7 @@ mod tests { #[test] fn prompt_result_response_rejects_extra_prompt_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let duplicate = payload.message.content[0].clone(); payload.message.content.push(duplicate); @@ -631,7 +635,7 @@ mod tests { #[test] fn prompt_result_rejection_reports_the_plugin_error_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let result = prompt_result_mut(&mut payload); result.is_error = true; result.error_message = Some("blocked by policy".to_owned()); @@ -642,7 +646,7 @@ mod tests { #[test] fn prompt_result_rejection_falls_back_when_the_plugin_gives_no_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); prompt_result_mut(&mut payload).is_error = true; assert_eq!(Some("Plugin rejected the rendered prompt".to_owned()), prompt_result_rejection(&payload)); @@ -651,7 +655,7 @@ mod tests { #[test] fn prompt_result_rejection_is_absent_for_a_normal_result() { let original = text_prompt(); - let payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let payload = prompt_result_payload(&original, "review", "prompt-1"); assert_eq!(None, prompt_result_rejection(&payload)); } @@ -720,7 +724,7 @@ mod tests { #[test] fn prompt_result_response_rejects_envelope_content_edit() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); prompt_result_mut(&mut payload).content = Some("[REDACTED]".to_owned()); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -729,7 +733,7 @@ mod tests { #[test] fn prompt_result_response_rejects_renamed_prompt() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); prompt_result_mut(&mut payload).prompt_name = "other".to_owned(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -738,7 +742,7 @@ mod tests { #[test] fn prompt_result_response_rejects_recorrelated_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); prompt_result_mut(&mut payload).prompt_request_id = "prompt-2".to_owned(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -747,7 +751,7 @@ mod tests { #[test] fn prompt_result_response_rejects_error_message_without_error_flag() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); prompt_result_mut(&mut payload).error_message = Some("blocked".to_owned()); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -763,7 +767,7 @@ mod tests { #[test] fn prompt_result_response_rejects_resource_type_edit() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -775,7 +779,7 @@ mod tests { #[test] fn prompt_result_response_rejects_dropped_resource_metadata() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -791,7 +795,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_an_image_edit() { let original = media_prompt(ContentBlock::image("aW1hZ2U=", "image/png")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -807,7 +811,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_an_audio_edit() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("audio reaches the plugin as a CMF audio part"); }; @@ -823,7 +827,7 @@ mod tests { #[test] fn prompt_result_response_rejects_url_sourced_audio() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an audio part"); }; @@ -836,7 +840,7 @@ mod tests { #[test] fn prompt_result_response_rejects_audio_without_media_type() { let original = media_prompt(ContentBlock::audio("YXVkaW8=", "audio/mp3")); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Audio { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an audio part"); }; @@ -848,7 +852,7 @@ mod tests { #[test] fn prompt_result_response_round_trips_a_resource_link_edit() { let original = media_prompt(ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env"))); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -864,7 +868,7 @@ mod tests { #[test] fn prompt_result_response_rejects_resource_with_removed_text() { let original = resource_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource part"); }; @@ -876,7 +880,7 @@ mod tests { #[test] fn prompt_result_response_rejects_multiple_content_parts() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); edited_messages(&mut payload)[0].content.push(ContentPart::Text { text: "extra".to_owned() }); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -885,7 +889,7 @@ mod tests { #[test] fn prompt_result_response_rejects_a_cmf_only_content_part() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); edited_messages(&mut payload)[0].content = vec![ContentPart::Thinking { text: "reasoning".to_owned() }]; assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -894,7 +898,7 @@ mod tests { #[test] fn prompt_result_response_rejects_a_payload_without_a_prompt_result() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); payload.message.content.clear(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -912,7 +916,7 @@ mod tests { fn prompt_result_response_rejects_url_sourced_image() { let original = GetPromptResult::new(vec![PromptMessage::new(McpRole::User, ContentBlock::image("aW1hZ2U=", "image/png"))]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -926,7 +930,7 @@ mod tests { fn prompt_result_response_rejects_image_without_media_type() { let original = GetPromptResult::new(vec![PromptMessage::new(McpRole::User, ContentBlock::image("aW1hZ2U=", "image/png"))]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Image { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected an image part"); }; @@ -941,7 +945,7 @@ mod tests { McpRole::User, ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -956,7 +960,7 @@ mod tests { McpRole::User, ContentBlock::ResourceLink(McpResource::new("file:///app.env", "app-env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::ResourceRef { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("expected a resource reference part"); }; @@ -968,7 +972,7 @@ mod tests { #[test] fn prompt_result_response_rejects_removed_message() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); edited_messages(&mut payload).clear(); assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); @@ -977,28 +981,16 @@ mod tests { #[test] fn prompt_result_response_rejects_unmappable_role() { let original = text_prompt(); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); edited_messages(&mut payload)[0].role = Role::System; assert!(prompt_result_response(original, &payload, "review", "prompt-1").is_none()); } - #[test] - fn prompt_result_payload_rejects_invalid_embedded_blob() { - let original = GetPromptResult::new(vec![ - PromptMessage::new_text(McpRole::User, "ordinary text"), - PromptMessage::new( - McpRole::User, - ContentBlock::resource(ResourceContents::blob("not base64!", "file:///blocked.env")), - ), - ]); - assert!(prompt_result_payload(&original, "review", "prompt-1").is_none()); - } - #[test] fn prompt_result_response_preserves_unmodified_messages() { let original = text_prompt(); - let payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let payload = prompt_result_payload(&original, "review", "prompt-1"); let result = prompt_result_response(original.clone(), &payload, "review", "prompt-1") .expect("unmodified payload applies"); @@ -1015,7 +1007,7 @@ mod tests { McpRole::User, ContentBlock::resource(ResourceContents::text("token=secret", "file:///app.env")), )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); + let mut payload = prompt_result_payload(&original, "review", "prompt-1"); let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { panic!("embedded resource reaches the plugin as a CMF resource part"); @@ -1035,31 +1027,6 @@ mod tests { assert_eq!("file:///app.env", uri); } - #[test] - fn prompt_result_response_round_trips_embedded_blob_resource() { - let original = GetPromptResult::new(vec![PromptMessage::new( - McpRole::User, - ContentBlock::resource(ResourceContents::blob(BASE64_STANDARD.encode(b"token=secret"), "file:///app.bin")), - )]); - let mut payload = prompt_result_payload(&original, "review", "prompt-1").expect("supported prompt content"); - let ContentPart::Resource { content } = &mut edited_messages(&mut payload)[0].content[0] else { - panic!("embedded resource reaches the plugin as a CMF resource part"); - }; - assert_eq!(Some(b"token=secret".as_slice()), content.blob.as_deref()); - content.blob = Some(b"token=[REDACTED]".to_vec()); - - let result = prompt_result_response(original, &payload, "review", "prompt-1").expect("resource edit applies"); - - let ContentBlock::Resource(resource) = &result.messages[0].content else { - panic!("expected an embedded resource"); - }; - let ResourceContents::BlobResourceContents { blob, uri, .. } = &resource.resource else { - panic!("expected blob resource contents"); - }; - assert_eq!(b"token=[REDACTED]", BASE64_STANDARD.decode(blob).expect("valid base64").as_slice()); - assert_eq!("file:///app.bin", uri); - } - #[test] fn tool_result_response_uses_cmf_error_flag_for_nested_mcp_result() { let original = CallToolResult::success(vec![ContentBlock::text("original")]); diff --git a/crates/contextforge-data-plane-cpex/src/handle.rs b/crates/contextforge-data-plane-cpex/src/handle.rs index 5ebebb1a..f6e8e507 100644 --- a/crates/contextforge-data-plane-cpex/src/handle.rs +++ b/crates/contextforge-data-plane-cpex/src/handle.rs @@ -53,11 +53,18 @@ struct RegistryResourceCallState { } /// Captures the resource post-hook decision and runtime for one request. -pub struct ResourceHookState(Option); +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.0 { + match self.call { Some(call) => call.runtime.after_read_resource(response, call.state).await, None => Ok(response), } @@ -304,12 +311,11 @@ impl GatewayPluginRuntimeHandle { let RuntimeState::Active(runtime) = state.as_ref() else { return Err(runtime_failed_error(state.as_ref())); }; - Ok(ResourceHookState( - runtime - .before_read_resource(resource_uri) - .await? - .map(|state| RegistryResourceCallState { runtime: Arc::clone(runtime), state }), - )) + 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( @@ -827,7 +833,7 @@ mod tests { 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.0.is_none(), "no post-hook state allocation"); + 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"); diff --git a/crates/contextforge-data-plane-cpex/src/pipeline.rs b/crates/contextforge-data-plane-cpex/src/pipeline.rs index a03081ca..93abb715 100644 --- a/crates/contextforge-data-plane-cpex/src/pipeline.rs +++ b/crates/contextforge-data-plane-cpex/src/pipeline.rs @@ -10,8 +10,8 @@ use tracing::warn; use crate::{ PromptArgumentsUpdate, ToolArgumentsUpdate, cmf::{ - prompt_request_arguments, prompt_result_rejection, prompt_result_response, resource_request_matches, - resource_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,30 +97,22 @@ pub(crate) fn effective_post_prompt_result( }) } -pub(crate) fn validate_pre_resource_result( - result: &PipelineResult, - resource_uri: &str, - resource_request_id: &str, -) -> Result<(), ErrorData> { - let Some(payload) = modified_message_payload(result) else { - return Ok(()); +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)); }; - if resource_request_matches(payload, resource_uri, resource_request_id) { - Ok(()) - } else { - Err(ErrorData::internal_error("Plugin attempted to modify the canonical resource route", None)) - } + Ok(Some(content.uri.clone())) } pub(crate) fn effective_post_resource_result( original: ReadResourceResult, result: &PipelineResult, - resource_request_id: &str, ) -> Result { let Some(payload) = modified_message_payload(result) else { return Ok(original); }; - resource_result_response(original, payload, resource_request_id) + resource_result_response(original, payload) .ok_or_else(|| ErrorData::internal_error("Plugin returned a resource result the gateway cannot apply", None)) } diff --git a/crates/contextforge-data-plane-cpex/src/runtime.rs b/crates/contextforge-data-plane-cpex/src/runtime.rs index 65ef1e4e..2d4b4914 100644 --- a/crates/contextforge-data-plane-cpex/src/runtime.rs +++ b/crates/contextforge-data-plane-cpex/src/runtime.rs @@ -29,8 +29,8 @@ use crate::{ hooks::{PromptPreFetchResult, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, pipeline::{ effective_post_json, effective_post_prompt_result, effective_post_resource_result, effective_post_result, - effective_pre_args, effective_pre_prompt_args, log_pipeline_errors, plugin_denied_error, - validate_pre_resource_result, + effective_pre_args, effective_pre_prompt_args, effective_pre_resource_uri, log_pipeline_errors, + plugin_denied_error, }, }; @@ -248,14 +248,17 @@ impl GatewayPluginRuntime { pub(crate) async fn before_read_resource( &self, resource_uri: &str, - ) -> Result, ErrorData> { + ) -> Result<(Option, Option), ErrorData> { if !self.hooks.resource.pre && !self.hooks.resource.post { - return Ok(None); + return Ok((None, None)); } let resource_request_id = next_resource_request_id(); if !self.hooks.resource.pre { - return Ok(Some(ResourceCallState { context_table: PluginContextTable::default(), resource_request_id })); + return Ok(( + None, + Some(ResourceCallState { context_table: PluginContextTable::default(), resource_request_id }), + )); } let payload = resource_request_payload(resource_uri, &resource_request_id); @@ -263,12 +266,14 @@ impl GatewayPluginRuntime { if pre_result.is_denied() { return Err(plugin_denied_error("resource", pre_result)); } - validate_pre_resource_result(&pre_result, resource_uri, &resource_request_id)?; - Ok(self - .hooks - .resource - .post - .then_some(ResourceCallState { context_table: pre_result.context_table, resource_request_id })) + 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( @@ -284,9 +289,7 @@ impl GatewayPluginRuntime { let state = state.and_then(|state| state.downcast::().ok()); let Some(state) = state else { return Ok(response) }; - let payload = prompt_result_payload(&response, prompt_name, &state.prompt_request_id).ok_or_else(|| { - ErrorData::internal_error("Prompt response contains unsupported or invalid content", None) - })?; + let payload = prompt_result_payload(&response, prompt_name, &state.prompt_request_id); 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() { @@ -308,7 +311,7 @@ impl GatewayPluginRuntime { if post_result.is_denied() { return Err(plugin_denied_error("resource", post_result)); } - effective_post_resource_result(response, &post_result, &state.resource_request_id) + effective_post_resource_result(response, &post_result) } pub(crate) async fn after_tool_call( 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 88db4948..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,20 +30,31 @@ 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(), data: None, })?; - let resource_hook = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { - Some(plugin_runtime.before_read_resource(&resource_uri).await?) - } else { - None - }; let mut backend_service = connect_backend_for_request(mcp_service, &backend_name, backend, &cx).await?; let mut routed_request = request; 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 0eb9ee80..5b9b5c1c 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/plugin.rs @@ -40,6 +40,7 @@ pub(crate) struct Observations { #[derive(Clone, Copy, Default)] pub(crate) enum PreBehavior { + ResourceUri(&'static str), #[default] Allow, Rewrite, @@ -50,6 +51,7 @@ pub(crate) enum PreBehavior { #[derive(Clone, Copy, Default)] pub(crate) enum PostBehavior { + ResourceText, #[default] Allow, Rewrite, @@ -82,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 @@ -180,10 +192,17 @@ 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); } @@ -260,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(); 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 a329ae5e..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 @@ -124,13 +124,6 @@ impl ServerHandler for TestBackend { .push(BackendObservation { tool_name: request.name.clone(), args: request.arguments.clone() }); self.state.events.lock().expect("backend events lock poisoned").push("backend"); - if request.name == "review_invalid_blob" { - return Ok(GetPromptResult::new(vec![PromptMessage::new( - Role::User, - ContentBlock::resource(ResourceContents::blob("not base64!", "file:///blocked.env")), - )]) - .into()); - } let topic = request .arguments .as_ref() @@ -261,9 +254,10 @@ pub const TOOL_NAMES: &[&str] = &[ "wait_for_cancellation", ]; pub const RESOURCE_URIS: &[&str] = &["file:///password.env", "file:///password.bin"]; -pub const PROMPT_NAMES: &[&str] = &["review_bundle", "review", "review_invalid_blob"]; +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, @@ -466,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], @@ -476,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 22f1ad91..c97cefc6 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/plugins.rs @@ -930,18 +930,81 @@ async fn resource_post_hook_denies_the_backend_response() { } #[tokio::test] -async fn invalid_prompt_blob_is_rejected_before_post_policy() { - let plugin = Arc::new(PromptTestPlugin::new("prompt", vec![cmf_hook_names::PROMPT_POST_FETCH])); +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_prompt_plugin(plugin).await).await; + 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 - .get_prompt(GetPromptRequestParams::new("review_invalid_blob")) + .read_resource(ReadResourceRequestParams::new("file:///password.env")) .await - .expect_err("invalid content fails closed"); + .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::INTERNAL_ERROR, code); - assert_eq!("Prompt response contains unsupported or invalid content", message); - assert_eq!(0, observations.lock().expect("observations lock").post_calls); + 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()); }