Fix token summing in ClickHouse queries (#105)

.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:

.tool-versions

@@ -1 +1 @@
-nodejs 20.18.1
+nodejs 20.19.2

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"

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",

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<string, string | number | string[]> = {
     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

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');
+  });
+});

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';

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)}`;
+};

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'],