From be0983237e2c8a3d22b23851f3e0d5b44190f278 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 19 Sep 2026 02:00:53 -0400 Subject: [PATCH] fix(mcp): an abandoned tool call tells the server to stop `try_call_tool` wrapped `call_tool` in a bare `tokio::time::timeout`. On a timeout or a cancelled turn the future dropped and the server kept working: no `notifications/cancelled` ever left the client. `call_tool_cancel_aware` sends the request through `send_cancellable_request` and holds a guard armed with the request id. A settled call disarms it. A dropped future (cancelled turn) spawns the cancel from the guard's drop. A timeout sends the cancel inline, before returning: the HTTP timeout arm resets the transport next, and dropping the last service handle killed a spawned cancel before it left. The inline send waits 500 ms at most. rmcp 2.2.0 awaits its own timeout-cancel without a bound, and its streamable-HTTP worker awaits each POST in turn, so against a JSON-mode server the cancel queues behind the hung call: with rmcp's timeout path the call returned after 3.0 s against a 1 s budget. That server class cannot receive the cancel until rmcp 3.x; the budget elapses, a warn is logged, and AGENTS.md records the limitation. stdio and SSE-mode servers receive it. Tests: dropped call, timed-out call followed by a service drop, settled call; the HTTP outer-timeout test now asserts the return time. Ported from upstream (grok-build mcp/servers.rs call_tool_cancel_aware, 1.0.33); upstream leans on rmcp 3.2's timeout path. --- AGENTS.md | 8 + crates/codegen/kigi-mcp/src/servers.rs | 64 ++--- .../kigi-mcp/src/servers/cancel_aware.rs | 224 ++++++++++++++++++ 3 files changed, 267 insertions(+), 29 deletions(-) create mode 100644 crates/codegen/kigi-mcp/src/servers/cancel_aware.rs diff --git a/AGENTS.md b/AGENTS.md index a63a950..52b5406 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -30,6 +30,14 @@ that stops reading still blocks exit-path prints behind std's stderr lock builder re-checks `Path::exists("/data")` and can omit the devbox bind on a stat error that the profile classification treated as present. +Known limitation left by the 0.1.16 sync: `notifications/cancelled` for a +timed-out MCP call cannot reach a streamable-HTTP server that answers +`tools/call` with a plain JSON body. rmcp 2.2.0's HTTP worker awaits each +POST in turn, so the cancel queues behind the hung call; +`servers/cancel_aware.rs` waits 500 ms, warns, and the transport reset +then drops the connection. stdio and SSE-mode servers get the cancel. +The cure is rmcp 3.x (upstream runs 3.2), a major bump not yet taken. + ## Hard constraints - **Zero egress**: only `auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`, diff --git a/crates/codegen/kigi-mcp/src/servers.rs b/crates/codegen/kigi-mcp/src/servers.rs index 4e7ccb4..bd1577c 100644 --- a/crates/codegen/kigi-mcp/src/servers.rs +++ b/crates/codegen/kigi-mcp/src/servers.rs @@ -32,6 +32,9 @@ use rmcp::{ use crate::oauth_config::McpOAuthConfig; +mod cancel_aware; +use cancel_aware::call_tool_cancel_aware; + use kigi_tools::types::{ output::{MCPOutput, MCPOutputDetails, ToolOutput}, tool::{ToolKind, ToolNamespace}, @@ -1599,12 +1602,24 @@ impl McpErasedTool { let mut params = CallToolRequestParams::new(self.tool.name.clone()); params.arguments = raw.as_object().cloned(); - let result = - tokio::time::timeout(timeout_duration, mcp_service.call_tool(params.clone())).await; - - match result { - Ok(Ok(call_result)) => Ok(call_result), - Ok(Err(service_err)) + match call_tool_cancel_aware(&mcp_service, params.clone(), timeout_duration).await { + Ok(call_result) => Ok(call_result), + Err(ServiceError::Timeout { .. }) => { + *is_timeout = true; + // Reset for the next call but don't retry — a slow side-effecting tool must not run twice. + if client.is_http() && !*reconnect_attempted { + client.reset_transport().await; + *reconnect_attempted = true; + } + Err(kigi_tool_runtime::ToolError::custom( + "process_manager", + format!( + "MCP tool '{}' timed out after {} seconds", + self.tool.name, tool_timeout + ), + )) + } + Err(service_err) if should_recover_service_error( &service_err, client.is_http(), @@ -1623,25 +1638,10 @@ impl McpErasedTool { ) .await } - Ok(Err(e)) => Err(kigi_tool_runtime::ToolError::custom( + Err(e) => Err(kigi_tool_runtime::ToolError::custom( "process_manager", e.to_string(), )), - Err(_) => { - *is_timeout = true; - // Reset for the next call but don't retry — a slow side-effecting tool must not run twice. - if client.is_http() && !*reconnect_attempted { - client.reset_transport().await; - *reconnect_attempted = true; - } - Err(kigi_tool_runtime::ToolError::custom( - "process_manager", - format!( - "MCP tool '{}' timed out after {} seconds", - self.tool.name, tool_timeout - ), - )) - } } } @@ -1692,13 +1692,9 @@ impl McpErasedTool { )); } }; - match tokio::time::timeout(timeout_duration, mcp_service.call_tool(params)).await { - Ok(Ok(call_result)) => Ok(call_result), - Ok(Err(retry_err)) => Err(kigi_tool_runtime::ToolError::custom( - "process_manager", - retry_err.to_string(), - )), - Err(_) => { + match call_tool_cancel_aware(&mcp_service, params, timeout_duration).await { + Ok(call_result) => Ok(call_result), + Err(ServiceError::Timeout { .. }) => { *is_timeout = true; Err(kigi_tool_runtime::ToolError::custom( "process_manager", @@ -1708,6 +1704,10 @@ impl McpErasedTool { ), )) } + Err(retry_err) => Err(kigi_tool_runtime::ToolError::custom( + "process_manager", + retry_err.to_string(), + )), } } } @@ -6223,11 +6223,17 @@ mod tests { let mut reconnect = false; let mut is_timeout = false; let raw = serde_json::json!({}); + let started = std::time::Instant::now(); let err = tool .try_call_tool(&client, &raw, &mut reconnect, &mut is_timeout, &ew) .await .expect_err("call must time out"); + assert!( + started.elapsed() < std::time::Duration::from_millis(2500), + "cancel must not wait out the 3s hang: {:?}", + started.elapsed() + ); assert!(err.to_string().contains("timed out"), "got: {err}"); assert!(is_timeout, "is_timeout must be set"); assert!(reconnect, "timeout arm flags the reconnect after resetting"); diff --git a/crates/codegen/kigi-mcp/src/servers/cancel_aware.rs b/crates/codegen/kigi-mcp/src/servers/cancel_aware.rs new file mode 100644 index 0000000..bad22da --- /dev/null +++ b/crates/codegen/kigi-mcp/src/servers/cancel_aware.rs @@ -0,0 +1,224 @@ +//! `tools/call` that tells the server when the host stops waiting. + +use rmcp::model::{ + CallToolRequest, CallToolRequestParams, CallToolResult, CancelledNotificationParam, + ClientRequest, RequestId, ServerResult, +}; +use rmcp::service::{Peer, PeerRequestOptions}; +use rmcp::{RoleClient, ServiceError}; + +use super::McpService; + +/// rmcp 2.2.0's HTTP worker queues the cancel behind a hung POST, so the wait is bounded. +const CANCEL_DELIVERY_BUDGET: std::time::Duration = std::time::Duration::from_millis(500); + +/// A timeout cancels inline, before the caller can reset the transport; a drop cancels from a task. +pub(super) async fn call_tool_cancel_aware( + service: &McpService, + params: CallToolRequestParams, + timeout: std::time::Duration, +) -> Result { + let handle = service + .peer() + .send_cancellable_request( + ClientRequest::CallToolRequest(CallToolRequest::new(params)), + PeerRequestOptions::no_options(), + ) + .await?; + let mut guard = CancelOnDrop { + peer: service.peer().clone(), + request_id: Some(handle.id.clone()), + }; + let Ok(response) = tokio::time::timeout(timeout, handle.await_response()).await else { + guard.cancel_within_budget().await; + return Err(ServiceError::Timeout { timeout }); + }; + guard.request_id = None; + match response? { + ServerResult::CallToolResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } +} + +/// Armed while the request is in flight; disarmed once it settles. +struct CancelOnDrop { + peer: Peer, + request_id: Option, +} + +impl CancelOnDrop { + async fn cancel_within_budget(&mut self) { + let Some(request_id) = self.request_id.clone() else { + return; + }; + let delivery = notify_cancelled(self.peer.clone(), request_id.clone()); + if tokio::time::timeout(CANCEL_DELIVERY_BUDGET, delivery) + .await + .is_err() + { + tracing::warn!( + ?request_id, + "notifications/cancelled still queued after budget" + ); + } + self.request_id = None; + } +} + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + let Some(request_id) = self.request_id.take() else { + return; + }; + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + tracing::warn!( + ?request_id, + "MCP call dropped outside a runtime; cancel not sent" + ); + return; + }; + runtime.spawn(notify_cancelled(self.peer.clone(), request_id)); + } +} + +async fn notify_cancelled(peer: Peer, request_id: RequestId) { + let params = CancelledNotificationParam::new( + Some(request_id.clone()), + Some("client cancelled".to_owned()), + ); + if let Err(e) = peer.notify_cancelled(params).await { + tracing::warn!(?request_id, error = %e, "notifications/cancelled not delivered"); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use rmcp::ServiceExt; + use serde_json::Value; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + + use super::super::{KigiClientHandler, McpClient}; + use super::*; + + type Received = Arc>>; + + /// Duplex MCP server: answers `initialize`, records all, answers `tools/call` only when `reply`. + async fn recording_service(reply: bool) -> (McpService, Received) { + let received = Received::default(); + let (client_read, mut server_write) = tokio::io::duplex(64 * 1024); + let (server_read, client_write) = tokio::io::duplex(64 * 1024); + let sink = received.clone(); + tokio::spawn(async move { + let mut lines = BufReader::new(server_read).lines(); + while let Ok(Some(line)) = lines.next_line().await { + let msg: Value = serde_json::from_str(&line).expect("client sends JSON"); + sink.lock().push(msg.clone()); + let result = match msg["method"].as_str() { + Some("initialize") => serde_json::json!({ + "protocolVersion": msg["params"]["protocolVersion"], + "capabilities": { "tools": {} }, + "serverInfo": { "name": "recording", "version": "0.0.0" }, + }), + Some("tools/call") if reply => serde_json::json!({ + "content": [{ "type": "text", "text": "done" }], + "isError": false, + }), + _ => continue, + }; + let response = + serde_json::json!({ "jsonrpc": "2.0", "id": msg["id"], "result": result }); + let encoded = format!("{response}\n"); + server_write + .write_all(encoded.as_bytes()) + .await + .expect("client reads replies"); + } + }); + let handler = KigiClientHandler { + info: McpClient::make_client_info("recording"), + server_name: "recording".to_string(), + notify_tx: Arc::default(), + }; + let transport = rmcp::transport::async_rw::AsyncRwTransport::::new( + client_read, + client_write, + ); + let service = Arc::new(handler.serve(transport).await.expect("handshake")); + (service, received) + } + + fn values_at(received: &Received, method: &str, pointer: &str) -> Vec { + received + .lock() + .iter() + .filter(|m| m["method"] == method) + .map(|m| m.pointer(pointer).cloned().unwrap_or(Value::Null)) + .collect() + } + + fn call_ids(received: &Received) -> Vec { + values_at(received, "tools/call", "/id") + } + + fn cancelled_ids(received: &Received) -> Vec { + values_at(received, "notifications/cancelled", "/params/requestId") + } + + #[tokio::test] + async fn dropped_call_cancels_its_request() { + let (service, received) = recording_service(false).await; + let mut call = Box::pin(call_tool_cancel_aware( + &service, + CallToolRequestParams::new("slow"), + Duration::from_secs(30), + )); + let pending = tokio::time::timeout(Duration::from_millis(100), &mut call).await; + assert!(pending.is_err(), "server never replies"); + drop(call); + tokio::time::sleep(Duration::from_millis(100)).await; + + let ids = call_ids(&received); + assert_eq!(ids.len(), 1, "{ids:?}"); + assert_eq!(cancelled_ids(&received), ids); + } + + #[tokio::test] + async fn timed_out_call_cancels_its_request_once() { + let (service, received) = recording_service(false).await; + let round = call_tool_cancel_aware( + &service, + CallToolRequestParams::new("slow"), + Duration::from_millis(50), + ) + .await; + assert!( + matches!(round, Err(ServiceError::Timeout { .. })), + "{round:?}" + ); + // The HTTP timeout arm resets the transport, which drops the last service handle. + drop(service); + tokio::time::sleep(Duration::from_millis(100)).await; + + let ids = call_ids(&received); + assert_eq!(ids.len(), 1, "{ids:?}"); + assert_eq!(cancelled_ids(&received), ids); + } + + #[tokio::test] + async fn completed_call_sends_no_cancellation() { + let (service, received) = recording_service(true).await; + let result = call_tool_cancel_aware( + &service, + CallToolRequestParams::new("fast"), + Duration::from_secs(5), + ) + .await + .expect("server replies"); + assert_eq!(result.is_error, Some(false)); + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(cancelled_ids(&received), Vec::::new()); + } +}