feat: Implement a planning subgraph (#12)

langgraph.json

@@ -1,7 +1,7 @@
 {
   "node_version": "20",
   "graphs": {
-    "agent": "./src/index.ts:graph"
+    "open-codex": "./src/index.ts:graph"
   },
   "env": ".env",
   "dependencies": ["."]

src/constants.ts

@@ -1,4 +1,5 @@
+export const TIMEOUT_MS = 300_000; // 5 minutes
 // Always extend timeouts by 5 min after each command execution.
 export const TIMEOUT_EXTENSION_OPT = {
-  timeoutMs: 300_000, // 5 minutes
+  timeoutMs: TIMEOUT_MS,
 };

src/index.ts

@@ -1,12 +1,6 @@
 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,
   generateAction,
   takeAction,
@@ -15,7 +9,7 @@ import {
   progressPlanStep,
 } from "./nodes/index.js";
 import { isAIMessage } from "@langchain/core/messages";
-import { pauseSandbox } from "./utils/sandbox.js";
+import { plannerGraph } from "./subgraphs/index.js";
 
 /**
  * @param {GraphState} state - The current graph state.
@@ -42,7 +36,6 @@ function routeAfterPlan(state: GraphState): "interrupt-plan" | typeof END {
  */
 async function takeActionOrEnd(
   state: GraphState,
-  config: GraphConfig,
 ): Promise<typeof END | "take-action"> {
   const { messages } = state;
   const lastMessage = messages[messages.length - 1];
@@ -52,35 +45,33 @@ async function takeActionOrEnd(
     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;
 }
 
 const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
-  .addNode("generate-plan", generatePlan)
+  .addNode("initialize", initialize)
+  .addNode("generate-plan-subgraph", plannerGraph)
   .addNode("rewrite-plan", rewritePlan)
   .addNode("interrupt-plan", interruptPlan, {
     // TODO: Hookup `Command` in interruptPlan node so this actually works.
-    ends: [END, "rewrite-plan", "initialize"],
+    ends: [END, "rewrite-plan", "generate-action"],
   })
-  .addNode("initialize", initialize)
   .addNode("generate-action", generateAction)
   .addNode("take-action", takeAction)
   .addNode("progress-plan-step", progressPlanStep)
-  .addEdge(START, "generate-plan")
+  .addEdge(START, "initialize")
+  .addEdge("initialize", "generate-plan-subgraph")
   // TODO: Update routing to work w/ new interrupt node.
-  .addConditionalEdges("generate-plan", routeAfterPlan, ["interrupt-plan", END])
+  .addConditionalEdges("generate-plan-subgraph", routeAfterPlan, [
+    "interrupt-plan",
+    END,
+  ])
   // Always interrupt after rewriting the plan.
   .addEdge("rewrite-plan", "interrupt-plan")
-  .addEdge("initialize", "generate-action")
   .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END])
   .addEdge("take-action", "progress-plan-step")
   .addEdge("progress-plan-step", "generate-action");
 
 // Zod types are messed up
 export const graph = workflow.compile() as any;
-graph.name = "LangGraph ReAct MCP";
+graph.name = "Open Codex";

src/nodes/generate-message.ts

@@ -2,6 +2,7 @@ import { GraphState, GraphConfig, GraphUpdate, PlanItem } from "../types.js";
 import { loadModel, Task } from "../utils/load-model.js";
 import { shellTool, applyPatchTool } from "../tools/index.js";
 import { formatPlanPrompt } from "../utils/plan-prompt.js";
+import { pauseSandbox } from "../utils/sandbox.js";
 
 const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful.
 
@@ -48,7 +49,9 @@ You MUST adhere to the following criteria when executing the task:
 - When your task involves writing or modifying files:
     - Do NOT tell the user to "save the file" or "copy the code into a file" if you already created or modified the file using \`apply_patch\`. Instead, reference the file as already saved.
     - Do NOT show the full contents of large files you have already written, unless the user explicitly asks for them.
-- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore.`;
+- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore.
+  - Always use glob patterns when searching with \`rg\` for specific file types. For example, to search for all TSX files, use \`rg -i star -g **/*.tsx project-directory/\`. This is because \`rg\` does not have built in file types for every language.
+`;
 
 const formatPrompt = (plan: PlanItem[]): string => {
   return systemPrompt.replace("{PLAN_PROMPT}", formatPlanPrompt(plan));
@@ -70,7 +73,16 @@ export async function generateAction(
     ...state.messages,
   ]);
 
+  const hasToolCalls = !!response.tool_calls?.length;
+  // No tool calls means the graph is going to end. Pause the sandbox.
+  let newSandboxSessionId: string | undefined;
+  if (!hasToolCalls && state.sandboxSessionId) {
+    console.log("No tool calls found. Pausing sandbox...");
+    newSandboxSessionId = await pauseSandbox(state.sandboxSessionId);
+  }
+
   return {
     messages: [response],
+    ...(newSandboxSessionId && { sandboxSessionId: newSandboxSessionId }),
   };
 }

