From c3b5bcb84c328d6cd4de45fabafde9150d0e7f62 Mon Sep 17 00:00:00 2001
From: Adam Weidman <adamfweidman@gmail.com>
Date: Fri, 20 Mar 2026 11:19:34 -0400
Subject: [PATCH] fix(cli): consume the stream returned by send

---
 packages/cli/src/nonInteractiveCli.test.ts | 28 ++++++++++++++++++++++
 packages/cli/src/nonInteractiveCli.ts      |  4 ++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts
index f506ec9d60..3fc827d497 100644
--- a/packages/cli/src/nonInteractiveCli.test.ts
+++ b/packages/cli/src/nonInteractiveCli.test.ts
@@ -277,6 +277,34 @@ describe('runNonInteractive', () => {
     // so we no longer expect shutdownTelemetry to be called directly here
   });
 
+  it('should stream the specific stream started by send', async () => {
+    const { LegacyAgentSession } = await import(
+      '../../core/src/agent/legacy-agent-session.js'
+    );
+    const streamSpy = vi.spyOn(LegacyAgentSession.prototype, 'stream');
+    const events: ServerGeminiStreamEvent[] = [
+      { type: GeminiEventType.Content, value: 'Hello again' },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
+    ];
+    mockGeminiClient.sendMessageStream.mockReturnValue(
+      createStreamFromEvents(events),
+    );
+
+    await runNonInteractive({
+      config: mockConfig,
+      settings: mockSettings,
+      input: 'Test input',
+      prompt_id: 'prompt-id-stream',
+    });
+
+    expect(streamSpy).toHaveBeenCalledWith({
+      streamId: expect.any(String),
+    });
+  });
+
   it('should register activity logger when GEMINI_CLI_ACTIVITY_LOG_TARGET is set', async () => {
     vi.stubEnv('GEMINI_CLI_ACTIVITY_LOG_TARGET', '/tmp/test.jsonl');
     const events: ServerGeminiStreamEvent[] = [
diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts
index 5bdf03d4b3..11e3fd2cc2 100644
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/packages/cli/src/nonInteractiveCli.ts
@@ -294,7 +294,7 @@ export async function runNonInteractive({
       });
 
       // Start the agentic loop (runs in background)
-      await session.send({
+      const { streamId } = await session.send({
         message: geminiPartsToContentParts(query),
       });
 
@@ -351,7 +351,7 @@ export async function runNonInteractive({
       // Consume AgentEvents for output formatting
       let responseText = '';
       let streamEnded = false;
-      for await (const event of session.stream()) {
+      for await (const event of session.stream({ streamId })) {
         if (streamEnded) break;
         switch (event.type) {
           case 'message': {