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

cmd_usage/ocuse/
mcp.rs
