diff --git a/src/main/worker/runtime/index.ts b/src/main/worker/runtime/index.ts index 044c9c74..43283bea 100644 --- a/src/main/worker/runtime/index.ts +++ b/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, diff --git a/src/main/worker/runtime/pi-agent-factory.ts b/src/main/worker/runtime/pi-agent-factory.ts index 2a6fb521..8e77760f 100644 --- a/src/main/worker/runtime/pi-agent-factory.ts +++ b/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 ?? {}) }, diff --git a/src/main/worker/runtime/pi-safe-stream.ts b/src/main/worker/runtime/pi-safe-stream.ts new file mode 100644 index 00000000..a353ad4b --- /dev/null +++ b/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, + context: Context, + options?: SimpleStreamOptions, +) => AssistantMessageEventStream | Promise; + +type SafePiStreamSource = ( + model: Model, + context: Context, + options?: SimpleStreamOptions, +) => AssistantMessageEventStream; + +type ErrorAssistantMessage = AssistantMessage & { + stopReason: Extract; +}; + +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, error: unknown, signal?: AbortSignal): ErrorAssistantMessage { + const stopReason: Extract = 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, + 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); diff --git a/tests/unit/pi-safe-stream.test.ts b/tests/unit/pi-safe-stream.test.ts new file mode 100644 index 00000000..5549fe65 --- /dev/null +++ b/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"); + }); +});