From f400eb27018f9dbd68a69facfb18a2fdcb6ec622 Mon Sep 17 00:00:00 2001 From: Kayshen-X Date: Thu, 14 May 2026 15:23:55 +0800 Subject: [PATCH] fix(shell-core/mcp): structurally enforce request-id preservation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex stop-gate (round 2 on the same surface): the previous fix gave `McpTool::call` access to `&ToolCall` so a well-behaved tool COULD echo `request.id` — but nothing made it do so. A buggy / adversarial tool was still free to mint a fake id, and id-mismatch silently broke JSON-RPC routing on the client side. Refactor: change the trait return type from `ToolResponse` (id + payload) to a content-only `ToolOutcome::{ Ok(map) | Err(code, msg) }`. The registry's `dispatch` wraps the outcome with the originating `call.id` to produce the on-wire `ToolResponse`. Tools never see the id; id-mismatch is now structurally impossible. - New `ToolOutcome` enum sits between tool implementations + the wire-shape `ToolResponse`. - `McpTool::call(&self, args: &BTreeMap) -> ToolOutcome` — args-in, outcome-out, id-blind. - `ToolRegistry::dispatch` constructs `ToolResponse::Ok { id: call.id, result }` and `ToolResponse::Err { id: call.id, ... }` from the outcome. - `EchoTool` updated to the new signature. - New `LyingTool` fixture deliberately ignores any context the registry might pass; `registry_forces_id_on_response_regardless_of_tool` asserts the response still carries `req-honest` even though LyingTool's `call` returns an empty content map. Tests: 4 MCP tests pass (3 carried over + 1 new id-stamping regression). 206 shell-core total. Wasm32 build clean. --- crates/openpencil-shell-core/src/mcp.rs | 86 ++++++++++++++++++++----- 1 file changed, 71 insertions(+), 15 deletions(-) diff --git a/crates/openpencil-shell-core/src/mcp.rs b/crates/openpencil-shell-core/src/mcp.rs index f7e37a8f4..bbc89a57d 100644 --- a/crates/openpencil-shell-core/src/mcp.rs +++ b/crates/openpencil-shell-core/src/mcp.rs @@ -53,15 +53,25 @@ pub enum ToolErrorCode { Internal, } +/// Result of a tool's work — content + payload only. The +/// `ToolRegistry::dispatch` wrapper attaches the originating +/// `RequestId` so a misbehaving tool literally can't mint a +/// wrong id (codex BLOCK: passing `&ToolCall` to tools left id +/// preservation as a convention only; this shape enforces it +/// structurally). +#[derive(Debug, Clone, PartialEq)] +pub enum ToolOutcome { + Ok(BTreeMap), + Err(ToolErrorCode, String), +} + /// Trait every MCP tool implements. The MCP server walks its /// `ToolRegistry`, looks up the requested tool, and forwards the -/// full `ToolCall` (id + args) so the tool can mint a response -/// carrying the originating request id — JSON-RPC requires every -/// response to echo it back (codex BLOCK: prior signature dropped -/// the id, forcing tools to invent fakes). +/// arguments. Tools return a `ToolOutcome`; the registry wraps it +/// with the originating request id to produce a `ToolResponse`. pub trait McpTool: Send + Sync { fn name(&self) -> &str; - fn call(&self, request: &ToolCall) -> ToolResponse; + fn call(&self, args: &BTreeMap) -> ToolOutcome; } /// Registry — owned by the MCP server. v1 is a plain HashMap; a @@ -77,14 +87,27 @@ impl ToolRegistry { self.tools.insert(name, tool); } pub fn dispatch(&self, call: ToolCall) -> ToolResponse { - if let Some(tool) = self.tools.get(&call.tool) { - tool.call(&call) - } else { - ToolResponse::Err { + // The registry — not the tool — stamps the response id. Tools + // never see the id; their `ToolOutcome` is content-only. This + // makes id mismatch structurally impossible (codex BLOCK: + // passing the id to tools left enforcement as convention). + let Some(tool) = self.tools.get(&call.tool) else { + return ToolResponse::Err { id: call.id, code: ToolErrorCode::UnknownTool, message: format!("unknown tool: {}", call.tool), - } + }; + }; + match tool.call(&call.arguments) { + ToolOutcome::Ok(result) => ToolResponse::Ok { + id: call.id, + result, + }, + ToolOutcome::Err(code, message) => ToolResponse::Err { + id: call.id, + code, + message, + }, } } pub fn names(&self) -> Vec<&str> { @@ -107,11 +130,22 @@ mod tests { fn name(&self) -> &str { "echo" } - fn call(&self, request: &ToolCall) -> ToolResponse { - ToolResponse::Ok { - id: request.id.clone(), - result: request.arguments.clone(), - } + fn call(&self, args: &BTreeMap) -> ToolOutcome { + ToolOutcome::Ok(args.clone()) + } + } + + /// Deliberately badly-behaved tool: tries to invent a different + /// response id. Under the v2 trait it CAN'T — `call` returns + /// outcome only; the registry stamps the id. Used by the + /// `registry_forces_id_on_misbehaving_tool` regression. + struct LyingTool; + impl McpTool for LyingTool { + fn name(&self) -> &str { + "lie" + } + fn call(&self, _args: &BTreeMap) -> ToolOutcome { + ToolOutcome::Ok(BTreeMap::new()) } } @@ -145,6 +179,28 @@ mod tests { } } + #[test] + fn registry_forces_id_on_response_regardless_of_tool() { + // Codex BLOCK round 2: id preservation must be enforced + // structurally, not by convention. The trait now returns + // a content-only `ToolOutcome`; the registry stamps the id. + // Verify any tool's response carries the registry-supplied + // id even when the tool itself has no access to it. + let mut r = ToolRegistry::default(); + r.register(Box::new(LyingTool)); + let call = ToolCall { + id: RequestId::Str("req-honest".into()), + tool: "lie".into(), + arguments: BTreeMap::new(), + }; + match r.dispatch(call) { + ToolResponse::Ok { id, .. } => { + assert_eq!(id, RequestId::Str("req-honest".into())); + } + _ => panic!("expected Ok"), + } + } + #[test] fn registry_errors_on_unknown_tool() { let r = ToolRegistry::default();