mirror of
https://github.com/Ark0N/Codeman.git
synced 2026-10-08 08:29:42 +02:00
fix: add validation to prevent token counter corruption
- Add max 500k tokens per session limit (Claude's context is ~200k) - Add max 100k tokens per update to prevent sudden jumps from parsing errors - Reject "M" suffix values > 0.5M to prevent false matches in output text - Add validation to restoreTokens() to reject corrupted values from state.json - Add validation to addToGlobalStats() for negative and absurd values - Fix daily sessions count: track unique session IDs instead of incrementing on every 5-minute recording interval - Pass sessionId to recordDailyUsage() for accurate session counting Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
+51
-4
@@ -626,6 +626,18 @@ export class Session extends EventEmitter {
|
|||||||
* Called when recovering sessions after server restart.
|
* Called when recovering sessions after server restart.
|
||||||
*/
|
*/
|
||||||
restoreTokens(inputTokens: number, outputTokens: number, totalCost: number): void {
|
restoreTokens(inputTokens: number, outputTokens: number, totalCost: number): void {
|
||||||
|
// Sanity check: reject absurdly large values (max 500k tokens per session)
|
||||||
|
const MAX_SESSION_TOKENS = 500_000;
|
||||||
|
if (inputTokens > MAX_SESSION_TOKENS || outputTokens > MAX_SESSION_TOKENS) {
|
||||||
|
console.warn(`[Session ${this.id}] Rejected absurd restored tokens: input=${inputTokens}, output=${outputTokens}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Reject negative values
|
||||||
|
if (inputTokens < 0 || outputTokens < 0 || totalCost < 0) {
|
||||||
|
console.warn(`[Session ${this.id}] Rejected negative restored tokens: input=${inputTokens}, output=${outputTokens}, cost=${totalCost}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
this._totalInputTokens = inputTokens;
|
this._totalInputTokens = inputTokens;
|
||||||
this._totalOutputTokens = outputTokens;
|
this._totalOutputTokens = outputTokens;
|
||||||
this._totalCost = totalCost;
|
this._totalCost = totalCost;
|
||||||
@@ -1298,10 +1310,19 @@ export class Session extends EventEmitter {
|
|||||||
this._textOutput.append(block.text);
|
this._textOutput.append(block.text);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Track tokens from usage
|
// Track tokens from usage (with validation)
|
||||||
if (msg.message.usage) {
|
if (msg.message.usage) {
|
||||||
this._totalInputTokens += msg.message.usage.input_tokens || 0;
|
const inputDelta = msg.message.usage.input_tokens || 0;
|
||||||
this._totalOutputTokens += msg.message.usage.output_tokens || 0;
|
const outputDelta = msg.message.usage.output_tokens || 0;
|
||||||
|
|
||||||
|
// Sanity check: max 100k tokens per message (generous limit)
|
||||||
|
const MAX_TOKENS_PER_MESSAGE = 100_000;
|
||||||
|
if (inputDelta > 0 && inputDelta <= MAX_TOKENS_PER_MESSAGE) {
|
||||||
|
this._totalInputTokens += inputDelta;
|
||||||
|
}
|
||||||
|
if (outputDelta > 0 && outputDelta <= MAX_TOKENS_PER_MESSAGE) {
|
||||||
|
this._totalOutputTokens += outputDelta;
|
||||||
|
}
|
||||||
|
|
||||||
// Check if we should auto-compact or auto-clear
|
// Check if we should auto-compact or auto-clear
|
||||||
this.checkAutoCompact();
|
this.checkAutoCompact();
|
||||||
@@ -1423,6 +1444,11 @@ export class Session extends EventEmitter {
|
|||||||
|
|
||||||
// Parse token count from Claude's status line in interactive mode
|
// Parse token count from Claude's status line in interactive mode
|
||||||
// Matches patterns like "123.4k tokens", "5234 tokens", "1.2M tokens"
|
// Matches patterns like "123.4k tokens", "5234 tokens", "1.2M tokens"
|
||||||
|
//
|
||||||
|
// SAFETY LIMITS:
|
||||||
|
// - Max tokens per session: 500k (Claude's context is ~200k)
|
||||||
|
// - Max delta per update: 100k (prevents sudden jumps from parsing errors)
|
||||||
|
// - Rejects "M" suffix values > 0.5 (500k) to prevent false matches
|
||||||
private parseTokensFromStatusLine(data: string): void {
|
private parseTokensFromStatusLine(data: string): void {
|
||||||
// Quick pre-check: skip expensive regex if "token" not present (performance optimization)
|
// Quick pre-check: skip expensive regex if "token" not present (performance optimization)
|
||||||
if (!data.includes('token')) return;
|
if (!data.includes('token')) return;
|
||||||
@@ -1442,16 +1468,37 @@ export class Session extends EventEmitter {
|
|||||||
if (suffix === 'k') {
|
if (suffix === 'k') {
|
||||||
tokenCount *= 1000;
|
tokenCount *= 1000;
|
||||||
} else if (suffix === 'm') {
|
} else if (suffix === 'm') {
|
||||||
|
// Safety: Reject M values that would result in > 500k tokens
|
||||||
|
// Claude's context window is ~200k, so anything claiming millions is likely a false match
|
||||||
|
if (tokenCount > 0.5) {
|
||||||
|
console.warn(`[Session ${this.id}] Rejected suspicious M token value: ${tokenMatch[0]} (would be ${tokenCount * 1000000} tokens)`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
tokenCount *= 1000000;
|
tokenCount *= 1000000;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Safety: Absolute maximum of 500k tokens per session
|
||||||
|
const MAX_SESSION_TOKENS = 500_000;
|
||||||
|
if (tokenCount > MAX_SESSION_TOKENS) {
|
||||||
|
console.warn(`[Session ${this.id}] Rejected token count exceeding max: ${tokenCount} > ${MAX_SESSION_TOKENS}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
// Only update if the new count is higher (tokens only increase within a session)
|
// Only update if the new count is higher (tokens only increase within a session)
|
||||||
// We use total tokens as an estimate - Claude shows combined input+output
|
// We use total tokens as an estimate - Claude shows combined input+output
|
||||||
const currentTotal = this._totalInputTokens + this._totalOutputTokens;
|
const currentTotal = this._totalInputTokens + this._totalOutputTokens;
|
||||||
if (tokenCount > currentTotal) {
|
if (tokenCount > currentTotal) {
|
||||||
|
const delta = tokenCount - currentTotal;
|
||||||
|
|
||||||
|
// Safety: Reject suspiciously large jumps (max 100k per update)
|
||||||
|
const MAX_DELTA_PER_UPDATE = 100_000;
|
||||||
|
if (delta > MAX_DELTA_PER_UPDATE) {
|
||||||
|
console.warn(`[Session ${this.id}] Rejected suspicious token jump: ${currentTotal} -> ${tokenCount} (delta: ${delta})`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
// Estimate: split roughly 60% input, 40% output (common ratio)
|
// Estimate: split roughly 60% input, 40% output (common ratio)
|
||||||
// This is an approximation since interactive mode doesn't give us the breakdown
|
// This is an approximation since interactive mode doesn't give us the breakdown
|
||||||
const delta = tokenCount - currentTotal;
|
|
||||||
this._totalInputTokens += Math.round(delta * 0.6);
|
this._totalInputTokens += Math.round(delta * 0.6);
|
||||||
this._totalOutputTokens += Math.round(delta * 0.4);
|
this._totalOutputTokens += Math.round(delta * 0.4);
|
||||||
|
|
||||||
|
|||||||
+38
-2
@@ -223,6 +223,18 @@ export class StateStore {
|
|||||||
* Call when a session is deleted to preserve its usage in lifetime stats.
|
* Call when a session is deleted to preserve its usage in lifetime stats.
|
||||||
*/
|
*/
|
||||||
addToGlobalStats(inputTokens: number, outputTokens: number, cost: number): void {
|
addToGlobalStats(inputTokens: number, outputTokens: number, cost: number): void {
|
||||||
|
// Sanity check: reject absurdly large values (max 500k tokens per session)
|
||||||
|
const MAX_SESSION_TOKENS = 500_000;
|
||||||
|
if (inputTokens > MAX_SESSION_TOKENS || outputTokens > MAX_SESSION_TOKENS) {
|
||||||
|
console.warn(`[StateStore] Rejected absurd global stats: input=${inputTokens}, output=${outputTokens}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Reject negative values
|
||||||
|
if (inputTokens < 0 || outputTokens < 0 || cost < 0) {
|
||||||
|
console.warn(`[StateStore] Rejected negative global stats: input=${inputTokens}, output=${outputTokens}, cost=${cost}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
const stats = this.getGlobalStats();
|
const stats = this.getGlobalStats();
|
||||||
stats.totalInputTokens += inputTokens;
|
stats.totalInputTokens += inputTokens;
|
||||||
stats.totalOutputTokens += outputTokens;
|
stats.totalOutputTokens += outputTokens;
|
||||||
@@ -309,18 +321,37 @@ export class StateStore {
|
|||||||
return inputCost + outputCost;
|
return inputCost + outputCost;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Track unique sessions per day for accurate session count
|
||||||
|
private dailySessionIds: Set<string> = new Set();
|
||||||
|
private dailySessionDate: string = '';
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Record token usage for today.
|
* Record token usage for today.
|
||||||
* Accumulates tokens to today's entry, creating it if needed.
|
* Accumulates tokens to today's entry, creating it if needed.
|
||||||
* @param inputTokens Input tokens to add
|
* @param inputTokens Input tokens to add
|
||||||
* @param outputTokens Output tokens to add
|
* @param outputTokens Output tokens to add
|
||||||
|
* @param sessionId Optional session ID for unique session counting
|
||||||
*/
|
*/
|
||||||
recordDailyUsage(inputTokens: number, outputTokens: number): void {
|
recordDailyUsage(inputTokens: number, outputTokens: number, sessionId?: string): void {
|
||||||
if (inputTokens <= 0 && outputTokens <= 0) return;
|
if (inputTokens <= 0 && outputTokens <= 0) return;
|
||||||
|
|
||||||
|
// Sanity check: reject absurdly large values (max 1M tokens per recording)
|
||||||
|
// Claude's context window is ~200k, so 1M per recording is already very generous
|
||||||
|
const MAX_TOKENS_PER_RECORDING = 1_000_000;
|
||||||
|
if (inputTokens > MAX_TOKENS_PER_RECORDING || outputTokens > MAX_TOKENS_PER_RECORDING) {
|
||||||
|
console.warn(`[StateStore] Rejected absurd token values: input=${inputTokens}, output=${outputTokens}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
const stats = this.getTokenStats();
|
const stats = this.getTokenStats();
|
||||||
const today = this.getTodayDateString();
|
const today = this.getTodayDateString();
|
||||||
|
|
||||||
|
// Reset daily session tracking on date change
|
||||||
|
if (this.dailySessionDate !== today) {
|
||||||
|
this.dailySessionIds.clear();
|
||||||
|
this.dailySessionDate = today;
|
||||||
|
}
|
||||||
|
|
||||||
// Find or create today's entry
|
// Find or create today's entry
|
||||||
let todayEntry = stats.daily.find(e => e.date === today);
|
let todayEntry = stats.daily.find(e => e.date === today);
|
||||||
if (!todayEntry) {
|
if (!todayEntry) {
|
||||||
@@ -341,7 +372,12 @@ export class StateStore {
|
|||||||
todayEntry.inputTokens,
|
todayEntry.inputTokens,
|
||||||
todayEntry.outputTokens
|
todayEntry.outputTokens
|
||||||
);
|
);
|
||||||
todayEntry.sessions += 1;
|
|
||||||
|
// Only increment session count for unique sessions
|
||||||
|
if (sessionId && !this.dailySessionIds.has(sessionId)) {
|
||||||
|
this.dailySessionIds.add(sessionId);
|
||||||
|
todayEntry.sessions = this.dailySessionIds.size;
|
||||||
|
}
|
||||||
|
|
||||||
// Prune old entries (keep last 30 days)
|
// Prune old entries (keep last 30 days)
|
||||||
if (stats.daily.length > StateStore.MAX_DAILY_HISTORY) {
|
if (stats.daily.length > StateStore.MAX_DAILY_HISTORY) {
|
||||||
|
|||||||
+2
-2
@@ -2167,7 +2167,7 @@ export class WebServer extends EventEmitter {
|
|||||||
const deltaInput = session.inputTokens - lastRecorded.input;
|
const deltaInput = session.inputTokens - lastRecorded.input;
|
||||||
const deltaOutput = session.outputTokens - lastRecorded.output;
|
const deltaOutput = session.outputTokens - lastRecorded.output;
|
||||||
if (deltaInput > 0 || deltaOutput > 0) {
|
if (deltaInput > 0 || deltaOutput > 0) {
|
||||||
this.store.recordDailyUsage(deltaInput, deltaOutput);
|
this.store.recordDailyUsage(deltaInput, deltaOutput, sessionId);
|
||||||
}
|
}
|
||||||
this.lastRecordedTokens.delete(sessionId);
|
this.lastRecordedTokens.delete(sessionId);
|
||||||
console.log(`[Server] Added to global stats: ${session.inputTokens + session.outputTokens} tokens, $${session.totalCost.toFixed(4)} from session ${sessionId}`);
|
console.log(`[Server] Added to global stats: ${session.inputTokens + session.outputTokens} tokens, $${session.totalCost.toFixed(4)} from session ${sessionId}`);
|
||||||
@@ -3036,7 +3036,7 @@ export class WebServer extends EventEmitter {
|
|||||||
const deltaOutput = session.outputTokens - last.output;
|
const deltaOutput = session.outputTokens - last.output;
|
||||||
|
|
||||||
if (deltaInput > 0 || deltaOutput > 0) {
|
if (deltaInput > 0 || deltaOutput > 0) {
|
||||||
this.store.recordDailyUsage(deltaInput, deltaOutput);
|
this.store.recordDailyUsage(deltaInput, deltaOutput, sessionId);
|
||||||
this.lastRecordedTokens.set(sessionId, {
|
this.lastRecordedTokens.set(sessionId, {
|
||||||
input: session.inputTokens,
|
input: session.inputTokens,
|
||||||
output: session.outputTokens,
|
output: session.outputTokens,
|
||||||
|
|||||||
Reference in New Issue
Block a user