fix: address code review issues

- Add buffer cleanup after all consumers complete (memory leak fix)
- Use lazy initialization for broadcaster (race condition fix)
- Add try-catch in onPreliminaryResult callback (error handling)
- Add tests for completion between consumer iterations

Author: mattapperson

src/lib/model-result.ts

@@ -125,6 +125,21 @@ export class ModelResult<TTools extends readonly Tool[]> {
     this.options = options;
   }
 
+  /**
+   * Get or create the tool event broadcaster (lazy initialization).
+   * Ensures only one broadcaster exists for the lifetime of this ModelResult.
+   */
+  private ensureBroadcaster(): ToolEventBroadcaster<{
+    type: 'preliminary_result';
+    toolCallId: string;
+    result: InferToolEventsUnion<TTools>;
+  }> {
+    if (!this.toolEventBroadcaster) {
+      this.toolEventBroadcaster = new ToolEventBroadcaster();
+    }
+    return this.toolEventBroadcaster;
+  }
+
   /**
    * Type guard to check if a value is a non-streaming response
    */
@@ -332,11 +347,15 @@ export class ModelResult<TTools extends readonly Tool[]> {
           // Create callback for real-time preliminary results
           const onPreliminaryResult = this.toolEventBroadcaster
             ? (callId: string, resultValue: unknown) => {
-                this.toolEventBroadcaster!.push({
-                  type: 'preliminary_result' as const,
-                  toolCallId: callId,
-                  result: resultValue as InferToolEventsUnion<TTools>,
-                });
+                try {
+                  this.toolEventBroadcaster!.push({
+                    type: 'preliminary_result' as const,
+                    toolCallId: callId,
+                    result: resultValue as InferToolEventsUnion<TTools>,
+                  });
+                } catch {
+                  // Don't crash tool execution if broadcasting fails
+                }
               }
             : undefined;
 
@@ -499,13 +518,13 @@ export class ModelResult<TTools extends readonly Tool[]> {
         throw new Error('Stream not initialized');
       }
 
-      // Create broadcaster for real-time tool events
-      this.toolEventBroadcaster = new ToolEventBroadcaster();
-      const toolEventConsumer = this.toolEventBroadcaster.createConsumer();
+      // Get or create broadcaster for real-time tool events (lazy init prevents race conditions)
+      const broadcaster = this.ensureBroadcaster();
+      const toolEventConsumer = broadcaster.createConsumer();
 
       // Start tool execution in background (completes broadcaster when done)
       const executionPromise = this.executeToolsIfNeeded().finally(() => {
-        this.toolEventBroadcaster?.complete();
+        broadcaster.complete();
       });
 
       const consumer = this.reusableStream.createConsumer();
@@ -615,13 +634,13 @@ export class ModelResult<TTools extends readonly Tool[]> {
         throw new Error('Stream not initialized');
       }
 
-      // Create broadcaster for real-time tool events
-      this.toolEventBroadcaster = new ToolEventBroadcaster();
-      const toolEventConsumer = this.toolEventBroadcaster.createConsumer();
+      // Get or create broadcaster for real-time tool events (lazy init prevents race conditions)
+      const broadcaster = this.ensureBroadcaster();
+      const toolEventConsumer = broadcaster.createConsumer();
 
       // Start tool execution in background (completes broadcaster when done)
       const executionPromise = this.executeToolsIfNeeded().finally(() => {
-        this.toolEventBroadcaster?.complete();
+        broadcaster.complete();
       });
 
       // Yield tool deltas from API stream

src/lib/tool-event-broadcaster.ts

@@ -28,11 +28,25 @@ export class ToolEventBroadcaster<T> {
   /**
    * Mark the broadcaster as complete - no more events will be pushed.
    * Optionally pass an error to signal failure to all consumers.
+   * Cleans up buffer and consumers after completion.
    */
   complete(error?: Error): void {
     this.isComplete = true;
     this.completionError = error ?? null;
     this.notifyWaitingConsumers();
+    // Schedule cleanup after consumers have processed completion
+    queueMicrotask(() => this.cleanup());
+  }
+
+  /**
+   * Clean up resources after all consumers have finished.
+   * Called automatically after complete(), but can be called manually.
+   */
+  private cleanup(): void {
+    // Only cleanup if complete and all consumers are done
+    if (this.isComplete && this.consumers.size === 0) {
+      this.buffer = [];
+    }
   }
 
   /**
@@ -73,6 +87,7 @@ export class ToolEventBroadcaster<T> {
         // If complete and caught up, we're done
         if (self.isComplete) {
           self.consumers.delete(consumerId);
+          self.cleanup();
           if (self.completionError) {
             throw self.completionError;
           }
@@ -105,6 +120,7 @@ export class ToolEventBroadcaster<T> {
         if (consumer) {
           consumer.cancelled = true;
           self.consumers.delete(consumerId);
+          self.cleanup();
         }
         return { done: true, value: undefined };
       },
@@ -114,6 +130,7 @@ export class ToolEventBroadcaster<T> {
         if (consumer) {
           consumer.cancelled = true;
           self.consumers.delete(consumerId);
+          self.cleanup();
         }
         throw e;
       },

tests/unit/tool-event-broadcaster.test.ts

@@ -221,6 +221,51 @@ describe('ToolEventBroadcaster', () => {
     });
   });
 
+  describe('completion between iterations', () => {
+    it('should handle completion between consumer iterations', async () => {
+      const broadcaster = new ToolEventBroadcaster<number>();
+      const consumer = broadcaster.createConsumer();
+
+      broadcaster.push(1);
+      const first = await consumer.next();
+      expect(first.done).toBe(false);
+      expect(first.value).toBe(1);
+
+      // Complete while consumer is between iterations
+      broadcaster.complete();
+
+      const second = await consumer.next();
+      expect(second.done).toBe(true);
+    });
+
+    it('should handle completion with remaining buffered events', async () => {
+      const broadcaster = new ToolEventBroadcaster<number>();
+      const consumer = broadcaster.createConsumer();
+
+      broadcaster.push(1);
+      broadcaster.push(2);
+      broadcaster.push(3);
+
+      // Read first event
+      const first = await consumer.next();
+      expect(first.value).toBe(1);
+
+      // Complete with events still in buffer
+      broadcaster.complete();
+
+      // Should still get remaining buffered events
+      const second = await consumer.next();
+      expect(second.value).toBe(2);
+
+      const third = await consumer.next();
+      expect(third.value).toBe(3);
+
+      // Now should be done
+      const fourth = await consumer.next();
+      expect(fourth.done).toBe(true);
+    });
+  });
+
   describe('typed events', () => {
     it('should work with typed tool events', async () => {
       type ToolEvent =