diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 808ab5f50..5cbd59a1e 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -6,6 +6,10 @@ - Added `durationMs` to tool result messages, `AgentToolCallOutcome`, and `tool_execution_end` events: how long `execute()` took, measured with a monotonic clock and excluding hooks. Calls that did not run have none ([#10549](https://github.com/earendil-works/pi/issues/10549)) +### Changed + +- `streamProxy()` returns an `AssistantMessageEventStream`, so proxied responses get `durationMs` like direct ones + ## [1.0.4] - 2026-10-05 ## [1.0.3] - 2026-10-05 diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 0a4711c55..c5f6388ed 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -7,7 +7,8 @@ import { type AssistantMessage, type AssistantMessageEvent, - EventStream, + type AssistantMessageEventStream, + createAssistantMessageEventStream, type Model, parseStreamingJson, type SimpleStreamOptions, @@ -16,20 +17,6 @@ import { type TranscriptContext, } from "@earendil-works/pi-ai"; -// Create stream class matching ProxyMessageEventStream -class ProxyMessageEventStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - /** * Proxy event types - server sends these with partial field stripped to reduce bandwidth. */ @@ -121,8 +108,8 @@ export function streamProxy( model: Model, context: TranscriptContext, options: ProxyStreamOptions, -): ProxyMessageEventStream { - const stream = new ProxyMessageEventStream(); +): AssistantMessageEventStream { + const stream = createAssistantMessageEventStream(); (async () => { // Initialize the partial message that we'll build up from events diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 2ac2ada6a..4a9386a13 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -1,7 +1,6 @@ import { type AssistantMessage, - type AssistantMessageEvent, - EventStream, + createAssistantMessageEventStream, type Message, type Model, type ToolResultMessage, @@ -20,20 +19,6 @@ import type { AgentToolCall, } from "../src/types.ts"; -// Mock stream for testing - mimics MockAssistantStream -class MockAssistantStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - function createUsage() { return { input: 0, @@ -96,7 +81,7 @@ describe("default stream function compatibility", () => { let calls = 0; setDefaultStreamFn(() => { calls++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", @@ -140,7 +125,7 @@ describe("agentLoop with AgentMessage", () => { }; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Hi there!" }]); stream.push({ type: "done", reason: "stop", message }); @@ -196,7 +181,7 @@ describe("agentLoop with AgentMessage", () => { // The provider receives a transcript: no top-level prompt or tool fields. expect(Object.keys(providerContext)).toEqual(["messages"]); expect(providerContext.messages[0]).toBe(initialSystem); - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -245,7 +230,7 @@ describe("agentLoop with AgentMessage", () => { }; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); @@ -297,7 +282,7 @@ describe("agentLoop with AgentMessage", () => { }; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); @@ -370,7 +355,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { // First call: return tool call @@ -437,7 +422,7 @@ describe("agentLoop with AgentMessage", () => { }; let callIndex = 0; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = callIndex === 0 @@ -497,7 +482,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { // Output hit the token limit mid tool call. The salvage parser can @@ -576,7 +561,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -648,7 +633,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -723,7 +708,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -826,7 +811,7 @@ describe("agentLoop with AgentMessage", () => { ); } - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { // First call: return two tool calls @@ -923,7 +908,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -1016,7 +1001,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -1090,7 +1075,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -1156,7 +1141,7 @@ describe("agentLoop with AgentMessage", () => { }, undefined, () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", @@ -1207,7 +1192,7 @@ describe("agentLoop with AgentMessage", () => { undefined, () => { providerCalls++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "error", @@ -1264,7 +1249,7 @@ describe("agentLoop with AgentMessage", () => { const stream = agentLoop([createUserMessage("run")], { messages: [], tools: [tool] }, config, undefined, () => { providerCalls++; - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -1302,7 +1287,7 @@ describe("agentLoop with AgentMessage", () => { undefined, () => { providerCalls++; - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -1343,7 +1328,7 @@ describe("agentLoop with AgentMessage", () => { const stream = agentLoop([createUserMessage("run")], { messages: [], tools: [tool] }, config, undefined, () => { providerCalls++; - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { const message = providerCalls === 1 @@ -1398,7 +1383,7 @@ describe("agentLoop with AgentMessage", () => { ), ); } - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -1455,7 +1440,7 @@ describe("agentLoop with AgentMessage", () => { expect(model).toBe(replacementModel); expect(context.messages).toEqual([canonicalMessage]); expect(options?.reasoning).toBe("high"); - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -1496,7 +1481,7 @@ describe("agentLoop with AgentMessage", () => { undefined, (_model, context) => { requestIncludedSteering.push(context.messages.includes(lateSteering)); - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { response.push({ type: "done", @@ -1561,7 +1546,7 @@ describe("agentLoop with AgentMessage", () => { (message) => message.role === "system" && message.content === "updated guidance", ); } - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (llmCalls === 1) { mockStream.push({ @@ -1625,7 +1610,7 @@ describe("agentLoop with AgentMessage", () => { (_model, context) => { providerCalls++; if (providerCalls === 2) secondRequestIncludedSteering = context.messages.includes(lateSteering); - const response = new MockAssistantStream(); + const response = createAssistantMessageEventStream(); queueMicrotask(() => { const message = providerCalls === 1 @@ -1693,7 +1678,7 @@ describe("agentLoop with AgentMessage", () => { let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (llmCalls === 1) { const message = createAssistantMessage( @@ -1773,7 +1758,7 @@ describe("agentLoop with AgentMessage", () => { let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], @@ -1824,7 +1809,7 @@ describe("agentLoop with AgentMessage", () => { let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = llmCalls === 1 @@ -1886,7 +1871,7 @@ describe("agentLoop with AgentMessage", () => { let llmCalls = 0; const stream = agentLoop([createUserMessage("echo both")], context, config, undefined, () => { llmCalls++; - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = llmCalls === 1 @@ -1940,7 +1925,7 @@ describe("agentLoop with AgentMessage", () => { let callIndex = 0; const stream = agentLoop([createUserMessage("echo both")], context, config, undefined, () => { - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( @@ -2005,7 +1990,7 @@ describe("agentLoop with AgentMessage", () => { let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; - const mockStream = new MockAssistantStream(); + const mockStream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], @@ -2057,7 +2042,7 @@ describe("agentLoopContinue with AgentMessage", () => { }; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); @@ -2123,7 +2108,7 @@ describe("agentLoopContinue with AgentMessage", () => { }; const streamFn = () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response to custom message" }]); stream.push({ type: "done", reason: "stop", message }); diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 2a339944e..8ada745d5 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -1,7 +1,6 @@ import { type AssistantMessage, - type AssistantMessageEvent, - EventStream, + createAssistantMessageEventStream, getCurrentSystemMessage, getModel, toToolDeclaration, @@ -18,20 +17,6 @@ import { setDefaultStreamFn, } from "../src/index.ts"; -// Mock stream that mimics AssistantMessageEventStream -class MockAssistantStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - function createUserMessage(text: string): UserMessage { return { role: "user", content: text, timestamp: Date.now() }; } @@ -108,7 +93,7 @@ describe("Agent", () => { let calls = 0; setDefaultStreamFn(() => { calls++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage("fallback"); stream.push({ type: "done", reason: "stop", message }); @@ -192,7 +177,7 @@ describe("Agent", () => { : [], ), ); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); }); @@ -235,7 +220,7 @@ describe("Agent", () => { initialState: { systemPrompt: "You are helpful." }, streamFn: (_model, context) => { expect(context.messages.filter((message) => message.role === "system")).toHaveLength(2); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); }); @@ -262,7 +247,7 @@ describe("Agent", () => { const agent = new Agent({ initialState: { systemPrompt: "You are helpful.", tools: [createTool("first")] }, streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); }); @@ -377,7 +362,7 @@ describe("Agent", () => { const barrier = createDeferred(); const agent = new Agent({ streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") }); }); @@ -415,7 +400,7 @@ describe("Agent", () => { const barrier = createDeferred(); const agent = new Agent({ streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") }); }); @@ -450,7 +435,7 @@ describe("Agent", () => { let receivedSignal: AbortSignal | undefined; const agent = new Agent({ streamFn: (_model, _context, options) => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); const checkAbort = () => { @@ -513,7 +498,7 @@ describe("Agent", () => { const agent = new Agent({ initialState: { tools: [tool] }, streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", @@ -588,7 +573,7 @@ describe("Agent", () => { const agent = new Agent({ initialState: { tools: [settledTool, slowTool] }, streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", @@ -692,7 +677,7 @@ describe("Agent", () => { const releaseResponse = createDeferred(); const agent = new Agent({ streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(async () => { stream.push({ type: "start", partial: createAssistantMessage("") }); streamStarted.resolve(); @@ -727,7 +712,7 @@ describe("Agent", () => { // Use a stream function that responds to abort streamFn: (_model, _context, options) => { abortSignal = options?.signal; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); // Check abort signal periodically @@ -766,7 +751,7 @@ describe("Agent", () => { const agent = new Agent({ streamFn: (_model, _context, options) => { abortSignal = options?.signal; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); const checkAbort = () => { @@ -800,7 +785,7 @@ describe("Agent", () => { it("continue() should process queued follow-up messages after an assistant turn", async () => { const agent = new Agent({ streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") }); }); @@ -848,7 +833,7 @@ describe("Agent", () => { message.role === "user" && typeof message.content === "string" ? [message.content] : [], ), ); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") }); }); @@ -890,7 +875,7 @@ describe("Agent", () => { }, streamFn: () => { requestCount++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (requestCount === 1) { const message = createAssistantToolUseMessage([ @@ -933,7 +918,7 @@ describe("Agent", () => { }, streamFn: () => { requestCount++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (requestCount === 1) { const message = createAssistantToolUseMessage([ @@ -1002,7 +987,7 @@ describe("Agent", () => { message.role === "user" && typeof message.content === "string" ? [message.content] : [], ), ); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), ); @@ -1032,7 +1017,7 @@ describe("Agent", () => { message.role === "user" && typeof message.content === "string" ? [message.content] : [], ), ); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), ); @@ -1064,7 +1049,7 @@ describe("Agent", () => { message.role === "user" && typeof message.content === "string" ? [message.content] : [], ), ); - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), ); @@ -1090,7 +1075,7 @@ describe("Agent", () => { const agent = new Agent({ finishTurn: () => ({ action: "continue" }), streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "error", @@ -1126,7 +1111,7 @@ describe("Agent", () => { const agent = new Agent({ finishTurn: () => ({ action: "end" }), streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), ); @@ -1151,7 +1136,7 @@ describe("Agent", () => { const agent = new Agent({ steeringMode: "one-at-a-time", followUpMode: "all", - streamFn: () => new MockAssistantStream(), + streamFn: () => createAssistantMessageEventStream(), }); const first = createUserMessage("first steering"); const second = createUserMessage("second steering"); @@ -1173,7 +1158,7 @@ describe("Agent", () => { providerEvents.push(data); }, streamFn: (model, _context, options) => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(async () => { await options?.onProviderStreamEvent?.({ request_cost: 0.01 }, model); const message = createAssistantMessage("ok"); @@ -1194,7 +1179,7 @@ describe("Agent", () => { sessionId: "session-abc", streamFn: (_model, _context, options) => { receivedSessionId = options?.sessionId; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = createAssistantMessage("ok"); stream.push({ type: "done", reason: "stop", message }); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index d15bac1ce..efc4386d8 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Breaking Changes + +- A stream function must return an `AssistantMessageEventStream`, for example from `createAssistantMessageEventStream()`; a hand-written `EventStream` subclass no longer type-checks in its place + ### Added - Added `durationMs` to `AssistantMessage`: `AssistantMessageEventStream` measures each response with a monotonic clock from the start of the request to its final message, for every API implementation, including direct calls. Deferred results fetched later stay untimed ([#10549](https://github.com/earendil-works/pi/issues/10549)) diff --git a/packages/ai/src/utils/event-stream.ts b/packages/ai/src/utils/event-stream.ts index 83faa6f63..eb968c016 100644 --- a/packages/ai/src/utils/event-stream.ts +++ b/packages/ai/src/utils/event-stream.ts @@ -88,9 +88,6 @@ export class EventStream implements AsyncIterable { } } -/** Start of each assistant stream, kept outside the class so it stays structurally a plain `EventStream`. */ -const streamStarts = new WeakMap(); - /** * Event stream of one assistant response. It also times the response: the final message (`done` or `error` event, or * the result passed to `end()`) gets `durationMs`, measured with a monotonic clock from the stream's creation, unless @@ -98,6 +95,9 @@ const streamStarts = new WeakMap { + readonly #startedAt = Date.now(); + readonly #startedAtMonotonic = performance.now(); + constructor() { super( (event) => event.type === "done" || event.type === "error", @@ -110,25 +110,23 @@ export class AssistantMessageEventStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - function createAssistantMessage(text: string): AssistantMessage { return { role: "assistant", @@ -96,7 +81,7 @@ describe("AgentSession concurrent prompt guard", () => { }, streamFn: (_model, _context, options) => { abortSignal = options?.signal; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); const checkAbort = () => { @@ -201,7 +186,7 @@ describe("AgentSession concurrent prompt guard", () => { }, streamFn: (_model, context, options) => { abortSignal = options?.signal; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const userTexts = context.messages .filter((message) => message.role === "user") @@ -306,7 +291,7 @@ describe("AgentSession concurrent prompt guard", () => { tools: [], }, streamFn: () => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") }); @@ -367,7 +352,7 @@ describe("AgentSession concurrent prompt guard", () => { tools: [tool], }, streamFn: async (_model, context) => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const toolResultCount = context.messages.filter((message) => message.role === "toolResult").length; if (toolResultCount > 0) { @@ -515,7 +500,7 @@ describe("AgentSession concurrent prompt guard", () => { tools: [tool], }, streamFn: async (_model, context) => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const hasToolResult = context.messages.some((message) => message.role === "toolResult"); diff --git a/packages/coding-agent/test/agent-session-retry.test.ts b/packages/coding-agent/test/agent-session-retry.test.ts index 9d80eb946..206b25b60 100644 --- a/packages/coding-agent/test/agent-session-retry.test.ts +++ b/packages/coding-agent/test/agent-session-retry.test.ts @@ -2,7 +2,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent, type AgentEvent, type AgentTool } from "@earendil-works/pi-agent-core"; -import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@earendil-works/pi-ai/compat"; +import { type AssistantMessage, createAssistantMessageEventStream, getModel } from "@earendil-works/pi-ai/compat"; import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.ts"; @@ -12,19 +12,6 @@ import { SettingsManager } from "../src/core/settings-manager.ts"; import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader } from "./utilities.ts"; -class MockAssistantStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - function createAssistantMessage(text: string, overrides?: Partial): AssistantMessage { return { role: "assistant", @@ -86,7 +73,7 @@ describe("AgentSession retry", () => { initialState: { model, systemPrompt: "Test", tools: [] }, streamFn: () => { callCount++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callCount <= failCount) { const msg = createAssistantMessage("", { @@ -194,7 +181,7 @@ describe("AgentSession retry", () => { let callCount = 0; const streamFn = () => { callCount++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callCount === 1) { const msg = createAssistantMessage("", { @@ -272,7 +259,7 @@ describe("AgentSession retry", () => { initialState: { model, systemPrompt: "Test", tools: [] }, streamFn: () => { callCount++; - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { if (callCount === 1) { // First call: overloaded error diff --git a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts index 98325b19d..18c40646c 100644 --- a/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts +++ b/packages/coding-agent/test/rpc-prompt-response-semantics.test.ts @@ -4,8 +4,7 @@ import { join } from "node:path"; import { Agent } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, - type AssistantMessageEvent, - EventStream, + createAssistantMessageEventStream, getModel, type Model, } from "@earendil-works/pi-ai/compat"; @@ -44,19 +43,6 @@ vi.mock("../src/modes/rpc/jsonl.js", () => ({ serializeJsonLine: (value: unknown) => `${JSON.stringify(value)}\n`, })); -class MockAssistantStream extends EventStream { - constructor() { - super( - (event) => event.type === "done" || event.type === "error", - (event) => { - if (event.type === "done") return event.message; - if (event.type === "error") return event.error; - throw new Error("Unexpected event type"); - }, - ); - } -} - function createAssistantMessage(text: string): AssistantMessage { return { role: "assistant", @@ -121,7 +107,7 @@ async function createRuntimeHost(options: { tools: [], }, streamFn: (_model, _context, _options) => { - const stream = new MockAssistantStream(); + const stream = createAssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); setTimeout(() => {