From 424124c286b156c92dcb6859de1e6183fe62765e Mon Sep 17 00:00:00 2001
From: Brace Sproul <braceasproul@gmail.com>
Date: Thu, 22 May 2025 11:40:54 -0700
Subject: [PATCH] fix: Initialize sandbox after planning (#6)

* fix: Initialize sandbox after planning

* fix routing in interrupt node

* cr
---
 src/index.ts                |  9 ++++-----
 src/nodes/initialize.ts     | 29 ++++++++++++++---------------
 src/nodes/interrupt-plan.ts |  8 ++++----
 3 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 031fef13..859676e8 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -45,22 +45,21 @@ function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
 }
 
 const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
-  .addNode("initialize", initialize)
   .addNode("generate-plan", generatePlan)
   .addNode("rewrite-plan", rewritePlan)
   .addNode("interrupt-plan", interruptPlan, {
     // TODO: Hookup `Command` in interruptPlan node so this actually works.
-    ends: [END, "rewrite-plan", "generate-action"],
+    ends: [END, "rewrite-plan", "initialize"],
   })
+  .addNode("initialize", initialize)
   .addNode("generate-action", generateAction)
   .addNode("take-action", takeAction)
-  .addEdge(START, "initialize")
-  .addEdge("initialize", "generate-plan")
+  .addEdge(START, "generate-plan")
   // TODO: Update routing to work w/ new interrupt node.
   .addConditionalEdges("generate-plan", routeAfterPlan, ["interrupt-plan", END])
   // Always interrupt after rewriting the plan.
   .addEdge("rewrite-plan", "interrupt-plan")
-  .addEdge("generate-plan", "generate-action")
+  .addEdge("initialize", "generate-action")
   .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END])
   .addEdge("take-action", "generate-action");
 
diff --git a/src/nodes/initialize.ts b/src/nodes/initialize.ts
index 6bcace87..010c74b8 100644
--- a/src/nodes/initialize.ts
+++ b/src/nodes/initialize.ts
@@ -46,8 +46,14 @@ export async function initialize(
   const { sandbox_session_id, target_repository, sandbox_language } =
     config.configurable;
   if (sandbox_session_id) {
-    // Session already initialized. Return early.
-    return {};
+    try {
+      // Resume the sandbox if the session ID is in the config.
+      await Sandbox.resume(sandbox_session_id, TIMEOUT_EXTENSION_OPT);
+      return {};
+    } catch (e) {
+      // Error thrown, log it and continue. Will create a new sandbox session since the resumption failed.
+      console.error("Failed to get sandbox session.", e);
+    }
   }
 
   if (!sandbox_language || !target_repository) {
@@ -57,23 +63,16 @@ export async function initialize(
   }
 
   if (sandbox_language === "js") {
-    const sandbox = await Sandbox.create(JS_SANDBOX_TEMPLATE_ID);
+    const sandbox = await Sandbox.create(
+      JS_SANDBOX_TEMPLATE_ID,
+      TIMEOUT_EXTENSION_OPT,
+    );
     config.configurable.sandbox_session_id = sandbox.sandboxId;
 
     const res = await cloneRepo(sandbox, target_repository);
     if (res.error) {
-      let errorMessage = "Failed to clone repository.";
-      if (res.error) {
-        // Basic redaction attempt for the token in the URL if present in the error
-        errorMessage +=
-          ": " + res.error.replace(/:\/\/[^@]+@/, "://<REDACTED_TOKEN>@");
-      } else if (typeof res.error === "string") {
-        errorMessage +=
-          ": " + res.error.replace(/:\/\/[^@]+@/, "://<REDACTED_TOKEN>@");
-      } else {
-        errorMessage += " An unknown error occurred.";
-      }
-      throw new Error(errorMessage);
+      // TODO: This should probably be an interrupt.
+      throw new Error(`Failed to clone repository.\n${res.error}`);
     }
     return {};
   }
diff --git a/src/nodes/interrupt-plan.ts b/src/nodes/interrupt-plan.ts
index 6094da73..00aa8c94 100644
--- a/src/nodes/interrupt-plan.ts
+++ b/src/nodes/interrupt-plan.ts
@@ -31,9 +31,9 @@ export function interruptPlan(state: GraphState): Command {
   })[0];
 
   if (interruptRes.type === "accept") {
-    // Plan was accepted, route to the generate action node.
+    // Plan was accepted, route to the initialize node.
     return new Command({
-      goto: "generate-action",
+      goto: "initialize",
       update: {
         plan: proposedPlan.map((p) => ({
           id: uuidv4(),
@@ -45,12 +45,12 @@ export function interruptPlan(state: GraphState): Command {
   }
 
   if (interruptRes.type === "edit") {
-    // Plan was edited, route to the generate action node.
+    // Plan was edited, route to the initialize node.
     const editedPlan = (interruptRes.args as ActionRequest).args.plan
       .split(":::")
       .map((step: string) => step.trim());
     return new Command({
-      goto: "generate-action",
+      goto: "initialize",
       update: {
         plan: editedPlan.map((p: string) => ({
           id: uuidv4(),