From 3ba22ce17abaabd1f49f35d9f18121ea4622af44 Mon Sep 17 00:00:00 2001
From: Mario Zechner <badlogicgames@gmail.com>
Date: Tue, 6 Oct 2026 20:08:35 +0200
Subject: [PATCH] fix(ai): keep AssistantMessageEventStream assignable from
 plain event streams

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.
---
 packages/ai/src/api/lazy.ts           |  3 ++-
 packages/ai/src/utils/event-stream.ts | 27 ++++++++++++---------------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/packages/ai/src/api/lazy.ts b/packages/ai/src/api/lazy.ts
index 2f9cf0a35..33ca2ac25 100644
--- a/packages/ai/src/api/lazy.ts
+++ b/packages/ai/src/api/lazy.ts
@@ -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);
 		});
diff --git a/packages/ai/src/utils/event-stream.ts b/packages/ai/src/utils/event-stream.ts
index 669deeaf3..83faa6f63 100644
--- a/packages/ai/src/utils/event-stream.ts
+++ b/packages/ai/src/utils/event-stream.ts
@@ -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) */