feat: Set max num of times reviewer can run (#418)

* feat: Set max num of times reviewer can run

* cr

apps/open-swe/src/graphs/programmer/index.ts

@@ -1,6 +1,7 @@
-import { END, Send, START, StateGraph } from "@langchain/langgraph";
+import { Command, END, Send, START, StateGraph } from "@langchain/langgraph";
 import {
   GraphAnnotation,
+  GraphConfig,
   GraphConfiguration,
   GraphState,
 } from "@open-swe/shared/open-swe/types";
@@ -39,12 +40,12 @@ function lastMessagesMissingToolCalls(
  * Otherwise, it ends the process.
  *
  * @param {GraphState} state - The current graph state.
- * @returns {"reviewer-subgraph" | "take-action" | "request-help" | "generate-action" | Send} The next node to execute, or END if the process should stop.
+ * @returns {"route-to-review-or-conclusion" | "take-action" | "request-help" | "generate-action" | Send} The next node to execute, or END if the process should stop.
  */
 function routeGeneratedAction(
   state: GraphState,
 ):
-  | "reviewer-subgraph"
+  | "route-to-review-or-conclusion"
   | "take-action"
   | "request-help"
   | "generate-action"
@@ -82,7 +83,7 @@ function routeGeneratedAction(
   }
 
   // No tool calls, route to reviewer subgraph
-  return "reviewer-subgraph";
+  return "route-to-review-or-conclusion";
 }
 
 /**
@@ -101,6 +102,22 @@ function routeGenerateActionsOrEnd(
   return "generate-action";
 }
 
+function routeToReviewOrConclusion(
+  state: GraphState,
+  config: GraphConfig,
+): Command {
+  const maxAllowedReviews = config.configurable?.maxReviewCount ?? 3;
+  if (state.reviewsCount >= maxAllowedReviews) {
+    return new Command({
+      goto: "generate-conclusion",
+    });
+  }
+
+  return new Command({
+    goto: "reviewer-subgraph",
+  });
+}
+
 const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addNode("initialize", initializeSandbox)
   .addNode("generate-action", generateAction)
@@ -109,12 +126,19 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   })
   .addNode("update-plan", updatePlan)
   .addNode("progress-plan-step", progressPlanStep, {
-    ends: ["summarize-history", "generate-action", "reviewer-subgraph"],
+    ends: [
+      "summarize-history",
+      "generate-action",
+      "route-to-review-or-conclusion",
+    ],
   })
   .addNode("generate-conclusion", generateConclusion)
   .addNode("request-help", requestHelp, {
     ends: ["generate-action", END],
   })
+  .addNode("route-to-review-or-conclusion", routeToReviewOrConclusion, {
+    ends: ["generate-conclusion", "route-to-review-or-conclusion"],
+  })
   .addNode("reviewer-subgraph", reviewerGraph)
   .addNode("open-pr", openPullRequest)
   .addNode("diagnose-error", diagnoseError)
@@ -124,7 +148,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addConditionalEdges("generate-action", routeGeneratedAction, [
     "take-action",
     "request-help",
-    "reviewer-subgraph",
+    "route-to-review-or-conclusion",
     "update-plan",
     "generate-action",
   ])

apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts

@@ -187,7 +187,7 @@ Once you've determined the status of the current task, call either the \`mark_ta
       taskPlan: updatedPlanTasks,
     };
     return new Command({
-      goto: "reviewer-subgraph",
+      goto: "route-to-review-or-conclusion",
       update: commandUpdate,
     });
   }

apps/open-swe/src/graphs/reviewer/nodes/final-review.ts

@@ -157,5 +157,6 @@ export async function finalReview(
     taskPlan: updatedTaskPlan,
     messages: messagesUpdate,
     internalMessages: messagesUpdate,
+    reviewsCount: (state.reviewsCount || 0) + 1,
   };
 }

packages/shared/src/open-swe/reviewer/types.ts

@@ -100,6 +100,16 @@ export const ReviewerGraphStateObj = MessagesZodState.extend({
       fn: (_state, update) => update,
     },
   }),
+  /**
+   * The number of times the reviewer subgraph has been executed.
+   */
+  reviewsCount: withLangGraph(z.custom<number>(), {
+    reducer: {
+      schema: z.custom<number>(),
+      fn: (_state, update) => update,
+    },
+    default: () => 0,
+  }),
 });
 
 export type ReviewerGraphState = z.infer<typeof ReviewerGraphStateObj>;

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

@@ -241,14 +241,14 @@ export const GraphAnnotation = MessagesZodState.extend({
     },
   }),
   /**
-   * The review generated by the reviewer subgraph
+   * The number of times the reviewer subgraph has been executed.
    */
-  review: withLangGraph(z.custom<string>(), {
+  reviewsCount: withLangGraph(z.custom<number>(), {
     reducer: {
-      schema: z.custom<string>(),
+      schema: z.custom<number>(),
       fn: (_state, update) => update,
     },
-    default: () => "",
+    default: () => 0,
   }),
 
   // ---NOT USED---
@@ -292,6 +292,11 @@ export const GraphConfigurationMetadata: {
         "Maximum number of review actions allowed during the review phase. An action consists of a tool call.",
     },
   },
+  maxReviewCount: {
+    x_open_swe_ui_config: {
+      type: "hidden",
+    },
+  },
   programmerModelName: {
     x_open_swe_ui_config: {
       type: "select",
@@ -514,6 +519,12 @@ export const GraphConfiguration = z.object({
   mcpServers: withLangGraph(z.string().optional(), {
     metadata: GraphConfigurationMetadata.mcpServers,
   }),
+  /**
+   * The maxium number of times the reviewer subgraph can be executed.
+   */
+  maxReviewCount: withLangGraph(z.number().optional(), {
+    metadata: GraphConfigurationMetadata.maxReviewCount,
+  }),
 });
 
 export type GraphConfig = LangGraphRunnableConfig<