fix: Improve rewrite plan and planner subgraph (#25)

apps/open-swe/package.json

@@ -31,6 +31,7 @@
     "@langchain/openai": "^0.5.10",
     "diff": "^8.0.1",
     "langchain": "^0.3.26",
+    "langsmith": "^0.3.29",
     "uuid": "^11.0.5",
     "zod": "^3.23.8"
   },

apps/open-swe/scripts/run-e2e.ts

@@ -15,8 +15,16 @@ async function runE2E() {
 
   const threadId = uuidv4();
 
-  const userRequest =
-    "The system message contains a REPO_DIRECTORY variable. But nothing is formatted in that variable. Please make sure to format the REPO_DIRECTORY variable in the system message. You will probably want to move the repo directory info from config to the State.";
+  const userRequest = `The project is structured as a monorepo, with some apps located inside the /apps directory. In this directory, there is an /auth directory. This directory only contains the scaffolding for a new app in the monorepo, but is not yet implemented. Please take the following plan/task description and implement it in the /auth directory:
+This monorepo is for an AI coding agent. The app runs and edits the code in the cloud in a sandboxed environment. Right now, we require users to generate a GitHub PAT, which we store in a .env file and can use to authenticate with GitHub. This is not idea, and instead we want to have a github oauth app which users can authenticate with.
+Please implement a new auth server inside the /auth directory which can do this.
+You will not have any access to secrets, so you will not be able to run the server to test it.
+I want the server to be able to authenticate users with GitHub, such that we will be able to take the following actions:
+1. clone repositories they give us access to
+2. checkout existing and create new branches on the repositories they give us access to
+3. make pull requests and push changes to the repositories they give us access to
+Once you're done, ensure you've documented the development process in the readme of this new app.`;
+
   const configurable: Omit<
     GraphConfig["configurable"],
     "thread_id" | "assistant_id"

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

@@ -38,7 +38,10 @@ You MUST adhere to the following criteria when executing the task:
 - Analyzing code for vulnerabilities is allowed.
 - Showing user code and tool call details is allowed.
 - Remember to always properly format and quote your shell commands.
-- Take advantage of the condensed context tool call messages in the conversation history. These contain summarized/condensed context from previously completed steps. Ensure you always read these messages to avoid duplicate work (e.g.: searching for file paths).
+- Take advantage of the condensed context tool call messages in the conversation history (under the names \`condense_task_context\` and \`condense_planning_context\`). These contain summarized/condensed context from previously completed steps. Ensure you always read these messages to avoid duplicate work (e.g.: searching for file paths).
+  - These summary messages may include a section called 'Codebase files and descriptions' which contains a list of files, and descriptions of the files' contents. If you need context on a file, or directory, ensure you first check this section of the summary messages to avoid duplicate work.
+  - The summary messages may also include a section called 'Key repository insights and learnings'. This contains key insights, learnings, and facts the model discovered while completing a task.
+  - Each summary message will also include a short description of the task it completed, how it did so, and every change it made to the codebase during this task. This section will be titled 'Repository modifications summary'.
 - All changes are automatically committed, so you should not worry about creating backups, or committing changes.
 - Use \`apply_patch\` to edit files. This tool accepts diffs and file paths. It will then apply the given diff to the file.
 - 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.

apps/open-swe/src/nodes/rewrite-plan.ts

@@ -1,33 +1,235 @@
-import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
+import { GraphState, GraphConfig, GraphUpdate, PlanItem } from "../types.js";
 import { loadModel, Task } from "../utils/load-model.js";
-import { sessionPlanTool } from "../tools/index.js";
+import { isHumanMessage } from "@langchain/core/messages";
+import { getMessageContentString } from "../utils/message/content.js";
+import { z } from "zod";
+import { tool } from "@langchain/core/tools";
+import { ConfigurableModel } from "langchain/chat_models/universal";
+import { traceable } from "langsmith/traceable";
 
-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.
+const systemPromptIdentifyChanges = `You are operating as an agentic coding assistant built by LangChain. You've previously been given a task to generate a plan of action for, to address the user's initial request.
 
-In this step, you are expected to rewrite a high-level plan to address the user's initial request. In a previous step you generated a plan, however the user has requested some changes:
-## User Request
+Here is the user's initial request:
+{USER_INITIAL_REQUEST}
+
+After generating that plan, the user has submitted some feedback/change requests. You should now identify exactly which tasks in the plan should be modified based on their request.
+
+Here is their request:
 {USER_REQUEST}
 
-Here is the previous plan:
-## Previous Plan
-{PREVIOUS_PLAN}
+The plan you generated originally, which they submitted the above request for is as follows:
+{PLAN}
 
-The plan must 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.
+Please read over the generated plan, and the user's request, and identify exactly which tasks in the plan should be modified/removed. Call the 'identify_plan_changes' tool and use the indices of the tasks listed above when calling the tool.`;
+
+const systemPrompt = `You are operating as an agentic coding assistant built by LangChain. You've previously been given a task to generate a plan of action for, to address the user's initial request.
+
+In this step, the user has requested you rewrite/modify parts of a high-level plan. You have already identified the specific tasks in the plan that should be modified/removed based on the user's request.
+
+Here is the user's initial request which you used to generate the initial plan:
+{USER_INITIAL_REQUEST}
+
+Here is the full plan you generated:
+{PLAN}
+
+Here is the request the user has just made which you should use to rewrite/modify the plan:
+{USER_REQUEST}
+
+And here are the specific tasks in the plan which were identified as tasks the user wants to modify/remove:
+{TASKS_TO_MODIFY}
+
+Given this context, please address the user's request to modify/remove/add the tasks in the plan.
 
 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.
-- Make as few changes as possible to the previous plan, while still addressing the user's request.
-- When you are ready to generate the plan, ensure you call the 'session_plan' tool.
-- Ensure you generate the full plan in this tool call, not just the changes.
+- Make as few changes as possible to the tasks, while still following the users request.
+- You should NOT make ANY changes to the tasks in the plan that are NOT listed as tasks to modify/remove.
+- Do NOT modify tasks in the plan not listed as tasks to modify/remove.
+- When responding, ensure you include the unmodified tasks in the plan, as well as the modified/new tasks.
+- To remove a specific task, simply do NOT include it in the response.
+- To add a new task, simply include it in the response.
 `;
 
-const formatSysPrompt = (userRequest: string, previousPlan: string) => {
-  return systemPrompt
+const formatSysPromptIdentifyTasks = (
+  userInitialRequest: string,
+  userRequest: string,
+  previousPlan: string[],
+) => {
+  return systemPromptIdentifyChanges
+    .replace("{USER_INITIAL_REQUEST}", userInitialRequest)
     .replace("{USER_REQUEST}", userRequest)
-    .replace("{PREVIOUS_PLAN}", previousPlan);
+    .replace(
+      "{PLAN}",
+      previousPlan.map((plan, index) => `${index}: ${plan}`).join("\n"),
+    );
 };
 
+const formatSysPromptRewritePlan = (
+  userInitialRequest: string,
+  userRequest: string,
+  previousPlan: string[],
+  tasksToModify: PlanItem[],
+) => {
+  return systemPrompt
+    .replace("{USER_INITIAL_REQUEST}", userInitialRequest)
+    .replace("{USER_REQUEST}", userRequest)
+    .replace(
+      "{PLAN}",
+      previousPlan.map((plan, index) => `${index}: ${plan}`).join("\n"),
+    )
+    .replace(
+      "{TASKS_TO_MODIFY}",
+      tasksToModify.map((p) => `${p.index}: ${p.plan}`).join("\n"),
+    );
+};
+
+async function identifyTasksToModifyFunc(
+  state: GraphState,
+  model: ConfigurableModel,
+): Promise<PlanItem[]> {
+  if (!state.planChangeRequest) {
+    throw new Error("No plan change request found.");
+  }
+
+  const identifyPlanChangesSchema = z.object({
+    task_change_indices: z
+      .array(z.number())
+      .describe(
+        "The indices of the tasks in the plan that should be modified/removed.",
+      ),
+  });
+
+  const identifyPlanChangesTool = tool(
+    (input): PlanItem[] => {
+      const { task_change_indices } = input;
+      const tasksToModify = state.proposedPlan.flatMap((plan, planIndex) => {
+        const planItem = task_change_indices.some(
+          (changeIndex) => changeIndex === planIndex,
+        );
+        if (!planItem) {
+          return [];
+        }
+        return {
+          index: planIndex,
+          plan: plan,
+          completed: false,
+        };
+      });
+
+      return tasksToModify;
+    },
+    {
+      name: "identify_plan_changes",
+      schema: identifyPlanChangesSchema,
+      description:
+        "Identify which tasks in the plan should be modified/removed based on the user's request.",
+    },
+  );
+
+  const modelWithIdentifyChangesTool = model.bindTools(
+    [identifyPlanChangesTool],
+    {
+      // The model should always call the tool when identifying plan changes.
+      tool_choice: identifyPlanChangesTool.name,
+    },
+  );
+
+  const firstUserMessage = state.messages.find(isHumanMessage);
+
+  const response = await modelWithIdentifyChangesTool.invoke([
+    {
+      role: "user",
+      content: formatSysPromptIdentifyTasks(
+        getMessageContentString(
+          firstUserMessage?.content ?? "No user message found",
+        ),
+        state.planChangeRequest,
+        state.proposedPlan,
+      ),
+    },
+  ]);
+
+  const toolCall = response.tool_calls?.[0];
+  if (!toolCall) {
+    throw new Error(
+      "Tool call not returned when attempting to identify plan changes.",
+    );
+  }
+
+  const tasksToModify = await identifyPlanChangesTool.invoke(
+    toolCall.args as z.infer<typeof identifyPlanChangesSchema>,
+  );
+  return tasksToModify;
+}
+
+const identifyTasksToModify = traceable(identifyTasksToModifyFunc, {
+  name: "identify_tasks_to_modify",
+});
+
+async function updatePlanTasksFunc(
+  state: GraphState,
+  tasksToModify: PlanItem[],
+  model: ConfigurableModel,
+): Promise<string[]> {
+  if (!state.planChangeRequest) {
+    throw new Error("No plan change request found.");
+  }
+
+  const updatePlanTasksSchema = z.object({
+    updated_plan_tasks: z
+      .array(
+        z
+          .string()
+          .describe(
+            "The updated or unmodified plan for the task. Do NOT include the task index.",
+          ),
+      )
+      .describe(
+        "The updated plan tasks. Must be in the order of which they should be executed in.",
+      ),
+  });
+  const updatePlanTasksTool = {
+    name: "update_plan_tasks",
+    description: "Call this tool to respond with the updated plan.",
+    schema: updatePlanTasksSchema,
+  };
+
+  const modelWithUpdatePlanTasksTool = model.bindTools([updatePlanTasksTool], {
+    // The model should always call the tool when identifying plan changes.
+    tool_choice: updatePlanTasksTool.name,
+  });
+
+  const firstUserMessage = state.messages.find(isHumanMessage);
+
+  const response = await modelWithUpdatePlanTasksTool.invoke([
+    {
+      role: "user",
+      content: formatSysPromptRewritePlan(
+        getMessageContentString(
+          firstUserMessage?.content ?? "No user message found",
+        ),
+        state.planChangeRequest,
+        state.proposedPlan,
+        tasksToModify,
+      ),
+    },
+  ]);
+
+  const toolCall = response.tool_calls?.[0];
+  if (!toolCall) {
+    throw new Error(
+      "Tool call not returned when attempting to update plan tasks.",
+    );
+  }
+
+  return (
+    toolCall.args as z.infer<typeof updatePlanTasksSchema>
+  ).updated_plan_tasks.map((p) => p);
+}
+
+const updatePlanTasks = traceable(updatePlanTasksFunc, {
+  name: "update_plan_tasks",
+});
+
 export async function rewritePlan(
   state: GraphState,
   config: GraphConfig,
@@ -37,28 +239,11 @@ export async function rewritePlan(
   }
 
   const model = await loadModel(config, Task.PLANNER);
-  const modelWithTools = model.bindTools([sessionPlanTool], {
-    // The model should always call the tool when rewriting the plan.
-    tool_choice: sessionPlanTool.name,
-  });
+  const tasksToModify = await identifyTasksToModify(state, model);
+  const updatedPlanTasks = await updatePlanTasks(state, tasksToModify, model);
 
-  const response = await modelWithTools.invoke([
-    {
-      role: "system",
-      content: formatSysPrompt(
-        state.planChangeRequest,
-        "  - " + state.plan.join("\n  - "),
-      ),
-    },
-    ...state.messages,
-  ]);
-
-  if (response.tool_calls?.length) {
-    return {
-      proposedPlan: response.tool_calls[0].args.plan,
-      plan: [],
-    };
-  }
-
-  throw new Error("Failed to rewrite plan.");
+  return {
+    plan: [],
+    proposedPlan: updatedPlanTasks,
+  };
 }

apps/open-swe/src/nodes/summarize-task-steps.ts

@@ -25,7 +25,12 @@ You do not want to keep the entire conversation history, but instead you want to
 {PLAN_PROMPT}
 
 You MUST adhere to the following criteria when summarizing the conversation history:
-- Retain context such as file paths, versions, and installed software.
+- Retain context such as file paths, versions, and installed software which future iterations will find useful.
+  - It is very important to include the file paths of files you've already searched for, along with a description of the file's contents, inside a 'Codebase files and descriptions' section, so that future steps can reuse this information, and will not need to search through the codebase for files again.
+  - Consider including a section titled 'Key repository insights and learnings' which may include information, insights and learnings you've discovered while completing the task.
+    - This section should be concise, but still including enough information so following steps will not repeat any mistakes or go down rabbit holes which you already know about.
+  - If changes were made to the repository during this task, ensure you include a section titled 'Repository modifications summary' which should include a short description of the task it completed, how it did so, and every change it made to the codebase during this task.
+    - Do not include the actual changes you made, but rather high level bullet points containing context and descriptions on the modifications made.
 - Do not retain any full code snippets.
 - Do not retain any full file contents.
 - Ensure your summary is concise, but useful for future context.

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

@@ -29,7 +29,7 @@ export async function generateAction(
   const firstUserMessage = state.messages.find(isHumanMessage);
 
   const response = await modelWithTools
-    .bind({ tags: ["langsmith:nostream"] })
+    .withConfig({ tags: ["langsmith:nostream"] })
     .invoke([
       {
         role: "system",

apps/open-swe/src/subgraphs/planner/nodes/generate-plan.ts

@@ -14,9 +14,11 @@ const systemPrompt = `You are operating as a terminal-based agentic coding assis
 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.
+- You have already gathered context from the repository the user has requested you take actions on, and are now ready to generate a plan based on it.
+  - 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.
+- Ensure your plan is as concise as possible. Omit any unnecessary details or steps the user did not request, or are not required to complete the task.
+  - Your goal is to complete the task outlined by the user 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.
@@ -45,7 +47,7 @@ export async function generatePlan(
   }
 
   const response = await modelWithTools
-    .bind({ tags: ["langsmith:nostream"] })
+    .withConfig({ tags: ["langsmith:nostream"] })
     .invoke([
       {
         role: "system",

apps/open-swe/src/subgraphs/planner/nodes/summarizer.ts

@@ -19,6 +19,9 @@ You do not want to keep the entire conversation history, but instead you want to
 
 You MUST adhere to the following criteria when summarizing the conversation history:
 - Retain context such as file paths, versions, and installed software.
+  - It is very important to include the file paths of files you've already searched for, along with a description of the file's contents, inside a 'Codebase files and descriptions' section, so that future steps can reuse this information, and will not need to search through the codebase for files again.
+- Consider including a section titled 'Key repository insights and learnings' which may include information, insights and learnings you've discovered while gathering context for the user's request.
+  - This section should be concise, but still including enough information so following steps will not repeat any mistakes or go down rabbit holes which you already know about.
 - Do not retain any full code snippets.
 - Do not retain any full file contents.
 - Ensure your summary is concise, but useful for future context.

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

@@ -12,11 +12,13 @@ const logger = createLogger(LogLevel.INFO, "ShellTool");
 const DEFAULT_COMMAND_TIMEOUT = 60_000; // 1 minute
 
 const shellToolSchema = z.object({
-  command: z.array(z.string()).describe("The command to run"),
   workdir: z
     .string()
     .optional()
-    .describe("The working directory for the command."),
+    .describe(
+      "The working directory for the command. Ensure this path is NOT included in any command arguments, as it will be added automatically.",
+    ),
+  command: z.array(z.string()).describe("The command to run"),
   timeout: z
     .number()
     .optional()

yarn.lock

@@ -2172,6 +2172,7 @@ __metadata:
     eslint-plugin-prettier: ^4.2.1
     jest: ^29.7.0
     langchain: ^0.3.26
+    langsmith: ^0.3.29
     prettier: ^3.5.2
     ts-jest: ^29.1.0
     typescript: ~5.7.2