Merge pull request #38 from XSpoonAi/fix/agent_loop

Fix: reset per-request think wrapper to stop stale task prompts leaki…

Author: veithly

spoon_bot/agent/loop.py

@@ -477,6 +477,7 @@ async def initialize(self) -> None:
         # and can cause the model to summarize instead of continue).
         self._agent.next_step_prompt = self.DEFAULT_NEXT_STEP_PROMPT
         self._agent._custom_next_step_prompt = True
+        self._agent._spoon_bot_base_think = self._agent.think
 
         # Build dynamic-tools prompt now that all tools (native + skill) are registered
         inactive_tools = self.tools.get_inactive_tools()
@@ -1036,6 +1037,7 @@ async def process(
                     raise
         finally:
             # Always ensure agent is back in IDLE state after processing
+            self._restore_agent_think()
             if hasattr(self._agent, 'state') and self._agent.state != AgentState.IDLE:
                 logger.warning(
                     f"Post-run cleanup: resetting agent from {self._agent.state} to IDLE"
@@ -1309,7 +1311,12 @@ def _install_anti_loop_tracker(self, base_prompt: str) -> None:
             return
 
         agent_loop = self
-        original_think = agent.think
+        original_think = getattr(agent, "_spoon_bot_base_think", None)
+        if original_think is None:
+            original_think = agent.think
+            setattr(agent, "_spoon_bot_base_think", original_think)
+        else:
+            agent.think = original_think
         call_tracker: Counter = Counter()
         detail_tracker: Counter = Counter()
         read_files: set = set()
@@ -1488,6 +1495,15 @@ def _log_tool_calls():
 
         agent.think = _tracked_think
 
+    def _restore_agent_think(self) -> None:
+        """Restore the agent's base think() implementation after a request."""
+        agent = self._agent
+        if agent is None:
+            return
+        original_think = getattr(agent, "_spoon_bot_base_think", None)
+        if original_think is not None:
+            agent.think = original_think
+
     def _filter_execution_steps(self, content: str) -> str:
         """
         Filter out technical execution steps from agent output.
@@ -1578,6 +1594,9 @@ async def stream(
             logger.warning(f"Failed to load memory context: {e}")
 
         full_content = ""
+        stream_completed = False
+        stream_cancelled = False
+        bg_task: asyncio.Task[None] | None = None
 
         # Trim and inject persisted history into runtime memory
         await self._prepare_request_context()
@@ -1658,8 +1677,9 @@ async def _run_and_signal() -> None:
                 except asyncio.TimeoutError:
                     continue
                 except asyncio.CancelledError:
-                    logger.warning("Streaming cancelled")
-                    break
+                    stream_cancelled = True
+                    logger.warning("Streaming cancelled while waiting for output")
+                    raise
                 except Exception as e:
                     logger.warning(f"Queue get error: {type(e).__name__}: {e}")
                     continue
@@ -1759,15 +1779,29 @@ async def _run_and_signal() -> None:
                 }
 
             # Emit done
+            stream_completed = True
             yield {"type": "done", "delta": "", "metadata": {"content": full_content}}
 
+        except asyncio.CancelledError:
+            stream_cancelled = True
+            logger.warning("Streaming cancelled")
+            raise
         except Exception as e:
             logger.error(f"Streaming error: {e}")
+            stream_completed = True
             yield {"type": "error", "delta": str(e), "metadata": {"error": str(e)}}
             yield {"type": "done", "delta": "", "metadata": {"error": str(e)}}
+        finally:
+            self._restore_agent_think()
+            if bg_task is not None and not bg_task.done():
+                bg_task.cancel()
+                try:
+                    await asyncio.wait_for(bg_task, timeout=5.0)
+                except (asyncio.CancelledError, asyncio.TimeoutError, Exception):
+                    pass
 
         # Save to session only if we got actual content
-        if full_content:
+        if full_content and stream_completed and not stream_cancelled:
             try:
                 self._session.add_message("user", message)
                 self._session.add_message("assistant", full_content)
@@ -1834,6 +1868,8 @@ async def process_with_thinking(
         except Exception as e:
             logger.error(f"Agent processing error: {e}")
             raise
+        finally:
+            self._restore_agent_think()
 
         # Save to session
         try:

tests/test_streaming_thinking.py

@@ -576,6 +576,89 @@ async def mock_stream(message, **kwargs):
         agent._session.add_message.assert_any_call("assistant", "hello")
         agent.sessions.save.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_stream_close_cancels_background_run_and_skips_session_save(self):
+        """Closing the stream should stop the background run and avoid persisting stale output."""
+        from spoon_bot.agent.loop import AgentLoop
+
+        run_cancelled = asyncio.Event()
+
+        async def mock_run(request):
+            await agent._agent.output_queue.put({"content": "hello"})
+            try:
+                await asyncio.Future()
+            except asyncio.CancelledError:
+                run_cancelled.set()
+                raise
+
+        agent = MagicMock(spec=AgentLoop)
+        agent._initialized = True
+        agent._agent = MagicMock()
+        agent._agent.output_queue = asyncio.Queue()
+        agent._agent.task_done = asyncio.Event()
+        agent._agent.run = mock_run
+        agent._agent.state = "idle"
+        agent._session = MagicMock()
+        agent._session.add_message = MagicMock()
+        agent.sessions = MagicMock()
+        agent.sessions.save = MagicMock()
+        agent.memory = MagicMock()
+        agent.memory.get_memory_context = MagicMock(return_value=None)
+        agent.context = MagicMock()
+        agent._prepare_request_context = AsyncMock()
+        agent._build_step_prompt = MagicMock(return_value="prompt")
+        agent._install_anti_loop_tracker = MagicMock()
+
+        stream = AgentLoop.stream(agent, message="test message")
+        first_chunk = await stream.__anext__()
+
+        assert first_chunk["type"] == "content"
+        assert first_chunk["delta"] == "hello"
+
+        await stream.aclose()
+        await asyncio.wait_for(run_cancelled.wait(), timeout=1.0)
+
+        agent.sessions.save.assert_not_called()
+        agent._session.add_message.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_install_anti_loop_tracker_does_not_stack_previous_request_prompt(self):
+        """A new request should not inherit the previous request's anti-loop wrapper."""
+        from pathlib import Path
+        from spoon_bot.agent.loop import AgentLoop
+
+        seen_prompts = []
+
+        async def base_think():
+            seen_prompts.append(agent._agent.next_step_prompt)
+            return True
+
+        tool_call = MagicMock()
+        tool_call.function = MagicMock()
+        tool_call.function.name = "shell"
+        tool_call.function.arguments = '{"command":"cd /workspace && ls -la .agents/skills/pdf"}'
+
+        agent = MagicMock(spec=AgentLoop)
+        agent.workspace = Path("/workspace")
+        agent._agent = MagicMock()
+        agent._agent.think = base_think
+        agent._agent._spoon_bot_base_think = base_think
+        agent._agent.next_step_prompt = ""
+        agent._agent.tool_calls = [tool_call]
+        agent._agent.memory = MagicMock()
+        agent._agent.memory.messages = []
+        agent._compress_runtime_context = MagicMock(return_value=0)
+
+        AgentLoop._install_anti_loop_tracker(agent, "prompt one")
+        agent._agent.next_step_prompt = "prompt one"
+        await agent._agent.think()
+
+        AgentLoop._install_anti_loop_tracker(agent, "prompt two")
+        agent._agent.next_step_prompt = "prompt two"
+        await agent._agent.think()
+
+        assert seen_prompts[-1] == "prompt two"
+
 
 @pytest.mark.requires_spoon_core
 class TestAgentLoopProcessWithThinking: