lattice_ai/mcp/
dispatch.rs

1//! MCP method dispatch: decode an incoming JSON-RPC frame, route it to its
2//! handler, and produce the outgoing frame(s).
3//!
4//! Pure + panic-free: every path returns a value; malformed input yields a
5//! JSON-RPC parse-error response. This is the unit the connection loop
6//! (`server::serve_connection`) drives; keeping it pure makes the protocol
7//! contract testable without a socket.
8//!
9//! I1 handles `initialize` / `tools/list` / `prompts/list` and emits
10//! `notifications/tools/list_changed` after the client's `initialized`
11//! notification. `tools/call` is stubbed (reads I2, writes I3, diff I4).
12
13use serde_json::{Value, json};
14
15use lattice_protocol::jsonrpc::{
16    Message, Notification, Request, RequestId, Response, ResponseError, error_codes,
17};
18
19use crate::mcp::diff;
20use crate::mcp::protocol;
21use crate::mcp::reads;
22use crate::mcp::writes;
23use lattice_diff::ProgrammaticDiffBus;
24
25/// The state the dispatcher needs to answer `tools/call`: the read tools'
26/// [`ReadContext`](crate::mcp::reads::ReadContext) + the write bus. Built once at
27/// server spawn and shared (behind an `Arc`) across connections. The I1
28/// methods (`initialize` / `tools/list` / `prompts/list`) ignore it.
29#[derive(Clone)]
30pub struct DispatchContext {
31    /// D-fix.6: the originating connection id — assigned per WS connection by
32    /// the server (`serve_connection` clones the shared context and stamps
33    /// this). Carried into `openDiff` (tags the diff's `origin_session`) and
34    /// the close tools (scopes teardown to THIS connection's diffs), so one
35    /// agent session can never tear down another's diffs. `0` for the shared
36    /// boot context + dispatch tests (no real connection).
37    pub conn_id: u64,
38    /// Read-tool services (cache + generic buffer-store / diagnostics +
39    /// workspace config). BC.3b / AG-2b: the write bus now rides inside
40    /// `reads.editor` (a `lattice_agent::EditorAccess`) — write tools call
41    /// `ctx.reads.editor.open_file(...)` etc. directly, no separate bus field.
42    pub reads: reads::ReadContext,
43    /// I4 `openDiff` host-drained bus ([`ProgrammaticDiffBus`]). `None` until
44    /// boot wires it (and in dispatch tests) — `openDiff` then returns a
45    /// graceful `isError: true`. Separate from `writes` because `openDiff` is
46    /// blocking + carries lattice-diff types that can't cross the `Effect`
47    /// boundary (so it can't ride the I3 handler bus).
48    pub diff: Option<ProgrammaticDiffBus>,
49    /// D-fix.6 follow-up: the shared pending-review tracker. `openDiff` brackets
50    /// its blocking `await` with `review.begin()`, so the modeline shows a
51    /// `◆ review` badge while the agent is blocked on the user — derived
52    /// entirely from claude-code's own openDiff lifecycle (no host signal).
53    pub review: crate::mcp::status::ReviewHandle,
54}
55
56/// A frame the server should send back to the agent in response to an
57/// incoming frame.
58#[derive(Debug)]
59pub enum Outgoing {
60    /// A reply to a request (carries the request's id).
61    Response(Response),
62    /// A server-initiated notification (no id, no reply expected).
63    Notification(Notification),
64}
65
66/// Decode + route one incoming frame, returning the frames to send back.
67///
68/// Never panics. Malformed bytes yield a single parse-error response with
69/// a null id (the id can't be recovered from unparseable input).
70/// Client→server responses are ignored (lattice is the server; it issues
71/// no requests in I1). A notification may produce zero or more follow-ups
72/// (e.g. `initialized` → `tools/list_changed`).
73pub async fn dispatch_frame(bytes: &[u8], ctx: &DispatchContext) -> Vec<Outgoing> {
74    match Message::from_json(bytes) {
75        Ok(Message::Request(req)) => vec![Outgoing::Response(handle_request(&req, ctx).await)],
76        Ok(Message::Notification(note)) => handle_notification(&note),
77        Ok(Message::Response(_)) => Vec::new(),
78        Err(e) => vec![Outgoing::Response(Response::err(
79            RequestId::Null,
80            ResponseError {
81                code: error_codes::PARSE_ERROR,
82                message: format!("parse error: {e}"),
83                data: None,
84            },
85        ))],
86    }
87}
88
89/// True if `frame` is a `tools/call` request for the blocking `openDiff`
90/// tool — the only tool whose handler ([`crate::mcp::diff::open_diff`]) awaits an
91/// unbounded, user-paced verdict with no timeout. The connection dispatches
92/// these on their OWN task so the per-connection read loop is never blocked by
93/// a pending review: it keeps polling the socket + shutdown signal (so a
94/// dropped connection or `:claude-code-stop` is observed promptly and the
95/// pending review can be rejected), and other tool calls still flow. Every
96/// non-blocking tool stays inline + ordered. Malformed / non-matching frames
97/// return `false` (dispatched inline, where the parse error is reported).
98pub fn is_blocking_tool_call(frame: &[u8]) -> bool {
99    let Ok(value) = serde_json::from_slice::<Value>(frame) else {
100        return false;
101    };
102    value.get("method").and_then(|m| m.as_str()) == Some("tools/call")
103        && value
104            .get("params")
105            .and_then(|p| p.get("name"))
106            .and_then(|n| n.as_str())
107            == Some("openDiff")
108}
109
110/// Route one request to its MCP response.
111pub async fn handle_request(req: &Request, ctx: &DispatchContext) -> Response {
112    match req.method.as_str() {
113        "initialize" => Response::ok(req.id.clone(), protocol::initialize_result()),
114        "tools/list" => Response::ok(req.id.clone(), protocol::tools_list_result()),
115        "prompts/list" => Response::ok(req.id.clone(), protocol::prompts_list_result()),
116        "tools/call" => handle_tools_call(req, ctx).await,
117        other => Response::err(
118            req.id.clone(),
119            ResponseError {
120                code: error_codes::METHOD_NOT_FOUND,
121                message: format!("method not found: {other}"),
122                data: None,
123            },
124        ),
125    }
126}
127
128/// Route `tools/call` to the read tools (I2). The tool's structured result
129/// is wrapped in the MCP `CallToolResult` content envelope. An unknown tool
130/// name → `METHOD_NOT_FOUND`. Writes (I3) + `openDiff` (I4) extend this.
131async fn handle_tools_call(req: &Request, ctx: &DispatchContext) -> Response {
132    let params = req.params.as_ref();
133    let name = params
134        .and_then(|p| p.get("name"))
135        .and_then(|v| v.as_str())
136        .unwrap_or("");
137    let empty = json!({});
138    let arguments = params.and_then(|p| p.get("arguments")).unwrap_or(&empty);
139    let editor = &ctx.reads.editor;
140
141    // I4: `openDiff` is blocking and returns its OWN `CallToolResult` envelope
142    // (the FILE_SAVED / DIFF_REJECTED markers are the contract), so it bypasses
143    // the `tool_text_result` JSON-blob wrapper the reads/writes use.
144    if name == "openDiff" {
145        return Response::ok(
146            req.id.clone(),
147            // D-fix.6: tag the opened diff with THIS connection's id so a later
148            // session-scoped close tears down only this session's diffs.
149            // The `&ctx.review` guard marks the review pending while we block.
150            diff::open_diff(ctx.diff.as_ref(), arguments, ctx.conn_id, &ctx.review).await,
151        );
152    }
153
154    let result = match name {
155        // Reads (sync, off-thread cache + ArcSwap snapshots).
156        "getCurrentSelection" => Some(reads::get_current_selection(&ctx.reads)),
157        "getOpenEditors" => Some(reads::get_open_editors(&ctx.reads)),
158        "getWorkspaceFolders" => Some(reads::get_workspace_folders(&ctx.reads)),
159        "getDiagnostics" => Some(reads::get_diagnostics(&ctx.reads, arguments)),
160        "checkDocumentDirty" => Some(reads::check_document_dirty(&ctx.reads, arguments)),
161        // Writes (I3, async: send on the bus → wake the actor → await the
162        // oneshot the per-tick drain resolves).
163        "openFile" => Some(writes::open_file(editor, arguments).await),
164        "saveDocument" => Some(writes::save_document(editor, arguments).await),
165        // D-fix.6: both close paths are scoped to THIS connection (`ctx.conn_id`)
166        // — they tear down the diff session(s) this agent session opened,
167        // regardless of how/where the diff is displayed, never another session's.
168        "close_tab" => Some(writes::close_tab(editor, arguments, ctx.conn_id).await),
169        "closeAllDiffTabs" => Some(writes::close_all_diff_tabs(editor, ctx.conn_id).await),
170        _ => None,
171    };
172
173    match result {
174        Some(data) => Response::ok(req.id.clone(), tool_text_result(&data)),
175        None => Response::err(
176            req.id.clone(),
177            ResponseError {
178                code: error_codes::METHOD_NOT_FOUND,
179                message: format!("unknown tool: {name}"),
180                data: None,
181            },
182        ),
183    }
184}
185
186/// Wrap a read-tool result in the MCP `CallToolResult` envelope: structured
187/// data is serialized into a single `text` content block.
188fn tool_text_result(data: &Value) -> Value {
189    json!({
190        "content": [{ "type": "text", "text": serde_json::to_string(data).unwrap_or_default() }],
191        "isError": false,
192    })
193}
194
195/// Route one notification, returning any server-initiated follow-ups.
196fn handle_notification(note: &Notification) -> Vec<Outgoing> {
197    match note.method.as_str() {
198        // Post-init: advertise the tool list once the client signals ready.
199        "notifications/initialized" => vec![Outgoing::Notification(Notification::new(
200            "notifications/tools/list_changed",
201            None,
202        ))],
203        _ => Vec::new(),
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    #![allow(clippy::unwrap_used, clippy::panic)]
210    use super::*;
211    use serde_json::json;
212
213    fn req(id: i64, method: &str) -> Request {
214        Request::new(RequestId::from_u64(id as u64), method, None)
215    }
216
217    /// An empty dispatch context: empty read cache, no buffer-store /
218    /// diagnostics services, one workspace folder. The I1 methods ignore
219    /// it; the read tools degrade to empty results.
220    fn test_ctx() -> DispatchContext {
221        DispatchContext {
222            conn_id: 0,
223            reads: crate::mcp::reads::ReadContext {
224                editor: lattice_agent::EditorAccess::new(
225                    std::sync::Arc::new(std::sync::Mutex::new(
226                        lattice_agent::EditorStateCache::default(),
227                    )),
228                    None,
229                    vec!["/work".to_string()],
230                    None,
231                ),
232                diagnostics: None,
233            },
234            diff: None,
235            review: crate::mcp::status::ReviewState::new(std::sync::Arc::new(
236                tokio::sync::Notify::new(),
237            )),
238        }
239    }
240
241    #[tokio::test]
242    async fn initialize_handshake_returns_protocol_version_and_capabilities() {
243        let r = handle_request(&req(1, "initialize"), &test_ctx()).await;
244        assert_eq!(r.id, RequestId::Number(1));
245        let result = r.result.expect("ok result");
246        assert_eq!(result["protocolVersion"], protocol::MCP_PROTOCOL_VERSION);
247        assert_eq!(result["capabilities"]["tools"]["listChanged"], json!(true));
248        assert_eq!(result["serverInfo"]["name"], protocol::SERVER_NAME);
249        assert!(r.error.is_none());
250    }
251
252    #[tokio::test]
253    async fn tools_list_enumerates_the_full_catalog() {
254        let r = handle_request(&req(2, "tools/list"), &test_ctx()).await;
255        let result = r.result.expect("ok result");
256        let tools = result["tools"].as_array().expect("tools array");
257        let names: Vec<&str> = tools.iter().filter_map(|t| t["name"].as_str()).collect();
258        for expected in [
259            "getCurrentSelection",
260            "getOpenEditors",
261            "getWorkspaceFolders",
262            "getDiagnostics",
263            "checkDocumentDirty",
264            "openFile",
265            "saveDocument",
266            "close_tab",
267            "closeAllDiffTabs",
268            "openDiff",
269        ] {
270            assert!(names.contains(&expected), "missing tool {expected}");
271        }
272    }
273
274    #[tokio::test]
275    async fn prompts_list_is_empty() {
276        let r = handle_request(&req(3, "prompts/list"), &test_ctx()).await;
277        let result = r.result.expect("ok result");
278        assert_eq!(result["prompts"].as_array().map(|a| a.len()), Some(0));
279    }
280
281    #[tokio::test]
282    async fn unknown_method_is_method_not_found() {
283        let r = handle_request(&req(4, "no/such/method"), &test_ctx()).await;
284        assert!(r.result.is_none());
285        let e = r.error.expect("error");
286        assert_eq!(e.code, error_codes::METHOD_NOT_FOUND);
287    }
288
289    #[tokio::test]
290    async fn tools_call_unknown_tool_is_method_not_found() {
291        // `req(5, "tools/call")` carries no `name` → unknown tool.
292        let r = handle_request(&req(5, "tools/call"), &test_ctx()).await;
293        let e = r.error.expect("unknown-tool error");
294        assert_eq!(e.code, error_codes::METHOD_NOT_FOUND);
295    }
296
297    #[tokio::test]
298    async fn tools_call_known_read_tool_returns_content_envelope() {
299        let call = Request::new(
300            RequestId::from_u64(6),
301            "tools/call",
302            Some(json!({ "name": "getWorkspaceFolders", "arguments": {} })),
303        );
304        let r = handle_request(&call, &test_ctx()).await;
305        let result = r.result.expect("ok result");
306        assert_eq!(result["isError"], json!(false));
307        let text = result["content"][0]["text"].as_str().expect("text block");
308        // The workspace result is JSON-stringified into the text block.
309        assert!(text.contains("folders"), "got {text}");
310    }
311
312    #[tokio::test]
313    async fn malformed_frame_yields_parse_error_and_does_not_panic() {
314        let out = dispatch_frame(b"{ this is not json", &test_ctx()).await;
315        assert_eq!(out.len(), 1);
316        match &out[0] {
317            Outgoing::Response(r) => {
318                let e = r.error.as_ref().expect("parse error");
319                assert_eq!(e.code, error_codes::PARSE_ERROR);
320                assert_eq!(r.id, RequestId::Null);
321            }
322            other => panic!("expected a response, got {other:?}"),
323        }
324    }
325
326    #[tokio::test]
327    async fn initialized_notification_triggers_tools_list_changed() {
328        let note = Notification::new("notifications/initialized", None);
329        let bytes = serde_json::to_vec(&note).unwrap();
330        let out = dispatch_frame(&bytes, &test_ctx()).await;
331        assert_eq!(out.len(), 1);
332        match &out[0] {
333            Outgoing::Notification(n) => {
334                assert_eq!(n.method, "notifications/tools/list_changed")
335            }
336            other => panic!("expected a notification, got {other:?}"),
337        }
338    }
339
340    #[tokio::test]
341    async fn dispatch_frame_routes_a_request_preserving_id() {
342        let bytes = serde_json::to_vec(&req(9, "tools/list")).unwrap();
343        let out = dispatch_frame(&bytes, &test_ctx()).await;
344        assert_eq!(out.len(), 1);
345        assert!(matches!(&out[0], Outgoing::Response(r) if r.id == RequestId::Number(9)));
346    }
347
348    #[test]
349    fn is_blocking_tool_call_matches_only_open_diff() {
350        let open_diff = json!({
351            "jsonrpc": "2.0", "id": 1, "method": "tools/call",
352            "params": { "name": "openDiff", "arguments": {} },
353        });
354        assert!(is_blocking_tool_call(open_diff.to_string().as_bytes()));
355
356        // Other tools/call names are NOT blocking (dispatched inline).
357        let save = json!({
358            "jsonrpc": "2.0", "id": 1, "method": "tools/call",
359            "params": { "name": "saveDocument", "arguments": {} },
360        });
361        assert!(!is_blocking_tool_call(save.to_string().as_bytes()));
362
363        // Non-tools/call methods are not blocking.
364        let list = json!({ "jsonrpc": "2.0", "id": 1, "method": "tools/list" });
365        assert!(!is_blocking_tool_call(list.to_string().as_bytes()));
366
367        // Malformed input never matches (it falls through to inline parse-error).
368        assert!(!is_blocking_tool_call(b"{ not json"));
369        assert!(!is_blocking_tool_call(b""));
370    }
371}
lattice_ai/mcp/dispatch.rs

lattice_ai/mcp/
dispatch.rs
