fix(agent): await subscribed event handlers

packages/agent/CHANGELOG.md

@@ -28,6 +28,11 @@
   - `agent.getSteeringMode()` -> `agent.steeringMode`
   - `agent.setFollowUpMode(mode)` -> `agent.followUpMode = mode`
   - `agent.getFollowUpMode()` -> `agent.followUpMode`
+- `Agent.subscribe()` listeners are now awaited and receive the active `AbortSignal`:
+  - `agent.subscribe((event) => { ... })` -> `agent.subscribe(async (event, signal) => { ... })`
+  - `agent_end` is now the final emitted event for a run, but not the idle boundary
+  - `agent.waitForIdle()`, `agent.prompt(...)`, and `agent.continue()` now settle only after awaited `agent_end` listeners finish
+  - `agent.state.isStreaming` remains `true` until that settlement completes
 
 ## [0.64.0] - 2026-03-29
 

packages/agent/README.md

@@ -124,7 +124,7 @@ The last message in context must be `user` or `toolResult` (not `assistant`).
 | Event | Description |
 |-------|-------------|
 | `agent_start` | Agent begins processing |
-| `agent_end` | Agent completes with all new messages |
+| `agent_end` | Final event for the run. Awaited subscribers for this event still count toward settlement |
 | `turn_start` | New turn begins (one LLM call + tool executions) |
 | `turn_end` | Turn completes with assistant message and tool results |
 | `message_start` | Any message begins (user, assistant, toolResult) |
@@ -133,6 +133,8 @@ The last message in context must be `user` or `toolResult` (not `assistant`).
 | `tool_execution_start` | Tool begins |
 | `tool_execution_update` | Tool streams progress |
 | `tool_execution_end` | Tool completes |
++
++`Agent.subscribe()` listeners are awaited in registration order. `agent_end` means no more loop events will be emitted, but `await agent.waitForIdle()` and `await agent.prompt(...)` only settle after awaited `agent_end` listeners finish.
 
 ## Agent Options
 
@@ -217,6 +219,8 @@ Assigning `agent.state.tools = [...]` or `agent.state.messages = [...]` copies t
 
 During streaming, `agent.state.streamingMessage` contains the current partial assistant message.
 
+`agent.state.isStreaming` remains `true` until the run fully settles, including awaited `agent_end` subscribers.
+
 ## Methods
 
 ### Prompting
@@ -275,8 +279,11 @@ await agent.waitForIdle(); // Wait for completion
 ### Events
 
 ```typescript
-const unsubscribe = agent.subscribe((event) => {
-  console.log(event.type);
+const unsubscribe = agent.subscribe(async (event, signal) => {
+  if (event.type === "agent_end") {
+    // Final barrier work for the run
+    await flushSessionState(signal);
+  }
 });
 unsubscribe();
 ```

packages/agent/src/agent.ts

