pi-long-task 0.2.0 → 0.3.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-long-task",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "type": "module",
   "description": "Pi extension for breaking down and running long coding tasks safely.",
   "keywords": [

package/src/coordinator.ts

@@ -80,6 +80,7 @@ export interface CoordinatorProgressUpdate {
   workerEventType?: string;
   isError?: boolean;
   totalTasks?: number;
+  workerCostTotal: number;
   currentTask?: CoordinatorProgressTask;
   subtasks?: CoordinatorProgressSubtask[];
   taskProgress?: TaskProgressModel;
@@ -145,10 +146,18 @@ export interface CoordinatorResult {
   commits: CoordinatorCommitSummary[];
   attempts: TaskAttemptSummary[];
   taskProgress: TaskProgressModel;
+  workerCostTotal: number;
   commit: boolean;
   error?: string;
 }
 
+interface WorkerCostState {
+  total: number;
+  finalizedByWorker: Map<string, number>;
+  liveByWorker: Map<string, number>;
+  liveByMessage: Map<string, number>;
+}
+
 interface RuntimeOptions {
   cwd: string;
   runId: string;
@@ -167,6 +176,7 @@ interface RuntimeOptions {
   todoSessionFactory?: WorkerSessionFactory;
   now: () => Date;
   onProgress?: CoordinatorProgressHandler;
+  workerCostState: WorkerCostState;
 }
 
 export async function runCoordinator(options: RunCoordinatorOptions): Promise<CoordinatorResult> {
@@ -237,6 +247,7 @@ export async function runCoordinator(options: RunCoordinatorOptions): Promise<Co
         onEvent: (event) => emitWorkerEventProgress(runtime, tasksBeforeAttempt, nextTask, attempts, attempt, event),
       });
       outcomes.push(outcome);
+      finalizeWorkerCost(runtime.workerCostState, outcome);
 
       if (outcome.done) {
         todoMarkdown = markTaskDone(todoMarkdown, nextTask.taskId);
@@ -346,6 +357,7 @@ export async function runCoordinator(options: RunCoordinatorOptions): Promise<Co
       commits,
       attempts,
       taskProgress,
+      workerCostTotal: runtime.workerCostState.total,
       commit: options.commit,
       error: failure,
     };
@@ -385,6 +397,7 @@ export async function runCoordinator(options: RunCoordinatorOptions): Promise<Co
       commits,
       attempts,
       taskProgress: buildTaskProgressModel({ tasks: [], attempts }),
+      workerCostTotal: runtime.workerCostState.total,
       commit: options.commit,
       error: message,
     };
@@ -467,23 +480,101 @@ function buildRuntimeOptions(options: RunCoordinatorOptions): RuntimeOptions {
     todoSessionFactory: options.todoSessionFactory,
     now: options.now ?? (() => new Date()),
     onProgress: options.onProgress,
+    workerCostState: createWorkerCostState(),
   };
 }
 
 function emitProgress(
   runtime: RuntimeOptions,
   message: string,
-  update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath">,
+  update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath" | "workerCostTotal">,
 ): void {
   runtime.onProgress?.({
     message,
     runId: runtime.runId,
     todoPath: runtime.todoPath,
     resultPath: runtime.taskResultPath,
+    workerCostTotal: runtime.workerCostState.total,
     ...update,
   });
 }
 
