feat: implement accept and dont ask again feature

agent/open-code/index.ts

@@ -1,14 +1,45 @@
-import { END, START, StateGraph } from "@langchain/langgraph";
-import { OpenCodeAnnotation } from "./types";
+import {
+  END,
+  LangGraphRunnableConfig,
+  START,
+  StateGraph,
+} from "@langchain/langgraph";
+import { OpenCodeAnnotation, OpenCodeState } from "./types";
 import { planner } from "./nodes/planner";
-import { executor } from "./nodes/executor";
+import {
+  executor,
+  SUCCESSFULLY_COMPLETED_STEPS_CONTENT,
+} from "./nodes/executor";
+import { AIMessage } from "@langchain/langgraph-sdk";
+
+function conditionallyEnd(
+  state: OpenCodeState,
+  config: LangGraphRunnableConfig,
+): typeof END | "planner" {
+  const fullWriteAccess = !!config.configurable?.permissions?.full_write_access;
+  const lastAiMessage = state.messages.findLast(
+    (m) => m.getType() === "ai",
+  ) as unknown as AIMessage;
+
+  // If the user did not grant full write access, or the last AI message is the success message, end
+  // otherwise, loop back to the start.
+  if (
+    (typeof lastAiMessage.content === "string" &&
+      lastAiMessage.content === SUCCESSFULLY_COMPLETED_STEPS_CONTENT) ||
+    !fullWriteAccess
+  ) {
+    return END;
+  }
+
+  return "planner";
+}
 
 const workflow = new StateGraph(OpenCodeAnnotation)
   .addNode("planner", planner)
   .addNode("executor", executor)
   .addEdge(START, "planner")
   .addEdge("planner", "executor")
-  .addEdge("executor", END);
+  .addConditionalEdges("executor", conditionallyEnd, ["planner", END]);
 
 export const graph = workflow.compile();
 graph.name = "Open Code Graph";

agent/open-code/nodes/executor.ts

@@ -6,6 +6,9 @@ import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import ComponentMap from "../../uis";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 
+export const SUCCESSFULLY_COMPLETED_STEPS_CONTENT =
+  "Successfully completed all the steps in the plan. Please let me know if you need anything else!";
+
 export async function executor(
   state: OpenCodeState,
   config: LangGraphRunnableConfig,
@@ -24,21 +27,24 @@ export async function executor(
   const nextPlanItem = planToolCallArgs?.remainingPlans?.[0] as
     | string
     | undefined;
-  const numOfExecutedPlanItems = planToolCallArgs?.executedPlans?.length ?? 0;
+  const numSeenPlans =
+    [
+      ...(planToolCallArgs?.executedPlans ?? []),
+      ...(planToolCallArgs?.rejectedPlans ?? []),
+    ]?.length ?? 0;
 
   if (!nextPlanItem) {
     // All plans have been executed
     const successfullyFinishedMsg: AIMessage = {
       type: "ai",
       id: uuidv4(),
-      content:
-        "Successfully completed all the steps in the plan. Please let me know if you need anything else!",
+      content: SUCCESSFULLY_COMPLETED_STEPS_CONTENT,
     };
     return { messages: [successfullyFinishedMsg] };
   }
 
   let updateFileContents = "";
-  switch (numOfExecutedPlanItems) {
+  switch (numSeenPlans) {
     case 0:
       updateFileContents = await fs.readFile(
         "agent/open-code/nodes/plan-code/step-1.txt",
@@ -101,10 +107,13 @@ export async function executor(
     ],
   };
 
+  const fullWriteAccess = !!config.configurable?.permissions?.full_write_access;
+
   const msg = ui.create("proposed-change", {
     toolCallId,
     change: updateFileContents,
     planItem: nextPlanItem,
+    fullWriteAccess,
   });
   msg.additional_kwargs["message_id"] = aiMessage.id;
 

agent/open-code/nodes/planner.ts

@@ -28,28 +28,46 @@ export async function planner(
         (tc) => tc.name === "update_file",
       ),
   ) as AIMessage | undefined;
+  const lastUpdateToolCallResponse = state.messages.findLast(
+    (m) =>
+      m.getType() === "tool" &&
+      (m as unknown as ToolMessage).tool_call_id ===
+        lastUpdateCodeToolCall?.tool_calls?.[0]?.id,
+  ) as ToolMessage | undefined;
   const lastPlanToolCall = state.messages.findLast(
     (m) =>
       m.getType() === "ai" &&
       (m as unknown as AIMessage).tool_calls?.some((tc) => tc.name === "plan"),
   ) as AIMessage | undefined;
 
+  const wasPlanRejected = (
+    lastUpdateToolCallResponse?.content as string | undefined
+  )
+    ?.toLowerCase()
+    .includes("rejected");
+
   const planToolCallArgs = lastPlanToolCall?.tool_calls?.[0]?.args as Record<
     string,
     any
   >;
   const executedPlans: string[] = planToolCallArgs?.executedPlans ?? [];
+  const rejectedPlans: string[] = planToolCallArgs?.rejectedPlans ?? [];
   let remainingPlans: string[] = planToolCallArgs?.remainingPlans ?? PLAN;
 
-  const executedPlanItem = lastUpdateCodeToolCall?.tool_calls?.[0]?.args
+  const proposedChangePlanItem = lastUpdateCodeToolCall?.tool_calls?.[0]?.args
     ?.executed_plan_item as string | undefined;
-  if (executedPlanItem) {
-    executedPlans.push(executedPlanItem);
-    remainingPlans = remainingPlans.filter((p) => p !== executedPlanItem);
+  if (proposedChangePlanItem) {
+    if (wasPlanRejected) {
+      rejectedPlans.push(proposedChangePlanItem);
+    } else {
+      executedPlans.push(proposedChangePlanItem);
+    }
+
+    remainingPlans = remainingPlans.filter((p) => p !== proposedChangePlanItem);
   }
 
-  const content = executedPlanItem
-    ? `I've updated the plan list based on the executed plans.`
+  const content = proposedChangePlanItem
+    ? `I've updated the plan list based on the last proposed change.`
     : `I've come up with a detailed plan for building the todo app.`;
 
   const toolCallId = uuidv4();
@@ -62,6 +80,7 @@ export async function planner(
         name: "plan",
         args: {
           executedPlans,
+          rejectedPlans,
           remainingPlans,
         },
         id: toolCallId,
@@ -73,6 +92,7 @@ export async function planner(
   const msg = ui.create("code-plan", {
     toolCallId,
     executedPlans,
+    rejectedPlans,
     remainingPlans,
   });
   msg.additional_kwargs["message_id"] = aiMessage.id;