mirror of
https://github.com/siteboon/claudecodeui.git
synced 2026-05-30 00:05:33 +08:00
fix: refine token usage reporting
The old token UI mixed context pressure, cache counters, and dollar estimates. That made the percentage look precise even when provider data was incomplete or different. The composer and /cost view now show concrete counts instead of a pie percentage. Token payloads now share a smaller shape: used, inputTokens, outputTokens, and breakdown. Claude uses per-step usage where available and Codex reads total_token_usage events. Gemini reads its tokens object without inventing a context window. OpenCode reads opencode.db session totals and includes all token columns in used. The /cost backend no longer returns cache display fields or input/output dollar estimates. This avoids derived values that look reliable but are not comparable across providers. Verification: npm run typecheck; targeted eslint; OpenCode session provider test.
This commit is contained in:
@@ -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' }));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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() {
|
||||
|
||||
159
server/index.js
159
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) {
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
};
|
||||
|
||||
|
||||
@@ -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 });
|
||||
|
||||
@@ -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' }));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user