@ryanfw/prompt-orchestration-pipeline 1.3.2 → 1.3.4

package/src/core/task-runner.ts

@@ -15,10 +15,11 @@ import { loadEnvironment } from "./environment";
 import { createJobLogger } from "./logger";
 import { createHighLevelLLM, createLLMWithOverride, getLLMEvents } from "../llm/index";
 import { executeAgent } from "./agent-step";
+import { withConsoleCapture, withTaskTelemetry } from "./task-telemetry";
 import { TaskState } from "../config/statuses";
 import { LogEvent, LogFileExtension } from "../config/log-events";
 // LLMClient: src/llm/index.ts exports HighLevelLLM (no LLMClient type exists there).
-import type { HighLevelLLM } from "../providers/types";
+import type { HighLevelLLM, UsageSource } from "../providers/types";
 import type { AgentStepResult, TaskAgentOptions, TaskAgentRunner } from "./agent-types";
 
 /** Opaque alias — the LLM client passed into execution contexts. */
@@ -165,7 +166,15 @@ export interface LLMMetricRecord {
 }
 
 /** Token usage tuple written to the job status file. */
-export type TokenUsageTuple = [modelKey: string, inputTokens: number, outputTokens: number, cost: number];
+export type TokenUsageTuple = [
+  modelKey: string,
+  inputTokens: number,
+  outputTokens: number,
+  cost: number,
+  source: UsageSource,
+  failed?: boolean,
+  costEstimated?: boolean,
+];
 
 // ─── Initial context ─────────────────────────────────────────────────────────
 
@@ -251,22 +260,84 @@ export function deriveModelKeyAndTokens(metric: Record<string, unknown>): TokenU
   const rawCost = metric["cost"];
   const cost = typeof rawCost === "number" && Number.isFinite(rawCost) ? rawCost : 0;
 
-  return [modelKey, inputTokens, outputTokens, cost];
+  const rawSource = metric["usageSource"];
+  const source: UsageSource =
+    rawSource === "reported" ||
+    rawSource === "estimated" ||
+    rawSource === "unavailable" ||
+    rawSource === "zero"
+      ? rawSource
+      : "reported";
+
+  const rawCostEstimated = metric["costEstimated"];
+  const costEstimated = typeof rawCostEstimated === "boolean"
+    ? rawCostEstimated
+    : source !== "reported";
+
+  return [modelKey, inputTokens, outputTokens, cost, source, false, costEstimated];
 }
 
