Plugin hooks for Prompts/Resources MCP paths

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<dyn std::error::Error + Send + Sync + 'static>;
@@ -31,3 +33,88 @@ impl ToolPreCallResult {
         Self { arguments: ToolArgumentsUpdate::Unchanged, state: None }
     }
 }
+
+#[derive(Debug)]
+pub enum PromptArgumentsUpdate {
+    Unchanged,
+    Replace(Option<Map<String, Value>>),
+}
+
+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<RuntimeHookState>,
+}
+
+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<String, String>>,
+}
+
+/// 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<RuntimeHookState>,
+}
+
+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,
+        }
+    }
+}

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,
+};

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<String, serde_json::Value>>,
+    pre_result: &PipelineResult,
+) -> Result<PromptArgumentsUpdate, ErrorData> {
+    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<GetPromptResult, ErrorData> {
+    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<CompleteResult, ErrorData> {
+    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<ReadResourceResult, ErrorData> {
+    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<ResourceUriUpdate, ErrorData> {
+    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<T>(original: T, result: &PipelineResult) -> Result<T, ErrorData>
 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) {