diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 53451fb1..b73ea918 100644 --- a/packages/agent/CHANGELOG.md +++ b/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 diff --git a/packages/agent/README.md b/packages/agent/README.md index 1b19d3f2..befdd9a1 100644 --- a/packages/agent/README.md +++ b/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(); ``` diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index f4faa496..b7d69546 100644 --- a/packages/agent/src/agent.ts +++ b/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>(); 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 { 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 { 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 { 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 { 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); } } } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 8d4e2561..b1214c9a 100644 --- a/packages/agent/src/types.ts +++ b/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 diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index c7ef1d87..94be88cb 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -36,6 +36,17 @@ function createAssistantMessage(text: string): AssistantMessage { }; } +function createDeferred(): { + promise: Promise; + resolve: () => void; +} { + let resolve = () => {}; + const promise = new Promise((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();