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    // `--tz` per call: the bucket offset is the one thing a caller routinely
106    // wants different from the server's flag.
107    let tz = match args.get("tz").and_then(|v| v.as_str()) {
108        None => tz_secs,
109        Some(s) => crate::cli::parse_tz(s).ok_or((
110            -32602,
111            "tz needs an offset like +05:30 or -08:00".to_string(),
112        ))?,
113    };
114    // Rows to list: `rows` (the CLI flag's name), `limit` still accepted because
115    // it shipped first, and the server returns every row when neither is given.
116    let rows_cap = |fallback: usize| -> usize {
117        match args
118            .get("rows")
119            .or_else(|| args.get("limit"))
120            .and_then(|v| v.as_u64())
121        {
122            Some(0) | None => fallback,
123            Some(n) => n as usize,
124        }
125    };
126    // `since` pins the read window exactly, as in the CLI.
127    let since_ms = match args.get("since").and_then(|v| v.as_str()) {
128        None => None,
129        Some(s) => Some(cmduse_core::dates::parse_iso_utc(s).ok_or((
130            -32602,
131            format!("since needs an ISO-8601 instant, got '{s}'"),
132        ))? as i64),
133    };
134    // The tier may be overridden per call; it defaults to the server's flag.
135    let plan_id = args
136        .get("plan")
137        .and_then(|v| v.as_str())
138        .unwrap_or(tier)
139        .to_string();
140    let catalog = ark::catalog();
141
142    let rows = || {
143        let cutoff =
144            since_ms.unwrap_or_else(|| window::now_ms() - number("days", days) * 86_400_000);
145        db::read_rows(db_path, cutoff)
146    };
147    let text = match name {
148        "usage" => rows().map(|rows| {
149            // The same period the dashboard uses: the meter names the billing
150            // anniversary, so take its month rather than the calendar one, and
151            // carry the resolved end so the header can print it.
152            let mut report = window::build(
153                &rows,
154                &catalog,
155                &plan_id,
156                window::now_ms(),
157                account_month.map(|(start, _)| start).or(period_start),
158                account,
159            );
160            if account_month.is_some() {
161                report.period_label = "Agent Plan billing month".into();
162                report.period_end_ms = account_month.map(|(_, reset)| reset);
163            }
164            report.account = account.cloned();
165            report.account_error = account_error.map(str::to_string);
166            let _ = colour;
167            render::render_text(&report, db_path, 15, colour, render::DEFAULT_BAR_WIDTH)
168        }),
169        "plans" => Ok(render::plans_text(colour)),
170        "models" => {
171            if catalog.plan(&plan_id).is_none() {
172                return Err((
173                    -32602,
174                    format!("plan must be one of: {}", catalog.plan_ids().join(", ")),
175                ));
176            }
177            let gated = args.get("gated").and_then(|v| v.as_bool()).unwrap_or(false);
178            Ok(render::models_text(&plan_id, gated, colour))
179        }
180        "daily" => rows().map(|rows| {
181            render::bucket_text(
182                &rows,
183                &catalog,
184                window::now_ms(),
185                number("days", days) * 86_400_000,
186                86_400_000,
187                tz,
188                colour,
189                crate::report_render::Fmt::Table,
190            )
191        }),
192        "hourly" => rows().map(|rows| {
193            render::bucket_text(
194                &rows,
195                &catalog,
196                window::now_ms(),
197                number("hours", 24) * 3_600_000,
198                3_600_000,
199                tz,
200                colour,
201                crate::report_render::Fmt::Table,
202            )
203        }),
204        "session" => rows().map(|rows| {
205            render::session_text(
206                &rows,
207                &catalog,
208                rows_cap(usize::MAX),
209                colour,
210                crate::report_render::Fmt::Table,
211            )
212        }),
213        "model" => rows().map(|rows| {
214            let report = window::build(
215                &rows,
216                &catalog,
217                &plan_id,
218                window::now_ms(),
219                period_start,
220                account,
221            );
222            match args.get("id").and_then(|v| v.as_str()) {
223                Some(wanted) => render::model_text(&report, wanted, colour),
224                None => render::models_period_text(&report, rows_cap(usize::MAX), colour),
225            }
226        }),
227        _ => {
228            return Ok(json!({
229                "content": [{ "type": "text", "text": format!("unknown tool: {name}") }],
230                "isError": true,
231            }))
232        }
233    };
234    Ok(match text {
235        Ok(body) => json!({ "content": [{ "type": "text", "text": body }] }),
236        Err(error) => json!({
237            "content": [{ "type": "text", "text": format!("error: {error}") }],
238            "isError": true,
239        }),
240    })
241}
242
243fn tools() -> Value {
244    json!([
245        { "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" } } } },
246        { "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": {} } },
247        { "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" } } } },
248        { "name": "daily", "description": "AFP, requests and tokens per day.", "inputSchema": { "type": "object", "properties": { "days": { "type": "integer", "description": "days of history (default 120)" }, "tz": { "type": "string", "description": "UTC offset for the day buckets, e.g. +05:30" } } } },
249        { "name": "hourly", "description": "AFP, requests and tokens per hour.", "inputSchema": { "type": "object", "properties": { "hours": { "type": "integer", "description": "hours of history (default 24)" }, "tz": { "type": "string", "description": "UTC offset for the hour buckets, e.g. -08:00" } } } },
250        { "name": "session", "description": "AFP per opencode session, newest first.", "inputSchema": { "type": "object", "properties": { "since": { "type": "string", "description": "ISO-8601 instant to count from, e.g. 2026-08-27T00:00:00Z" }, "days": { "type": "integer", "description": "days back (used when since is absent)" }, "rows": { "type": "integer", "description": "sessions to list (0 = every row)" } } } },
251        { "name": "model", "description": "One model's coefficient, tokens and AFP for the period against the given tier; without id, the whole period table.", "inputSchema": { "type": "object", "properties": { "id": { "type": "string", "description": "model id, e.g. deepseek-v4.1-flash (omit for the table)" }, "plan": { "type": "string", "description": "ARK tier: small, medium (default), large, or max" }, "since": { "type": "string", "description": "ISO-8601 instant to count from" }, "days": { "type": "integer", "description": "days back (used when since is absent)" }, "rows": { "type": "integer", "description": "table rows to list (0 = every row)" } } } },
252    ])
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    fn call(method: &str, params: Value, id: i64) -> Value {
260        handle(
261            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
262            &Settings {
263                db_path: "/nonexistent/opencode.db",
264                days: 120,
265                tz_secs: 0,
266                colour: false,
267                tier: ark::DEFAULT_PLAN,
268                period_start: None,
269                account_month: None,
270                // No meters: the tests never shell out, and the local estimate is
271                // what the fixtures are written against.
272                account: None,
273                account_error: None,
274            },
275        )
276        .expect("notification")
277    }
278
279    #[test]
280    fn initialize_negotiates_the_client_version() {
281        let r = call("initialize", json!({ "protocolVersion": "2025-03-26" }), 1);
282        assert_eq!(r["result"]["protocolVersion"], "2025-03-26");
283        assert_eq!(r["result"]["serverInfo"]["name"], "arkuse");
284    }
285
286    #[test]
287    fn unknown_method_is_jsonrpc_error() {
288        let r = call("resources/list", json!({}), 2);
289        assert_eq!(r["error"]["code"], -32601);
290    }
291
292    #[test]
293    fn unknown_tool_is_iserror_result_not_rpc_error() {
294        let r = call("tools/call", json!({ "name": "nope", "arguments": {} }), 3);
295        assert!(r.get("error").is_none());
296        assert_eq!(r["result"]["isError"], true);
297    }
298
299    #[test]
300    fn a_missing_store_reports_an_error_result() {
301        let r = call("tools/call", json!({ "name": "usage", "arguments": {} }), 4);
302        assert_eq!(r["result"]["isError"], true, "{r}");
303        assert!(r["result"]["content"][0]["text"]
304            .as_str()
305            .unwrap()
306            .contains("error"));
307    }
308
309    #[test]
310    fn tools_list_covers_every_command() {
311        let r = call("tools/list", json!({}), 5);
312        // Sorted set: fails if any MCP server's tool list drifts (cmduse and
313        // ocuse assert the same seven).
314        let mut names: Vec<&str> = r["result"]["tools"]
315            .as_array()
316            .unwrap()
317            .iter()
318            .map(|t| t["name"].as_str().unwrap())
319            .collect();
320        names.sort_unstable();
321        assert_eq!(
322            names,
323            ["daily", "hourly", "model", "models", "plans", "session", "usage"]
324        );
325    }
326
327    #[test]
328    fn tool_arguments_are_the_family_surface() {
329        let r = call("tools/list", json!({}), 40);
330        let tools = &r["result"]["tools"];
331        let params = |name: &str| crate::mcp_tests::tool_params(tools, name);
332        // The ARK source forces `plan` (the tier); the rest is the family
333        // surface cmduse and ocuse carry.
334        assert_eq!(params("usage"), "plan");
335        assert_eq!(params("plans"), "");
336        assert_eq!(params("models"), "gated plan");
337        assert_eq!(params("daily"), "days tz");
338        assert_eq!(params("hourly"), "hours tz");
339        assert_eq!(params("session"), "days rows since");
340        assert_eq!(params("model"), "days id plan rows since");
341    }
342
343    #[test]
344    fn models_tool_is_the_catalogue_and_needs_no_store() {
345        let r = call(
346            "tools/call",
347            json!({ "name": "models", "arguments": {} }),
348            6,
349        );
350        assert!(r.get("error").is_none(), "{r}");
351        let text = r["result"]["content"][0]["text"].as_str().unwrap();
352        assert!(text.contains("deepseek-v4.1-flash"), "{text}");
353        assert!(text.contains("GRANTED"), "{text}");
354    }
355
356    #[test]
357    fn models_tool_rejects_an_unknown_plan() {
358        let r = call(
359            "tools/call",
360            json!({ "name": "models", "arguments": { "plan": "nope" } }),
361            7,
362        );
363        assert_eq!(r["error"]["code"], -32602, "{r}");
364    }
365}
cmd_usage/arkuse/mcp.rs

cmd_usage/arkuse/
mcp.rs
