pi-long-task 0.1.3 → 0.3.0

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,