refactor: removed the token calculator logic

The information given was not accurate for the model providers
This commit is contained in:
Haileyesus
2026-04-10 16:21:01 +03:00
parent a8dab0edcf
commit cc7f652044
19 changed files with 20 additions and 561 deletions

View File

@@ -274,46 +274,6 @@ function transformMessage(sdkMessage) {
return sdkMessage;
}
/**
* Extracts token usage from SDK result messages
* @param {Object} resultMessage - SDK result message
* @returns {Object|null} Token budget object or null
*/
function extractTokenBudget(resultMessage) {
if (resultMessage.type !== 'result' || !resultMessage.modelUsage) {
return null;
}
// Get the first model's usage data
const modelKey = Object.keys(resultMessage.modelUsage)[0];
const modelData = resultMessage.modelUsage[modelKey];
if (!modelData) {
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;
// Total used = input + output + cache tokens
const totalUsed = inputTokens + outputTokens + cacheReadTokens + cacheCreationTokens;
// 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
return {
used: totalUsed,
total: contextWindow
};
}
/**
* Handles image processing for SDK queries
* Saves base64 images to temporary files and returns modified prompt with file paths
@@ -657,18 +617,6 @@ 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' }));
}
}
}
// Clean up session on completion

View File

@@ -2218,194 +2218,6 @@ app.post('/api/projects/:projectName/upload-images', authenticateToken, async (r
}
});
// Get token usage for a specific session
app.get('/api/projects/:projectName/sessions/:sessionId/token-usage', authenticateToken, async (req, res) => {
try {
const { projectName, sessionId } = req.params;
const { provider = 'claude' } = req.query;
const homeDir = os.homedir();
// Allow only safe characters in sessionId
const safeSessionId = String(sessionId).replace(/[^a-zA-Z0-9._-]/g, '');
if (!safeSessionId || safeSessionId !== String(sessionId)) {
return res.status(400).json({ error: 'Invalid sessionId' });
}
// Handle Cursor sessions - they use SQLite and don't have token usage info
if (provider === 'cursor') {
return res.json({
used: 0,
total: 0,
breakdown: { input: 0, cacheCreation: 0, cacheRead: 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') {
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'
});
}
// Handle Codex sessions
if (provider === 'codex') {
const codexSessionsDir = path.join(homeDir, '.codex', 'sessions');
// Find the session file by searching for the session ID
const findSessionFile = async (dir) => {
try {
const entries = await fsPromises.readdir(dir, { withFileTypes: true });
for (const entry of entries) {
const fullPath = path.join(dir, entry.name);
if (entry.isDirectory()) {
const found = await findSessionFile(fullPath);
if (found) return found;
} else if (entry.name.includes(safeSessionId) && entry.name.endsWith('.jsonl')) {
return fullPath;
}
}
} catch (error) {
// Skip directories we can't read
}
return null;
};
const sessionFilePath = await findSessionFile(codexSessionsDir);
if (!sessionFilePath) {
return res.status(404).json({ error: 'Codex session file not found', sessionId: safeSessionId });
}
// Read and parse the Codex JSONL file
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 totalTokens = 0;
let contextWindow = 200000; // Default for Codex/OpenAI
// Find the latest token_count event with info (scan from end)
for (let i = lines.length - 1; i >= 0; i--) {
try {
const entry = JSON.parse(lines[i]);
// Codex stores token info in event_msg with type: "token_count"
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;
}
if (tokenInfo.model_context_window) {
contextWindow = tokenInfo.model_context_window;
}
break; // Stop after finding the latest token count
}
} catch (parseError) {
// Skip lines that can't be parsed
continue;
}
}
return res.json({
used: totalTokens,
total: contextWindow
});
}
// Handle Claude sessions (default)
// Extract actual project path
let projectPath;
try {
projectPath = await extractProjectDirectory(projectName);
} catch (error) {
console.error('Error extracting project directory:', error);
return res.status(500).json({ error: 'Failed to determine project path' });
}
// Construct the JSONL file path
// Claude stores session files in ~/.claude/projects/[encoded-project-path]/[session-id].jsonl
// The encoding replaces any non-alphanumeric character (except -) with -
const encodedPath = projectPath.replace(/[^a-zA-Z0-9-]/g, '-');
const projectDir = path.join(homeDir, '.claude', 'projects', encodedPath);
const jsonlPath = path.join(projectDir, `${safeSessionId}.jsonl`);
// Constrain to projectDir
const rel = path.relative(path.resolve(projectDir), path.resolve(jsonlPath));
if (rel.startsWith('..') || path.isAbsolute(rel)) {
return res.status(400).json({ error: 'Invalid path' });
}
// Read and parse the JSONL file
let fileContent;
try {
fileContent = await fsPromises.readFile(jsonlPath, 'utf8');
} catch (error) {
if (error.code === 'ENOENT') {
return res.status(404).json({ error: 'Session file not found', path: jsonlPath });
}
throw error; // Re-throw other errors to be caught by outer try-catch
}
const lines = fileContent.trim().split('\n');
const parsedContextWindow = parseInt(process.env.CONTEXT_WINDOW, 10);
const contextWindow = Number.isFinite(parsedContextWindow) ? parsedContextWindow : 160000;
let inputTokens = 0;
let cacheCreationTokens = 0;
let cacheReadTokens = 0;
// Find the latest assistant message with usage data (scan from end)
for (let i = lines.length - 1; i >= 0; i--) {
try {
const entry = JSON.parse(lines[i]);
// Only count assistant messages which have usage data
if (entry.type === 'assistant' && entry.message?.usage) {
const usage = entry.message.usage;
// 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;
break; // Stop after finding the latest assistant message
}
} catch (parseError) {
// Skip lines that can't be parsed
continue;
}
}
// Calculate total context usage (excluding output_tokens, as per ccusage)
const totalUsed = inputTokens + cacheCreationTokens + cacheReadTokens;
res.json({
used: totalUsed,
total: contextWindow,
breakdown: {
input: inputTokens,
cacheCreation: cacheCreationTokens,
cacheRead: cacheReadTokens
}
});
} catch (error) {
console.error('Error reading session token usage:', error);
res.status(500).json({ error: 'Failed to read session token usage' });
}
});
// Serve React app for all other routes (excluding static files)
app.get('*', (req, res) => {
// Skip requests for static assets (files with extensions)

View File

@@ -129,8 +129,7 @@ function transformCodexEvent(event) {
case 'turn.completed':
return {
type: 'turn_complete',
usage: event.usage
type: 'turn_complete'
};
case 'turn.failed':
@@ -279,12 +278,6 @@ export async function queryCodex(command, options = {}, ws) {
error: terminalFailure
});
}
// 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: currentSessionId, provider: 'codex' }));
}
}
// Send completion event

View File

@@ -1618,7 +1618,6 @@ async function getCodexSessionMessages(sessionId, limit = null, offset = 0) {
}
const messages = [];
let tokenUsage = null;
const fileStream = fsSync.createReadStream(sessionFilePath);
const rl = readline.createInterface({
input: fileStream,
@@ -1647,17 +1646,6 @@ async function getCodexSessionMessages(sessionId, limit = null, offset = 0) {
try {
const entry = JSON.parse(line);
// Extract token usage from token_count events (keep latest)
if (entry.type === 'event_msg' && entry.payload?.type === 'token_count' && entry.payload?.info) {
const info = entry.payload.info;
if (info.total_token_usage) {
tokenUsage = {
used: info.total_token_usage.total_tokens || 0,
total: info.model_context_window || 200000
};
}
}
// Use event_msg.user_message for user-visible inputs.
if (entry.type === 'event_msg' && isVisibleCodexUserMessage(entry.payload)) {
messages.push({
@@ -1820,11 +1808,10 @@ async function getCodexSessionMessages(sessionId, limit = null, offset = 0) {
hasMore,
offset,
limit,
tokenUsage
};
}
return { messages, tokenUsage };
return { messages };
} catch (error) {
console.error(`Error reading Codex session messages for ${sessionId}:`, error);

View File

@@ -214,7 +214,6 @@ export const codexAdapter = {
const rawMessages = Array.isArray(result) ? result : (result.messages || []);
const total = Array.isArray(result) ? rawMessages.length : (result.total || 0);
const hasMore = Array.isArray(result) ? false : Boolean(result.hasMore);
const tokenUsage = result.tokenUsage || null;
const normalized = [];
for (const raw of rawMessages) {
@@ -242,7 +241,6 @@ export const codexAdapter = {
hasMore,
offset,
limit,
tokenUsage,
};
},
};

View File

@@ -53,14 +53,7 @@ export function normalizeMessage(raw, sessionId) {
}
if (raw.type === 'result') {
const msgs = [createNormalizedMessage({ sessionId, timestamp: ts, provider: PROVIDER, kind: 'stream_end' })];
if (raw.stats?.total_tokens) {
msgs.push(createNormalizedMessage({
sessionId, timestamp: ts, provider: PROVIDER,
kind: 'status', text: 'Complete', tokens: raw.stats.total_tokens, canInterrupt: false,
}));
}
return msgs;
return [createNormalizedMessage({ sessionId, timestamp: ts, provider: PROVIDER, kind: 'stream_end' })];
}
if (raw.type === 'error') {

View File

@@ -41,7 +41,7 @@
* - stream_end: (no extra fields)
* - error: content
* - complete: (no extra fields)
* - status: text, tokens?, canInterrupt?
* - status: text, canInterrupt?
* - permission_request: requestId, toolName, input, context?
* - permission_cancelled: requestId
* - session_created: newSessionId
@@ -66,7 +66,6 @@
* @property {boolean} hasMore - Whether more messages exist before the current page
* @property {number} offset - Current offset
* @property {number|null} limit - Page size used
* @property {object} [tokenUsage] - Token usage data (provider-specific)
*/
// ─── Provider Adapter Interface ──────────────────────────────────────────────

View File

@@ -546,7 +546,12 @@ class ResponseCollector {
const parsed = JSON.parse(msg);
// Only include claude-response messages with assistant type
if (parsed.type === 'claude-response' && parsed.data && parsed.data.type === 'assistant') {
assistantMessages.push(parsed.data);
const assistantMessage = { ...parsed.data };
if (assistantMessage.message?.usage) {
assistantMessage.message = { ...assistantMessage.message };
delete assistantMessage.message.usage;
}
assistantMessages.push(assistantMessage);
}
} catch (e) {
// Not JSON, skip
@@ -556,49 +561,6 @@ class ResponseCollector {
return assistantMessages;
}
/**
* Calculate total tokens from all messages
*/
getTotalTokens() {
let totalInput = 0;
let totalOutput = 0;
let totalCacheRead = 0;
let totalCacheCreation = 0;
for (const msg of this.messages) {
let data = msg;
// Parse if string
if (typeof msg === 'string') {
try {
data = JSON.parse(msg);
} catch (e) {
continue;
}
}
// Extract usage from claude-response messages
if (data && data.type === 'claude-response' && data.data) {
const msgData = data.data;
if (msgData.message && msgData.message.usage) {
const usage = msgData.message.usage;
totalInput += usage.input_tokens || 0;
totalOutput += usage.output_tokens || 0;
totalCacheRead += usage.cache_read_input_tokens || 0;
totalCacheCreation += usage.cache_creation_input_tokens || 0;
}
}
}
return {
inputTokens: totalInput,
outputTokens: totalOutput,
cacheReadTokens: totalCacheRead,
cacheCreationTokens: totalCacheCreation,
totalTokens: totalInput + totalOutput + totalCacheRead + totalCacheCreation
};
}
}
// ===============================
@@ -789,13 +751,6 @@ class ResponseCollector {
* success: true,
* sessionId: "session-123",
* messages: [...], // Assistant messages only (filtered)
* tokens: {
* inputTokens: 150,
* outputTokens: 50,
* cacheReadTokens: 0,
* cacheCreationTokens: 0,
* totalTokens: 200
* },
* projectPath: "/path/to/project",
* branch: { // Only if createBranch=true
* name: "feature/xyz",
@@ -1173,15 +1128,13 @@ router.post('/', validateExternalApiKey, async (req, res) => {
// Streaming mode: end the SSE stream
writer.end();
} else {
// Non-streaming mode: send filtered messages and token summary as JSON
// Non-streaming mode: send filtered messages as JSON
const assistantMessages = writer.getAssistantMessages();
const tokenSummary = writer.getTotalTokens();
const response = {
success: true,
sessionId: writer.getSessionId(),
messages: assistantMessages,
tokens: tokenSummary,
projectPath: finalProjectPath
};

View File

@@ -97,12 +97,6 @@ const builtInCommands = [
namespace: 'builtin',
metadata: { type: 'builtin' }
},
{
name: '/cost',
description: 'Display token usage and cost information',
namespace: 'builtin',
metadata: { type: 'builtin' }
},
{
name: '/memory',
description: 'Open CLAUDE.md memory file for editing',
@@ -209,86 +203,6 @@ Custom commands can be created in:
};
},
'/cost': async (args, context) => {
const tokenUsage = context?.tokenUsage || {};
const provider = context?.provider || 'claude';
const model =
context?.model ||
(provider === 'cursor'
? CURSOR_MODELS.DEFAULT
: provider === 'codex'
? CODEX_MODELS.DEFAULT
: CLAUDE_MODELS.DEFAULT);
const used = Number(tokenUsage.used ?? tokenUsage.totalUsed ?? tokenUsage.total_tokens ?? 0) || 0;
const total =
Number(
tokenUsage.total ??
tokenUsage.contextWindow ??
parseInt(process.env.CONTEXT_WINDOW || '160000', 10),
) || 160000;
const percentage = total > 0 ? Number(((used / total) * 100).toFixed(1)) : 0;
const inputTokensRaw =
Number(
tokenUsage.inputTokens ??
tokenUsage.input ??
tokenUsage.cumulativeInputTokens ??
tokenUsage.promptTokens ??
0,
) || 0;
const outputTokens =
Number(
tokenUsage.outputTokens ??
tokenUsage.output ??
tokenUsage.cumulativeOutputTokens ??
tokenUsage.completionTokens ??
0,
) || 0;
const cacheTokens =
Number(
tokenUsage.cacheReadTokens ??
tokenUsage.cacheCreationTokens ??
tokenUsage.cacheTokens ??
tokenUsage.cachedTokens ??
0,
) || 0;
// If we only have total used tokens, treat them as input for display/estimation.
const inputTokens =
inputTokensRaw > 0 || outputTokens > 0 || cacheTokens > 0 ? inputTokensRaw + cacheTokens : 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',
data: {
tokenUsage: {
used,
total,
percentage,
},
cost: {
input: inputCost.toFixed(4),
output: outputCost.toFixed(4),
total: totalCost.toFixed(4),
},
model,
},
};
},
'/status': async (args, context) => {
// Read version from package.json
const packageJsonPath = path.join(path.dirname(__dirname), '..', 'package.json');