From c59248bda3746fe87f9280dfacd8147b5ed9b4ec Mon Sep 17 00:00:00 2001 From: Chris Estreich Date: Wed, 18 Jun 2025 11:36:58 -0700 Subject: [PATCH] Fix token summing in ClickHouse queries (#105) --- .github/workflows/CI.yml | 2 +- .tool-versions | 2 +- docker-compose.yml | 1 + package.json | 2 + src/actions/analytics/events.ts | 13 ++- src/lib/__tests__/query-utils.test.ts | 127 ++++++++++++++++++++++++++ src/lib/index.ts | 10 ++ src/lib/query-utils.ts | 28 ++++++ vitest.config.mts | 3 + 9 files changed, 181 insertions(+), 7 deletions(-) create mode 100644 src/lib/__tests__/query-utils.test.ts create mode 100644 src/lib/index.ts create mode 100644 src/lib/query-utils.ts diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 8aa542d1e0..27f486b4e1 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -7,7 +7,7 @@ on: branches: [main] env: - NODE_VERSION: 20.18.1 + NODE_VERSION: 20.19.2 PNPM_VERSION: 10.8.1 jobs: diff --git a/.tool-versions b/.tool-versions index e8fc3f8ea0..269cea0b28 100644 --- a/.tool-versions +++ b/.tool-versions @@ -1 +1 @@ -nodejs 20.18.1 +nodejs 20.19.2 diff --git a/docker-compose.yml b/docker-compose.yml index 25dd62e19d..9a92e89985 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -12,6 +12,7 @@ services: - POSTGRES_PASSWORD=password - POSTGRES_DATABASES=roo_code_development,roo_code_test clickhouse: + container_name: clickhouse image: clickhouse/clickhouse-server ports: - "8123:8123" diff --git a/package.json b/package.json index b484c67c69..4114a96524 100644 --- a/package.json +++ b/package.json @@ -19,6 +19,8 @@ "db:test:migrate": "dotenvx run -f .env.test -- drizzle-kit migrate", "db:test:push": "dotenvx run -f .env.test -- drizzle-kit push", "db:test:studio": "dotenvx run -f .env.test -- drizzle-kit studio", + "db:up": "docker compose up -d", + "db:down": "docker compose down", "storybook": "storybook dev -p 6006", "storybook:build": "storybook build", "storybook:serve": "http-server storybook-static --port 6006 --silent", diff --git a/src/actions/analytics/events.ts b/src/actions/analytics/events.ts index 3eb85321f8..067dc1ce28 100644 --- a/src/actions/analytics/events.ts +++ b/src/actions/analytics/events.ts @@ -9,6 +9,7 @@ import { import type { AnyTimePeriod } from '@/types'; import { analytics } from '@/lib/server'; +import { tokenSumSql } from '@/lib'; import { type User, getUsersById } from '@/db/server'; import { validateAnalyticsAccess } from '@/actions/auth'; @@ -108,7 +109,7 @@ export const getUsage = async ({ type, COUNT(1) as events, COUNT(distinct userId) as users, - SUM(COALESCE(inputTokens, 0) + COALESCE(outputTokens, 0)) AS tokens, + SUM(${tokenSumSql()}) AS tokens, SUM(COALESCE(cost, 0)) AS cost FROM events WHERE @@ -186,7 +187,7 @@ export const getDeveloperUsage = async ({ userId, SUM(CASE WHEN type = '${TelemetryEventName.TASK_CREATED}' THEN 1 ELSE 0 END) AS tasksStarted, SUM(CASE WHEN type = '${TelemetryEventName.TASK_COMPLETED}' THEN 1 ELSE 0 END) AS tasksCompleted, - SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(inputTokens, 0) + COALESCE(outputTokens, 0) ELSE 0 END) AS tokens, + SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN ${tokenSumSql()} ELSE 0 END) AS tokens, SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(cost, 0) ELSE 0 END) AS cost, MAX(timestamp) AS lastEventTimestamp FROM events @@ -266,7 +267,7 @@ export const getModelUsage = async ({ apiProvider as provider, modelId as model, SUM(CASE WHEN type = '${TelemetryEventName.TASK_CREATED}' THEN 1 ELSE 0 END) AS tasks, - SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(inputTokens, 0) + COALESCE(outputTokens, 0) ELSE 0 END) AS tokens, + SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN ${tokenSumSql()} ELSE 0 END) AS tokens, SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(cost, 0) ELSE 0 END) AS cost FROM events WHERE @@ -370,7 +371,7 @@ export const getTasks = async ({ argMin(e.modelId, e.timestamp) as model, any(fm.mode) AS mode, MAX(CASE WHEN e.type = 'Task Completed' THEN 1 ELSE 0 END) AS completed, - SUM(CASE WHEN e.type = 'LLM Completion' THEN COALESCE(e.inputTokens, 0) + COALESCE(e.outputTokens, 0) ELSE 0 END) AS tokens, + SUM(CASE WHEN e.type = 'LLM Completion' THEN ${tokenSumSql('e')} ELSE 0 END) AS tokens, SUM(CASE WHEN e.type = 'LLM Completion' THEN COALESCE(e.cost, 0) ELSE 0 END) AS cost, MIN(e.timestamp) AS timestamp, any(fm.title) AS title @@ -433,6 +434,7 @@ export const getHourlyUsageByUser = async ({ } const userFilter = effectiveUserId ? 'AND userId = {userId: String}' : ''; + const queryParams: Record = { orgId: orgId!, timePeriod, @@ -442,6 +444,7 @@ export const getHourlyUsageByUser = async ({ TelemetryEventName.LLM_COMPLETION, ], }; + if (effectiveUserId) { queryParams.userId = effectiveUserId; } @@ -452,7 +455,7 @@ export const getHourlyUsageByUser = async ({ toString(toStartOfHour(fromUnixTimestamp(timestamp))) as hour_utc, userId, SUM(CASE WHEN type = '${TelemetryEventName.TASK_CREATED}' THEN 1 ELSE 0 END) AS tasks, - SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(inputTokens, 0) + COALESCE(outputTokens, 0) ELSE 0 END) AS tokens, + SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN ${tokenSumSql()} ELSE 0 END) AS tokens, SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(cost, 0) ELSE 0 END) AS cost FROM events WHERE diff --git a/src/lib/__tests__/query-utils.test.ts b/src/lib/__tests__/query-utils.test.ts new file mode 100644 index 0000000000..388915061d --- /dev/null +++ b/src/lib/__tests__/query-utils.test.ts @@ -0,0 +1,127 @@ +// pnpm test src/lib/__tests__/query-utils.test.ts + +import { + inputTokenSumSql, + outputTokenSumSql, + tokenSumSql, +} from '../query-utils'; + +describe('inputTokenSumSql', () => { + it('should generate SQL for input tokens without table prefix', () => { + const result = inputTokenSumSql(); + expect(result).toBe( + 'COALESCE(inputTokens, 0) + COALESCE(cacheReadTokens, 0)', + ); + }); + + it('should generate SQL for input tokens with table prefix', () => { + const result = inputTokenSumSql('e'); + expect(result).toBe( + 'COALESCE(e.inputTokens, 0) + COALESCE(e.cacheReadTokens, 0)', + ); + }); + + it('should handle different table prefixes', () => { + const result = inputTokenSumSql('events'); + expect(result).toBe( + 'COALESCE(events.inputTokens, 0) + COALESCE(events.cacheReadTokens, 0)', + ); + }); + + it('should handle empty string table prefix', () => { + const result = inputTokenSumSql(''); + expect(result).toBe( + 'COALESCE(inputTokens, 0) + COALESCE(cacheReadTokens, 0)', + ); + }); +}); + +describe('outputTokenSumSql', () => { + it('should generate SQL for output tokens without table prefix', () => { + const result = outputTokenSumSql(); + expect(result).toBe( + 'COALESCE(outputTokens, 0) + COALESCE(cacheWriteTokens, 0)', + ); + }); + + it('should generate SQL for output tokens with table prefix', () => { + const result = outputTokenSumSql('e'); + expect(result).toBe( + 'COALESCE(e.outputTokens, 0) + COALESCE(e.cacheWriteTokens, 0)', + ); + }); + + it('should handle different table prefixes', () => { + const result = outputTokenSumSql('events'); + expect(result).toBe( + 'COALESCE(events.outputTokens, 0) + COALESCE(events.cacheWriteTokens, 0)', + ); + }); + + it('should handle empty string table prefix', () => { + const result = outputTokenSumSql(''); + expect(result).toBe( + 'COALESCE(outputTokens, 0) + COALESCE(cacheWriteTokens, 0)', + ); + }); +}); + +describe('tokenSumSql', () => { + it('should generate SQL for total tokens without table prefix', () => { + const result = tokenSumSql(); + const expected = + 'COALESCE(inputTokens, 0) + COALESCE(cacheReadTokens, 0) + COALESCE(outputTokens, 0) + COALESCE(cacheWriteTokens, 0)'; + expect(result).toBe(expected); + }); + + it('should generate SQL for total tokens with table prefix', () => { + const result = tokenSumSql('e'); + const expected = + 'COALESCE(e.inputTokens, 0) + COALESCE(e.cacheReadTokens, 0) + COALESCE(e.outputTokens, 0) + COALESCE(e.cacheWriteTokens, 0)'; + expect(result).toBe(expected); + }); + + it('should handle different table prefixes', () => { + const result = tokenSumSql('events'); + const expected = + 'COALESCE(events.inputTokens, 0) + COALESCE(events.cacheReadTokens, 0) + COALESCE(events.outputTokens, 0) + COALESCE(events.cacheWriteTokens, 0)'; + expect(result).toBe(expected); + }); + + it('should handle empty string table prefix', () => { + const result = tokenSumSql(''); + const expected = + 'COALESCE(inputTokens, 0) + COALESCE(cacheReadTokens, 0) + COALESCE(outputTokens, 0) + COALESCE(cacheWriteTokens, 0)'; + expect(result).toBe(expected); + }); + + it('should combine input and output token sums correctly', () => { + const table = 'test'; + const inputSql = inputTokenSumSql(table); + const outputSql = outputTokenSumSql(table); + const totalSql = tokenSumSql(table); + + expect(totalSql).toBe(`${inputSql} + ${outputSql}`); + }); +}); + +describe('Query utils integration', () => { + it('should maintain consistency between individual and combined functions', () => { + const tables = [undefined, '', 'e', 'events', 'messages']; + + tables.forEach((table) => { + const inputSql = inputTokenSumSql(table); + const outputSql = outputTokenSumSql(table); + const totalSql = tokenSumSql(table); + + expect(totalSql).toBe(`${inputSql} + ${outputSql}`); + }); + }); + + it('should generate valid SQL identifiers', () => { + const result = tokenSumSql('my_table'); + expect(result).toMatch(/^COALESCE\(my_table\./); + expect(result).not.toContain('undefined'); + expect(result).not.toContain('null'); + }); +}); diff --git a/src/lib/index.ts b/src/lib/index.ts new file mode 100644 index 0000000000..d9f5026de8 --- /dev/null +++ b/src/lib/index.ts @@ -0,0 +1,10 @@ +export * from './clipboard'; +export * from './constants'; +export * from './formatters'; +export * from './metadata'; +export * from './providers'; +export * from './query-utils'; +export * from './task-sharing'; +export * from './task-utils'; +export * from './timezone-utils'; +export * from './utils'; diff --git a/src/lib/query-utils.ts b/src/lib/query-utils.ts new file mode 100644 index 0000000000..ec772666a0 --- /dev/null +++ b/src/lib/query-utils.ts @@ -0,0 +1,28 @@ +/** + * Generates SQL for summing input tokens (inputTokens + cacheReadTokens) + * @param table Optional table prefix (e.g., 'e' for 'e.inputTokens') + * @returns SQL fragment for input token sum + */ +export const inputTokenSumSql = (table?: string) => { + const t = table ? `${table}.` : ''; + return `COALESCE(${t}inputTokens, 0) + COALESCE(${t}cacheReadTokens, 0)`; +}; + +/** + * Generates SQL for summing output tokens (outputTokens + cacheWriteTokens) + * @param table Optional table prefix (e.g., 'e' for 'e.outputTokens') + * @returns SQL fragment for output token sum + */ +export const outputTokenSumSql = (table?: string) => { + const t = table ? `${table}.` : ''; + return `COALESCE(${t}outputTokens, 0) + COALESCE(${t}cacheWriteTokens, 0)`; +}; + +/** + * Generates SQL for summing all tokens (input + output) + * @param table Optional table prefix (e.g., 'e' for 'e.inputTokens') + * @returns SQL fragment for total token sum + */ +export const tokenSumSql = (table?: string) => { + return `${inputTokenSumSql(table)} + ${outputTokenSumSql(table)}`; +}; diff --git a/vitest.config.mts b/vitest.config.mts index 8f9b8aee73..ef40c8edf0 100644 --- a/vitest.config.mts +++ b/vitest.config.mts @@ -2,6 +2,9 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ test: { + watch: false, + reporters: ['dot'], + silent: true, coverage: { include: ['src/**/*'], exclude: ['src/**/*.stories.{js,jsx,ts,tsx}', '**/*.d.ts'],