From 68ab03647c9961da4ecf4209308be5d70d1210ec Mon Sep 17 00:00:00 2001
From: bracesproul <braceasproul@gmail.com>
Date: Thu, 22 May 2025 11:50:06 -0700
Subject: [PATCH] feat: Pause sandboxes before ending the graph

---
 src/index.ts         | 19 +++++++++++++++++--
 src/utils/sandbox.ts | 17 +++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 src/utils/sandbox.ts

diff --git a/src/index.ts b/src/index.ts
index 859676e8..9bb3acfa 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,5 +1,10 @@
 import { END, START, StateGraph } from "@langchain/langgraph";
-import { GraphAnnotation, GraphConfiguration, GraphState } from "./types.js";
+import {
+  GraphAnnotation,
+  GraphConfig,
+  GraphConfiguration,
+  GraphState,
+} from "./types.js";
 import {
   generatePlan,
   initialize,
@@ -9,6 +14,7 @@ import {
   interruptPlan,
 } from "./nodes/index.js";
 import { isAIMessage } from "@langchain/core/messages";
+import { pauseSandbox } from "./utils/sandbox.js";
 
 /**
  * @param {GraphState} state - The current graph state.
@@ -33,7 +39,10 @@ function routeAfterPlan(state: GraphState): "interrupt-plan" | typeof END {
  * @param {GraphState} state - The current graph state.
  * @returns {typeof END | "take-action"} The next node to execute, or END if the process should stop.
  */
-function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
+async function takeActionOrEnd(
+  state: GraphState,
+  config: GraphConfig,
+): Promise<typeof END | "take-action"> {
   const { messages } = state;
   const lastMessage = messages[messages.length - 1];
 
@@ -41,6 +50,12 @@ function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
   if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) {
     return "take-action";
   }
+
+  // First, pause the sandbox before ending the graph.
+  if (config.configurable?.sandbox_session_id) {
+    await pauseSandbox(config.configurable.sandbox_session_id);
+  }
+
   return END;
 }
 
diff --git a/src/utils/sandbox.ts b/src/utils/sandbox.ts
new file mode 100644
index 00000000..85629bdd
--- /dev/null
+++ b/src/utils/sandbox.ts
@@ -0,0 +1,17 @@
+import { Sandbox } from "@e2b/code-interpreter";
+
+/**
+ * Pauses the sandbox. Either pass an existing sandbox client, or a sandbox session ID.
+ * If no sandbox client is provided, the sandbox will be connected to.
+ 
+ * @param sandboxSessionId The ID of the sandbox to pause.
+ * @param sandbox The sandbox client to pause. If not provided, the sandbox will be connected to.
+ * @returns The sandbox session ID.
+ */
+export async function pauseSandbox(
+  sandboxSessionId: string,
+  sandbox?: Sandbox,
+): Promise<string> {
+  const sandboxClient = sandbox ?? (await Sandbox.connect(sandboxSessionId));
+  return await sandboxClient.pause();
+}