From 004f9fd9c90d0cb3398bfbda47b724c96b6a9feb Mon Sep 17 00:00:00 2001
From: Yunxiao Xu <xuyunxiao2001@gmail.com>
Date: Wed, 11 Feb 2026 16:50:48 -0800
Subject: [PATCH] refactor(graph): Add iteration limit to prevent infinite
 loops in code correction

---
 src/ea_chatbot/graph/nodes/error_corrector.py | 6 +++++-
 src/ea_chatbot/graph/nodes/query_analyzer.py  | 3 ++-
 src/ea_chatbot/graph/state.py                 | 3 +++
 src/ea_chatbot/graph/workflow.py              | 5 +++++
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/ea_chatbot/graph/nodes/error_corrector.py b/src/ea_chatbot/graph/nodes/error_corrector.py
index 165119d..a2aed0b 100644
--- a/src/ea_chatbot/graph/nodes/error_corrector.py
+++ b/src/ea_chatbot/graph/nodes/error_corrector.py
@@ -31,9 +31,13 @@ def error_corrector_node(state: AgentState) -> dict:
     try:
         response = structured_llm.invoke(messages)
         logger.info("[bold green]Correction generated.[/bold green]")
+        
+        current_iterations = state.get("iterations", 0)
+        
         return {
             "code": response.parsed_code,
-            "error": None # Clear error after fix attempt
+            "error": None, # Clear error after fix attempt
+            "iterations": current_iterations + 1
         }
     except Exception as e:
         logger.error(f"Failed to correct code: {str(e)}")
diff --git a/src/ea_chatbot/graph/nodes/query_analyzer.py b/src/ea_chatbot/graph/nodes/query_analyzer.py
index 4b511f8..349bcc5 100644
--- a/src/ea_chatbot/graph/nodes/query_analyzer.py
+++ b/src/ea_chatbot/graph/nodes/query_analyzer.py
@@ -66,7 +66,8 @@ def query_analyzer_node(state: AgentState) -> dict:
     
     return {
         "analysis": analysis_dict,
-        "next_action": next_action
+        "next_action": next_action,
+        "iterations": 0
     }
         
     
\ No newline at end of file
diff --git a/src/ea_chatbot/graph/state.py b/src/ea_chatbot/graph/state.py
index a5659c7..65f2863 100644
--- a/src/ea_chatbot/graph/state.py
+++ b/src/ea_chatbot/graph/state.py
@@ -31,3 +31,6 @@ class AgentState(AS):
     
     # Routing hint: "clarify", "plan", "research", "end"
     next_action: str
+
+    # Number of execution attempts
+    iterations: int
diff --git a/src/ea_chatbot/graph/workflow.py b/src/ea_chatbot/graph/workflow.py
index 64db3dd..a312638 100644
--- a/src/ea_chatbot/graph/workflow.py
+++ b/src/ea_chatbot/graph/workflow.py
@@ -10,6 +10,8 @@ from ea_chatbot.graph.nodes.researcher import researcher_node
 from ea_chatbot.graph.nodes.clarification import clarification_node
 from ea_chatbot.graph.nodes.summarize_conversation import summarize_conversation_node
 
+MAX_ITERATIONS = 3
+
 def router(state: AgentState) -> str:
     """Route to the next node based on the analysis."""
     next_action = state.get("next_action")
@@ -59,6 +61,9 @@ def create_workflow():
     # Executor routing
     def executor_router(state: AgentState) -> str:
         if state.get("error"):
+            # Check for iteration limit to prevent infinite loops
+            if state.get("iterations", 0) >= MAX_ITERATIONS:
+                return "summarizer"
             return "error_corrector"
         return "summarizer"