From f5c0977e96b05f973d664772a6d8962dd12577ba Mon Sep 17 00:00:00 2001
From: Coco Sheng <cocosheng@google.com>
Date: Tue, 5 May 2026 15:19:50 -0400
Subject: [PATCH] fix(core): retry on ERR_STREAM_PREMATURE_CLOSE errors
 (#26519)

---
 .../src/core/geminiChat_network_retry.test.ts | 62 +++++++++++++++++++
 packages/core/src/utils/retry.ts              |  1 +
 2 files changed, 63 insertions(+)

diff --git a/packages/core/src/core/geminiChat_network_retry.test.ts b/packages/core/src/core/geminiChat_network_retry.test.ts
index 49013f6461..a2b2ee6e9f 100644
--- a/packages/core/src/core/geminiChat_network_retry.test.ts
+++ b/packages/core/src/core/geminiChat_network_retry.test.ts
@@ -587,4 +587,66 @@ describe('GeminiChat Network Retries', () => {
       }),
     );
   });
+
+  it('should retry on premature stream closure (ERR_STREAM_PREMATURE_CLOSE)', async () => {
+    mockConfig.getRetryFetchErrors = vi.fn().mockReturnValue(true);
+
+    const prematureCloseError = new Error('Premature close');
+    Object.defineProperty(prematureCloseError, 'code', {
+      value: 'ERR_STREAM_PREMATURE_CLOSE',
+    });
+
+    vi.mocked(mockContentGenerator.generateContentStream)
+      .mockResolvedValueOnce(
+        (async function* () {
+          yield {
+            candidates: [{ content: { parts: [{ text: 'Incomplete part' }] } }],
+          } as unknown as GenerateContentResponse;
+          throw prematureCloseError;
+        })(),
+      )
+      .mockResolvedValueOnce(
+        (async function* () {
+          yield {
+            candidates: [
+              {
+                content: { parts: [{ text: 'Complete response after retry' }] },
+                finishReason: 'STOP',
+              },
+            ],
+          } as unknown as GenerateContentResponse;
+        })(),
+      );
+
+    const stream = await chat.sendMessageStream(
+      { model: 'test-model' },
+      'test message',
+      'prompt-id-premature-close',
+      new AbortController().signal,
+      LlmRole.MAIN,
+    );
+
+    const events: StreamEvent[] = [];
+    for await (const event of stream) {
+      events.push(event);
+    }
+
+    const retryEvent = events.find((e) => e.type === StreamEventType.RETRY);
+    expect(retryEvent).toBeDefined();
+
+    const successChunk = events.find(
+      (e) =>
+        e.type === StreamEventType.CHUNK &&
+        e.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+          'Complete response after retry',
+    );
+    expect(successChunk).toBeDefined();
+
+    expect(mockLogNetworkRetryAttempt).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({
+        error_type: 'ERR_STREAM_PREMATURE_CLOSE',
+      }),
+    );
+  });
 });
diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts
index 404b9cf0b2..a45ba0c0b0 100644
--- a/packages/core/src/utils/retry.ts
+++ b/packages/core/src/utils/retry.ts
@@ -58,6 +58,7 @@ const RETRYABLE_NETWORK_CODES = [
   'UND_ERR_HEADERS_TIMEOUT',
   'UND_ERR_BODY_TIMEOUT',
   'UND_ERR_CONNECT_TIMEOUT',
+  'ERR_STREAM_PREMATURE_CLOSE',
 ];
 
 // Node.js builds SSL error codes by prepending ERR_SSL_ to the uppercased