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): keep AssistantMessageEventStream assignable from plain event streams

Browse source 
The response timing added private fields, which made test and host streams built on
EventStream no longer assignable to StreamFn. Keep the start in a WeakMap instead.
This commit is contained in:
Mario Zechner 2026-10-06 20:08:35 +02:00
parent 4dd2af42c8
commit 3ba22ce17a
2 changed files with 14 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

3
packages/ai/src/api/lazy.ts
View file

@ -47,12 +47,13 @@ export function lazyStream(
model: Model<Api>,
setup: () => Promise<AsyncIterable<AssistantMessageEvent>>,
): AssistantMessageEventStream {
const startedAt = Date.now();
const outer = new AssistantMessageEventStream();
setup()
.then((inner) => forwardStream(outer, inner))
.catch((error) => {
const message = createSetupErrorMessage(model, error, outer.startedAt);
const message = createSetupErrorMessage(model, error, startedAt);
outer.push({ type: "error", reason: "error", error: message });
outer.end(message);
});

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

@ -88,6 +88,9 @@ 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
@ -95,9 +98,6 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
* 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,28 +110,25 @@ export class AssistantMessageEventStream extends EventStream<AssistantMessageEve
throw new Error("Unexpected event type for final result");
},
);
}
/** Wall-clock start of this stream, for messages created before the response could start. */
get startedAt(): number {
return this.#startedAt;
streamStarts.set(this, { wall: Date.now(), monotonic: performance.now() });
}
override push(event: AssistantMessageEvent): void {
if (event.type === "done") this.#time(event.message);
else if (event.type === "error") this.#time(event.error);
if (!this.done && event.type === "done") timeResponse(this, event.message);
else if (!this.done && event.type === "error") timeResponse(this, event.error);
super.push(event);
}
override end(result?: AssistantMessage): void {
if (result !== undefined) this.#time(result);
if (!this.done && result !== undefined) timeResponse(this, result);
super.end(result);
}
}
#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));
}
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));
}
/** Factory function for AssistantMessageEventStream (for use in extensions) */