diff --git a/crates/contextforge-gateway-rs-cpex/src/cmf.rs b/crates/contextforge-gateway-rs-cpex/src/cmf.rs index b4de6f6..08f68e2 100644 --- a/crates/contextforge-gateway-rs-cpex/src/cmf.rs +++ b/crates/contextforge-gateway-rs-cpex/src/cmf.rs @@ -1,5 +1,17 @@ -use cpex::cpex_core::cmf::{ContentPart, Message, MessagePayload, Role, ToolCall, ToolResult}; -use rmcp::model::{CallToolRequestParams, CallToolResult, ContentBlock}; +use std::collections::HashMap; + +use cpex::cpex_core::cmf::{ + ContentPart, Message, MessagePayload, PromptRequest, PromptResult, Resource, ResourceReference, ResourceType, Role, + ToolCall, ToolResult, +}; +use rmcp::{ + ErrorData, + model::{ + ArgumentInfo, CallToolRequestParams, CallToolResult, CompleteResult, ContentBlock, ErrorCode, + GetPromptRequestParams, GetPromptResult, Prompt, PromptMessage, ReadResourceResult, Resource as McpResource, + ResourceContents, ResourceTemplate, Role as McpRole, + }, +}; use serde_json::{Map, Value}; pub(crate) fn tool_call_payload( @@ -110,6 +122,525 @@ fn raw_error_tool_result(value: Value) -> CallToolResult { } } +/// Builds the `cmf.prompt_pre_fetch` payload. Mirrors [`tool_call_payload`]: the hook runs after +/// routing, so `prompt_name` is already backend-local and `backend_name` carries the namespace. +pub(crate) fn prompt_request_payload( + request: &GetPromptRequestParams, + prompt_name: &str, + backend_name: &str, + prompt_request_id: &str, +) -> MessagePayload { + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::User, + content: vec![ContentPart::PromptRequest { + content: PromptRequest { + prompt_request_id: prompt_request_id.to_owned(), + name: prompt_name.to_owned(), + arguments: request.arguments.clone().unwrap_or_default().into_iter().collect(), + server_id: Some(backend_name.to_owned()), + }, + }], + channel: None, + }, + } +} + +/// Builds the `cmf.prompt_post_fetch` payload. +/// +/// CMF `PromptResult` is typed rather than free-form JSON, so MCP messages are mapped into CMF +/// messages one-for-one. Text blocks are exposed as [`ContentPart::Text`]; other block kinds +/// (image, audio, embedded resource, resource link) become a message with empty content. That +/// keeps index alignment with the original response without claiming to expose content the +/// mapping cannot write back. +pub(crate) fn prompt_result_payload( + prompt_name: &str, + result: &GetPromptResult, + prompt_request_id: &str, +) -> MessagePayload { + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::Assistant, + content: vec![ContentPart::PromptResult { + content: PromptResult { + prompt_request_id: prompt_request_id.to_owned(), + prompt_name: prompt_name.to_owned(), + messages: result.messages.iter().map(prompt_message_to_cmf).collect(), + content: None, + is_error: false, + error_message: None, + }, + }], + channel: None, + }, + } +} + +fn prompt_message_to_cmf(message: &PromptMessage) -> Message { + let content = match &message.content { + ContentBlock::Text(text) => vec![ContentPart::Text { text: text.text.clone() }], + _ => Vec::new(), + }; + Message { + schema_version: "2.0".to_owned(), + role: match message.role { + McpRole::User => Role::User, + McpRole::Assistant => Role::Assistant, + }, + content, + channel: None, + } +} + +/// Builds the `cmf.prompt_pre_fetch` payload for `prompts/list`. +/// +/// A listing has no single prompt and no single backend, so the request carries an **empty name** +/// and no `server_id`. Plugins discriminate a listing from a `prompts/get` on that empty name. +pub(crate) fn prompt_list_request_payload(cursor: Option<&str>, prompt_request_id: &str) -> MessagePayload { + let arguments = cursor + .map(|cursor| HashMap::from([("cursor".to_owned(), Value::String(cursor.to_owned()))])) + .unwrap_or_default(); + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::User, + content: vec![ContentPart::PromptRequest { + content: PromptRequest { + prompt_request_id: prompt_request_id.to_owned(), + name: String::new(), + arguments, + server_id: None, + }, + }], + channel: None, + }, + } +} + +/// Builds the `cmf.prompt_post_fetch` payload for `prompts/list`, over the merged and namespaced +/// listing the client is about to receive. +/// +/// Exposure is read-only: MCP `Prompt` is metadata (`title`, `arguments`, `icons`, `_meta`) with no +/// CMF equivalent, so a modified payload cannot be written back without silently dropping fields. +/// Each prompt contributes its exposed name, plus its description as text when it has one. +pub(crate) fn prompt_list_result_payload(prompts: &[Prompt], prompt_request_id: &str) -> MessagePayload { + let mut content = Vec::with_capacity(prompts.len()); + for prompt in prompts { + content.push(ContentPart::PromptRequest { + content: PromptRequest { + prompt_request_id: prompt_request_id.to_owned(), + name: prompt.name.clone(), + arguments: HashMap::new(), + server_id: None, + }, + }); + if let Some(description) = &prompt.description { + content.push(ContentPart::Text { text: description.clone() }); + } + } + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +/// Builds the `cmf.resource_pre_fetch` payload for the resource listings — `resources/list` and +/// `resources/templates/list`. +/// +/// Like the prompt listing, a resource listing has no single resource and no single backend, so the +/// reference carries an **empty URI**. Plugins discriminate a listing on that empty URI. The +/// pagination cursor is not exposed: CMF `ResourceReference` has no field for request arguments, +/// and overloading `selector` would misrepresent what that field means. +pub(crate) fn resource_list_request_payload(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: String::new(), + name: None, + resource_type: ResourceType::Uri, + range_start: None, + range_end: None, + selector: None, + }, + }], + channel: None, + }, + } +} + +/// Builds the `cmf.resource_post_fetch` payload for `resources/templates/list`, over the merged and +/// namespaced listing the client is about to receive. +/// +/// Read-only for the same reason as the prompt listing: MCP `ResourceTemplate` metadata (`title`, +/// `mime_type`, `icons`, `annotations`, `_meta`) has no CMF equivalent to be rebuilt from. Each +/// template contributes its URI template and name, plus its description as text when it has one. +pub(crate) fn resource_template_list_result_payload( + templates: &[ResourceTemplate], + resource_request_id: &str, +) -> MessagePayload { + let mut content = Vec::with_capacity(templates.len()); + for template in templates { + content.push(ContentPart::ResourceRef { + content: ResourceReference { + resource_request_id: resource_request_id.to_owned(), + uri: template.uri_template.clone(), + name: Some(template.name.clone()), + resource_type: ResourceType::Uri, + range_start: None, + range_end: None, + selector: None, + }, + }); + if let Some(description) = &template.description { + content.push(ContentPart::Text { text: description.clone() }); + } + } + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +/// CMF annotation key carrying the backend a resource operation routes to. +/// +/// Neither CMF resource type has a backend field the way `ToolCall.namespace` and +/// `PromptRequest.server_id` do, so the backend is exposed through `Resource.annotations` — the +/// open metadata map CMF provides for exactly this kind of attribution. +pub(crate) const RESOURCE_BACKEND_ANNOTATION: &str = "backend"; + +/// Builds the `cmf.resource_pre_fetch` payload for the routed resource paths — `resources/read`, +/// `resources/subscribe`, and `resources/unsubscribe`. The hook runs after routing, so +/// `resource_uri` is backend-local and `backend_name` carries the namespace, matching the +/// `call_tool` contract. +/// +/// `subscribe` and `unsubscribe` have no matching post hook: both return an empty +/// acknowledgement, so a post hook would receive nothing to inspect or rewrite. +pub(crate) fn resource_request_payload( + resource_uri: &str, + backend_name: &str, + resource_request_id: &str, +) -> MessagePayload { + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::User, + content: vec![ContentPart::Resource { + content: Resource { + resource_request_id: resource_request_id.to_owned(), + uri: resource_uri.to_owned(), + resource_type: ResourceType::Uri, + annotations: HashMap::from([( + RESOURCE_BACKEND_ANNOTATION.to_owned(), + Value::String(backend_name.to_owned()), + )]), + ..Resource::default() + }, + }], + channel: None, + }, + } +} + +/// Builds the completion pre-hook payload for a prompt reference. The argument being completed is +/// merged into the prompt arguments alongside any already-resolved context arguments. +/// +/// Read-only: the argument value is a partial user keystroke, so rewriting it would return +/// completions for text the user never typed. Denial is supported. +pub(crate) fn prompt_completion_payload( + prompt_name: &str, + backend_name: &str, + argument: &ArgumentInfo, + context: Option<&HashMap>, + prompt_request_id: &str, +) -> MessagePayload { + let mut arguments = completion_arguments(context); + arguments.insert(argument.name.clone(), Value::String(argument.value.clone())); + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::User, + content: vec![ContentPart::PromptRequest { + content: PromptRequest { + prompt_request_id: prompt_request_id.to_owned(), + name: prompt_name.to_owned(), + arguments, + server_id: Some(backend_name.to_owned()), + }, + }], + channel: None, + }, + } +} + +/// Builds the completion pre-hook payload for a resource or template reference. CMF resource types +/// have no argument map, so the completion argument and context ride in `Resource.annotations` +/// alongside the backend. Read-only, same as the prompt form. +pub(crate) fn resource_completion_payload( + resource_uri: &str, + backend_name: &str, + argument: &ArgumentInfo, + context: Option<&HashMap>, + resource_request_id: &str, +) -> MessagePayload { + let annotations = HashMap::from([ + (RESOURCE_BACKEND_ANNOTATION.to_owned(), Value::String(backend_name.to_owned())), + ( + COMPLETION_ARGUMENT_ANNOTATION.to_owned(), + serde_json::json!({ "name": argument.name, "value": argument.value }), + ), + (COMPLETION_CONTEXT_ANNOTATION.to_owned(), Value::Object(completion_arguments(context).into_iter().collect())), + ]); + MessagePayload { + message: Message { + schema_version: "2.0".to_owned(), + role: Role::User, + content: vec![ContentPart::Resource { + content: Resource { + resource_request_id: resource_request_id.to_owned(), + uri: resource_uri.to_owned(), + resource_type: ResourceType::Uri, + annotations, + ..Resource::default() + }, + }], + channel: None, + }, + } +} + +/// CMF annotation keys carrying the completion argument and its resolved context on resource +/// references, which have no argument map of their own. +pub(crate) const COMPLETION_ARGUMENT_ANNOTATION: &str = "completion_argument"; +pub(crate) const COMPLETION_CONTEXT_ANNOTATION: &str = "completion_context"; + +fn completion_arguments(context: Option<&HashMap>) -> HashMap { + context + .map(|context| context.iter().map(|(name, value)| (name.clone(), Value::String(value.clone()))).collect()) + .unwrap_or_default() +} + +/// Builds the completion post-hook payload for a prompt reference: the reference itself, followed +/// by one text part per completion value. +pub(crate) fn prompt_completion_result_payload( + prompt_name: &str, + values: &[String], + prompt_request_id: &str, +) -> MessagePayload { + let mut content = vec![ContentPart::PromptResult { + content: PromptResult { + prompt_request_id: prompt_request_id.to_owned(), + prompt_name: prompt_name.to_owned(), + messages: Vec::new(), + content: None, + is_error: false, + error_message: None, + }, + }]; + content.extend(values.iter().map(|value| ContentPart::Text { text: value.clone() })); + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +/// Builds the completion post-hook payload for a resource or template reference. +pub(crate) fn resource_completion_result_payload( + resource_uri: &str, + values: &[String], + resource_request_id: &str, +) -> MessagePayload { + let mut content = vec![ContentPart::Resource { + content: Resource { + resource_request_id: resource_request_id.to_owned(), + uri: resource_uri.to_owned(), + resource_type: ResourceType::Uri, + ..Resource::default() + }, + }]; + content.extend(values.iter().map(|value| ContentPart::Text { text: value.clone() })); + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +/// Applies a plugin-modified completion payload back onto the backend result. +/// +/// Only the completion values are written back, and the count must match what the plugin was +/// given, so the backend's `total` and `has_more` stay accurate. +pub(crate) fn completion_result_response( + mut original: CompleteResult, + payload: &MessagePayload, +) -> Result { + let values: Vec = payload + .message + .content + .iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + if values.len() != original.completion.values.len() { + return Err(invalid_payload("Plugin modified completion payload value count")); + } + original.completion.values = values; + Ok(original) +} + +/// Builds the `cmf.resource_post_fetch` payload for `resources/list`, over the merged and +/// namespaced listing. Read-only for the same reason as the template listing: MCP `Resource` +/// listing entries are metadata with no CMF equivalent to be rebuilt from. +pub(crate) fn resource_list_result_payload(resources: &[McpResource], resource_request_id: &str) -> MessagePayload { + let mut content = Vec::with_capacity(resources.len()); + for resource in resources { + content.push(ContentPart::ResourceRef { + content: ResourceReference { + resource_request_id: resource_request_id.to_owned(), + uri: resource.uri.clone(), + name: Some(resource.name.clone()), + resource_type: ResourceType::Uri, + range_start: None, + range_end: None, + selector: None, + }, + }); + if let Some(description) = &resource.description { + content.push(ContentPart::Text { text: description.clone() }); + } + } + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +/// Builds the `cmf.resource_post_fetch` payload for `resources/read`, one CMF resource per returned +/// content item. +/// +/// Text contents are exposed in full and can be rewritten. Blob contents are exposed as a reference +/// — URI and MIME type — but their bytes are withheld: MCP carries blobs base64-encoded while CMF +/// wants raw bytes, and cloning arbitrary binary payloads through the plugin pipeline on the hot +/// path is a memory hazard. Blob contents are never mutated. +pub(crate) fn read_resource_result_payload(contents: &[ResourceContents], resource_request_id: &str) -> MessagePayload { + let content = contents + .iter() + .map(|item| ContentPart::Resource { content: resource_contents_to_cmf(item, resource_request_id) }) + .collect(); + MessagePayload { + message: Message { schema_version: "2.0".to_owned(), role: Role::Assistant, content, channel: None }, + } +} + +fn resource_contents_to_cmf(contents: &ResourceContents, resource_request_id: &str) -> Resource { + match contents { + ResourceContents::TextResourceContents { uri, mime_type, text, .. } => Resource { + resource_request_id: resource_request_id.to_owned(), + uri: uri.clone(), + resource_type: ResourceType::Uri, + content: Some(text.clone()), + mime_type: mime_type.clone(), + ..Resource::default() + }, + ResourceContents::BlobResourceContents { uri, mime_type, .. } => Resource { + resource_request_id: resource_request_id.to_owned(), + uri: uri.clone(), + resource_type: ResourceType::Blob, + mime_type: mime_type.clone(), + ..Resource::default() + }, + _ => Resource { resource_request_id: resource_request_id.to_owned(), ..Resource::default() }, + } +} + +/// Applies a plugin-modified `resources/read` payload back onto the backend response. +/// +/// Only text content is written back, and the content count must match what the plugin was given. +/// `uri`, `mimeType`, `_meta`, and blob bytes always come from the backend response, so a plugin +/// that only observes gets a byte-identical passthrough. +pub(crate) fn read_resource_response( + mut original: ReadResourceResult, + payload: &MessagePayload, +) -> Result { + let resources = payload.message.get_resources(); + if resources.len() != original.contents.len() { + return Err(invalid_payload("Plugin modified resource payload content count")); + } + + for (resource, original_contents) in resources.iter().zip(original.contents.iter_mut()) { + match original_contents { + ResourceContents::TextResourceContents { text: original_text, .. } => match &resource.content { + Some(text) => original_text.clone_from(text), + // Text contents are always exposed, so `None` means the plugin cleared them. + // Restoring the backend text would fail open on redaction. + None => return Err(invalid_payload("Plugin removed text from a resource content")), + }, + // Blob bytes are never exposed, so having no text back is expected. + _ if resource.content.is_none() => {}, + _ => return Err(invalid_payload("Plugin added text to a non-text resource content")), + } + } + + Ok(original) +} + +pub(crate) fn resource_uri(payload: &MessagePayload) -> Option { + payload.message.get_resources().first().map(|resource| resource.uri.clone()) +} + +pub(crate) fn prompt_request_arguments(payload: &MessagePayload) -> Option> { + payload + .message + .get_prompt_requests() + .first() + .map(|request| request.arguments.clone().into_iter().collect::>()) +} + +/// Applies a plugin-modified prompt payload back onto the backend response. +/// +/// Only text content is written back. The message count must match the payload the plugin was +/// given; anything else is a structural change the gateway cannot apply safely. `description`, +/// `_meta`, message roles, and non-text blocks are always taken from the original backend +/// response, so a plugin that only observes gets a byte-identical passthrough. +pub(crate) fn prompt_result_response( + mut original: GetPromptResult, + payload: &MessagePayload, +) -> Result { + let prompt_results = payload.message.get_prompt_results(); + let Some(result) = prompt_results.first() else { + return Err(invalid_payload("Plugin modified prompt payload without a prompt result")); + }; + if result.messages.len() != original.messages.len() { + return Err(invalid_payload("Plugin modified prompt payload message count")); + } + + for (message, original_message) in result.messages.iter().zip(original.messages.iter_mut()) { + let text = message.content.iter().find_map(|part| match part { + ContentPart::Text { text } => Some(text), + _ => None, + }); + match &mut original_message.content { + ContentBlock::Text(original_text) => match text { + Some(text) => original_text.text.clone_from(text), + // A text block is always exposed as a text part, so its absence means the plugin + // removed it. Falling back to the backend text here would return content a + // redaction plugin explicitly stripped, so this fails closed instead. + None => return Err(invalid_payload("Plugin removed text from a prompt message")), + }, + // Non-text blocks are never exposed, so having no text back is expected. + _ if text.is_none() => {}, + _ => return Err(invalid_payload("Plugin added text to a non-text prompt message")), + } + } + + Ok(original) +} + +fn invalid_payload(message: &'static str) -> ErrorData { + ErrorData { code: ErrorCode::INVALID_PARAMS, message: message.into(), data: None } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/contextforge-gateway-rs-cpex/src/handle.rs b/crates/contextforge-gateway-rs-cpex/src/handle.rs index b63d631..96d8a72 100644 --- a/crates/contextforge-gateway-rs-cpex/src/handle.rs +++ b/crates/contextforge-gateway-rs-cpex/src/handle.rs @@ -13,7 +13,10 @@ use cpex::cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult, ErrorCode}, + model::{ + CallToolRequestParams, CallToolResult, CompleteResult, ErrorCode, GetPromptRequestParams, GetPromptResult, + Prompt, ReadResourceResult, Resource as McpResource, ResourceTemplate, + }, serde::{Serialize, de::DeserializeOwned}, }; use tokio::task::JoinHandle; @@ -21,7 +24,10 @@ use tokio::task::JoinHandle; use crate::{ config::{LoadedRuntimePluginConfig, RedisRuntimePluginConfigStore, RuntimePluginConfigStore, cpex_config}, error::GatewayPluginRuntimeError, - hooks::{RuntimeHookError, RuntimeHookState, ToolPreCallResult}, + hooks::{ + CompletionRequest, CompletionTarget, PromptPreFetchResult, ResourcePreFetchResult, ResourceUriUpdate, + RuntimeHookError, RuntimeHookState, ToolPreCallResult, + }, runtime::GatewayPluginRuntime, }; @@ -40,7 +46,9 @@ pub struct GatewayPluginRuntimeHandle { runtime: Arc>, } -struct RegistryToolCallState { +/// Pins the runtime that ran a pre hook to its matching post hook, so a config reload mid-call +/// cannot hand an in-flight request to a different plugin set. +struct RegistryCallState { runtime: Arc, state: Option, } @@ -251,9 +259,9 @@ impl GatewayPluginRuntimeHandle { return Err(runtime_failed_error(state.as_ref())); }; let mut result = runtime.before_tool_call(request, tool_name, backend_name).await?; - if runtime.has_post_hook() { + if runtime.has_tool_post_hook() { let state = result.state.take(); - result.state = Some(Arc::new(RegistryToolCallState { runtime: Arc::clone(runtime), state })); + result.state = Some(Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state })); } else { result.state = None; } @@ -266,12 +274,188 @@ impl GatewayPluginRuntimeHandle { response: CallToolResult, state: Option, ) -> Result { - match state.and_then(|state| state.downcast::().ok()) { + match state.and_then(|state| state.downcast::().ok()) { Some(state) => state.runtime.after_tool_call(tool_name, response, state.state.clone()).await, None => Ok(response), } } + pub async fn before_complete( + &self, + request: &CompletionRequest<'_>, + ) -> Result, ErrorData> { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + let pre_state = runtime.before_complete(request).await?; + Ok(runtime.has_completion_post_hook(request.target).then(|| { + Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state: pre_state }) as RuntimeHookState + })) + } + + pub async fn after_complete( + &self, + target: CompletionTarget, + identifier: &str, + response: CompleteResult, + state: Option, + ) -> Result { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_complete(target, identifier, response, state.state.clone()).await, + None => Ok(response), + } + } + + pub async fn before_subscribe(&self, uri: &str, backend_name: &str) -> Result { + self.before_resource_subscription(uri, backend_name, "resource subscribe").await + } + + pub async fn before_unsubscribe(&self, uri: &str, backend_name: &str) -> Result { + self.before_resource_subscription(uri, backend_name, "resource unsubscribe").await + } + + async fn before_resource_subscription( + &self, + uri: &str, + backend_name: &str, + denied: &str, + ) -> Result { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + runtime.before_resource_subscription(uri, backend_name, denied).await + } + + pub async fn before_read_resource( + &self, + uri: &str, + backend_name: &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(uri, backend_name).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_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 before_list_resources(&self) -> Result, ErrorData> { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + let pre_state = runtime.before_list_resources().await?; + Ok(runtime.has_resource_post_hook().then(|| { + Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state: pre_state }) as RuntimeHookState + })) + } + + pub async fn after_list_resources( + &self, + resources: &[McpResource], + state: Option, + ) -> Result<(), ErrorData> { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_list_resources(resources, state.state.clone()).await, + None => Ok(()), + } + } + + pub async fn before_list_resource_templates(&self) -> Result, ErrorData> { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + let pre_state = runtime.before_list_resource_templates().await?; + Ok(runtime.has_resource_post_hook().then(|| { + Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state: pre_state }) as RuntimeHookState + })) + } + + pub async fn after_list_resource_templates( + &self, + templates: &[ResourceTemplate], + state: Option, + ) -> Result<(), ErrorData> { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_list_resource_templates(templates, state.state.clone()).await, + None => Ok(()), + } + } + + pub async fn before_list_prompts(&self, cursor: Option<&str>) -> Result, ErrorData> { + let state = self.current(); + let RuntimeState::Active(runtime) = state.as_ref() else { + return Err(runtime_failed_error(state.as_ref())); + }; + let pre_state = runtime.before_list_prompts(cursor).await?; + Ok(runtime.has_prompt_post_hook().then(|| { + Arc::new(RegistryCallState { runtime: Arc::clone(runtime), state: pre_state }) as RuntimeHookState + })) + } + + pub async fn after_list_prompts( + &self, + prompts: &[Prompt], + state: Option, + ) -> Result<(), ErrorData> { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_list_prompts(prompts, state.state.clone()).await, + None => Ok(()), + } + } + + pub async fn before_get_prompt( + &self, + request: &GetPromptRequestParams, + prompt_name: &str, + backend_name: &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_get_prompt(request, prompt_name, backend_name).await?; + if runtime.has_prompt_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, + response: GetPromptResult, + state: Option, + ) -> Result { + match state.and_then(|state| state.downcast::().ok()) { + Some(state) => state.runtime.after_get_prompt(prompt_name, response, state.state.clone()).await, + None => Ok(response), + } + } + /// Runs the tool post hooks over a streamed tool event (progress or logging /// notification). Returns `None` when a plugin denies the event. pub async fn after_stream_event( @@ -283,7 +467,7 @@ impl GatewayPluginRuntimeHandle { where T: Serialize + DeserializeOwned, { - match state.and_then(|state| state.downcast::().ok()) { + match state.and_then(|state| state.downcast::().ok()) { Some(state) => state.runtime.after_tool_event(tool_name, event, state.state.clone()).await, None => Ok(Some(event)), } @@ -292,7 +476,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 } } @@ -320,6 +504,7 @@ mod tests { }; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, NumberOrString, ProgressNotificationParam, ProgressToken, + PromptMessage, Role as McpRole, }; use serde_json::{Value, json}; use tokio::sync::Mutex as TokioMutex; @@ -329,11 +514,13 @@ mod tests { }; use crate::config::LoadedRuntimePluginConfig; - use crate::{CmfPluginFactory, ToolArgumentsUpdate}; + use crate::{CmfPluginFactory, PromptArgumentsUpdate, ToolArgumentsUpdate}; use super::*; const TEST_MISSING_CONTEXT_ERROR_CODE: i64 = -32003; + const TEST_PROMPT_PRE_DENY_ERROR_CODE: i32 = -32011; + const TEST_PROMPT_POST_DENY_ERROR_CODE: i32 = -32012; const TEST_REWRITTEN_SUM_A: i64 = 10; const TEST_REWRITTEN_SUM_B: i64 = 20; const TEST_SHUTDOWN_RETRY_COUNT: usize = 20; @@ -611,6 +798,331 @@ mod tests { } } + /// Prompt hooks share the CMF pipeline with tool hooks but carry `PromptRequest` / + /// `PromptResult` content parts, so this fixture discriminates on content part rather than + /// on message role. + #[derive(Clone, Copy, Default)] + enum PromptBehavior { + #[default] + Allow, + RewriteArguments, + DenyPre, + RewriteText, + DenyPost, + DropMessage, + } + + #[derive(Default)] + struct PromptObservations { + pre_calls: usize, + post_calls: usize, + pre_name: Option, + pre_server_id: Option, + pre_request_id: Option, + post_request_id: Option, + post_texts: Vec, + } + + struct PromptTestPlugin { + config: PluginConfig, + observations: Arc>, + behavior: PromptBehavior, + } + + impl PromptTestPlugin { + fn new(behavior: PromptBehavior) -> Self { + Self { + config: PluginConfig { + name: "prompt".to_owned(), + kind: "prompt-test".to_owned(), + hooks: vec![ + cmf_hook_names::PROMPT_PRE_FETCH.to_owned(), + cmf_hook_names::PROMPT_POST_FETCH.to_owned(), + ], + ..Default::default() + }, + observations: Arc::new(Mutex::new(PromptObservations::default())), + behavior, + } + } + + fn observations(&self) -> Arc> { + Arc::clone(&self.observations) + } + } + + #[async_trait] + impl Plugin for PromptTestPlugin { + fn config(&self) -> &PluginConfig { + &self.config + } + } + + impl HookHandler for PromptTestPlugin { + async fn handle( + &self, + payload: &MessagePayload, + _extensions: &Extensions, + _ctx: &mut PluginContext, + ) -> PluginResult { + let is_post = payload.message.role == Role::Assistant; + let mut observations = self.observations.lock().expect("observations lock poisoned"); + if is_post { + observations.post_calls += 1; + if let Some(result) = payload.message.get_prompt_results().first() { + observations.post_request_id = Some(result.prompt_request_id.clone()); + observations.post_texts = result + .messages + .iter() + .flat_map(|message| message.content.iter()) + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + } + } else { + observations.pre_calls += 1; + if let Some(request) = payload.message.get_prompt_requests().first() { + observations.pre_name = Some(request.name.clone()); + observations.pre_server_id.clone_from(&request.server_id); + observations.pre_request_id = Some(request.prompt_request_id.clone()); + } + } + drop(observations); + + match (is_post, self.behavior) { + (false, PromptBehavior::RewriteArguments) => { + let mut modified = payload.clone(); + if let Some(ContentPart::PromptRequest { content }) = modified + .message + .content + .iter_mut() + .find(|part| matches!(part, ContentPart::PromptRequest { .. })) + { + content.arguments = HashMap::from([("topic".to_owned(), json!("rewritten"))]); + } + PluginResult::modify_payload(modified) + }, + (false, PromptBehavior::DenyPre) => PluginResult::deny( + PluginViolation::new("pre_denied", "prompt pre denied") + .with_proto_error_code(i64::from(TEST_PROMPT_PRE_DENY_ERROR_CODE)), + ), + (true, PromptBehavior::RewriteText) => { + let mut modified = payload.clone(); + if let Some(ContentPart::PromptResult { content }) = modified + .message + .content + .iter_mut() + .find(|part| matches!(part, ContentPart::PromptResult { .. })) + { + for message in &mut content.messages { + for part in &mut message.content { + if let ContentPart::Text { text } = part { + *text = format!("redacted:{text}"); + } + } + } + } + PluginResult::modify_payload(modified) + }, + (true, PromptBehavior::DropMessage) => { + let mut modified = payload.clone(); + if let Some(ContentPart::PromptResult { content }) = modified + .message + .content + .iter_mut() + .find(|part| matches!(part, ContentPart::PromptResult { .. })) + { + content.messages.clear(); + } + PluginResult::modify_payload(modified) + }, + (true, PromptBehavior::DenyPost) => PluginResult::deny( + PluginViolation::new("post_denied", "prompt post denied") + .with_proto_error_code(i64::from(TEST_PROMPT_POST_DENY_ERROR_CODE)), + ), + _ => PluginResult::allow(), + } + } + } + + struct PromptTestPluginFactory { + observations: Arc>, + behavior: PromptBehavior, + } + + impl PluginFactory for PromptTestPluginFactory { + fn create(&self, config: &PluginConfig) -> Result> { + let plugin = Arc::new(PromptTestPlugin { + config: config.clone(), + observations: Arc::clone(&self.observations), + behavior: self.behavior, + }); + let handlers = config + .hooks + .iter() + .filter_map(|hook| { + let hook = match hook.as_str() { + cmf_hook_names::PROMPT_PRE_FETCH => cmf_hook_names::PROMPT_PRE_FETCH, + cmf_hook_names::PROMPT_POST_FETCH => cmf_hook_names::PROMPT_POST_FETCH, + _ => return None, + }; + Some(( + hook, + Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) + as Arc, + )) + }) + .collect(); + let plugin: Arc = plugin; + Ok(PluginInstance { plugin, handlers }) + } + } + + async fn prompt_runtime(behavior: PromptBehavior) -> (CpexRuntimeRegistry, Arc>) { + let plugin = PromptTestPlugin::new(behavior); + let observations = plugin.observations(); + let config = config_document(json!({ + "plugins": [{ + "name": plugin.config.name.clone(), + "kind": plugin.config.kind.clone(), + "hooks": plugin.config.hooks.clone(), + }] + })); + let mut runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); + runtime + .register_factory( + "prompt-test", + Box::new(PromptTestPluginFactory { observations: Arc::clone(&observations), behavior }), + ) + .expect("prompt factory registers"); + runtime.initialize().await.expect("prompt runtime initializes"); + (runtime, observations) + } + + fn prompt_request(topic: &str) -> GetPromptRequestParams { + GetPromptRequestParams::new("review") + .with_arguments(serde_json::Map::from_iter([("topic".to_owned(), json!(topic))])) + } + + fn prompt_result(text: &str) -> GetPromptResult { + GetPromptResult::new(vec![PromptMessage::new_text(McpRole::User, text)]) + } + + fn prompt_text(result: &GetPromptResult) -> String { + result + .messages + .iter() + .filter_map(|message| match &message.content { + ContentBlock::Text(text) => Some(text.text.clone()), + _ => None, + }) + .collect() + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn prompt_pre_hook_sees_routed_name_and_backend() { + let (runtime, observations) = prompt_runtime(PromptBehavior::Allow).await; + + let result = runtime + .handle() + .before_get_prompt(&prompt_request("weather"), "review", "backend-a") + .await + .expect("prompt pre hook runs"); + + assert!(matches!(result.arguments, PromptArgumentsUpdate::Unchanged)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(Some("review"), observations.pre_name.as_deref()); + assert_eq!(Some("backend-a"), observations.pre_server_id.as_deref()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn prompt_pre_hook_replaces_arguments() { + let (runtime, _) = prompt_runtime(PromptBehavior::RewriteArguments).await; + + let result = runtime + .handle() + .before_get_prompt(&prompt_request("weather"), "review", "backend-a") + .await + .expect("prompt pre hook runs"); + + let PromptArgumentsUpdate::Replace(Some(arguments)) = result.arguments else { + panic!("prompt arguments are replaced"); + }; + assert_eq!(Some(&json!("rewritten")), arguments.get("topic")); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn denied_prompt_pre_hook_maps_violation_code() { + let (runtime, _) = prompt_runtime(PromptBehavior::DenyPre).await; + + let error = expect_prompt_denied( + runtime.handle().before_get_prompt(&prompt_request("weather"), "review", "backend-a").await, + ); + + assert_eq!(ErrorCode(TEST_PROMPT_PRE_DENY_ERROR_CODE), error.code); + assert_eq!("Plugin denied prompt fetch", error.message); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn prompt_post_hook_rewrites_text_and_correlates_with_pre() { + let (runtime, observations) = prompt_runtime(PromptBehavior::RewriteText).await; + + let pre = runtime + .handle() + .before_get_prompt(&prompt_request("weather"), "review", "backend-a") + .await + .expect("prompt pre hook runs"); + let response = runtime + .handle() + .after_get_prompt("review", prompt_result("secret"), pre.state) + .await + .expect("prompt post hook runs"); + + assert_eq!("redacted:secret", prompt_text(&response)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(observations.pre_request_id, observations.post_request_id); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn denied_prompt_post_hook_maps_violation_code() { + let (runtime, _) = prompt_runtime(PromptBehavior::DenyPost).await; + + let pre = runtime + .handle() + .before_get_prompt(&prompt_request("weather"), "review", "backend-a") + .await + .expect("prompt pre hook runs"); + let error = runtime + .handle() + .after_get_prompt("review", prompt_result("secret"), pre.state) + .await + .expect_err("prompt post deny is an error"); + + assert_eq!(ErrorCode(TEST_PROMPT_POST_DENY_ERROR_CODE), error.code); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn prompt_post_hook_rejects_structural_message_change() { + let (runtime, _) = prompt_runtime(PromptBehavior::DropMessage).await; + + let pre = runtime + .handle() + .before_get_prompt(&prompt_request("weather"), "review", "backend-a") + .await + .expect("prompt pre hook runs"); + let error = runtime + .handle() + .after_get_prompt("review", prompt_result("secret"), pre.state) + .await + .expect_err("dropping messages is rejected"); + + assert_eq!(ErrorCode::INVALID_PARAMS, error.code); + } + fn sum_request(a: i64, b: i64) -> CallToolRequestParams { CallToolRequestParams::new("sum") .with_arguments(serde_json::Map::from_iter([("a".to_owned(), json!(a)), ("b".to_owned(), json!(b))])) @@ -645,6 +1157,15 @@ mod tests { })) } + /// `PromptPreFetchResult` carries an opaque `Arc` state, so it cannot derive `Debug` + /// and `expect_err` is unavailable. + fn expect_prompt_denied(result: Result) -> ErrorData { + match result { + Ok(_) => panic!("prompt pre hook should deny"), + Err(error) => error, + } + } + fn expect_runtime_failed(result: Result) -> ErrorData { match result { Ok(_) => panic!("runtime should be failed"), @@ -714,6 +1235,25 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn non_tool_mcp_hooks_are_accepted_config() { + for hook in [ + cmf_hook_names::PROMPT_PRE_FETCH, + cmf_hook_names::PROMPT_POST_FETCH, + cmf_hook_names::RESOURCE_PRE_FETCH, + cmf_hook_names::RESOURCE_POST_FETCH, + ] { + let plugin = Arc::new(TestPlugin::new("non-tool", vec![hook])); + let config = plugin_config(&[Arc::clone(&plugin)]); + let mut runtime = CpexRuntimeRegistry::with_config_store(Arc::new(MemoryConfigStore::with_config(config))); + runtime + .register_factory("test", Box::new(TestPluginFactory::from_plugin(&plugin))) + .expect("test factory registers"); + + runtime.initialize().await.expect("non-tool MCP hook config initializes"); + } + } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn runtime_config_loads_registered_factory_plugin() { let plugin = diff --git a/crates/contextforge-gateway-rs-cpex/src/hooks.rs b/crates/contextforge-gateway-rs-cpex/src/hooks.rs index 0e6ce65..bd9b119 100644 --- a/crates/contextforge-gateway-rs-cpex/src/hooks.rs +++ b/crates/contextforge-gateway-rs-cpex/src/hooks.rs @@ -1,6 +1,8 @@ use std::{any::Any, sync::Arc}; -use rmcp::model::CallToolRequestParams; +use std::collections::HashMap; + +use rmcp::model::{ArgumentInfo, CallToolRequestParams, GetPromptRequestParams}; use serde_json::{Map, Value}; pub type RuntimeHookError = Box; @@ -31,3 +33,88 @@ impl ToolPreCallResult { Self { arguments: ToolArgumentsUpdate::Unchanged, state: None } } } + +#[derive(Debug)] +pub enum PromptArgumentsUpdate { + Unchanged, + Replace(Option>), +} + +impl PromptArgumentsUpdate { + pub fn apply_to_request(self, request: &mut GetPromptRequestParams, routed_prompt_name: &str) { + routed_prompt_name.clone_into(&mut request.name); + if let Self::Replace(arguments) = self { + request.arguments = arguments; + } + } +} + +pub struct PromptPreFetchResult { + pub arguments: PromptArgumentsUpdate, + pub state: Option, +} + +impl PromptPreFetchResult { + pub fn unchanged() -> Self { + Self { arguments: PromptArgumentsUpdate::Unchanged, state: None } + } +} + +/// Which CPEX hook family a `completion/complete` request dispatches to. +/// +/// CPEX defines no completion hook, so a completion runs the hooks of the thing being completed: +/// a prompt reference uses the prompt hooks, a resource or template reference uses the resource +/// hooks. Plugins tell a completion apart from a plain fetch by the `gateway-completion-*` +/// correlation id. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompletionTarget { + Prompt, + Resource, +} + +/// A routed `completion/complete` request as the hooks see it. +pub struct CompletionRequest<'a> { + /// Which hook family this reference dispatches to. + pub target: CompletionTarget, + /// Backend-local prompt name or resource URI, with the namespace prefix already stripped. + pub identifier: &'a str, + /// Backend that owns the reference. + pub backend_name: &'a str, + /// The argument being completed. + pub argument: &'a ArgumentInfo, + /// Argument values the client already resolved, when it supplied any. + pub context: Option<&'a HashMap>, +} + +/// Outcome of the resource pre hook on `resources/subscribe` and `resources/unsubscribe`. These +/// paths return an empty acknowledgement, so there is no post hook and no state to carry. +#[derive(Debug)] +pub enum ResourceUriUpdate { + Unchanged, + Replace(String), +} + +/// Outcome of the resource pre hook on `resources/read`, which unlike the subscription paths does +/// have a post hook, so it carries state forward. +pub struct ResourcePreFetchResult { + pub uri: ResourceUriUpdate, + pub state: Option, +} + +impl ResourcePreFetchResult { + pub fn unchanged() -> Self { + Self { uri: ResourceUriUpdate::Unchanged, state: None } + } +} + +impl ResourceUriUpdate { + /// Resolves the URI the gateway should forward to the backend and track the subscription + /// under. Both must be the same value, or a rewritten subscription would be tracked under a + /// URI the backend never notifies about. + pub fn apply_to(self, routed_uri: String) -> String { + match self { + Self::Unchanged => routed_uri, + Self::Replace(uri) => uri, + } + } +} diff --git a/crates/contextforge-gateway-rs-cpex/src/lib.rs b/crates/contextforge-gateway-rs-cpex/src/lib.rs index 6a57cb5..92f3487 100644 --- a/crates/contextforge-gateway-rs-cpex/src/lib.rs +++ b/crates/contextforge-gateway-rs-cpex/src/lib.rs @@ -10,4 +10,7 @@ mod runtime; pub use error::GatewayPluginRuntimeError; pub use factory::CmfPluginFactory; pub use handle::{CpexRuntimeRegistry, GatewayPluginRuntimeHandle}; -pub use hooks::{RuntimeHookError, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}; +pub use hooks::{ + CompletionRequest, CompletionTarget, PromptArgumentsUpdate, PromptPreFetchResult, ResourcePreFetchResult, + ResourceUriUpdate, RuntimeHookError, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, +}; diff --git a/crates/contextforge-gateway-rs-cpex/src/pipeline.rs b/crates/contextforge-gateway-rs-cpex/src/pipeline.rs index 5d6cdd9..6ee0c53 100644 --- a/crates/contextforge-gateway-rs-cpex/src/pipeline.rs +++ b/crates/contextforge-gateway-rs-cpex/src/pipeline.rs @@ -2,14 +2,17 @@ use cpex::cpex_core::cmf::MessagePayload; use cpex::cpex_core::executor::PipelineResult; use rmcp::{ ErrorData, - model::{CallToolResult, ErrorCode}, + model::{CallToolResult, CompleteResult, ErrorCode, GetPromptResult, ReadResourceResult}, serde::de::DeserializeOwned, }; use tracing::warn; use crate::{ - ToolArgumentsUpdate, - cmf::{tool_call_arguments, tool_result_content, tool_result_response}, + PromptArgumentsUpdate, ResourceUriUpdate, ToolArgumentsUpdate, + cmf::{ + completion_result_response, prompt_request_arguments, prompt_result_response, read_resource_response, + resource_uri, tool_call_arguments, tool_result_content, tool_result_response, + }, }; pub(crate) fn modified_message_payload(result: &PipelineResult) -> Option<&MessagePayload> { @@ -46,6 +49,78 @@ pub(crate) fn effective_post_result(original: CallToolResult, result: &PipelineR } } +pub(crate) fn effective_pre_prompt_args( + original_args: Option<&serde_json::Map>, + pre_result: &PipelineResult, +) -> Result { + let Some(modified_payload) = modified_message_payload(pre_result) else { + return Ok(PromptArgumentsUpdate::Unchanged); + }; + + let Some(arguments) = prompt_request_arguments(modified_payload) else { + return Err(ErrorData { + code: ErrorCode::INVALID_PARAMS, + message: "Plugin modified prompt payload without a prompt request".into(), + data: None, + }); + }; + + if original_args == Some(&arguments) || (original_args.is_none() && arguments.is_empty()) { + Ok(PromptArgumentsUpdate::Unchanged) + } else { + Ok(PromptArgumentsUpdate::Replace(Some(arguments))) + } +} + +pub(crate) fn effective_post_prompt_result( + original: GetPromptResult, + result: &PipelineResult, +) -> Result { + match modified_message_payload(result) { + Some(payload) => prompt_result_response(original, payload), + None => Ok(original), + } +} + +pub(crate) fn effective_post_completion_result( + original: CompleteResult, + result: &PipelineResult, +) -> Result { + match modified_message_payload(result) { + Some(payload) => completion_result_response(original, payload), + None => Ok(original), + } +} + +pub(crate) fn effective_post_read_resource( + original: ReadResourceResult, + result: &PipelineResult, +) -> Result { + match modified_message_payload(result) { + Some(payload) => read_resource_response(original, payload), + None => Ok(original), + } +} + +pub(crate) fn effective_pre_resource_uri( + original_uri: &str, + pre_result: &PipelineResult, +) -> Result { + let Some(modified_payload) = modified_message_payload(pre_result) else { + return Ok(ResourceUriUpdate::Unchanged); + }; + + let Some(uri) = resource_uri(modified_payload) else { + return Err(ErrorData { + code: ErrorCode::INVALID_PARAMS, + message: "Plugin modified resource payload without a resource".into(), + data: None, + }); + }; + + if uri == original_uri { Ok(ResourceUriUpdate::Unchanged) } else { Ok(ResourceUriUpdate::Replace(uri)) } +} + pub(crate) fn effective_post_json(original: T, result: &PipelineResult) -> Result where T: DeserializeOwned, @@ -67,16 +142,18 @@ where }) } -pub(crate) fn plugin_denied_error(result: PipelineResult) -> ErrorData { +/// Maps a plugin denial onto an MCP error. `denied` names the denied operation for the client-facing +/// message, e.g. `"tool call"` or `"prompt fetch"`. +pub(crate) fn plugin_denied_error(result: PipelineResult, denied: &str) -> ErrorData { let code = result .violation .and_then(|violation| { - warn!("Plugin denied tool call: code={} plugin={:?}", violation.code, violation.plugin_name); + warn!("Plugin denied {denied}: code={} plugin={:?}", violation.code, violation.plugin_name); violation.proto_error_code.and_then(|code| i32::try_from(code).ok()).map(ErrorCode) }) .unwrap_or(ErrorCode::INVALID_REQUEST); - ErrorData { code, message: "Plugin denied tool call".into(), data: None } + ErrorData { code, message: format!("Plugin denied {denied}").into(), data: None } } pub(crate) fn log_pipeline_errors(hook: &'static str, result: &PipelineResult) { diff --git a/crates/contextforge-gateway-rs-cpex/src/runtime.rs b/crates/contextforge-gateway-rs-cpex/src/runtime.rs index 21dc39c..9ce23ac 100644 --- a/crates/contextforge-gateway-rs-cpex/src/runtime.rs +++ b/crates/contextforge-gateway-rs-cpex/src/runtime.rs @@ -14,50 +14,105 @@ use cpex::cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult}, + model::{ + CallToolRequestParams, CallToolResult, CompleteResult, GetPromptRequestParams, GetPromptResult, Prompt, + ReadResourceResult, Resource as McpResource, ResourceTemplate, + }, serde::{Serialize, de::DeserializeOwned}, }; use tokio::sync::Mutex; use crate::{ - cmf::{tool_call_payload, tool_json_result_payload, tool_result_payload}, + cmf::{ + prompt_completion_payload, prompt_completion_result_payload, prompt_list_request_payload, + prompt_list_result_payload, prompt_request_payload, prompt_result_payload, read_resource_result_payload, + resource_completion_payload, resource_completion_result_payload, resource_list_request_payload, + resource_list_result_payload, resource_request_payload, resource_template_list_result_payload, + tool_call_payload, tool_json_result_payload, tool_result_payload, + }, error::GatewayPluginRuntimeError, - hooks::{RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, + hooks::{ + CompletionRequest, CompletionTarget, PromptArgumentsUpdate, PromptPreFetchResult, ResourcePreFetchResult, + ResourceUriUpdate, RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult, + }, pipeline::{ - effective_post_json, effective_post_result, effective_pre_args, log_pipeline_errors, plugin_denied_error, + effective_post_completion_result, effective_post_json, effective_post_prompt_result, + effective_post_read_resource, effective_post_result, effective_pre_args, effective_pre_prompt_args, + effective_pre_resource_uri, log_pipeline_errors, plugin_denied_error, }, }; +/// Whether the loaded plugin config declares the pre and post hook of one MCP path. +#[derive(Default, Clone, Copy)] +struct HookPair { + pre: bool, + post: bool, +} + +impl HookPair { + fn from_config(config: &CpexConfig, pre: &str, post: &str) -> Self { + let declares = |name: &str| config.plugins.iter().any(|plugin| plugin.hooks.iter().any(|hook| hook == name)); + Self { pre: declares(pre), post: declares(post) } + } +} + +/// Which hooks the loaded plugin config actually declares. Each MCP path checks only its own +/// pair, so a prompt-only plugin config leaves the tool hot path untouched. +#[derive(Default, Clone, Copy)] +struct HookPresence { + tool: HookPair, + prompt: HookPair, + resource: HookPair, +} + #[derive(Default)] pub(crate) struct GatewayPluginRuntime { manager: PluginManager, - has_pre_hook: bool, - has_post_hook: bool, + hooks: HookPresence, } -struct ToolCallState { +/// Per-call CPEX state carried from a pre hook to its matching post hook. `correlation_id` is the +/// CMF `tool_call_id` / `prompt_request_id` that links the two payloads. +struct CallState { context_table: PluginContextTable, - tool_call_id: String, + correlation_id: String, } -type SharedToolCallState = Mutex; +type SharedCallState = Mutex; -static TOOL_CALL_ID: AtomicU64 = AtomicU64::new(1); +static CORRELATION_ID: AtomicU64 = AtomicU64::new(1); -fn next_tool_call_id() -> String { - format!("gateway-tool-call-{}", TOOL_CALL_ID.fetch_add(1, Ordering::Relaxed)) +fn next_correlation_id(kind: &str) -> String { + format!("gateway-{kind}-{}", CORRELATION_ID.fetch_add(1, Ordering::Relaxed)) } -fn new_tool_call_state() -> RuntimeHookState { - Arc::new(Mutex::new(ToolCallState { +fn new_call_state(kind: &str) -> RuntimeHookState { + Arc::new(Mutex::new(CallState { context_table: PluginContextTable::default(), - tool_call_id: next_tool_call_id(), + correlation_id: next_correlation_id(kind), })) } +const TOOL_CALL_KIND: &str = "tool-call"; +const PROMPT_REQUEST_KIND: &str = "prompt-request"; +const PROMPT_LIST_KIND: &str = "prompt-list"; +const RESOURCE_TEMPLATE_LIST_KIND: &str = "resource-template-list"; +const RESOURCE_SUBSCRIPTION_KIND: &str = "resource-subscription"; +const COMPLETION_KIND: &str = "completion"; +const RESOURCE_READ_KIND: &str = "resource-read"; +const RESOURCE_LIST_KIND: &str = "resource-list"; + impl GatewayPluginRuntime { - pub(crate) fn has_post_hook(&self) -> bool { - self.has_post_hook + pub(crate) fn has_tool_post_hook(&self) -> bool { + self.hooks.tool.post + } + + pub(crate) fn has_prompt_post_hook(&self) -> bool { + self.hooks.prompt.post + } + + pub(crate) fn has_resource_post_hook(&self) -> bool { + self.hooks.resource.post } pub(crate) async fn from_config( @@ -66,16 +121,19 @@ impl GatewayPluginRuntime { ) -> Result { validate_gateway_supported_config(&config)?; - let has_pre_hook = - config.plugins.iter().any(|plugin| plugin.hooks.iter().any(|hook| hook == cmf_hook_names::TOOL_PRE_INVOKE)); - let has_post_hook = config - .plugins - .iter() - .any(|plugin| plugin.hooks.iter().any(|hook| hook == cmf_hook_names::TOOL_POST_INVOKE)); + let hooks = HookPresence { + tool: HookPair::from_config(&config, cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE), + prompt: HookPair::from_config(&config, cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH), + resource: HookPair::from_config( + &config, + cmf_hook_names::RESOURCE_PRE_FETCH, + cmf_hook_names::RESOURCE_POST_FETCH, + ), + }; let manager = PluginManager::from_config(config, factories) .map_err(|source| GatewayPluginRuntimeError::Configuration { hook: "config", source })?; manager.initialize().await.map_err(|source| GatewayPluginRuntimeError::Initialization { source })?; - Ok(Self { manager, has_pre_hook, has_post_hook }) + Ok(Self { manager, hooks }) } } @@ -93,6 +151,18 @@ impl Drop for GatewayPluginRuntime { } } +/// CPEX hooks the gateway is able to run. Tool hooks wrap `call_tool`; prompt and resource hooks +/// wrap the non-tool MCP pass-through paths. Any other hook is rejected before the runtime loads, +/// because the gateway has no defined behavior for it on the hot path. +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 validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayPluginRuntimeError> { if config.routing_enabled() || config.plugin_settings.fail_on_plugin_error @@ -109,11 +179,7 @@ fn validate_gateway_supported_config(config: &CpexConfig) -> Result<(), GatewayP return Err(GatewayPluginRuntimeError::ConfigUnsupported); } - if plugin - .hooks - .iter() - .any(|hook| hook != cmf_hook_names::TOOL_PRE_INVOKE && hook != cmf_hook_names::TOOL_POST_INVOKE) - { + if plugin.hooks.iter().any(|hook| !SUPPORTED_HOOKS.contains(&hook.as_str())) { return Err(GatewayPluginRuntimeError::ConfigUnsupported); } } @@ -146,26 +212,317 @@ impl GatewayPluginRuntime { result } + 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 + } + + 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 + } + + fn completion_hooks(&self, target: CompletionTarget) -> HookPair { + match target { + CompletionTarget::Prompt => self.hooks.prompt, + CompletionTarget::Resource => self.hooks.resource, + } + } + + pub(crate) fn has_completion_post_hook(&self, target: CompletionTarget) -> bool { + self.completion_hooks(target).post + } + + /// Runs the pre hook for `completion/complete`, after routing. CPEX has no completion hook, so + /// this dispatches to the prompt or resource family depending on what is being completed. + /// + /// Deny-only: the argument value is a partial user keystroke, so rewriting it would return + /// completions for text the user never typed. + pub(crate) async fn before_complete( + &self, + request: &CompletionRequest<'_>, + ) -> Result, ErrorData> { + let CompletionRequest { target, identifier, backend_name, argument, context } = *request; + let hooks = self.completion_hooks(target); + if !hooks.pre { + return Ok(hooks.post.then(|| new_call_state(COMPLETION_KIND))); + } + + let correlation_id = next_correlation_id(COMPLETION_KIND); + let pre_result = match target { + CompletionTarget::Prompt => { + let payload = prompt_completion_payload(identifier, backend_name, argument, context, &correlation_id); + self.invoke_prompt_pre(payload).await + }, + CompletionTarget::Resource => { + let payload = resource_completion_payload(identifier, backend_name, argument, context, &correlation_id); + self.invoke_resource_pre(payload).await + }, + }; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, "completion")); + } + + Ok(Some(Arc::new(Mutex::new(CallState { context_table: pre_result.context_table, correlation_id })))) + } + + /// Runs the post hook for `completion/complete`. A plugin may rewrite the completion values; + /// the count must match, so the backend's `total` and `has_more` stay accurate. + pub(crate) async fn after_complete( + &self, + target: CompletionTarget, + identifier: &str, + response: CompleteResult, + state: Option, + ) -> Result { + if !self.completion_hooks(target).post { + return Ok(response); + } + + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(response) }; + + let mut state = state.lock().await; + let values = &response.completion.values; + let post_result = match target { + CompletionTarget::Prompt => { + let payload = prompt_completion_result_payload(identifier, values, &state.correlation_id); + self.invoke_prompt_post(payload, Some(state.context_table.clone())).await + }, + CompletionTarget::Resource => { + let payload = resource_completion_result_payload(identifier, values, &state.correlation_id); + self.invoke_resource_post(payload, Some(state.context_table.clone())).await + }, + }; + if post_result.is_denied() { + return Err(plugin_denied_error(post_result, "completion")); + } + + state.context_table = post_result.context_table.clone(); + effective_post_completion_result(response, &post_result) + } + + /// Runs the resource pre hook for `resources/read`, after routing. A plugin can deny the read + /// or rewrite the target URI, and state is carried forward to the post hook. + pub(crate) async fn before_read_resource( + &self, + uri: &str, + backend_name: &str, + ) -> Result { + if !self.hooks.resource.pre { + let state = self.hooks.resource.post.then(|| new_call_state(RESOURCE_READ_KIND)); + return Ok(ResourcePreFetchResult { uri: ResourceUriUpdate::Unchanged, state }); + } + + let resource_request_id = next_correlation_id(RESOURCE_READ_KIND); + let pre_result = + self.invoke_resource_pre(resource_request_payload(uri, backend_name, &resource_request_id)).await; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, "resource read")); + } + + let uri = effective_pre_resource_uri(uri, &pre_result)?; + let state = + Mutex::new(CallState { context_table: pre_result.context_table, correlation_id: resource_request_id }); + Ok(ResourcePreFetchResult { uri, state: Some(Arc::new(state)) }) + } + + /// Runs the resource post hook over the contents `resources/read` returned. A plugin may + /// rewrite text contents; blob bytes are never exposed and never modified. + 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 mut state = state.lock().await; + let post_result = self + .invoke_resource_post( + read_resource_result_payload(&response.contents, &state.correlation_id), + Some(state.context_table.clone()), + ) + .await; + if post_result.is_denied() { + return Err(plugin_denied_error(post_result, "resource read")); + } + + state.context_table = post_result.context_table.clone(); + effective_post_read_resource(response, &post_result) + } + + /// Runs the resource pre hook for `resources/list`, before any backend is contacted. + /// Deny-only, matching the template listing. + pub(crate) async fn before_list_resources(&self) -> Result, ErrorData> { + self.before_resource_listing(RESOURCE_LIST_KIND, "resource list").await + } + + /// Runs the resource post hook over the merged, namespaced resource listing. Read-only: MCP + /// `Resource` listing entries are metadata that cannot be rebuilt from CMF content parts. + pub(crate) async fn after_list_resources( + &self, + resources: &[McpResource], + state: Option, + ) -> Result<(), ErrorData> { + self.after_resource_listing(state, "resource list", |correlation_id| { + resource_list_result_payload(resources, correlation_id) + }) + .await + } + + /// Runs the resource pre hook for `resources/subscribe` and `resources/unsubscribe`, after + /// routing. A plugin can deny the request or rewrite the target URI. There is no post hook: + /// both paths return an empty acknowledgement. `denied` names the operation in the error a + /// denial produces. + pub(crate) async fn before_resource_subscription( + &self, + uri: &str, + backend_name: &str, + denied: &str, + ) -> Result { + if !self.hooks.resource.pre { + return Ok(ResourceUriUpdate::Unchanged); + } + + let resource_request_id = next_correlation_id(RESOURCE_SUBSCRIPTION_KIND); + let pre_result = + self.invoke_resource_pre(resource_request_payload(uri, backend_name, &resource_request_id)).await; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, denied)); + } + + effective_pre_resource_uri(uri, &pre_result) + } + + /// Runs the resource pre hook for `resources/templates/list`, before any backend is contacted. + /// Deny-only: the payload carries no routable resource, so a modified payload has nothing to + /// apply to. + pub(crate) async fn before_list_resource_templates(&self) -> Result, ErrorData> { + self.before_resource_listing(RESOURCE_TEMPLATE_LIST_KIND, "resource template list").await + } + + /// Runs the resource post hook over the merged, namespaced template listing. Read-only: a + /// modified payload is ignored because MCP `ResourceTemplate` metadata cannot be rebuilt from + /// CMF content parts. Denial is still honored and surfaces as an MCP error. + pub(crate) async fn after_list_resource_templates( + &self, + templates: &[ResourceTemplate], + state: Option, + ) -> Result<(), ErrorData> { + self.after_resource_listing(state, "resource template list", |correlation_id| { + resource_template_list_result_payload(templates, correlation_id) + }) + .await + } + + /// Shared pre-hook body for the resource listings, which differ only in correlation kind and + /// the operation named in a denial. + async fn before_resource_listing(&self, kind: &str, denied: &str) -> Result, ErrorData> { + if !self.hooks.resource.pre { + return Ok(self.hooks.resource.post.then(|| new_call_state(kind))); + } + + let resource_request_id = next_correlation_id(kind); + let pre_result = self.invoke_resource_pre(resource_list_request_payload(&resource_request_id)).await; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, denied)); + } + + Ok(Some(Arc::new(Mutex::new(CallState { + context_table: pre_result.context_table, + correlation_id: resource_request_id, + })))) + } + + /// Shared post-hook body for the resource listings. Read-only: the modified payload is + /// discarded, but a denial still surfaces as an MCP error. + async fn after_resource_listing( + &self, + state: Option, + denied: &str, + payload: impl FnOnce(&str) -> MessagePayload, + ) -> Result<(), ErrorData> { + if !self.hooks.resource.post { + return Ok(()); + } + + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(()) }; + + let mut state = state.lock().await; + let post_result = + self.invoke_resource_post(payload(&state.correlation_id), Some(state.context_table.clone())).await; + if post_result.is_denied() { + return Err(plugin_denied_error(post_result, denied)); + } + + state.context_table = post_result.context_table.clone(); + Ok(()) + } + pub(crate) async fn before_tool_call( &self, request: &CallToolRequestParams, tool_name: &str, backend_name: &str, ) -> Result { - if !self.has_pre_hook { - let state = self.has_post_hook.then(new_tool_call_state); + if !self.hooks.tool.pre { + let state = self.hooks.tool.post.then(|| new_call_state(TOOL_CALL_KIND)); return Ok(ToolPreCallResult { arguments: ToolArgumentsUpdate::Unchanged, state }); } - let tool_call_id = next_tool_call_id(); + let tool_call_id = next_correlation_id(TOOL_CALL_KIND); let original_payload = tool_call_payload(request, tool_name, backend_name, &tool_call_id); let pre_result = self.invoke_tool_pre(original_payload).await; if pre_result.is_denied() { - return Err(plugin_denied_error(pre_result)); + return Err(plugin_denied_error(pre_result, "tool call")); } let arguments = effective_pre_args(request.arguments.as_ref(), &pre_result)?; - let state = Mutex::new(ToolCallState { context_table: pre_result.context_table, tool_call_id }); + let state = Mutex::new(CallState { context_table: pre_result.context_table, correlation_id: tool_call_id }); Ok(ToolPreCallResult { arguments, state: Some(Arc::new(state)) }) } @@ -175,28 +532,132 @@ impl GatewayPluginRuntime { response: CallToolResult, state: Option, ) -> Result { - if !self.has_post_hook { + if !self.hooks.tool.post { return Ok(response); } - let state = state.and_then(|state| state.downcast::().ok()); + let state = state.and_then(|state| state.downcast::().ok()); let Some(state) = state else { return Ok(response) }; let mut state = state.lock().await; let post_result = self .invoke_tool_post( - tool_result_payload(tool_name, &response, &state.tool_call_id), + tool_result_payload(tool_name, &response, &state.correlation_id), Some(state.context_table.clone()), ) .await; if post_result.is_denied() { - return Err(plugin_denied_error(post_result)); + return Err(plugin_denied_error(post_result, "tool call")); } state.context_table = post_result.context_table.clone(); Ok(effective_post_result(response, &post_result)) } + /// Runs the prompt pre hook for `prompts/list`, before any backend is contacted. Deny-only: + /// the payload carries no routable prompt, so a modified payload has nothing to apply to. + pub(crate) async fn before_list_prompts( + &self, + cursor: Option<&str>, + ) -> Result, ErrorData> { + if !self.hooks.prompt.pre { + return Ok(self.hooks.prompt.post.then(|| new_call_state(PROMPT_LIST_KIND))); + } + + let prompt_request_id = next_correlation_id(PROMPT_LIST_KIND); + let pre_result = self.invoke_prompt_pre(prompt_list_request_payload(cursor, &prompt_request_id)).await; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, "prompt list")); + } + + Ok(Some(Arc::new(Mutex::new(CallState { + context_table: pre_result.context_table, + correlation_id: prompt_request_id, + })))) + } + + /// Runs the prompt post hook over the merged, namespaced listing. Read-only: a modified payload + /// is ignored because MCP `Prompt` metadata cannot be rebuilt from CMF content parts. Denial is + /// still honored and surfaces as an MCP error. + pub(crate) async fn after_list_prompts( + &self, + prompts: &[Prompt], + state: Option, + ) -> Result<(), ErrorData> { + if !self.hooks.prompt.post { + return Ok(()); + } + + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(()) }; + + let mut state = state.lock().await; + let post_result = self + .invoke_prompt_post( + prompt_list_result_payload(prompts, &state.correlation_id), + Some(state.context_table.clone()), + ) + .await; + if post_result.is_denied() { + return Err(plugin_denied_error(post_result, "prompt list")); + } + + state.context_table = post_result.context_table.clone(); + Ok(()) + } + + pub(crate) async fn before_get_prompt( + &self, + request: &GetPromptRequestParams, + prompt_name: &str, + backend_name: &str, + ) -> Result { + if !self.hooks.prompt.pre { + let state = self.hooks.prompt.post.then(|| new_call_state(PROMPT_REQUEST_KIND)); + return Ok(PromptPreFetchResult { arguments: PromptArgumentsUpdate::Unchanged, state }); + } + + let prompt_request_id = next_correlation_id(PROMPT_REQUEST_KIND); + let original_payload = prompt_request_payload(request, prompt_name, backend_name, &prompt_request_id); + let pre_result = self.invoke_prompt_pre(original_payload).await; + if pre_result.is_denied() { + return Err(plugin_denied_error(pre_result, "prompt fetch")); + } + + let arguments = effective_pre_prompt_args(request.arguments.as_ref(), &pre_result)?; + let state = + Mutex::new(CallState { context_table: pre_result.context_table, correlation_id: prompt_request_id }); + Ok(PromptPreFetchResult { arguments, state: Some(Arc::new(state)) }) + } + + pub(crate) async fn after_get_prompt( + &self, + prompt_name: &str, + response: GetPromptResult, + state: Option, + ) -> Result { + if !self.hooks.prompt.post { + return Ok(response); + } + + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(response) }; + + let mut state = state.lock().await; + let post_result = self + .invoke_prompt_post( + prompt_result_payload(prompt_name, &response, &state.correlation_id), + Some(state.context_table.clone()), + ) + .await; + if post_result.is_denied() { + return Err(plugin_denied_error(post_result, "prompt fetch")); + } + + state.context_table = post_result.context_table.clone(); + effective_post_prompt_result(response, &post_result) + } + pub(crate) async fn after_tool_event( &self, tool_name: &str, @@ -206,18 +667,18 @@ impl GatewayPluginRuntime { where T: Serialize + DeserializeOwned, { - if !self.has_post_hook { + if !self.hooks.tool.post { return Ok(Some(event)); } - let state = state.and_then(|state| state.downcast::().ok()); + let state = state.and_then(|state| state.downcast::().ok()); let Some(state) = state else { return Ok(Some(event)) }; 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( - tool_json_result_payload(tool_name, content, false, &state.tool_call_id), + tool_json_result_payload(tool_name, content, false, &state.correlation_id), Some(state.context_table.clone()), ) .await; diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs index 7d9b2eb..6e8eff3 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs @@ -23,7 +23,16 @@ pub(crate) struct GatewayBackendClient { initialize_request: InitializeRequestParams, plugin_runtime: Option, in_flight_calls: Arc>>>, - resource_subscriptions: Arc>>>, + /// Keyed by the URI the client subscribed to; see [`Self::track_resource_subscription`]. + resource_subscriptions: Arc>>, +} + +/// A tracked subscription. `upstream_uri` is what the backend was actually subscribed to, which +/// differs from the map key whenever a resource pre hook rewrote it. +#[derive(Debug)] +struct ResourceSubscription { + upstream_uri: String, + downstream: Peer, } #[derive(Debug)] @@ -90,21 +99,43 @@ impl GatewayBackendClient { calls.get(progress_token).cloned() } - pub(crate) async fn track_resource_subscription(&self, resource_uri: &str, downstream: Peer) { - debug!("track_resource_subscription backend {} uri {resource_uri}", self.backend_name); + /// Records a subscription keyed by the URI the backend will notify about, remembering the URI + /// the client actually subscribed to. + /// + /// A resource pre hook may rewrite the URI, so the two can differ. Both directions of that + /// mapping have to be persisted here rather than recomputed: rerunning the hook on unsubscribe + /// can produce a different URI for a stateful plugin, and notifications must be reported to the + /// client under the URI it asked for, not the rewritten one. + pub(crate) async fn track_resource_subscription( + &self, + upstream_uri: &str, + client_uri: &str, + downstream: Peer, + ) { + debug!( + "track_resource_subscription backend {} upstream_uri {upstream_uri} client_uri {client_uri}", + self.backend_name + ); let mut subscriptions = self.resource_subscriptions.lock().await; - subscriptions.insert(resource_uri.to_owned(), downstream); + subscriptions + .insert(client_uri.to_owned(), ResourceSubscription { upstream_uri: upstream_uri.to_owned(), downstream }); } - pub(crate) async fn stop_tracking_resource_subscription(&self, resource_uri: &str) { - debug!("stop_tracking_resource_subscription backend {} uri {resource_uri}", self.backend_name); + /// Drops the subscription the client registered for `client_uri`, returning the URI the + /// backend was actually subscribed to so the caller can unsubscribe upstream with it. + pub(crate) async fn stop_tracking_resource_subscription(&self, client_uri: &str) -> Option { + debug!("stop_tracking_resource_subscription backend {} client_uri {client_uri}", self.backend_name); let mut subscriptions = self.resource_subscriptions.lock().await; - subscriptions.remove(resource_uri); + subscriptions.remove(client_uri).map(|subscription| subscription.upstream_uri) } - async fn resource_subscription(&self, resource_uri: &str) -> Option> { + /// Resolves a backend notification URI back to the subscribing peer and the URI that peer + /// subscribed to. + async fn resource_subscription(&self, upstream_uri: &str) -> Option<(Peer, String)> { let subscriptions = self.resource_subscriptions.lock().await; - subscriptions.get(resource_uri).cloned() + subscriptions.iter().find_map(|(client_uri, subscription)| { + (subscription.upstream_uri == upstream_uri).then(|| (subscription.downstream.clone(), client_uri.clone())) + }) } async fn stream_event_post_hook(&self, call: &InFlightToolCall, event: T) -> Option @@ -160,12 +191,14 @@ impl ClientHandler for GatewayBackendClient { mut params: ResourceUpdatedNotificationParam, _context: NotificationContext, ) { - let Some(downstream) = self.resource_subscription(¶ms.uri).await else { + let Some((downstream, client_uri)) = self.resource_subscription(¶ms.uri).await else { debug!("resource_updated: dropping backend notification for unsubscribed uri {}", params.uri); return; }; - params.uri = resource_uri_for_downstream(&self.backend_name, params.uri, self.namespace_identifiers); + // Report the update under the URI the client subscribed to. A pre hook may have rewritten + // the URI sent upstream, and the client never saw that value. + params.uri = resource_uri_for_downstream(&self.backend_name, client_uri, self.namespace_identifiers); if let Err(error) = downstream.notify_resource_updated(params).await { warn!("resource_updated: unable to forward backend notification downstream: {error:?}"); } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/completion.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/completion.rs index 7fe5211..36bb34e 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/completion.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/completion.rs @@ -1,3 +1,4 @@ +use contextforge_gateway_rs_cpex::{CompletionRequest, CompletionTarget}; use rmcp::{ ErrorData, RoleServer, model::{CompleteRequestParams, CompleteResult, Reference}, @@ -25,9 +26,10 @@ where let (virtual_host, session_id, claims) = mcp_call_validator.validate()?; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); - let identifier = match &request.r#ref { - Reference::Prompt(prompt) => prompt.name.as_str(), - Reference::Resource(resource) => resource.uri.as_str(), + // CPEX has no completion hook, so a completion runs the hooks of whatever is being completed. + let (identifier, target) = match &request.r#ref { + Reference::Prompt(prompt) => (prompt.name.as_str(), CompletionTarget::Prompt), + Reference::Resource(resource) => (resource.uri.as_str(), CompletionTarget::Resource), _ => return Err(ErrorData::invalid_params("Unsupported completion reference", None)), }; @@ -39,16 +41,41 @@ where ) .await?; + // Hooks run after routing, so plugins see the backend-local identifier and the backend that + // owns it. The pre hook is deny-only: the argument value is a partial keystroke, so rewriting + // it would return completions for text the user never typed. + let post_state = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => { + let context = request.context.as_ref().and_then(|context| context.arguments.as_ref()); + plugin_runtime + .before_complete(&CompletionRequest { + target, + identifier: &routed_identifier, + backend_name: &service_name, + argument: &request.argument, + context, + }) + .await? + }, + None => None, + }; + let mut routed_request = request; match &mut routed_request.r#ref { - Reference::Prompt(prompt) => prompt.name = routed_identifier, - Reference::Resource(resource) => resource.uri = routed_identifier, + Reference::Prompt(prompt) => prompt.name.clone_from(&routed_identifier), + Reference::Resource(resource) => resource.uri.clone_from(&routed_identifier), _ => return Err(ErrorData::invalid_params("Unsupported completion reference", None)), } let response = service .complete(routed_request) .await .map_err(|error| backend_forward_error("complete", &service_name, &error))?; + let response = match (&mcp_service.plugin_runtime, post_state) { + (Some(plugin_runtime), Some(post_state)) => { + plugin_runtime.after_complete(target, &routed_identifier, response, Some(post_state)).await? + }, + _ => response, + }; info!("complete: backend {service_name} returned {} values", response.completion.values.len()); Ok(response) } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index 14c0654..0d4c561 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -1,3 +1,4 @@ +use contextforge_gateway_rs_cpex::PromptPreFetchResult; use rmcp::{ ErrorData, RoleServer, model::{GetPromptRequestParams, GetPromptResponse, ListPromptsResult, PaginatedRequestParams}, @@ -27,6 +28,17 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); + + // A listing has no single backend, so the pre hook runs once before fan-out as a deny gate and + // the post hook sees the merged, namespaced page. Both are read-only for this path. + let post_state = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => { + let cursor = request.as_ref().and_then(|request| request.cursor.as_deref()); + plugin_runtime.before_list_prompts(cursor).await? + }, + None => None, + }; + let all_transports: Vec<_> = session_manager.borrow_transports().await; let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_prompts")?; @@ -59,6 +71,10 @@ where .await; let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers, &gateway_cursor, "list_prompts"); + if let (Some(plugin_runtime), Some(post_state)) = (&mcp_service.plugin_runtime, post_state) { + plugin_runtime.after_list_prompts(&prompts, Some(post_state)).await?; + } + let mut result = ListPromptsResult::with_all_items(prompts); result.next_cursor = next_cursor; Ok(result) @@ -84,12 +100,27 @@ where ) .await?; + // Hooks run after routing, so plugins see the backend-local prompt name and the backend that + // owns it, matching the `call_tool` contract. + let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime { + plugin_runtime.before_get_prompt(&request, &prompt_name, &service_name).await? + } else { + PromptPreFetchResult::unchanged() + }; + let post_state = pre_result.state; let mut routed_request = request; - routed_request.name = prompt_name; + pre_result.arguments.apply_to_request(&mut routed_request, &prompt_name); + let response = service .get_prompt(routed_request) .await .map_err(|error| backend_forward_error("get_prompt", &service_name, &error))?; + let response = match (&mcp_service.plugin_runtime, post_state) { + (Some(plugin_runtime), Some(post_state)) => { + plugin_runtime.after_get_prompt(&prompt_name, response, Some(post_state)).await? + }, + _ => response, + }; info!("get_prompt: backend {service_name} returned {} messages", response.messages.len()); Ok(response.into()) } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index d75dea0..16f67ce 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -1,3 +1,4 @@ +use contextforge_gateway_rs_cpex::ResourcePreFetchResult; use rmcp::{ ErrorData, RoleServer, model::{ @@ -6,7 +7,7 @@ use rmcp::{ }, service::RequestContext, }; -use tracing::info; +use tracing::{debug, info}; use super::McpService; use crate::gateway::{ @@ -30,6 +31,14 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); + + // A listing has no single backend, so the pre hook runs once before fan-out as a deny gate and + // the post hook sees the merged, namespaced page. Both are read-only for this path. + let post_state = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => plugin_runtime.before_list_resources().await?, + None => None, + }; + let all_transports: Vec<_> = session_manager.borrow_transports().await; let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_resources")?; @@ -62,6 +71,10 @@ where .await; let (resources, next_cursor) = merge_resources(responses, namespace_identifiers, &gateway_cursor, "list_resources"); + if let (Some(plugin_runtime), Some(post_state)) = (&mcp_service.plugin_runtime, post_state) { + plugin_runtime.after_list_resources(&resources, Some(post_state)).await?; + } + let mut result = ListResourcesResult::with_all_items(resources); result.next_cursor = next_cursor; Ok(result) @@ -87,12 +100,26 @@ where ) .await?; + // Hooks run after routing, so plugins see the backend-local URI and the backend that owns it. + let pre_result = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => plugin_runtime.before_read_resource(&resource_uri, &service_name).await?, + None => ResourcePreFetchResult::unchanged(), + }; + let post_state = pre_result.state; + let resource_uri = pre_result.uri.apply_to(resource_uri); + let mut routed_request = request; routed_request.uri = resource_uri; let response = service .read_resource(routed_request) .await .map_err(|error| backend_forward_error("read_resource", &service_name, &error))?; + let response = match (&mcp_service.plugin_runtime, post_state) { + (Some(plugin_runtime), Some(post_state)) => { + plugin_runtime.after_read_resource(response, Some(post_state)).await? + }, + _ => response, + }; info!("read_resource: backend {service_name} returned {} contents", response.contents.len()); Ok(response.into()) } @@ -110,6 +137,14 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); + + // A listing has no single backend, so the pre hook runs once before fan-out as a deny gate and + // the post hook sees the merged, namespaced page. Both are read-only for this path. + let post_state = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => plugin_runtime.before_list_resource_templates().await?, + None => None, + }; + let all_transports: Vec<_> = session_manager.borrow_transports().await; let gateway_cursor = @@ -144,6 +179,10 @@ where let (resource_templates, next_cursor) = merge_resource_templates(responses, namespace_identifiers, &gateway_cursor, "list_resource_templates"); + if let (Some(plugin_runtime), Some(post_state)) = (&mcp_service.plugin_runtime, post_state) { + plugin_runtime.after_list_resource_templates(&resource_templates, Some(post_state)).await?; + } + let mut result = ListResourceTemplatesResult::with_all_items(resource_templates); result.next_cursor = next_cursor; Ok(result) @@ -170,12 +209,26 @@ where ) .await?; + // Hooks run after routing, so plugins see the backend-local URI and the backend that owns it. + // There is no post hook: `subscribe` returns an empty acknowledgement. + // + // A rewritten URI is recorded against the URI the client subscribed to. The mapping has to be + // stored rather than recomputed later: a stateful plugin can return a different URI on the next + // call, which would leave the backend subscription and the tracking entry orphaned. + let client_uri = resource_uri.clone(); + let upstream_uri = match &mcp_service.plugin_runtime { + Some(plugin_runtime) => { + plugin_runtime.before_subscribe(&resource_uri, &service_name).await?.apply_to(resource_uri) + }, + None => resource_uri, + }; + let mut routed_request = request; - routed_request.uri = resource_uri.clone(); - service.service().track_resource_subscription(&resource_uri, cx.peer.clone()).await; + routed_request.uri = upstream_uri.clone(); + service.service().track_resource_subscription(&upstream_uri, &client_uri, cx.peer.clone()).await; if let Err(error) = service.subscribe(routed_request).await { - service.service().stop_tracking_resource_subscription(&resource_uri).await; + service.service().stop_tracking_resource_subscription(&client_uri).await; return Err(backend_forward_error("subscribe", &service_name, &error)); } info!("subscribe: backend {service_name} completed"); @@ -203,13 +256,26 @@ where ) .await?; + // Same contract as `subscribe`: pre hook only. The hook still runs so a plugin can deny an + // unsubscribe, but its URI rewrite is *not* used to address the backend — the URI recorded at + // subscribe time is. Recomputing it here would unsubscribe the wrong resource whenever a + // plugin's rewrite is stateful or non-deterministic. + if let Some(plugin_runtime) = &mcp_service.plugin_runtime { + plugin_runtime.before_unsubscribe(&resource_uri, &service_name).await?; + } + + let upstream_uri = + service.service().stop_tracking_resource_subscription(&resource_uri).await.unwrap_or_else(|| { + debug!("unsubscribe: no tracked subscription for {resource_uri}, forwarding it unchanged"); + resource_uri.clone() + }); + let mut routed_request = request; - routed_request.uri = resource_uri.clone(); + routed_request.uri = upstream_uri; service .unsubscribe(routed_request) .await .map_err(|error| backend_forward_error("unsubscribe", &service_name, &error))?; - service.service().stop_tracking_resource_subscription(&resource_uri).await; info!("unsubscribe: backend {service_name} completed"); Ok(()) } diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_completion_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_completion_plugins.rs new file mode 100644 index 0000000..6fd1488 --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_completion_plugins.rs @@ -0,0 +1,211 @@ +//! End-to-end CPEX hook coverage for `completion/complete`. +//! +//! CPEX defines no completion hook, so a completion runs the hooks of whatever is being completed: +//! a prompt reference dispatches to the prompt hooks, a resource or template reference to the +//! resource hooks. Plugins tell a completion apart from a plain fetch by the `gateway-completion-*` +//! correlation id. +//! +//! The pre hook is deny-only. The argument value is a partial user keystroke, so rewriting it would +//! return completions for text the user never typed. The post hook may rewrite the returned values, +//! but the count must match so the backend's `total` and `has_more` stay accurate. + +mod support; + +use std::sync::Arc; + +use cpex::cpex_core::hooks::types::cmf_hook_names; +use rmcp::model::{ + ArgumentInfo, CompleteRequestParams, ErrorCode, PromptReference, Reference, ResourceTemplateReference, +}; + +use support::{ + COMPLETION_VALUES, PROMPT_PRE_DENY_ERROR_CODE, PromptBehavior, PromptTestPlugin, RESOURCE_POST_DENY_ERROR_CODE, + RESOURCE_PRE_DENY_ERROR_CODE, ResourceBehavior, ResourceTestPlugin, TEMPLATE_URI, TEST_USER_ID, error_code, + runtime_with_prompt_plugin, runtime_with_resource_plugin, start_gateway, +}; + +fn prompt_completion() -> CompleteRequestParams { + CompleteRequestParams::new(Reference::Prompt(PromptReference::new("review")), ArgumentInfo::new("topic", "wea")) +} + +fn resource_completion() -> CompleteRequestParams { + CompleteRequestParams::new( + Reference::Resource(ResourceTemplateReference::new(TEMPLATE_URI)), + ArgumentInfo::new("id", "4"), + ) +} + +fn expected_values() -> Vec { + COMPLETION_VALUES.iter().map(|value| (*value).to_owned()).collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_reference_completion_runs_the_prompt_hooks() { + let plugin = Arc::new(PromptTestPlugin::new( + "completion-prompt", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.complete(prompt_completion()).await.expect("completion succeeds"); + + assert_eq!(expected_values(), result.completion.values); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + assert_eq!(Some("review"), observations.pre_name.as_deref()); + assert_eq!(Some(gateway.backend_name.as_str()), observations.pre_server_id.as_deref()); + // The correlation id marks this as a completion rather than a `prompts/get`. + let request_id = observations.pre_request_id.clone().expect("pre hook ran"); + assert!(request_id.starts_with("gateway-completion-"), "unexpected correlation id {request_id}"); + assert_eq!(observations.pre_request_id, observations.post_request_id); + assert_eq!(expected_values(), observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_reference_completion_post_hook_rewrites_values() { + let plugin = Arc::new(PromptTestPlugin::new( + "completion-prompt-rewrite", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::RewriteText, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.complete(prompt_completion()).await.expect("completion succeeds"); + + assert_eq!(vec!["redacted:alpha".to_owned(), "redacted:beta".to_owned()], result.completion.values); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_prompt_reference_completion_never_reaches_the_backend() { + let plugin = Arc::new(PromptTestPlugin::new( + "completion-prompt-deny", + vec![cmf_hook_names::PROMPT_PRE_FETCH], + PromptBehavior::DenyPre, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.complete(prompt_completion()).await.expect_err("completion deny is an error"); + + assert_eq!(ErrorCode(PROMPT_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.completions.lock().expect("backend completions lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_reference_completion_runs_the_resource_hooks() { + let plugin = Arc::new(ResourceTestPlugin::new( + "completion-resource", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.complete(resource_completion()).await.expect("completion succeeds"); + + assert_eq!(expected_values(), result.completion.values); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + assert_eq!(Some(TEMPLATE_URI), observations.pre_uri.as_deref()); + assert_eq!(Some(gateway.backend_name.as_str()), observations.pre_backend.as_deref()); + assert_eq!(expected_values(), observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_reference_completion_post_hook_rewrites_values() { + let plugin = Arc::new(ResourceTestPlugin::new( + "completion-resource-rewrite", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::RewriteText, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.complete(resource_completion()).await.expect("completion succeeds"); + + assert_eq!(vec!["redacted:alpha".to_owned(), "redacted:beta".to_owned()], result.completion.values); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_resource_reference_completion_never_reaches_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "completion-resource-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.complete(resource_completion()).await.expect_err("completion deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.completions.lock().expect("backend completions lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_resource_reference_completion_post_hook_errors_after_the_backend_ran() { + let plugin = Arc::new(ResourceTestPlugin::new( + "completion-resource-post-deny", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::DenyPost, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.complete(resource_completion()).await.expect_err("completion post deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_POST_DENY_ERROR_CODE), error_code(error)); + assert_eq!(1, gateway.backend_state.completions.lock().expect("backend completions lock poisoned").len()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_reference_completion_does_not_run_resource_hooks() { + let plugin = Arc::new(ResourceTestPlugin::new( + "completion-resource-only", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.complete(prompt_completion()).await.expect("completion succeeds"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_reference_completion_does_not_run_prompt_hooks() { + let plugin = Arc::new(PromptTestPlugin::new( + "completion-prompt-only", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.complete(resource_completion()).await.expect("completion succeeds"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs index 8273937..d354404 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs @@ -421,7 +421,7 @@ async fn secrets_detection_clean_tool_payload_passes_through_unchanged() { assert!(!result_text.contains("[redacted]")); let backend_calls = gateway.backend_state.calls.lock().expect("backend calls lock poisoned"); assert_eq!(1, backend_calls.len()); - assert_eq!("sum", backend_calls[0].tool_name); + assert_eq!("sum", backend_calls[0].name); let args = backend_calls[0].args.as_ref().expect("backend call has args"); assert_eq!(2, args.len()); assert_eq!(Some(&Value::from(1)), args.get("a")); @@ -516,7 +516,7 @@ async fn pre_hook_modifies_backend_arguments_without_rerouting_tool() { assert_eq!((REWRITTEN_SUM_A + REWRITTEN_SUM_B).to_string(), text(&result)); let backend_calls = gateway.backend_state.calls.lock().expect("backend calls lock poisoned"); - assert_eq!("sum", backend_calls[0].tool_name); + assert_eq!("sum", backend_calls[0].name); assert_eq!(Some(&Value::from(REWRITTEN_SUM_A)), backend_calls[0].args.as_ref().and_then(|args| args.get("a"))); let observations = observations.lock().expect("observations lock poisoned"); diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_prompt_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_prompt_plugins.rs new file mode 100644 index 0000000..ff180ee --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_prompt_plugins.rs @@ -0,0 +1,322 @@ +//! End-to-end CPEX prompt hook coverage for `prompts/get` and `prompts/list`. +//! +//! The tool equivalents live in `gateway_plugins.rs`; these prove the same contract holds on the +//! non-tool pass-through paths: hooks run after routing, a denial stops the request before the +//! backend sees it, and mutations reach the backend (pre) and the client (post). +//! +//! `prompts/list` differs deliberately. It fans out to every backend, so its pre hook runs once +//! before fan-out as a deny gate and its post hook sees the merged, namespaced listing. Exposure +//! there is read-only: MCP `Prompt` metadata has no CMF equivalent to be rebuilt from, so a +//! modified payload is ignored rather than partially applied. + +mod support; + +use std::sync::Arc; + +use cpex::cpex_core::hooks::types::cmf_hook_names; +use rmcp::model::{ContentBlock, ErrorCode, GetPromptRequestParams, GetPromptResult}; +use serde_json::{Value, json}; + +use support::{ + PROMPT_DESCRIPTION, PROMPT_POST_DENY_ERROR_CODE, PROMPT_PRE_DENY_ERROR_CODE, PromptBehavior, PromptTestPlugin, + REWRITTEN_PROMPT_TOPIC, TEST_USER_ID, error_code, runtime_with_prompt_plugin, start_gateway, +}; + +fn review_request(topic: &str) -> GetPromptRequestParams { + GetPromptRequestParams::new("review") + .with_arguments(serde_json::Map::from_iter([("topic".to_owned(), json!(topic))])) +} + +fn prompt_text(result: &GetPromptResult) -> String { + result + .messages + .iter() + .filter_map(|message| match &message.content { + ContentBlock::Text(text) => Some(text.text.clone()), + _ => None, + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_hooks_are_skipped_when_no_plugin_is_configured() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-allow", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + // Runtime plugins disabled: the gateway is built without a plugin handle at all. + let gateway = start_gateway(TEST_USER_ID, false, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.get_prompt(review_request("weather")).await.expect("prompt is returned"); + + assert_eq!("review of weather", prompt_text(&result)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_pre_hook_sees_routed_name_and_backend() { + let plugin = + Arc::new(PromptTestPlugin::new("prompt-pre", vec![cmf_hook_names::PROMPT_PRE_FETCH], PromptBehavior::Allow)); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.get_prompt(review_request("weather")).await.expect("prompt is returned"); + + assert_eq!("review of weather", prompt_text(&result)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(Some("review"), observations.pre_name.as_deref()); + assert_eq!(Some(gateway.backend_name.as_str()), observations.pre_server_id.as_deref()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_pre_hook_rewrites_arguments_reaching_the_backend() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-pre-rewrite", + vec![cmf_hook_names::PROMPT_PRE_FETCH], + PromptBehavior::RewriteArguments, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.get_prompt(review_request("weather")).await.expect("prompt is returned"); + + assert_eq!(format!("review of {REWRITTEN_PROMPT_TOPIC}"), prompt_text(&result)); + let prompt_calls = gateway.backend_state.prompts.lock().expect("backend prompts lock poisoned"); + assert_eq!("review", prompt_calls[0].name); + assert_eq!( + Some(&Value::from(REWRITTEN_PROMPT_TOPIC)), + prompt_calls[0].args.as_ref().and_then(|args| args.get("topic")) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_post_hook_rewrites_text_returned_to_the_client() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-post", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::RewriteText, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.get_prompt(review_request("weather")).await.expect("prompt is returned"); + + assert_eq!("redacted:review of weather", prompt_text(&result)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(vec!["review of weather".to_owned()], observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_pre_and_post_hooks_share_one_request_id() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-both", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.get_prompt(review_request("weather")).await.expect("prompt is returned"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + assert!(observations.pre_request_id.is_some()); + assert_eq!(observations.pre_request_id, observations.post_request_id); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_prompt_pre_hook_never_reaches_the_backend() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-pre-deny", + vec![cmf_hook_names::PROMPT_PRE_FETCH], + PromptBehavior::DenyPre, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.get_prompt(review_request("weather")).await.expect_err("prompt pre deny is an error"); + + assert_eq!(ErrorCode(PROMPT_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.prompts.lock().expect("backend prompts lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_prompt_post_hook_returns_error_after_the_backend_ran() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-post-deny", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::DenyPost, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.get_prompt(review_request("weather")).await.expect_err("prompt post deny is an error"); + + assert_eq!(ErrorCode(PROMPT_POST_DENY_ERROR_CODE), error_code(error)); + assert_eq!(1, gateway.backend_state.prompts.lock().expect("backend prompts lock poisoned").len()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_list_pre_hook_marks_a_listing_with_an_empty_name() { + let plugin = + Arc::new(PromptTestPlugin::new("list-pre", vec![cmf_hook_names::PROMPT_PRE_FETCH], PromptBehavior::Allow)); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.list_prompts(None).await.expect("prompts are listed"); + + assert_eq!(1, result.prompts.len()); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + // Empty name is the listing marker that separates this from a `prompts/get`. + assert_eq!(Some(""), observations.pre_name.as_deref()); + assert_eq!(None, observations.pre_server_id); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_list_post_hook_sees_merged_names_and_descriptions() { + let plugin = + Arc::new(PromptTestPlugin::new("list-post", vec![cmf_hook_names::PROMPT_POST_FETCH], PromptBehavior::Allow)); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_prompts(None).await.expect("prompts are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(vec!["review".to_owned()], observations.post_prompt_names); + assert_eq!(vec![PROMPT_DESCRIPTION.to_owned()], observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_list_post_hook_modifications_are_not_applied() { + let plugin = Arc::new(PromptTestPlugin::new( + "list-mutate", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::RewriteListNames, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.list_prompts(None).await.expect("prompts are listed"); + + // `prompts/list` exposure is read-only: MCP prompt metadata cannot be rebuilt from CMF. + assert_eq!(vec!["review".to_owned()], result.prompts.iter().map(|p| p.name.clone()).collect::>()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_prompt_list_pre_hook_never_fans_out_to_backends() { + let plugin = Arc::new(PromptTestPlugin::new( + "list-pre-deny", + vec![cmf_hook_names::PROMPT_PRE_FETCH], + PromptBehavior::DenyPre, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.list_prompts(None).await.expect_err("prompt list pre deny is an error"); + + assert_eq!(ErrorCode(PROMPT_PRE_DENY_ERROR_CODE), error_code(error)); + assert_eq!(0, *gateway.backend_state.prompt_lists.lock().expect("backend prompt lists lock poisoned")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_prompt_list_post_hook_returns_error_after_fan_out() { + let plugin = Arc::new(PromptTestPlugin::new( + "list-post-deny", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::DenyPost, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.list_prompts(None).await.expect_err("prompt list post deny is an error"); + + assert_eq!(ErrorCode(PROMPT_POST_DENY_ERROR_CODE), error_code(error)); + assert_eq!(1, *gateway.backend_state.prompt_lists.lock().expect("backend prompt lists lock poisoned")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_list_pre_and_post_hooks_share_one_request_id() { + let plugin = Arc::new(PromptTestPlugin::new( + "list-both", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_prompts(None).await.expect("prompts are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + assert!(observations.pre_request_id.is_some()); + assert_eq!(observations.pre_request_id, observations.post_request_id); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn prompt_post_hook_removing_text_fails_closed() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-delete", + vec![cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::DeleteText, + )); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service + .get_prompt(review_request("weather")) + .await + .expect_err("removing text must not fall back to the backend message"); + + // Restoring the original here would return text a redaction plugin stripped. + assert_eq!(ErrorCode::INVALID_PARAMS, error_code(error)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn tool_hooks_do_not_run_for_prompt_requests() { + let plugin = Arc::new(PromptTestPlugin::new( + "prompt-only", + vec![cmf_hook_names::PROMPT_PRE_FETCH, cmf_hook_names::PROMPT_POST_FETCH], + PromptBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_prompt_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.call_tool(support::sum_request("sum", 1, 2)).await.expect("tool call succeeds"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_resource_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_resource_plugins.rs new file mode 100644 index 0000000..35a6b98 --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_resource_plugins.rs @@ -0,0 +1,509 @@ +//! End-to-end CPEX resource hook coverage for `resources/templates/list`. +//! +//! Template listings fan out to every backend, so the pre hook runs once before fan-out as a deny +//! gate and the post hook sees the merged, namespaced listing. Exposure is read-only: MCP +//! `ResourceTemplate` metadata (`title`, `mime_type`, `icons`, `annotations`, `_meta`) has no CMF +//! equivalent to be rebuilt from, so a modified payload is ignored rather than partially applied. +//! +//! `resources/subscribe` and `resources/unsubscribe` are routed paths instead: the pre hook runs +//! after routing and may deny or rewrite the target URI. Neither has a post hook, because both +//! return an empty acknowledgement with nothing to inspect. + +mod support; + +use std::sync::Arc; + +use cpex::cpex_core::hooks::types::cmf_hook_names; +use rmcp::model::{ + ErrorCode, ReadResourceRequestParams, ReadResourceResult, ResourceContents, SubscribeRequestParams, + UnsubscribeRequestParams, +}; + +use support::{ + RESOURCE_DESCRIPTION, RESOURCE_POST_DENY_ERROR_CODE, RESOURCE_PRE_DENY_ERROR_CODE, RESOURCE_TEXT, RESOURCE_URI, + REWRITTEN_RESOURCE_URI, ResourceBehavior, ResourceTestPlugin, TEMPLATE_DESCRIPTION, TEMPLATE_URI, TEST_USER_ID, + error_code, runtime_with_resource_plugin, start_gateway, +}; + +const SUBSCRIBED_URI: &str = "report://42/summary"; + +fn resource_texts(result: &ReadResourceResult) -> Vec { + result + .contents + .iter() + .filter_map(|contents| match contents { + ResourceContents::TextResourceContents { text, .. } => Some(text.clone()), + _ => None, + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_hooks_are_skipped_when_no_plugin_is_configured() { + let plugin = Arc::new(ResourceTestPlugin::new( + "resource-allow", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, false, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.list_resource_templates(None).await.expect("templates are listed"); + + assert_eq!(1, result.resource_templates.len()); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn template_list_pre_hook_marks_a_listing_with_an_empty_uri() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-pre", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_resource_templates(None).await.expect("templates are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + // Empty URI is the listing marker that separates this from a single-resource fetch. + assert_eq!(Some(""), observations.pre_uri.as_deref()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn template_list_post_hook_sees_merged_uris_names_and_descriptions() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-post", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_resource_templates(None).await.expect("templates are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(vec![TEMPLATE_URI.to_owned()], observations.post_uris); + assert_eq!(vec!["report".to_owned()], observations.post_names); + assert_eq!(vec![TEMPLATE_DESCRIPTION.to_owned()], observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn template_list_post_hook_modifications_are_not_applied() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-mutate", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::RewriteListUris, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.list_resource_templates(None).await.expect("templates are listed"); + + // Read-only exposure: MCP template metadata cannot be rebuilt from CMF content parts. + assert_eq!( + vec![TEMPLATE_URI.to_owned()], + result.resource_templates.iter().map(|template| template.uri_template.clone()).collect::>() + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_template_list_pre_hook_never_fans_out_to_backends() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-pre-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.list_resource_templates(None).await.expect_err("template list pre deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert_eq!(0, *gateway.backend_state.template_lists.lock().expect("backend template lists lock poisoned")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_template_list_post_hook_returns_error_after_fan_out() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-post-deny", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::DenyPost, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.list_resource_templates(None).await.expect_err("template list post deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_POST_DENY_ERROR_CODE), error_code(error)); + assert_eq!(1, *gateway.backend_state.template_lists.lock().expect("backend template lists lock poisoned")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn template_list_pre_and_post_hooks_share_one_request_id() { + let plugin = Arc::new(ResourceTestPlugin::new( + "template-both", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_resource_templates(None).await.expect("templates are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(1, observations.post_calls); + assert!(observations.pre_request_id.is_some()); + assert_eq!(observations.pre_request_id, observations.post_request_id); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn subscribe_pre_hook_sees_routed_uri_and_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "subscribe-pre", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("subscribe succeeds"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(Some(SUBSCRIBED_URI), observations.pre_uri.as_deref()); + assert_eq!(Some(gateway.backend_name.as_str()), observations.pre_backend.as_deref()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn subscribe_has_no_post_hook() { + let plugin = Arc::new(ResourceTestPlugin::new( + "subscribe-post", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("subscribe succeeds"); + service.unsubscribe(UnsubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("unsubscribe succeeds"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(2, observations.pre_calls); + // Both paths return an empty acknowledgement, so there is nothing for a post hook to see. + assert_eq!(0, observations.post_calls); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn subscribe_pre_hook_rewrites_the_uri_reaching_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "subscribe-rewrite", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::RewriteUri, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("subscribe succeeds"); + + let subscriptions = gateway.backend_state.subscriptions.lock().expect("backend subscriptions lock poisoned"); + assert_eq!(vec![REWRITTEN_RESOURCE_URI.to_owned()], *subscriptions); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn unsubscribe_pre_hook_rewrites_the_uri_reaching_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "unsubscribe-rewrite", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::RewriteUri, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("subscribe succeeds"); + service.unsubscribe(UnsubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("unsubscribe succeeds"); + + // Subscribe and unsubscribe rewrite to the same URI, so the tracking key matches on both sides. + let unsubscriptions = gateway.backend_state.unsubscriptions.lock().expect("backend lock poisoned"); + assert_eq!(vec![REWRITTEN_RESOURCE_URI.to_owned()], *unsubscriptions); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn denied_subscribe_pre_hook_never_reaches_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "subscribe-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect_err("subscribe deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.subscriptions.lock().expect("backend subscriptions lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn denied_unsubscribe_pre_hook_never_reaches_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "unsubscribe-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service + .unsubscribe(UnsubscribeRequestParams::new(SUBSCRIBED_URI)) + .await + .expect_err("unsubscribe deny is an error"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.unsubscriptions.lock().expect("backend lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn read_resource_pre_hook_sees_routed_uri_and_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-pre", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.read_resource(ReadResourceRequestParams::new(RESOURCE_URI)).await.expect("read succeeds"); + + assert_eq!(1, result.contents.len()); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.pre_calls); + assert_eq!(Some(RESOURCE_URI), observations.pre_uri.as_deref()); + assert_eq!(Some(gateway.backend_name.as_str()), observations.pre_backend.as_deref()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn read_resource_post_hook_rewrites_text_returned_to_the_client() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-post", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::RewriteContent, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.read_resource(ReadResourceRequestParams::new(RESOURCE_URI)).await.expect("read succeeds"); + + assert_eq!(vec![format!("redacted:{RESOURCE_TEXT}")], resource_texts(&result)); + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(vec![RESOURCE_TEXT.to_owned()], observations.post_contents); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn read_resource_pre_hook_rewrites_the_uri_reaching_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-rewrite-uri", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::RewriteUri, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.read_resource(ReadResourceRequestParams::new(RESOURCE_URI)).await.expect("read succeeds"); + + let reads = gateway.backend_state.reads.lock().expect("backend reads lock poisoned"); + assert_eq!(vec![REWRITTEN_RESOURCE_URI.to_owned()], *reads); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_read_resource_pre_hook_never_reaches_the_backend() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-pre-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = + service.read_resource(ReadResourceRequestParams::new(RESOURCE_URI)).await.expect_err("read deny errors"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert!(gateway.backend_state.reads.lock().expect("backend reads lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_read_resource_post_hook_errors_after_the_backend_ran() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-post-deny", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::DenyPost, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = + service.read_resource(ReadResourceRequestParams::new(RESOURCE_URI)).await.expect_err("read post deny errors"); + + assert_eq!(ErrorCode(RESOURCE_POST_DENY_ERROR_CODE), error_code(error)); + assert_eq!(1, gateway.backend_state.reads.lock().expect("backend reads lock poisoned").len()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_list_post_hook_sees_merged_uris_and_descriptions() { + let plugin = Arc::new(ResourceTestPlugin::new( + "resource-list-post", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.list_resources(None).await.expect("resources are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(1, observations.post_calls); + assert_eq!(vec![RESOURCE_URI.to_owned()], observations.post_uris); + assert_eq!(vec![RESOURCE_DESCRIPTION.to_owned()], observations.post_texts); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_list_post_hook_modifications_are_not_applied() { + let plugin = Arc::new(ResourceTestPlugin::new( + "resource-list-mutate", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::RewriteListUris, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let result = service.list_resources(None).await.expect("resources are listed"); + + assert_eq!( + vec![RESOURCE_URI.to_owned()], + result.resources.iter().map(|resource| resource.uri.clone()).collect::>() + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn denied_resource_list_pre_hook_never_fans_out_to_backends() { + let plugin = Arc::new(ResourceTestPlugin::new( + "resource-list-deny", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::DenyPre, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service.list_resources(None).await.expect_err("resource list deny errors"); + + assert_eq!(ErrorCode(RESOURCE_PRE_DENY_ERROR_CODE), error_code(error)); + assert_eq!(0, *gateway.backend_state.resource_lists.lock().expect("backend resource lists lock poisoned")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn read_resource_post_hook_removing_text_fails_closed() { + let plugin = Arc::new(ResourceTestPlugin::new( + "read-delete", + vec![cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::DeleteContent, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + let error = service + .read_resource(ReadResourceRequestParams::new(RESOURCE_URI)) + .await + .expect_err("removing text must not fall back to the backend content"); + + // Restoring the original here would return content a redaction plugin stripped. + assert_eq!(ErrorCode::INVALID_PARAMS, error_code(error)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[expect(deprecated, reason = "legacy RMCP subscription coverage; listen migration is deferred")] +async fn unsubscribe_uses_the_uri_recorded_at_subscribe_not_a_fresh_rewrite() { + let plugin = Arc::new(ResourceTestPlugin::new( + "sub-stateful", + vec![cmf_hook_names::RESOURCE_PRE_FETCH], + ResourceBehavior::RewriteUriPerCall, + )); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.subscribe(SubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("subscribe succeeds"); + service.unsubscribe(UnsubscribeRequestParams::new(SUBSCRIBED_URI)).await.expect("unsubscribe succeeds"); + + // The plugin rewrites to a different URI on every call. Unsubscribe must use the URI recorded + // at subscribe time, or the backend subscription is orphaned. + let subscribed = gateway.backend_state.subscriptions.lock().expect("backend lock poisoned").clone(); + let unsubscribed = gateway.backend_state.unsubscriptions.lock().expect("backend lock poisoned").clone(); + assert_eq!(1, subscribed.len()); + assert_eq!(subscribed, unsubscribed, "unsubscribe addressed a different URI than subscribe"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn resource_hooks_do_not_run_for_prompt_or_tool_requests() { + let plugin = Arc::new(ResourceTestPlugin::new( + "resource-only", + vec![cmf_hook_names::RESOURCE_PRE_FETCH, cmf_hook_names::RESOURCE_POST_FETCH], + ResourceBehavior::Allow, + )); + let observations = plugin.observations(); + let runtime = runtime_with_resource_plugin(plugin).await; + + let gateway = start_gateway(TEST_USER_ID, true, runtime).await; + let service = gateway.connect(TEST_USER_ID).await; + service.call_tool(support::sum_request("sum", 1, 2)).await.expect("tool call succeeds"); + service.list_prompts(None).await.expect("prompts are listed"); + + let observations = observations.lock().expect("observations lock poisoned"); + assert_eq!(0, observations.pre_calls); + assert_eq!(0, observations.post_calls); +} diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs index 54f6898..1a5c098 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs @@ -24,9 +24,18 @@ pub(crate) use list_tools_gateway::{ create_tls_gateway_with_four_tls_counters, plaintext_config, }; pub(crate) use plugin::{ - POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, TestPlugin, TestPluginFactory, + POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, PROMPT_POST_DENY_ERROR_CODE, PROMPT_PRE_DENY_ERROR_CODE, PromptBehavior, + PromptTestPlugin, PromptTestPluginFactory, RESOURCE_POST_DENY_ERROR_CODE, RESOURCE_PRE_DENY_ERROR_CODE, + REWRITTEN_PROMPT_TOPIC, REWRITTEN_RESOURCE_URI, REWRITTEN_SUM_A, REWRITTEN_SUM_B, ResourceBehavior, + ResourceTestPlugin, ResourceTestPluginFactory, TestPlugin, TestPluginFactory, +}; +pub(crate) use plugin_gateway::{ + COMPLETION_VALUES, PROMPT_DESCRIPTION, RESOURCE_DESCRIPTION, RESOURCE_TEXT, RESOURCE_URI, RunningGateway, + TEMPLATE_DESCRIPTION, TEMPLATE_URI, start_gateway, start_gateway_with_json_backend_responses, +}; +pub(crate) use runtime::{ + runtime_with_post, runtime_with_pre, runtime_with_pre_and_post, runtime_with_prompt_plugin, + runtime_with_resource_plugin, }; -pub(crate) use plugin_gateway::{RunningGateway, start_gateway, start_gateway_with_json_backend_responses}; -pub(crate) use runtime::{runtime_with_post, runtime_with_pre, runtime_with_pre_and_post}; pub(crate) use tool::{error_code, sum_request, text}; pub(crate) use user_config_store::MemoryUserConfigStore; diff --git a/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs b/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs index ddc0210..fe2ef7d 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs @@ -298,6 +298,463 @@ fn cmf_result_text(payload: &MessagePayload) -> String { .map_or_else(|| payload.message.get_text_content(), |result| text(&result)) } +pub(crate) const PROMPT_PRE_DENY_ERROR_CODE: i32 = -32011; +pub(crate) const PROMPT_POST_DENY_ERROR_CODE: i32 = -32012; +pub(crate) const REWRITTEN_PROMPT_TOPIC: &str = "rewritten-topic"; + +#[derive(Default)] +pub(crate) struct PromptObservations { + pub(crate) pre_calls: usize, + pub(crate) post_calls: usize, + pub(crate) pre_name: Option, + pub(crate) pre_server_id: Option, + pub(crate) pre_request_id: Option, + pub(crate) post_request_id: Option, + pub(crate) post_texts: Vec, + pub(crate) post_prompt_names: Vec, +} + +#[derive(Clone, Copy, Default)] +pub(crate) enum PromptBehavior { + #[default] + Allow, + RewriteArguments, + DenyPre, + RewriteText, + DenyPost, + /// Strips every text part, as a redaction plugin dropping sensitive data would. + DeleteText, + /// Renames every listed prompt. `prompts/list` exposure is read-only, so the client must still + /// receive the original names. + RewriteListNames, +} + +/// Prompt hooks carry `PromptRequest` / `PromptResult` content parts rather than the tool parts +/// [`TestPlugin`] handles. Pre and post are told apart by the CMF role — `User` on the request +/// side, `Assistant` on the response side — which holds for both `prompts/get` and `prompts/list`. +/// Content part alone does not work: a `prompts/list` post payload is also made of +/// `PromptRequest` parts. +pub(crate) struct PromptTestPlugin { + pub(crate) config: PluginConfig, + pub(crate) observations: Arc>, + behavior: PromptBehavior, +} + +impl PromptTestPlugin { + pub(crate) fn new(name: &str, hooks: Vec<&'static str>, behavior: PromptBehavior) -> Self { + Self { + config: PluginConfig { + name: name.to_owned(), + kind: "prompt-test".to_owned(), + hooks: hooks.into_iter().map(str::to_owned).collect(), + ..Default::default() + }, + observations: Arc::new(Mutex::new(PromptObservations::default())), + behavior, + } + } + + pub(crate) fn observations(&self) -> Arc> { + Arc::clone(&self.observations) + } +} + +#[async_trait] +impl Plugin for PromptTestPlugin { + fn config(&self) -> &PluginConfig { + &self.config + } +} + +impl HookHandler for PromptTestPlugin { + async fn handle( + &self, + payload: &MessagePayload, + _extensions: &Extensions, + _ctx: &mut PluginContext, + ) -> PluginResult { + let is_post = payload.message.role == Role::Assistant; + let mut observations = self.observations.lock().expect("observations lock poisoned"); + if is_post { + observations.post_calls += 1; + if let Some(result) = payload.message.get_prompt_results().first() { + observations.post_request_id = Some(result.prompt_request_id.clone()); + // `prompts/get` carries text inside the rendered messages; a prompt-reference + // completion carries its values as top-level text parts. + observations.post_texts = result + .messages + .iter() + .flat_map(|message| message.content.iter()) + .chain(payload.message.content.iter()) + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + } else { + // `prompts/list` post: the listing arrives as one `PromptRequest` per prompt, with + // each description following as text. + let listed = payload.message.get_prompt_requests(); + observations.post_request_id = listed.first().map(|request| request.prompt_request_id.clone()); + observations.post_prompt_names = listed.iter().map(|request| request.name.clone()).collect(); + observations.post_texts = payload + .message + .content + .iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + } + } else { + observations.pre_calls += 1; + if let Some(request) = payload.message.get_prompt_requests().first() { + observations.pre_name = Some(request.name.clone()); + observations.pre_server_id.clone_from(&request.server_id); + observations.pre_request_id = Some(request.prompt_request_id.clone()); + } + } + drop(observations); + + match (is_post, self.behavior) { + (false, PromptBehavior::RewriteArguments) => { + let mut modified = payload.clone(); + if let Some(ContentPart::PromptRequest { content }) = + modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::PromptRequest { .. })) + { + content.arguments = HashMap::from([("topic".to_owned(), json!(REWRITTEN_PROMPT_TOPIC))]); + } + PluginResult::modify_payload(modified) + }, + (false, PromptBehavior::DenyPre) => PluginResult::deny( + PluginViolation::new("prompt_pre_denied", "prompt pre denied") + .with_proto_error_code(i64::from(PROMPT_PRE_DENY_ERROR_CODE)), + ), + (true, PromptBehavior::RewriteText) => { + let mut modified = payload.clone(); + // `prompts/get` nests text inside the rendered messages; a prompt-reference + // completion carries its values as top-level text parts. + for part in &mut modified.message.content { + match part { + ContentPart::PromptResult { content } => { + for message in &mut content.messages { + for part in &mut message.content { + if let ContentPart::Text { text } = part { + *text = format!("redacted:{text}"); + } + } + } + }, + ContentPart::Text { text } => *text = format!("redacted:{text}"), + _ => {}, + } + } + PluginResult::modify_payload(modified) + }, + (true, PromptBehavior::DeleteText) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::PromptResult { content } = part { + for message in &mut content.messages { + message.content.retain(|part| !matches!(part, ContentPart::Text { .. })); + } + } + } + PluginResult::modify_payload(modified) + }, + (true, PromptBehavior::RewriteListNames) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::PromptRequest { content } = part { + content.name = format!("mutated-{}", content.name); + } + } + PluginResult::modify_payload(modified) + }, + (true, PromptBehavior::DenyPost) => PluginResult::deny( + PluginViolation::new("prompt_post_denied", "prompt post denied") + .with_proto_error_code(i64::from(PROMPT_POST_DENY_ERROR_CODE)), + ), + _ => PluginResult::allow(), + } + } +} + +pub(crate) struct PromptTestPluginFactory { + observations: Arc>, + behavior: PromptBehavior, +} + +impl PromptTestPluginFactory { + pub(crate) fn from_plugin(plugin: &PromptTestPlugin) -> Self { + Self { observations: Arc::clone(&plugin.observations), behavior: plugin.behavior } + } +} + +impl PluginFactory for PromptTestPluginFactory { + fn create(&self, config: &PluginConfig) -> Result> { + let plugin = Arc::new(PromptTestPlugin { + config: config.clone(), + observations: Arc::clone(&self.observations), + behavior: self.behavior, + }); + let handlers = config + .hooks + .iter() + .filter_map(|hook| { + let hook = match hook.as_str() { + cmf_hook_names::PROMPT_PRE_FETCH => cmf_hook_names::PROMPT_PRE_FETCH, + cmf_hook_names::PROMPT_POST_FETCH => cmf_hook_names::PROMPT_POST_FETCH, + _ => return None, + }; + Some(( + hook, + Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) + as Arc, + )) + }) + .collect(); + let plugin: Arc = plugin; + Ok(PluginInstance { plugin, handlers }) + } +} + +pub(crate) const RESOURCE_PRE_DENY_ERROR_CODE: i32 = -32021; +pub(crate) const RESOURCE_POST_DENY_ERROR_CODE: i32 = -32022; +pub(crate) const REWRITTEN_RESOURCE_URI: &str = "rewritten://resource"; + +#[derive(Default)] +pub(crate) struct ResourceObservations { + pub(crate) pre_calls: usize, + pub(crate) post_calls: usize, + pub(crate) pre_uri: Option, + pub(crate) pre_backend: Option, + pub(crate) pre_request_id: Option, + pub(crate) post_request_id: Option, + pub(crate) post_uris: Vec, + pub(crate) post_names: Vec, + pub(crate) post_texts: Vec, + pub(crate) post_contents: Vec, +} + +#[derive(Clone, Copy, Default)] +pub(crate) enum ResourceBehavior { + #[default] + Allow, + DenyPre, + DenyPost, + /// Rewrites the target URI of a subscribe/unsubscribe request. + RewriteUri, + /// Rewrites completion values, which arrive as top-level text parts. + RewriteText, + /// Rewrites `resources/read` text contents. + RewriteContent, + /// Clears text contents, as a redaction plugin dropping sensitive data would. + DeleteContent, + /// Rewrites the target URI to a different value on each call, as a stateful plugin would. + RewriteUriPerCall, + /// Rewrites every listed template URI. `resources/templates/list` exposure is read-only, so the + /// client must still receive the original URI templates. + RewriteListUris, +} + +/// Resource hooks carry `ResourceRef` content parts. Pre and post are told apart by the CMF role, +/// the same convention the prompt hooks use. +pub(crate) struct ResourceTestPlugin { + pub(crate) config: PluginConfig, + pub(crate) observations: Arc>, + behavior: ResourceBehavior, +} + +impl ResourceTestPlugin { + pub(crate) fn new(name: &str, hooks: Vec<&'static str>, behavior: ResourceBehavior) -> Self { + Self { + config: PluginConfig { + name: name.to_owned(), + kind: "resource-test".to_owned(), + hooks: hooks.into_iter().map(str::to_owned).collect(), + ..Default::default() + }, + observations: Arc::new(Mutex::new(ResourceObservations::default())), + behavior, + } + } + + pub(crate) fn observations(&self) -> Arc> { + Arc::clone(&self.observations) + } +} + +#[async_trait] +impl Plugin for ResourceTestPlugin { + fn config(&self) -> &PluginConfig { + &self.config + } +} + +impl HookHandler for ResourceTestPlugin { + async fn handle( + &self, + payload: &MessagePayload, + _extensions: &Extensions, + _ctx: &mut PluginContext, + ) -> PluginResult { + let is_post = payload.message.role == Role::Assistant; + let references = payload.message.get_resource_refs(); + let mut observations = self.observations.lock().expect("observations lock poisoned"); + if is_post { + observations.post_calls += 1; + let resources = payload.message.get_resources(); + observations.post_request_id = references + .first() + .map(|item| item.resource_request_id.clone()) + .or_else(|| resources.first().map(|item| item.resource_request_id.clone())); + observations.post_uris = references + .iter() + .map(|item| item.uri.clone()) + .chain(resources.iter().map(|item| item.uri.clone())) + .collect(); + observations.post_names = references.iter().filter_map(|item| item.name.clone()).collect(); + observations.post_texts = payload + .message + .content + .iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + observations.post_contents = resources.iter().filter_map(|item| item.content.clone()).collect(); + } else { + observations.pre_calls += 1; + // Subscribe/unsubscribe payloads carry a `Resource`; listings carry `ResourceRef`. + let resources = payload.message.get_resources(); + if let Some(resource) = resources.first() { + observations.pre_uri = Some(resource.uri.clone()); + observations.pre_request_id = Some(resource.resource_request_id.clone()); + observations.pre_backend = + resource.annotations.get("backend").and_then(Value::as_str).map(str::to_owned); + } else if let Some(reference) = references.first() { + observations.pre_uri = Some(reference.uri.clone()); + observations.pre_request_id = Some(reference.resource_request_id.clone()); + } + } + drop(observations); + + match (is_post, self.behavior) { + (false, ResourceBehavior::DenyPre) => PluginResult::deny( + PluginViolation::new("resource_pre_denied", "resource pre denied") + .with_proto_error_code(i64::from(RESOURCE_PRE_DENY_ERROR_CODE)), + ), + (true, ResourceBehavior::DenyPost) => PluginResult::deny( + PluginViolation::new("resource_post_denied", "resource post denied") + .with_proto_error_code(i64::from(RESOURCE_POST_DENY_ERROR_CODE)), + ), + (false, ResourceBehavior::RewriteUri) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part { + REWRITTEN_RESOURCE_URI.clone_into(&mut content.uri); + } + } + PluginResult::modify_payload(modified) + }, + (true, ResourceBehavior::DeleteContent) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part { + content.content = None; + } + } + PluginResult::modify_payload(modified) + }, + (false, ResourceBehavior::RewriteUriPerCall) => { + let call = { + let observations = self.observations.lock().expect("observations lock poisoned"); + observations.pre_calls + }; + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part { + content.uri = format!("{REWRITTEN_RESOURCE_URI}-{call}"); + } + } + PluginResult::modify_payload(modified) + }, + (true, ResourceBehavior::RewriteContent) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Resource { content } = part + && let Some(text) = &content.content + { + content.content = Some(format!("redacted:{text}")); + } + } + PluginResult::modify_payload(modified) + }, + (true, ResourceBehavior::RewriteText) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::Text { text } = part { + *text = format!("redacted:{text}"); + } + } + PluginResult::modify_payload(modified) + }, + (true, ResourceBehavior::RewriteListUris) => { + let mut modified = payload.clone(); + for part in &mut modified.message.content { + if let ContentPart::ResourceRef { content } = part { + content.uri = format!("mutated://{}", content.uri); + } + } + PluginResult::modify_payload(modified) + }, + _ => PluginResult::allow(), + } + } +} + +pub(crate) struct ResourceTestPluginFactory { + observations: Arc>, + behavior: ResourceBehavior, +} + +impl ResourceTestPluginFactory { + pub(crate) fn from_plugin(plugin: &ResourceTestPlugin) -> Self { + Self { observations: Arc::clone(&plugin.observations), behavior: plugin.behavior } + } +} + +impl PluginFactory for ResourceTestPluginFactory { + fn create(&self, config: &PluginConfig) -> Result> { + let plugin = Arc::new(ResourceTestPlugin { + config: config.clone(), + observations: Arc::clone(&self.observations), + behavior: self.behavior, + }); + let handlers = config + .hooks + .iter() + .filter_map(|hook| { + let hook = match hook.as_str() { + 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(( + hook, + Arc::new(TypedHandlerAdapter::::new(Arc::clone(&plugin))) + as Arc, + )) + }) + .collect(); + let plugin: Arc = plugin; + Ok(PluginInstance { plugin, handlers }) + } +} + pub(crate) struct TestPluginFactory { pub(crate) observations: Arc>, pub(crate) pre_behavior: PreBehavior, diff --git a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs index edf75a5..023e118 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs @@ -15,9 +15,12 @@ use http::{HeaderMap, HeaderValue}; use rmcp::{ ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ErrorCode, Implementation, - InitializeRequestParams, InitializeResult, NumberOrString, ProgressNotificationParam, ProgressToken, - ServerCapabilities, + CallToolRequestParams, CallToolResponse, CallToolResult, CompleteRequestParams, CompleteResult, CompletionInfo, + ContentBlock, ErrorCode, GetPromptRequestParams, GetPromptResponse, GetPromptResult, Implementation, + InitializeRequestParams, InitializeResult, ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, + NumberOrString, PaginatedRequestParams, ProgressNotificationParam, ProgressToken, Prompt, PromptMessage, + ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, Reference, Resource, ResourceContents, + ResourceTemplate, Role, ServerCapabilities, SubscribeRequestParams, UnsubscribeRequestParams, }, service::{RequestContext, Service}, transport::{ @@ -33,18 +36,33 @@ use super::{MemoryUserConfigStore, token}; static GATEWAY_PORT_LOCK: OnceLock>> = OnceLock::new(); const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +pub(crate) const PROMPT_DESCRIPTION: &str = "Reviews a topic"; +pub(crate) const TEMPLATE_URI: &str = "report://{id}/summary"; +pub(crate) const TEMPLATE_DESCRIPTION: &str = "A generated report"; +pub(crate) const COMPLETION_VALUES: [&str; 2] = ["alpha", "beta"]; +pub(crate) const RESOURCE_URI: &str = "report://42/summary"; +pub(crate) const RESOURCE_DESCRIPTION: &str = "A stored report"; +pub(crate) const RESOURCE_TEXT: &str = "quarterly numbers"; const GATEWAY_PORT_READY_TIMEOUT: Duration = Duration::from_secs(10); const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); #[derive(Clone)] pub(crate) struct BackendObservation { - pub(crate) tool_name: String, + pub(crate) name: String, pub(crate) args: Option>, } #[derive(Clone, Default)] pub(crate) struct BackendState { pub(crate) calls: Arc>>, + pub(crate) prompts: Arc>>, + pub(crate) prompt_lists: Arc>, + pub(crate) template_lists: Arc>, + pub(crate) completions: Arc>>, + pub(crate) reads: Arc>>, + pub(crate) resource_lists: Arc>, + pub(crate) subscriptions: Arc>>, + pub(crate) unsubscriptions: Arc>>, pub(crate) cancellations: Arc>>, } @@ -59,8 +77,118 @@ impl ServerHandler for TestBackend { _request: InitializeRequestParams, _cx: RequestContext, ) -> Result { - Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("test-backend", "0.1.0"))) + Ok(InitializeResult::new( + ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .enable_resources() + .enable_resources_subscribe() + .build(), + ) + .with_server_info(Implementation::new("test-backend", "0.1.0"))) + } + + async fn list_prompts( + &self, + _request: Option, + _cx: RequestContext, + ) -> Result { + *self.state.prompt_lists.lock().expect("backend prompt lists lock poisoned") += 1; + Ok(ListPromptsResult::with_all_items(vec![Prompt::new("review", Some(PROMPT_DESCRIPTION), None)])) + } + + async fn subscribe( + &self, + request: SubscribeRequestParams, + _cx: RequestContext, + ) -> Result<(), ErrorData> { + self.state.subscriptions.lock().expect("backend subscriptions lock poisoned").push(request.uri.clone()); + Ok(()) + } + + async fn unsubscribe( + &self, + request: UnsubscribeRequestParams, + _cx: RequestContext, + ) -> Result<(), ErrorData> { + self.state.unsubscriptions.lock().expect("backend unsubscriptions lock poisoned").push(request.uri.clone()); + Ok(()) + } + + async fn list_resources( + &self, + _request: Option, + _cx: RequestContext, + ) -> Result { + *self.state.resource_lists.lock().expect("backend resource lists lock poisoned") += 1; + Ok(ListResourcesResult::with_all_items(vec![ + Resource::new(RESOURCE_URI, "report").with_description(RESOURCE_DESCRIPTION), + ])) + } + + async fn read_resource( + &self, + request: ReadResourceRequestParams, + _cx: RequestContext, + ) -> Result { + self.state.reads.lock().expect("backend reads lock poisoned").push(request.uri.clone()); + Ok(ReadResourceResult::new(vec![ResourceContents::text(RESOURCE_TEXT, request.uri)]).into()) + } + + async fn complete( + &self, + request: CompleteRequestParams, + _cx: RequestContext, + ) -> Result { + let identifier = match &request.r#ref { + Reference::Prompt(prompt) => prompt.name.clone(), + Reference::Resource(resource) => resource.uri.clone(), + _ => String::new(), + }; + self.state.completions.lock().expect("backend completions lock poisoned").push(identifier); + let completion = CompletionInfo::new(COMPLETION_VALUES.iter().map(|value| (*value).to_owned()).collect()) + .expect("completion values are within the MCP limit"); + Ok(CompleteResult::new(completion)) + } + + async fn list_resource_templates( + &self, + _request: Option, + _cx: RequestContext, + ) -> Result { + *self.state.template_lists.lock().expect("backend template lists lock poisoned") += 1; + Ok(ListResourceTemplatesResult::with_all_items(vec![ + ResourceTemplate::new(TEMPLATE_URI, "report").with_description(TEMPLATE_DESCRIPTION), + ])) + } + + /// Renders `review` from its `topic` argument so tests can prove a pre-hook argument rewrite + /// reached the backend, and prove a post-hook rewrite changed what the client received. + async fn get_prompt( + &self, + request: GetPromptRequestParams, + _cx: RequestContext, + ) -> Result { + self.state + .prompts + .lock() + .expect("backend prompts lock poisoned") + .push(BackendObservation { name: request.name.clone(), args: request.arguments.clone() }); + + if request.name != "review" { + return Err(ErrorData { + code: ErrorCode::METHOD_NOT_FOUND, + message: format!("unknown prompt {}", request.name).into(), + data: None, + }); + } + let topic = request + .arguments + .as_ref() + .and_then(|arguments| arguments.get("topic")) + .and_then(Value::as_str) + .unwrap_or("nothing"); + Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, format!("review of {topic}"))]).into()) } async fn call_tool( @@ -72,7 +200,7 @@ impl ServerHandler for TestBackend { .calls .lock() .expect("backend calls lock poisoned") - .push(BackendObservation { tool_name: request.name.to_string(), args: request.arguments.clone() }); + .push(BackendObservation { name: request.name.to_string(), args: request.arguments.clone() }); let result: Result = match request.name.as_ref() { "sum" => { diff --git a/crates/contextforge-gateway-rs-lib/tests/support/runtime.rs b/crates/contextforge-gateway-rs-lib/tests/support/runtime.rs index 5a6f463..6ec099b 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/runtime.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/runtime.rs @@ -4,7 +4,44 @@ use contextforge_gateway_rs_cpex::CpexRuntimeRegistry; use cpex::cpex_core::config::CpexConfig; use serde_json::json; -use super::{TestPlugin, TestPluginFactory}; +use super::{ + PromptTestPlugin, PromptTestPluginFactory, ResourceTestPlugin, ResourceTestPluginFactory, TestPlugin, + TestPluginFactory, +}; + +pub(crate) async fn runtime_with_resource_plugin(plugin: Arc) -> Arc { + let mut runtime = CpexRuntimeRegistry::default(); + runtime + .register_factory("resource-test", Box::new(ResourceTestPluginFactory::from_plugin(&plugin))) + .expect("resource test factory registers"); + let config = serde_json::from_value(json!({ + "plugins": [{ + "name": plugin.config.name.clone(), + "kind": plugin.config.kind.clone(), + "hooks": plugin.config.hooks.clone(), + }] + })) + .expect("resource CPEX config parses"); + runtime.apply_config(Some(config)).await.expect("resource runtime config applies"); + Arc::new(runtime) +} + +pub(crate) async fn runtime_with_prompt_plugin(plugin: Arc) -> Arc { + let mut runtime = CpexRuntimeRegistry::default(); + runtime + .register_factory("prompt-test", Box::new(PromptTestPluginFactory::from_plugin(&plugin))) + .expect("prompt test factory registers"); + let config = serde_json::from_value(json!({ + "plugins": [{ + "name": plugin.config.name.clone(), + "kind": plugin.config.kind.clone(), + "hooks": plugin.config.hooks.clone(), + }] + })) + .expect("prompt CPEX config parses"); + runtime.apply_config(Some(config)).await.expect("prompt runtime config applies"); + Arc::new(runtime) +} pub(crate) async fn runtime_with_pre(plugin: Arc) -> Arc { runtime_with_plugins(&[plugin]).await diff --git a/docs/book/src/mcp-method-reference.md b/docs/book/src/mcp-method-reference.md index d378a85..fd114d0 100644 --- a/docs/book/src/mcp-method-reference.md +++ b/docs/book/src/mcp-method-reference.md @@ -43,10 +43,15 @@ single-backend pass-through, or a multi-backend prefix: | Method | Behavior | | --- | --- | | `call_tool` | Resolves an exact control-plane alias first, then falls back to the single-versus-multi rule. It runs the optional plugin hooks, tracks progress, forwards the backend-local tool name, and propagates downstream cancellation. | -| `read_resource` | Preserves a single-backend URI or strips a multi-backend prefix, then returns the selected backend's result. | -| `subscribe`, `unsubscribe` | Apply the same resource-URI routing, track the downstream subscription, and forward or stop forwarding matching resource-update notifications. | -| `get_prompt` | Preserves a single-backend prompt name or strips a multi-backend prefix, then returns the selected backend's result. | -| `complete` | Applies the same routing to the prompt name or resource URI in `ref` and returns the selected backend's completion result. | +| `read_resource` | Preserves a single-backend URI or strips a multi-backend prefix, then returns the selected backend's result. Runs the optional resource hooks, which may deny, rewrite the URI, or rewrite returned text contents. | +| `subscribe`, `unsubscribe` | Apply the same resource-URI routing, track the downstream subscription, and forward or stop forwarding matching resource-update notifications. The optional resource pre hook may deny or rewrite the URI; the rewritten URI is both forwarded and tracked. Neither has a post hook, because both return an empty acknowledgement. | +| `get_prompt` | Preserves a single-backend prompt name or strips a multi-backend prefix, then returns the selected backend's result. Runs the optional prompt hooks, which may deny, replace arguments, or rewrite message text. | +| `complete` | Applies the same routing to the prompt name or resource URI in `ref` and returns the selected backend's completion result. Runs the optional prompt or resource hooks depending on the reference kind; the pre hook is deny-only and the post hook may rewrite completion values. | + +Fan-out methods also run hooks: `list_prompts`, `list_resources`, and `list_resource_templates` +run their family's pre hook once before fan-out and the post hook once on the +merged result. See [Plugins And Policy](plugins-and-policy.md) for the full +hook contract. Routed failures are JSON-RPC errors: an identifier that cannot select a backend or an unavailable backend returns an internal error, and duplicate backend diff --git a/docs/book/src/plugins-and-policy.md b/docs/book/src/plugins-and-policy.md index 9ca71fe..fb148e4 100644 --- a/docs/book/src/plugins-and-policy.md +++ b/docs/book/src/plugins-and-policy.md @@ -24,37 +24,160 @@ The supported surface is deliberately narrow: ```text cmf.tool_pre_invoke cmf.tool_post_invoke +cmf.prompt_pre_fetch +cmf.prompt_post_fetch +cmf.resource_pre_fetch +cmf.resource_post_fetch ``` The gateway rejects route-based plugin selection, plugin directories, global -policies/defaults, non-tool hooks, and plugin conditions. Those features need -clear behavior for streaming, failures, timeouts, backpressure, context -propagation, and observability before they belong on the hot path. - -## Tool Call Behavior - -For `call_tool`, the pre hook runs after backend routing has selected the -backend and stripped the public prefix. The hook sees the backend name, routed -tool name, and arguments. It can: - -- leave arguments unchanged -- replace arguments -- deny the call - -After the upstream backend returns, the post hook can: - -- leave the result unchanged -- rewrite the result payload -- deny the response - -Hook state is carried across the upstream call so pre and post hooks can share -CPEX context for the same logical tool call. +policies/defaults, LLM hooks, and plugin conditions. Those features need clear +behavior for streaming, failures, timeouts, backpressure, context propagation, +and observability before they belong on the hot path. + +Each MCP path checks only its own hook pair, so a plugin config that declares +only prompt hooks leaves the tool hot path untouched, and vice versa. + +## Covered MCP Paths + +| MCP path | Hook family | Pre hook | Post hook | +| --- | --- | --- | --- | +| `tools/call` | tool | deny, replace arguments | deny, rewrite result | +| `prompts/get` | prompt | deny, replace arguments | deny, rewrite message text | +| `resources/read` | resource | deny, rewrite URI | deny, rewrite text contents | +| `resources/subscribe` | resource | deny, rewrite URI | none | +| `resources/unsubscribe` | resource | deny, rewrite URI | none | +| `completion/complete` | prompt or resource | deny | deny, rewrite values | +| `prompts/list` | prompt | deny | deny, read-only | +| `resources/list` | resource | deny | deny, read-only | +| `resources/templates/list` | resource | deny | deny, read-only | + +Every MCP method the gateway routes now has hook coverage. + +CPEX defines no completion hook, so `completion/complete` runs the hooks of +whatever is being completed: a prompt reference uses the prompt hooks, a +resource or template reference uses the resource hooks. + +## Hook Ordering + +For routed methods — `tools/call`, `prompts/get`, `resources/read`, +`resources/subscribe`, `resources/unsubscribe`, and `completion/complete` — the +pre hook runs **after** backend routing has selected the backend and stripped the +public prefix. The hook sees the backend-local identifier and the backend name +separately, so plugins never have to understand the gateway's namespace scheme. + +For fan-out methods — `prompts/list`, `resources/list`, and +`resources/templates/list` — the pre +hook runs **once before fan-out**, so a denial costs no backend traffic. The +post hook runs **once on the merged, namespaced page**, so plugins see exactly +what the client will receive. Mutating per backend before the merge would let a +plugin break the routing contract. + +These methods are paginated, so the post hook observes one page per request, not +the complete set. A plugin that needs to reason across the whole listing must +accumulate across calls using CPEX context; the pre hook receives the incoming +cursor so pages can be correlated. + +A denial becomes an MCP error and the backend is never called. Hook state is +carried across the upstream call so pre and post hooks share CPEX context for +the same logical request. + +## Payload Shapes + +| MCP path | Pre payload | Post payload | +| --- | --- | --- | +| `tools/call` | `ToolCall`, backend in `namespace` | `ToolResult` with the serialized MCP result | +| `prompts/get` | `PromptRequest`, backend in `server_id` | `PromptResult` with mapped messages | +| `prompts/list` | `PromptRequest`, **empty name**, no `server_id` | one `PromptRequest` per prompt, descriptions as text | +| `resources/templates/list` | `ResourceRef`, **empty URI** | one `ResourceRef` per template, descriptions as text | +| `resources/list` | `ResourceRef`, **empty URI** | one `ResourceRef` per resource, descriptions as text | +| `resources/read` | `Resource`, backend in `annotations.backend` | one `Resource` per content item | +| `resources/subscribe`, `resources/unsubscribe` | `Resource`, backend in `annotations.backend` | — | +| `completion/complete` (prompt ref) | `PromptRequest`, argument merged into arguments | `PromptResult` plus one text part per value | +| `completion/complete` (resource ref) | `Resource`, argument in `annotations.completion_argument` | `Resource` plus one text part per value | + +CMF resource types have no backend field of their own, unlike `ToolCall` and +`PromptRequest`, so resource paths carry the backend in `Resource.annotations` +under the `backend` key. + +## Telling Payloads Apart + +A plugin that registers one handler for both hooks of a family needs to know +what it is looking at. Three rules cover it: + +- **Pre versus post.** Prompt and resource payloads use CMF role `User` on the + request side and `Assistant` on the response side. Tool payloads use + `Assistant` for the call and `Tool` for the result. Do not discriminate on + content part alone: a `prompts/list` post payload is built from + `PromptRequest` parts, the same part a pre payload uses. +- **Listing versus routed.** Listings carry an empty identifier — an empty + prompt name or an empty resource URI — and no backend. +- **Which MCP path.** The CPEX correlation id is prefixed by originating path: + + ```text + gateway-tool-call-N gateway-prompt-request-N + gateway-prompt-list-N gateway-resource-template-list-N + gateway-resource-subscription-N gateway-completion-N + gateway-resource-read-N gateway-resource-list-N + ``` + + This is how a prompt-reference completion is told apart from a `prompts/get`, + which otherwise produce identically shaped payloads. + +## Mutation Limits + +Mutation is supported only where the gateway can apply a change without +silently dropping data: + +- **Prompt messages** write back text content only. The message count must + match what the plugin was given. `description`, `_meta`, message roles, and + non-text blocks always come from the backend response, so a plugin that only + observes gets a byte-identical passthrough. +- **Removing exposed text fails closed.** A text block is always exposed as a + text part, so its absence in the returned payload means the plugin deleted it. + The gateway rejects that with `INVALID_PARAMS` rather than restoring the + backend's text, which would return content a redaction plugin stripped. + Explicit deletion semantics are not supported yet; rewrite to an empty string + instead. +- **Completion values** write back in full, but the value count must match, so + the backend's `total` and `has_more` stay accurate. +- **Listings are read-only.** MCP `Prompt` and `ResourceTemplate` are metadata — + `title`, `arguments`, `mime_type`, `icons`, `annotations`, `_meta` — with no + CMF equivalent to rebuild from. A modified listing payload is ignored rather + than partially applied. Changing which items are exposed is filtering, which + belongs to the control plane. +- **Resource contents** write back text only, with a matching content count. Blob + contents are exposed as a reference — URI and MIME type — but their bytes are + withheld and never modified: MCP carries blobs base64-encoded while CMF wants + raw bytes, and cloning arbitrary binary payloads through the plugin pipeline + on the hot path is a memory hazard. +- **Subscription URIs** may be rewritten by the pre hook. The mapping from the + client's URI to the rewritten upstream URI is recorded at subscribe time and + reused afterwards — it is never recomputed. `unsubscribe` still runs its pre + hook so a plugin can deny it, but addresses the backend with the stored URI, + and resource-update notifications are reported to the client under the URI it + subscribed to. Recomputing would misroute both whenever a plugin's rewrite is + stateful or non-deterministic. +- **Completion arguments are read-only.** The argument value is a partial user + keystroke, so rewriting it would return completions for text the user never + typed. + +A structural change the gateway cannot apply — a changed message count, text +aimed at a non-text message, a changed completion value count — is an +`INVALID_PARAMS` error rather than a partial write. ## Boundary Rules Plugin execution must not poison shared gateway state. A plugin denial becomes -an MCP error. Soft plugin errors are logged. Unsupported plugin configuration -fails validation before the runtime is accepted. +an MCP error, carrying the plugin's `proto_error_code` when it set one and +`INVALID_REQUEST` otherwise. Soft plugin errors are logged. Unsupported plugin +configuration fails validation before the runtime is accepted. + +An in-flight request is pinned to the runtime that ran its pre hook, so a +configuration reload mid-request cannot hand the post hook to a different plugin +set. Future hook expansion should define behavior for streaming/SSE, cancellation, -timeouts, backpressure, and telemetry before adding new hook points. +timeouts, backpressure, and telemetry before adding new hook points. One gap is +known and deliberate: resource-update notifications have no post-hook equivalent +to the tool path's streamed-event hook.