From 6e4564abd18a14b79ed0f25dc8d88e7652cba7fa Mon Sep 17 00:00:00 2001 From: arkon Date: Mon, 26 Jan 2026 12:22:05 +0100 Subject: [PATCH] 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 --- src/session.ts | 55 ++++++++++++++++++++++++++++++++++++++++++---- src/state-store.ts | 40 +++++++++++++++++++++++++++++++-- src/web/server.ts | 4 ++-- 3 files changed, 91 insertions(+), 8 deletions(-) diff --git a/src/session.ts b/src/session.ts index 36070d3f..f256d588 100644 --- a/src/session.ts +++ b/src/session.ts @@ -626,6 +626,18 @@ export class Session extends EventEmitter { * Called when recovering sessions after server restart. */ 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._totalOutputTokens = outputTokens; this._totalCost = totalCost; @@ -1298,10 +1310,19 @@ export class Session extends EventEmitter { this._textOutput.append(block.text); } } - // Track tokens from usage + // Track tokens from usage (with validation) if (msg.message.usage) { - this._totalInputTokens += msg.message.usage.input_tokens || 0; - this._totalOutputTokens += msg.message.usage.output_tokens || 0; + const inputDelta = msg.message.usage.input_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 this.checkAutoCompact(); @@ -1423,6 +1444,11 @@ export class Session extends EventEmitter { // Parse token count from Claude's status line in interactive mode // 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 { // Quick pre-check: skip expensive regex if "token" not present (performance optimization) if (!data.includes('token')) return; @@ -1442,16 +1468,37 @@ export class Session extends EventEmitter { if (suffix === 'k') { tokenCount *= 1000; } 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; } + // 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) // We use total tokens as an estimate - Claude shows combined input+output const currentTotal = this._totalInputTokens + this._totalOutputTokens; 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) // 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._totalOutputTokens += Math.round(delta * 0.4); diff --git a/src/state-store.ts b/src/state-store.ts index 6ada0178..222f61fc 100644 --- a/src/state-store.ts +++ b/src/state-store.ts @@ -223,6 +223,18 @@ export class StateStore { * Call when a session is deleted to preserve its usage in lifetime stats. */ 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(); stats.totalInputTokens += inputTokens; stats.totalOutputTokens += outputTokens; @@ -309,18 +321,37 @@ export class StateStore { return inputCost + outputCost; } + // Track unique sessions per day for accurate session count + private dailySessionIds: Set = new Set(); + private dailySessionDate: string = ''; + /** * Record token usage for today. * Accumulates tokens to today's entry, creating it if needed. * @param inputTokens Input 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; + // 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 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 let todayEntry = stats.daily.find(e => e.date === today); if (!todayEntry) { @@ -341,7 +372,12 @@ export class StateStore { todayEntry.inputTokens, 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) if (stats.daily.length > StateStore.MAX_DAILY_HISTORY) { diff --git a/src/web/server.ts b/src/web/server.ts index 4efe2835..d4c09cd3 100644 --- a/src/web/server.ts +++ b/src/web/server.ts @@ -2167,7 +2167,7 @@ export class WebServer extends EventEmitter { const deltaInput = session.inputTokens - lastRecorded.input; const deltaOutput = session.outputTokens - lastRecorded.output; if (deltaInput > 0 || deltaOutput > 0) { - this.store.recordDailyUsage(deltaInput, deltaOutput); + this.store.recordDailyUsage(deltaInput, deltaOutput, sessionId); } this.lastRecordedTokens.delete(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; if (deltaInput > 0 || deltaOutput > 0) { - this.store.recordDailyUsage(deltaInput, deltaOutput); + this.store.recordDailyUsage(deltaInput, deltaOutput, sessionId); this.lastRecordedTokens.set(sessionId, { input: session.inputTokens, output: session.outputTokens,