feat: Pause sandboxes before ending the graph

This commit is contained in:
bracesproul 2025-05-22 11:50:06 -07:00
parent 424124c286
commit 68ab03647c
2 changed files with 34 additions and 2 deletions

View file

@ -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;
}

17
src/utils/sandbox.ts Normal file
View file

@ -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();
}