@@ -156,7 +156,7 @@ type ActiveRun = {
  */
 export class Agent {
 	private _state: MutableAgentState;
-	private readonly listeners = new Set<(event: AgentEvent) => void>();
+	private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise<void> | void>();
 	private readonly steeringQueue: PendingMessageQueue;
 	private readonly followUpQueue: PendingMessageQueue;
 
@@ -203,8 +203,17 @@ export class Agent {
 		this.toolExecution = options.toolExecution ?? "parallel";
 	}
 
-	/** Subscribe to agent lifecycle events. Returns an unsubscribe function. */
-	subscribe(listener: (event: AgentEvent) => void): () => void {
+	/**
+	 * Subscribe to agent lifecycle events.
+	 *
+	 * Listener promises are awaited in subscription order and are included in
+	 * the current run's settlement. Listeners also receive the active abort
+	 * signal for the current run.
+	 *
+	 * `agent_end` is the final emitted event for a run, but the agent does not
+	 * become idle until all awaited listeners for that event have settled.
+	 */
+	subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise<void> | void): () => void {
 		this.listeners.add(listener);
 		return () => this.listeners.delete(listener);
 	}
@@ -277,7 +286,11 @@ export class Agent {
 		this.activeRun?.abortController.abort();
 	}
 
-	/** Resolve when the current run has finished. */
+	/**
+	 * Resolve when the current run and all awaited event listeners have finished.
+	 *
+	 * This resolves after `agent_end` listeners settle.
+	 */
 	waitForIdle(): Promise<void> {
 		return this.activeRun?.promise ?? Promise.resolve();
 	}
@@ -437,13 +450,13 @@ export class Agent {
 		try {
 			await executor(abortController.signal);
 		} catch (error) {
-			this.handleRunFailure(error, abortController.signal.aborted);
+			await this.handleRunFailure(error, abortController.signal.aborted);
 		} finally {
 			this.finishRun();
 		}
 	}
 
-	private handleRunFailure(error: unknown, aborted: boolean): void {
+	private async handleRunFailure(error: unknown, aborted: boolean): Promise<void> {
 		const failureMessage = {
 			role: "assistant",
 			content: [{ type: "text", text: "" }],
@@ -457,7 +470,7 @@ export class Agent {
 		} satisfies AgentMessage;
 		this._state.messages.push(failureMessage);
 		this._state.errorMessage = failureMessage.errorMessage;
-		this.processEvents({ type: "agent_end", messages: [failureMessage] });
+		await this.processEvents({ type: "agent_end", messages: [failureMessage] });
 	}
 
 	private finishRun(): void {
@@ -468,7 +481,14 @@ export class Agent {
 		this.activeRun = undefined;
 	}
 
-	private processEvents(event: AgentEvent): void {
+	/**
+	 * Reduce internal state for a loop event, then await listeners.
+	 *
+	 * `agent_end` only means no further loop events will be emitted. The run is
+	 * considered idle later, after all awaited listeners for `agent_end` finish
+	 * and `finishRun()` clears runtime-owned state.
+	 */
+	private async processEvents(event: AgentEvent): Promise<void> {
 		switch (event.type) {
 			case "message_start":
 				this._state.streamingMessage = event.message;
@@ -504,13 +524,16 @@ export class Agent {
 				break;
 
 			case "agent_end":
-				this._state.isStreaming = false;
 				this._state.streamingMessage = undefined;
 				break;
 		}
 
+		const signal = this.activeRun?.abortController.signal;
+		if (!signal) {
+			throw new Error("Agent listener invoked outside active run");
+		}
 		for (const listener of this.listeners) {
-			listener(event);
+			await listener(event, signal);
 		}
 	}
 }

packages/agent/src/types.ts

@@ -263,7 +263,11 @@ export interface AgentState {
 	/** Conversation transcript. Assigning a new array copies the top-level array. */
 	set messages(messages: AgentMessage[]);
 	get messages(): AgentMessage[];
-	/** True while the agent is processing a prompt or continuation. */
+	/**
+	 * True while the agent is processing a prompt or continuation.
+	 *
+	 * This remains true until awaited `agent_end` listeners settle.
+	 */
 	readonly isStreaming: boolean;
 	/** Partial assistant message for the current streamed response, if any. */
 	readonly streamingMessage?: AgentMessage;
@@ -314,7 +318,10 @@ export interface AgentContext {
 
 /**
  * Events emitted by the Agent for UI updates.
- * These events provide fine-grained lifecycle information for messages, turns, and tool executions.
+ *
+ * `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()`
+ * listeners for that event are still part of run settlement. The agent becomes
+ * idle only after those listeners finish.
  */
 export type AgentEvent =
 	// Agent lifecycle

packages/agent/test/agent.test.ts

@@ -36,6 +36,17 @@ function createAssistantMessage(text: string): AssistantMessage {
 	};
 }
 
+function createDeferred(): {
+	promise: Promise<void>;
+	resolve: () => void;
+} {
+	let resolve = () => {};
+	const promise = new Promise<void>((resolvePromise) => {
+		resolve = resolvePromise;
+	});
+	return { promise, resolve };
+}
+
 describe("Agent", () => {
 	it("should create an agent instance with default state", () => {
 		const agent = new Agent();
@@ -89,6 +100,117 @@ describe("Agent", () => {
 		expect(eventCount).toBe(0); // Should not increase
 	});
 
+	it("should await async subscribers before prompt resolves", async () => {
+		const barrier = createDeferred();
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
+				});
+				return stream;
+			},
+		});
+
+		let listenerFinished = false;
+		agent.subscribe(async (event) => {
+			if (event.type === "agent_end") {
+				await barrier.promise;
+				listenerFinished = true;
+			}
+		});
+
+		let promptResolved = false;
+		const promptPromise = agent.prompt("hello").then(() => {
+			promptResolved = true;
+		});
+
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(promptResolved).toBe(false);
+		expect(listenerFinished).toBe(false);
+		expect(agent.state.isStreaming).toBe(true);
+
+		barrier.resolve();
+		await promptPromise;
+
+		expect(listenerFinished).toBe(true);
+		expect(promptResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("waitForIdle should wait for async subscribers", async () => {
+		const barrier = createDeferred();
+		const agent = new Agent({
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
+				});
+				return stream;
+			},
+		});
+
+		agent.subscribe(async (event) => {
+			if (event.type === "message_end" && event.message.role === "assistant") {
+				await barrier.promise;
+			}
+		});
+
+		const promptPromise = agent.prompt("hello");
+		let idleResolved = false;
+		const idlePromise = agent.waitForIdle().then(() => {
+			idleResolved = true;
+		});
+
+		await new Promise((resolve) => setTimeout(resolve, 10));
+		expect(idleResolved).toBe(false);
+		expect(agent.state.isStreaming).toBe(true);
+
+		barrier.resolve();
+		await Promise.all([promptPromise, idlePromise]);
+
+		expect(idleResolved).toBe(true);
+		expect(agent.state.isStreaming).toBe(false);
+	});
+
+	it("should pass the active abort signal to subscribers", async () => {
+		let receivedSignal: AbortSignal | undefined;
+		const agent = new Agent({
+			streamFn: (_model, _context, options) => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "start", partial: createAssistantMessage("") });
+					const checkAbort = () => {
+						if (options?.signal?.aborted) {
+							stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
+						} else {
+							setTimeout(checkAbort, 5);
+						}
+					};
+					checkAbort();
+				});
+				return stream;
+			},
+		});
+
+		agent.subscribe((event, signal) => {
+			if (event.type === "agent_start") {
+				receivedSignal = signal;
+			}
+		});
+
+		const promptPromise = agent.prompt("hello");
+		await new Promise((resolve) => setTimeout(resolve, 10));
+
+		expect(receivedSignal).toBeDefined();
+		expect(receivedSignal?.aborted).toBe(false);
+
+		agent.abort();
+		await promptPromise;
+
+		expect(receivedSignal?.aborted).toBe(true);
+	});
+
 	it("should update state with mutators", () => {
 		const agent = new Agent();