cmd_usage/arkuse/
mcp.rs

1// MCP stdio server for arkuse — the same JSON-RPC profile cmduse and ocuse
2// speak (initialize / ping / tools/list / tools/call), and the same seven tools:
3// usage, plans, models, daily, hourly, session, model. A host that drives either
4// other binary drives this one with no change but the plan names.
5//
6// Read-only: it reads opencode.db and the hand-maintained catalogue. The control
7// plane's meters are fetched once by the caller (the CLI's rule: one fetch per
8// run, not one per tool call) and handed in, so a `usage` call shows the same
9// rows the dashboard does — and the same reason when `arkcli` could not answer.
10use crate::arkuse::api::Usage;
11use crate::arkuse::{ark, db, render, window};
12use serde_json::{json, Value};
13
14const PROTOCOL_VERSION: &str = "2025-06-18";
15
16/// Everything a tool call needs, resolved once by the caller.
17#[derive(Clone, Copy)]
18pub struct Settings<'a> {
19    pub db_path: &'a str,
20    pub days: i64,
21    pub tz_secs: i64,
22    pub colour: bool,
23    pub tier: &'a str,
24    pub period_start: Option<i64>,
25    /// The control plane's billing month (start, reset), already resolved by the
26    /// caller; `None` when it did not answer or the period was pinned.
27    pub account_month: Option<(i64, i64)>,
28    pub account: Option<&'a Usage>,
29    /// Why the meters are missing — the dashboard's own line, so an MCP `usage`
30    /// call explains itself instead of showing a bare local estimate.
31    pub account_error: Option<&'a str>,
32}
33
34/// Serve stdin -> stdout until the host closes the pipe.
35pub fn run(settings: Settings) {
36    use std::io::{BufRead, Write};
37    let stdin = std::io::stdin();
38    let mut out = std::io::stdout().lock();
39    for line in stdin.lock().lines() {
40        let Ok(line) = line else { break };
41        if line.trim().is_empty() {
42            continue;
43        }
44        let Ok(msg) = serde_json::from_str::<Value>(&line) else {
45            continue; // not JSON: ignore, a host bug shouldn't kill the server
46        };
47        if let Some(response) = handle(&msg, &settings) {
48            if writeln!(out, "{response}").is_err() {
49                break; // stdout closed: host is gone
50            }
51            out.flush().ok();
52        }
53    }
54}
55
56/// One message in, one response out; `None` for notifications (no `id`).
57pub fn handle(msg: &Value, settings: &Settings) -> Option<Value> {
58    let method = msg.get("method")?.as_str()?;
59    let id = msg.get("id")?;
60    let result = match method {
61        "initialize" => Ok(json!({
62            "protocolVersion": msg
63                .pointer("/params/protocolVersion")
64                .and_then(|v| v.as_str())
65                .unwrap_or(PROTOCOL_VERSION),
66            "capabilities": { "tools": {} },
67            "serverInfo": { "name": "arkuse", "version": env!("CARGO_PKG_VERSION") },
68        })),
69        "ping" => Ok(json!({})),
70        "tools/list" => Ok(json!({ "tools": tools() })),
71        "tools/call" => tool_call(msg, settings),
72        other => Err((-32601, format!("method not found: {other}"))),
73    };
74    Some(match result {
75        Ok(value) => json!({ "jsonrpc": "2.0", "id": id, "result": value }),
76        Err((code, message)) => json!({
77            "jsonrpc": "2.0", "id": id,
78            "error": { "code": code, "message": message },
79        }),
80    })
81}
82
83fn tool_call(msg: &Value, settings: &Settings) -> Result<Value, (i64, String)> {
84    let Settings {
85        db_path,
86        days,
87        tz_secs,
88        colour,
89        tier,
90        period_start,
91        account_month,
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 tier 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(tier)
110        .to_string();
111    let catalog = ark::catalog();
112
113    let rows = || db::read_rows(db_path, window::now_ms() - days * 86_400_000);
114    let text = match name {
115        "usage" => rows().map(|rows| {
116            // The same period the dashboard uses: the meter names the billing
117            // anniversary, so take its month rather than the calendar one, and
118            // carry the resolved end so the header can print it.
119            let mut report = window::build(
120                &rows,
121                &catalog,
122                &plan_id,
123                window::now_ms(),
124                account_month.map(|(start, _)| start).or(period_start),
125                account,
126            );
127            if account_month.is_some() {
128                report.period_label = "Agent Plan billing month".into();
129                report.period_end_ms = account_month.map(|(_, reset)| reset);
130            }
131            report.account = account.cloned();
132            report.account_error = account_error.map(str::to_string);
133            let _ = colour;
134            render::render_text(&report, db_path, 15, colour, render::DEFAULT_BAR_WIDTH)
135        }),
136        "plans" => Ok(render::plans_text(colour)),
137        "models" => {
138            if catalog.plan(&plan_id).is_none() {
139                return Err((
140                    -32602,
141                    format!("plan must be one of: {}", catalog.plan_ids().join(", ")),
142                ));
143            }
144            let gated = args.get("gated").and_then(|v| v.as_bool()).unwrap_or(false);
145            Ok(render::models_text(&plan_id, gated, colour))
146        }
147        "daily" => rows().map(|rows| {
148            render::bucket_text(
149                &rows,
150                &catalog,
151                window::now_ms(),
152                number("days", days) * 86_400_000,
153                86_400_000,
154                tz_secs,
155                colour,
156                crate::report_render::Fmt::Table,
157            )
158        }),
159        "hourly" => rows().map(|rows| {
160            render::bucket_text(
161                &rows,
162                &catalog,
163                window::now_ms(),
164                number("hours", 24) * 3_600_000,
165                3_600_000,
166                tz_secs,
167                colour,
168                crate::report_render::Fmt::Table,
169            )
170        }),
171        "session" => rows().map(|rows| {
172            render::session_text(
173                &rows,
174                &catalog,
175                number("limit", 20) as usize,
176                colour,
177                crate::report_render::Fmt::Table,
178            )
179        }),
180        "model" => {
181            let wanted = args
182                .get("id")
183                .and_then(|v| v.as_str())
184                .unwrap_or_default()
185                .to_string();
186            rows().map(|rows| {
187                let report = window::build(
188                    &rows,
189                    &catalog,
190                    &plan_id,
191                    window::now_ms(),
192                    period_start,
193                    account,
194                );
195                render::model_text(&report, &wanted, colour)
196            })
197        }
198        _ => {
199            return Ok(json!({
200                "content": [{ "type": "text", "text": format!("unknown tool: {name}") }],
201                "isError": true,
202            }))
203        }
204    };
205    Ok(match text {
206        Ok(body) => json!({ "content": [{ "type": "text", "text": body }] }),
207        Err(error) => json!({
208            "content": [{ "type": "text", "text": format!("error: {error}") }],
209            "isError": true,
210        }),
211    })
212}
213
214fn tools() -> Value {
215    json!([
216        { "name": "usage", "description": "Volcengine ARK Agent Plan usage in AFP: the plan's own 5-hour / weekly / monthly meters, this machine's per-model AFP, and the reconciliation between them.", "inputSchema": { "type": "object", "properties": { "plan": { "type": "string", "description": "ARK tier: small, medium (default), large, or max" } } } },
217        { "name": "plans", "description": "The ARK Agent Plan catalogue: each personal tier's price and monthly AFP quota, the window shares, and the enterprise price list.", "inputSchema": { "type": "object", "properties": {} } },
218        { "name": "models", "description": "ARK model coefficients: AFP per 10,000 tokens, the yuan rate that implies, the minimum tier, and whether the given tier grants it.", "inputSchema": { "type": "object", "properties": { "plan": { "type": "string", "description": "ARK tier: small, medium (default), large, or max" }, "gated": { "type": "boolean", "description": "Only models the tier grants" } } } },
219        { "name": "daily", "description": "AFP, requests and tokens per day.", "inputSchema": { "type": "object", "properties": { "days": { "type": "integer", "description": "days of history (default 120)" } } } },
220        { "name": "hourly", "description": "AFP, requests and tokens per hour.", "inputSchema": { "type": "object", "properties": { "hours": { "type": "integer", "description": "hours of history (default 24)" } } } },
221        { "name": "session", "description": "AFP per opencode session, newest first.", "inputSchema": { "type": "object", "properties": { "limit": { "type": "integer", "description": "sessions to list (default 20)" } } } },
222        { "name": "model", "description": "One model's coefficient, tokens and AFP for the period, against the given tier.", "inputSchema": { "type": "object", "properties": { "id": { "type": "string", "description": "model id, e.g. deepseek-v4.1-flash" }, "plan": { "type": "string", "description": "ARK tier: small, medium (default), large, or max" } }, "required": ["id"] } },
223    ])
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    fn call(method: &str, params: Value, id: i64) -> Value {
231        handle(
232            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
233            &Settings {
234                db_path: "/nonexistent/opencode.db",
235                days: 120,
236                tz_secs: 0,
237                colour: false,
238                tier: ark::DEFAULT_PLAN,
239                period_start: None,
240                account_month: None,
241                // No meters: the tests never shell out, and the local estimate is
242                // what the fixtures are written against.
243                account: None,
244                account_error: None,
245            },
246        )
247        .expect("notification")
248    }
249
250    #[test]
251    fn initialize_negotiates_the_client_version() {
252        let r = call("initialize", json!({ "protocolVersion": "2025-03-26" }), 1);
253        assert_eq!(r["result"]["protocolVersion"], "2025-03-26");
254        assert_eq!(r["result"]["serverInfo"]["name"], "arkuse");
255    }
256
257    #[test]
258    fn unknown_method_is_jsonrpc_error() {
259        let r = call("resources/list", json!({}), 2);
260        assert_eq!(r["error"]["code"], -32601);
261    }
262
263    #[test]
264    fn unknown_tool_is_iserror_result_not_rpc_error() {
265        let r = call("tools/call", json!({ "name": "nope", "arguments": {} }), 3);
266        assert!(r.get("error").is_none());
267        assert_eq!(r["result"]["isError"], true);
268    }
269
270    #[test]
271    fn a_missing_store_reports_an_error_result() {
272        let r = call("tools/call", json!({ "name": "usage", "arguments": {} }), 4);
273        assert_eq!(r["result"]["isError"], true, "{r}");
274        assert!(r["result"]["content"][0]["text"]
275            .as_str()
276            .unwrap()
277            .contains("error"));
278    }
279
280    #[test]
281    fn tools_list_covers_every_command() {
282        let r = call("tools/list", json!({}), 5);
283        // Sorted set: fails if any MCP server's tool list drifts (cmduse and
284        // ocuse assert the same seven).
285        let mut names: Vec<&str> = r["result"]["tools"]
286            .as_array()
287            .unwrap()
288            .iter()
289            .map(|t| t["name"].as_str().unwrap())
290            .collect();
291        names.sort_unstable();
292        assert_eq!(
293            names,
294            ["daily", "hourly", "model", "models", "plans", "session", "usage"]
295        );
296    }
297
298    #[test]
299    fn models_tool_is_the_catalogue_and_needs_no_store() {
300        let r = call(
301            "tools/call",
302            json!({ "name": "models", "arguments": {} }),
303            6,
304        );
305        assert!(r.get("error").is_none(), "{r}");
306        let text = r["result"]["content"][0]["text"].as_str().unwrap();
307        assert!(text.contains("deepseek-v4.1-flash"), "{text}");
308        assert!(text.contains("GRANTED"), "{text}");
309    }
310
311    #[test]
312    fn models_tool_rejects_an_unknown_plan() {
313        let r = call(
314            "tools/call",
315            json!({ "name": "models", "arguments": { "plan": "nope" } }),
316            7,
317        );
318        assert_eq!(r["error"]["code"], -32602, "{r}");
319    }
320}
cmd_usage/arkuse/mcp.rs

cmd_usage/arkuse/
mcp.rs
