sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(ai): time responses with private fields; use the real stream in proxy and test mocks

Browse source 
Replaces the WeakMap workaround. streamProxy and the agent and coding-agent test
mocks reimplemented AssistantMessageEventStream as plain EventStream subclasses;
they now use createAssistantMessageEventStream().
This commit is contained in:
Mario Zechner 2026-10-06 20:26:27 +02:00
parent 3ba22ce17a
commit f284a2460c
9 changed files with 93 additions and 172 deletions
Download patch file Download diff file Expand all files Collapse all files

4
packages/agent/CHANGELOG.md
View file

@ -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

21
packages/agent/src/proxy.ts
View file

@ -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<AssistantMessageEvent, AssistantMessage> {
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<any>,
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

83
packages/agent/test/agent-loop.test.ts
View file

@ -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<AssistantMessageEvent, AssistantMessage> {
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 });

65
packages/agent/test/agent.test.ts
View file

@ -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<AssistantMessageEvent, AssistantMessage> {
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 });

4
packages/ai/CHANGELOG.md
View file

@ -2,6 +2,10 @@
## [Unreleased]
### Breaking Changes
- A stream function must return an `AssistantMessageEventStream`, for example from `createAssistantMessageEventStream()`; a hand-written `EventStream<AssistantMessageEvent, AssistantMessage>` 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))

22
packages/ai/src/utils/event-stream.ts
View file

@ -88,9 +88,6 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
}
}
/** Start of each assistant stream, kept outside the class so it stays structurally a plain `EventStream`. */
const streamStarts = new WeakMap<AssistantMessageEventStream, { readonly wall: number; readonly monotonic: number }>();
/**
* 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<AssistantMessageEventStream, { readonly wall: n
* elsewhere, such as a deferred result fetched later, therefore leaves it untimed.
*/
export class AssistantMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
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<AssistantMessageEve
throw new Error("Unexpected event type for final result");
},
);
streamStarts.set(this, { wall: Date.now(), monotonic: performance.now() });
}
override push(event: AssistantMessageEvent): void {
if (!this.done && event.type === "done") timeResponse(this, event.message);
else if (!this.done && event.type === "error") timeResponse(this, event.error);
if (event.type === "done") this.#time(event.message);
else if (event.type === "error") this.#time(event.error);
super.push(event);
}
override end(result?: AssistantMessage): void {
if (!this.done && result !== undefined) timeResponse(this, result);
if (result !== undefined) this.#time(result);
super.end(result);
}
}
function timeResponse(stream: AssistantMessageEventStream, message: AssistantMessage): void {
const start = streamStarts.get(stream);
if (start === undefined || message.durationMs !== undefined || message.timestamp < start.wall) return;
message.durationMs = Math.max(0, Math.round(performance.now() - start.monotonic));
#time(message: AssistantMessage): void {
if (this.done || message.durationMs !== undefined || message.timestamp < this.#startedAt) return;
message.durationMs = Math.max(0, Math.round(performance.now() - this.#startedAtMonotonic));
}
}
/** Factory function for AssistantMessageEventStream (for use in extensions) */

27
packages/coding-agent/test/agent-session-concurrent.test.ts
View file

@ -9,8 +9,7 @@ import { join } from "node:path";
import { Agent } from "@earendil-works/pi-agent-core";
import {
type AssistantMessage,
type AssistantMessageEvent,
EventStream,
createAssistantMessageEventStream,
getModel,
type ImageContent,
type TextContent,
@ -28,20 +27,6 @@ import {
} from "../src/core/system-prompt.ts";
import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.ts";
// Mock stream that mimics AssistantMessageEventStream
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
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");

21
packages/coding-agent/test/agent-session-retry.test.ts
View file

@ -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<AssistantMessageEvent, AssistantMessage> {
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>): 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

18
packages/coding-agent/test/rpc-prompt-response-semantics.test.ts
View file

@ -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<AssistantMessageEvent, AssistantMessage> {
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(() => {