diff --git a/prisma/migrations/20260728120000_add_cache_token_columns_to_shared_conversations/migration.sql b/prisma/migrations/20260728120000_add_cache_token_columns_to_shared_conversations/migration.sql new file mode 100644 index 0000000000..04cf86f234 --- /dev/null +++ b/prisma/migrations/20260728120000_add_cache_token_columns_to_shared_conversations/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "shared_conversations" ADD COLUMN "cache_read_tokens" INTEGER NOT NULL DEFAULT 0, +ADD COLUMN "cache_write_tokens" INTEGER NOT NULL DEFAULT 0; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 669c72039f..d8f7d52ecb 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -640,6 +640,9 @@ model SharedConversation { // weights it accordingly when summing recent spend. inputTokens Int @default(0) @map("input_tokens") outputTokens Int @default(0) @map("output_tokens") + // Cache portions of `inputTokens`, which is the inclusive total. + cacheReadTokens Int @default(0) @map("cache_read_tokens") + cacheWriteTokens Int @default(0) @map("cache_write_tokens") createdAt DateTime @default(now()) @map("created_at") updatedAt DateTime @updatedAt @map("updated_at") isShared Boolean @default(false) @map("is_shared") diff --git a/src/__tests__/integration/api/ask/quick.test.ts b/src/__tests__/integration/api/ask/quick.test.ts index d9ae0731c9..fd872a0434 100644 --- a/src/__tests__/integration/api/ask/quick.test.ts +++ b/src/__tests__/integration/api/ask/quick.test.ts @@ -28,10 +28,16 @@ vi.mock('next/server', async (importOriginal) => { }); // Mock the AI streaming service -vi.mock('ai', () => ({ +vi.mock('ai', async (importOriginal) => ({ + ...(await importOriginal()), streamText: vi.fn(), })); +// The route merges `result.toUIMessageStream()` into its own +// `createUIMessageStream`, so every mocked result needs one. +const emptyUIMessageStream = () => + new ReadableStream({ start: (c) => c.close() }); + // Mock the AI provider module (wrapper around aieo) vi.mock('@/lib/ai/provider', () => ({ getModel: vi.fn(), @@ -478,6 +484,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test', { headers: { 'Content-Type': 'text/plain' }, @@ -581,6 +588,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test', { headers: { 'Content-Type': 'text/plain' }, @@ -666,6 +674,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test response', { headers: { 'Content-Type': 'text/plain' }, @@ -725,6 +734,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test response', { headers: { 'Content-Type': 'text/plain' }, @@ -781,6 +791,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test response', { headers: { 'Content-Type': 'text/plain' }, @@ -851,6 +862,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn( () => new Response('ok', { headers: { 'Content-Type': 'text/plain' } }), ), @@ -951,6 +963,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn( () => new Response('ok', { headers: { 'Content-Type': 'text/plain' } }), ), @@ -1045,6 +1058,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { await setupWorkspaceWithConversation(storedMessages); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test', { headers: { 'Content-Type': 'text/plain' }, @@ -1208,6 +1222,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test', { headers: { 'Content-Type': 'text/plain' } }) ), @@ -1310,6 +1325,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { await setupWorkspaceWithConversation(storedMessages); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn(() => new Response('test', { headers: { 'Content-Type': 'text/plain' }, @@ -1483,6 +1499,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn( () => new Response('ok', { headers: { 'Content-Type': 'text/plain' } }), ), @@ -1537,6 +1554,7 @@ describe('POST /api/ask/quick - Quick Ask Integration Tests', () => { }); const mockStream = { + toUIMessageStream: vi.fn(() => emptyUIMessageStream()), toUIMessageStreamResponse: vi.fn( () => new Response('ok', { headers: { 'Content-Type': 'text/plain' } }), ), diff --git a/src/__tests__/unit/components/agent-logs/TurnTokenUsage.test.tsx b/src/__tests__/unit/components/agent-logs/TurnTokenUsage.test.tsx index ff03884bcc..1938299f8e 100644 --- a/src/__tests__/unit/components/agent-logs/TurnTokenUsage.test.tsx +++ b/src/__tests__/unit/components/agent-logs/TurnTokenUsage.test.tsx @@ -98,4 +98,36 @@ describe("TurnTokenUsage", () => { expect(screen.queryByText(/read:/)).toBeNull(); expect(screen.queryByText(/write:/)).toBeNull(); }); + it("renders elapsed time alongside the token counts", () => { + render( + React.createElement(TurnTokenUsage, { + usage: { inputTokens: 3411, outputTokens: 102 }, + elapsedMs: 4912, + }), + ); + expect(screen.getByText("4.9s")).toBeTruthy(); + }); + + it("formats sub-second, second and minute durations", () => { + const cases: Array<[number, string]> = [ + [820, "820ms"], + [4912, "4.9s"], + [72_400, "1m 12s"], + ]; + for (const [ms, label] of cases) { + const { unmount } = render( + React.createElement(TurnTokenUsage, { + usage: { inputTokens: 10 }, + elapsedMs: ms, + }), + ); + expect(screen.getByText(label)).toBeTruthy(); + unmount(); + } + }); + + it("omits the duration when elapsedMs is not provided", () => { + render(React.createElement(TurnTokenUsage, { usage: { inputTokens: 10 } })); + expect(screen.queryByText(/ms$|s$/)).toBeNull(); + }); }); diff --git a/src/__tests__/unit/lib/ai/publicChatBudget.test.ts b/src/__tests__/unit/lib/ai/publicChatBudget.test.ts index d84df666a4..629ef0690e 100644 --- a/src/__tests__/unit/lib/ai/publicChatBudget.test.ts +++ b/src/__tests__/unit/lib/ai/publicChatBudget.test.ts @@ -1,12 +1,40 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi, beforeEach } from "vitest"; import { NextRequest } from "next/server"; + +vi.mock("@/lib/db", () => ({ + db: { + sharedConversation: { + aggregate: vi.fn(), + update: vi.fn(), + }, + }, +})); + +import { db } from "@/lib/db"; import { ANON_DAILY_TOKEN_CAP, OUTPUT_TOKEN_WEIGHT, + CACHE_READ_TOKEN_WEIGHT, + CACHE_WRITE_TOKEN_WEIGHT, WORKSPACE_PUBLIC_DAILY_TOKEN_CAP, deriveAnonymousId, + checkPublicChatBudget, + recordTurnTokens, } from "@/lib/ai/publicChatBudget"; +const aggregate = db.sharedConversation.aggregate as ReturnType; +const update = db.sharedConversation.update as ReturnType; + +/** Both the anon and the workspace aggregate resolve to these sums. */ +function withSpend(sum: Record) { + aggregate.mockResolvedValue({ _sum: sum }); +} + +beforeEach(() => { + vi.clearAllMocks(); + update.mockResolvedValue({}); +}); + function makeReq(headers: Record): NextRequest { return new NextRequest("http://localhost/api/ask/quick", { method: "POST", @@ -70,4 +98,115 @@ describe("token budget constants", () => { ANON_DAILY_TOKEN_CAP, ); }); + + it("prices cache reads below and cache writes above fresh input", () => { + // Anthropic: $0.30/Mtok cache read, $3.75/Mtok cache write, + // $3/Mtok fresh input. + expect(CACHE_READ_TOKEN_WEIGHT).toBe(0.1); + expect(CACHE_WRITE_TOKEN_WEIGHT).toBe(1.25); + }); +}); + +describe("checkPublicChatBudget — cache repricing", () => { + const args = { workspaceId: "ws-1", anonymousId: "anon-1" }; + + it("treats a row with no cache data exactly as before", async () => { + // inputTokens alone must still price at 1x, so historical rows and + // non-cached providers are unaffected. + withSpend({ inputTokens: ANON_DAILY_TOKEN_CAP - 1, outputTokens: 0 }); + await expect(checkPublicChatBudget(args)).resolves.toMatchObject({ + allowed: true, + }); + + withSpend({ inputTokens: ANON_DAILY_TOKEN_CAP, outputTokens: 0 }); + await expect(checkPublicChatBudget(args)).resolves.toMatchObject({ + allowed: false, + reason: "anon", + }); + }); + + it("does not double-count cache tokens, which are already inside inputTokens", async () => { + // A cache-heavy conversation: 200k input of which 199k was a cache + // read. Adding the cache on top of inputTokens would price this at + // ~220k and block the visitor; repricing correctly puts it at + // 1k + 199k*0.1 ≈ 20.9k, comfortably under the 100k cap. + withSpend({ + inputTokens: 200_000, + outputTokens: 0, + cacheReadTokens: 199_000, + cacheWriteTokens: 0, + }); + + await expect(checkPublicChatBudget(args)).resolves.toMatchObject({ + allowed: true, + }); + }); + + it("still blocks when the repriced cost genuinely exceeds the cap", async () => { + withSpend({ + inputTokens: 100_000, + outputTokens: 50_000, + cacheReadTokens: 90_000, + cacheWriteTokens: 0, + }); + // fresh 10k + reads 9k + output 250k = well past the cap. + await expect(checkPublicChatBudget(args)).resolves.toMatchObject({ + allowed: false, + }); + }); + + it("charges a cache write more than the same number of fresh tokens", async () => { + // 80k of pure cache writes → 100k weighted, exactly at the cap. + withSpend({ + inputTokens: 80_000, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 80_000, + }); + await expect(checkPublicChatBudget(args)).resolves.toMatchObject({ + allowed: false, + }); + }); +}); + +describe("recordTurnTokens", () => { + it("increments all four counters", async () => { + await recordTurnTokens({ + conversationId: "conv-1", + inputTokens: 6873, + outputTokens: 47, + cacheReadTokens: 6059, + cacheWriteTokens: 812, + }); + + expect(update).toHaveBeenCalledWith({ + where: { id: "conv-1" }, + data: { + inputTokens: { increment: 6873 }, + outputTokens: { increment: 47 }, + cacheReadTokens: { increment: 6059 }, + cacheWriteTokens: { increment: 812 }, + }, + }); + }); + + it("records a cache-only turn rather than skipping it", async () => { + await recordTurnTokens({ + conversationId: "conv-1", + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 4096, + cacheWriteTokens: 0, + }); + expect(update).toHaveBeenCalledTimes(1); + }); + + it("skips the write when the turn spent nothing at all", async () => { + await recordTurnTokens({ + conversationId: "conv-1", + inputTokens: 0, + outputTokens: 0, + }); + expect(update).not.toHaveBeenCalled(); + }); }); diff --git a/src/__tests__/unit/lib/streaming/useStreamProcessor.test.ts b/src/__tests__/unit/lib/streaming/useStreamProcessor.test.ts index 3a94c16e15..df2e488444 100644 --- a/src/__tests__/unit/lib/streaming/useStreamProcessor.test.ts +++ b/src/__tests__/unit/lib/streaming/useStreamProcessor.test.ts @@ -1294,4 +1294,110 @@ describe("useStreamProcessor", () => { expect(finalMessage.usage?.cacheWriteTokens).toBe(50); }); }); + describe("Per-step usage (data-usage parts)", () => { + const usagePart = (stepNumber: number, usage: Record) => + TestDataFactories.createSSEEvent("data-usage", { + data: { stepNumber, timestamp: "2026-07-28T10:00:00.000Z", usage }, + }); + + test("accumulates a running total across steps", async () => { + const { result } = renderHook(() => useStreamProcessor()); + const onUpdate = TestUtils.createOnUpdateSpy(); + + const response = TestDataFactories.createMockResponse([ + TestDataFactories.createTextStartEvent("t1"), + TestDataFactories.createTextDeltaEvent("t1", "hi"), + usagePart(0, { + inputTokens: 3411, + outputTokens: 102, + cacheReadTokens: 0, + cacheWriteTokens: 3408, + }), + usagePart(1, { + inputTokens: 3594, + outputTokens: 59, + cacheReadTokens: 3408, + cacheWriteTokens: 179, + }), + TestDataFactories.createFinishEvent("stop"), + ]); + + await result.current.processStream(response, "msg-1", onUpdate); + + const finalMessage = onUpdate.mock.calls[onUpdate.mock.calls.length - 1][0]; + expect(finalMessage.usage).toEqual({ + inputTokens: 7005, + outputTokens: 161, + cacheReadTokens: 3408, + cacheWriteTokens: 3587, + }); + }); + + test("surfaces usage without any finish event", async () => { + // The UI-protocol `finish` chunk carries no usage, so the running + // total must not depend on one arriving. + const { result } = renderHook(() => useStreamProcessor()); + const onUpdate = TestUtils.createOnUpdateSpy(); + + const response = TestDataFactories.createMockResponse([ + TestDataFactories.createTextStartEvent("t1"), + TestDataFactories.createTextDeltaEvent("t1", "hi"), + usagePart(0, { inputTokens: 500, outputTokens: 20 }), + ]); + + await result.current.processStream(response, "msg-1", onUpdate); + + const finalMessage = onUpdate.mock.calls[onUpdate.mock.calls.length - 1][0]; + expect(finalMessage.usage).toEqual({ inputTokens: 500, outputTokens: 20 }); + }); + + test("exposes elapsedMs from the most recent part", async () => { + const { result } = renderHook(() => useStreamProcessor()); + const onUpdate = TestUtils.createOnUpdateSpy(); + + const response = TestDataFactories.createMockResponse([ + TestDataFactories.createTextStartEvent("t1"), + TestDataFactories.createTextDeltaEvent("t1", "hi"), + TestDataFactories.createSSEEvent("data-usage", { + data: { + stepNumber: 0, + timestamp: "2026-07-28T10:00:00.000Z", + elapsedMs: 2700, + usage: { inputTokens: 100 }, + }, + }), + TestDataFactories.createSSEEvent("data-usage", { + data: { + stepNumber: 1, + timestamp: "2026-07-28T10:00:05.000Z", + elapsedMs: 7200, + usage: { inputTokens: 50 }, + }, + }), + ]); + + await result.current.processStream(response, "msg-1", onUpdate); + + const finalMessage = onUpdate.mock.calls[onUpdate.mock.calls.length - 1][0]; + // Latest wins — it is a turn duration, not a per-step delta to sum. + expect(finalMessage.elapsedMs).toBe(7200); + expect(finalMessage.usage).toEqual({ inputTokens: 150 }); + }); + + test("ignores a malformed part with no usage payload", async () => { + const { result } = renderHook(() => useStreamProcessor()); + const onUpdate = TestUtils.createOnUpdateSpy(); + + const response = TestDataFactories.createMockResponse([ + TestDataFactories.createTextStartEvent("t1"), + TestDataFactories.createTextDeltaEvent("t1", "hi"), + TestDataFactories.createSSEEvent("data-usage", { data: {} }), + ]); + + await result.current.processStream(response, "msg-1", onUpdate); + + const finalMessage = onUpdate.mock.calls[onUpdate.mock.calls.length - 1][0]; + expect(finalMessage.usage).toBeUndefined(); + }); + }); }); diff --git a/src/__tests__/unit/services/canvas-agent-log.test.ts b/src/__tests__/unit/services/canvas-agent-log.test.ts new file mode 100644 index 0000000000..162373e916 --- /dev/null +++ b/src/__tests__/unit/services/canvas-agent-log.test.ts @@ -0,0 +1,221 @@ +import { describe, test, expect, vi, beforeEach } from "vitest"; + +const { put } = vi.hoisted(() => ({ put: vi.fn() })); + +vi.mock("@vercel/blob", () => ({ put })); + +vi.mock("@/lib/db", () => ({ + db: { + sharedConversation: { findUnique: vi.fn() }, + agentLog: { findFirst: vi.fn(), create: vi.fn(), update: vi.fn() }, + }, +})); + +import { db } from "@/lib/db"; +import { parseAgentLogStats } from "@/lib/utils/agent-log-stats"; +import { + transcriptFromStoredMessages, + writeCanvasAgentLog, + CANVAS_AGENT_NAME, +} from "@/services/canvas-agent-log"; +import type { StoredMessage } from "@/services/canvas-turn-persistence"; + +const findUnique = db.sharedConversation.findUnique as ReturnType; +const findFirst = db.agentLog.findFirst as ReturnType; +const create = db.agentLog.create as ReturnType; +const update = db.agentLog.update as ReturnType; + +const CONVERSATION: StoredMessage[] = [ + { id: "t1-u", role: "user", content: "Why is CI failing?" }, + { + id: "t1-a0", + role: "assistant", + content: "Let me look.", + timestamp: "2026-07-28T10:00:00.000Z", + usage: { + inputTokens: 3411, + outputTokens: 102, + cacheReadTokens: 0, + cacheWriteTokens: 3408, + }, + }, + { + id: "t1-a1", + role: "assistant", + content: "", + timestamp: "2026-07-28T10:00:00.000Z", + toolCalls: [ + { + id: "tc1", + toolName: "repo_agent", + input: { prompt: "check CI" }, + output: { answer: "a flaky test" }, + status: "output-available", + }, + { + id: "tc2", + toolName: "learn_concept", + input: { id: "c1" }, + output: { concept: "CI" }, + status: "output-available", + }, + ], + }, + { + id: "t1-a2", + role: "assistant", + content: "A flaky test is failing.", + timestamp: "2026-07-28T10:00:05.000Z", + usage: { + inputTokens: 3594, + outputTokens: 59, + cacheReadTokens: 3408, + cacheWriteTokens: 179, + }, + }, +]; + +beforeEach(() => { + vi.clearAllMocks(); + put.mockResolvedValue({ url: "https://blob.example/canvas.json" }); +}); + +describe("transcriptFromStoredMessages", () => { + test("turns stored toolCalls into tool-call content parts the parser counts", () => { + const transcript = transcriptFromStoredMessages(CONVERSATION); + const toolRow = transcript[2]; + + expect(Array.isArray(toolRow.content)).toBe(true); + const parts = toolRow.content as Array<{ type: string; toolName?: string }>; + expect(parts).toHaveLength(2); + expect(parts.map((p) => p.type)).toEqual(["tool-call", "tool-call"]); + expect(parts.map((p) => p.toolName)).toEqual(["repo_agent", "learn_concept"]); + }); + + test("preserves per-message timestamps and usage", () => { + const transcript = transcriptFromStoredMessages(CONVERSATION); + expect(transcript[1].timestamp).toBe("2026-07-28T10:00:00.000Z"); + expect(transcript[3].timestamp).toBe("2026-07-28T10:00:05.000Z"); + expect(transcript[1].usage).toEqual({ + inputTokens: 3411, + outputTokens: 102, + cacheReadTokens: 0, + cacheWriteTokens: 3408, + }); + }); + + test("round-trips through parseAgentLogStats into real stats", () => { + const blob = JSON.stringify({ + sessionId: "conv-1", + messages: transcriptFromStoredMessages(CONVERSATION), + }); + + const { stats } = parseAgentLogStats(blob); + + expect(stats.totalMessages).toBe(4); + expect(stats.totalToolCalls).toBe(2); + expect(stats.toolFrequency).toEqual({ repo_agent: 1, learn_concept: 1 }); + // Per-message usage accumulates into the session totals the UI shows. + expect(stats.actualInputTokens).toBe(7005); + expect(stats.actualOutputTokens).toBe(161); + expect(stats.actualCacheReadTokens).toBe(3408); + expect(stats.actualCacheWriteTokens).toBe(3587); + expect(stats.estimatedTokens).toBeGreaterThan(0); + }); +}); + +describe("writeCanvasAgentLog", () => { + const args = { + workspaceId: "ws-1", + conversationId: "conv-1", + model: "claude-sonnet-5", + provider: "anthropic", + startedAt: new Date("2026-07-28T10:00:00.000Z"), + completedAt: new Date("2026-07-28T10:00:06.000Z"), + }; + + test("creates a row with the transcript, stats and timing on first turn", async () => { + findUnique.mockResolvedValue({ messages: CONVERSATION }); + findFirst.mockResolvedValue(null); + + await writeCanvasAgentLog(args); + + expect(put).toHaveBeenCalledTimes(1); + const [blobPath, blobBody] = put.mock.calls[0]; + expect(blobPath).toBe(`agent-logs/ws-1/conv-1/${CANVAS_AGENT_NAME}.json`); + + // The blob is the real transcript, not an empty array. + const parsed = JSON.parse(blobBody as string); + expect(parsed.sessionId).toBe("conv-1"); + expect(parsed.messages).toHaveLength(4); + expect(parsed.config).toMatchObject({ + model: "claude-sonnet-5", + provider: "anthropic", + source: "canvas_chat", + }); + + expect(update).not.toHaveBeenCalled(); + expect(create).toHaveBeenCalledTimes(1); + const row = create.mock.calls[0][0].data; + expect(row).toMatchObject({ + agent: CANVAS_AGENT_NAME, + workspaceId: "ws-1", + sessionId: "conv-1", + provider: "anthropic", + source: "canvas_chat", + startedAt: args.startedAt, + completedAt: args.completedAt, + blobUrl: "https://blob.example/canvas.json", + }); + expect(row.stats).toMatchObject({ + totalToolCalls: 2, + actualInputTokens: 7005, + actualCacheReadTokens: 3408, + }); + }); + + test("updates the existing row on later turns instead of adding another", async () => { + findUnique.mockResolvedValue({ messages: CONVERSATION }); + findFirst.mockResolvedValue({ id: "log-1" }); + + await writeCanvasAgentLog(args); + + expect(create).not.toHaveBeenCalled(); + expect(update).toHaveBeenCalledTimes(1); + expect(update.mock.calls[0][0]).toMatchObject({ where: { id: "log-1" } }); + // startedAt belongs to the run, so it is not rewritten each turn. + expect(update.mock.calls[0][0].data.startedAt).toBeUndefined(); + expect(update.mock.calls[0][0].data.completedAt).toBe(args.completedAt); + }); + + test("keys the lookup on agent + workspace + conversation", async () => { + findUnique.mockResolvedValue({ messages: CONVERSATION }); + findFirst.mockResolvedValue(null); + + await writeCanvasAgentLog(args); + + expect(findFirst.mock.calls[0][0].where).toEqual({ + agent: CANVAS_AGENT_NAME, + workspaceId: "ws-1", + sessionId: "conv-1", + }); + }); + + test("writes nothing for a conversation with no messages", async () => { + findUnique.mockResolvedValue({ messages: [] }); + + await writeCanvasAgentLog(args); + + expect(put).not.toHaveBeenCalled(); + expect(create).not.toHaveBeenCalled(); + expect(update).not.toHaveBeenCalled(); + }); + + test("swallows a blob failure so the finished stream is unaffected", async () => { + findUnique.mockResolvedValue({ messages: CONVERSATION }); + put.mockRejectedValue(new Error("blob down")); + + await expect(writeCanvasAgentLog(args)).resolves.toBeUndefined(); + expect(create).not.toHaveBeenCalled(); + }); +}); diff --git a/src/__tests__/unit/services/canvas-turn-persistence.test.ts b/src/__tests__/unit/services/canvas-turn-persistence.test.ts index 99bd7adc00..ddfe870d69 100644 --- a/src/__tests__/unit/services/canvas-turn-persistence.test.ts +++ b/src/__tests__/unit/services/canvas-turn-persistence.test.ts @@ -35,6 +35,7 @@ import { appendTurnMessages, fetchStoredConversationMessages, normalizeStoredAttachments, + tokenUsageFrom, } from "@/services/canvas-turn-persistence"; const queryRaw = db.$queryRaw as ReturnType; @@ -121,6 +122,182 @@ describe("messagesFromSteps", () => { errorText: "Tool call failed", }); }); + + test("stamps each row with its own step's response time, not persist time", () => { + // A turn whose two model calls were 5s apart. `messagesFromSteps` + // runs in after(), long after both — the stamps must still be the + // step times, not one shared persist-time value. + const steps = [ + { + text: "Looking it up.", + toolCalls: [{ toolCallId: "tc1", toolName: "search", input: {} }], + toolResults: [{ toolCallId: "tc1", output: { ok: true } }], + response: { timestamp: new Date("2026-07-28T10:00:00.000Z") }, + }, + { + text: "Here's the answer.", + response: { timestamp: new Date("2026-07-28T10:00:05.000Z") }, + }, + ]; + + const rows = messagesFromSteps(steps, "turn-1-a"); + + // Step 0 produced a text row and a tool row. Same model call, so + // they legitimately share a timestamp; step 1 must not. + expect(rows[0].timestamp).toBe("2026-07-28T10:00:00.000Z"); + expect(rows[1].timestamp).toBe("2026-07-28T10:00:00.000Z"); + expect(rows[2].timestamp).toBe("2026-07-28T10:00:05.000Z"); + + const spread = + new Date(rows[2].timestamp!).getTime() - + new Date(rows[0].timestamp!).getTime(); + expect(spread).toBe(5000); + }); + + test("falls back to now when a step carries no response metadata", () => { + const before = Date.now(); + const rows = messagesFromSteps([{ text: "No metadata here." }], "turn-1-a"); + const stamped = new Date(rows[0].timestamp!).getTime(); + + expect(stamped).toBeGreaterThanOrEqual(before); + expect(stamped).toBeLessThanOrEqual(Date.now()); + }); + + test("attaches each step's usage to that step's first row only", () => { + const steps = [ + { + text: "Working.", + toolCalls: [{ toolCallId: "tc1", toolName: "search", input: {} }], + toolResults: [{ toolCallId: "tc1", output: { ok: true } }], + usage: { + inputTokens: 3411, + outputTokens: 102, + inputTokenDetails: { cacheReadTokens: 0, cacheWriteTokens: 3408 }, + }, + }, + { + text: "Done.", + usage: { + inputTokens: 3594, + outputTokens: 59, + inputTokenDetails: { cacheReadTokens: 3408, cacheWriteTokens: 179 }, + }, + }, + ]; + + const rows = messagesFromSteps(steps, "turn-1-a"); + + expect(rows[0].usage).toEqual({ + inputTokens: 3411, + outputTokens: 102, + cacheReadTokens: 0, + cacheWriteTokens: 3408, + }); + // Same model call as rows[0] — must not be counted twice. + expect(rows[1].usage).toBeUndefined(); + expect(rows[2].usage).toEqual({ + inputTokens: 3594, + outputTokens: 59, + cacheReadTokens: 3408, + cacheWriteTokens: 179, + }); + + // Summing the rows must reproduce the turn total exactly. + const sum = (k: "inputTokens" | "cacheReadTokens" | "cacheWriteTokens") => + rows.reduce((a, r) => a + (r.usage?.[k] ?? 0), 0); + expect(sum("inputTokens")).toBe(7005); + expect(sum("cacheReadTokens")).toBe(3408); + expect(sum("cacheWriteTokens")).toBe(3587); + }); + + test("omits usage entirely for steps that report none", () => { + const rows = messagesFromSteps([{ text: "No usage." }], "turn-1-a"); + expect(rows[0].usage).toBeUndefined(); + }); + + test("carries usage forward when a step is stripped down to nothing", () => { + // The auto-turn's control tool is `stay_silent`, so a step whose only + // output is stripped still cost real tokens. Those must not vanish. + const steps = [ + { + toolCalls: [{ toolCallId: "s1", toolName: "stay_silent", input: {} }], + usage: { + inputTokens: 5000, + outputTokens: 20, + inputTokenDetails: { cacheReadTokens: 4900, cacheWriteTokens: 0 }, + }, + }, + { + text: "Actually, here's something.", + usage: { inputTokens: 100, outputTokens: 10, inputTokenDetails: {} }, + }, + ]; + + const rows = messagesFromSteps(steps, "turn-1-a", new Set(["stay_silent"])); + + expect(rows).toHaveLength(1); + expect(rows[0].usage).toEqual({ + inputTokens: 5100, + outputTokens: 30, + cacheReadTokens: 4900, + cacheWriteTokens: 0, + }); + }); + + test("carries usage onto the last row when the final step emits nothing", () => { + const steps = [ + { + text: "Done.", + usage: { inputTokens: 100, outputTokens: 10, inputTokenDetails: {} }, + }, + { + toolCalls: [{ toolCallId: "s1", toolName: "stay_silent", input: {} }], + usage: { inputTokens: 200, outputTokens: 5, inputTokenDetails: {} }, + }, + ]; + + const rows = messagesFromSteps(steps, "turn-1-a", new Set(["stay_silent"])); + + expect(rows).toHaveLength(1); + expect(rows[0].usage).toMatchObject({ inputTokens: 300, outputTokens: 15 }); + }); +}); + +describe("tokenUsageFrom", () => { + test("reads the AI SDK v6 usage shape", () => { + // Verbatim shape produced by @ai-sdk/anthropic through streamText. + // The cache counts live under `inputTokenDetails`; reading + // `cacheReadInputTokens` / `cacheWriteInputTokens` off the top level + // silently yields undefined, which is what this guards against. + const usage = { + inputTokens: 6873, + inputTokenDetails: { + noCacheTokens: 2, + cacheReadTokens: 6059, + cacheWriteTokens: 812, + }, + outputTokens: 47, + outputTokenDetails: {}, + totalTokens: 6920, + cachedInputTokens: 6059, + }; + + expect(tokenUsageFrom(usage)).toEqual({ + inputTokens: 6873, + outputTokens: 47, + cacheReadTokens: 6059, + cacheWriteTokens: 812, + }); + }); + + test("returns undefined when there is nothing to record", () => { + expect(tokenUsageFrom(undefined)).toBeUndefined(); + expect(tokenUsageFrom({})).toBeUndefined(); + }); + + test("keeps whichever fields are present", () => { + expect(tokenUsageFrom({ outputTokens: 12 })).toEqual({ outputTokens: 12 }); + }); }); describe("appendTurnMessages", () => { diff --git a/src/app/api/ask/quick/route.ts b/src/app/api/ask/quick/route.ts index 71d4c3c361..25fe62a33c 100644 --- a/src/app/api/ask/quick/route.ts +++ b/src/app/api/ask/quick/route.ts @@ -1,7 +1,12 @@ import { NextRequest, NextResponse, after } from "next/server"; import { validationError, serverError, forbiddenError, isApiError } from "@/types/errors"; import { validateUserBelongsToOrg, validateWorkspaceAccess } from "@/services/workspace"; -import { ModelMessage } from "ai"; +import { + ModelMessage, + createUIMessageStream, + createUIMessageStreamResponse, + type UIMessageStreamWriter, +} from "ai"; import { getMiddlewareContext } from "@/lib/middleware/utils"; import { getBaseUrl } from "@/lib/utils"; import { resolveWorkspaceAccess } from "@/lib/auth/workspace-access"; @@ -26,9 +31,11 @@ import { appendTurnMessages, fetchStoredConversationMessages, normalizeStoredAttachments, + tokenUsageFrom, type StoredMessage, type StoredAttachment, } from "@/services/canvas-turn-persistence"; +import { writeCanvasAgentLog } from "@/services/canvas-agent-log"; import { buildDeferredCheckTools } from "@/lib/ai/deferredCheckTools"; import { resolveOrgConversationRowId, @@ -487,6 +494,10 @@ export async function POST(request: NextRequest) { // internally-wired dispatch_graph_walk tool; consumed in after() // to schedule one graph-walk sub-agent worker per dispatched intent. const dispatchedGraphWalks: DispatchedGraphWalkIntent[] = []; + const turnStartedAt = new Date(); + // Set in `createUIMessageStream`'s execute, before any step finishes. + let usageWriter: UIMessageStreamWriter | undefined; + let stepNumber = 0; const tAgent = Date.now(); const { @@ -574,6 +585,26 @@ export async function POST(request: NextRequest) { onStepFinish: (sf) => { const conceptIds = extractConceptIdsFromStep(sf.content); conceptIds.forEach((id) => learnedConceptIds.add(id)); + + // Live ticker only; the durable copy is written per row by + // `messagesFromSteps` at the end of the turn. + const usage = tokenUsageFrom( + sf.usage as Parameters[0], + ); + if (usage && usageWriter) { + usageWriter.write({ + type: "data-usage", + transient: true, + data: { + stepNumber: stepNumber++, + timestamp: ( + sf.response?.timestamp ?? new Date() + ).toISOString(), + elapsedMs: Date.now() - turnStartedAt.getTime(), + usage, + }, + }); + } }, onFinish: async ({ usage }) => { // Persist the turn's token usage to the conversation row so @@ -581,15 +612,15 @@ export async function POST(request: NextRequest) { // the next request. Best-effort; failures are logged but // do not surface to the user — the stream already finished. if (!tokenAttributionRowId) return; - const u = usage as - | { inputTokens?: number; outputTokens?: number } - | undefined; - const inputTokens = Number(u?.inputTokens ?? 0); - const outputTokens = Number(u?.outputTokens ?? 0); + const u = tokenUsageFrom( + usage as Parameters[0], + ); await recordTurnTokens({ conversationId: tokenAttributionRowId, - inputTokens, - outputTokens, + inputTokens: u?.inputTokens ?? 0, + outputTokens: u?.outputTokens ?? 0, + cacheReadTokens: u?.cacheReadTokens ?? 0, + cacheWriteTokens: u?.cacheWriteTokens ?? 0, }); }, }, @@ -675,6 +706,15 @@ export async function POST(request: NextRequest) { idPrefix: assistantPrefix, reason: "user-turn", }); + if (primaryWorkspaceId) { + await writeCanvasAgentLog({ + workspaceId: primaryWorkspaceId, + conversationId: rowId, + model: chatAgentModel, + startedAt: turnStartedAt, + completedAt: new Date(), + }); + } } catch (err) { console.error("❌ [quick-ask] Turn persist failed:", err); // Persist a trailing error row so a reopened tab sees the @@ -769,7 +809,35 @@ export async function POST(request: NextRequest) { }); console.log("[quick-ask] timing", { stage: "setup-to-stream", ms: Date.now() - t0, workspaces: slugs, orgId: orgId ?? null }); - return result.toUIMessageStreamResponse({ + + // By default the AI SDK masks mid-stream errors as the literal + // string "An error occurred." — useless for diagnosis and + // indistinguishable from a clean finish on the client. Forward + // the real message instead. `runCanvasAgent`'s `onError` logs the + // full error + stack server-side; this surfaces a readable + // message to the chat so the user sees *why* it failed rather + // than a generic fallback. + const onStreamError = (error: unknown) => { + const message = error instanceof Error ? error.message : String(error); + console.error("❌ [quick-ask] Mid-stream error:", { + workspaces: slugs, + message, + }); + return message; + }; + + const stream = createUIMessageStream({ + execute: ({ writer }) => { + usageWriter = writer; + writer.merge( + result.toUIMessageStream({ onError: onStreamError }), + ); + }, + onError: onStreamError, + }); + + return createUIMessageStreamResponse({ + stream, // Hand the server-created/validated org-canvas row id back to the // client (same pattern as `X-Approval-Result`) so it can stamp // `serverConversationId` on the first turn without a separate POST. @@ -781,22 +849,6 @@ export async function POST(request: NextRequest) { }, } : {}), - // By default the AI SDK masks mid-stream errors as the literal - // string "An error occurred." — useless for diagnosis and - // indistinguishable from a clean finish on the client. Forward - // the real message instead. `runCanvasAgent`'s `onError` logs the - // full error + stack server-side; this surfaces a readable - // message to the chat so the user sees *why* it failed rather - // than a generic fallback. - onError: (error) => { - const message = - error instanceof Error ? error.message : String(error); - console.error("❌ [quick-ask] Mid-stream error:", { - workspaces: slugs, - message, - }); - return message; - }, }); } catch (streamError) { // Preserve typed ApiError statuses (forbidden, notFound, diff --git a/src/app/org/[githubLogin]/_components/SidebarChat.tsx b/src/app/org/[githubLogin]/_components/SidebarChat.tsx index cb964f4781..c4910f687e 100644 --- a/src/app/org/[githubLogin]/_components/SidebarChat.tsx +++ b/src/app/org/[githubLogin]/_components/SidebarChat.tsx @@ -550,6 +550,7 @@ export function SidebarChat({ githubLogin }: SidebarChatProps) { timeline: filteredTimeline, isStreaming: isMessageStreaming, usage: message.usage, + elapsedMs: message.elapsedMs, }} /> diff --git a/src/app/org/[githubLogin]/_state/canvasChatStore.ts b/src/app/org/[githubLogin]/_state/canvasChatStore.ts index 6915f25689..83eb10e02b 100644 --- a/src/app/org/[githubLogin]/_state/canvasChatStore.ts +++ b/src/app/org/[githubLogin]/_state/canvasChatStore.ts @@ -208,11 +208,12 @@ export interface CanvasChatMessage { /** Files attached by the user before send. Persisted in SharedConversation JSON. */ attachments?: CanvasAttachment[]; /** - * Per-turn aggregated token usage, stamped by `useSendCanvasChatMessage` - * onto the last tool-call batch message when the stream `finish` event - * carries usage data. Live-stream only — not persisted to `SharedConversation`. + * Running token usage, stamped by `useSendCanvasChatMessage` onto the last + * tool-call batch message as `data-usage` parts arrive. Live-stream only — + * the durable per-step copy lives on `SharedConversation.messages`. */ usage?: TokenUsage; + elapsedMs?: number; /** * Synthetic assistant message describing an approval outcome. Set by * `/api/ask/quick` after `handleApproval` creates the DB row; carries diff --git a/src/app/org/[githubLogin]/_state/useSendCanvasChatMessage.ts b/src/app/org/[githubLogin]/_state/useSendCanvasChatMessage.ts index e3a60b905b..32d8254eb4 100644 --- a/src/app/org/[githubLogin]/_state/useSendCanvasChatMessage.ts +++ b/src/app/org/[githubLogin]/_state/useSendCanvasChatMessage.ts @@ -495,16 +495,20 @@ export function useSendCanvasChatMessage() { }); } - // When streaming is done and the finish event carried usage data, - // stamp it onto the last tool-call batch row so StreamingMessage - // can render TurnTokenUsage + Complete pill. No text-only fallback: - // text-only turns render via SidebarChatMessage which has no - // TurnTokenUsage render path. - if (!updatedMessage.isStreaming && updatedMessage.usage) { + // Stamp the running usage total onto the last tool-call batch row + // so StreamingMessage can render TurnTokenUsage + Complete pill. + // Mid-stream too, so the figure ticks as the turn progresses. No + // text-only fallback: those render via SidebarChatMessage, which + // has no TurnTokenUsage render path. + if (updatedMessage.usage) { for (let i = timelineMessages.length - 1; i >= 0; i--) { const m = timelineMessages[i]; if (m.role === "assistant" && !!m.timeline?.length) { - timelineMessages[i] = { ...m, usage: updatedMessage.usage }; + timelineMessages[i] = { + ...m, + usage: updatedMessage.usage, + elapsedMs: updatedMessage.elapsedMs, + }; break; } } diff --git a/src/components/agent-logs/TurnTokenUsage.tsx b/src/components/agent-logs/TurnTokenUsage.tsx index 91bf65278c..58c670e500 100644 --- a/src/components/agent-logs/TurnTokenUsage.tsx +++ b/src/components/agent-logs/TurnTokenUsage.tsx @@ -20,13 +20,23 @@ function hasAnyToken(usage: UsageData): usage is NonNullable { interface TurnTokenUsageProps { usage: UsageData; + /** Turn duration so far, rendered alongside the token counts. */ + elapsedMs?: number; +} + +/** "820ms" under a second, else "4.9s", else "1m 12s". */ +function formatElapsed(ms: number): string { + if (ms < 1000) return `${Math.round(ms)}ms`; + const secs = ms / 1000; + if (secs < 60) return `${secs.toFixed(1)}s`; + return `${Math.floor(secs / 60)}m ${Math.round(secs % 60)}s`; } /** * Compact per-turn token usage footer for assistant messages in agent logs. * Collapsed by default (single line); click expands to show cache read/write split. */ -export function TurnTokenUsage({ usage }: TurnTokenUsageProps) { +export function TurnTokenUsage({ usage, elapsedMs }: TurnTokenUsageProps) { const [expanded, setExpanded] = useState(false); if (!hasAnyToken(usage)) return null; @@ -52,6 +62,9 @@ export function TurnTokenUsage({ usage }: TurnTokenUsageProps) { {cacheTotal > 0 && ( · cache: {formatTokens(cacheTotal)} )} + {elapsedMs != null && ( + · {formatElapsed(elapsedMs)} + )} {hasCacheSplit && ( expanded diff --git a/src/components/streaming/StreamingMessage.tsx b/src/components/streaming/StreamingMessage.tsx index 357c215019..4a1ea2367b 100644 --- a/src/components/streaming/StreamingMessage.tsx +++ b/src/components/streaming/StreamingMessage.tsx @@ -94,8 +94,8 @@ export function StreamingMessage({ {/* Render timeline items in order */} {regularTimeline?.map((item) => renderTimelineItem(item))} - {!message.isStreaming && message.usage && ( - + {message.usage && ( + )} {shouldShowThinking() && ( diff --git a/src/lib/ai/publicChatBudget.ts b/src/lib/ai/publicChatBudget.ts index d0f8fe68ec..c4e88abb40 100644 --- a/src/lib/ai/publicChatBudget.ts +++ b/src/lib/ai/publicChatBudget.ts @@ -27,6 +27,10 @@ import { getClientIp } from "@/lib/rate-limit"; // output. 5:1 weighting keeps the cap close to actual dollar cost. export const OUTPUT_TOKEN_WEIGHT = 5; +// Relative to fresh input: cache reads ~$0.30/Mtok, writes ~$3.75/Mtok. +export const CACHE_READ_TOKEN_WEIGHT = 0.1; +export const CACHE_WRITE_TOKEN_WEIGHT = 1.25; + // Per-visitor daily cap (~$3-5 of Anthropic spend at Sonnet rates). export const ANON_DAILY_TOKEN_CAP = 100_000; @@ -49,8 +53,21 @@ export function deriveAnonymousId(req: NextRequest): string { return createHash("sha256").update(`${ip}|${ua}`).digest("hex").slice(0, 16); } -function weightedCost(input: number, output: number): number { - return input + output * OUTPUT_TOKEN_WEIGHT; +// `input` already includes the cached reads and writes, so they are +// subtracted out and repriced rather than added on top. +function weightedCost( + input: number, + output: number, + cacheRead = 0, + cacheWrite = 0, +): number { + const fresh = Math.max(0, input - cacheRead - cacheWrite); + return ( + fresh + + cacheRead * CACHE_READ_TOKEN_WEIGHT + + cacheWrite * CACHE_WRITE_TOKEN_WEIGHT + + output * OUTPUT_TOKEN_WEIGHT + ); } interface BudgetResult { @@ -79,7 +96,12 @@ export async function checkPublicChatBudget(args: { anonymousId: args.anonymousId, createdAt: { gte: since }, }, - _sum: { inputTokens: true, outputTokens: true }, + _sum: { + inputTokens: true, + outputTokens: true, + cacheReadTokens: true, + cacheWriteTokens: true, + }, }), db.sharedConversation.aggregate({ where: { @@ -87,13 +109,20 @@ export async function checkPublicChatBudget(args: { userId: null, createdAt: { gte: since }, }, - _sum: { inputTokens: true, outputTokens: true }, + _sum: { + inputTokens: true, + outputTokens: true, + cacheReadTokens: true, + cacheWriteTokens: true, + }, }), ]); const anonCost = weightedCost( anonAgg._sum.inputTokens ?? 0, anonAgg._sum.outputTokens ?? 0, + anonAgg._sum.cacheReadTokens ?? 0, + anonAgg._sum.cacheWriteTokens ?? 0, ); if (anonCost >= ANON_DAILY_TOKEN_CAP) { return { allowed: false, reason: "anon", retryAfterSecs: 24 * 60 * 60 }; @@ -102,6 +131,8 @@ export async function checkPublicChatBudget(args: { const wsCost = weightedCost( wsAgg._sum.inputTokens ?? 0, wsAgg._sum.outputTokens ?? 0, + wsAgg._sum.cacheReadTokens ?? 0, + wsAgg._sum.cacheWriteTokens ?? 0, ); if (wsCost >= WORKSPACE_PUBLIC_DAILY_TOKEN_CAP) { return { allowed: false, reason: "workspace", retryAfterSecs: 24 * 60 * 60 }; @@ -120,14 +151,27 @@ export async function recordTurnTokens(args: { conversationId: string; inputTokens: number; outputTokens: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; }): Promise { - if (!args.inputTokens && !args.outputTokens) return; + const cacheReadTokens = args.cacheReadTokens ?? 0; + const cacheWriteTokens = args.cacheWriteTokens ?? 0; + if ( + !args.inputTokens && + !args.outputTokens && + !cacheReadTokens && + !cacheWriteTokens + ) { + return; + } try { await db.sharedConversation.update({ where: { id: args.conversationId }, data: { inputTokens: { increment: Math.max(0, args.inputTokens) }, outputTokens: { increment: Math.max(0, args.outputTokens) }, + cacheReadTokens: { increment: Math.max(0, cacheReadTokens) }, + cacheWriteTokens: { increment: Math.max(0, cacheWriteTokens) }, }, }); } catch (error) { diff --git a/src/lib/ai/runCanvasAgent.ts b/src/lib/ai/runCanvasAgent.ts index 5ca0861342..99a8e85f09 100644 --- a/src/lib/ai/runCanvasAgent.ts +++ b/src/lib/ai/runCanvasAgent.ts @@ -147,7 +147,11 @@ export interface CanvasAgentHooks { * uses this for token-attribution prep work; programmatic callers * typically don't need it. */ - onStepFinish?: (sf: { content: unknown }) => void | Promise; + onStepFinish?: (sf: { + content: unknown; + usage?: unknown; + response?: { timestamp?: Date }; + }) => void | Promise; /** * Called once when the stream finishes successfully. Receives the * final `usage` so the caller can record token spend. diff --git a/src/lib/streaming/useStreamProcessor.ts b/src/lib/streaming/useStreamProcessor.ts index 638b9af798..7e8f52f95f 100644 --- a/src/lib/streaming/useStreamProcessor.ts +++ b/src/lib/streaming/useStreamProcessor.ts @@ -1,6 +1,6 @@ import { useCallback } from "react"; import type { BaseStreamingMessage, StreamProcessorConfig, StreamEvent, ToolCallStatus } from "@/types/streaming"; -import type { TokenUsage } from "@/types/usage"; +import { addUsage, type TokenUsage } from "@/types/usage"; import { DEFAULT_DEBOUNCE_MS } from "./constants"; import { parseSSELine } from "./helpers"; @@ -72,6 +72,7 @@ export function useStreamProcessor = []; // Unified timeline let error: string | undefined; let capturedUsage: TokenUsage | undefined; + let capturedElapsedMs: number | undefined; // Track text part sequence to generate unique IDs when stream reuses IDs let textPartSequence = 0; @@ -130,6 +131,7 @@ export function useStreamProcessor { + const base: ParsedMessage = { + role: m.role, + timestamp: m.timestamp ?? null, + }; + if (m.usage) base.usage = m.usage; + + if (m.toolCalls && m.toolCalls.length > 0) { + base.content = m.toolCalls.map( + (tc): ToolCallContent => ({ + type: "tool-call", + toolCallId: tc.id, + toolName: tc.toolName, + input: tc.input, + }), + ); + return base; + } + + base.content = m.content; + return base; + }); +} + +export async function writeCanvasAgentLog(args: { + workspaceId: string; + conversationId: string; + model?: string; + startedAt: Date; + completedAt: Date; +}): Promise { + const { workspaceId, conversationId, model, startedAt, completedAt } = args; + + try { + const row = await db.sharedConversation.findUnique({ + where: { id: conversationId }, + select: { messages: true }, + }); + const messages = Array.isArray(row?.messages) + ? (row.messages as unknown as StoredMessage[]) + : []; + if (messages.length === 0) return; + + const config = { + ...(model ? { model } : {}), + provider: CANVAS_AGENT_PROVIDER, + source: CANVAS_AGENT_SOURCE, + }; + const logContent = JSON.stringify({ + sessionId: conversationId, + messages: transcriptFromStoredMessages(messages), + config, + }); + + const { stats } = parseAgentLogStats(logContent); + + const blob = await put( + `agent-logs/${workspaceId}/${conversationId}/${CANVAS_AGENT_NAME}.json`, + logContent, + { + access: "private", + contentType: "application/json", + addRandomSuffix: false, + allowOverwrite: true, + }, + ); + + const existing = await db.agentLog.findFirst({ + where: { + agent: CANVAS_AGENT_NAME, + workspaceId, + sessionId: conversationId, + }, + select: { id: true }, + }); + + const data = { + blobUrl: blob.url, + config: config as Prisma.InputJsonValue, + stats: stats as unknown as Prisma.InputJsonValue, + provider: CANVAS_AGENT_PROVIDER, + source: CANVAS_AGENT_SOURCE, + completedAt, + }; + + if (existing) { + await db.agentLog.update({ where: { id: existing.id }, data }); + } else { + await db.agentLog.create({ + data: { + ...data, + agent: CANVAS_AGENT_NAME, + workspaceId, + sessionId: conversationId, + startedAt, + }, + }); + } + } catch (error) { + console.error("[canvas-agent-log] write failed", { + conversationId, + error: error instanceof Error ? error.message : String(error), + }); + } +} diff --git a/src/services/canvas-turn-persistence.ts b/src/services/canvas-turn-persistence.ts index 3042822368..38ca64a780 100644 --- a/src/services/canvas-turn-persistence.ts +++ b/src/services/canvas-turn-persistence.ts @@ -24,6 +24,7 @@ */ import { db } from "@/lib/db"; +import { addUsage, type TokenUsage } from "@/types/usage"; import { notifyCanvasConversationUpdated, type CanvasConversationUpdateReason, @@ -94,6 +95,7 @@ export interface StoredMessage { role: "user" | "assistant"; content: string; timestamp?: string; + usage?: TokenUsage; toolCalls?: StoredToolCall[]; attachments?: StoredAttachment[]; source?: { kind: string; featureId?: string; plannerMessageId?: string }; @@ -115,14 +117,50 @@ export interface StoredMessage { }; } +type ModelUsageLike = + | { + inputTokens?: number; + outputTokens?: number; + inputTokenDetails?: { + cacheReadTokens?: number; + cacheWriteTokens?: number; + }; + } + | undefined; + type StepLike = { text?: string; toolCalls?: Array<{ toolCallId: string; toolName: string; input?: unknown }>; toolResults?: Array<{ toolCallId: string; output?: unknown; result?: unknown }>; + usage?: ModelUsageLike; + response?: { timestamp?: Date }; }; const NO_STRIP: ReadonlySet = new Set(); +// Steps synthesised by hand (sub-agent paths) carry no response metadata. +function stepTimestamp(step: StepLike): string { + const t = step.response?.timestamp; + if (t instanceof Date && !Number.isNaN(t.getTime())) return t.toISOString(); + return new Date().toISOString(); +} + +export function tokenUsageFrom(usage: ModelUsageLike): TokenUsage | undefined { + if (!usage) return undefined; + const out: TokenUsage = {}; + if (typeof usage.inputTokens === "number") out.inputTokens = usage.inputTokens; + if (typeof usage.outputTokens === "number") { + out.outputTokens = usage.outputTokens; + } + if (typeof usage.inputTokenDetails?.cacheReadTokens === "number") { + out.cacheReadTokens = usage.inputTokenDetails.cacheReadTokens; + } + if (typeof usage.inputTokenDetails?.cacheWriteTokens === "number") { + out.cacheWriteTokens = usage.inputTokenDetails.cacheWriteTokens; + } + return Object.keys(out).length > 0 ? out : undefined; +} + /** * Reconstruct the agent's output as `CanvasChatMessage`-shaped rows from * the finished stream's `steps`. Mirrors the client-side timeline split @@ -146,12 +184,17 @@ export function messagesFromSteps( const rows: StoredMessage[] = []; let idx = 0; const nextId = () => `${idPrefix}${idx++}`; - const now = new Date().toISOString(); + // Usage from steps that emitted no rows (a stripped control tool is the + // whole step), held over so those tokens land on the next row instead of + // being dropped. + let carriedUsage: TokenUsage | undefined; for (const step of steps) { // Extract any schedule_check result from this step so it can be // attached to the text row as `deferredCheck` metadata. const deferredCheck = extractDeferredCheckFromStep(step); + const now = stepTimestamp(step); + const firstRowOfStep = rows.length; if (step.text && step.text.trim()) { const textRow: StoredMessage = { @@ -167,51 +210,65 @@ export function messagesFromSteps( } const calls = step.toolCalls ?? []; - if (calls.length === 0) continue; + if (calls.length > 0) { + const resultByCallId = new Map( + (step.toolResults ?? []).map((r) => [r.toolCallId, r] as const), + ); - const resultByCallId = new Map( - (step.toolResults ?? []).map((r) => [r.toolCallId, r] as const), - ); + const toolCalls: StoredToolCall[] = calls + .filter((tc) => !stripToolNames.has(tc.toolName)) + .map((tc) => { + const r = resultByCallId.get(tc.toolCallId); + const output = r ? (r.output ?? r.result) : undefined; + const isError = + !!output && + typeof output === "object" && + "error" in (output as Record); + return { + id: tc.toolCallId, + toolName: tc.toolName, + input: tc.input, + output, + status: + output === undefined + ? "input-available" + : isError + ? "output-error" + : "output-available", + ...(isError ? { errorText: "Tool call failed" } : {}), + }; + }); - const toolCalls: StoredToolCall[] = calls - .filter((tc) => !stripToolNames.has(tc.toolName)) - .map((tc) => { - const r = resultByCallId.get(tc.toolCallId); - const output = r ? (r.output ?? r.result) : undefined; - const isError = - !!output && - typeof output === "object" && - "error" in (output as Record); - return { - id: tc.toolCallId, - toolName: tc.toolName, - input: tc.input, - output, - status: - output === undefined - ? "input-available" - : isError - ? "output-error" - : "output-available", - ...(isError ? { errorText: "Tool call failed" } : {}), + if (toolCalls.length > 0) { + const toolRow: StoredMessage = { + id: nextId(), + role: "assistant", + content: "", + timestamp: now, + toolCalls, }; - }); - - if (toolCalls.length > 0) { - const toolRow: StoredMessage = { - id: nextId(), - role: "assistant", - content: "", - timestamp: now, - toolCalls, - }; - // If there was no text in this step, attach deferredCheck to the - // tool-call row instead so the card is always anchored somewhere. - if (deferredCheck && rows[rows.length - 1]?.deferredCheck == null) { - toolRow.deferredCheck = deferredCheck; + // If there was no text in this step, attach deferredCheck to the + // tool-call row instead so the card is always anchored somewhere. + if (deferredCheck && rows[rows.length - 1]?.deferredCheck == null) { + toolRow.deferredCheck = deferredCheck; + } + rows.push(toolRow); } - rows.push(toolRow); } + + // One model call per step: attach its usage once, not per row. + const usage = addUsage(carriedUsage, tokenUsageFrom(step.usage)); + if (usage && rows.length > firstRowOfStep) { + rows[firstRowOfStep].usage = usage; + carriedUsage = undefined; + } else { + carriedUsage = usage; + } + } + + if (carriedUsage && rows.length > 0) { + const last = rows[rows.length - 1]; + last.usage = addUsage(last.usage, carriedUsage); } return rows; diff --git a/src/types/streaming.ts b/src/types/streaming.ts index febb2013fa..615716d56f 100644 --- a/src/types/streaming.ts +++ b/src/types/streaming.ts @@ -17,7 +17,8 @@ export type StreamEventType = | "tool-error" | "start" | "finish" - | "error"; + | "error" + | "data-usage"; export type ToolCallStatus = | "input-start" @@ -144,7 +145,21 @@ export interface FinishEvent extends BaseStreamEvent { }; } +/** Per-step token usage, emitted as each model call in the turn completes. */ +export interface StepUsageEvent extends BaseStreamEvent { + type: "data-usage"; + data: { + stepNumber: number; + /** When this model call started, per the provider. */ + timestamp: string; + /** Wall-clock since the turn began, so the UI needn't mix clocks. */ + elapsedMs: number; + usage: TokenUsage; + }; +} + export type StreamEvent = + | StepUsageEvent | TextStartEvent | TextDeltaEvent | ReasoningStartEvent @@ -204,6 +219,7 @@ export interface BaseStreamingMessage { error?: string; /** Per-turn aggregated token usage, populated from the `finish` SSE event. */ usage?: TokenUsage; + elapsedMs?: number; } // Tool processor function type diff --git a/src/types/usage.ts b/src/types/usage.ts index 2c1503c2c1..de623666d4 100644 --- a/src/types/usage.ts +++ b/src/types/usage.ts @@ -11,3 +11,25 @@ export interface TokenUsage { cacheReadTokens?: number; cacheWriteTokens?: number; } + +const USAGE_KEYS = [ + "inputTokens", + "outputTokens", + "cacheReadTokens", + "cacheWriteTokens", +] as const; + +/** Field-wise sum. Here, not in a service, so the client can import it. */ +export function addUsage( + a: TokenUsage | undefined, + b: TokenUsage | undefined, +): TokenUsage | undefined { + if (!a) return b; + if (!b) return a; + const out: TokenUsage = { ...a }; + for (const key of USAGE_KEYS) { + const add = b[key]; + if (typeof add === "number") out[key] = (out[key] ?? 0) + add; + } + return out; +}