Capture more events when backfilling (#133)

This commit is contained in:
Chris Estreich 2025-06-24 11:35:25 -07:00 committed by GitHub
parent a989449eb2
commit 3b2f9ee24a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 188 additions and 0 deletions

View file

@ -239,6 +239,7 @@ describe('/api/events/backfill', () => {
}>;
expect(eventData).toHaveLength(1);
expect(eventData[0]).toMatchObject({
taskId: 'test-task-from-file',
type: 'Task Created',
@ -246,6 +247,7 @@ describe('/api/events/backfill', () => {
userId: 'test-user-id',
orgId: 'test-org-id',
});
expect(typeof eventData[0]?.timestamp).toBe('number');
const dbResults = await analytics.query({
@ -276,6 +278,7 @@ describe('/api/events/backfill', () => {
expect(dbData).toHaveLength(messages.length);
const textMessage = dbData.find((msg) => msg.say === 'text');
expect(textMessage).toMatchObject({
taskId: 'test-task-from-file',
mode: 'code',
@ -285,6 +288,7 @@ describe('/api/events/backfill', () => {
});
const apiMessage = dbData.find((msg) => msg.say === 'api_req_started');
expect(apiMessage).toMatchObject({
taskId: 'test-task-from-file',
mode: 'code',
@ -301,6 +305,134 @@ describe('/api/events/backfill', () => {
expect(messageTypes).toContain('checkpoint_saved');
expect(messageTypes).toContain('reasoning');
expect(messageTypes).toContain('completion_result');
const llmEventsResults = await analytics.query({
query: `
SELECT
taskId,
type,
mode,
inputTokens,
outputTokens,
cacheReadTokens,
cacheWriteTokens,
cost,
userId,
orgId,
timestamp
FROM events
WHERE taskId = 'test-task-from-file' AND type = 'LLM Completion'
ORDER BY timestamp ASC
`,
format: 'JSONEachRow',
});
const llmEventData = (await llmEventsResults.json()) as Array<{
taskId: string;
type: string;
mode: string;
inputTokens: number;
outputTokens: number;
cacheReadTokens?: number;
cacheWriteTokens?: number;
cost?: number;
userId: string;
orgId: string;
timestamp: number;
}>;
expect(llmEventData).toHaveLength(3);
expect(llmEventData[0]).toMatchObject({
taskId: 'test-task-from-file',
type: 'LLM Completion',
mode: 'code',
inputTokens: 12990,
outputTokens: 559,
cacheReadTokens: 0,
cacheWriteTokens: 0,
cost: 0.02246925,
userId: 'test-user-id',
orgId: 'test-org-id',
});
expect(llmEventData[1]).toMatchObject({
taskId: 'test-task-from-file',
type: 'LLM Completion',
mode: 'code',
inputTokens: 13788,
outputTokens: 487,
cacheReadTokens: 0,
cacheWriteTokens: 0,
cost: 0.0142215,
userId: 'test-user-id',
orgId: 'test-org-id',
});
expect(llmEventData[2]).toMatchObject({
taskId: 'test-task-from-file',
type: 'LLM Completion',
mode: 'code',
inputTokens: 14399,
outputTokens: 466,
cacheReadTokens: 0,
cacheWriteTokens: 0,
cost: 0.01344465,
userId: 'test-user-id',
orgId: 'test-org-id',
});
llmEventData.forEach((event) => {
expect(typeof event.timestamp).toBe('number');
});
const taskCompletedEventsResults = await analytics.query({
query: `
SELECT
taskId,
type,
mode,
userId,
orgId,
timestamp
FROM events
WHERE taskId = 'test-task-from-file' AND type = 'Task Completed'
ORDER BY timestamp ASC
`,
format: 'JSONEachRow',
});
const taskCompletedEventData =
(await taskCompletedEventsResults.json()) as Array<{
taskId: string;
type: string;
mode: string;
userId: string;
orgId: string;
timestamp: number;
}>;
expect(taskCompletedEventData).toHaveLength(2);
expect(taskCompletedEventData[0]).toMatchObject({
taskId: 'test-task-from-file',
type: 'Task Completed',
mode: 'code',
userId: 'test-user-id',
orgId: 'test-org-id',
});
expect(taskCompletedEventData[1]).toMatchObject({
taskId: 'test-task-from-file',
type: 'Task Completed',
mode: 'code',
userId: 'test-user-id',
orgId: 'test-org-id',
});
taskCompletedEventData.forEach((event) => {
expect(typeof event.timestamp).toBe('number');
});
});
it('should return 401 if authentication fails', async () => {

View file

@ -109,6 +109,54 @@ export async function POST(request: NextRequest) {
const timestamp = Math.round(message.ts / 1000);
const mode = extractMode(message.text) || defaultMode;
if (message.say === 'api_req_started') {
try {
const result = apiReqStartedSchema.safeParse(
JSON.parse(message.text || '{}'),
);
if (
result.success &&
(result.data.tokensIn ||
result.data.tokensOut ||
result.data.cost)
) {
await captureEvent({
id: uuidv4(),
orgId,
userId,
timestamp,
event: {
type: TelemetryEventName.LLM_COMPLETION,
properties: {
taskId,
...properties,
mode,
inputTokens: result.data.tokensIn ?? 0,
outputTokens: result.data.tokensOut ?? 0,
cacheReadTokens: result.data.cacheReads,
cacheWriteTokens: result.data.cacheWrites,
cost: result.data.cost,
},
},
});
}
} catch {
// Ignore JSON parsing and validation errors.
}
} else if (message.say === 'completion_result') {
await captureEvent({
id: uuidv4(),
orgId,
userId,
timestamp,
event: {
type: TelemetryEventName.TASK_COMPLETED,
properties: { taskId, ...properties, mode },
},
});
}
const event = {
type: TelemetryEventName.TASK_MESSAGE as const,
properties: { taskId, message, ...properties, mode },
@ -133,6 +181,14 @@ export async function POST(request: NextRequest) {
}
}
const apiReqStartedSchema = z.object({
tokensIn: z.number().optional(),
tokensOut: z.number().optional(),
cacheReads: z.number().optional(),
cacheWrites: z.number().optional(),
cost: z.number().optional(),
});
/**
* Extracts the mode from a message text if it contains a <slug>mode</slug> pattern.
* @param text - The message text to parse