zcode-token-usage-statusbar/usage_mcp.py at main · xhwxt/zcode-token-usage-statusbar · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# -*- coding: utf-8 -*-
"""ZCode token 用量 MCP server（stdio, 零依赖；注册名 zcode-token-usage-statusbar）。

提供工具 token_usage(scope)：current(默认)/today/days:N/sessions[:N]/models[:days]/session:<id前缀>/workspace:<目录关键词>。
协议仅实现 initialize / tools/list / tools/call；调试信息只走 stderr。
"""
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
import zusage  # noqa: E402

PROTOCOL_VERSION = "2024-11-05"

L = zusage.L   # 输出语言与 CLI 同源（config.json 的 "lang"）


def tool_token_usage(scope: str = "current") -> str:
    con = zusage.connect()
    try:
        if scope in ("current", "now"):
            return zusage.render_current(con)
        if scope == "today":
            return zusage.render_today(con)
        if scope.startswith("days:"):
            return zusage.render_days(con, int(scope.split(":", 1)[1]))
        if scope == "week":
            return zusage.render_days(con, 7)
        if scope.startswith("sessions"):
            n = int(scope.split(":", 1)[1]) if ":" in scope else 10
            return zusage.render_sessions(con, n)
        if scope.startswith("models"):
            d = int(scope.split(":", 1)[1]) if ":" in scope else 7
            return zusage.render_models(con, d)
        if scope.startswith("session:"):
            return zusage.render_session_detail(con, scope.split(":", 1)[1].strip())
        if scope.startswith("workspace:"):
            return zusage.render_workspace(con, scope.split(":", 1)[1].strip())
        return L(f"未知 scope: {scope!r}。可用: current | today | week | days:N | sessions[:N] | models[:days] | session:<id前缀> | workspace:<目录关键词>",
                 f"Unknown scope: {scope!r}. Available: current | today | week | days:N | sessions[:N] | models[:days] | session:<id prefix> | workspace:<dir keyword>")
    finally:
        con.close()


TOOLS = [
    {
        "name": "token_usage",
        "description": (
            L("查询 ZCode 的 token 用量（数据来自本地 db.sqlite，只读）。"
              "scope 可选: current(当前会话+今日,默认), today, week, days:N(近N天每日), "
              "sessions:N(最近N个会话), models:days(按模型), session:<id前缀>(指定会话,含按模型与子代理明细), "
              "workspace:<目录关键词>(按工作区聚合主会话+子代理,含按会话与按模型明细)。",
              "Query ZCode token usage (from the local db.sqlite, read-only). "
              "scope: current(current session+today, default), today, week, days:N(last N days per day), "
              "sessions:N(last N sessions), models:days(by model), session:<id prefix>(specific session, with per-model and subagent details), "
              "workspace:<dir keyword>(aggregate a workspace's main sessions+subagents, with per-session and per-model details).")
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "scope": {
                    "type": "string",
                    "description": L("查询范围，默认 current", "Query scope, default current"),
                }
            },
        },
    }
]


def reply(req_id, result):
    print(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": result}), flush=True)


def main():
    for raw in sys.stdin:
        raw = raw.strip()
        if not raw:
            continue
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            continue
        method = msg.get("method", "")
        req_id = msg.get("id")
        if method == "initialize":
            reply(
                req_id,
                {
                    "protocolVersion": msg.get("params", {}).get(
                        "protocolVersion", PROTOCOL_VERSION
                    ),
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": "zcode-token-usage", "version": "1.0.0"},
                },
            )
        elif method == "tools/list":
            reply(req_id, {"tools": TOOLS})
        elif method == "tools/call":
            params = msg.get("params", {})
            args = params.get("arguments") or {}
            try:
                text = tool_token_usage(**args)
                reply(req_id, {"content": [{"type": "text", "text": text}], "isError": False})
            except Exception as e:  # noqa: BLE001
                reply(
                    req_id,
                    {
                        "content": [{"type": "text", "text": L(f"查询失败: {e}", f"Query failed: {e}")}],
                        "isError": True,
                    },
                )
        elif req_id is not None:
            print(
                json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": req_id,
                        "error": {"code": -32601, "message": f"unknown method {method}"},
                    }
                ),
                flush=True,
            )


if __name__ == "__main__":
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass
    main()