src/nodes/generate-plan.ts

@@ -1,48 +0,0 @@
-import { sessionPlanTool } from "../tools/index.js";
-import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
-import { loadModel, Task } from "../utils/load-model.js";
-
-const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful.
-
-In this step, you are expected to generate a high-level plan to address the user's request. The plan should be a list of actions to take, in order, to address the user's request. You should not include any code in the plan, only a list of actions to take.
-
-You MUST adhere to the following criteria when generating the plan:
-- You do not have access to the codebase yet, so you cannot inspect it or make assumptions about it.
-- Your plan should be high-level in nature, but should still be specific enough to be actionable.
-- If you can not generate a plan due to a lack of context, you are permitted to ask the user followup questions.
-  - If asking followup questions, ensure every question is asked in a single message to avoid back and forth.
-  - Your questions should be concise and to the point. Remember that you are not including code or technical details in your plan, so your questions should be focused on high-level issues.
-- When you are ready to generate the plan, ensure you call the 'session_plan' tool.
-`;
-
-export async function generatePlan(
-  state: GraphState,
-  config: GraphConfig,
-): Promise<GraphUpdate> {
-  const model = await loadModel(config, Task.PLANNER);
-  const modelWithTools = model.bindTools([sessionPlanTool], {
-    tool_choice: "auto",
-  });
-
-  const response = await modelWithTools.invoke([
-    {
-      role: "system",
-      content: systemPrompt,
-    },
-    ...state.messages,
-  ]);
-
-  if (response.tool_calls?.length) {
-    return {
-      proposedPlan: response.tool_calls[0].args.plan,
-      plan: [],
-    };
-  }
-
-  // No tool calls generated, instead we should just return the messages.
-  return {
-    messages: [response],
-    proposedPlan: [],
-    plan: [],
-  };
-}

src/nodes/index.ts

@@ -1,4 +1,3 @@
-export * from "./generate-plan.js";
 export * from "./initialize.js";
 export * from "./generate-message.js";
 export * from "./take-action.js";

src/nodes/initialize.ts

@@ -24,6 +24,9 @@ async function cloneRepo(sandbox: Sandbox, targetRepository: TargetRepository) {
     gitCloneCommand.push(repoUrlWithToken);
   }
 
+  console.log("Cloning repository...", {
+    command: gitCloneCommand.join(" "),
+  });
   return await sandbox.commands.run(
     gitCloneCommand.join(" "),
     TIMEOUT_EXTENSION_OPT,
@@ -37,25 +40,35 @@ async function cloneRepo(sandbox: Sandbox, targetRepository: TargetRepository) {
  * branch.
  */
 export async function initialize(
-  _state: GraphState,
+  state: GraphState,
   config: GraphConfig,
 ): Promise<GraphUpdate> {
   if (!config.configurable) {
     throw new Error("Configuration object not found.");
   }
-  const { sandbox_session_id, target_repository, sandbox_language } =
-    config.configurable;
-  if (sandbox_session_id) {
+  const { sandboxSessionId } = state;
+
+  if (sandboxSessionId) {
     try {
+      console.log("Sandbox session ID exists. Resuming...", {
+        sandboxSessionId,
+      });
       // Resume the sandbox if the session ID is in the config.
-      await Sandbox.resume(sandbox_session_id, TIMEOUT_EXTENSION_OPT);
-      return {};
+      const newSandbox = await Sandbox.resume(
+        sandboxSessionId,
+        TIMEOUT_EXTENSION_OPT,
+      );
+      return {
+        sandboxSessionId: newSandbox.sandboxId,
+      };
     } 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);
     }
   }
 
+  const { target_repository, sandbox_language } = config.configurable;
+
   if (!sandbox_language || !target_repository) {
     throw new Error(
       "Missing required configuration. Please provide a sandbox language and git repository URL.",
@@ -63,18 +76,22 @@ export async function initialize(
   }
 
   if (sandbox_language === "js") {
+    console.log("Creating JS sandbox...");
     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) {
       // TODO: This should probably be an interrupt.
+      console.error("Failed to clone repository.", res.error);
       throw new Error(`Failed to clone repository.\n${res.error}`);
     }
-    return {};
+    console.log("Repository cloned successfully.");
+    return {
+      sandboxSessionId: sandbox.sandboxId,
+    };
   }
 
   if (sandbox_language === "python") {

src/nodes/interrupt-plan.ts

@@ -5,9 +5,9 @@ import {
   HumanInterrupt,
   HumanResponse,
 } from "@langchain/langgraph/prebuilt";
-import { v4 as uuidv4 } from "uuid";
+import { resumeSandbox } from "../utils/sandbox.js";
 
-export function interruptPlan(state: GraphState): Command {
+export async function interruptPlan(state: GraphState): Promise<Command> {
   const { proposedPlan } = state;
   if (!proposedPlan.length) {
     throw new Error("No proposed plan found.");
@@ -30,33 +30,46 @@ export function interruptPlan(state: GraphState): Command {
     If editing the plan, ensure each step in the plan is separated by ":::".`,
   })[0];
 
