feat: Planner notes (#305)

* feat: Planner notes

* update frontend

* cr

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

@@ -18,6 +18,7 @@ import { filterHiddenMessages } from "../../../../utils/message/filter-hidden.js
 import { getTaskPlanFromIssue } from "../../../../utils/github/issue-task.js";
 import { createRgTool } from "../../../../tools/rg.js";
 import { formatCustomRulesPrompt } from "../../../../utils/custom-rules.js";
+import { createPlannerNotesTool } from "../../../../tools/planner-notes.js";
 
 const logger = createLogger(LogLevel.INFO, "GeneratePlanningMessageNode");
 
@@ -46,7 +47,11 @@ export async function generateAction(
   config: GraphConfig,
 ): Promise<PlannerGraphUpdate> {
   const model = await loadModel(config, Task.ACTION_GENERATOR);
-  const tools = [createRgTool(state), createShellTool(state)];
+  const tools = [
+    createRgTool(state),
+    createShellTool(state),
+    createPlannerNotesTool(),
+  ];
   const modelWithTools = model.bindTools(tools, {
     tool_choice: "auto",
     parallel_tool_calls: true,

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

@@ -15,6 +15,13 @@ import { stopSandbox } from "../../../utils/sandbox.js";
 import { filterHiddenMessages } from "../../../utils/message/filter-hidden.js";
 import { z } from "zod";
 import { formatCustomRulesPrompt } from "../../../utils/custom-rules.js";
+import { getPlannerNotes } from "../utils/get-notes.js";
+
+const PLANNER_NOTES_PROMPT = `Here is a collection of technical notes you took while gathering context for the plan. Ensure you take these into account when writing your plan.
+
+<planner_notes>
+{PLANNER_NOTES}
+</planner_notes>`;
 
 const systemPrompt = `You are a terminal-based agentic coding assistant built by LangChain, designed to enable natural language interaction with local codebases through wrapped LLM models.
 
@@ -65,13 +72,17 @@ Structure your plan items as clear directives, for example:
 
 {CUSTOM_RULES}
 
+{PLANNER_NOTES}
+
 Remember: Your goal is to create a focused, executable plan that efficiently accomplishes the user's request using the context you've already gathered.`;
 
 function formatSystemPrompt(state: PlannerGraphState): string {
   // It's a followup if there's more than one human message.
   const isFollowup = isFollowupRequest(state.taskPlan, state.proposedPlan);
   const userRequest = getUserRequest(state.messages);
-
+  const plannerNotes = getPlannerNotes(state.messages)
+    .map((n) => `- ${n}`)
+    .join("\n");
   return systemPrompt
     .replace(
       "{FOLLOWUP_MESSAGE_PROMPT}",
@@ -82,7 +93,13 @@ function formatSystemPrompt(state: PlannerGraphState): string {
         : "",
     )
     .replace("{USER_REQUEST}", userRequest)
-    .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(state.customRules));
+    .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(state.customRules))
+    .replaceAll(
+      "{PLANNER_NOTES}",
+      plannerNotes.length
+        ? PLANNER_NOTES_PROMPT.replace("{PLANNER_NOTES}", plannerNotes)
+        : "",
+    );
 }
 
 export async function generatePlan(

apps/open-swe/src/graphs/planner/nodes/notetaker.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { CustomRules, GraphConfig } from "@open-swe/shared/open-swe/types";
+import { GraphConfig } from "@open-swe/shared/open-swe/types";
 import {
   PlannerGraphState,
   PlannerGraphUpdate,
@@ -7,8 +7,16 @@ import {
 import { loadModel, Task } from "../../../utils/load-model.js";
 import { getMessageString } from "../../../utils/message/content.js";
 import { getUserRequest } from "../../../utils/user-request.js";
-import { BaseMessage } from "@langchain/core/messages";
 import { formatCustomRulesPrompt } from "../../../utils/custom-rules.js";
+import { getPlannerNotes } from "../utils/get-notes.js";
+
+const PLANNER_NOTES_PROMPT = `You've also taken technical notes throughout the context gathering process. Ensure you include/incorporate these notes, or the highest quality parts of these notes in your conclusion notes.
+
+<planner_notes>
+{PLANNER_NOTES}
+</planner_notes>`;
+const CUSTOM_RULES_EXTRA_CONTEXT =
+  "- Carefully read over the user's custom rules to ensure you don't duplicate or repeat information found in that section, as you will always have access to it (even after the planning step!).";
 
 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.
 
@@ -19,12 +27,15 @@ The notes you extract should be thoughtful, and should include technical details
 These notes should not be overly verbose, as you'll be able to gather additional context when executing.
 Your goal is to generate notes on all of the low-hanging fruit from the conversation history, to speed up the execution so that you don't need to duplicate work to gather context.
 
+{CUSTOM_RULES}
+
+{PLANNER_NOTES}
+
 You MUST adhere to the following criteria when generating your notes:
 - Do not retain any full code snippets.
 - Do not retain any full file contents.
 - Only take notes on the context provided below, and do not make up, or attempt to infer any information/context which is not explicitly provided.
-
-{CUSTOM_RULES}
+{EXTRA_RULES}
 
 Here is the user's request
 ## User request:
@@ -41,20 +52,41 @@ And here is the plan you just generated:
 With all of this in mind, please carefully inspect the conversation history, and the plan you generated. Then, determine which actions and context from the conversation history will be most useful to you when you execute the plan. After you're done analyzing, call the \`write_technical_notes\` tool.
 `;
 
-const formatPrompt = (
-  userRequest: string,
-  conversationHistory: BaseMessage[],
-  proposedPlan: string[],
-  customRules?: CustomRules,
-): string =>
-  systemPrompt
+const formatPrompt = (state: PlannerGraphState): string => {
+  const userRequest =
+    getUserRequest(state.messages) || "No user request provided.";
+  const plannerNotes = getPlannerNotes(state.messages)
+    .map((n) => `  - ${n}`)
+    .join("\n");
+
+  return systemPrompt
     .replace("{USER_REQUEST}", userRequest)
     .replace(
       "{CONVERSATION_HISTORY}",
-      conversationHistory.map(getMessageString).join("\n"),
+      state.messages.map(getMessageString).join("\n"),
     )
-    .replace("{PROPOSED_PLAN}", `  - ${proposedPlan.join("\n  - ")}`)
-    .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(customRules));
+    .replace(
+      "{PROPOSED_PLAN}",
+      state.proposedPlan.map((p) => `  - ${p}`).join("\n"),
+    )
+    .replaceAll(
+      "{CUSTOM_RULES}",
+      formatCustomRulesPrompt(
+        state.customRules,
+        "Keep in mind these user provided rules will always be available to you, so any context present here should NOT be included in your notes as to not duplicate information.",
+      ),
+    )
+    .replaceAll(
+      "{PLANNER_NOTES}",
+      plannerNotes.length
+        ? PLANNER_NOTES_PROMPT.replace("{PLANNER_NOTES}", plannerNotes)
+        : "",
+    )
+    .replaceAll(
+      "{EXTRA_RULES}",
+      state.customRules ? CUSTOM_RULES_EXTRA_CONTEXT : "",
+    );
+};
 
 const condenseContextToolSchema = z.object({
   notes: z
@@ -78,7 +110,6 @@ export async function notetaker(
     parallel_tool_calls: false,
   });
 
-  const userRequest = getUserRequest(state.messages);
   const conversationHistoryStr = `Here is the full conversation history:
 
 ${state.messages.map(getMessageString).join("\n")}`;
@@ -86,12 +117,7 @@ ${state.messages.map(getMessageString).join("\n")}`;
   const response = await modelWithTools.invoke([
     {
       role: "system",
-      content: formatPrompt(
-        userRequest || "No user request provided.",
-        state.messages,
-        state.proposedPlan,
-        state.customRules,
-      ),
+      content: formatPrompt(state),
     },
     {
       role: "user",

apps/open-swe/src/graphs/planner/nodes/take-action.ts

@@ -10,6 +10,13 @@ import { zodSchemaToString } from "../../../utils/zod-to-string.js";
 import { formatBadArgsError } from "../../../utils/zod-to-string.js";
 import { truncateOutput } from "../../../utils/truncate-outputs.js";
 import { createRgTool } from "../../../tools/rg.js";
+import {
+  getChangedFilesStatus,
+  stashAndClearChanges,
+} from "../../../utils/github/git.js";
+import { getRepoAbsolutePath } from "@open-swe/shared/git";
+import { daytonaClient } from "../../../utils/sandbox.js";
+import { createPlannerNotesTool } from "../../../tools/planner-notes.js";
 
 const logger = createLogger(LogLevel.INFO, "TakeAction");
 
@@ -26,9 +33,11 @@ export async function takeActions(
 
   const shellTool = createShellTool(state);
   const rgTool = createRgTool(state);
+  const plannerNotesTool = createPlannerNotesTool();
   const toolsMap = {
     [shellTool.name]: shellTool,
     [rgTool.name]: rgTool,
+    [plannerNotesTool.name]: plannerNotesTool,
   };
 
   const toolCalls = lastMessage.tool_calls;
@@ -96,7 +105,34 @@ export async function takeActions(
     return toolMessage;
   });
 
-  const toolCallResults = await Promise.all(toolCallResultsPromise);
+  let toolCallResults = await Promise.all(toolCallResultsPromise);
+  const sandbox = await daytonaClient().get(state.sandboxSessionId);
+  const repoPath = getRepoAbsolutePath(state.targetRepository);
+  const changedFiles = await getChangedFilesStatus(repoPath, sandbox);
+  if (changedFiles?.length > 0) {
+    logger.warn(
+      "Changes found in the codebase after taking action. Reverting.",
+      {
+        changedFiles,
+      },
+    );
+    await stashAndClearChanges(repoPath, sandbox);
+
+    // Rewrite the tool call contents to include a changed files warning.
+    toolCallResults = toolCallResults.map(
+      (tc) =>
+        new ToolMessage({
+          ...tc,
+          content: `**WARNING**: THIS TOOL, OR A PREVIOUS TOOL HAS CHANGED FILES IN THE REPO.
+Remember that you are only permitted to take **READ** actions during the planning step. The changes have been reverted.
+
+Please ensure you only take read actions during the planning step to gather context. You may also call the \`take_notes\` tool at any time to record important information for the programmer step.
+
+Command Output:\n
+${tc.content}`,
+        }),
+    );
+  }
 
   logger.info("Completed planner tool action", {
     ...toolCallResults.map((tc) => ({

apps/open-swe/src/graphs/planner/utils/get-notes.ts

@@ -0,0 +1,22 @@
+import { BaseMessage, isAIMessage } from "@langchain/core/messages";
+import { createTakePlannerNotesFields } from "@open-swe/shared/open-swe/tools";
+import z from "zod";
+
+export function getPlannerNotes(messages: BaseMessage[]): string[] {
+  const plannerNotesFields = createTakePlannerNotesFields();
+  const plannerNotes = messages.flatMap((m) => {
+    if (!isAIMessage(m)) {
+      return [];
+    }
+    const notesToolCalls = m.tool_calls?.filter(
+      (tc) => tc.name === plannerNotesFields.name,
+    );
+    if (!notesToolCalls?.length) {
+      return [];
+    }
+    return notesToolCalls.map(
+      (tc) => (tc.args as z.infer<typeof plannerNotesFields.schema>).notes,
+    );
+  });
+  return plannerNotes.flat();
+}

apps/open-swe/src/tools/planner-notes.ts

@@ -0,0 +1,19 @@
+import { tool } from "@langchain/core/tools";
+import { createTakePlannerNotesFields } from "@open-swe/shared/open-swe/tools";
+
+export function createPlannerNotesTool() {
+  const plannerNotesTool = tool(
+    async (
+      _input,
+    ): Promise<{ result: string; status: "success" | "error" }> => {
+      // TODO: This should write to saved state once that feature is released in LangGraph.
+      return {
+        result: "Successfully saved notes. Thank you!",
+        status: "success",
+      };
+    },
+    createTakePlannerNotesFields(),
+  );
+
+  return plannerNotesTool;
+}

apps/open-swe/src/utils/custom-rules.ts

@@ -26,7 +26,9 @@ const ALL_TAGS = [
   TESTING_INSTRUCTIONS_CLOSE_TAG,
 ];
 
-export function parseCustomRulesFromString(contents: string): CustomRules {
+export function parseCustomRulesFromString(
+  contents: string,
+): CustomRules | undefined {
   if (ALL_TAGS.every((tag) => !contents.includes(tag))) {
     // Text file has no custom rules. Return all as general rules
     return {
@@ -78,6 +80,15 @@ export function parseCustomRulesFromString(contents: string): CustomRules {
     );
   }
 
+  if (
+    !generalRules &&
+    !repositoryStructure &&
+    !dependenciesAndInstallation &&
+    !testingInstructions
+  ) {
+    return undefined;
+  }
+
   return {
     generalRules,
     repositoryStructure,
@@ -129,20 +140,28 @@ export async function getCustomRules(
 
 export const CUSTOM_RULES_PROMPT = `<custom_rules>
 The following are custom rules provided by the user.
+{EXTRA_CONTEXT}
 {GENERAL_RULES}
 {REPOSITORY_STRUCTURE}
 {TESTING_INSTRUCTIONS}
 {DEPENDENCIES_AND_INSTALLATION}
 </custom_rules>`;
 
-export function formatCustomRulesPrompt(customRules?: CustomRules): string {
+export function formatCustomRulesPrompt(
+  customRules?: CustomRules,
+  extraContextStr?: string,
+): string {
   if (!customRules) return "";
   return CUSTOM_RULES_PROMPT.replace(
-    "{GENERAL_RULES}",
-    customRules.generalRules
-      ? `<general_rules>\n${customRules.generalRules}\n</general_rules>`
-      : "",
+    "{EXTRA_CONTEXT}",
+    extraContextStr ? extraContextStr : "",
   )
+    .replace(
+      "{GENERAL_RULES}",
+      customRules.generalRules
+        ? `<general_rules>\n${customRules.generalRules}\n</general_rules>`
+        : "",
+    )
     .replace(
       "{REPOSITORY_STRUCTURE}",
       customRules.repositoryStructure

apps/open-swe/src/utils/github/git.ts

@@ -380,6 +380,38 @@ export async function getChangedFilesStatus(
     .filter((line) => line !== "");
 }
 
+export async function stashAndClearChanges(
+  absoluteRepoDir: string,
+  sandbox: Sandbox,
+): Promise<ExecuteResponse | false> {
+  try {
+    const gitStashOutput = await sandbox.process.executeCommand(
+      "git stash && git reset --hard",
+      absoluteRepoDir,
+      undefined,
+      TIMEOUT_SEC,
+    );
+
+    if (gitStashOutput.exitCode !== 0) {
+      logger.error(`Failed to stash and clear changes`, {
+        gitStashOutput,
+      });
+    }
+    return gitStashOutput;
+  } catch (e) {
+    const errorFields = getSandboxErrorFields(e);
+    logger.error(`Failed to stash and clear changes`, {
+      ...(errorFields && { errorFields }),
+      ...(e instanceof Error && {
+        name: e.name,
+        message: e.message,
+        stack: e.stack,
+      }),
+    });
+    return errorFields ?? false;
+  }
+}
+
 export async function checkoutBranchAndCommit(
   config: GraphConfig,
   targetRepository: TargetRepository,

apps/web/src/components/gen-ui/action-step.tsx

@@ -18,6 +18,7 @@ import {
   createApplyPatchToolFields,
   createShellToolFields,
   createInstallDependenciesToolFields,
+  createTakePlannerNotesFields,
   formatRgCommand,
   RipgrepCommand,
 } from "@open-swe/shared/open-swe/tools";
@@ -33,6 +34,8 @@ const installDependenciesTool = createInstallDependenciesToolFields(dummyRepo);
 type InstallDependenciesToolArgs = z.infer<
   typeof installDependenciesTool.schema
 >;
+const plannerNotesTool = createTakePlannerNotesFields();
+type PlannerNotesToolArgs = z.infer<typeof plannerNotesTool.schema>;
 
 // Common props for all action types
 type BaseActionProps = {
@@ -73,12 +76,18 @@ type InstallDependenciesActionProps = BaseActionProps &
     errorCode?: number;
   };
 
+type PlannerNotesActionProps = BaseActionProps &
+  Partial<PlannerNotesToolArgs> & {
+    actionType: "planner_notes";
+  };
+
 export type ActionItemProps =
   | (BaseActionProps & { status: "loading" })
   | ShellActionProps
   | PatchActionProps
   | RgActionProps
-  | InstallDependenciesActionProps;
+  | InstallDependenciesActionProps
+  | PlannerNotesActionProps;
 
 export type ActionStepProps = {
   actions: ActionItemProps[];
@@ -126,6 +135,8 @@ function ActionItem(props: ActionItemProps) {
         return props.success ? "Search completed" : "Search failed";
       } else if (props.actionType === "install_dependencies") {
         return props.success ? "Dependencies installed" : "Installation failed";
+      } else if (props.actionType === "planner_notes") {
+        return props.success ? "Notes saved" : "Failed to save notes";
       }
     }
 
@@ -144,6 +155,8 @@ function ActionItem(props: ActionItemProps) {
       return !!props.output;
     } else if (props.actionType === "apply-patch") {
       return !!props.diff;
+    } else if (props.actionType === "planner_notes") {
+      return !!(props.notes && props.notes.length > 0);
     }
 
     return false;
@@ -156,7 +169,9 @@ function ActionItem(props: ActionItemProps) {
       return <Loader2 className="text-muted-foreground mr-2 size-3.5" />;
     }
 
-    if (props.actionType === "install_dependencies") {
+    if (props.actionType === "planner_notes") {
+      return <FileText className="text-muted-foreground mr-2 size-3.5" />;
+    } else if (props.actionType === "install_dependencies") {
       return <CloudDownload className="text-muted-foreground mr-2 size-3.5" />;
     } else if (props.actionType === "apply-patch") {
       return <FileCode className="text-muted-foreground mr-2 size-3.5" />;
@@ -177,6 +192,16 @@ function ActionItem(props: ActionItemProps) {
       );
     }
 
+    if (props.actionType === "planner_notes") {
+      return (
+        <div className="flex-1">
+          <span className="text-foreground/80 text-xs font-normal">
+            Planner Notes
+          </span>
+        </div>
+      );
+    }
+
     if (
       props.actionType === "shell" ||
       props.actionType === "install_dependencies"
@@ -289,6 +314,25 @@ function ActionItem(props: ActionItemProps) {
           )}
         </div>
       );
+    } else if (
+      props.actionType === "planner_notes" &&
+      props.notes &&
+      props.notes.length > 0
+    ) {
+      return (
+        <div className="bg-muted overflow-x-auto p-2 dark:bg-gray-900">
+          <ul className="list-disc pl-5 text-xs font-normal">
+            {props.notes.map((note, i) => (
+              <li
+                key={i}
+                className="text-foreground/90 mb-1 whitespace-pre-wrap"
+              >
+                {note}
+              </li>
+            ))}
+          </ul>
+        </div>
+      );
     }
 
     return null;

apps/web/src/components/thread/messages/ai.tsx

@@ -29,6 +29,7 @@ import {
   createRgToolFields,
   createOpenPrToolFields,
   createInstallDependenciesToolFields,
+  createTakePlannerNotesFields,
 } from "@open-swe/shared/open-swe/tools";
 import { z } from "zod";
 import { isAIMessageSDK, isToolMessageSDK } from "@/lib/langchain-messages";
@@ -50,6 +51,8 @@ const installDependenciesTool = createInstallDependenciesToolFields(dummyRepo);
 type InstallDependenciesToolArgs = z.infer<
   typeof installDependenciesTool.schema
 >;
+const plannerNotesTool = createTakePlannerNotesFields();
+type PlannerNotesToolArgs = z.infer<typeof plannerNotesTool.schema>;
 
 function CustomComponent({
   message,
@@ -168,6 +171,15 @@ export function mapToolMessageToActionStepProps(
       output: getContentString(message.content),
       reasoningText,
     };
+  } else if (toolCall?.name === plannerNotesTool.name) {
+    const args = toolCall.args as PlannerNotesToolArgs;
+    return {
+      actionType: "planner_notes",
+      status,
+      success,
+      notes: args.notes || [],
+      reasoningText,
+    };
   }
   return {
     status: "loading",
@@ -231,7 +243,8 @@ export function AssistantMessage({
           tc.name === shellTool.name ||
           tc.name === applyPatchTool.name ||
           tc.name === rgTool.name ||
-          tc.name === installDependenciesTool.name,
+          tc.name === installDependenciesTool.name ||
+          tc.name === plannerNotesTool.name,
       )
     : [];
 
@@ -350,6 +363,13 @@ export function AssistantMessage({
           workdir: args?.workdir || "",
           output: "",
         } as ActionItemProps;
+      } else if (toolCall.name === plannerNotesTool.name) {
+        const args = toolCall.args as PlannerNotesToolArgs;
+        return {
+          actionType: "planner_notes",
+          status: "generating",
+          notes: args?.notes || [],
+        } as ActionItemProps;
       } else {
         if (isShellTool) {
           const args = toolCall.args as ShellToolArgs;

packages/shared/src/open-swe/tools.ts

@@ -231,3 +231,22 @@ export function createOpenPrToolFields() {
     description: "Use this tool to open a pull request.",
   };
 }
+
+export function createTakePlannerNotesFields() {
+  const plannerContextNotesSchema = z.object({
+    notes: z
+      .array(z.string())
+      .describe(
+        "The high quality, concise and technical notes you deem important to save for the programmer to use when implementing the plan.",
+      ),
+  });
+
+  return {
+    name: "take_notes",
+    schema: plannerContextNotesSchema,
+    description:
+      "Use this tool to write & save technical notes on the planner context.\n" +
+      "This should be called if you come across context you think will be highly useful to the programmer when they are actually implementing the plan, and you want to ensure it's not lost.\n" +
+      "Do not duplicate any information present in the user provided 'custom rules', as we want to avoid duplicating context.",
+  };
+}