From ebcc15c68d228bdf2151e74baf9386e92b00672b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Apr 2026 16:18:03 -0400 Subject: [PATCH] fix(runtime): derive unattributed cost from usage status --- src/main/run-usage-status.ts | 4 ++ src/main/worker/entry/index.ts | 1 - src/main/worker/entry/types.ts | 3 +- .../worker/handlers/run-finished-helpers.ts | 7 ++- src/main/worker/runtime/helpers.ts | 1 - src/main/worker/runtime/index.ts | 1 - tests/unit/run-usage-status.test.ts | 8 ++++ tests/unit/worker-message-schema.test.ts | 43 +++++++++++++++++++ 8 files changed, 61 insertions(+), 7 deletions(-) create mode 100644 tests/unit/worker-message-schema.test.ts diff --git a/src/main/run-usage-status.ts b/src/main/run-usage-status.ts index 096707cd..5d32a51d 100644 --- a/src/main/run-usage-status.ts +++ b/src/main/run-usage-status.ts @@ -56,6 +56,10 @@ export function classifyRunUsage(input: RunUsageClassificationInput): RunUsageSt return "pre_first_message"; } +export function isUsageStatusUnattributed(status: RunUsageStatus): boolean { + return status !== "ok"; +} + export function deriveLegacyRunUsageStatus(input: { runStatus: RunStatus; unattributedCost: boolean; diff --git a/src/main/worker/entry/index.ts b/src/main/worker/entry/index.ts index 320c6ad3..f36a0364 100644 --- a/src/main/worker/entry/index.ts +++ b/src/main/worker/entry/index.ts @@ -108,7 +108,6 @@ function cancelRun(runId: string) { signal: null, model: null, usageStatus: "aborted_before_usage", - unattributedCost: true, inputTokens: null, outputTokens: null, }); diff --git a/src/main/worker/entry/types.ts b/src/main/worker/entry/types.ts index 9945ea20..b88535f6 100644 --- a/src/main/worker/entry/types.ts +++ b/src/main/worker/entry/types.ts @@ -53,10 +53,9 @@ export const workerMessageSchema = z.union([ signal: z.string().nullable(), model: z.string().nullable(), usageStatus: runUsageStatusSchema, - unattributedCost: z.boolean(), inputTokens: z.number().nullable(), outputTokens: z.number().nullable(), - }), + }).strict(), z.object({ type: z.literal("agent-event"), runId: z.string().uuid(), diff --git a/src/main/worker/handlers/run-finished-helpers.ts b/src/main/worker/handlers/run-finished-helpers.ts index 5fab1382..0094f2fc 100644 --- a/src/main/worker/handlers/run-finished-helpers.ts +++ b/src/main/worker/handlers/run-finished-helpers.ts @@ -4,6 +4,7 @@ import { handleFailureEscalation, handleSuccessReset } from "../observability/fa import { persistRunEvent } from "../observability/run-event-recorder"; import type { workerMessageSchema } from "../entry/types"; import type { HandlerContext, HandlerDeps, WorkerMessageEvent } from "./handler-context"; +import { isUsageStatusUnattributed } from "../../run-usage-status"; type RunFinishedEvent = Extract; type RunRecord = ReturnType; @@ -33,6 +34,8 @@ export function finalizeRun( return null; } + const unattributedCost = isUsageStatusUnattributed(event.usageStatus); + persistRunEvent(db, logger, { companyId: run.companyId, runId: run.id, @@ -48,7 +51,7 @@ export function finalizeRun( signal: event.signal, model: event.model, usageStatus: event.usageStatus, - unattributedCost: event.unattributedCost, + unattributedCost, inputTokens: event.inputTokens, outputTokens: event.outputTokens, }, @@ -67,7 +70,7 @@ export function finalizeRun( } leaseRenewalMap.delete(event.runId); - const runOutcome = db.runs.finishRun(event); + const runOutcome = db.runs.finishRun({ ...event, unattributedCost }); if (run.workspaceId) { db.workspacesMod.releaseWorkspaceLock(run.workspaceId, run.id); } diff --git a/src/main/worker/runtime/helpers.ts b/src/main/worker/runtime/helpers.ts index 13baf684..773add88 100644 --- a/src/main/worker/runtime/helpers.ts +++ b/src/main/worker/runtime/helpers.ts @@ -65,7 +65,6 @@ export type ChildMessage = signal: string | null; model: string | null; usageStatus: RunUsageStatus; - unattributedCost: boolean; inputTokens: number | null; outputTokens: number | null; } diff --git a/src/main/worker/runtime/index.ts b/src/main/worker/runtime/index.ts index 388b0e64..45621112 100644 --- a/src/main/worker/runtime/index.ts +++ b/src/main/worker/runtime/index.ts @@ -290,7 +290,6 @@ export class WorkerRunner { signal: null, model: resolvedModelId, usageStatus, - unattributedCost: usageStatus !== "ok", inputTokens: finalUsage.inputTokens, outputTokens: finalUsage.outputTokens, }); diff --git a/tests/unit/run-usage-status.test.ts b/tests/unit/run-usage-status.test.ts index 0210a534..68a9b76a 100644 --- a/tests/unit/run-usage-status.test.ts +++ b/tests/unit/run-usage-status.test.ts @@ -3,6 +3,7 @@ import { classifyRunUsage, deriveLegacyRunUsageStatus, hasAttributableAssistantUsage, + isUsageStatusUnattributed, } from "../../src/main/run-usage-status"; describe("run usage status", () => { @@ -81,4 +82,11 @@ describe("run usage status", () => { ); expect(deriveLegacyRunUsageStatus({ runStatus: "failed", unattributedCost: true })).toBe("pre_first_message"); }); + + it("derives the legacy unattributed-cost flag from usage status", () => { + expect(isUsageStatusUnattributed("ok")).toBe(false); + expect(isUsageStatusUnattributed("pre_first_message")).toBe(true); + expect(isUsageStatusUnattributed("error_stop")).toBe(true); + expect(isUsageStatusUnattributed("aborted_before_usage")).toBe(true); + }); }); diff --git a/tests/unit/worker-message-schema.test.ts b/tests/unit/worker-message-schema.test.ts new file mode 100644 index 00000000..cf641392 --- /dev/null +++ b/tests/unit/worker-message-schema.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from "vitest"; +import { workerMessageSchema } from "../../src/main/worker/entry/types"; + +const RUN_ID = "00000000-0000-4000-8000-000000000001"; + +describe("worker message schema", () => { + it("accepts run-finished messages keyed by usage status", () => { + expect( + workerMessageSchema.safeParse({ + type: "run-finished", + runId: RUN_ID, + status: "failed", + summary: null, + errorMessage: "Agent stopped with error.", + exitCode: null, + signal: null, + model: "gpt-5.4", + usageStatus: "error_stop", + inputTokens: 12, + outputTokens: 1, + }).success, + ).toBe(true); + }); + + it("rejects legacy unattributed-cost truth in run-finished messages", () => { + expect( + workerMessageSchema.safeParse({ + type: "run-finished", + runId: RUN_ID, + status: "failed", + summary: null, + errorMessage: "Agent stopped with error.", + exitCode: null, + signal: null, + model: "gpt-5.4", + usageStatus: "error_stop", + unattributedCost: false, + inputTokens: 12, + outputTokens: 1, + }).success, + ).toBe(false); + }); +});