fix: correct _interrupted state management in async API

- _setup_run() now resets _interrupted=False — the canonical place for
  per-run resets, shared by run()/arun()/run_stream()
- Removed try/except from arun(): _setup_run() already handles the reset
- Removed try/except from aresume()/areject(): clearing _interrupted on a
  stream crash was wrong — the LangGraph checkpointer still holds the
  pending interrupt, so _interrupted=True is correct and lets the user retry
- Renamed _post_stream_result() param final_state -> last_values to avoid
  confusion with LangGraph's StateSnapshot
- Deduplicated _interrupted_agent() test helper into make_interrupted_agent()
- Updated test comment for test_arun_propagates_stream_error
- Added test_aresume_allows_retry_after_stream_error
- Added test_areject_allows_retry_after_stream_error

Author: BinglanLi

BaseAgent/base_agent.py

@@ -1179,11 +1179,14 @@ def _select_skills_for_prompt(self, prompt: str) -> None:
     def _setup_run(self, prompt: str, thread_id: str | None) -> tuple[dict, dict]:
         """Reset per-run counters and build the LangGraph inputs/config dicts.
 
-        Shared by ``run()``, ``arun()``, and ``run_stream()``.  Does *not*
-        reset ``self.log`` — callers that maintain a log do so explicitly.
+        Shared by ``run()``, ``arun()``, and ``run_stream()``.  Resets
+        ``_interrupted`` because any prior interrupt is irrelevant to a new run.
+        Does *not* reset ``self.log`` — ``resume()``/``reject()`` append to the
+        existing log rather than starting fresh, so they manage it themselves.
         """
         self.critic_count = 0
         self.user_task = prompt
+        self._interrupted = False
         # Reset per-run counters; mark the start index in usage_metrics for
         # per-run cost budget checks (existing metrics are preserved for
         # lifetime cost tracking via agent._usage_metrics)
@@ -1204,15 +1207,15 @@ def _setup_run(self, prompt: str, thread_id: str | None) -> tuple[dict, dict]:
         self._run_config = config
         return inputs, config
 
-    def _post_stream_result(self, graph_state, final_state, last_msg):
+    def _post_stream_result(self, graph_state, last_values, last_msg):
         """Set conversation state, detect interrupt, and return ``(log, content)``.
 
         Shared by ``run()``, ``arun()``, ``resume()``, ``aresume()``,
         ``reject()``, and ``areject()``.  The caller fetches ``graph_state``
         via ``get_state()`` (sync) or ``await aget_state()`` (async) and
         passes it in, keeping this helper synchronous.
         """
-        self._conversation_state = final_state
+        self._conversation_state = last_values
         for task in (graph_state.tasks or []):
             if hasattr(task, "interrupts") and task.interrupts:
                 self._interrupted = True
@@ -1242,13 +1245,13 @@ def run(self, prompt: str, thread_id: str | None = None):
         inputs, config = self._setup_run(prompt, thread_id)
         self.log = []
         last_msg = None
-        final_state = None
+        last_values = None
         for state in self.app.stream(inputs, stream_mode="values", config=config):
             last_msg = state["input"][-1]
             self.log.append(pretty_print(last_msg))
-            final_state = state
+            last_values = state
         graph_state = self.app.get_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     def resume(self):
         """Approve the pending code block and continue execution.
@@ -1268,13 +1271,13 @@ def resume(self):
 
         config = self._run_config
         last_msg = None
-        final_state = None
+        last_values = None
         for state in self.app.stream(Command(resume=True), stream_mode="values", config=config):
             last_msg = state["input"][-1]
             self.log.append(pretty_print(last_msg))
-            final_state = state
+            last_values = state
         graph_state = self.app.get_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     def reject(self, feedback: str = "User rejected this code. Try a different approach."):
         """Reject the pending code block and provide feedback to the agent.
