Forward backend RMCP notifications

crates/contextforge-gateway-rs-cpex/src/handle.rs

@@ -13,7 +13,8 @@ use cpex_core::{
 };
 use rmcp::{
     ErrorData,
-    model::{CallToolRequestParams, CallToolResult, ErrorCode, ProgressNotificationParam},
+    model::{CallToolRequestParams, CallToolResult, ErrorCode},
+    serde::{Serialize, de::DeserializeOwned},
 };
 use tokio::task::JoinHandle;
 
@@ -271,17 +272,20 @@ impl GatewayPluginRuntimeHandle {
         }
     }
 
-    /// Runs the tool post hooks over a streamed progress notification. Returns
-    /// `None` when a plugin denies the notification.
-    pub async fn after_progress_notification(
+    /// 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<T>(
         &self,
         tool_name: &str,
-        progress: ProgressNotificationParam,
+        event: T,
         state: Option<RuntimeHookState>,
-    ) -> Result<Option<ProgressNotificationParam>, ErrorData> {
+    ) -> Result<Option<T>, ErrorData>
+    where
+        T: Serialize + DeserializeOwned,
+    {
         match state.and_then(|state| state.downcast::<RegistryToolCallState>().ok()) {
-            Some(state) => state.runtime.after_progress_notification(tool_name, progress, state.state.clone()).await,
-            None => Ok(Some(progress)),
+            Some(state) => state.runtime.after_tool_event(tool_name, event, state.state.clone()).await,
+            None => Ok(Some(event)),
         }
     }
 }
