feat(agent): add wait() and stream cancel on break

Author: marslavish

packages/agent/__tests__/run-handle.test.ts

@@ -304,6 +304,23 @@ describe('AgentRunHandle', () => {
       expect(getSignal()?.aborted).toBe(true);
       expect(agent.state.isStreaming).toBe(false);
     });
+
+    it('aborts streamFn when events() iteration breaks early', async () => {
+      const { streamFn, getSignal } = makeAbortableStreamFn();
+      const agent = new Agent({
+        initialState: { model: makeFakeModel() },
+        streamFn,
+      });
+
+      const handle = agent.prompt('go');
+      for await (const _event of handle.events()) {
+        break;
+      }
+      await agent.waitForIdle();
+
+      expect(getSignal()?.aborted).toBe(true);
+      expect(agent.state.isStreaming).toBe(false);
+    });
   });
 
   describe('single-use enforcement', () => {
@@ -328,6 +345,29 @@ describe('AgentRunHandle', () => {
       expect(() => handle.events()).toThrow(/already consumed/);
     });
 
+    it('throws on a second prompt() while a handle from the first is still unconsumed', () => {
+      const provider = createScriptedProvider({
+        responses: [
+          makeFakeAssistantMessage({
+            stopReason: 'stop',
+            content: [{ type: 'text', text: 'x' }],
+          }),
+        ],
+      });
+      const agent = new Agent({
+        initialState: { model: makeFakeModel() },
+        streamFn: provider.stream,
+      });
+
+      const first = agent.prompt('hi');
+      expect(() => agent.prompt('hi again')).toThrow(/unconsumed run handle/);
+
+      // abort frees the agent state; first remains a dangling handle reference
+      agent.abort();
+      expect(() => agent.prompt('after abort')).not.toThrow();
+      void first;
+    });
+
     it('throws when toResponse() is called twice', () => {
       const provider = createScriptedProvider({
         responses: [
@@ -349,8 +389,8 @@ describe('AgentRunHandle', () => {
     });
   });
 
-  describe('PromiseLike auto-sink', () => {
-    it('await on the handle drives the run to completion without an explicit consumer', async () => {
+  describe('wait()', () => {
+    it('drives the run to completion without an explicit event consumer', async () => {
       const provider = createScriptedProvider({
         responses: [
           makeFakeAssistantMessage({
@@ -373,15 +413,15 @@ describe('AgentRunHandle', () => {
       expect(agent.state.isStreaming).toBe(false);
     });
 
-    it('rejects the awaited handle when the binder rejects (e.g. streamFn throws)', async () => {
+    it('rejects when the binder rejects (e.g. streamFn throws)', async () => {
       const agent = new Agent({
         initialState: { model: makeFakeModel() },
         streamFn: () => {
           throw new Error('binder failure');
         },
       });
 
-      await expect(agent.prompt('hi')).rejects.toThrow(/binder failure/);
+      await expect(agent.prompt('hi').wait()).rejects.toThrow(/binder failure/);
       expect(agent.state.isStreaming).toBe(false);
     });
 

packages/agent/src/run-handle.ts

@@ -11,6 +11,13 @@ export interface AgentRunHandle extends PromiseLike<void> {
   events(): AsyncIterable<AgentEvent>;
   toReadableStream(): ReadableStream<AgentEvent>;
   toResponse(init?: ResponseInit): Response;
+  /**
+   * Run to completion without observing events. Equivalent to `await handle`
+   * (the handle is `PromiseLike<void>`), but explicit. Use this if you want
+   * to avoid accidental thenable assimilation in code paths where the handle
+   * might be passed through generic wrappers.
+   */
+  wait(): Promise<void>;
 }
 
 const DEFAULT_HIGH_WATER_MARK = 8;
@@ -59,14 +66,18 @@ export class DefaultAgentRunHandle implements AgentRunHandle {
     return new Response(sse, responseInit);
   }
 
+  wait(): Promise<void> {
+    if (!this.startedAs) {
+      this.startSink();
+    }
+    return this.completion!;
+  }
+
   then<TResult1 = void, TResult2 = never>(
     onfulfilled?: ((value: void) => TResult1 | PromiseLike<TResult1>) | null,
     onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null
   ): Promise<TResult1 | TResult2> {
-    if (!this.startedAs) {
-      this.startSink();
-    }
-    return this.completion!.then(onfulfilled, onrejected);
+    return this.wait().then(onfulfilled, onrejected);
   }
 
   private ensureNotStarted(via: NonNullable<DefaultAgentRunHandle['startedAs']>): void {
@@ -175,15 +186,24 @@ async function* readableStreamToAsyncIterable<T>(
   stream: ReadableStream<T>
 ): AsyncIterableIterator<T> {
   const reader = stream.getReader();
+  let drained = false;
   try {
     while (true) {
       const { done, value } = await reader.read();
       if (done) {
+        drained = true;
         return;
       }
       yield value;
     }
   } finally {
+    if (!drained) {
+      try {
+        await reader.cancel();
+      } catch {
+        // cancel can reject if the stream already errored — safe to ignore
+      }
+    }
     reader.releaseLock();
   }
 }