From 741ff416d684c145e5b20f89346ec2e790396500 Mon Sep 17 00:00:00 2001 From: amosblomqvist Date: Fri, 11 Sep 2026 04:16:05 +0300 Subject: [PATCH] Fix turn-usage keying: idempotent per request turn id, not messageId MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live smoke exposed two bugs in the turn usage log: - onFinish messages carry no ids (withMessageIds fills copies inside persistTurn), so fresh turns keyed messageId "" — provenance lost. - The SDK coalesces continuation turns into ONE response message (a whole probe conversation lands in a single assistant message), so per-messageId idempotency swallowed every continuation turn's usage. TurnUsage gains id: string (unique per request, the append idempotency key); messageId stays as transcript provenance, taken from the id-filled list. appendTurnUsage dedupes by id; continuation turns with the same messageId each keep their usage. --- src/core/contracts/chat.ts | 9 +++++++-- src/core/ports.ts | 5 +++-- src/server/chat-handler.ts | 23 ++++++++++++++--------- src/store/chats.ts | 7 ++++--- tests/render-stats-bar.test.tsx | 1 + tests/usage.test.ts | 14 +++++++++----- 6 files changed, 38 insertions(+), 21 deletions(-) diff --git a/src/core/contracts/chat.ts b/src/core/contracts/chat.ts index 3fd08c9..8baed80 100644 --- a/src/core/contracts/chat.ts +++ b/src/core/contracts/chat.ts @@ -93,8 +93,13 @@ export interface CallUsage { * agent steps, and each step bills separately. Capture lands at the persist * seam (server); cost is INFERRED at read time from core/prices.ts. */ export interface TurnUsage { - /** The assistant message the turn produced — keys this entry to the - * transcript (the block this turn "belongs" to). */ + /** Unique per REQUEST (turn): one invocation of the handler — the + * idempotency key for the append-only log. Continuation turns share a + * response message id (the SDK coalesces them); keying on messageId + * would swallow their usage. */ + id: string; + /** The assistant message the turn produced — provenance, keys this + * entry to the transcript (the block this turn "belongs" to). */ messageId: string; at: string; /** The OpenRouter model id that ran the turn. */ diff --git a/src/core/ports.ts b/src/core/ports.ts index b612877..c3b347b 100644 --- a/src/core/ports.ts +++ b/src/core/ports.ts @@ -39,8 +39,9 @@ export interface ChatStore { /** Persist the session plan tree (plan tools call this in execute). */ saveTree(id: string, tree: PlanTree): void; /** Append one usage entry to the record's turn log (the money channel, - * 2026-09): idempotent per messageId — a turn persisted twice (abort - * replays, wake-turn retries) appends exactly once. */ + * 2026-09): idempotent per turn id — each request carries its own, so + * a turn persisted twice appends once while continuation turns (which + * share a response message id) each keep their usage. */ appendTurnUsage(id: string, entry: TurnUsage): void; /** Append one extraction-run usage entry — idempotent per startedAt. */ appendExtractionUsage(id: string, entry: ExtractionRunUsage): void; diff --git a/src/server/chat-handler.ts b/src/server/chat-handler.ts index 4e9809f..69b1d33 100644 --- a/src/server/chat-handler.ts +++ b/src/server/chat-handler.ts @@ -3,12 +3,10 @@ import { createTeacherAgent, type TeacherAgent } from "../agent/teacher-agent.ts import { createJobRunner } from "../agent/subagents/runner.ts"; import { createExtractionPipeline } from "../agent/subagents/extraction.ts"; import { createTeacherTools } from "../agent/tools.ts"; -import { modelFor } from "../agent/models.ts"; -import { MODELS } from "../../config.ts"; -import { claimSteers, deliverJobResult, setJobRespawn, setTurnRunner } from "./jobs.ts"; +import { MODELS } from "../../config.ts";import { claimSteers, deliverJobResult, setJobRespawn, setTurnRunner } from "./jobs.ts"; import { stores } from "./stores.ts"; import { beginTurn, endTurn } from "./turn-liveness.ts"; -import { buildTeacherContext, persistTurn, triggerExtraction } from "./turn-service.ts"; +import { buildTeacherContext, persistTurn, triggerExtraction, withMessageIds } from "./turn-service.ts"; import type { CallUsage } from "../core/contracts.ts"; /** @@ -98,8 +96,10 @@ export async function handleChatRequest(req: Request, agent?: TeacherAgent): Pro // MONEY CHANNEL (wave 4): one collector per request, fed by the agent's // step callback. Every model call's usage lands in the chat record's - // append-only turns log at onFinish — idempotent per response message id, - // so abort replays and wake retries append exactly once. + // append-only turns log at onFinish — idempotent per turn id, so a + // replayed persist appends once while continuation turns (which share a + // response message id — the SDK coalesces them) each keep their usage. + const turnUsageId = crypto.randomUUID(); const steps: CallUsage[] = []; const collectUsage = (step: { usage: LanguageModelUsage }) => { steps.push({ @@ -140,11 +140,16 @@ export async function handleChatRequest(req: Request, agent?: TeacherAgent): Pro }, onFinish: ({ messages, isAborted }) => { persistTurn(stores.chatStore, chatId, messages as UIMessage[], isAborted); - // usage log append (idempotent per response message id) — the - // stats bar refolds this on the client after the turn settles - const response = [...messages].reverse().find((m) => m.role === "assistant"); + // usage log append (idempotent per turn id) — the stats bar + // refolds this on the client after the turn settles. messageId + // comes from the id-FILLED list: onFinish messages carry no ids + // (withMessageIds fills copies in persistTurn, not these); the + // re-fill here is idempotent, so the key is stable. + const filled = withMessageIds(messages as UIMessage[]); + const response = [...filled].reverse().find((m) => m.role === "assistant"); if (response && steps.length > 0) { stores.chatStore.appendTurnUsage(chatId, { + id: turnUsageId, messageId: response.id, at: new Date().toISOString(), model: MODELS.teacher.model, diff --git a/src/store/chats.ts b/src/store/chats.ts index 9f26f2b..5e851d9 100644 --- a/src/store/chats.ts +++ b/src/store/chats.ts @@ -80,12 +80,13 @@ function saveMessages(id: string, messages: UIMessage[]): void { } /** Append one usage entry to the record's turn log — idempotent per - * messageId: abort replays and wake-turn retries persist the same turn - * twice; the second append is a no-op. */ + * turn id: each REQUEST carries its own id, so a replay of the same turn's + * persist appends once, while continuation turns (which share a response + * message id) each keep their usage. */ function appendTurnUsage(id: string, entry: TurnUsage): void { const record = getChat(id); if (!record) throw new Error(`Chat not found: ${id}`); - record.turns = record.turns?.some((t) => t.messageId === entry.messageId) + record.turns = record.turns?.some((t) => t.id === entry.id) ? record.turns : [...(record.turns ?? []), entry]; fs.writeFileSync(chatFile(id), JSON.stringify(record, null, "\t"), "utf-8"); diff --git a/tests/render-stats-bar.test.tsx b/tests/render-stats-bar.test.tsx index b12e388..4b01db3 100644 --- a/tests/render-stats-bar.test.tsx +++ b/tests/render-stats-bar.test.tsx @@ -15,6 +15,7 @@ function step(inputTokens: number, outputTokens: number, over: Record = {}): TurnUsage { return { + id: "t1", messageId: "m1", at: "2026-09-09T12:00:00Z", model: "z-ai/glm-5.3-flash", diff --git a/tests/usage.test.ts b/tests/usage.test.ts index e71839c..6474c9d 100644 --- a/tests/usage.test.ts +++ b/tests/usage.test.ts @@ -62,6 +62,7 @@ test("callCost: unknown model and absent fields are null, never zero", () => { test("turnCost sums steps; sessionCost reports unpriced turns", () => { const turn: TurnUsage = { + id: "t1", messageId: "m1", at: "2026-09-09T00:00:00Z", model: "z-ai/glm-5.3-flash", @@ -69,7 +70,7 @@ test("turnCost sums steps; sessionCost reports unpriced turns", () => { }; assert.ok(turnCost(turn) !== null); const totals = sessionCost({ - turns: [turn, { messageId: "m2", at: "x", model: "not/a-model", steps: [call({ inputTokens: 1, outputTokens: 1 })] }], + turns: [turn, { id: "t2", messageId: "m2", at: "x", model: "not/a-model", steps: [call({ inputTokens: 1, outputTokens: 1 })] }], }); assert.equal(totals.unpriced, 1); assert.ok(totals.cost > 0); @@ -78,8 +79,8 @@ test("turnCost sums steps; sessionCost reports unpriced turns", () => { test("contextWindow: last call's footprint against the model's limit", () => { const cw = contextWindow({ turns: [ - { messageId: "m1", at: "x", model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 9_000, outputTokens: 1_000 })] }, - { messageId: "m2", at: "x", model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 90_000, outputTokens: 10_000 })] }, + { id: "t1", messageId: "m1", at: "x", model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 9_000, outputTokens: 1_000 })] }, + { id: "t2", messageId: "m2", at: "x", model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 90_000, outputTokens: 10_000 })] }, ], }); assert.ok(cw); @@ -108,9 +109,10 @@ test("jobCost: null without usage, computed with it", () => { // ─── the store seam ────────────────────────────────────────────────────────── -test("appendTurnUsage: appends, is idempotent per messageId, survives saveMessages", () => { +test("appendTurnUsage: appends, idempotent per turn id, continuation turns (same message) all kept, survives saveMessages", () => { const record = chatStore.createChat("usage test chat"); const entry: TurnUsage = { + id: "turn_1", messageId: "msg_a", at: "2026-09-09T00:00:00Z", model: "z-ai/glm-5.3-flash", @@ -118,7 +120,9 @@ test("appendTurnUsage: appends, is idempotent per messageId, survives saveMessag }; chatStore.appendTurnUsage(record.id, entry); chatStore.appendTurnUsage(record.id, entry); // replay (abort/retry) → no-op - chatStore.appendTurnUsage(record.id, { ...entry, messageId: "msg_b" }); + // continuation turn: SAME response message (SDK coalesces), new turn id — + // usage must survive, not collapse into the first entry + chatStore.appendTurnUsage(record.id, { ...entry, id: "turn_2" }); let reread = chatStore.getChat(record.id); assert.equal(reread?.turns?.length, 2);