lex_api/
mcp.rs

1//! Model Context Protocol server (#171).
2//!
3//! Wraps the JSON API at `/v1/*` as MCP tools so any MCP-speaking
4//! host (Claude Code, Cursor, Codex, etc.) can invoke Lex actions
5//! natively. Stdio transport, JSON-RPC 2.0, hand-rolled — no SDK
6//! dependency, ~250 lines.
7//!
8//! Run with `lex serve --mcp` (sticks to the same `State` shape
9//! as the HTTP server, just speaks a different protocol on
10//! stdin/stdout instead of HTTP).
11//!
12//! Tools shipped in v1:
13//!
14//! * `lex_check` — POST /v1/check
15//! * `lex_publish` — POST /v1/publish
16//! * `lex_run` — POST /v1/run (effect policy passes through)
17//! * `lex_stage_get` — GET /v1/stage/<id>
18//! * `lex_stage_attestations` — GET /v1/stage/<id>/attestations
19//!
20//! Merge endpoints + trace/diff/replay land in v2; the v1 set
21//! covers the common agent loop (check → publish → run → read
22//! attestations).
23
24use serde::{Deserialize, Serialize};
25use serde_json::{json, Value};
26use std::io::{BufRead, BufReader, Write};
27use std::sync::Arc;
28
29use crate::handlers::State;
30
31/// JSON-RPC 2.0 envelope. We deserialize into `Value` first so a
32/// missing `id` (notification) doesn't error out — notifications
33/// are valid in MCP and shouldn't crash the loop.
34#[derive(Debug, Deserialize)]
35struct RpcRequest {
36    #[serde(default)]
37    id: Option<Value>,
38    method: String,
39    #[serde(default)]
40    params: Value,
41}
42
43#[derive(Debug, Serialize)]
44struct RpcResponse {
45    jsonrpc: &'static str,
46    id: Value,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    result: Option<Value>,
49    #[serde(skip_serializing_if = "Option::is_none")]
50    error: Option<RpcError>,
51}
52
53#[derive(Debug, Serialize)]
54struct RpcError {
55    code: i32,
56    message: String,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    data: Option<Value>,
59}
60
61/// Run the MCP server. Reads JSON-RPC requests from stdin (one
62/// per line), writes responses to stdout. Stops cleanly on EOF.
63///
64/// `eprintln!` is fine for diagnostics — MCP hosts read stdout
65/// for the protocol and surface stderr in their UI separately.
66pub fn serve_mcp(state: Arc<State>) -> std::io::Result<()> {
67    eprintln!("lex MCP server ready (stdio); v1 tools: lex_check, lex_publish, lex_run, lex_stage_get, lex_stage_attestations");
68    let stdin = std::io::stdin();
69    let reader = BufReader::new(stdin.lock());
70    let stdout = std::io::stdout();
71    let mut out = stdout.lock();
72
73    for line in reader.lines() {
74        let line = line?;
75        if line.trim().is_empty() { continue; }
76        match serde_json::from_str::<RpcRequest>(&line) {
77            Ok(req) => {
78                if let Some(resp) = dispatch(&state, req) {
79                    let body = serde_json::to_string(&resp).unwrap_or_else(|_| "{}".into());
80                    writeln!(out, "{body}")?;
81                    out.flush()?;
82                }
83                // None = notification (no `id`); MCP says don't reply.
84            }
85            Err(e) => {
86                // Parse error per JSON-RPC spec: id is null.
87                let resp = RpcResponse {
88                    jsonrpc: "2.0",
89                    id: Value::Null,
90                    result: None,
91                    error: Some(RpcError {
92                        code: -32700,
93                        message: format!("parse error: {e}"),
94                        data: None,
95                    }),
96                };
97                writeln!(out, "{}", serde_json::to_string(&resp).unwrap())?;
98                out.flush()?;
99            }
100        }
101    }
102    Ok(())
103}
104
105/// Returns `Some(response)` for requests that need a reply, `None`
106/// for notifications. The dispatch is small enough to be flat.
107fn dispatch(state: &State, req: RpcRequest) -> Option<RpcResponse> {
108    let id = req.id?;          // notification → drop
109    let method = req.method.as_str();
110    let result = match method {
111        "initialize" => Ok(json!({
112            "protocolVersion": "2024-11-05",
113            "capabilities": { "tools": { "listChanged": false } },
114            "serverInfo": {
115                "name": "lex",
116                "version": env!("CARGO_PKG_VERSION"),
117            }
118        })),
119        "tools/list" => Ok(json!({ "tools": tool_definitions() })),
120        "tools/call" => call_tool(state, &req.params),
121        // Hosts ping; respond with empty object.
122        "ping" => Ok(json!({})),
123        other => Err(RpcError {
124            code: -32601,
125            message: format!("method not found: {other}"),
126            data: None,
127        }),
128    };
129    Some(match result {
130        Ok(v) => RpcResponse { jsonrpc: "2.0", id, result: Some(v), error: None },
131        Err(e) => RpcResponse { jsonrpc: "2.0", id, result: None, error: Some(e) },
132    })
133}
134
135/// MCP `tools/list` response. Each entry is `{name, description,
136/// inputSchema}`. Schemas mirror the JSON request bodies the
137/// HTTP handlers already accept.
138fn tool_definitions() -> Value {
139    json!([
140        {
141            "name": "lex_check",
142            "description": "Type-check a Lex source string. Returns ok or a list of TypeErrors with structured detail.",
143            "inputSchema": {
144                "type": "object",
145                "properties": { "source": { "type": "string" } },
146                "required": ["source"]
147            }
148        },
149        {
150            "name": "lex_publish",
151            "description": "Publish a Lex source to the store. Type-check gate runs first; rejected sources don't advance the branch head. Returns the typed ops produced.",
152            "inputSchema": {
153                "type": "object",
154                "properties": {
155                    "source": { "type": "string" },
156                    "activate": { "type": "boolean", "default": false }
157                },
158                "required": ["source"]
159            }
160        },
161        {
162            "name": "lex_run",
163            "description": "Execute a Lex function under an effect policy. Pure programs run with no policy; effectful ones need allow_effects / allow_fs_read / allow_fs_write / allow_net_host grants.",
164            "inputSchema": {
165                "type": "object",
166                "properties": {
167                    "source": { "type": "string" },
168                    "fn": { "type": "string" },
169                    "args": { "type": "array", "items": {} },
170                    "policy": {
171                        "type": "object",
172                        "properties": {
173                            "allow_effects": { "type": "array", "items": { "type": "string" } },
174                            "allow_fs_read": { "type": "array", "items": { "type": "string" } },
175                            "allow_fs_write": { "type": "array", "items": { "type": "string" } },
176                            "budget": { "type": "integer" }
177                        }
178                    }
179                },
180                "required": ["source", "fn"]
181            }
182        },
183        {
184            "name": "lex_stage_get",
185            "description": "Fetch a stage's metadata + canonical AST + status by stage_id (lowercase-hex SHA-256).",
186            "inputSchema": {
187                "type": "object",
188                "properties": { "stage_id": { "type": "string" } },
189                "required": ["stage_id"]
190            }
191        },
192        {
193            "name": "lex_stage_attestations",
194            "description": "List every attestation persisted against a stage (TypeCheck / Spec / Examples / DiffBody / EffectAudit / SandboxRun). Newest-first.",
195            "inputSchema": {
196                "type": "object",
197                "properties": { "stage_id": { "type": "string" } },
198                "required": ["stage_id"]
199            }
200        }
201    ])
202}
203
204/// Dispatch a `tools/call` request. The MCP shape is `{name,
205/// arguments}`; we route on `name` and forward `arguments` as
206/// the JSON body the corresponding HTTP handler already accepts
207/// — keeps the two surfaces in lockstep without duplicating
208/// business logic.
209fn call_tool(state: &State, params: &Value) -> Result<Value, RpcError> {
210    let name = params.get("name").and_then(|v| v.as_str()).ok_or_else(|| RpcError {
211        code: -32602, message: "tools/call: missing `name`".into(), data: None,
212    })?;
213    let args = params.get("arguments").cloned().unwrap_or_else(|| json!({}));
214
215    // The HTTP handlers all take `&str` body and return a typed
216    // `Response<Cursor<Vec<u8>>>`. We extract the body bytes +
217    // status, then wrap into MCP's `content` shape: success →
218    // text content with the JSON body; error → isError + same
219    // content. Lets the host see the structured error envelope
220    // the JSON API already produces.
221    let body = serde_json::to_string(&args).unwrap_or_else(|_| "{}".into());
222    let (status, response_body): (u16, String) = match name {
223        "lex_check" => http_to_string(crate::handlers::check_handler(&body)),
224        "lex_publish" => http_to_string(crate::handlers::publish_handler(state, &body)),
225        "lex_run" => http_to_string(crate::handlers::run_handler(state, &body, false)),
226        "lex_stage_get" => {
227            let id = args.get("stage_id").and_then(|v| v.as_str()).ok_or_else(|| RpcError {
228                code: -32602, message: "lex_stage_get: missing stage_id".into(), data: None,
229            })?;
230            http_to_string(crate::handlers::stage_handler(state, id))
231        }
232        "lex_stage_attestations" => {
233            let id = args.get("stage_id").and_then(|v| v.as_str()).ok_or_else(|| RpcError {
234                code: -32602, message: "lex_stage_attestations: missing stage_id".into(), data: None,
235            })?;
236            http_to_string(crate::handlers::stage_attestations_handler(state, id))
237        }
238        other => return Err(RpcError {
239            code: -32602, message: format!("unknown tool: {other}"), data: None,
240        }),
241    };
242
243    let is_error = !(200..300).contains(&status);
244    Ok(json!({
245        "content": [{ "type": "text", "text": response_body }],
246        "isError": is_error,
247    }))
248}
249
250/// Drain an HTTP `Response` (the type all handlers return) into
251/// `(status, body_string)` so the MCP wrapper can repackage it.
252fn http_to_string(
253    resp: tiny_http::Response<std::io::Cursor<Vec<u8>>>,
254) -> (u16, String) {
255    // tiny_http's Response doesn't expose body bytes after construction
256    // without consuming the reader; we use our caller-side knowledge
257    // that handlers built the body from a Vec<u8>. The status code
258    // is on `status_code()`. For the body we re-render via the
259    // public iterator the response exposes.
260    let status = resp.status_code().0;
261    let mut buf = Vec::new();
262    let mut reader = resp.into_reader();
263    let _ = std::io::copy(&mut reader, &mut buf);
264    let body = String::from_utf8(buf).unwrap_or_default();
265    (status, body)
266}
lex_api/mcp.rs

lex_api/
mcp.rs
