diff --git a/src/utils/incrementalTokenCounter.test.ts b/src/utils/incrementalTokenCounter.test.ts index 26c42d905a..203674d4eb 100644 --- a/src/utils/incrementalTokenCounter.test.ts +++ b/src/utils/incrementalTokenCounter.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, beforeEach } from 'bun:test' +import { roughTokenCountEstimationForMessages } from '../services/tokenEstimation.js' import { IncrementalTokenCounter, CounterFactory } from './incrementalTokenCounter.js' import type { Message } from '../types/message.js' @@ -247,6 +248,32 @@ describe('IncrementalTokenCounter', () => { expect(count2).toBeGreaterThan(count1) }) }) + + describe('attachment cache invalidation', () => { + it('recalculates when attachment-only messages differ', () => { + const counter = new IncrementalTokenCounter() + const smallAttachment = { + type: 'attachment', + attachment: { + type: 'opened_file_in_ide', + filename: 'a', + }, + } as unknown as Message + const largeAttachment = { + type: 'attachment', + attachment: { + type: 'opened_file_in_ide', + filename: 'x'.repeat(2000), + }, + } as unknown as Message + + const smallCount = counter.getCount([smallAttachment]) + const largeCount = counter.getCount([largeAttachment]) + + expect(largeCount).toBe(roughTokenCountEstimationForMessages([largeAttachment])) + expect(largeCount).toBeGreaterThan(smallCount) + }) + }) }) describe('CounterFactory', () => { @@ -270,4 +297,4 @@ describe('CounterFactory', () => { expect(counter).toBeInstanceOf(IncrementalTokenCounter) }) }) -}) \ No newline at end of file +}) diff --git a/src/utils/incrementalTokenCounter.ts b/src/utils/incrementalTokenCounter.ts index c88ddaf4fa..690b099459 100644 --- a/src/utils/incrementalTokenCounter.ts +++ b/src/utils/incrementalTokenCounter.ts @@ -30,16 +30,16 @@ export interface CounterStats { function getMessageHash(messages: readonly Message[]): string { if (messages.length === 0) return 'empty' - const fullContent = messages.map(m => { - const c = typeof m.message?.content === 'string' - ? m.message.content - : Array.isArray(m.message?.content) - ? JSON.stringify(m.message.content) - : '' - return c - }).join('|') + const tokenRelevantInput = messages.map(m => ({ + type: m.type, + content: m.message?.content ?? null, + attachment: m.attachment ?? null, + })) - return createHash('sha256').update(fullContent).digest('hex').slice(0, 16) + return createHash('sha256') + .update(JSON.stringify(tokenRelevantInput)) + .digest('hex') + .slice(0, 16) } /** @@ -254,4 +254,4 @@ export const CounterFactory = { estimationMultiplier: 1.2, }) }, -} \ No newline at end of file +}