From 1e54967d9766784702b0fd2c070c805ca8d17afd Mon Sep 17 00:00:00 2001 From: iamtoruk Date: Tue, 2 Jun 2026 02:16:10 -0700 Subject: [PATCH] feat(mcp): get_usage + get_savings tools with annotations, schemas, coalescing --- src/mcp/server.ts | 130 +++++++++++++++++++++++++++++++++++++++ tests/mcp-server.test.ts | 60 ++++++++++++++++++ 2 files changed, 190 insertions(+) create mode 100644 src/mcp/server.ts create mode 100644 tests/mcp-server.test.ts diff --git a/src/mcp/server.ts b/src/mcp/server.ts new file mode 100644 index 0000000..b55cb23 --- /dev/null +++ b/src/mcp/server.ts @@ -0,0 +1,130 @@ +import { z } from 'zod' +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' +import { getDateRange } from '../cli-date.js' +import { loadPricing } from '../models.js' +import { buildMenubarPayloadForRange, type PeriodInfo } from '../usage-aggregator.js' +import type { MenubarPayload } from '../menubar-json.js' +import { redactProjectNames } from './redact.js' +import { renderSummaryTable, renderBreakdownTable, renderSavingsTable, type BreakdownBy } from './tables.js' + +const PERIOD = { today: 'today', last_7_days: 'week', last_30_days: '30days', month_to_date: 'month', last_6_months: 'all' } as const +type McpPeriod = keyof typeof PERIOD +const periodSchema = z.enum(['today', 'last_7_days', 'last_30_days', 'month_to_date', 'last_6_months']) + +type Aggregate = (periodInfo: PeriodInfo, opts: { provider?: string; optimize?: boolean }) => Promise + +const INSTRUCTIONS = + 'CodeBurn exposes local AI-coding spend data. Use get_usage for spend/usage and breakdowns (fast); ' + + 'use get_savings to find cost reductions (slower — runs a deeper analysis). Project names are pseudonymized ' + + 'unless include_project_names is true. All data is read locally from this machine; last_6_months is the widest ' + + 'window. Numbers reflect the most recent scan and may lag the current session by up to a few minutes.' + +function breakdownRows(p: MenubarPayload, by: BreakdownBy, limit: number): Array<{ name: string; costUSD: number }> { + const c = p.current + if (by === 'model') return c.topModels.slice(0, limit).map(m => ({ name: m.name, costUSD: m.cost })) + if (by === 'project') return c.topProjects.slice(0, limit).map(x => ({ name: x.name, costUSD: x.cost })) + if (by === 'task') return c.topActivities.slice(0, limit).map(a => ({ name: a.name, costUSD: a.cost })) + return Object.entries(c.providers).sort(([, a], [, b]) => b - a).slice(0, limit).map(([name, cost]) => ({ name, costUSD: cost })) +} + +export function createServer(deps: { version: string; aggregate?: Aggregate }): McpServer { + const aggregate = deps.aggregate ?? buildMenubarPayloadForRange + const inflight = new Map>() + + const getPayload = (period: McpPeriod, optimize: boolean): Promise => { + const key = `${optimize ? 'sav' : 'use'}:${period}` + const existing = inflight.get(key) + if (existing) return existing + const { range, label } = getDateRange(PERIOD[period]) + const p = aggregate({ range, label }, { provider: 'all', optimize }).finally(() => inflight.delete(key)) + inflight.set(key, p) + return p + } + + const server = new McpServer({ name: 'codeburn', version: deps.version }, { instructions: INSTRUCTIONS }) + + server.registerTool( + 'get_usage', + { + title: 'CodeBurn — usage & cost', + description: + 'Show AI coding token spend and usage for a period. Omit `by` for a headline summary; set `by` to break ' + + 'it down by project, model, task, or provider (Claude Code / Cursor / Codex). Fast. Local to this machine.', + inputSchema: { + period: periodSchema.default('today'), + by: z.enum(['project', 'model', 'task', 'provider']).optional(), + limit: z.number().int().min(1).max(100).default(20), + include_project_names: z.boolean().default(false), + }, + outputSchema: { + period: z.string(), + empty: z.boolean(), + totals: z.object({ costUSD: z.number(), calls: z.number(), sessions: z.number(), cacheHitPercent: z.number(), oneShotRate: z.number().nullable() }), + breakdown: z.array(z.object({ name: z.string(), costUSD: z.number() })).nullable(), + }, + annotations: { title: 'CodeBurn — usage & cost', readOnlyHint: true, openWorldHint: false, idempotentHint: true }, + }, + async ({ period, by, limit, include_project_names }) => { + try { + const payload = redactProjectNames(await getPayload(period, false), include_project_names) + const c = payload.current + const totals = { costUSD: c.cost, calls: c.calls, sessions: c.sessions, cacheHitPercent: c.cacheHitPercent, oneShotRate: c.oneShotRate } + if (c.calls === 0) { + return { + content: [{ type: 'text' as const, text: `No usage recorded for ${c.label} yet — run some coding sessions and try again.` }], + structuredContent: { period: c.label, empty: true, totals, breakdown: null }, + } + } + const text = by ? renderBreakdownTable(payload, by, limit) : renderSummaryTable(payload) + const breakdown = by ? breakdownRows(payload, by, limit) : null + return { + content: [{ type: 'text' as const, text }], + structuredContent: { period: c.label, empty: false, totals, breakdown }, + } + } catch (err) { + return { content: [{ type: 'text' as const, text: `codeburn: failed to read usage — ${err instanceof Error ? err.message : String(err)}` }], isError: true } + } + }, + ) + + server.registerTool( + 'get_savings', + { + title: 'CodeBurn — savings opportunities', + description: + 'Find ways to reduce AI coding cost for a period: optimization findings, retry tax (money spent re-doing ' + + 'work), and routing waste (what you would have saved on a cheaper model). Slower than get_usage.', + inputSchema: { period: periodSchema.default('last_7_days'), include_project_names: z.boolean().default(false) }, + outputSchema: { + period: z.string(), + optimize: z.object({ findingCount: z.number(), savingsUSD: z.number(), topFindings: z.array(z.object({ title: z.string(), impact: z.string(), savingsUSD: z.number() })) }), + retryTaxUSD: z.number(), + routingWasteUSD: z.number(), + }, + annotations: { title: 'CodeBurn — savings opportunities', readOnlyHint: true, openWorldHint: false, idempotentHint: true }, + }, + async ({ period, include_project_names }) => { + try { + const payload = redactProjectNames(await getPayload(period, true), include_project_names) + const c = payload.current + return { + content: [{ type: 'text' as const, text: renderSavingsTable(payload) }], + structuredContent: { period: c.label, optimize: payload.optimize, retryTaxUSD: c.retryTax.totalUSD, routingWasteUSD: c.routingWaste.totalSavingsUSD }, + } + } catch (err) { + return { content: [{ type: 'text' as const, text: `codeburn: failed to compute savings — ${err instanceof Error ? err.message : String(err)}` }], isError: true } + } + }, + ) + + return server +} + +export async function startStdioServer(version: string): Promise { + await loadPricing() + const server = createServer({ version }) + // Pre-warm the parser cache for the common case; ignore failures. + void buildMenubarPayloadForRange(getDateRange('today'), { provider: 'all', optimize: false }).catch(() => {}) + await server.connect(new StdioServerTransport()) +} diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts new file mode 100644 index 0000000..e022af2 --- /dev/null +++ b/tests/mcp-server.test.ts @@ -0,0 +1,60 @@ +import { describe, expect, it } from 'vitest' +import { Client } from '@modelcontextprotocol/sdk/client/index.js' +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js' +import { createServer } from '../src/mcp/server.js' +import type { MenubarPayload } from '../src/menubar-json.js' + +function fakePayload(calls = 100): MenubarPayload { + return { + generated: '', optimize: { findingCount: 1, savingsUSD: 2, topFindings: [{ title: 'X', impact: 'high', savingsUSD: 2 }] }, history: { daily: [] }, + current: { + label: 'Today', cost: 9, calls, sessions: 1, oneShotRate: 0.5, inputTokens: 10, outputTokens: 5, cacheHitPercent: 50, + topActivities: [{ name: 'feature', cost: 9, turns: 5, oneShotRate: 0.5 }], topModels: [{ name: 'Opus 4.8', cost: 9, calls }], + providers: { 'claude code': 9 }, topProjects: [{ name: 'real-repo', cost: 9, sessions: 1, avgCostPerSession: 9, sessionDetails: [] }], + modelEfficiency: [], topSessions: [{ project: 'real-repo', cost: 9, calls, date: '2026-06-01' }], + retryTax: { totalUSD: 1, retries: 2, editTurns: 5, byModel: [{ name: 'Opus 4.8', taxUSD: 1, retries: 2, retriesPerEdit: 0.4 }] }, + routingWaste: { totalSavingsUSD: 1, baselineModel: 'Haiku 4.5', baselineCostPerEdit: 0.01, byModel: [] }, + tools: [], skills: [], subagents: [], mcpServers: [], + }, + } as MenubarPayload +} + +async function connect(aggregate: (p: unknown, o: unknown) => Promise) { + const server = createServer({ version: 'test', aggregate: aggregate as never }) + const [a, b] = InMemoryTransport.createLinkedPair() + const client = new Client({ name: 'test', version: '1' }) + await Promise.all([server.connect(a), client.connect(b)]) + return client +} + +describe('mcp server', () => { + it('exposes exactly two read-only tools', async () => { + const client = await connect(async () => fakePayload()) + const { tools } = await client.listTools() + expect(tools.map(t => t.name).sort()).toEqual(['get_savings', 'get_usage']) + expect(tools.find(t => t.name === 'get_usage')!.annotations?.readOnlyHint).toBe(true) + }) + it('get_usage hashes project names by default', async () => { + const client = await connect(async () => fakePayload()) + const res = await client.callTool({ name: 'get_usage', arguments: { period: 'today', by: 'project' } }) + expect(JSON.stringify(res)).not.toContain('real-repo') + expect(JSON.stringify(res)).toMatch(/project-[0-9a-f]{6}/) + expect(res.isError).toBeFalsy() + }) + it('get_usage reveals names when opted in', async () => { + const client = await connect(async () => fakePayload()) + const res = await client.callTool({ name: 'get_usage', arguments: { period: 'today', by: 'project', include_project_names: true } }) + expect(JSON.stringify(res)).toContain('real-repo') + }) + it('empty data returns a friendly message, not a zero table', async () => { + const client = await connect(async () => fakePayload(0)) + const res = await client.callTool({ name: 'get_usage', arguments: { period: 'today' } }) + expect(String((res.content as Array<{ text: string }>)[0].text).toLowerCase()).toContain('no usage') + }) + it('aggregator failure surfaces as isError', async () => { + const client = await connect(async () => { throw new Error('boom') }) + const res = await client.callTool({ name: 'get_savings', arguments: {} }) + expect(res.isError).toBe(true) + expect(JSON.stringify(res)).toContain('boom') + }) +})