feat: Implement progress plan node to check if tasks are completed (#8)

src/index.ts

@@ -12,6 +12,7 @@ import {
   takeAction,
   rewritePlan,
   interruptPlan,
+  progressPlanStep,
 } from "./nodes/index.js";
 import { isAIMessage } from "@langchain/core/messages";
 import { pauseSandbox } from "./utils/sandbox.js";
@@ -69,6 +70,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addNode("initialize", initialize)
   .addNode("generate-action", generateAction)
   .addNode("take-action", takeAction)
+  .addNode("progress-plan-step", progressPlanStep)
   .addEdge(START, "generate-plan")
   // TODO: Update routing to work w/ new interrupt node.
   .addConditionalEdges("generate-plan", routeAfterPlan, ["interrupt-plan", END])
@@ -76,7 +78,8 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addEdge("rewrite-plan", "interrupt-plan")
   .addEdge("initialize", "generate-action")
   .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END])
-  .addEdge("take-action", "generate-action");
+  .addEdge("take-action", "progress-plan-step")
+  .addEdge("progress-plan-step", "take-action");
 
 export const graph = workflow.compile();
 graph.name = "LangGraph ReAct MCP";

src/nodes/generate-message.ts

@@ -1,16 +1,9 @@
-import { spawnSync } from "node:child_process";
-import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
+import { GraphState, GraphConfig, GraphUpdate, PlanItem } from "../types.js";
 import { loadModel } from "../utils/load-model.js";
 import { shellTool, applyPatchTool } from "../tools/index.js";
+import { formatPlanPrompt } from "../utils/plan-prompt.js";
 
-const dynamicLines: Array<string> = [];
-if (spawnSync("rg", ["--version"], { stdio: "ignore" }).status === 0) {
-  dynamicLines.push(
-    "- Always use rg instead of grep/ls -R because it is much faster and respects gitignore",
-  );
-}
-const dynamicPrefix = dynamicLines.join("\n");
-const prefix = `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 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.
 
 You can:
 - Receive user prompts, project context, and files.
@@ -18,15 +11,9 @@ You can:
 - Apply patches, run commands, and manage user approvals based on policy.
 - Work inside a sandboxed, git-backed workspace with rollback support.
 
-You work based on a plan which was generated in a previous step.
-The following tasks have already been completed:
-{COMPLETED_TASKS}
+You work based on a plan which was generated in a previous step. The plan items are as follows:
 
-The following tasks remain to be completed:
-{PENDING_TASKS}
-
-You are currently working on the following task:
-{CURRENT_TASK}
+{PLAN_PROMPT}
 
 You are an agent - please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. If you are not sure about file content or codebase structure pertaining to the user's request, use your tools to read files and gather the relevant information: do NOT guess or make up an answer.
 
@@ -61,8 +48,11 @@ 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.`;
 