@@ -1298,17 +1301,17 @@ def reject(self, feedback: str = "User rejected this code. Try a different appro
 
         config = self._run_config
         last_msg = None
-        final_state = None
+        last_values = None
         for state in self.app.stream(
             Command(resume={"approved": False, "feedback": feedback}),
             stream_mode="values",
             config=config,
         ):
             last_msg = state["input"][-1]
             self.log.append(pretty_print(last_msg))
-            final_state = state
+            last_values = state
         graph_state = self.app.get_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     # ------------------------------------------------------------------
     # Async public API
@@ -1332,17 +1335,13 @@ async def arun(self, prompt: str, thread_id: str | None = None):
         inputs, config = self._setup_run(prompt, thread_id)
         self.log = []
         last_msg = None
-        final_state = None
-        try:
-            async for state in self.app.astream(inputs, stream_mode="values", config=config):
-                last_msg = state["input"][-1]
-                self.log.append(pretty_print(last_msg))
-                final_state = state
-        except Exception:
-            self._interrupted = False  # prevent stale interrupt state on error
-            raise
+        last_values = None
+        async for state in self.app.astream(inputs, stream_mode="values", config=config):
+            last_msg = state["input"][-1]
+            self.log.append(pretty_print(last_msg))
+            last_values = state
         graph_state = await self.app.aget_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     async def aresume(self):
         """Async equivalent of :meth:`resume`.
@@ -1358,19 +1357,15 @@ async def aresume(self):
 
         config = self._run_config
         last_msg = None
-        final_state = None
-        try:
-            async for state in self.app.astream(
-                Command(resume=True), stream_mode="values", config=config
-            ):
-                last_msg = state["input"][-1]
-                self.log.append(pretty_print(last_msg))
-                final_state = state
-        except Exception:
-            self._interrupted = False
-            raise
+        last_values = None
+        async for state in self.app.astream(
+            Command(resume=True), stream_mode="values", config=config
+        ):
+            last_msg = state["input"][-1]
+            self.log.append(pretty_print(last_msg))
+            last_values = state
         graph_state = await self.app.aget_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     async def areject(self, feedback: str = "User rejected this code. Try a different approach."):
         """Async equivalent of :meth:`reject`.
@@ -1389,21 +1384,17 @@ async def areject(self, feedback: str = "User rejected this code. Try a differen
 
         config = self._run_config
         last_msg = None
-        final_state = None
-        try:
-            async for state in self.app.astream(
-                Command(resume={"approved": False, "feedback": feedback}),
-                stream_mode="values",
-                config=config,
-            ):
-                last_msg = state["input"][-1]
-                self.log.append(pretty_print(last_msg))
-                final_state = state
-        except Exception:
-            self._interrupted = False
-            raise
+        last_values = None
+        async for state in self.app.astream(
+            Command(resume={"approved": False, "feedback": feedback}),
+            stream_mode="values",
+            config=config,
+        ):
+            last_msg = state["input"][-1]
+            self.log.append(pretty_print(last_msg))
+            last_values = state
         graph_state = await self.app.aget_state(config)
-        return self._post_stream_result(graph_state, final_state, last_msg)
+        return self._post_stream_result(graph_state, last_values, last_msg)
 
     async def run_stream(
         self,

BaseAgent/tests/test_async_api.py

@@ -36,6 +36,18 @@ def make_agent():
     return agent
 
 
+def make_interrupted_agent():
+    """Return a BaseAgent in an interrupted state, ready for aresume()/areject()."""
+    agent = make_agent()
+    agent.app = MagicMock()
+    agent.app.astream = MagicMock(return_value=_async_iter([_make_state()]))
+    agent.app.aget_state = AsyncMock(return_value=_no_interrupt_graph_state())
+    agent._interrupted = True
+    agent._run_config = {"configurable": {"thread_id": "t1"}}
+    agent.log = []  # normally set by arun(); set here for isolation
+    return agent
+
+
 def _make_state(content: str = "<solution>done</solution>"):
     """One LangGraph state snapshot with a single AIMessage."""
     return {"input": [AIMessage(content=content)], "next_step": "end"}
@@ -144,7 +156,8 @@ async def _failing_stream(*args, **kwargs):
         with pytest.raises(RuntimeError, match="stream exploded"):
             asyncio.run(agent.arun("test task"))
 
-        # stale interrupt must be cleared on error
+        # _setup_run() cleared _interrupted before the stream started, so it
+        # stays False even though the stream raised before _post_stream_result ran
         assert agent._interrupted is False
 
     def test_arun_calls_astream_not_stream(self):
@@ -167,18 +180,8 @@ def test_arun_calls_astream_not_stream(self):
 
 @pytest.mark.unit
 class TestAresume:
-    def _interrupted_agent(self):
-        agent = make_agent()
-        agent.app = MagicMock()
-        agent.app.astream = MagicMock(return_value=_async_iter([_make_state()]))
-        agent.app.aget_state = AsyncMock(return_value=_no_interrupt_graph_state())
-        agent._interrupted = True
-        agent._run_config = {"configurable": {"thread_id": "t1"}}
-        agent.log = []  # normally set by arun(); set here for isolation
-        return agent
-
     def test_aresume_returns_result(self):
-        agent = self._interrupted_agent()
+        agent = make_interrupted_agent()
 
         log, content = asyncio.run(agent.aresume())
 
@@ -194,14 +197,33 @@ def test_aresume_raises_if_not_interrupted(self):
             asyncio.run(agent.aresume())
 
     def test_aresume_calls_astream_not_stream(self):
-        agent = self._interrupted_agent()
+        agent = make_interrupted_agent()
         agent.app.stream = MagicMock()
 
         asyncio.run(agent.aresume())
 
         agent.app.astream.assert_called_once()
         agent.app.stream.assert_not_called()
 
+    def test_aresume_allows_retry_after_stream_error(self):
+        """A stream crash in aresume() must not clear _interrupted.
+
+        The LangGraph checkpointer still holds the pending interrupt, so the
+        user should be able to call aresume() again to retry.
+        """
+        agent = make_interrupted_agent()
+
+        async def _failing_stream(*args, **kwargs):
+            yield _make_state()
+            raise RuntimeError("transient error")
+
+        agent.app.astream = MagicMock(return_value=_failing_stream())
+
+        with pytest.raises(RuntimeError, match="transient error"):
+            asyncio.run(agent.aresume())
+
+        assert agent._interrupted is True  # still retryable
+
 
 # ---------------------------------------------------------------------------
 # areject() tests
@@ -210,18 +232,8 @@ def test_aresume_calls_astream_not_stream(self):
 
 @pytest.mark.unit
 class TestAreject:
-    def _interrupted_agent(self):
-        agent = make_agent()
-        agent.app = MagicMock()
-        agent.app.astream = MagicMock(return_value=_async_iter([_make_state()]))
-        agent.app.aget_state = AsyncMock(return_value=_no_interrupt_graph_state())
-        agent._interrupted = True
-        agent._run_config = {"configurable": {"thread_id": "t1"}}
-        agent.log = []  # normally set by arun(); set here for isolation
-        return agent
-
     def test_areject_returns_result(self):
-        agent = self._interrupted_agent()
+        agent = make_interrupted_agent()
 
         log, content = asyncio.run(agent.areject("try again"))
 
@@ -236,7 +248,7 @@ def test_areject_raises_if_not_interrupted(self):
             asyncio.run(agent.areject())
 
     def test_areject_passes_correct_command(self):
-        agent = self._interrupted_agent()
+        agent = make_interrupted_agent()
 
         asyncio.run(agent.areject("do it differently"))
 
@@ -246,14 +258,33 @@ def test_areject_passes_correct_command(self):
         assert command.resume == {"approved": False, "feedback": "do it differently"}
 
     def test_areject_calls_astream_not_stream(self):
-        agent = self._interrupted_agent()
+        agent = make_interrupted_agent()
         agent.app.stream = MagicMock()
 
         asyncio.run(agent.areject("nope"))
 
         agent.app.astream.assert_called_once()
         agent.app.stream.assert_not_called()
 
+    def test_areject_allows_retry_after_stream_error(self):
+        """A stream crash in areject() must not clear _interrupted.
+
+        The LangGraph checkpointer still holds the pending interrupt, so the
+        user should be able to call areject() again to retry.
+        """
+        agent = make_interrupted_agent()
+
+        async def _failing_stream(*args, **kwargs):
+            yield _make_state()
+            raise RuntimeError("transient error")
+
+        agent.app.astream = MagicMock(return_value=_failing_stream())
+
+        with pytest.raises(RuntimeError, match="transient error"):
+            asyncio.run(agent.areject("nope"))
+
+        assert agent._interrupted is True  # still retryable
+
 
 # ---------------------------------------------------------------------------
 # Async generator helper