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

cmd_usage/ocuse/
mcp.rs
