Skip to content

EventStream: push in, async iterator out

源码版本v0.73.1

utils/event-stream.ts is pi-ai's streaming primitive — a generic EventStream<T, R>. Producers push events in via push(), consumers pull events out via for await, and call result() to get the final value. Underneath there is no WebSocket, no Observable — just a hand-written queue + waiter array that translates a push-based producer API into a pull-based async iterator. All providers ultimately push events into AssistantMessageEventStream (the concrete subclass of EventStream in this file), and callers see no difference between them.

Responsibilities

  1. Queue + waiter: on push, if a consumer is waiting, resolve it directly; otherwise push to queue, see packages/ai/src/utils/event-stream.ts:20-35.
  2. Async iterator: [Symbol.asyncIterator] implements AsyncIterator<T> — drain the queue first, then wait on the waiter list, see packages/ai/src/utils/event-stream.ts:49-61.
  3. Completion signal: isComplete(event) decides which event counts as a terminal state, triggering resolveFinalResult(extractResult(event)), see packages/ai/src/utils/event-stream.ts:23-26.
  4. Final result Promise: result(): Promise<R> returns the result extracted from the terminal event, decoupled from for await, see packages/ai/src/utils/event-stream.ts:63-65.
  5. End fallback: end(result?) forces completion and clears all waiters, see packages/ai/src/utils/event-stream.ts:37-47.
  6. AssistantMessage subclass: AssistantMessageEventStream treats done / error as terminal states, and done.message / error.error as the result, see packages/ai/src/utils/event-stream.ts:68-82.

Design motivation

Why not just use Node's EventEmitter or RxJS Observable? Because the consumer wants for await syntax — the pull pace is controlled by the consumer, so backpressure comes for free. EventEmitter is push without backpressure, and converting an Observable to an async iterator needs a from adapter. An 80-line class collapses both concerns into one, and throws in "terminal event triggers the result Promise" for free — the complete function can return an AssistantMessage with a single stream(...).result(), without the consumer manually accumulating.

Why is AssistantMessageEventStream a subclass rather than a generic parameter? Because "which event counts as terminal, and how to extract the result from it" is a provider-agnostic policy — every provider pushes AssistantMessageEvent, and done.message is the final result. The subclass bakes the policy in, so providers do not each have to pass isComplete / extractResult.

Key files

push is the producer's only entry point — two paths: resolve a waiter or enqueue:

typescript
// packages/ai/src/utils/event-stream.ts:20-35
push(event: T): void {
	if (this.done) return;

	if (this.isComplete(event)) {
		this.done = true;
		this.resolveFinalResult(this.extractResult(event));
	}

	// Deliver to waiting consumer or queue it
	const waiter = this.waiting.shift();
	if (waiter) {
		waiter({ value: event, done: false });
	} else {
		this.queue.push(event);
	}
}

The async iterator drains enqueued events first, then waits:

typescript
// packages/ai/src/utils/event-stream.ts:49-61
async *[Symbol.asyncIterator](): AsyncIterator<T> {
	while (true) {
		if (this.queue.length > 0) {
			yield this.queue.shift()!;
		} else if (this.done) {
			return;
		} else {
			const result = await new Promise<IteratorResult<T>>((resolve) => this.waiting.push(resolve));
			if (result.done) return;
			yield result.value;
		}
	}
}

AssistantMessageEventStream bakes the policy in, so providers do not have to care about terminal-state detection:

typescript
// packages/ai/src/utils/event-stream.ts:68-82
export class AssistantMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
	constructor() {
		super(
			(event) => event.type === "done" || event.type === "error",
			(event) => {
				if (event.type === "done") {
					return event.message;
				} else if (event.type === "error") {
					return event.error;
				}
				throw new Error("Unexpected event type for final result");
			},
		);
	}
}

Data flow

The interaction between producer and consumer:

Boundaries and failure modes

  • Push after done: if (this.done) return silently ignores; a provider pushing after error is harmless, see packages/ai/src/utils/event-stream.ts:21-21.
  • end before a terminal event: when end(result?) is called without a result, finalResultPromise may never resolve. The caller should ensure a terminal event has been pushed first or pass a result, see packages/ai/src/utils/event-stream.ts:37-41.
  • Error as terminal state: AssistantMessageEventStream's extractResult returns event.error on error, so result() on an error stream still yields the error object instead of hanging, see packages/ai/src/utils/event-stream.ts:75-77.
  • Multiple consumers: multiple concurrent for await loops will contend for the waiter queue, and events will be distributed across consumers — but the real-world usage is always single-consumer; the code does not constrain multi-consumer behavior.

Summary

EventStream is an 80-line hand-written push-to-async-iterator adapter, and the AssistantMessageEventStream subclass bakes in the terminal-state policy. Providers all push AssistantMessageEvent, and callers all use for await + result(). For how providers use it, see Anthropic SSE implementation; for the top-level dispatch, see stream/complete facade.