-function deriveAgentUsageTuple(options: TaskAgentOptions, result: AgentStepResult): TokenUsageTuple | null {
+/**
+ * Builds a source-aware {@link TokenUsageTuple} from an agent step result.
+ *
+ * Classification (#18 `UsageSource` model — the agent lane never estimates
+ * tokens, so it keys on the result shape rather than a provider/harness table):
+ * - `result.usage` object present -> `"reported"` (even when token counts are
+ *   zero); `costUsd` (when finite) is preserved as the cost, `costEstimated`
+ *   stays `false`.
+ * - no `result.usage` but a finite `costUsd` -> `"unavailable"`: tokens are 0,
+ *   the harness-reported `costUsd` is preserved as the cost, `costEstimated` is
+ *   `true` (agent spend stops lying as measured).
+ * - neither `usage` nor `costUsd` -> `null` (nothing to record).
+ *
+ * `opts.failed` sets the tuple's `failed` flag (index 5).
+ */
+export function deriveAgentUsageTuple(
+  options: { harness: string; model?: TaskAgentOptions["model"] },
+  result: AgentStepResult,
+  opts?: { failed?: boolean },
+): TokenUsageTuple | null {
   const costUsd =
     typeof result.costUsd === "number" && Number.isFinite(result.costUsd)
       ? result.costUsd
       : undefined;
   if (!result.usage && costUsd === undefined) return null;
 
-  return [
-    `${options.harness}:${options.model ?? "default"}`,
-    result.usage?.inputTokens ?? 0,
-    result.usage?.outputTokens ?? 0,
-    costUsd ?? 0,
-  ];
+  const modelId =
+    typeof options.model === "string" ? options.model : options.model?.model;
+  const modelKey = `${options.harness}:${modelId ?? "default"}`;
+  const failed = opts?.failed === true;
+
+  if (result.usage) {
+    return [
+      modelKey,
+      result.usage.inputTokens ?? 0,
+      result.usage.outputTokens ?? 0,
+      costUsd ?? 0,
+      "reported",
+      failed,
+      false,
+    ];
+  }
+
+  return [modelKey, 0, 0, costUsd ?? 0, "unavailable", failed, true];
+}
+
+/**
+ * Records agent usage by building a {@link TokenUsageTuple} via
+ * {@link deriveAgentUsageTuple} and pushing it through the injected `push`
+ * sink only when a tuple exists. The skip-on-null decision lives here, so
+ * callers never branch on null. Owns the recording decision, not the storage
+ * mechanism (the persistence primitive is injected as `push`).
+ */
+export function recordAgentUsage(
+  options: { harness: string; model?: TaskAgentOptions["model"] },
+  result: AgentStepResult,
+  push: (tuple: TokenUsageTuple) => void,
+  opts?: { failed?: boolean },
+): void {
+  const tuple = deriveAgentUsageTuple(options, result, opts);
+  if (tuple) push(tuple);
 }
 
 // ─── Safe clone ─────────────────────────────────────────────────────────────
@@ -417,47 +488,8 @@ export function checkFlagTypeConflicts(
 
 // ─── Console capture ──────────────────────────────────────────────────────────
 
-/**
- * Replaces console.log/error/warn/info/debug with functions that write to a
- * string buffer with appropriate prefixes. Returns a restore function that
- * flushes the buffer to `logPath` via Bun.write() and restores originals.
- * Creates the log directory before capturing.
- */
-export function captureConsoleOutput(logPath: string): () => Promise<void> {
-  const origLog = console.log;
-  const origError = console.error;
-  const origWarn = console.warn;
-  const origInfo = console.info;
-  const origDebug = console.debug;
-
-  let buffer = "";
-
-  const append = (prefix: string, args: unknown[]) => {
-    buffer += prefix + args.map(String).join(" ") + "\n";
-  };
-
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  console.log = (...args: any[]) => append("", args);
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  console.error = (...args: any[]) => append("[ERROR] ", args);
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  console.warn = (...args: any[]) => append("[WARN] ", args);
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  console.info = (...args: any[]) => append("[INFO] ", args);
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  console.debug = (...args: any[]) => append("[DEBUG] ", args);
-
-  return async () => {
-    console.log = origLog;
-    console.error = origError;
-    console.warn = origWarn;
-    console.info = origInfo;
-    console.debug = origDebug;
-
-    await mkdir(dirname(logPath), { recursive: true });
-    await Bun.write(logPath, buffer);
-  };
-}
+// `captureConsoleOutput` moved to ./task-telemetry; re-exported below so existing
+// importers keep working. The per-stage usage at line ~733 is replaced in step 4.
 
 // ─── Pipeline stage factory ───────────────────────────────────────────────────
 
@@ -532,6 +564,11 @@ export async function runPipeline(
   let tokenWriteQueue: Promise<void> = Promise.resolve();
   const events = getLLMEvents();
 
+  // `context` is built inside the withTaskTelemetry body below, but the LLM
+  // handlers (defined here so they can be passed into the wrapper) read
+  // `context.currentStage`. Declare it at this scope so both see the same value.
+  let context: ExecutionContext;
+
   const onComplete = (metric: Record<string, unknown>) => {
     const record: LLMMetricRecord = {
       ...metric,
@@ -564,304 +601,303 @@ export async function runPipeline(
       stage: context.currentStage ?? undefined,
       failed: true,
     });
-  };
-
-  events.on("llm:request:complete", onComplete);
-  events.on("llm:request:error", onError);
-
-  // 7. Load task module
-  const taskModule = await loadFreshModule(pathToFileURL(modulePath));
 
-  // 8. Create per-invocation stage array
-  const stages = createPipelineStages(taskModule, initialContext.tasksOverride);
-
-  // 9. Create file I/O adapter
-  let currentStageRef: StageName | null = null;
-  const io = createTaskFileIO({
-    workDir,
-    taskName,
-    getStage: () => currentStageRef ?? "",
-    statusPath,
-  });
-
-  // 9b. Build the CLI-agent harness runner injected into stages. It reuses the
-  // task's own `io`, so an agent reads and writes the same artifacts the task
-  // sees. A caller may override it (tests/custom hosts).
-  const runAgentCore: TaskAgentRunner =
-    initialContext.runAgent ??
-    ((options: TaskAgentOptions) =>
-      executeAgent({ io, entry: { ...options, name: taskName } }));
-  const runAgent: TaskAgentRunner = async (options: TaskAgentOptions) => {
-    const result = await runAgentCore(options);
-    const tuple = deriveAgentUsageTuple(options, result);
-    if (tuple) enqueueTokenUsage(tuple);
-    return result;
+    const tuple = deriveModelKeyAndTokens(metric);
+    tuple[5] = true;
+    enqueueTokenUsage(tuple);
   };
 
-  // 10. Build execution context
-  const pipelineTasks = initialContext.pipelineTasks ?? initialContext.meta?.pipelineTasks;
-
-  // Seed handling: if initialContext.seed is falsy, use initialContext itself.
-  // JSON round-trip strips non-cloneable fields (functions, proxies) which matches
-  // the original JS behavior and ensures structuredClone works later.
-  const seedData = initialContext.seed || JSON.parse(JSON.stringify(initialContext));
-
-  const context: ExecutionContext = {
-    io,
-    llm,
-    runAgent,
-    meta: {
-      taskName,
+  // INVARIANT: ONE runPipeline at a time per process (sequential stages, one
+  // runner per job; see task-telemetry.ts). withTaskTelemetry attaches the two
+  // LLM listeners, runs the body below, and in `finally` removes both listeners
+  // and flushes the latest tokenWriteQueue (thunked) on every exit — return or
+  // throw. The body owns module-load through the final return.
+  return withTaskTelemetry(events, { onComplete, onError }, () => tokenWriteQueue, async () => {
+    // 7. Load task module
+    const taskModule = await loadFreshModule(pathToFileURL(modulePath));
+
+    // 8. Create per-invocation stage array
+    const stages = createPipelineStages(taskModule, initialContext.tasksOverride);
+
+    // 9. Create file I/O adapter
+    let currentStageRef: StageName | null = null;
+    const io = createTaskFileIO({
       workDir,
+      taskName,
+      getStage: () => currentStageRef ?? "",
       statusPath,
-      jobId: initialContext.jobId,
-      envLoaded: true,
-      modelConfig: initialContext.modelConfig,
-      pipelineTasks,
-    },
-    data: {
-      seed: seedData,
-    },
-    flags: {},
-    logs: [],
-    currentStage: null,
-    validators: { validateWithSchema },
-  };
-
-  // 11. Ensure log directory exists
-  const logsDir = join(workDir, "files", "logs");
-  await mkdir(logsDir, { recursive: true });
-  const registerStageLog = async (fileName: string): Promise<void> => {
-    await trackFile(jobDir, "logs", fileName, taskName, true);
-  };
-
-  // Track returned logs (separate from context.logs)
-  const returnedLogs: AuditLogEntry[] = [];
-  let lastStageOutput: unknown = context.data["seed"];
-  let previousStage = "seed";
-
-  // 12. Execute each stage sequentially
-  for (const stage of stages) {
-    const stageName = stage.name;
-
-    // Check skipIf -> skip if true (log to context.logs only)
-    if (stage.skipIf && stage.skipIf(context.flags)) {
-      context.logs.push({ stage: stageName, skipped: true });
-      continue;
-    }
-
-    // Check handler -> skip if null (log to returned logs)
-    if (stage.handler === null) {
-      returnedLogs.push({ stage: stageName, skipped: true });
-      continue;
-    }
-
-    // Capture console output
-    const logFileName = generateLogName(taskName, stageName, LogEvent.START, LogFileExtension.TEXT);
-    const logPath = join(logsDir, logFileName);
-    const restoreConsole = captureConsoleOutput(logPath);
+    });
 
-    let stageError: { err: unknown } | null = null;
+    // 9b. Build the CLI-agent harness runner injected into stages. It reuses the
+    // task's own `io`, so an agent reads and writes the same artifacts the task
+    // sees. A caller may override it (tests/custom hosts).
+    const runAgentCore: TaskAgentRunner =
+      initialContext.runAgent ??
+      ((options: TaskAgentOptions) =>
+        executeAgent({ io, entry: { ...options, name: taskName } }));
+    const runAgent: TaskAgentRunner = async (options: TaskAgentOptions) => {
+      const result = await runAgentCore(options);
+      recordAgentUsage(options, result, enqueueTokenUsage);
+      return result;
+    };
 
-    try {
-      // Set currentStage
-      context.currentStage = stageName;
-      currentStageRef = stageName;
+    // 10. Build execution context
+    const pipelineTasks = initialContext.pipelineTasks ?? initialContext.meta?.pipelineTasks;
 
-      // Write stage-start status (swallow errors)
-      try {
-        await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
-          if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
-          snapshot.tasks[taskName]!.state = TaskState.RUNNING;
-          snapshot.tasks[taskName]!.currentStage = stageName;
-        });
-      } catch {
-        // Best-effort: swallow status write errors
-      }
+    // Seed handling: if initialContext.seed is falsy, use initialContext itself.
+    // JSON round-trip strips non-cloneable fields (functions, proxies) which matches
+    // the original JS behavior and ensures structuredClone works later.
+    const seedData = initialContext.seed || JSON.parse(JSON.stringify(initialContext));
 
-      // Clone data, flags, output into StageContext.
-      // structuredClone throws on non-cloneable values (functions, streams,
-      // class instances with internal slots). Fall back to JSON round-trip
-      // which silently strips those fields.
-      const stageContext: StageContext = {
-        io,
-        llm,
-        runAgent,
-        meta: context.meta,
-        data: safeClone(context.data),
-        flags: safeClone(context.flags),
-        currentStage: stageName,
-        output: safeClone(lastStageOutput),
-        previousStage,
-        validators: context.validators,
-      };
-
-      // Write context snapshot
-      const snapshotFileName = generateLogName(
+    context = {
+      io,
+      llm,
+      runAgent,
+      meta: {
         taskName,
-        stageName,
-        LogEvent.CONTEXT,
-        LogFileExtension.JSON,
-      );
-      const snapshotPath = join(logsDir, snapshotFileName);
-      try {
-        await Bun.write(snapshotPath, JSON.stringify(stageContext, null, 2));
-        await registerStageLog(snapshotFileName);
-      } catch {
-        // Best-effort
-      }
-
-      // Validate prerequisite flags
-      validateFlagTypes(stageName, context.flags, "requires");
-
-      // Execute handler, time with performance.now()
-      const t0 = performance.now();
-      const result = await stage.handler(stageContext);
-      const elapsed = performance.now() - t0;
+        workDir,
+        statusPath,
+        jobId: initialContext.jobId,
+        envLoaded: true,
+        modelConfig: initialContext.modelConfig,
+        pipelineTasks,
+      },
+      data: {
+        seed: seedData,
+      },
+      flags: {},
+      logs: [],
+      currentStage: null,
+      validators: { validateWithSchema },
+    };
 
-      // Validate result shape
-      assertStageResult(result, stageName);
+    // 11. Ensure log directory exists
+    const logsDir = join(workDir, "files", "logs");
+    await mkdir(logsDir, { recursive: true });
+    const registerStageLog = async (fileName: string): Promise<void> => {
+      await trackFile(jobDir, "logs", fileName, taskName, true);
+    };
 
-      // Validate produced flag types
-      validateFlagTypes(stageName, result.flags, "produces");
+    // Track returned logs (separate from context.logs)
+    const returnedLogs: AuditLogEntry[] = [];
+    let lastStageOutput: unknown = context.data["seed"];
+    let previousStage = "seed";
 
-      // Check flag type conflicts
-      checkFlagTypeConflicts(context.flags, result.flags);
+    // 12. Execute each stage sequentially
+    for (const stage of stages) {
+      const stageName = stage.name;
 
-      // Store output in context.data[stageName]
-      context.data[stageName] = result.output;
+      // Check skipIf -> skip if true (log to context.logs only)
+      if (stage.skipIf && stage.skipIf(context.flags)) {
+        context.logs.push({ stage: stageName, skipped: true });
+        continue;
+      }
 
-      // Update lastStageOutput if not a validation stage
-      if (!VALIDATION_STAGES.has(stageName)) {
-        lastStageOutput = result.output;
-        previousStage = stageName;
+      // Check handler -> skip if null (log to returned logs)
+      if (stage.handler === null) {
+        returnedLogs.push({ stage: stageName, skipped: true });
+        continue;
       }
 
-      // Merge flags
-      Object.assign(context.flags, result.flags);
+      // Capture console output
+      const logFileName = generateLogName(taskName, stageName, LogEvent.START, LogFileExtension.TEXT);
+      const logPath = join(logsDir, logFileName);
 
-      // Log completion audit entry
-      const auditEntry: AuditLogEntry = { stage: stageName, ok: true, ms: elapsed };
-      context.logs.push(auditEntry);
-      returnedLogs.push(auditEntry);
+      let stageError: { err: unknown } | null = null;
 
-      // Write stage-completion status (swallow errors)
+      // withConsoleCapture restores every captured console method (and writes the
+      // buffered log file) in its own finally on both success and throw, then
+      // re-throws — so a stage throw is captured here as stageError. registerStageLog
+      // then runs once on every path, after console restore, matching the original
+      // "finally restores console then registers log" ordering.
       try {
-        await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
-          if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
-          snapshot.tasks[taskName]!.currentStage = stageName;
+        await withConsoleCapture(logPath, async () => {
+          // Set currentStage
+          context.currentStage = stageName;
+          currentStageRef = stageName;
+
+          // Write stage-start status (swallow errors)
+          try {
+            await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
+              if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
+              snapshot.tasks[taskName]!.state = TaskState.RUNNING;
+              snapshot.tasks[taskName]!.currentStage = stageName;
+            });
+          } catch {
+            // Best-effort: swallow status write errors
+          }
+
+          // Clone data, flags, output into StageContext.
+          // structuredClone throws on non-cloneable values (functions, streams,
+          // class instances with internal slots). Fall back to JSON round-trip
+          // which silently strips those fields.
+          const stageContext: StageContext = {
+            io,
+            llm,
+            runAgent,
+            meta: context.meta,
+            data: safeClone(context.data),
+            flags: safeClone(context.flags),
+            currentStage: stageName,
+            output: safeClone(lastStageOutput),
+            previousStage,
+            validators: context.validators,
+          };
+
+          // Write context snapshot
+          const snapshotFileName = generateLogName(
+            taskName,
+            stageName,
+            LogEvent.CONTEXT,
+            LogFileExtension.JSON,
+          );
+          const snapshotPath = join(logsDir, snapshotFileName);
+          try {
+            await Bun.write(snapshotPath, JSON.stringify(stageContext, null, 2));
+            await registerStageLog(snapshotFileName);
+          } catch {
+            // Best-effort
+          }
+
+          // Validate prerequisite flags
+          validateFlagTypes(stageName, context.flags, "requires");
+
+          // Execute handler, time with performance.now()
+          const t0 = performance.now();
+          const result = await stage.handler!(stageContext);
+          const elapsed = performance.now() - t0;
+
+          // Validate result shape
+          assertStageResult(result, stageName);
+
+          // Validate produced flag types
+          validateFlagTypes(stageName, result.flags, "produces");
+
+          // Check flag type conflicts
+          checkFlagTypeConflicts(context.flags, result.flags);
+
+          // Store output in context.data[stageName]
+          context.data[stageName] = result.output;
+
+          // Update lastStageOutput if not a validation stage
+          if (!VALIDATION_STAGES.has(stageName)) {
+            lastStageOutput = result.output;
+            previousStage = stageName;
+          }
+
+          // Merge flags
+          Object.assign(context.flags, result.flags);
+
+          // Log completion audit entry
+          const auditEntry: AuditLogEntry = { stage: stageName, ok: true, ms: elapsed };
+          context.logs.push(auditEntry);
+          returnedLogs.push(auditEntry);
+
+          // Write stage-completion status (swallow errors)
+          try {
+            await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
+              if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
+              snapshot.tasks[taskName]!.currentStage = stageName;
+            });
+          } catch {
+            // Best-effort
+          }
+
+          // Write completion log marker
+          const completeFileName = generateLogName(
+            taskName,
+            stageName,
+            LogEvent.COMPLETE,
+            LogFileExtension.TEXT,
+          );
+          try {
+            await Bun.write(join(logsDir, completeFileName), `Stage ${stageName} completed in ${elapsed.toFixed(1)}ms\n`);
+            await registerStageLog(completeFileName);
+          } catch {
+            // Best-effort
+          }
         });
-      } catch {
-        // Best-effort
+      } catch (err: unknown) {
+        // 13. On stage error — capture error; withConsoleCapture already restored console
+        stageError = { err };
       }
 
-      // Write completion log marker
-      const completeFileName = generateLogName(
-        taskName,
-        stageName,
-        LogEvent.COMPLETE,
-        LogFileExtension.TEXT,
-      );
-      try {
-        await Bun.write(join(logsDir, completeFileName), `Stage ${stageName} completed in ${elapsed.toFixed(1)}ms\n`);
-        await registerStageLog(completeFileName);
-      } catch {
-        // Best-effort
-      }
-    } catch (err: unknown) {
-      // 13. On stage error — capture error, let finally restore console
-      stageError = { err };
-    } finally {
-      // Always restore console, regardless of success or failure
-      await restoreConsole();
+      // Register the stage start log on every path, after console restore.
       await registerStageLog(logFileName).catch(() => {});
-    }
 
-    if (stageError !== null) {
-      const { err } = stageError;
-      const base = normalizeError(err);
-      const logFileName2 = generateLogName(taskName, stageName, LogEvent.START, LogFileExtension.TEXT);
-      const snapshotFileName2 = generateLogName(taskName, stageName, LogEvent.CONTEXT, LogFileExtension.JSON);
-
-      const normalized: NormalizedError = {
-        ...base,
-        debug: {
-          stage: stageName,
-          previousStage,
-          logPath: join(logsDir, logFileName2),
-          snapshotPath: join(logsDir, snapshotFileName2),
-          dataHasSeed: "seed" in context.data,
-          seedHasData: context.data["seed"] != null && typeof context.data["seed"] === "object" && "data" in (context.data["seed"] as Record<string, unknown>),
-          flagsKeys: Object.keys(context.flags),
-        },
-      };
-
-      const errorEntry: AuditLogEntry = { stage: stageName, ok: false, ms: 0, error: normalized };
-      context.logs.push(errorEntry);
-      returnedLogs.push(errorEntry);
-
-      // Write failure status (swallow errors)
-      try {
-        await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
-          if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
-          snapshot.tasks[taskName]!.state = TaskState.FAILED;
-          snapshot.tasks[taskName]!.failedStage = stageName;
-          snapshot.tasks[taskName]!.error = normalized.message;
-        });
-      } catch {
-        // Best-effort
+      if (stageError !== null) {
+        const { err } = stageError;
+        const base = normalizeError(err);
+        const logFileName2 = generateLogName(taskName, stageName, LogEvent.START, LogFileExtension.TEXT);
+        const snapshotFileName2 = generateLogName(taskName, stageName, LogEvent.CONTEXT, LogFileExtension.JSON);
+
+        const normalized: NormalizedError = {
+          ...base,
+          debug: {
+            stage: stageName,
+            previousStage,
+            logPath: join(logsDir, logFileName2),
+            snapshotPath: join(logsDir, snapshotFileName2),
+            dataHasSeed: "seed" in context.data,
+            seedHasData: context.data["seed"] != null && typeof context.data["seed"] === "object" && "data" in (context.data["seed"] as Record<string, unknown>),
+            flagsKeys: Object.keys(context.flags),
+          },
+        };
+
+        const errorEntry: AuditLogEntry = { stage: stageName, ok: false, ms: 0, error: normalized };
+        context.logs.push(errorEntry);
+        returnedLogs.push(errorEntry);
+
+        // Write failure status (swallow errors)
+        try {
+          await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
+            if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
+            snapshot.tasks[taskName]!.state = TaskState.FAILED;
+            snapshot.tasks[taskName]!.failedStage = stageName;
+            snapshot.tasks[taskName]!.error = normalized.message;
+          });
+        } catch {
+          // Best-effort
+        }
+
+        return {
+          ok: false,
+          failedStage: stageName,
+          error: normalized,
+          logs: returnedLogs,
+          context,
+        };
       }
