Per-step timestamps, token usage, and AgentLog rows for the canvas agent

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;

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

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<typeof import('ai')>()),
   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' } }),
         ),

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

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<typeof vi.fn>;
+const update = db.sharedConversation.update as ReturnType<typeof vi.fn>;
+
+/** Both the anon and the workspace aggregate resolve to these sums. */
+function withSpend(sum: Record<string, number>) {
+  aggregate.mockResolvedValue({ _sum: sum });
+}
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  update.mockResolvedValue({});
+});
+
 function makeReq(headers: Record<string, string>): 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();
+  });
 });