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, models, daily,
4// hourly, session, model — the same seven cmduse exposes.
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                render::DEFAULT_BAR_WIDTH,
112            )
113        }),
114        "plans" => Ok(render::plans_text(colour)),
115        "models" => {
116            // The catalogue is static, but a per-call plan override must be real.
117            let catalog = zen::catalog();
118            if catalog.plan(&plan_id).is_none() {
119                return Err((
120                    -32602,
121                    format!("plan must be one of: {}", catalog.plan_ids().join(", ")),
122                ));
123            }
124            let gated = args.get("gated").and_then(|v| v.as_bool()).unwrap_or(false);
125            Ok(render::models_text(&plan_id, gated, colour))
126        }
127        "daily" => rows().map(|rows| {
128            render::bucket_text(
129                &rows,
130                window::now_ms(),
131                number("days", days) * 86_400_000,
132                86_400_000,
133                tz_secs,
134                colour,
135                crate::report_render::Fmt::Table,
136            )
137        }),
138        "hourly" => rows().map(|rows| {
139            render::bucket_text(
140                &rows,
141                window::now_ms(),
142                number("hours", 24) * 3_600_000,
143                3_600_000,
144                tz_secs,
145                colour,
146                crate::report_render::Fmt::Table,
147            )
148        }),
149        "session" => rows().map(|rows| {
150            render::session_text(
151                &rows,
152                number("limit", 20) as usize,
153                colour,
154                crate::report_render::Fmt::Table,
155            )
156        }),
157        "model" => {
158            let wanted = args
159                .get("id")
160                .and_then(|v| v.as_str())
161                .unwrap_or_default()
162                .to_string();
163            rows().map(|rows| {
164                let report = window::build(
165                    &rows,
166                    &zen::catalog(),
167                    &plan_id,
168                    window::now_ms(),
169                    period_start,
170                );
171                render::model_text(&report, &wanted, colour)
172            })
173        }
174        _ => {
175            return Ok(json!({
176                "content": [{ "type": "text", "text": format!("unknown tool: {name}") }],
177                "isError": true,
178            }))
179        }
180    };
181    Ok(match text {
182        Ok(body) => json!({ "content": [{ "type": "text", "text": body }] }),
183        Err(error) => json!({
184            "content": [{ "type": "text", "text": format!("error: {error}") }],
185            "isError": true,
186        }),
187    })
188}
189
190fn tools() -> Value {
191    json!([
192        { "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" } } } },
193        { "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": {} } },
194        { "name": "models", "description": "The Go catalogue for a plan: every model's $/1M in/out/cache-read rates and its monthly limit, marked granted or not.", "inputSchema": { "type": "object", "properties": { "plan": { "type": "string", "description": "Go plan: go (default) or go-plus" }, "gated": { "type": "boolean", "description": "Only models the plan grants" } } } },
195        { "name": "daily", "description": "Spend, requests and tokens per day.", "inputSchema": { "type": "object", "properties": { "days": { "type": "integer", "description": "days of history (default 120)" } } } },
196        { "name": "hourly", "description": "Spend, requests and tokens per hour.", "inputSchema": { "type": "object", "properties": { "hours": { "type": "integer", "description": "hours of history (default 24)" } } } },
197        { "name": "session", "description": "Spend per OpenCode session, newest first.", "inputSchema": { "type": "object", "properties": { "limit": { "type": "integer", "description": "sessions to list (default 20)" } } } },
198        { "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"] } },
199    ])
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    fn call(method: &str, params: Value, id: i64) -> Value {
207        handle(
208            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
209            "/nonexistent/opencode.db",
210            120,
211            0,
212            false,
213            crate::ocuse::zen::DEFAULT_PLAN,
214            None,
215        )
216        .expect("notification")
217    }
218
219    #[test]
220    fn initialize_negotiates_the_client_version() {
221        let r = call("initialize", json!({ "protocolVersion": "2025-03-26" }), 1);
222        assert_eq!(r["result"]["protocolVersion"], "2025-03-26");
223        assert_eq!(r["result"]["serverInfo"]["name"], "ocuse");
224    }
225
226    #[test]
227    fn unknown_method_is_jsonrpc_error() {
228        let r = call("resources/list", json!({}), 2);
229        assert_eq!(r["error"]["code"], -32601);
230    }
231
232    #[test]
233    fn unknown_tool_is_iserror_result_not_rpc_error() {
234        let r = call("tools/call", json!({ "name": "nope", "arguments": {} }), 3);
235        assert!(r.get("error").is_none());
236        assert_eq!(r["result"]["isError"], true);
237    }
238
239    #[test]
240    fn a_missing_store_reports_an_error_result() {
241        let r = call("tools/call", json!({ "name": "usage", "arguments": {} }), 4);
242        assert_eq!(r["result"]["isError"], true, "{r}");
243        assert!(r["result"]["content"][0]["text"]
244            .as_str()
245            .unwrap()
246            .contains("error"));
247    }
248
249    #[test]
250    fn tools_list_covers_every_command() {
251        let r = call("tools/list", json!({}), 5);
252        // Sorted set: fails if either MCP server's tool list drifts (the cmduse
253        // side asserts the same seven).
254        let mut names: Vec<&str> = r["result"]["tools"]
255            .as_array()
256            .unwrap()
257            .iter()
258            .map(|t| t["name"].as_str().unwrap())
259            .collect();
260        names.sort_unstable();
261        assert_eq!(
262            names,
263            ["daily", "hourly", "model", "models", "plans", "session", "usage"]
264        );
265    }
266
267    #[test]
268    fn models_tool_is_the_catalogue_and_needs_no_store() {
269        let r = call(
270            "tools/call",
271            json!({ "name": "models", "arguments": {} }),
272            6,
273        );
274        assert!(r.get("error").is_none(), "{r}");
275        let text = r["result"]["content"][0]["text"].as_str().unwrap();
276        assert!(text.contains("glm-5.3-flash"), "{text}");
277        assert!(text.contains("GRANTED"), "{text}");
278    }
279
280    #[test]
281    fn models_tool_rejects_an_unknown_plan() {
282        let r = call(
283            "tools/call",
284            json!({ "name": "models", "arguments": { "plan": "nope" } }),
285            7,
286        );
287        assert_eq!(r["error"]["code"], -32602, "{r}");
288    }
289}
cmd_usage/ocuse/mcp.rs

cmd_usage/ocuse/
mcp.rs
