From 63a6211fe029c41d7fcb9b38cbe5dc6a41f4e677 Mon Sep 17 00:00:00 2001
From: Spencer <spencertang@google.com>
Date: Fri, 10 Apr 2026 15:47:53 +0000
Subject: [PATCH] fix(core): ensure binary shell output files are still written
 to disk for 20MB files, and wait for stream close

---
 .../run_shell_command_file_stream.test.ts     | 46 +++++++++++++++----
 packages/core/src/tools/shell.ts              | 11 +++--
 2 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/integration-tests/run_shell_command_file_stream.test.ts b/integration-tests/run_shell_command_file_stream.test.ts
index 4e2c15b8f2..b862f9bcdd 100644
--- a/integration-tests/run_shell_command_file_stream.test.ts
+++ b/integration-tests/run_shell_command_file_stream.test.ts
@@ -70,12 +70,27 @@ describe('run_shell_command streaming to file regression', () => {
       files.sort((a, b) => fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs);
 
       for (const p of files) {
-        const stat = fs.statSync(p);
-        if (stat.size >= 20000000) {
-          savedFilePath = p;
-          break;
+        try {
+          const stat = fs.statSync(p);
+          if (stat.size >= 20000000) {
+            savedFilePath = p;
+            break;
+          }
+        } catch {
+          // ignore
         }
       }
+
+      if (!savedFilePath) {
+        const fileStats = files.map((p) => {
+          try {
+            return { p, size: fs.statSync(p).size };
+          } catch {
+            return { p, size: 'error' };
+          }
+        });
+        console.error('Available files:', JSON.stringify(fileStats, null, 2));
+      }
     }
 
     expect(
@@ -144,12 +159,27 @@ describe('run_shell_command streaming to file regression', () => {
       files.sort((a, b) => fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs);
 
       for (const p of files) {
-        const stat = fs.statSync(p);
-        if (stat.size >= 20000000) {
-          savedFilePath = p;
-          break;
+        try {
+          const stat = fs.statSync(p);
+          if (stat.size >= 20000000) {
+            savedFilePath = p;
+            break;
+          }
+        } catch {
+          // ignore
         }
       }
+
+      if (!savedFilePath) {
+        const fileStats = files.map((p) => {
+          try {
+            return { p, size: fs.statSync(p).size };
+          } catch {
+            return { p, size: 'error' };
+          }
+        });
+        console.error('Available files:', JSON.stringify(fileStats, null, 2));
+      }
     }
 
     expect(
diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts
index 8dabc9296b..575f3ac4ab 100644
--- a/packages/core/src/tools/shell.ts
+++ b/packages/core/src/tools/shell.ts
@@ -532,10 +532,8 @@ export class ShellToolInvocation extends BaseToolInvocation<
                 // We rely on 'file_data' for the clean output stream.
                 break;
               case 'file_data':
-                if (!isBinaryStream) {
-                  totalBytesWritten += Buffer.byteLength(event.chunk);
-                  outputStream.write(event.chunk);
-                }
+                totalBytesWritten += Buffer.byteLength(event.chunk);
+                outputStream.write(event.chunk);
                 break;
               case 'data':
                 if (isBinaryStream) break;
@@ -676,6 +674,11 @@ export class ShellToolInvocation extends BaseToolInvocation<
         outputStream.end(resolve);
       });
 
+      // Ensure the stream is fully closed before we proceed
+      if (!outputStream.closed) {
+        await new Promise<void>((resolve) => outputStream.on('close', resolve));
+      }
+
       const backgroundPIDs: number[] = [];
       if (os.platform() !== 'win32') {
         let tempFileExists = false;