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

@@ -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');