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);

tests/unit/pi-safe-stream.test.ts

@@ -0,0 +1,94 @@
+import {
+  createAssistantMessageEventStream,
+  type AssistantMessage,
+  type AssistantMessageEvent,
+  type Context,
+  type Model,
+} from "@mariozechner/pi-ai";
+import { describe, expect, it } from "vitest";
+import { createSafePiStream } from "../../src/main/worker/runtime/pi-safe-stream";
+
+function createModel(): Model<"openai-responses"> {
+  return {
+    id: "mock",
+    name: "mock",
+    api: "openai-responses",
+    provider: "openai",
+    baseUrl: "https://example.invalid",
+    reasoning: false,
+    input: ["text"],
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+    contextWindow: 8192,
+    maxTokens: 2048,
+  };
+}
+
+function createAssistantMessage(): AssistantMessage {
+  return {
+    role: "assistant",
+    content: [{ type: "text", text: "done" }],
+    api: "openai-responses",
+    provider: "openai",
+    model: "mock",
+    usage: {
+      input: 1,
+      output: 1,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 2,
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+    },
+    stopReason: "stop",
+    timestamp: Date.now(),
+  };
+}
+
+const emptyContext: Context = { messages: [] };
+
+describe("pi safe stream", () => {
+  it("forwards successful stream events", async () => {
+    const safeStream = createSafePiStream(() => {
+      const stream = createAssistantMessageEventStream();
+      queueMicrotask(() => {
+        const message = createAssistantMessage();
+        stream.push({ type: "done", reason: "stop", message });
+        stream.end(message);
+      });
+      return stream;
+    });
+
+    const stream = safeStream(createModel(), emptyContext);
+    const events: AssistantMessageEvent[] = [];
+    for await (const event of stream) {
+      events.push(event);
+    }
+
+    expect(events.map((event) => event.type)).toEqual(["done"]);
+    await expect(stream.result()).resolves.toMatchObject({ stopReason: "stop" });
+  });
+
+  it("turns rejected stream setup into an error event", async () => {
+    const safeStream = createSafePiStream(async () => {
+      throw new Error("stream setup failed");
+    });
+
+    const stream = safeStream(createModel(), emptyContext);
+    const events: AssistantMessageEvent[] = [];
+    for await (const event of stream) {
+      events.push(event);
+    }
+    const result = await stream.result();
+
+    expect(events).toHaveLength(1);
+    expect(events[0]).toMatchObject({
+      type: "error",
+      reason: "error",
+      error: {
+        stopReason: "error",
+        errorMessage: "stream setup failed",
+      },
+    });
+    expect(result.stopReason).toBe("error");
+    expect(result.errorMessage).toBe("stream setup failed");
+  });
+});