diff --git a/apps/coder/src/services/token-analysis/__tests__/analyzer.test.ts b/apps/coder/src/services/token-analysis/__tests__/analyzer.test.ts new file mode 100644 index 0000000..8d83de7 --- /dev/null +++ b/apps/coder/src/services/token-analysis/__tests__/analyzer.test.ts @@ -0,0 +1,29 @@ +import { describe, it, expect } from 'vitest'; +import { analyzeMessages } from '../analyzer.js'; + +describe('analyzeMessages', () => { + it('classifies user messages', () => { + const breakdown = analyzeMessages([{ role: 'user', content: 'hello world' }]); + expect(breakdown.user).toBeGreaterThan(0); + expect(breakdown.total).toBe(breakdown.user); + }); + + it('counts tool calls', () => { + const parts = [ + { role: 'assistant', content: 'using grep', tool_calls: [{ id: '1', name: 'grep', arguments: '{}' }] }, + { role: 'tool', content: '{"files":[]}', tool_call_id: '1' }, + ]; + const breakdown = analyzeMessages(parts); + expect(breakdown.tools).toBeGreaterThan(0); + expect(breakdown.assistant).toBeGreaterThan(0); + }); + + it('separates reasoning tokens', () => { + const parts = [ + { role: 'assistant', content: 'short answer', reasoning_parts: [{ text: 'long chain of thought reasoning here' }] }, + ]; + const breakdown = analyzeMessages(parts); + expect(breakdown.reasoning).toBeGreaterThan(0); + expect(breakdown.assistant).toBeLessThan(breakdown.reasoning); + }); +}); diff --git a/apps/coder/src/services/token-analysis/__tests__/persist.test.ts b/apps/coder/src/services/token-analysis/__tests__/persist.test.ts new file mode 100644 index 0000000..9106d15 --- /dev/null +++ b/apps/coder/src/services/token-analysis/__tests__/persist.test.ts @@ -0,0 +1,10 @@ +import { describe, it, expect } from 'vitest'; + +describe('persistTaskBreakdown', () => { + it('exports functions', async () => { + const mod = await import('../persist.js'); + expect(typeof mod.persistTaskBreakdown).toBe('function'); + expect(typeof mod.getTaskBreakdown).toBe('function'); + expect(typeof mod.analyzeAndPersistTaskBreakdown).toBe('function'); + }); +}); diff --git a/apps/coder/src/services/token-analysis/analyzer.ts b/apps/coder/src/services/token-analysis/analyzer.ts new file mode 100644 index 0000000..74a54b3 --- /dev/null +++ b/apps/coder/src/services/token-analysis/analyzer.ts @@ -0,0 +1,60 @@ +// TokenScope analyzer — classifies message parts into category breakdown. +// Ported from opencode-tokenscope (MIT). + +export interface TokenBreakdown { + system: number; + user: number; + assistant: number; + tools: number; + reasoning: number; + total: number; +} + +const CHARS_PER_TOKEN = 4; + +function estimateTokens(text: string): number { + return Math.ceil(text.length / CHARS_PER_TOKEN); +} + +export function analyzeMessages(parts: any[]): TokenBreakdown { + const breakdown: TokenBreakdown = { system: 0, user: 0, assistant: 0, tools: 0, reasoning: 0, total: 0 }; + + for (const part of parts) { + const role = part.role ?? ''; + const content = part.content ?? ''; + const tokens = estimateTokens(content); + + switch (role) { + case 'system': + breakdown.system += tokens; + break; + case 'user': + breakdown.user += tokens; + break; + case 'assistant': + breakdown.assistant += tokens; + if (part.tool_calls) { + for (const tc of part.tool_calls) { + breakdown.tools += estimateTokens(JSON.stringify(tc)); + } + } + break; + case 'tool': + breakdown.tools += tokens; + break; + default: + breakdown.assistant += tokens; + } + + if (part.reasoning_parts) { + for (const rp of part.reasoning_parts) { + const rTokens = estimateTokens(rp.text ?? ''); + breakdown.reasoning += rTokens; + breakdown.assistant -= rTokens; + } + } + } + + breakdown.total = breakdown.system + breakdown.user + breakdown.assistant + breakdown.tools + breakdown.reasoning; + return breakdown; +} diff --git a/apps/coder/src/services/token-analysis/persist.ts b/apps/coder/src/services/token-analysis/persist.ts new file mode 100644 index 0000000..39ec79d --- /dev/null +++ b/apps/coder/src/services/token-analysis/persist.ts @@ -0,0 +1,35 @@ +// TokenScope persistence — writes breakdown to task records. +import type { Sql } from '../../db.js'; +import type { TokenBreakdown } from './analyzer.js'; + +export async function persistTaskBreakdown( + sql: Sql, + taskId: string, + breakdown: TokenBreakdown, +): Promise { + await sql` + UPDATE tasks SET token_breakdown = ${sql.json(breakdown as never)} + WHERE id = ${taskId} + `; +} + +export async function getTaskBreakdown( + sql: Sql, + taskId: string, +): Promise { + const rows = await sql<{ token_breakdown: any }[]>` + SELECT token_breakdown FROM tasks WHERE id = ${taskId} + `; + return rows[0]?.token_breakdown ?? null; +} + +export async function analyzeAndPersistTaskBreakdown( + sql: Sql, + taskId: string, + parts: any[], +): Promise { + const { analyzeMessages } = await import('./analyzer.js'); + const breakdown = analyzeMessages(parts); + await persistTaskBreakdown(sql, taskId, breakdown); + return breakdown; +}