-
-      // Flush token write queue
-      await tokenWriteQueue.catch(() => {});
-
-      // Remove LLM listeners
-      events.removeListener("llm:request:complete", onComplete);
-      events.removeListener("llm:request:error", onError);
-
-      return {
-        ok: false,
-        failedStage: stageName,
-        error: normalized,
-        logs: returnedLogs,
-        context,
-      };
     }
-  }
-
-  // 14. After all stages: flush token write queue
-  await tokenWriteQueue.catch(() => {});
 
-  // Remove LLM listeners
-  events.removeListener("llm:request:complete", onComplete);
-  events.removeListener("llm:request:error", onError);
-
-  // Write done status (best-effort)
-  try {
-    const lastStage = KNOWN_STAGES.at(-1)!;
-    const doneProgress = computeDeterministicProgress(
-      pipelineTasks ?? [taskName],
-      taskName,
-      lastStage,
-    );
-    await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
-      snapshot.progress = doneProgress;
-      if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
-      snapshot.tasks[taskName]!.state = TaskState.DONE;
-      snapshot.tasks[taskName]!.currentStage = null;
-    });
-  } catch {
-    // Best-effort
-  }
+    // 14. After all stages: write done status (best-effort).
+    // Token-queue flush + listener removal are handled by withTaskTelemetry's finally.
+    try {
+      const lastStage = KNOWN_STAGES.at(-1)!;
+      const doneProgress = computeDeterministicProgress(
+        pipelineTasks ?? [taskName],
+        taskName,
+        lastStage,
+      );
+      await writeJobStatus(jobDir, (snapshot: StatusSnapshot) => {
+        snapshot.progress = doneProgress;
+        if (!snapshot.tasks[taskName]) snapshot.tasks[taskName] = {};
+        snapshot.tasks[taskName]!.state = TaskState.DONE;
+        snapshot.tasks[taskName]!.currentStage = null;
+      });
+    } catch {
+      // Best-effort
+    }
 
-  context.currentStage = null;
+    context.currentStage = null;
 
-  return {
-    ok: true,
-    logs: returnedLogs,
-    context,
-    llmMetrics,
-  };
+    return {
+      ok: true,
+      logs: returnedLogs,
+      context,
+      llmMetrics,
+    };
+  });
 }
 
 // ─── runPipelineWithModelRouting ──────────────────────────────────────────────
@@ -884,5 +920,6 @@ export async function runPipelineWithModelRouting(
 
 // ─── Re-exports ───────────────────────────────────────────────────────────────
 
+export { captureConsoleOutput } from "./task-telemetry";
 export { decideTransition } from "./lifecycle-policy";
 export { computeDeterministicProgress, KNOWN_STAGES } from "./progress";