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: it reads opencode.db and the generated catalogue. The Go usage
7// endpoint's meters are not fetched here — the caller resolves them once (the
8// CLI's own rule: one fetch per run, not one per tool call) and hands them in,
9// so a `usage` call shows the same window rows the dashboard does.
10use crate::ocuse::api::Usage;
11use crate::ocuse::{db, render, window, zen};
12use serde_json::{json, Value};
13
14const PROTOCOL_VERSION: &str = "2025-06-18";
15
16/// What every tool call needs, resolved once by the caller — the per-run
17/// settings plus the account meters it already fetched. Bundled rather than
18/// passed positionally so adding a source (as `account` did) does not grow every
19/// signature in the file.
20#[derive(Clone, Copy)]
21pub struct Settings<'a> {
22    pub db_path: &'a str,
23    pub days: i64,
24    pub tz_secs: i64,
25    pub colour: bool,
26    pub plan: &'a str,
27    pub period_start: Option<i64>,
28    pub account: Option<&'a Usage>,
29    /// Why the account's meters are missing, when the CLI resolved it — the same
30    /// line the dashboard prints, so a `usage` call through MCP explains itself
31    /// too instead of showing bare local estimates.
32    pub account_error: Option<&'a str>,
33}
34
35/// Serve stdin -> stdout until the host closes the pipe.
36pub fn run(settings: Settings) {
37    use std::io::{BufRead, Write};
38    let stdin = std::io::stdin();
39    let mut out = std::io::stdout().lock();
40    for line in stdin.lock().lines() {
41        let Ok(line) = line else { break };
42        if line.trim().is_empty() {
43            continue;
44        }
45        let Ok(msg) = serde_json::from_str::<Value>(&line) else {
46            continue; // not JSON: ignore, a host bug shouldn't kill the server
47        };
48        if let Some(response) = handle(&msg, &settings) {
49            if writeln!(out, "{response}").is_err() {
50                break; // stdout closed: host is gone
51            }
52            out.flush().ok();
53        }
54    }
55}
56
57/// One message in, one response out; `None` for notifications (no `id`).
58pub fn handle(msg: &Value, settings: &Settings) -> Option<Value> {
59    let method = msg.get("method")?.as_str()?;
60    let id = msg.get("id")?;
61    let result = match method {
62        "initialize" => Ok(json!({
63            "protocolVersion": msg
64                .pointer("/params/protocolVersion")
65                .and_then(|v| v.as_str())
66                .unwrap_or(PROTOCOL_VERSION),
67            "capabilities": { "tools": {} },
68            "serverInfo": { "name": "ocuse", "version": env!("CARGO_PKG_VERSION") },
69        })),
70        "ping" => Ok(json!({})),
71        "tools/list" => Ok(json!({ "tools": tools() })),
72        "tools/call" => tool_call(msg, settings),
73        other => Err((-32601, format!("method not found: {other}"))),
74    };
75    Some(match result {
76        Ok(value) => json!({ "jsonrpc": "2.0", "id": id, "result": value }),
77        Err((code, message)) => json!({
78            "jsonrpc": "2.0", "id": id,
79            "error": { "code": code, "message": message },
80        }),
81    })
82}
83
84fn tool_call(msg: &Value, settings: &Settings) -> Result<Value, (i64, String)> {
85    let Settings {
86        db_path,
87        days,
88        tz_secs,
89        colour,
90        plan,
91        period_start,
92        account,
93        account_error,
94    } = *settings;
95    let name = msg
96        .pointer("/params/name")
97        .and_then(|v| v.as_str())
98        .unwrap_or_default();
99    let args = msg
100        .pointer("/params/arguments")
101        .cloned()
102        .unwrap_or(json!({}));
103    let number =
104        |key: &str, fallback: i64| args.get(key).and_then(|v| v.as_i64()).unwrap_or(fallback);
105    // The plan may be overridden per call; it defaults to the server's flag.
106    let plan_id = args
107        .get("plan")
108        .and_then(|v| v.as_str())
109        .unwrap_or(plan)
110        .to_string();
111
112    let rows = || db::read_rows(db_path, crate::ocuse::window::now_ms() - days * 86_400_000);
113    let text = match name {
114        "usage" => rows().map(|rows| {
115            let mut report = window::build(
116                &rows,
117                &zen::catalog(),
118                &plan_id,
119                window::now_ms(),
120                None,
121                account,
122            );
123            // The account's own meters when the caller resolved them, so a `usage`
124            // call renders the same window rows the dashboard does — and the same
125            // reason when a key was refused.
126            report.account = account.cloned();
127            report.account_error = account_error.map(str::to_string);
128            render::render_text(&report, db_path, 15, colour, render::DEFAULT_BAR_WIDTH)
129        }),
130        "plans" => Ok(render::plans_text(colour)),
131        "models" => {
132            // The catalogue is static, but a per-call plan override must be real.
133            let catalog = zen::catalog();
134            if catalog.plan(&plan_id).is_none() {
135                return Err((
136                    -32602,
137                    format!("plan must be one of: {}", catalog.plan_ids().join(", ")),
138                ));
139            }
140            let gated = args.get("gated").and_then(|v| v.as_bool()).unwrap_or(false);
141            Ok(render::models_text(&plan_id, gated, colour))
142        }
143        "daily" => rows().map(|rows| {
144            render::bucket_text(
145                &rows,
146                window::now_ms(),
147                number("days", days) * 86_400_000,
148                86_400_000,
149                tz_secs,
150                colour,
151                crate::report_render::Fmt::Table,
152            )
153        }),
154        "hourly" => rows().map(|rows| {
155            render::bucket_text(
156                &rows,
157                window::now_ms(),
158                number("hours", 24) * 3_600_000,
159                3_600_000,
160                tz_secs,
161                colour,
162                crate::report_render::Fmt::Table,
163            )
164        }),
165        "session" => rows().map(|rows| {
166            render::session_text(
167                &rows,
168                number("limit", 20) as usize,
169                colour,
170                crate::report_render::Fmt::Table,
171            )
172        }),
173        "model" => {
174            let wanted = args
175                .get("id")
176                .and_then(|v| v.as_str())
177                .unwrap_or_default()
178                .to_string();
179            rows().map(|rows| {
180                let report = window::build(
181                    &rows,
182                    &zen::catalog(),
183                    &plan_id,
184                    window::now_ms(),
185                    period_start,
186                    account,
187                );
188                render::model_text(&report, &wanted, colour)
189            })
190        }
191        _ => {
192            return Ok(json!({
193                "content": [{ "type": "text", "text": format!("unknown tool: {name}") }],
194                "isError": true,
195            }))
196        }
197    };
198    Ok(match text {
199        Ok(body) => json!({ "content": [{ "type": "text", "text": body }] }),
200        Err(error) => json!({
201            "content": [{ "type": "text", "text": format!("error: {error}") }],
202            "isError": true,
203        }),
204    })
205}
206
207fn tools() -> Value {
208    json!([
209        { "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" } } } },
210        { "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": {} } },
211        { "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" } } } },
212        { "name": "daily", "description": "Spend, requests and tokens per day.", "inputSchema": { "type": "object", "properties": { "days": { "type": "integer", "description": "days of history (default 120)" } } } },
213        { "name": "hourly", "description": "Spend, requests and tokens per hour.", "inputSchema": { "type": "object", "properties": { "hours": { "type": "integer", "description": "hours of history (default 24)" } } } },
214        { "name": "session", "description": "Spend per OpenCode session, newest first.", "inputSchema": { "type": "object", "properties": { "limit": { "type": "integer", "description": "sessions to list (default 20)" } } } },
215        { "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"] } },
216    ])
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    fn call(method: &str, params: Value, id: i64) -> Value {
224        handle(
225            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
226            &Settings {
227                db_path: "/nonexistent/opencode.db",
228                days: 120,
229                tz_secs: 0,
230                colour: false,
231                plan: crate::ocuse::zen::DEFAULT_PLAN,
232                period_start: None,
233                // No account: the tests never reach the network, and the local
234                // estimate is what the fixtures below are written against.
235                account: None,
236                account_error: None,
237            },
238        )
239        .expect("notification")
240    }
241
242    #[test]
243    fn initialize_negotiates_the_client_version() {
244        let r = call("initialize", json!({ "protocolVersion": "2025-03-26" }), 1);
245        assert_eq!(r["result"]["protocolVersion"], "2025-03-26");
246        assert_eq!(r["result"]["serverInfo"]["name"], "ocuse");
247    }
248
249    #[test]
250    fn unknown_method_is_jsonrpc_error() {
251        let r = call("resources/list", json!({}), 2);
252        assert_eq!(r["error"]["code"], -32601);
253    }
254
255    #[test]
256    fn unknown_tool_is_iserror_result_not_rpc_error() {
257        let r = call("tools/call", json!({ "name": "nope", "arguments": {} }), 3);
258        assert!(r.get("error").is_none());
259        assert_eq!(r["result"]["isError"], true);
260    }
261
262    #[test]
263    fn a_missing_store_reports_an_error_result() {
264        let r = call("tools/call", json!({ "name": "usage", "arguments": {} }), 4);
265        assert_eq!(r["result"]["isError"], true, "{r}");
266        assert!(r["result"]["content"][0]["text"]
267            .as_str()
268            .unwrap()
269            .contains("error"));
270    }
271
272    #[test]
273    fn tools_list_covers_every_command() {
274        let r = call("tools/list", json!({}), 5);
275        // Sorted set: fails if either MCP server's tool list drifts (the cmduse
276        // side asserts the same seven).
277        let mut names: Vec<&str> = r["result"]["tools"]
278            .as_array()
279            .unwrap()
280            .iter()
281            .map(|t| t["name"].as_str().unwrap())
282            .collect();
283        names.sort_unstable();
284        assert_eq!(
285            names,
286            ["daily", "hourly", "model", "models", "plans", "session", "usage"]
287        );
288    }
289
290    #[test]
291    fn models_tool_is_the_catalogue_and_needs_no_store() {
292        let r = call(
293            "tools/call",
294            json!({ "name": "models", "arguments": {} }),
295            6,
296        );
297        assert!(r.get("error").is_none(), "{r}");
298        let text = r["result"]["content"][0]["text"].as_str().unwrap();
299        assert!(text.contains("glm-5.3-flash"), "{text}");
300        assert!(text.contains("GRANTED"), "{text}");
301    }
302
303    #[test]
304    fn models_tool_rejects_an_unknown_plan() {
305        let r = call(
306            "tools/call",
307            json!({ "name": "models", "arguments": { "plan": "nope" } }),
308            7,
309        );
310        assert_eq!(r["error"]["code"], -32602, "{r}");
311    }
312}
cmd_usage/ocuse/mcp.rs

cmd_usage/ocuse/
mcp.rs
