firstpass_proxy/
mcp.rs

1//! Minimal MCP (Model Context Protocol) stdio server (SPEC §0.2/§7.4): lets an agent inspect and
2//! correct its own routing — read the audit traces, and submit feedback — over JSON-RPC 2.0.
3//!
4//! Hand-rolled over `serde_json` (newline-delimited JSON-RPC on stdin/stdout). The surface is three
5//! methods — `initialize`, `tools/list`, `tools/call` — so a dependency-free handler we can
6//! unit-test by value beats pulling in an async MCP framework. [`handle_rpc`] is the pure core;
7//! [`serve_stdio`] is the thin transport loop around it.
8
9use firstpass_core::{DeferredVerdict, Score, Verdict};
10use serde_json::{Value, json};
11
12use crate::store;
13
14/// MCP protocol revision this server implements.
15const PROTOCOL_VERSION: &str = "2024-11-05";
16
17/// The tools this server exposes, with their JSON-Schema input contracts (agent-first: an agent
18/// discovers them at runtime via `tools/list`).
19fn tool_schemas() -> Value {
20    json!([
21        {
22            "name": "list_traces",
23            "description": "List the most recent Firstpass audit traces (routing decisions + receipts).",
24            "inputSchema": {
25                "type": "object",
26                "properties": { "limit": { "type": "integer", "description": "Max traces to return (default 20)." } }
27            }
28        },
29        {
30            "name": "get_trace",
31            "description": "Fetch a single audit trace by id, with any deferred verdicts merged in.",
32            "inputSchema": {
33                "type": "object",
34                "properties": { "trace_id": { "type": "string" } },
35                "required": ["trace_id"]
36            }
37        },
38        {
39            "name": "submit_feedback",
40            "description": "Attach a downstream outcome (deferred verdict) to a past decision, closing the feedback loop.",
41            "inputSchema": {
42                "type": "object",
43                "properties": {
44                    "trace_id": { "type": "string" },
45                    "gate_id":  { "type": "string" },
46                    "verdict":  { "type": "string", "enum": ["pass", "fail", "abstain"] },
47                    "score":    { "type": "number" },
48                    "reporter": { "type": "string" }
49                },
50                "required": ["trace_id", "gate_id", "verdict", "reporter"]
51            }
52        }
53    ])
54}
55
56/// Handle one JSON-RPC message. Returns the response for a request, or `None` for a notification
57/// (a message with no `id`, which must not be answered). Tool calls touch the trace store at
58/// `db_path`; everything else is pure.
59#[must_use]
60pub fn handle_rpc(req: &Value, db_path: &str) -> Option<Value> {
61    let is_notification = req.get("id").is_none();
62    let id = req.get("id").cloned().unwrap_or(Value::Null);
63    let method = req
64        .get("method")
65        .and_then(Value::as_str)
66        .unwrap_or_default();
67
68    match method {
69        "initialize" => Some(ok(
70            id,
71            json!({
72                "protocolVersion": PROTOCOL_VERSION,
73                "capabilities": { "tools": {} },
74                "serverInfo": { "name": "firstpass", "version": env!("CARGO_PKG_VERSION") }
75            }),
76        )),
77        "tools/list" => Some(ok(id, json!({ "tools": tool_schemas() }))),
78        "tools/call" => Some(handle_tool_call(id, req.get("params"), db_path)),
79        "ping" => Some(ok(id, json!({}))),
80        // Notifications (e.g. `notifications/initialized`) and anything else without an id: silent.
81        _ if is_notification => None,
82        other => Some(err(id, -32601, &format!("method not found: {other}"))),
83    }
84}
85
86/// Dispatch a `tools/call`. Tool-level failures come back as an `isError` result (MCP convention),
87/// not a protocol error, so the agent can read the reason as content.
88fn handle_tool_call(id: Value, params: Option<&Value>, db_path: &str) -> Value {
89    let name = params
90        .and_then(|p| p.get("name"))
91        .and_then(Value::as_str)
92        .unwrap_or_default();
93    let args = params
94        .and_then(|p| p.get("arguments"))
95        .cloned()
96        .unwrap_or_else(|| json!({}));
97
98    let result = match name {
99        "list_traces" => tool_list_traces(&args, db_path),
100        "get_trace" => tool_get_trace(&args, db_path),
101        "submit_feedback" => tool_submit_feedback(&args, db_path),
102        other => Err(format!("unknown tool: {other}")),
103    };
104
105    match result {
106        Ok(text) => ok(id, json!({ "content": [{ "type": "text", "text": text }] })),
107        Err(msg) => ok(
108            id,
109            json!({ "content": [{ "type": "text", "text": msg }], "isError": true }),
110        ),
111    }
112}
113
114fn tool_list_traces(args: &Value, db_path: &str) -> Result<String, String> {
115    let limit = args
116        .get("limit")
117        .and_then(Value::as_u64)
118        .map_or(20, |n| n as usize);
119    // Mirror `firstpass trace`: an absent/unreadable store (no traffic yet) is "no traces", not an
120    // error — unlike get_trace/submit_feedback, which name a specific trace and should fail loudly.
121    let all = store::load_all_traces(db_path).unwrap_or_default();
122    let start = all.len().saturating_sub(limit);
123    let recent = &all[start..];
124    serde_json::to_string(recent).map_err(|e| format!("encode error: {e}"))
125}
126
127fn tool_get_trace(args: &Value, db_path: &str) -> Result<String, String> {
128    let trace_id = args
129        .get("trace_id")
130        .and_then(Value::as_str)
131        .ok_or("missing `trace_id`")?;
132    match store::load_trace_view(db_path, trace_id).map_err(|e| format!("store error: {e}"))? {
133        Some(trace) => serde_json::to_string(&trace).map_err(|e| format!("encode error: {e}")),
134        None => Err(format!("unknown trace_id {trace_id:?}")),
135    }
136}
137
138fn tool_submit_feedback(args: &Value, db_path: &str) -> Result<String, String> {
139    let trace_id = args
140        .get("trace_id")
141        .and_then(Value::as_str)
142        .ok_or("missing `trace_id`")?;
143    let gate_id = args
144        .get("gate_id")
145        .and_then(Value::as_str)
146        .ok_or("missing `gate_id`")?;
147    let reporter = args
148        .get("reporter")
149        .and_then(Value::as_str)
150        .ok_or("missing `reporter`")?;
151    let verdict = match args.get("verdict").and_then(Value::as_str) {
152        Some("pass") => Verdict::Pass,
153        Some("fail") => Verdict::Fail,
154        Some("abstain") => Verdict::Abstain,
155        other => return Err(format!("invalid verdict {other:?}")),
156    };
157    let score = match args.get("score") {
158        None | Some(Value::Null) => None,
159        Some(v) => {
160            let s = v.as_f64().ok_or("score must be a number")?;
161            Some(Score::new(s).map_err(|_| format!("score {s} out of range [0,1]"))?)
162        }
163    };
164
165    if !store::trace_exists(db_path, trace_id).map_err(|e| format!("store error: {e}"))? {
166        return Err(format!("unknown trace_id {trace_id:?}"));
167    }
168    let dv = DeferredVerdict {
169        gate_id: gate_id.to_owned(),
170        verdict,
171        score,
172        reported_at: jiff::Timestamp::now(),
173        reporter: reporter.to_owned(),
174    };
175    store::append_deferred(db_path, trace_id, &dv).map_err(|e| format!("store error: {e}"))?;
176    Ok(json!({ "status": "recorded", "trace_id": trace_id }).to_string())
177}
178
179fn ok(id: Value, result: Value) -> Value {
180    json!({ "jsonrpc": "2.0", "id": id, "result": result })
181}
182
183fn err(id: Value, code: i64, message: &str) -> Value {
184    json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
185}
186
187/// Serve MCP over stdio: read newline-delimited JSON-RPC from stdin, write responses to stdout.
188/// Blocks until stdin closes. Synchronous by design — one request in flight at a time.
189///
190/// # Errors
191/// Returns an [`std::io::Error`] if reading stdin or writing stdout fails.
192pub fn serve_stdio(db_path: &str) -> std::io::Result<()> {
193    use std::io::{BufRead, Write};
194    let stdin = std::io::stdin();
195    let mut stdout = std::io::stdout();
196    for line in stdin.lock().lines() {
197        let line = line?;
198        if line.trim().is_empty() {
199            continue;
200        }
201        let response = match serde_json::from_str::<Value>(&line) {
202            Ok(req) => handle_rpc(&req, db_path),
203            Err(_) => Some(err(Value::Null, -32700, "parse error")),
204        };
205        if let Some(resp) = response {
206            writeln!(stdout, "{resp}")?;
207            stdout.flush()?;
208        }
209    }
210    Ok(())
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    fn tmp_db() -> String {
218        std::env::temp_dir()
219            .join(format!("firstpass-mcp-{}.db", uuid::Uuid::now_v7()))
220            .to_string_lossy()
221            .into_owned()
222    }
223
224    #[test]
225    fn initialize_announces_protocol_and_server() {
226        let req = json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize" });
227        let resp = handle_rpc(&req, "unused.db").unwrap();
228        assert_eq!(resp["result"]["protocolVersion"], PROTOCOL_VERSION);
229        assert_eq!(resp["result"]["serverInfo"]["name"], "firstpass");
230        assert_eq!(resp["id"], 1);
231    }
232
233    #[test]
234    fn tools_list_exposes_the_three_tools() {
235        let req = json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list" });
236        let resp = handle_rpc(&req, "unused.db").unwrap();
237        let names: Vec<&str> = resp["result"]["tools"]
238            .as_array()
239            .unwrap()
240            .iter()
241            .filter_map(|t| t["name"].as_str())
242            .collect();
243        assert_eq!(names, ["list_traces", "get_trace", "submit_feedback"]);
244    }
245
246    #[test]
247    fn notifications_get_no_response() {
248        let req = json!({ "jsonrpc": "2.0", "method": "notifications/initialized" });
249        assert!(handle_rpc(&req, "unused.db").is_none());
250    }
251
252    #[test]
253    fn unknown_method_errors_for_requests_but_not_notifications() {
254        let request = json!({ "jsonrpc": "2.0", "id": 9, "method": "does/not/exist" });
255        let resp = handle_rpc(&request, "unused.db").unwrap();
256        assert_eq!(resp["error"]["code"], -32601);
257
258        let notification = json!({ "jsonrpc": "2.0", "method": "does/not/exist" });
259        assert!(handle_rpc(&notification, "unused.db").is_none());
260    }
261
262    #[test]
263    fn list_traces_on_empty_store_is_empty_not_error() {
264        let db = tmp_db();
265        let req = json!({ "jsonrpc": "2.0", "id": 3, "method": "tools/call",
266                          "params": { "name": "list_traces", "arguments": {} } });
267        let resp = handle_rpc(&req, &db).unwrap();
268        assert!(
269            resp["result"]["isError"].is_null(),
270            "empty store is not an error"
271        );
272        assert_eq!(resp["result"]["content"][0]["text"], "[]");
273    }
274
275    #[test]
276    fn get_trace_unknown_is_tool_error() {
277        let db = tmp_db();
278        // Touch the store so the file exists but has no such trace.
279        let _ = store::load_all_traces(&db);
280        let req = json!({ "jsonrpc": "2.0", "id": 4, "method": "tools/call",
281                          "params": { "name": "get_trace", "arguments": { "trace_id": "nope" } } });
282        let resp = handle_rpc(&req, &db).unwrap();
283        assert_eq!(resp["result"]["isError"], true);
284    }
285
286    #[test]
287    fn submit_feedback_validates_and_rejects_unknown_trace() {
288        let db = tmp_db();
289        let _ = store::load_all_traces(&db);
290
291        // Bad verdict → tool error.
292        let bad = json!({ "jsonrpc": "2.0", "id": 5, "method": "tools/call",
293            "params": { "name": "submit_feedback", "arguments": {
294                "trace_id": "t", "gate_id": "g", "verdict": "banana", "reporter": "ci" } } });
295        assert_eq!(handle_rpc(&bad, &db).unwrap()["result"]["isError"], true);
296
297        // Valid shape but unknown trace → tool error.
298        let unknown = json!({ "jsonrpc": "2.0", "id": 6, "method": "tools/call",
299            "params": { "name": "submit_feedback", "arguments": {
300                "trace_id": "missing", "gate_id": "g", "verdict": "pass", "reporter": "ci" } } });
301        assert_eq!(
302            handle_rpc(&unknown, &db).unwrap()["result"]["isError"],
303            true
304        );
305    }
306
307    #[test]
308    fn unknown_tool_is_tool_error() {
309        let req = json!({ "jsonrpc": "2.0", "id": 7, "method": "tools/call",
310                          "params": { "name": "no_such_tool", "arguments": {} } });
311        let resp = handle_rpc(&req, "unused.db").unwrap();
312        assert_eq!(resp["result"]["isError"], true);
313    }
314}
firstpass_proxy/mcp.rs

firstpass_proxy/
mcp.rs