+function createWorkerCostState(): WorkerCostState {
+  return {
+    total: 0,
+    finalizedByWorker: new Map(),
+    liveByWorker: new Map(),
+    liveByMessage: new Map(),
+  };
+}
+
+function recordLiveWorkerCost(
+  state: WorkerCostState,
+  worker: string,
+  event: { usageCostTotal?: number; usageCostKey?: string },
+): boolean {
+  if (state.finalizedByWorker.has(worker) || event.usageCostTotal === undefined || !event.usageCostKey) {
+    return false;
+  }
+
+  const cost = finiteNonNegativeNumber(event.usageCostTotal);
+  if (cost === undefined) {
+    return false;
+  }
+
+  const messageKey = `${worker}:${event.usageCostKey}`;
+  if (state.liveByMessage.get(messageKey) === cost) {
+    return false;
+  }
+
+  state.liveByMessage.set(messageKey, cost);
+  recomputeLiveWorkerCost(state, worker);
+  recomputeWorkerCostTotal(state);
+  return true;
+}
+
+function finalizeWorkerCost(
+  state: WorkerCostState,
+  outcome: Pick<SessionOutcome, "task" | "attempt" | "workerCostTotal">,
+): void {
+  const worker = workerKey(outcome.task.taskId, outcome.attempt);
+  state.finalizedByWorker.set(worker, finiteNonNegativeNumber(outcome.workerCostTotal) ?? 0);
+  state.liveByWorker.delete(worker);
+  for (const messageKey of state.liveByMessage.keys()) {
+    if (messageKey.startsWith(`${worker}:`)) {
+      state.liveByMessage.delete(messageKey);
+    }
+  }
+  recomputeWorkerCostTotal(state);
+}
+
+function recomputeLiveWorkerCost(state: WorkerCostState, worker: string): void {
+  let total = 0;
+  for (const [messageKey, cost] of state.liveByMessage) {
+    if (messageKey.startsWith(`${worker}:`)) {
+      total += cost;
+    }
+  }
+  state.liveByWorker.set(worker, total);
+}
+
+function recomputeWorkerCostTotal(state: WorkerCostState): void {
+  let total = 0;
+  for (const cost of state.finalizedByWorker.values()) {
+    total += cost;
+  }
+  for (const [worker, cost] of state.liveByWorker) {
+    if (!state.finalizedByWorker.has(worker)) {
+      total += cost;
+    }
+  }
+  state.total = total;
+}
+
+function workerKey(taskId: string, attempt: number): string {
+  return `${taskId}:${attempt}`;
+}
+
 function currentTaskProgress(
   task: Pick<Task, "taskId" | "title" | "statusItems">,
   status: CoordinatorProgressItemStatus,
@@ -521,13 +612,33 @@ function emitWorkerEventProgress(
   task: Pick<Task, "taskId" | "title" | "statusItems">,
   attempts: readonly TaskAttemptSummary[],
   attempt: number,
-  event: { type: string; toolName?: string; isError?: boolean },
+  event: { type: string; toolName?: string; isError?: boolean; usageCostTotal?: number; usageCostKey?: string },
 ): void {
+  if (event.usageCostTotal !== undefined) {
+    const changed = recordLiveWorkerCost(runtime.workerCostState, workerKey(task.taskId, attempt), event);
+    if (changed) {
+      emitProgress(
+        runtime,
+        `TODO ${task.taskId}: worker cost updated to ${formatCost(runtime.workerCostState.total)}.`,
+        {
+          phase: "worker_tool",
+          taskId: task.taskId,
+          title: task.title,
+          attempt,
+          status: "in_progress",
+          workerEventType: event.type,
+          ...currentTaskProgress(task, "in_progress"),
+          taskProgress: buildTaskProgressModel({ tasks, attempts, currentTaskId: task.taskId }),
+        },
+      );
+    }
+  }
+
   if (!event.toolName || (event.type !== "tool_execution_start" && event.type !== "tool_execution_end")) {
     return;
   }
   const action = event.type === "tool_execution_start" ? "started" : event.isError ? "failed" : "finished";
-  const update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath"> = {
+  const update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath" | "workerCostTotal"> = {
     phase: "worker_tool",
     taskId: task.taskId,
     title: task.title,
@@ -568,7 +679,7 @@ function emitTaskOutcomeProgress(
     : outcome.reportedStatus === "blocked"
       ? "task_blocked"
       : "task_failed";
-  const update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath"> = {
+  const update: Omit<CoordinatorProgressUpdate, "message" | "runId" | "todoPath" | "resultPath" | "workerCostTotal"> = {
     phase,
     taskId: task.taskId,
     title: task.title,
@@ -759,6 +870,20 @@ function positiveMilliseconds(value: number | undefined, fallback: number): numb
   return fallback;
 }
 
+function finiteNonNegativeNumber(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
+function formatCost(value: number): string {
+  if (value === 0) {
+    return "$0";
+  }
+  if (value < 0.01) {
+    return `$${value.toFixed(4)}`;
+  }
+  return `$${value.toFixed(2)}`;
+}
+
 function errorMessage(error: unknown): string {
   return error instanceof Error ? error.message : String(error);
 }

package/src/index.ts

@@ -1,10 +1,60 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { AssistantMessage } from "@earendil-works/pi-ai";
 
 import { runCoordinator, type CoordinatorProgressUpdate, type CoordinatorResult } from "./coordinator.ts";
 import { longTaskInputTransform } from "./input_router.ts";
 import { renderLongTaskToolCall, renderLongTaskToolResult } from "./render.ts";
 import { PiLongTaskParams } from "./types.ts";
 
+export function createWorkerCostAccumulator() {
+  let pendingWorkerCostTotal = 0;
+
+  return {
+    add(cost: number): void {
+      const value = finiteNonNegativeNumber(cost);
+      if (value && value > 0) {
+        pendingWorkerCostTotal += value;
+      }
+    },
+    applyToAssistantMessage(message: AssistantMessage): AssistantMessage | undefined {
+      if (pendingWorkerCostTotal <= 0) {
+        return undefined;
+      }
+
+      const replacement = addWorkerCostToAssistantMessage(message, pendingWorkerCostTotal);
+      if (replacement) {
+        pendingWorkerCostTotal = 0;
+      }
+      return replacement;
+    },
+  };
+}
+
+export function addWorkerCostToAssistantMessage(
+  message: AssistantMessage,
+  workerCostTotal: number,
+): AssistantMessage | undefined {
+  const workerCost = finiteNonNegativeNumber(workerCostTotal);
+  if (!workerCost || workerCost <= 0) {
+    return undefined;
+  }
+
+  return {
+    ...message,
+    usage: {
+      ...message.usage,
+      cost: {
+        ...message.usage.cost,
+        total: message.usage.cost.total + workerCost,
+      },
+    },
+  };
+}
+
+function finiteNonNegativeNumber(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
 function toolDetails(result: CoordinatorResult) {
   return {
     runId: result.runId,
@@ -19,11 +69,23 @@ function toolDetails(result: CoordinatorResult) {
     blockedTasks: result.blockedTasks,
     remainingTasks: result.remainingTasks,
     taskProgress: result.taskProgress,
+    workerCostTotal: result.workerCostTotal,
     summary: result.summary,
   };
 }
 
 export default function registerPiLongTaskExtension(pi: ExtensionAPI) {
+  const workerCostAccumulator = createWorkerCostAccumulator();
+
+  pi.on("message_end", (event) => {
+    if (event.message.role !== "assistant") {
+      return undefined;
+    }
+
+    const message = workerCostAccumulator.applyToAssistantMessage(event.message);
+    return message ? { message } : undefined;
+  });
+
   pi.on("input", (event) => {
     if (event.source === "extension") {
       return { action: "continue" as const };
@@ -64,6 +126,7 @@ export default function registerPiLongTaskExtension(pi: ExtensionAPI) {
         abortSignal: signal,
         onProgress: publishProgress,
       });
+      workerCostAccumulator.add(result.workerCostTotal);
 
       return {
         content: [

package/src/render.ts

@@ -17,6 +17,7 @@ export interface CoordinatorResultForRendering {
   commits?: CoordinatorCommitSummary[];
   remainingTasks?: CoordinatorRemainingTask[];
   taskProgress?: TaskProgressModel;
+  workerCostTotal?: number;
   error?: string;
 }
 
@@ -45,11 +46,13 @@ const SIDEBAR_GAP = 2;
 class LongTaskSidebarShell implements Component {
   private readonly mainText: string;
   private readonly taskProgress: TaskProgressModel;
+  private readonly workerCostTotal: number | undefined;
   private readonly theme: Theme;
 
-  constructor(mainText: string, taskProgress: TaskProgressModel, theme: Theme) {
+  constructor(mainText: string, taskProgress: TaskProgressModel, theme: Theme, workerCostTotal?: number) {
     this.mainText = mainText;
     this.taskProgress = taskProgress;
+    this.workerCostTotal = workerCostTotal;
     this.theme = theme;
   }
 
@@ -92,7 +95,7 @@ class LongTaskSidebarShell implements Component {
     }
 
     const innerWidth = Math.max(0, width - 2);
-    const rows = sidebarRows(this.taskProgress, this.theme);
+    const rows = sidebarRows(this.taskProgress, this.theme, this.workerCostTotal);
     return [
       sidebarBorder("Long Task", width, this.theme),
       ...rows.map((row) => sidebarRow(row, innerWidth, this.theme)),
@@ -113,6 +116,10 @@ export function formatCoordinatorResultMessage(result: CoordinatorResultForRende
     `TODO file: ${result.todoPath}`,
   ];
 
+  if (result.workerCostTotal) {
+    lines.push(`Worker spend: ${formatCost(result.workerCostTotal)}`);
+  }
+
   const commitLines = commits
     .filter((commit) => commit.hash || commit.error)
     .map((commit) => {
@@ -154,8 +161,9 @@ export function renderLongTaskToolResult(
   const details = recordOrUndefined(result.details);
   if (options.isPartial) {
     const taskProgress = taskProgressModel(details?.taskProgress);
+    const workerCostTotal = numberValue(details?.workerCostTotal);
     const main = renderLongTaskProgress(details, contentText(result), theme);
-    return taskProgress ? new LongTaskSidebarShell(main, taskProgress, theme) : new Text(main, 0, 0);
+    return taskProgress ? new LongTaskSidebarShell(main, taskProgress, theme, workerCostTotal) : new Text(main, 0, 0);
   }
 
   const finalDetails = longTaskDetails(details);
@@ -165,7 +173,7 @@ export function renderLongTaskToolResult(
 
   const main = renderLongTaskSummary(finalDetails, options.expanded, theme);
   return finalDetails.taskProgress
-    ? new LongTaskSidebarShell(main, finalDetails.taskProgress, theme)
+    ? new LongTaskSidebarShell(main, finalDetails.taskProgress, theme, finalDetails.workerCostTotal)
     : new Text(main, 0, 0);
 }
 
@@ -208,6 +216,7 @@ function renderLongTaskSummary(details: CoordinatorToolRenderDetails, expanded:
     details.blockedTasks ? theme.fg("warning", `${details.blockedTasks} blocked`) : undefined,
     remainingCount ? theme.fg("muted", `${remainingCount} remaining`) : undefined,
     commitCount ? theme.fg("success", `${commitCount} commit${commitCount === 1 ? "" : "s"}`) : undefined,
+    details.workerCostTotal ? theme.fg("muted", `worker ${formatCost(details.workerCostTotal)}`) : undefined,
   ].filter(Boolean);
 
   if (!expanded) {
@@ -252,9 +261,12 @@ function padLine(line: string, width: number): string {
   return truncateToWidth(line, width, "…", true);
 }
 
-function sidebarRows(taskProgress: TaskProgressModel, theme: Theme): string[] {
+function sidebarRows(taskProgress: TaskProgressModel, theme: Theme, workerCostTotal?: number): string[] {
   const summary = normalizedTaskProgressSummary(taskProgress);
   const rows = [theme.fg("toolTitle", theme.bold("Task sidebar")), theme.fg("dim", "Centered timeline")];
+  if (workerCostTotal) {
+    rows.push(theme.fg("muted", `Worker spend: ${formatCost(workerCostTotal)}`));
+  }
   if (summary.totalTasks === 0) {
     rows.push("", theme.fg("muted", "Waiting for TODO plan..."));
     return rows;
@@ -503,6 +515,7 @@ function longTaskDetails(details: Record<string, unknown> | undefined): Coordina
     commits: commitSummaries(details.commits),
     remainingTasks: remainingTaskSummaries(details.remainingTasks),
     taskProgress: taskProgressModel(details.taskProgress),
+    workerCostTotal: nonNegativeNumberValue(details.workerCostTotal),
     error: stringValue(details.error),
   };
 }
@@ -592,6 +605,20 @@ function numberValue(value: unknown): number | undefined {
   return typeof value === "number" && Number.isFinite(value) ? value : undefined;
 }
 
+function nonNegativeNumberValue(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
+function formatCost(value: number): string {
+  if (value === 0) {
+    return "$0";
+  }
+  if (value < 0.01) {
+    return `$${value.toFixed(4)}`;
+  }
+  return `$${value.toFixed(2)}`;
+}
+
 function recordOrUndefined(value: unknown): Record<string, unknown> | undefined {
   return typeof value === "object" && value !== null ? (value as Record<string, unknown>) : undefined;
 }

package/src/types.ts

@@ -62,6 +62,7 @@ export interface PiLongTaskResult {
     commitSkipped?: string;
   }>;
   taskProgress: TaskProgressModel;
+  workerCostTotal: number;
   commit: boolean;
   error?: string;
 }

package/src/worker_session.ts

@@ -201,7 +201,7 @@ export interface WorkerSessionLike {
   compact?(customInstructions?: string): Promise<unknown>;
   dispose?(): void;
   getLastAssistantText?(): string | undefined;
-  getSessionStats?(): unknown;
+  getSessionStats?(): unknown | Promise<unknown>;
   getContextUsage?(): unknown;
   sessionFile?: string;
   sessionId?: string;
@@ -246,6 +246,8 @@ export interface CapturedWorkerEvent {
   toolName?: string;
   isError?: boolean;
   note?: string;
+  usageCostTotal?: number;
+  usageCostKey?: string;
 }
 
 export interface SessionOutcome {
@@ -261,6 +263,8 @@ export interface SessionOutcome {
   contextObservations: string[];
   compactionEvents: string[];
   events: CapturedWorkerEvent[];
+  workerCostTotal: number;
+  workerCostSource?: string;
   shutdownRequested: boolean;
   timedOut: boolean;
   aborted: boolean;
@@ -345,6 +349,9 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
   let error: string | undefined;
   let finished = false;
   let turnCount = 0;
+  let messageUsageCostTotal = 0;
+  let hasMessageUsageCost = false;
+  let sessionStatsCostTotal: number | undefined;
 
   const prompt = buildTaskPrompt(options);
   const taskTimeoutSeconds = options.taskTimeoutSeconds ?? DEFAULT_TASK_TIMEOUT_SECONDS;
@@ -366,6 +373,22 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     timers.clear();
   };
 
+  const messageUsageCostsByKey = new Map<string, number>();
+  const recordWorkerUsageCost = (cost: number | undefined, key: string | undefined) => {
+    if (cost === undefined) {
+      return;
+    }
+    hasMessageUsageCost = true;
+    if (!key) {
+      messageUsageCostTotal += cost;
+      return;
+    }
+
+    const previousCost = messageUsageCostsByKey.get(key) ?? 0;
+    messageUsageCostsByKey.set(key, cost);
+    messageUsageCostTotal += cost - previousCost;
+  };
+
   const schedule = (fn: () => void | Promise<void>, ms: number) => {
     const timer = setTimeout(() => {
       timers.delete(timer);
@@ -458,6 +481,7 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
           if (messageText) {
             assistantText = messageText;
           }
+          recordWorkerUsageCost(workerUsageCostFromEvent(event), workerUsageCostKeyFromEvent(event));
           break;
         }
         case "turn_end": {
@@ -538,6 +562,7 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
       assistantText = latestAssistantText(session, assistantText);
       sessionFile = session.sessionFile ?? sessionFile;
       sessionId = session.sessionId ?? sessionId;
+      sessionStatsCostTotal = await workerUsageCostFromSessionStats(session);
       session.dispose?.();
     }
   }
@@ -547,6 +572,10 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
   }
 
   const reportedStatus = parseReportedStatus(assistantText);
+  const capturedWorkerCost = selectWorkerCostTotal({
+    messageCostTotal: hasMessageUsageCost ? messageUsageCostTotal : undefined,
+    statsCostTotal: sessionStatsCostTotal,
+  });
   return {
     task: options.task,
     attempt: options.attempt,
@@ -560,6 +589,8 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     contextObservations,
     compactionEvents,
     events,
+    workerCostTotal: capturedWorkerCost.total,
+    workerCostSource: capturedWorkerCost.source,
     shutdownRequested,
     timedOut,
     aborted: aborted || Boolean(options.abortSignal?.aborted),
@@ -723,9 +754,15 @@ function summarizeWorkerEvent(event: unknown): CapturedWorkerEvent | undefined {
     };
   }
 
+  if (event.type === "message_end") {
+    const usageCostTotal = workerUsageCostFromEvent(event);
+    return usageCostTotal === undefined
+      ? { type: event.type }
+      : { type: event.type, usageCostTotal, usageCostKey: workerUsageCostKeyFromEvent(event) };
+  }
+
   if (
     event.type === "turn_end" ||
-    event.type === "message_end" ||
     event.type === "compaction_start" ||
     event.type === "compaction_end" ||
     event.type === "agent_end" ||
@@ -738,6 +775,109 @@ function summarizeWorkerEvent(event: unknown): CapturedWorkerEvent | undefined {
   return undefined;
 }
 
+export function workerUsageCostFromEvent(event: unknown): number | undefined {
+  if (!isRecord(event)) {
+    return undefined;
+  }
+
+  for (const candidate of [event.assistantMessage, event.message, event]) {
+    const cost = workerUsageCostFromAssistantMessage(candidate);
+    if (cost !== undefined) {
+      return cost;
+    }
+  }
+  return undefined;
+}
+
+export function workerUsageCostFromAssistantMessage(message: unknown): number | undefined {
+  if (!isRecord(message)) {
+    return undefined;
+  }
+  return usageCostTotal(message.usage);
+}
+
+export function workerUsageCostKeyFromEvent(event: unknown): string | undefined {
+  if (!isRecord(event)) {
+    return undefined;
+  }
+
+  for (const candidate of [event.assistantMessage, event.message, event]) {
+    const key = workerUsageCostKeyFromAssistantMessage(candidate);
+    if (key) {
+      return key;
+    }
+  }
+  return undefined;
+}
+
+function workerUsageCostKeyFromAssistantMessage(message: unknown): string | undefined {
+  if (!isRecord(message)) {
+    return undefined;
+  }
+
+  for (const keyName of ["id", "messageId", "uuid"] as const) {
+    const value = message[keyName];
+    if (typeof value === "string" && value) {
+      return `${keyName}:${value}`;
+    }
+  }
+  return undefined;
+}
+
+export function workerUsageCostFromStats(stats: unknown): number | undefined {
+  if (!isRecord(stats)) {
+    return undefined;
+  }
+
+  const directCost = finiteNonNegativeNumber(stats.cost);
+  if (directCost !== undefined) {
+    return directCost;
+  }
+
+  return usageCostTotal(stats.usage) ?? usageCostTotal(stats);
+}
+
+async function workerUsageCostFromSessionStats(session: WorkerSessionLike): Promise<number | undefined> {
+  if (!session.getSessionStats) {
+    return undefined;
+  }
+
+  try {
+    return workerUsageCostFromStats(await session.getSessionStats());
+  } catch {
+    return undefined;
+  }
+}
+
+function usageCostTotal(usage: unknown): number | undefined {
+  if (!isRecord(usage)) {
+    return undefined;
+  }
+
+  const cost = usage.cost;
+  if (isRecord(cost)) {
+    return finiteNonNegativeNumber(cost.total);
+  }
+  return finiteNonNegativeNumber(cost);
+}
+
+function selectWorkerCostTotal(options: { messageCostTotal: number | undefined; statsCostTotal: number | undefined }): {
+  total: number;
+  source?: string;
+} {
+  if (options.statsCostTotal !== undefined) {
+    return { total: options.statsCostTotal, source: "session_stats" };
+  }
+  if (options.messageCostTotal !== undefined) {
+    return { total: options.messageCostTotal, source: "message_end" };
+  }
+  return { total: 0 };
+}
+
+function finiteNonNegativeNumber(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
 function captureContextUsage(
   session: WorkerSessionLike | undefined,
   turnCount: number,