Update graph.py

graph.py

@@ -893,15 +893,26 @@ def run_disclaimer_agent(state: AgentState):
     return {"draftResponse": final_response, "execution_path": path, "status_update": reason}
 
 def should_continue(state: AgentState):
-    # NEW: Check for budget excess first
+    # Budget check first
     if state.get("budget_exceeded"):
         return "disclaimer_agent"
+    try:
+        rework = int(state.get("rework_cycles", 0))
+        max_loops_allowed = int(state.get("max_loops", 0))
+    except Exception:
+        rework = state.get("rework_cycles", 0) or 0
+        max_loops_allowed = state.get("max_loops", 0) or 0
+
+    # If approved -> archive
     if state.get("approved"):
         return "archivist_agent"
-    if ensure_int(state, "rework_cycles", 0) > ensure_int(state, "max_loops", 0):
+    # If we have exceeded allowed reworks -> disclaimer
+    if rework > max_loops_allowed:
         return "disclaimer_agent"
+    # Default: return pm_agent so planner will create next plan
     return "pm_agent"
 
+
 def should_run_experiment(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"