@@ -853,8 +857,7 @@ mod tests {
         let observations = plugin.observations();
         let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await;
 
-        let event =
-            runtime.handle().after_progress_notification("sum", progress_event(), None).await.expect("event passes");
+        let event = runtime.handle().after_stream_event("sum", progress_event(), None).await.expect("event passes");
 
         assert_eq!(Some("step 1/2"), event.expect("event is kept").message.as_deref());
         assert_eq!(0, observations.lock().expect("observations lock poisoned").post_calls);
@@ -868,11 +871,8 @@ mod tests {
         let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await;
 
         let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created");
-        let event = runtime
-            .handle()
-            .after_progress_notification("sum", progress_event(), pre.state)
-            .await
-            .expect("event passes");
+        let event =
+            runtime.handle().after_stream_event("sum", progress_event(), pre.state).await.expect("event passes");
 
         assert_eq!(Some("plugin:step 1/2"), event.expect("event is kept").message.as_deref());
         assert_eq!(1, observations.lock().expect("observations lock poisoned").post_calls);
@@ -886,7 +886,7 @@ mod tests {
         let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created");
         let event = runtime
             .handle()
-            .after_progress_notification("sum", progress_event(), pre.state)
+            .after_stream_event("sum", progress_event(), pre.state)
             .await
             .expect("deny drops the event");
 
@@ -902,7 +902,7 @@ mod tests {
         let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created");
         let error = runtime
             .handle()
-            .after_progress_notification("sum", progress_event(), pre.state)
+            .after_stream_event("sum", progress_event(), pre.state)
             .await
             .expect_err("invalid rewrite is rejected");
 

crates/contextforge-gateway-rs-cpex/src/pipeline.rs

@@ -2,7 +2,8 @@ use cpex_core::cmf::MessagePayload;
 use cpex_core::executor::PipelineResult;
 use rmcp::{
     ErrorData,
-    model::{CallToolResult, ErrorCode, ProgressNotificationParam},
+    model::{CallToolResult, ErrorCode},
+    serde::de::DeserializeOwned,
 };
 use tracing::warn;
 
@@ -45,23 +46,23 @@ pub(crate) fn effective_post_result(original: CallToolResult, result: &PipelineR
     }
 }
 
-pub(crate) fn effective_post_progress(
-    original: ProgressNotificationParam,
-    result: &PipelineResult,
-) -> Result<ProgressNotificationParam, ErrorData> {
+pub(crate) fn effective_post_json<T>(original: T, result: &PipelineResult) -> Result<T, ErrorData>
+where
+    T: DeserializeOwned,
+{
     let Some(payload) = modified_message_payload(result) else {
         return Ok(original);
     };
     let Some(content) = tool_result_content(payload) else {
         return Err(ErrorData {
             code: ErrorCode::INVALID_PARAMS,
-            message: "Plugin modified progress notification payload without a tool result".into(),
+            message: "Plugin modified stream event payload without a tool result".into(),
             data: None,
         });
     };
-    serde_json::from_value(content).map_err(|error| ErrorData {
+    serde_json::from_value::<T>(content).map_err(|error| ErrorData {
         code: ErrorCode::INVALID_PARAMS,
-        message: format!("Plugin modified progress notification payload with invalid JSON: {error}").into(),
+        message: format!("Plugin modified stream event payload with invalid JSON: {error}").into(),
         data: None,
     })
 }

crates/contextforge-gateway-rs-cpex/src/runtime.rs

@@ -14,7 +14,8 @@ use cpex_core::{
 };
 use rmcp::{
     ErrorData,
-    model::{CallToolRequestParams, CallToolResult, ProgressNotificationParam},
+    model::{CallToolRequestParams, CallToolResult},
+    serde::{Serialize, de::DeserializeOwned},
 };
 use tokio::sync::Mutex;
 
@@ -23,7 +24,7 @@ use crate::{
     error::GatewayPluginRuntimeError,
     hooks::{RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult},
     pipeline::{
-        effective_post_progress, effective_post_result, effective_pre_args, log_pipeline_errors, plugin_denied_error,
+        effective_post_json, effective_post_result, effective_pre_args, log_pipeline_errors, plugin_denied_error,
     },
 };
 
@@ -145,24 +146,6 @@ impl GatewayPluginRuntime {
         result
     }
 
-    /// Runs the tool post pipeline for one event of an in-flight call: the
-    /// payload is built from the call's `tool_call_id`, the call's context table
-    /// seeds the pipeline, and the resulting context table is carried back so
-    /// later events in the same call observe it.
-    async fn run_tool_post(
-        &self,
-        state: &SharedToolCallState,
-        build_payload: impl FnOnce(&str) -> MessagePayload,
-    ) -> PipelineResult {
-        let mut state = state.lock().await;
-        let payload = build_payload(&state.tool_call_id);
-        let post_result = self.invoke_tool_post(payload, Some(state.context_table.clone())).await;
-        if !post_result.is_denied() {
-            state.context_table = post_result.context_table.clone();
-        }
-        post_result
-    }
-
     pub(crate) async fn before_tool_call(
         &self,
         request: &CallToolRequestParams,
@@ -196,37 +179,53 @@ impl GatewayPluginRuntime {
             return Ok(response);
         }
 
-        let Some(state) = state.and_then(|state| state.downcast::<SharedToolCallState>().ok()) else {
-            return Ok(response);
-        };
+        let state = state.and_then(|state| state.downcast::<SharedToolCallState>().ok());
+        let Some(state) = state else { return Ok(response) };
 
-        let post_result = self.run_tool_post(&state, |id| tool_result_payload(tool_name, &response, id)).await;
+        let mut state = state.lock().await;
+        let post_result = self
+            .invoke_tool_post(
+                tool_result_payload(tool_name, &response, &state.tool_call_id),
+                Some(state.context_table.clone()),
+            )
+            .await;
         if post_result.is_denied() {
             return Err(plugin_denied_error(post_result));
         }
+
+        state.context_table = post_result.context_table.clone();
         Ok(effective_post_result(response, &post_result))
     }
 
-    pub(crate) async fn after_progress_notification(
+    pub(crate) async fn after_tool_event<T>(
         &self,
         tool_name: &str,
-        progress: ProgressNotificationParam,
+        event: T,
         state: Option<RuntimeHookState>,
-    ) -> Result<Option<ProgressNotificationParam>, ErrorData> {
+    ) -> Result<Option<T>, ErrorData>
+    where
+        T: Serialize + DeserializeOwned,
+    {
         if !self.has_post_hook {
-            return Ok(Some(progress));
+            return Ok(Some(event));
         }
 
-        let Some(state) = state.and_then(|state| state.downcast::<SharedToolCallState>().ok()) else {
-            return Ok(Some(progress));
-        };
+        let state = state.and_then(|state| state.downcast::<SharedToolCallState>().ok());
+        let Some(state) = state else { return Ok(Some(event)) };
 
-        let content = serde_json::to_value(&progress).unwrap_or(serde_json::Value::Null);
-        let post_result =
-            self.run_tool_post(&state, |id| tool_json_result_payload(tool_name, content, false, id)).await;
+        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),
+                Some(state.context_table.clone()),
+            )
+            .await;
         if post_result.is_denied() {
             return Ok(None);
         }
-        Ok(Some(effective_post_progress(progress, &post_result)?))
+
+        state.context_table = post_result.context_table.clone();
+        Ok(Some(effective_post_json(event, &post_result)?))
     }
 }

crates/contextforge-gateway-rs-lib/Cargo.toml

@@ -16,6 +16,7 @@ serde.workspace = true
 serde_json.workspace = true
 tracing.workspace = true
 tokio.workspace = true
+tokio-util = "0.7"
 axum.workspace = true
 axum-otel-metrics.workspace = true
 tower-http.workspace = true