fix(runtime): guard pi stream setup failures

src/main/worker/runtime/index.ts

@@ -211,8 +211,18 @@ export class WorkerRunner {
               taskSessionId: job.taskSessionId,
               message: event.message,
             });
-            const maybeAssistant = event.message as { role?: string; usage?: { input?: number; output?: number; cost?: { total?: number } } };
-            if (maybeAssistant.role === "assistant" && maybeAssistant.usage) {
+            const maybeAssistant = event.message as {
+              role?: string;
+              stopReason?: string;
+              errorMessage?: string;
+              usage?: { input?: number; output?: number; cost?: { total?: number } };
+            };
+            if (
+              maybeAssistant.role === "assistant"
+              && (maybeAssistant.stopReason === "error" || maybeAssistant.stopReason === "aborted")
+            ) {
+              errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`;
+            } else if (maybeAssistant.role === "assistant" && maybeAssistant.usage) {
               finalUsage = {
                 inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
                 outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,

src/main/worker/runtime/pi-agent-factory.ts

@@ -18,10 +18,11 @@ import {
   type AgentMessage,
   type ThinkingLevel,
 } from "@mariozechner/pi-agent-core";
-import { streamSimple, type Api, type Model } from "@mariozechner/pi-ai";
+import type { Api, Model } from "@mariozechner/pi-ai";
 import { buildAgentHooks } from "@main/agent-runtime/agent-hooks";
 import type { ContextCompactionTelemetry } from "@main/agent-runtime/agent-hooks";
 import type { WorkerRpcClient } from "../entry/rpc";
+import { safeStreamSimple } from "./pi-safe-stream";
 // Side-effect import: drags the hook signature contract into the typecheck
 // graph so any drift in pi-agent-core hook shapes breaks the build here.
 import "./pi-hook-contract";
@@ -78,7 +79,7 @@ export function createPiAgent(input: PiAgentFactoryInput): Agent {
     },
     sessionId: input.taskSessionId,
     streamFn: (model, messages, options) =>
-      streamSimple(model, messages, {
+      safeStreamSimple(model, messages, {
         ...options,
         apiKey: input.apiKey,
         headers: { ...(options?.headers ?? {}), ...(input.extraHeaders ?? {}) },

src/main/worker/runtime/pi-safe-stream.ts

@@ -0,0 +1,88 @@
+import {
+  createAssistantMessageEventStream,
+  streamSimple,
+  type Api,
+  type AssistantMessage,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+  type StopReason,
+  type Usage,
+} from "@mariozechner/pi-ai";
+
+type PiStreamSource = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;
+
+type SafePiStreamSource = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+type ErrorAssistantMessage = AssistantMessage & {
+  stopReason: Extract<StopReason, "error" | "aborted">;
+};
+
+const EMPTY_USAGE: Usage = {
+  input: 0,
+  output: 0,
+  cacheRead: 0,
+  cacheWrite: 0,
+  totalTokens: 0,
+  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+function formatStreamError(error: unknown): string {
+  if (error instanceof Error) return error.message;
+  return String(error);
+}
+
+function createStreamErrorMessage(model: Model<Api>, error: unknown, signal?: AbortSignal): ErrorAssistantMessage {
+  const stopReason: Extract<StopReason, "error" | "aborted"> = signal?.aborted ? "aborted" : "error";
+  return {
+    role: "assistant",
+    content: [{ type: "text", text: "" }],
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    usage: EMPTY_USAGE,
+    stopReason,
+    errorMessage: formatStreamError(error),
+    timestamp: Date.now(),
+  };
+}
+
+function pushStreamError(
+  stream: AssistantMessageEventStream,
+  model: Model<Api>,
+  error: unknown,
+  signal?: AbortSignal,
+): void {
+  const message = createStreamErrorMessage(model, error, signal);
+  stream.push({ type: "error", reason: message.stopReason, error: message });
+  stream.end(message);
+}
+
+export function createSafePiStream(source: PiStreamSource): SafePiStreamSource {
+  return (model, context, options) => {
+    const target = createAssistantMessageEventStream();
+    void Promise.resolve()
+      .then(() => source(model, context, options))
+      .then(async (sourceStream) => {
+        for await (const event of sourceStream) {
+          target.push(event);
+        }
+        target.end();
+      })
+      .catch((error) => {
+        pushStreamError(target, model, error, options?.signal);
+      });
+    return target;
+  };
+}
+
+export const safeStreamSimple = createSafePiStream(streamSimple);