+  if (!state.sandboxSessionId) {
+    // TODO: This should prob just create a sandbox?
+    throw new Error("No sandbox session ID found.");
+  }
+
   if (interruptRes.type === "accept") {
-    // Plan was accepted, route to the initialize node.
+    const newSandboxSessionId = (await resumeSandbox(state.sandboxSessionId))
+      .sandboxId;
+
+    // Plan was accepted, route to the generate-action node to start taking actions.
     return new Command({
-      goto: "initialize",
+      goto: "generate-action",
       update: {
-        plan: proposedPlan.map((p) => ({
-          id: uuidv4(),
+        plan: proposedPlan.map((p, index) => ({
+          index,
           plan: p,
           completed: false,
         })),
+        sandboxSessionId: newSandboxSessionId,
       },
     });
   }
 
   if (interruptRes.type === "edit") {
-    // Plan was edited, route to the initialize node.
+    const newSandboxSessionId = (await resumeSandbox(state.sandboxSessionId))
+      .sandboxId;
+
+    // Plan was edited, route to the generate-action node to start taking actions.
     const editedPlan = (interruptRes.args as ActionRequest).args.plan
       .split(":::")
       .map((step: string) => step.trim());
     return new Command({
-      goto: "initialize",
+      goto: "generate-action",
       update: {
-        plan: editedPlan.map((p: string) => ({
-          id: uuidv4(),
+        plan: editedPlan.map((p: string, index: number) => ({
+          index,
           plan: p,
           completed: false,
         })),
+        sandboxSessionId: newSandboxSessionId,
       },
     });
   }

src/nodes/take-action.ts

@@ -28,11 +28,21 @@ export async function takeAction(
   if (!tool) {
     throw new Error(`Unknown tool: ${toolCall.name}`);
   }
+  if (!state.sandboxSessionId) {
+    throw new Error(
+      "Failed to take action: No sandbox session ID found in state.",
+    );
+  }
 
   // @ts-expect-error tool.invoke types are weird here...
-  const result: ToolMessage = await tool.invoke(toolCall.args);
+  const result: string = await tool.invoke(toolCall.args);
+  const toolMessage = new ToolMessage({
+    tool_call_id: toolCall.id ?? "",
+    content: result,
+    name: toolCall.name,
+  });
 
   return {
-    messages: [...state.messages, result],
+    messages: [toolMessage],
   };
 }

src/subgraphs/index.ts

@@ -0,0 +1 @@
+export * from "./planner/index.js";

src/subgraphs/planner/index.ts

@@ -0,0 +1,66 @@
+/**
+ * Subgraph for gathering context & generating a plan.
+ * pre-requisites:
+ * VM is already booted & repo is cloned (init node)
+ * Steps:
+ */
+
+import { END, Send, START, StateGraph } from "@langchain/langgraph";
+import { PlannerGraphState, PlannerGraphStateObj } from "./types.js";
+import { GraphConfiguration } from "../../types.js";
+import { generateAction, generatePlan, takeAction } from "./nodes/index.js";
+import { isAIMessage, ToolMessage } from "@langchain/core/messages";
+
+function takeActionOrGeneratePlan(
+  state: PlannerGraphState,
+): "take-plan-action" | "generate-plan" | Send {
+  const { plannerMessages } = state;
+  const lastMessage = plannerMessages[plannerMessages.length - 1];
+  // If the last message is a tool call, and we have executed less than 6 actions, take action.
+  // Max actions is 13, because that's 6 actions (2 messages per action, ai & tool) plus the input message.
+  const maxActionsCount = 13;
+  if (
+    isAIMessage(lastMessage) &&
+    lastMessage.tool_calls?.length &&
+    plannerMessages.length < maxActionsCount
+  ) {
+    return "take-plan-action";
+  }
+
+  if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) {
+    // If this is true, we need to return a `Command` adding a ToolMessage to the state
+    // so that the last AI message has a tool message pair
+    const lastMessageToolCall = lastMessage.tool_calls[0];
+    return new Send("generate-plan", {
+      ...state,
+      plannerMessages: [
+        ...state.plannerMessages,
+        new ToolMessage({
+          tool_call_id: lastMessageToolCall.id ?? "",
+          name: lastMessageToolCall.name,
+          content: "Tool call not executed. Max actions reached.",
+        }),
+      ],
+    });
+  }
+
+  // If the last message does not have tool calls, continue to generate plan without modifications.
+  return "generate-plan";
+}
+
+const workflow = new StateGraph(PlannerGraphStateObj, GraphConfiguration)
+  .addNode("generate-plan-context-action", generateAction)
+  .addNode("take-plan-action", takeAction)
+  .addNode("generate-plan", generatePlan)
+  .addEdge(START, "generate-plan-context-action")
+  .addConditionalEdges(
+    "generate-plan-context-action",
+    takeActionOrGeneratePlan,
+    ["take-plan-action", "generate-plan"],
+  )
+  .addEdge("take-plan-action", "generate-plan-context-action")
+  .addEdge("generate-plan", END);
+
+// TODO: Fix zod types
+export const plannerGraph = workflow.compile() as any;
+plannerGraph.name = "Planner";

src/subgraphs/planner/nodes/generate-message.ts

@@ -0,0 +1,45 @@
+import { loadModel, Task } from "../../../utils/load-model.js";
+import { shellTool, applyPatchTool } from "../../../tools/index.js";
+import { PlannerGraphState, PlannerGraphUpdate } from "../types.js";
+import { GraphConfig } from "../../../types.js";
+import { isHumanMessage } from "@langchain/core/messages";
+
+const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful.
+
+Your sole task is to gather context from the repository the user has provided which will be helpful when generating a plan to address the user's request.
+
+You MUST adhere to the following criteria when gathering context for the plan:
+- You must ONLY take read actions to gather context. Write actions are NOT allowed.
+- Keep in mind you are only permitted to make a maximum of 6 tool calls to gather all your context. Ensure each action is of high quality, and targeted to aid in generating a plan.
+- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore.
+  - Always use glob patterns when searching with \`rg\` for specific file types. For example, to search for all TSX files, use \`rg -i star -g **/*.tsx project-directory/\`. This is because \`rg\` does not have built in file types for every language.
+- If you determine you've gathered enough context to generate a plan, simply reply with 'done' and do NOT call any tools.
+- Not generating a tool call will be interpreted as an indication that you've gathered enough context to generate a plan.
+- The first user message in this conversation contains the user's request.
+`;
+
+export async function generateAction(
+  state: PlannerGraphState,
+  config: GraphConfig,
+): Promise<PlannerGraphUpdate> {
+  const model = await loadModel(config, Task.ACTION_GENERATOR);
+  const tools = [shellTool, applyPatchTool];
+  const modelWithTools = model.bindTools(tools, { tool_choice: "auto" });
+
+  const firstUserMessage = state.messages.find(isHumanMessage);
+
+  const response = await modelWithTools
+    .bind({ tags: ["langsmith:nostream"] })
+    .invoke([
+      {
+        role: "system",
+        content: systemPrompt,
+      },
+      ...(firstUserMessage ? [firstUserMessage] : []),
+      ...state.plannerMessages,
+    ]);
+
+  return {
+    plannerMessages: [response],
+  };
+}

src/subgraphs/planner/nodes/generate-plan.ts

@@ -0,0 +1,58 @@
+import { isHumanMessage } from "@langchain/core/messages";
+import { sessionPlanTool } from "../../../tools/index.js";
+import { GraphConfig } from "../../../types.js";
+import { loadModel, Task } from "../../../utils/load-model.js";
+import { PlannerGraphState, PlannerGraphUpdate } from "../types.js";
+import { pauseSandbox } from "../../../utils/sandbox.js";
+
+const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful.
+
+In this step, you are expected to generate a high-level plan to address the user's request. The plan should be a list of actions to take, in order, to address the user's request. You should not include any code in the plan, only a list of actions to take.
+
+You MUST adhere to the following criteria when generating the plan:
+- You have already gathered context from the repository the user has requested you take actions on. This context is provided in the conversation history below.
+- Your plan should be high-level in nature, but should still be specific enough to be actionable.
+- Ensure your plan is as concise as possible. Omit any unnecessary details or steps. Your goal is to complete the task in the least number of steps possible.
+- Do not pack multiple complex tasks into a single plan item. Each high level task you'll need to complete should have its own plan item.
+- When you are ready to generate the plan, ensure you call the 'session_plan' tool. You are REQUIRED to call this tool.
+- The first user message in this conversation contains the user's request.
+`;
+
+export async function generatePlan(
+  state: PlannerGraphState,
+  config: GraphConfig,
+): Promise<PlannerGraphUpdate> {
+  const model = await loadModel(config, Task.PLANNER);
+  const modelWithTools = model.bindTools([sessionPlanTool], {
+    tool_choice: sessionPlanTool.name,
+  });
+
+  const firstUserMessage = state.messages.find(isHumanMessage);
+
+  const response = await modelWithTools
+    .bind({ tags: ["langsmith:nostream"] })
+    .invoke([
+      {
+        role: "system",
+        content: systemPrompt,
+      },
+      ...(firstUserMessage ? [firstUserMessage] : []),
+      ...state.plannerMessages,
+    ]);
+
+  if (!response.tool_calls?.length) {
+    throw new Error("Failed to generate plan");
+  }
+
+  let newSessionId: string | undefined;
+  if (state.sandboxSessionId) {
+    // Pause before returning, as the next step will be to interrupt the graph.
+    newSessionId = await pauseSandbox(state.sandboxSessionId);
+  }
+
+  return {
+    proposedPlan: response.tool_calls[0].args.plan,
+    plan: [],
+    ...(newSessionId && { sandboxSessionId: newSessionId }),
+  };
+}

src/subgraphs/planner/nodes/index.ts

@@ -0,0 +1,3 @@
+export * from "./generate-message.js";
+export * from "./take-action.js";
+export * from "./generate-plan.js";

src/subgraphs/planner/nodes/take-action.ts

@@ -0,0 +1,45 @@
+import { isAIMessage, ToolMessage } from "@langchain/core/messages";
+import { applyPatchTool, shellTool } from "../../../tools/index.js";
+import { GraphConfig } from "../../../types.js";
+import { PlannerGraphState, PlannerGraphUpdate } from "../types.js";
+
+export async function takeAction(
+  state: PlannerGraphState,
+  _config: GraphConfig,
+): Promise<PlannerGraphUpdate> {
+  const { plannerMessages: messages } = state;
+  const lastMessage = messages[messages.length - 1];
+
+  if (!isAIMessage(lastMessage) || !lastMessage.tool_calls?.length) {
+    throw new Error("Last message is not an AI message with tool calls.");
+  }
+
+  const toolsMap = {
+    [applyPatchTool.name]: applyPatchTool,
+    [shellTool.name]: shellTool,
+  };
+
+  const toolCall = lastMessage.tool_calls[0];
+
+  if (!toolCall) {
+    throw new Error("No tool call found.");
+  }
+
+  const tool = toolsMap[toolCall.name];
+
+  if (!tool) {
+    throw new Error(`Unknown tool: ${toolCall.name}`);
+  }
+
+  // @ts-expect-error tool.invoke types are weird here...
+  const result: string = await tool.invoke(toolCall.args);
+  const toolMessage = new ToolMessage({
+    tool_call_id: toolCall.id ?? "",
+    content: result,
+    name: toolCall.name,
+  });
+
+  return {
+    plannerMessages: [toolMessage],
+  };
+}

src/subgraphs/planner/types.ts

@@ -0,0 +1,15 @@
+import "@langchain/langgraph/zod";
+import { z } from "zod";
+import { addMessages, Messages } from "@langchain/langgraph";
+import { BaseMessage } from "@langchain/core/messages";
+import { GraphAnnotation } from "../../types.js";
+
+export const PlannerGraphStateObj = GraphAnnotation.extend({
+  plannerMessages: z
+    .custom<BaseMessage[]>()
+    .default(() => [])
+    .langgraph.reducer<Messages>((state, update) => addMessages(state, update)),
+});
+
+export type PlannerGraphState = z.infer<typeof PlannerGraphStateObj>;
+export type PlannerGraphUpdate = Partial<PlannerGraphState>;

src/tools/apply-patch.ts

@@ -1,9 +1,10 @@
 import { tool } from "@langchain/core/tools";
 import { z } from "zod";
 import { applyPatch } from "diff";
-import { GraphConfig } from "../types.js";
+import { GraphState } from "../types.js";
 import { Sandbox } from "@e2b/code-interpreter";
 import { readFile, writeFile } from "../utils/read-write.js";
+import { getCurrentTaskInput } from "@langchain/langgraph";
 
 const applyPatchToolSchema = z.object({
   diff: z.string().describe("The diff to apply. Use a standard diff format."),
@@ -11,14 +12,19 @@ const applyPatchToolSchema = z.object({
 });
 
 export const applyPatchTool = tool(
-  async (input, config: GraphConfig) => {
-    const { diff, file_path } = input;
-    const sessionId = config.configurable?.sandbox_session_id;
-    if (!sessionId) {
-      return "FAILED TO RUN COMMAND: No sandbox session ID provided";
+  async (input) => {
+    const state = getCurrentTaskInput<GraphState>();
+    const { sandboxSessionId } = state;
+    if (!sandboxSessionId) {
+      console.error("FAILED TO RUN COMMAND: No sandbox session ID provided", {
+        input,
+      });
+      throw new Error("FAILED TO RUN COMMAND: No sandbox session ID provided");
     }
 
-    const sandbox = await Sandbox.connect(sessionId);
+    const { diff, file_path } = input;
+
+    const sandbox = await Sandbox.connect(sandboxSessionId);
 
     const { success: readFileSuccess, output: readFileOutput } = await readFile(
       sandbox,

src/tools/shell.ts

@@ -1,7 +1,9 @@
 import { tool } from "@langchain/core/tools";
 import { z } from "zod";
 import { Sandbox } from "@e2b/code-interpreter";
-import { GraphConfig } from "../types.js";
+import { GraphState } from "../types.js";
+import { getCurrentTaskInput } from "@langchain/langgraph";
+import { TIMEOUT_MS } from "../constants.js";
 
 const shellToolSchema = z.object({
   command: z.array(z.string()).describe("The command to run"),
@@ -18,22 +20,59 @@ const shellToolSchema = z.object({
 });
 
 export const shellTool = tool(
-  async (input, config: GraphConfig) => {
-    const sessionId = config.configurable?.sandbox_session_id;
-    if (!sessionId) {
-      return "FAILED TO RUN COMMAND: No sandbox session ID provided";
+  async (input) => {
+    try {
+      const state = getCurrentTaskInput<GraphState>();
+      const { sandboxSessionId } = state;
+      if (!sandboxSessionId) {
+        console.error("FAILED TO RUN COMMAND: No sandbox session ID provided", {
+          input,
+        });
+        throw new Error(
+          "FAILED TO RUN COMMAND: No sandbox session ID provided",
+        );
+      }
+
+      const sandbox = await Sandbox.connect(sandboxSessionId);
+      const { command, workdir, timeout } = input;
+      const result = await sandbox.commands.run(command.join(" "), {
+        timeoutMs: timeout,
+        cwd: workdir,
+      });
+      // Add an extra 5 min timeout to the sandbox.
+      await sandbox.setTimeout(TIMEOUT_MS);
+
+      if (result.error) {
+        console.error("Failed to run command", {
+          error: result.error,
+          error_result: result,
+          input,
+        });
+        return `Command failed. Exit code: ${result.exitCode}\nError: ${result.error}\nStderr:\n${result.stderr}`;
+      }
+
+      return result.stdout;
+    } catch (e: any) {
+      if (typeof e === "object" && "result" in e && e.result) {
+        console.error("Failed to run command", {
+          error: e.message,
+          error_result: e.result,
+          input,
+        });
+        return (
+          "FAILED TO RUN COMMAND: " +
+          e.message +
+          "\n" +
+          JSON.stringify(e.result, null, 2)
+        );
+      }
+
+      console.error("Failed to run command: " + e.message, {
+        error: e,
+        input,
+      });
+      throw new Error("FAILED TO RUN COMMAND: " + e.message);
     }
-
-    const sandbox = await Sandbox.connect(sessionId);
-    const { command, workdir, timeout } = input;
-    const result = await sandbox.commands.run(command.join(" "), {
-      timeoutMs: timeout,
-      cwd: workdir,
-    });
-
-    return `Command exited with code ${result.exitCode}:\n
-${result.error ? `Error: ${result.error}\n` : ""}
-Stdout:\n${result.stdout}\n\nStderr:\n${result.stderr}`;
   },
   {
     name: "shell",

src/types.ts

@@ -47,6 +47,13 @@ export const GraphAnnotation = z.object({
     .nullable()
     .default(() => null)
     .langgraph.reducer((_state, update) => update),
+  /**
+   * The session ID of the Sandbox to use.
+   */
+  sandboxSessionId: z
+    .string()
+    .optional()
+    .langgraph.reducer((_state, update) => update),
 });
 
 export type GraphState = z.infer<typeof GraphAnnotation>;
@@ -111,18 +118,12 @@ const MODEL_OPTIONS = [
   },
 ];
 
+const MODEL_OPTIONS_NO_THINKING = MODEL_OPTIONS.filter(
+  ({ value }) =>
+    !value.includes("extended-thinking") || !value.startsWith("openai:o"),
+);
+
 export const GraphConfiguration = z.object({
-  /**
-   * The session ID of the Sandbox to use.
-   */
-  sandbox_session_id: z
-    .string()
-    .optional()
-    .langgraph.metadata({
-      x_lg_ui_config: {
-        type: "hidden",
-      },
-    }),
   /**
    * The URL of the repository to clone.
    */
@@ -132,25 +133,54 @@ export const GraphConfiguration = z.object({
       repo: z.string(),
       branch: z.string().optional(),
     })
-    .langgraph.metadata({}),
+    .langgraph.metadata({
+      x_oap_ui_config: {
+        type: "json",
+        default: `{
+  "owner": "",
+  "repo": "",
+  "branch": ""
+}`,
+      },
+    }),
   /**
    * The language of the sandbox to use.
    */
-  sandbox_language: z.enum(["js", "python"]).optional().langgraph.metadata({}),
+  sandbox_language: z
+    .enum(["js", "python"])
+    .optional()
+    .langgraph.metadata({
+      x_oap_ui_config: {
+        type: "select",
+        default: "js",
+        description: "The primary language of the sandbox to use.",
+        options: [
+          {
+            label: "JavaScript/TypeScript",
+            value: "js",
+          },
+          {
+            label: "Python",
+            value: "python",
+          },
+        ],
+      },
+    }),
   /**
    * The model ID to use for the planning step.
    * This includes initial planning, and rewriting.
-   * @default "anthropic:extended-thinking:claude-sonnet-4-0"
+   * @default "anthropic:claude-sonnet-4-0"
    */
   plannerModelName: z
     .string()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
         type: "select",
-        default: "anthropic:extended-thinking:claude-sonnet-4-0",
+        default: "anthropic:claude-sonnet-4-0",
         description: "The model to use for planning",
-        options: MODEL_OPTIONS,
+        // Do not show extended thinking models
+        options: MODEL_OPTIONS_NO_THINKING,
       },
     }),
   /**
@@ -163,7 +193,42 @@ export const GraphConfiguration = z.object({
     .number()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
+        type: "slider",
+        default: 0,
+        min: 0,
+        max: 2,
+        step: 0.1,
+        description: "Controls randomness (0 = deterministic, 2 = creative)",
+      },
+    }),
+  /**
+   * The model ID to use for the planning step.
+   * This includes initial planning, and rewriting.
+   * @default "anthropic:claude-sonnet-4-0"
+   */
+  plannerContextModelName: z
+    .string()
+    .optional()
+    .langgraph.metadata({
+      x_oap_ui_config: {
+        type: "select",
+        default: "anthropic:claude-sonnet-4-0",
+        description: "The model to use for planning",
+        options: MODEL_OPTIONS,
+      },
+    }),
+  /**
+   * The temperature to use for the planning step.
+   * This includes initial planning, and rewriting.
+   * If selecting a reasoning model, this will be ignored.
+   * @default 0
+   */
+  plannerContextTemperature: z
+    .number()
+    .optional()
+    .langgraph.metadata({
+      x_oap_ui_config: {
         type: "slider",
         default: 0,
         min: 0,
@@ -181,7 +246,7 @@ export const GraphConfiguration = z.object({
     .string()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
         type: "select",
         default: "anthropic:claude-sonnet-4-0",
         description: "The model to use for action generation",
@@ -197,7 +262,7 @@ export const GraphConfiguration = z.object({
     .number()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
         type: "slider",
         default: 0,
         min: 0,
@@ -215,11 +280,11 @@ export const GraphConfiguration = z.object({
     .string()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
         type: "select",
         default: "anthropic:claude-sonnet-4-0",
         description: "The model to use for progress plan checking",
-        options: MODEL_OPTIONS,
+        options: MODEL_OPTIONS_NO_THINKING,
       },
     }),
   /**
@@ -231,7 +296,7 @@ export const GraphConfiguration = z.object({
     .number()
     .optional()
     .langgraph.metadata({
-      x_lg_ui_config: {
+      x_oap_ui_config: {
         type: "slider",
         default: 0,
         min: 0,

src/utils/load-model.ts

@@ -3,13 +3,18 @@ import { GraphConfig } from "../types.js";
 
 export enum Task {
   PLANNER = "planner",
+  PLANNER_CONTEXT = "plannerContext",
   ACTION_GENERATOR = "actionGenerator",
   PROGRESS_PLAN_CHECKER = "progressPlanChecker",
 }
 
 const TASK_TO_CONFIG_DEFAULTS_MAP = {
   [Task.PLANNER]: {
-    modelName: "anthropic:extended-thinking:claude-sonnet-4-0",
+    modelName: "anthropic:claude-sonnet-4-0",
+    temperature: 0,
+  },
+  [Task.PLANNER_CONTEXT]: {
+    modelName: "anthropic:claude-sonnet-4-0",
     temperature: 0,
   },
   [Task.ACTION_GENERATOR]: {
@@ -34,9 +39,11 @@ export async function loadModel(config: GraphConfig, task: Task) {
 
   let thinkingModel = false;
   if (modelNameParts[0] === "extended-thinking") {
+    console.log("THINKING ENABLED", modelNameParts);
     // Using a thinking model. Remove it from the model name.
     modelNameParts.shift();
     thinkingModel = true;
+    console.log("THINKING ENABLED", modelNameParts);
   }
 
   const modelName = modelNameParts.join(":");
@@ -44,11 +51,17 @@ export async function loadModel(config: GraphConfig, task: Task) {
     thinkingModel = true;
   }
 
+  const thinkingBudgetTokens = 5000;
+  const maxTokens = thinkingBudgetTokens * 4;
+
   const model = await initChatModel(modelName, {
     modelProvider,
     temperature: thinkingModel ? undefined : temperature,
     ...(thinkingModel && modelProvider === "anthropic"
-      ? { thinking: { budgetTokens: 5000, type: "enabled" } }
+      ? {
+          thinking: { budget_tokens: thinkingBudgetTokens, type: "enabled" },
+          maxTokens,
+        }
       : {}),
   });
 

src/utils/plan-prompt.ts

@@ -16,11 +16,15 @@ export function formatPlanPrompt(plan: PlanItem[]): string {
 
   return PLAN_PROMPT.replace(
     "{COMPLETED_TASKS}",
-    completedTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"),
+    completedTasks?.length
+      ? completedTasks.map((task) => `${task.index}. ${task.plan}`).join("\n")
+      : "No completed tasks.",
   )
     .replace(
       "{REMAINING_TASKS}",
-      remainingTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"),
+      remainingTasks?.length
+        ? remainingTasks.map((task) => `${task.index}. ${task.plan}`).join("\n")
+        : "No remaining tasks.",
     )
-    .replace("{CURRENT_TASK}", currentTask.plan);
+    .replace("{CURRENT_TASK}", currentTask?.plan || "No current task.");
 }

src/utils/read-write.ts

@@ -1,4 +1,5 @@
 import { Sandbox } from "@e2b/code-interpreter";
+import { TIMEOUT_MS } from "../constants.js";
 
 export async function readFile(
   sandbox: Sandbox,
@@ -9,6 +10,9 @@ export async function readFile(
 }> {
   try {
     const readOutput = await sandbox.commands.run(`cat "${filePath}"`);
+    // Add an extra 5 min timeout to the sandbox.
+    await sandbox.setTimeout(TIMEOUT_MS);
+
     if (readOutput.exitCode !== 0) {
       console.error(
         `Error reading file '${filePath}' from sandbox via cat:`,
@@ -51,6 +55,8 @@ export async function writeFile(
   try {
     const writeCommand = `printf '%s' '${content}' > "${filePath}"`;
     const writeOutput = await sandbox.commands.run(writeCommand);
+    // Add an extra 5 min timeout to the sandbox.
+    await sandbox.setTimeout(TIMEOUT_MS);
 
     if (writeOutput.exitCode !== 0) {
       console.error(

src/utils/sandbox.ts

@@ -15,3 +15,14 @@ export async function pauseSandbox(
   const sandboxClient = sandbox ?? (await Sandbox.connect(sandboxSessionId));
   return await sandboxClient.pause();
 }
+
+/**
+ * Resumes the sandbox.
+ * @param sandboxSessionId The ID of the sandbox to resume.
+ * @returns The sandbox client.
+ */
+export async function resumeSandbox(
+  sandboxSessionId: string,
+): Promise<Sandbox> {
+  return await Sandbox.resume(sandboxSessionId);
+}