fix(shell-core/mcp): structurally enforce request-id preservation
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<String, String>) ->
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.
This commit is contained in:
parent
883bd70a65
commit
f400eb2701
|
|
@ -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<String, String>),
|
||||
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<String, String>) -> 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<String, String>) -> 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<String, String>) -> 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();
|
||||
|
|
|
|||
Loading…
Reference in a new issue