-${dynamicPrefix}`;
+const formatPrompt = (plan: PlanItem[]): string => {
+  return systemPrompt.replace("{PLAN_PROMPT}", formatPlanPrompt(plan));
+};
 
 export async function generateAction(
   state: GraphState,
@@ -75,7 +65,7 @@ export async function generateAction(
   const response = await modelWithTools.invoke([
     {
       role: "system",
-      content: prefix,
+      content: formatPrompt(state.plan),
     },
     ...state.messages,
   ]);

src/nodes/index.ts

@@ -4,3 +4,4 @@ export * from "./generate-message.js";
 export * from "./take-action.js";
 export * from "./rewrite-plan.js";
 export * from "./interrupt-plan.js";
+export * from "./progress-plan-step.js";

src/nodes/progress-plan-step.ts

@@ -0,0 +1,85 @@
+import { z } from "zod";
+import { GraphConfig, GraphState, GraphUpdate, PlanItem } from "../types.js";
+import { loadModel } from "../utils/load-model.js";
+import { formatPlanPrompt } from "../utils/plan-prompt.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 your workflow, you generate a plan, then act on said plan. It may take many actions to complete a single step, or a single action to complete the step.
+
+Here is the plan:
+
+{PLAN_PROMPT}
+
+In this task, you will analyze the plan, the tasks you've completed, the tasks which are left, and the current task you just took an action on. In addition to this, you're also provided the full conversation history between you and the user. All of the messages in this conversation are from the previous steps/actions you've taken, and any user input.
+
+Take all of this information, and determine whether or not you have completed this task in the plan. To do this, you will call the \`confirm_task_completion\` tool.`;
+
+const confirmTaskCompletionToolSchema = z.object({
+  current_task_completed: z
+    .boolean()
+    .describe("Whether or not the current task has been completed."),
+});
+
+const confirmTaskCompletionTool = {
+  name: "confirm_task_completion",
+  description: "Whether or not the current task has been completed.",
+  schema: confirmTaskCompletionToolSchema,
+};
+
+const formatPrompt = (plan: PlanItem[]): string => {
+  return systemPrompt.replace("{PLAN_PROMPT}", formatPlanPrompt(plan));
+};
+
+export async function progressPlanStep(
+  state: GraphState,
+  config: GraphConfig,
+): Promise<GraphUpdate> {
+  const model = await loadModel(config);
+  const modelWithTools = model.bindTools([confirmTaskCompletionTool], {
+    tool_choice: confirmTaskCompletionTool.name,
+  });
+
+  const response = await modelWithTools.invoke([
+    {
+      role: "system",
+      content: formatPrompt(state.plan),
+    },
+    ...state.messages,
+  ]);
+  const toolCall = response.tool_calls?.[0];
+
+  if (!toolCall) {
+    throw new Error("Failed to check plan.");
+  }
+
+  const isCompleted = (
+    toolCall.args as z.infer<typeof confirmTaskCompletionToolSchema>
+  ).current_task_completed;
+
+  if (!isCompleted) {
+    // Not completed, no changes need to be made
+    return {};
+  }
+
+  const remainingTask = state.plan.find((p) => !p.completed);
+  if (!remainingTask) {
+    // No remaining tasks, end the process
+    console.log(
+      "Found no remaining tasks in the plan during the check plan step.",
+    );
+    return {};
+  }
+
+  return {
+    plan: state.plan.map((p) => {
+      if (p.index === remainingTask.index) {
+        return {
+          ...p,
+          completed: true,
+        };
+      }
+      return p;
+    }),
+  };
+}

src/types.ts

@@ -6,9 +6,19 @@ import {
   MessagesAnnotation,
 } from "@langchain/langgraph";
 
-type PlanItem = {
-  id: string;
+export type PlanItem = {
+  /**
+   * The index of the plan item. This is the order in which
+   * it should be executed.
+   */
+  index: number;
+  /**
+   * The actual task to perform.
+   */
   plan: string;
+  /**
+   * Whether or not the plan item has been completed.
+   */
   completed: boolean;
 };
 

src/utils/plan-prompt.ts

@@ -0,0 +1,26 @@
+import { PlanItem } from "../types.js";
+
+export const PLAN_PROMPT = `## Completed Tasks
+{COMPLETED_TASKS}
+
+## Remaining Tasks
+{REMAINING_TASKS}
+
+## Current Task
+{CURRENT_TASK}`;
+
+export function formatPlanPrompt(plan: PlanItem[]): string {
+  const completedTasks = plan.filter((p) => p.completed);
+  const remainingTasks = plan.filter((p) => !p.completed);
+  const currentTask = remainingTasks.sort((a, b) => a.index - b.index)[0];
+
+  return PLAN_PROMPT.replace(
+    "{COMPLETED_TASKS}",
+    completedTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"),
+  )
+    .replace(
+      "{REMAINING_TASKS}",
+      remainingTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"),
+    )
+    .replace("{CURRENT_TASK}", currentTask.plan);
+}