diff --git a/server/claude-sdk.js b/server/claude-sdk.js index da15142d..5d8a27d7 100644 --- a/server/claude-sdk.js +++ b/server/claude-sdk.js @@ -285,43 +285,68 @@ function transformMessage(sdkMessage) { return sdkMessage; } +function readNumber(value) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : 0; +} + /** - * Extracts token usage from SDK result messages - * @param {Object} resultMessage - SDK result message + * Extracts token usage from SDK messages. + * Prefers per-step `message.usage` (Claude message payload), then falls back + * to result-level usage/modelUsage for compatibility across SDK versions. + * @param {Object} sdkMessage - SDK stream message * @returns {Object|null} Token budget object or null */ -function extractTokenBudget(resultMessage) { - if (resultMessage.type !== 'result' || !resultMessage.modelUsage) { +function extractTokenBudget(sdkMessage) { + if (!sdkMessage || typeof sdkMessage !== 'object') { return null; } - // Get the first model's usage data - const modelKey = Object.keys(resultMessage.modelUsage)[0]; - const modelData = resultMessage.modelUsage[modelKey]; + const messageUsage = sdkMessage.message?.usage || sdkMessage.usage; + if (messageUsage && typeof messageUsage === 'object') { + const inputTokens = readNumber(messageUsage.input_tokens ?? messageUsage.inputTokens); + const outputTokens = readNumber(messageUsage.output_tokens ?? messageUsage.outputTokens); + const totalUsed = inputTokens + outputTokens; + const contextWindow = parseInt(process.env.CONTEXT_WINDOW, 10) || 160000; - if (!modelData) { + return { + used: totalUsed, + total: contextWindow, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens, + }, + }; + } + + if (!sdkMessage.modelUsage || typeof sdkMessage.modelUsage !== 'object') { return null; } - // Use cumulative tokens if available (tracks total for the session) - // Otherwise fall back to per-request tokens - const inputTokens = modelData.cumulativeInputTokens || modelData.inputTokens || 0; - const outputTokens = modelData.cumulativeOutputTokens || modelData.outputTokens || 0; - const cacheReadTokens = modelData.cumulativeCacheReadInputTokens || modelData.cacheReadInputTokens || 0; - const cacheCreationTokens = modelData.cumulativeCacheCreationInputTokens || modelData.cacheCreationInputTokens || 0; + // Fallback for older SDK messages with only modelUsage + const modelKey = Object.keys(sdkMessage.modelUsage)[0]; + const modelData = sdkMessage.modelUsage[modelKey]; - // Total used = input + output + cache tokens - const totalUsed = inputTokens + outputTokens + cacheReadTokens + cacheCreationTokens; + if (!modelData || typeof modelData !== 'object') { + return null; + } - // Use configured context window budget from environment (default 160000) - // This is the user's budget limit, not the model's context window - const contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 160000; - - // Token calc logged via token-budget WS event + const inputTokens = readNumber(modelData.cumulativeInputTokens ?? modelData.inputTokens); + const outputTokens = readNumber(modelData.cumulativeOutputTokens ?? modelData.outputTokens); + const totalUsed = inputTokens + outputTokens; + const contextWindow = parseInt(process.env.CONTEXT_WINDOW, 10) || 160000; return { used: totalUsed, - total: contextWindow + total: contextWindow, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens, + }, }; } @@ -684,16 +709,10 @@ async function queryClaudeSDK(command, options = {}, ws) { ws.send(msg); } - // Extract and send token budget updates from result messages - if (message.type === 'result') { - const models = Object.keys(message.modelUsage || {}); - if (models.length > 0) { - // Model info available in result message - } - const tokenBudgetData = extractTokenBudget(message); - if (tokenBudgetData) { - ws.send(createNormalizedMessage({ kind: 'status', text: 'token_budget', tokenBudget: tokenBudgetData, sessionId: capturedSessionId || sessionId || null, provider: 'claude' })); - } + // Extract and send token budget updates from assistant/result usage payloads + const tokenBudgetData = extractTokenBudget(message); + if (tokenBudgetData) { + ws.send(createNormalizedMessage({ kind: 'status', text: 'token_budget', tokenBudget: tokenBudgetData, sessionId: capturedSessionId || sessionId || null, provider: 'claude' })); } } diff --git a/server/gemini-response-handler.js b/server/gemini-response-handler.js index b0a17485..f57ad072 100644 --- a/server/gemini-response-handler.js +++ b/server/gemini-response-handler.js @@ -1,5 +1,29 @@ // Gemini Response Handler - JSON Stream processing import { sessionsService } from './modules/providers/services/sessions.service.js'; +import { createNormalizedMessage } from './shared/utils.js'; + +function buildGeminiTokenBudget(tokens) { + if (!tokens || typeof tokens !== 'object') { + return null; + } + + const inputTokens = Number(tokens.input || 0); + const outputTokens = Number(tokens.output || 0); + const used = Number(tokens.total || inputTokens + outputTokens || 0); + if (used <= 0) { + return null; + } + + return { + used, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens, + }, + }; +} class GeminiResponseHandler { constructor(ws, options = {}) { @@ -60,6 +84,17 @@ class GeminiResponseHandler { for (const msg of normalized) { this.ws.send(msg); } + + const tokenBudget = buildGeminiTokenBudget(event.tokens); + if (tokenBudget) { + this.ws.send(createNormalizedMessage({ + kind: 'status', + text: 'token_budget', + tokenBudget, + sessionId: sid, + provider: 'gemini', + })); + } } forceFlush() { diff --git a/server/index.js b/server/index.js index 8ef49461..a986148e 100755 --- a/server/index.js +++ b/server/index.js @@ -10,8 +10,9 @@ import { spawn } from 'child_process'; import express from 'express'; import cors from 'cors'; import mime from 'mime-types'; +import Database from 'better-sqlite3'; -import { AppError, WORKSPACES_ROOT, validateWorkspacePath } from '@/shared/utils.js'; +import { AppError, WORKSPACES_ROOT, getOpenCodeDatabasePath, validateWorkspacePath } from '@/shared/utils.js'; import { closeSessionsWatcher, initializeSessionsWatcher } from '@/modules/providers/index.js'; import { createWebSocketServer } from '@/modules/websocket/index.js'; @@ -72,7 +73,7 @@ import geminiRoutes from './routes/gemini.js'; import pluginsRoutes from './routes/plugins.js'; import providerRoutes from './modules/providers/provider.routes.js'; import { startEnabledPluginServers, stopAllPlugins, getPluginPort } from './utils/plugin-process-manager.js'; -import { initializeDatabase, projectsDb } from './modules/database/index.js'; +import { initializeDatabase, projectsDb, sessionsDb } from './modules/database/index.js'; import { configureWebPush } from './services/vapid-keys.js'; import { validateApiKey, authenticateToken, authenticateWebSocket } from './middleware/auth.js'; import { IS_PLATFORM } from './constants/config.js'; @@ -1141,33 +1142,127 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate return res.json({ used: 0, total: 0, - breakdown: { input: 0, cacheCreation: 0, cacheRead: 0 }, + inputTokens: 0, + outputTokens: 0, + breakdown: { input: 0, output: 0 }, unsupported: true, message: 'Token usage tracking not available for Cursor sessions' }); } - // Handle Gemini sessions - they are raw logs in our current setup if (provider === 'gemini') { + const session = sessionsDb.getSessionById(safeSessionId); + const sessionFilePath = session?.jsonl_path; + if (!sessionFilePath) { + return res.json({ + used: 0, + inputTokens: 0, + outputTokens: 0, + breakdown: { input: 0, output: 0 }, + unsupported: true, + message: 'Token usage tracking not available for this Gemini session' + }); + } + + let fileContent; + try { + fileContent = await fsPromises.readFile(sessionFilePath, 'utf8'); + } catch (error) { + if (error.code === 'ENOENT') { + return res.status(404).json({ error: 'Session file not found', path: sessionFilePath }); + } + throw error; + } + + const lines = fileContent.trim().split('\n'); + let inputTokens = 0; + let outputTokens = 0; + let totalTokens = 0; + + for (let i = lines.length - 1; i >= 0; i--) { + try { + const entry = JSON.parse(lines[i]); + if (!entry.tokens || typeof entry.tokens !== 'object') { + continue; + } + + inputTokens = Number(entry.tokens.input || 0); + outputTokens = Number(entry.tokens.output || 0); + totalTokens = Number(entry.tokens.total || inputTokens + outputTokens || 0); + break; + } catch { + continue; + } + } + return res.json({ - used: 0, - total: 0, - breakdown: { input: 0, cacheCreation: 0, cacheRead: 0 }, - unsupported: true, - message: 'Token usage tracking not available for Gemini sessions' + used: totalTokens, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens + } }); } - // OpenCode token totals are surfaced through provider history reads. - // This legacy endpoint only knows file-backed session formats. if (provider === 'opencode') { - return res.json({ - used: 0, - total: 0, - breakdown: { input: 0, cacheCreation: 0, cacheRead: 0 }, - unsupported: true, - message: 'Token usage tracking is available in OpenCode session history, not this legacy endpoint' - }); + const dbPath = getOpenCodeDatabasePath(); + if (!fs.existsSync(dbPath)) { + return res.status(404).json({ error: 'OpenCode database not found' }); + } + + const db = new Database(dbPath, { readonly: true, fileMustExist: true }); + try { + const columns = db.prepare('PRAGMA table_info(session)').all(); + const columnNames = new Set(columns.map((column) => column.name)); + const requiredColumns = ['tokens_input', 'tokens_output', 'tokens_reasoning', 'tokens_cache_read', 'tokens_cache_write']; + if (!requiredColumns.every((column) => columnNames.has(column))) { + return res.json({ + used: 0, + inputTokens: 0, + outputTokens: 0, + breakdown: { input: 0, output: 0 }, + unsupported: true, + message: 'Token usage tracking is not available in this OpenCode database schema' + }); + } + + const row = db.prepare(` + SELECT + tokens_input AS inputTokens, + tokens_output AS outputTokens, + tokens_reasoning AS reasoningTokens, + tokens_cache_read AS cacheReadTokens, + tokens_cache_write AS cacheWriteTokens + FROM session + WHERE id = ? + `).get(safeSessionId); + + if (!row) { + return res.status(404).json({ error: 'OpenCode session not found', sessionId: safeSessionId }); + } + + const inputTokens = Number(row.inputTokens || 0) + Number(row.cacheReadTokens || 0); + const outputTokens = Number(row.outputTokens || 0); + const totalUsed = Number(row.inputTokens || 0) + + outputTokens + + Number(row.reasoningTokens || 0) + + Number(row.cacheReadTokens || 0) + + Number(row.cacheWriteTokens || 0); + + return res.json({ + used: totalUsed, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens + } + }); + } finally { + db.close(); + } } // Handle Codex sessions @@ -1210,6 +1305,8 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate throw error; } const lines = fileContent.trim().split('\n'); + let inputTokens = 0; + let outputTokens = 0; let totalTokens = 0; let contextWindow = 200000; // Default for Codex/OpenAI @@ -1222,7 +1319,9 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate if (entry.type === 'event_msg' && entry.payload?.type === 'token_count' && entry.payload?.info) { const tokenInfo = entry.payload.info; if (tokenInfo.total_token_usage) { - totalTokens = tokenInfo.total_token_usage.total_tokens || 0; + inputTokens = tokenInfo.total_token_usage.input_tokens || 0; + outputTokens = tokenInfo.total_token_usage.output_tokens || 0; + totalTokens = tokenInfo.total_token_usage.total_tokens || inputTokens + outputTokens; } if (tokenInfo.model_context_window) { contextWindow = tokenInfo.model_context_window; @@ -1237,7 +1336,13 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate return res.json({ used: totalTokens, - total: contextWindow + total: contextWindow, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens + } }); } @@ -1280,8 +1385,7 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate const parsedContextWindow = parseInt(process.env.CONTEXT_WINDOW, 10); const contextWindow = Number.isFinite(parsedContextWindow) ? parsedContextWindow : 160000; let inputTokens = 0; - let cacheCreationTokens = 0; - let cacheReadTokens = 0; + let outputTokens = 0; // Find the latest assistant message with usage data (scan from end) for (let i = lines.length - 1; i >= 0; i--) { @@ -1294,8 +1398,7 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate // Use token counts from latest assistant message only inputTokens = usage.input_tokens || 0; - cacheCreationTokens = usage.cache_creation_input_tokens || 0; - cacheReadTokens = usage.cache_read_input_tokens || 0; + outputTokens = usage.output_tokens || 0; break; // Stop after finding the latest assistant message } @@ -1305,16 +1408,16 @@ app.get('/api/projects/:projectId/sessions/:sessionId/token-usage', authenticate } } - // Calculate total context usage (excluding output_tokens, as per ccusage) - const totalUsed = inputTokens + cacheCreationTokens + cacheReadTokens; + const totalUsed = inputTokens + outputTokens; res.json({ used: totalUsed, total: contextWindow, + inputTokens, + outputTokens, breakdown: { input: inputTokens, - cacheCreation: cacheCreationTokens, - cacheRead: cacheReadTokens + output: outputTokens } }); } catch (error) { diff --git a/server/modules/providers/list/gemini/gemini-sessions.provider.ts b/server/modules/providers/list/gemini/gemini-sessions.provider.ts index 98de12c7..32781c9d 100644 --- a/server/modules/providers/list/gemini/gemini-sessions.provider.ts +++ b/server/modules/providers/list/gemini/gemini-sessions.provider.ts @@ -88,22 +88,15 @@ function buildGeminiTokenUsage(tokens: unknown): AnyRecord | undefined { const record = tokens as AnyRecord; const input = Number(record.input || 0); const output = Number(record.output || 0); - const cached = Number(record.cached || 0); - const thoughts = Number(record.thoughts || 0); - const tool = Number(record.tool || 0); - - const totalFromFields = input + output + cached + thoughts + tool; - const total = Number(record.total || totalFromFields || 0); + const total = Number(record.total || input + output || 0); return { used: total, - total: total, + inputTokens: input, + outputTokens: output, breakdown: { input, output, - cached, - thoughts, - tool, }, }; } diff --git a/server/modules/providers/list/opencode/opencode-sessions.provider.ts b/server/modules/providers/list/opencode/opencode-sessions.provider.ts index f5b8ba84..5b7bcce7 100644 --- a/server/modules/providers/list/opencode/opencode-sessions.provider.ts +++ b/server/modules/providers/list/opencode/opencode-sessions.provider.ts @@ -28,9 +28,9 @@ type OpenCodeHistoryRow = { type OpenCodeTokenTotals = { inputTokens: number; outputTokens: number; - cacheReadTokens: number; - cacheCreationTokens: number; reasoningTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; }; const openOpenCodeDatabase = (): Database.Database | null => { @@ -106,11 +106,13 @@ const buildTokenUsage = (totals: OpenCodeTokenTotals | undefined): AnyRecord | u } const inputTokens = totals.inputTokens; + const displayInputTokens = inputTokens + totals.cacheReadTokens; const outputTokens = totals.outputTokens; - const cacheReadTokens = totals.cacheReadTokens; - const cacheCreationTokens = totals.cacheCreationTokens; - const reasoningTokens = totals.reasoningTokens; - const used = inputTokens + outputTokens + cacheReadTokens + cacheCreationTokens + reasoningTokens; + const used = inputTokens + + outputTokens + + totals.reasoningTokens + + totals.cacheReadTokens + + totals.cacheWriteTokens; if (used <= 0) { return undefined; @@ -118,14 +120,50 @@ const buildTokenUsage = (totals: OpenCodeTokenTotals | undefined): AnyRecord | u return { used, - total: used, - inputTokens, + inputTokens: displayInputTokens, outputTokens, - cacheReadTokens, - cacheCreationTokens, + breakdown: { + input: displayInputTokens, + output: outputTokens, + }, }; }; +const readOpenCodeSessionColumnTokenUsage = ( + db: Database.Database, + sessionId: string, +): AnyRecord | undefined => { + const columns = db.prepare('PRAGMA table_info(session)').all() as { name: string }[]; + const columnNames = new Set(columns.map((column) => column.name)); + const requiredColumns = ['tokens_input', 'tokens_output', 'tokens_reasoning', 'tokens_cache_read', 'tokens_cache_write']; + if (!requiredColumns.every((column) => columnNames.has(column))) { + return undefined; + } + + const row = db.prepare(` + SELECT + tokens_input AS inputTokens, + tokens_output AS outputTokens, + tokens_reasoning AS reasoningTokens, + tokens_cache_read AS cacheReadTokens, + tokens_cache_write AS cacheWriteTokens + FROM session + WHERE id = ? + `).get(sessionId) as OpenCodeTokenTotals | undefined; + + if (!row) { + return undefined; + } + + return buildTokenUsage({ + inputTokens: Number(row.inputTokens ?? 0), + outputTokens: Number(row.outputTokens ?? 0), + reasoningTokens: Number(row.reasoningTokens ?? 0), + cacheReadTokens: Number(row.cacheReadTokens ?? 0), + cacheWriteTokens: Number(row.cacheWriteTokens ?? 0), + }); +}; + /** * OpenCode stores per-message token counts on assistant `message.data` objects * (see MessageV2.Assistant). Older DBs also had session-level counters; this @@ -135,13 +173,18 @@ const aggregateOpenCodeSessionTokenUsage = ( db: Database.Database, sessionId: string, ): AnyRecord | undefined => { + const sessionColumnUsage = readOpenCodeSessionColumnTokenUsage(db, sessionId); + if (sessionColumnUsage) { + return sessionColumnUsage; + } + const rows = db.prepare('SELECT data FROM message WHERE session_id = ?').all(sessionId) as { data: string }[]; let inputTokens = 0; let outputTokens = 0; - let cacheReadTokens = 0; - let cacheCreationTokens = 0; let reasoningTokens = 0; + let cacheReadTokens = 0; + let cacheWriteTokens = 0; for (const row of rows) { const info = readJsonRecord(row.data); @@ -159,15 +202,15 @@ const aggregateOpenCodeSessionTokenUsage = ( reasoningTokens += Number(tokens.reasoning ?? 0); const cache = readObjectRecord(tokens.cache); cacheReadTokens += Number(cache?.read ?? 0); - cacheCreationTokens += Number(cache?.write ?? 0); + cacheWriteTokens += Number(cache?.write ?? 0); } return buildTokenUsage({ inputTokens, outputTokens, - cacheReadTokens, - cacheCreationTokens, reasoningTokens, + cacheReadTokens, + cacheWriteTokens, }); }; diff --git a/server/modules/providers/tests/opencode-sessions.test.ts b/server/modules/providers/tests/opencode-sessions.test.ts index 139a28aa..d5b65e4e 100644 --- a/server/modules/providers/tests/opencode-sessions.test.ts +++ b/server/modules/providers/tests/opencode-sessions.test.ts @@ -85,6 +85,12 @@ const createOpenCodeDatabase = async (homeDir: string, workspacePath: string): P path TEXT, agent TEXT, model TEXT, + cost REAL NOT NULL DEFAULT 0, + tokens_input INTEGER NOT NULL DEFAULT 0, + tokens_output INTEGER NOT NULL DEFAULT 0, + tokens_reasoning INTEGER NOT NULL DEFAULT 0, + tokens_cache_read INTEGER NOT NULL DEFAULT 0, + tokens_cache_write INTEGER NOT NULL DEFAULT 0, FOREIGN KEY (project_id) REFERENCES project(id) ON DELETE CASCADE ); @@ -124,9 +130,10 @@ const createOpenCodeDatabase = async (homeDir: string, workspacePath: string): P ); db.prepare(` INSERT INTO session ( - id, project_id, slug, directory, title, version, time_created, time_updated, time_archived + id, project_id, slug, directory, title, version, time_created, time_updated, time_archived, + tokens_input, tokens_output, tokens_reasoning, tokens_cache_read, tokens_cache_write ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `).run( 'open-session-1', 'project-1', @@ -137,6 +144,11 @@ const createOpenCodeDatabase = async (homeDir: string, workspacePath: string): P 1_700_000_000_000, 1_700_000_004_000, null, + 10, + 20, + 7, + 3, + 2, ); const userMessageData = JSON.stringify({ @@ -302,12 +314,13 @@ test('OpenCode sessions provider reads sqlite history and token usage', { concur assert.equal(history.messages[3]?.kind, 'tool_use'); assert.deepEqual(history.messages[3]?.toolResult, { content: 'ok', isError: false }); assert.deepEqual(history.tokenUsage, { - used: 35, - total: 35, - inputTokens: 10, + used: 42, + inputTokens: 13, outputTokens: 20, - cacheReadTokens: 3, - cacheCreationTokens: 2, + breakdown: { + input: 13, + output: 20, + }, }); const paged = await provider.fetchHistory('open-session-1', { limit: 2, offset: 0 }); diff --git a/server/openai-codex.js b/server/openai-codex.js index 95913586..8e14fcdf 100644 --- a/server/openai-codex.js +++ b/server/openai-codex.js @@ -23,6 +23,34 @@ import { createNormalizedMessage } from './shared/utils.js'; // Track active sessions const activeCodexSessions = new Map(); +function readUsageNumber(value) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : 0; +} + +function extractCodexTokenBudget(event) { + const info = event?.info || event?.payload?.info || event?.usage?.info; + const usage = info?.total_token_usage || event?.usage?.total_token_usage || event?.usage; + if (!usage || typeof usage !== 'object') { + return null; + } + + const inputTokens = readUsageNumber(usage.input_tokens); + const outputTokens = readUsageNumber(usage.output_tokens); + const used = readUsageNumber(usage.total_tokens) || inputTokens + outputTokens; + + return { + used, + total: readUsageNumber(info?.model_context_window || event?.usage?.model_context_window) || 200000, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens, + }, + }; +} + /** * Transform Codex SDK event to WebSocket message format * @param {object} event - SDK event @@ -316,9 +344,11 @@ export async function queryCodex(command, options = {}, ws) { } // Extract and send token usage if available (normalized to match Claude format) - if (event.type === 'turn.completed' && event.usage) { - const totalTokens = (event.usage.input_tokens || 0) + (event.usage.output_tokens || 0); - sendMessage(ws, createNormalizedMessage({ kind: 'status', text: 'token_budget', tokenBudget: { used: totalTokens, total: 200000 }, sessionId: capturedSessionId || sessionId || null, provider: 'codex' })); + if (event.type === 'turn.completed') { + const tokenBudget = extractCodexTokenBudget(event); + if (tokenBudget) { + sendMessage(ws, createNormalizedMessage({ kind: 'status', text: 'token_budget', tokenBudget, sessionId: capturedSessionId || sessionId || null, provider: 'codex' })); + } } } diff --git a/server/opencode-cli.js b/server/opencode-cli.js index c386d475..119c5ee4 100644 --- a/server/opencode-cli.js +++ b/server/opencode-cli.js @@ -1,12 +1,14 @@ import { spawn } from 'child_process'; +import fsSync from 'node:fs'; import crossSpawn from 'cross-spawn'; +import Database from 'better-sqlite3'; import { sessionsService } from './modules/providers/services/sessions.service.js'; import { providerAuthService } from './modules/providers/services/provider-auth.service.js'; import { providerModelsService } from './modules/providers/services/provider-models.service.js'; import { notifyRunFailed, notifyRunStopped } from './services/notification-orchestrator.js'; -import { createNormalizedMessage } from './shared/utils.js'; +import { createNormalizedMessage, getOpenCodeDatabasePath } from './shared/utils.js'; const spawnFunction = process.platform === 'win32' ? crossSpawn : spawn; @@ -20,6 +22,61 @@ function readOpenCodeSessionId(event) { return event.sessionID || event.sessionId || null; } +function readOpenCodeTokenUsage(sessionId) { + const dbPath = getOpenCodeDatabasePath(); + if (!sessionId || !fsSync.existsSync(dbPath)) { + return null; + } + + const db = new Database(dbPath, { readonly: true, fileMustExist: true }); + try { + const columns = db.prepare('PRAGMA table_info(session)').all(); + const columnNames = new Set(columns.map((column) => column.name)); + const requiredColumns = ['tokens_input', 'tokens_output', 'tokens_reasoning', 'tokens_cache_read', 'tokens_cache_write']; + if (!requiredColumns.every((column) => columnNames.has(column))) { + return null; + } + + const row = db.prepare(` + SELECT + tokens_input AS inputTokens, + tokens_output AS outputTokens, + tokens_reasoning AS reasoningTokens, + tokens_cache_read AS cacheReadTokens, + tokens_cache_write AS cacheWriteTokens + FROM session + WHERE id = ? + `).get(sessionId); + + if (!row) { + return null; + } + + const inputTokens = Number(row.inputTokens || 0) + Number(row.cacheReadTokens || 0); + const outputTokens = Number(row.outputTokens || 0); + const used = Number(row.inputTokens || 0) + + outputTokens + + Number(row.reasoningTokens || 0) + + Number(row.cacheReadTokens || 0) + + Number(row.cacheWriteTokens || 0); + if (used <= 0) { + return null; + } + + return { + used, + inputTokens, + outputTokens, + breakdown: { + input: inputTokens, + output: outputTokens, + }, + }; + } finally { + db.close(); + } +} + async function spawnOpenCode(command, options = {}, ws) { return new Promise((resolve, reject) => { const { sessionId, projectPath, cwd, model, sessionSummary } = options; @@ -183,6 +240,17 @@ async function spawnOpenCode(command, options = {}, ws) { stdoutLineBuffer = ''; } + const tokenBudget = readOpenCodeTokenUsage(finalSessionId); + if (tokenBudget) { + ws.send(createNormalizedMessage({ + kind: 'status', + text: 'token_budget', + tokenBudget, + sessionId: finalSessionId, + provider: 'opencode', + })); + } + ws.send(createNormalizedMessage({ kind: 'complete', exitCode: code, diff --git a/server/routes/commands.js b/server/routes/commands.js index fd197565..82a0e7e0 100644 --- a/server/routes/commands.js +++ b/server/routes/commands.js @@ -174,7 +174,7 @@ const builtInCommands = [ }, { name: "/cost", - description: "Display token usage and cost information", + description: "Display token usage information", namespace: "builtin", metadata: { type: "builtin" }, }, @@ -258,7 +258,7 @@ Custom commands can be created in: const catalog = (await providerModelsService.getProviderModels(provider)).models; const model = await resolveCommandModel(provider, catalog, context?.sessionId); - const used = + const reportedUsed = Number( tokenUsage.used ?? tokenUsage.totalUsed ?? tokenUsage.total_tokens ?? 0, ) || 0; @@ -266,16 +266,15 @@ Custom commands can be created in: Number( tokenUsage.total ?? tokenUsage.contextWindow ?? - parseInt(process.env.CONTEXT_WINDOW || "160000", 10), - ) || 160000; - const percentage = - total > 0 ? Number(((used / total) * 100).toFixed(1)) : 0; - + 0, + ) || 0; const inputTokensRaw = Number( tokenUsage.inputTokens ?? tokenUsage.input ?? + tokenUsage.input_tokens ?? tokenUsage.cumulativeInputTokens ?? + tokenUsage.breakdown?.input ?? tokenUsage.promptTokens ?? 0, ) || 0; @@ -283,37 +282,21 @@ Custom commands can be created in: Number( tokenUsage.outputTokens ?? tokenUsage.output ?? + tokenUsage.output_tokens ?? tokenUsage.cumulativeOutputTokens ?? + tokenUsage.breakdown?.output ?? tokenUsage.completionTokens ?? 0, ) || 0; - const cacheTokens = - Number( - tokenUsage.cacheReadTokens ?? - tokenUsage.cacheCreationTokens ?? - tokenUsage.cacheTokens ?? - tokenUsage.cachedTokens ?? - 0, - ) || 0; + const hasTokenBreakdown = inputTokensRaw > 0 || outputTokens > 0; + const used = reportedUsed || inputTokensRaw + outputTokens; - // If we only have total used tokens, treat them as input for display/estimation. + // If we only have total used tokens, keep the list populated without guessing output. const inputTokens = - inputTokensRaw > 0 || outputTokens > 0 || cacheTokens > 0 - ? inputTokensRaw + cacheTokens + hasTokenBreakdown + ? inputTokensRaw : used; - // Rough default rates by provider (USD / 1M tokens). - const pricingByProvider = { - claude: { input: 3, output: 15 }, - cursor: { input: 3, output: 15 }, - codex: { input: 1.5, output: 6 }, - }; - const rates = pricingByProvider[provider] || pricingByProvider.claude; - - const inputCost = (inputTokens / 1_000_000) * rates.input; - const outputCost = (outputTokens / 1_000_000) * rates.output; - const totalCost = inputCost + outputCost; - return { type: "builtin", action: "cost", @@ -321,17 +304,10 @@ Custom commands can be created in: tokenUsage: { used, total, - percentage, }, tokenBreakdown: { input: inputTokens, output: outputTokens, - cache: cacheTokens, - }, - cost: { - input: inputCost.toFixed(4), - output: outputCost.toFixed(4), - total: totalCost.toFixed(4), }, provider, model, diff --git a/src/components/chat/hooks/useChatComposerState.ts b/src/components/chat/hooks/useChatComposerState.ts index 49b610cc..6d3b93df 100644 --- a/src/components/chat/hooks/useChatComposerState.ts +++ b/src/components/chat/hooks/useChatComposerState.ts @@ -97,17 +97,10 @@ export type CostCommandData = { tokenUsage?: { used?: number; total?: number; - percentage?: number; - }; - cost?: { - input?: string; - output?: string; - total?: string; }; tokenBreakdown?: { input?: number; output?: number; - cache?: number; }; provider?: string; model?: string; diff --git a/src/components/chat/hooks/useChatSessionState.ts b/src/components/chat/hooks/useChatSessionState.ts index b96e6646..20f42551 100644 --- a/src/components/chat/hooks/useChatSessionState.ts +++ b/src/components/chat/hooks/useChatSessionState.ts @@ -624,19 +624,23 @@ export function useChatSessionState({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [chatMessages.length, isLoadingSessionMessages, searchTarget]); - // Token usage fetch for Claude + // Initial token usage fetch for providers with file-backed usage data. useEffect(() => { if (!selectedProject || !selectedSession?.id) { setTokenBudget(null); return; } const sessionProvider = selectedSession.__provider || 'claude'; - if (sessionProvider !== 'claude') return; + if (sessionProvider !== 'claude' && sessionProvider !== 'codex' && sessionProvider !== 'gemini' && sessionProvider !== 'opencode') { + setTokenBudget(null); + return; + } const fetchInitialTokenUsage = async () => { try { // Token usage endpoint is now keyed by the DB projectId. - const url = `/api/projects/${selectedProject.projectId}/sessions/${selectedSession.id}/token-usage`; + const params = new URLSearchParams({ provider: sessionProvider }); + const url = `/api/projects/${selectedProject.projectId}/sessions/${selectedSession.id}/token-usage?${params.toString()}`; const response = await authenticatedFetch(url); if (response.ok) { setTokenBudget(await response.json()); diff --git a/src/components/chat/view/subcomponents/ChatComposer.tsx b/src/components/chat/view/subcomponents/ChatComposer.tsx index 2701e7b1..25df7ecf 100644 --- a/src/components/chat/view/subcomponents/ChatComposer.tsx +++ b/src/components/chat/view/subcomponents/ChatComposer.tsx @@ -18,7 +18,7 @@ import ClaudeStatus from './ClaudeStatus'; import ImageAttachment from './ImageAttachment'; import PermissionRequestsBanner from './PermissionRequestsBanner'; import ThinkingModeSelector from './ThinkingModeSelector'; -import TokenUsagePie from './TokenUsagePie'; +import TokenUsageSummary from './TokenUsageSummary'; import { PromptInput, PromptInputHeader, @@ -60,7 +60,7 @@ interface ChatComposerProps { onModeSwitch: () => void; thinkingMode: string; setThinkingMode: Dispatch>; - tokenBudget: { used?: number; total?: number } | null; + tokenBudget: Record | null; slashCommandsCount: number; onToggleCommandMenu: () => void; hasInput: boolean; @@ -361,7 +361,7 @@ export default function ChatComposer({ {}} className="" /> )} - + = { const FALLBACK_COMMANDS: CommandEntry[] = [ { name: '/models', description: 'Browse available models for the active provider.' }, - { name: '/cost', description: 'Review context usage and estimated token spend.' }, + { name: '/cost', description: 'Review token usage for the active session.' }, { name: '/status', description: 'Inspect runtime, version, provider, and environment status.' }, { name: '/memory', description: 'Open the project CLAUDE.md memory file.' }, { name: '/config', description: 'Open settings and configuration.' }, @@ -99,13 +97,6 @@ const getProviderLabel = (provider: string | undefined, fallback = 'Unknown') => return PROVIDER_LABELS[provider] || provider; }; -const clampPercentage = (value: number) => { - if (!Number.isFinite(value)) { - return 0; - } - return Math.max(0, Math.min(100, value)); -}; - const formatNumber = (value: number) => { if (!Number.isFinite(value)) { return '0'; @@ -113,11 +104,6 @@ const formatNumber = (value: number) => { return value.toLocaleString(); }; -const formatCurrency = (value: number | string | undefined) => { - const numeric = Number(value ?? 0); - return `$${Number.isFinite(numeric) ? numeric.toFixed(4) : '0.0000'}`; -}; - function MetricCard({ label, value, @@ -507,62 +493,52 @@ function ModelsContent({ function CostContent({ data }: { data: CostCommandData }) { const used = Number(data.tokenUsage?.used ?? 0); const total = Number(data.tokenUsage?.total ?? 0); - const percentage = clampPercentage(Number(data.tokenUsage?.percentage ?? 0)); const model = data.model || 'Unknown'; const provider = getProviderLabel(data.provider, data.provider || 'Unknown'); const inputTokens = Number(data.tokenBreakdown?.input ?? 0); const outputTokens = Number(data.tokenBreakdown?.output ?? 0); - const cacheTokens = Number(data.tokenBreakdown?.cache ?? 0); - const totalCost = Number(data.cost?.total ?? 0); + const usageRows = [ + { label: 'Total tokens used', value: formatNumber(used), icon: Activity }, + { label: 'Input tokens', value: formatNumber(inputTokens), icon: TerminalSquare }, + { label: 'Output tokens', value: formatNumber(outputTokens), icon: Coins }, + ...(total > 0 + ? [{ label: 'Context window', value: formatNumber(total), icon: Gauge }] + : []), + ]; return ( -
-
-
-
-
-

{percentage.toFixed(1)}%

-

context

+
+
+ {usageRows.map((row) => { + const Icon = row.icon; + + return ( +
+
+ + + + {row.label} +
+ {row.value}
-
-
-

- {formatNumber(used)} of {formatNumber(total)} tokens used -

+ ); + })}
-
-
- - - -
- -
- - - -
- -
-
-
-

Provider

-

{provider}

-
-
-

Model

-

{model}

-
+
+
+
+

Provider

+

{provider}

+
+
+

Model

+

{model}

-

- Cost is an estimate based on the available token counters and default provider rates. -

@@ -636,8 +612,8 @@ export default function CommandResultModal({ }, cost: { eyebrow: 'Session telemetry', - title: 'Usage & Cost', - subtitle: 'Token budget, context pressure, and estimated spend for this session.', + title: 'Token Usage', + subtitle: 'Input, output, and total token counts for this session.', icon: Coins, }, status: { diff --git a/src/components/chat/view/subcomponents/TokenUsagePie.tsx b/src/components/chat/view/subcomponents/TokenUsagePie.tsx deleted file mode 100644 index 1e7adb17..00000000 --- a/src/components/chat/view/subcomponents/TokenUsagePie.tsx +++ /dev/null @@ -1,54 +0,0 @@ -type TokenUsagePieProps = { - used: number; - total: number; -}; - -export default function TokenUsagePie({ used, total }: TokenUsagePieProps) { - // Token usage visualization component - // Only bail out on missing values or non‐positive totals; allow used===0 to render 0% - if (used == null || total == null || total <= 0) return null; - - const percentage = Math.min(100, (used / total) * 100); - const radius = 10; - const circumference = 2 * Math.PI * radius; - const offset = circumference - (percentage / 100) * circumference; - - // Color based on usage level - const getColor = () => { - if (percentage < 50) return '#3b82f6'; // blue - if (percentage < 75) return '#f59e0b'; // orange - return '#ef4444'; // red - }; - - return ( -
- - {/* Background circle */} - - {/* Progress circle */} - - - - {percentage.toFixed(1)}% - -
- ); -} \ No newline at end of file diff --git a/src/components/chat/view/subcomponents/TokenUsageSummary.tsx b/src/components/chat/view/subcomponents/TokenUsageSummary.tsx new file mode 100644 index 00000000..5d9f575b --- /dev/null +++ b/src/components/chat/view/subcomponents/TokenUsageSummary.tsx @@ -0,0 +1,53 @@ +import { ActivityIcon } from 'lucide-react'; + +type TokenUsageSummaryProps = { + usage: Record | null; +}; + +const formatTokenCount = (value: number) => { + if (!Number.isFinite(value) || value <= 0) { + return '0'; + } + + if (value >= 1_000_000) { + return `${(value / 1_000_000).toFixed(value >= 10_000_000 ? 0 : 1)}M`; + } + + if (value >= 10_000) { + return `${Math.round(value / 1_000)}K`; + } + + if (value >= 1_000) { + return `${(value / 1_000).toFixed(1)}K`; + } + + return value.toLocaleString(); +}; + +const readUsageNumber = (value: unknown) => { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : 0; +}; + +export default function TokenUsageSummary({ usage }: TokenUsageSummaryProps) { + const breakdown = + usage?.breakdown && typeof usage.breakdown === 'object' + ? usage.breakdown as Record + : null; + const inputTokens = readUsageNumber(usage?.inputTokens ?? breakdown?.input); + const outputTokens = readUsageNumber(usage?.outputTokens ?? breakdown?.output); + const usedTokens = readUsageNumber(usage?.used) || inputTokens + outputTokens; + + return ( +
+ + + + {formatTokenCount(usedTokens)} + tokens +
+ ); +}