cmd_usage/ocuse/
mcp.rs

1// MCP stdio server for ocuse — the same JSON-RPC profile cmduse speaks
2// (initialize / ping / tools/list / tools/call), so a host that already drives
3// cmduse can drive this. Tools mirror the CLI: usage, plans, daily, hourly,
4// session, model.
5//
6// Read-only, like everything else here: it reads opencode.db and the generated
7// catalogue, never the network.
8use crate::ocuse::{db, render, window, zen};
9use serde_json::{json, Value};
10
11const PROTOCOL_VERSION: &str = "2025-06-18";
12
13/// Serve stdin -> stdout until the host closes the pipe.
14pub fn run(db_path: &str, days: i64, tz_secs: i64, colour: bool, period_start: Option<i64>) {
15    use std::io::{BufRead, Write};
16    let stdin = std::io::stdin();
17    let mut out = std::io::stdout().lock();
18    for line in stdin.lock().lines() {
19        let Ok(line) = line else { break };
20        if line.trim().is_empty() {
21            continue;
22        }
23        let Ok(msg) = serde_json::from_str::<Value>(&line) else {
24            continue; // not JSON: ignore, a host bug shouldn't kill the server
25        };
26        if let Some(response) = handle(&msg, db_path, days, tz_secs, colour, period_start) {
27            if writeln!(out, "{response}").is_err() {
28                break; // stdout closed: host is gone
29            }
30            out.flush().ok();
31        }
32    }
33}
34
35/// One message in, one response out; `None` for notifications (no `id`).
36pub fn handle(
37    msg: &Value,
38    db_path: &str,
39    days: i64,
40    tz_secs: i64,
41    colour: bool,
42    period_start: Option<i64>,
43) -> Option<Value> {
44    let method = msg.get("method")?.as_str()?;
45    let id = msg.get("id")?;
46    let result = match method {
47        "initialize" => Ok(json!({
48            "protocolVersion": msg
49                .pointer("/params/protocolVersion")
50                .and_then(|v| v.as_str())
51                .unwrap_or(PROTOCOL_VERSION),
52            "capabilities": { "tools": {} },
53            "serverInfo": { "name": "ocuse", "version": env!("CARGO_PKG_VERSION") },
54        })),
55        "ping" => Ok(json!({})),
56        "tools/list" => Ok(json!({ "tools": tools() })),
57        "tools/call" => tool_call(msg, db_path, days, tz_secs, colour, period_start),
58        other => Err((-32601, format!("method not found: {other}"))),
59    };
60    Some(match result {
61        Ok(value) => json!({ "jsonrpc": "2.0", "id": id, "result": value }),
62        Err((code, message)) => json!({
63            "jsonrpc": "2.0", "id": id,
64            "error": { "code": code, "message": message },
65        }),
66    })
67}
68
69fn tool_call(
70    msg: &Value,
71    db_path: &str,
72    days: i64,
73    tz_secs: i64,
74    colour: bool,
75    period_start: Option<i64>,
76) -> Result<Value, (i64, String)> {
77    let name = msg
78        .pointer("/params/name")
79        .and_then(|v| v.as_str())
80        .unwrap_or_default();
81    let args = msg
82        .pointer("/params/arguments")
83        .cloned()
84        .unwrap_or(json!({}));
85    let number =
86        |key: &str, fallback: i64| args.get(key).and_then(|v| v.as_i64()).unwrap_or(fallback);
87
88    let rows = || db::read_rows(db_path, crate::ocuse::window::now_ms() - days * 86_400_000);
89    let text = match name {
90        "usage" => rows().map(|rows| {
91            render::render_text(
92                &window::build(&rows, &zen::catalog(), window::now_ms(), None),
93                db_path,
94                15,
95                colour,
96            )
97        }),
98        "plans" => Ok(render::plans_text()),
99        "daily" => rows().map(|rows| {
100            render::bucket_text(
101                &rows,
102                window::now_ms(),
103                number("days", days),
104                86_400_000,
105                "day",
106                tz_secs,
107            )
108        }),
109        "hourly" => rows().map(|rows| {
110            render::bucket_text(
111                &rows,
112                window::now_ms(),
113                number("hours", 24),
114                3_600_000,
115                "hour",
116                tz_secs,
117            )
118        }),
119        "session" => rows().map(|rows| render::session_text(&rows, number("limit", 20) as usize)),
120        "model" => {
121            let wanted = args
122                .get("id")
123                .and_then(|v| v.as_str())
124                .unwrap_or_default()
125                .to_string();
126            rows().map(|rows| {
127                let report = window::build(&rows, &zen::catalog(), window::now_ms(), period_start);
128                render::model_text(&report, &wanted)
129            })
130        }
131        _ => {
132            return Ok(json!({
133                "content": [{ "type": "text", "text": format!("unknown tool: {name}") }],
134                "isError": true,
135            }))
136        }
137    };
138    Ok(match text {
139        Ok(body) => json!({ "content": [{ "type": "text", "text": body }] }),
140        Err(error) => json!({
141            "content": [{ "type": "text", "text": format!("error: {error}") }],
142            "isError": true,
143        }),
144    })
145}
146
147fn tools() -> Value {
148    json!([
149        { "name": "usage", "description": "OpenCode Go/Zen usage: period totals, rolling windows, per-model spend against each model's monthly limit.", "inputSchema": { "type": "object", "properties": {} } },
150        { "name": "plans", "description": "The Go/Zen catalogue from the docs: per-model monthly limits and $/1M rates.", "inputSchema": { "type": "object", "properties": {} } },
151        { "name": "daily", "description": "Spend, requests and tokens per day.", "inputSchema": { "type": "object", "properties": { "days": { "type": "integer", "description": "days of history (default 120)" } } } },
152        { "name": "hourly", "description": "Spend, requests and tokens per hour.", "inputSchema": { "type": "object", "properties": { "hours": { "type": "integer", "description": "hours of history (default 24)" } } } },
153        { "name": "session", "description": "Spend per OpenCode session, newest first.", "inputSchema": { "type": "object", "properties": { "limit": { "type": "integer", "description": "sessions to list (default 20)" } } } },
154        { "name": "model", "description": "One model's limit, spend and window figures.", "inputSchema": { "type": "object", "properties": { "id": { "type": "string", "description": "model id, e.g. glm-5.2" } }, "required": ["id"] } },
155    ])
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    fn call(method: &str, params: Value, id: i64) -> Value {
163        handle(
164            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
165            "/nonexistent/opencode.db",
166            120,
167            0,
168            false,
169            None,
170        )
171        .expect("notification")
172    }
173
174    #[test]
175    fn initialize_negotiates_the_client_version() {
176        let r = call("initialize", json!({ "protocolVersion": "2025-03-26" }), 1);
177        assert_eq!(r["result"]["protocolVersion"], "2025-03-26");
178        assert_eq!(r["result"]["serverInfo"]["name"], "ocuse");
179    }
180
181    #[test]
182    fn unknown_method_is_jsonrpc_error() {
183        let r = call("resources/list", json!({}), 2);
184        assert_eq!(r["error"]["code"], -32601);
185    }
186
187    #[test]
188    fn unknown_tool_is_iserror_result_not_rpc_error() {
189        let r = call("tools/call", json!({ "name": "nope", "arguments": {} }), 3);
190        assert!(r.get("error").is_none());
191        assert_eq!(r["result"]["isError"], true);
192    }
193
194    #[test]
195    fn a_missing_store_reports_an_error_result() {
196        let r = call("tools/call", json!({ "name": "usage", "arguments": {} }), 4);
197        assert_eq!(r["result"]["isError"], true, "{r}");
198        assert!(r["result"]["content"][0]["text"]
199            .as_str()
200            .unwrap()
201            .contains("error"));
202    }
203
204    #[test]
205    fn tools_list_covers_every_command() {
206        let r = call("tools/list", json!({}), 5);
207        let names: Vec<&str> = r["result"]["tools"]
208            .as_array()
209            .unwrap()
210            .iter()
211            .map(|t| t["name"].as_str().unwrap())
212            .collect();
213        for expected in ["usage", "plans", "daily", "hourly", "session", "model"] {
214            assert!(
215                names.contains(&expected),
216                "{expected} missing from {names:?}"
217            );
218        }
219    }
220}
cmd_usage/ocuse/mcp.rs

cmd_usage/ocuse/
mcp.rs
