fix: Better handling of updating sandbox timeout (#41)

apps/open-swe/src/nodes/generate-message.ts

@@ -54,6 +54,7 @@ You MUST adhere to the following criteria when executing the task:
 - When using the \`shell\` tool, always take advantage of the \`workdir\` parameter to run commands inside the repo directory. You should not try to generate a command with \`cd <some path>\` as passing that path to \`workdir\` is much more efficient.
 - Always use the correct package manager to install dependencies. If the package manager is not already installed in the sandbox, use the \`shell\` tool to install it.
   - If the package manager fails to install, or you have issues installing dependencies, do not try to use a different package manager. Instead, skip installing dependencies.
+  - If installing dependencies fails, it can be useful to try again passing in a much longer timeout than the default.
 - If completing the user's task requires writing or modifying files:
     - Your code and final answer should follow these *CODING GUIDELINES*:
         - Avoid writing to files which you have not already read.

apps/open-swe/src/tools/shell.ts

@@ -30,6 +30,7 @@ const shellToolSchema = z.object({
 
 export const shellTool = tool(
   async (input): Promise<{ result: string; status: "success" | "error" }> => {
+    let sandbox: Sandbox | undefined;
     try {
       const state = getCurrentTaskInput<GraphState>();
       const { sandboxSessionId } = state;
@@ -42,14 +43,12 @@ export const shellTool = tool(
         );
       }
 
-      const sandbox = await Sandbox.connect(sandboxSessionId);
+      sandbox = await Sandbox.connect(sandboxSessionId);
       const { command, workdir, timeout } = input;
       const result = await sandbox.commands.run(command.join(" "), {
         timeoutMs: timeout ?? DEFAULT_COMMAND_TIMEOUT,
         cwd: workdir,
       });
-      // Add an extra 5 min timeout to the sandbox.
-      await sandbox.setTimeout(TIMEOUT_MS);
 
       if (result.error) {
         logger.error("Failed to run command", {
@@ -92,6 +91,17 @@ export const shellTool = tool(
         "FAILED TO RUN COMMAND: " +
           (e instanceof Error ? e.message : "Unknown error"),
       );
+    } finally {
+      try {
+        if (sandbox) {
+          // Add an extra 5 min timeout to the sandbox.
+          await sandbox.setTimeout(TIMEOUT_MS);
+        }
+      } catch (_) {
+        logger.warn(
+          "Failed to set timeout for sandbox inside 'finally' block for shell tool.",
+        );
+      }
     }
   },
   {

apps/open-swe/src/utils/read-write.ts

@@ -47,8 +47,6 @@ async function readFileFunc(
     const readOutput = await sandbox.commands.run(`cat "${filePath}"`, {
       cwd: args?.workDir,
     });
-    // Add an extra 5 min timeout to the sandbox.
-    await sandbox.setTimeout(TIMEOUT_MS);
 
     if (readOutput.exitCode !== 0) {
       logger.error(`Error reading file '${filePath}' from sandbox via cat:`, {
@@ -102,6 +100,17 @@ async function readFileFunc(
       success: false,
       output: outputMessage,
     };
+  } finally {
+    try {
+      if (sandbox) {
+        // Add an extra 5 min timeout to the sandbox.
+        await sandbox.setTimeout(TIMEOUT_MS);
+      }
+    } catch (_) {
+      logger.warn(
+        "Failed to set timeout for sandbox inside 'finally' block for read file.",
+      );
+    }
   }
 }
 
@@ -128,8 +137,6 @@ ${delimiter}`;
     const writeOutput = await sandbox.commands.run(writeCommand, {
       cwd: args?.workDir,
     });
-    // Add an extra 5 min timeout to the sandbox.
-    await sandbox.setTimeout(TIMEOUT_MS);
 
     if (writeOutput.exitCode !== 0) {
       logger.error(`Error writing file '${filePath}' to sandbox via cat:`, {
@@ -171,6 +178,17 @@ ${delimiter}`;
       success: false,
       output: outputMessage,
     };
+  } finally {
+    try {
+      if (sandbox) {
+        // Add an extra 5 min timeout to the sandbox.
+        await sandbox.setTimeout(TIMEOUT_MS);
+      }
+    } catch (_) {
+      logger.warn(
+        "Failed to set timeout for sandbox inside 'finally' block for write file.",
+      );
+    }
   }
 }