From 384f6259183f7bf82b822b28f05a18e1e1c3d6d4 Mon Sep 17 00:00:00 2001
From: Sweets Sweetman <sweetmantech@gmail.com>
Date: Mon, 9 Feb 2026 09:24:48 -0500
Subject: [PATCH] fix: defer sandbox cleanup until stream completes

The `finally` block was calling `sandbox.stop()` immediately when
`createAgentUIStreamResponse` returned, before the client consumed
the stream. This killed the sandbox before any agent tool calls
could execute, causing all bash commands to fail with
"[Tool Error] Tool execution error".

Now uses a TransformStream wrapper with `pipeTo().finally()` to
ensure sandbox cleanup happens after streaming completes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 app/api/agent/route.ts | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/app/api/agent/route.ts b/app/api/agent/route.ts
index 4d55f387..022fa314 100644
--- a/app/api/agent/route.ts
+++ b/app/api/agent/route.ts
@@ -104,12 +104,30 @@ export async function POST(req: Request) {
       stopWhen: stepCountIs(20),
     });
 
-    return createAgentUIStreamResponse({
+    const response = await createAgentUIStreamResponse({
       agent,
       uiMessages: messages,
     });
-  } finally {
-    // Don't await — let it clean up in background so response isn't delayed
+
+    // Clean up sandbox after the stream finishes (not before).
+    // The original `finally` block killed the sandbox immediately when
+    // createAgentUIStreamResponse returned, before any tool calls ran.
+    const body = response.body;
+    if (body) {
+      const transform = new TransformStream();
+      body.pipeTo(transform.writable).finally(() => {
+        sandbox.stop().catch(() => {});
+      });
+      return new Response(transform.readable, {
+        headers: response.headers,
+        status: response.status,
+      });
+    }
+
+    sandbox.stop().catch(() => {});
+    return response;
+  } catch (error) {
     sandbox.stop().catch(() => {});
+    throw error;
   }
 }