Hermes-agent

tests/run_agent/test_stream_interrupt_retry.py

@@ -0,0 +1,162 @@
+"""Tests that /stop interrupts streaming retry loops immediately.
+
+When the agent is interrupted during a streaming API call, the outer poll
+loop closes the HTTP connection.  The inner `_call()` thread sees a
+connection error and enters its retry loop.  Before this fix, the retry
+loop would open a FRESH connection without checking `_interrupt_requested`,
+making /stop take multiple retry cycles × read-timeout to actually stop
+(510+ seconds observed on slow ollama-cloud providers).
+
+The fix adds an `_interrupt_requested` check at the top of the retry loop
+so the agent exits immediately instead of retrying.
+"""
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+def _make_agent(**kwargs):
+    """Create a minimal AIAgent for streaming tests."""
+    from run_agent import AIAgent
+
+    defaults = dict(
+        api_key="test-key",
+        base_url="https://example.com/v1",
+        model="test/model",
+        quiet_mode=True,
+        skip_context_files=True,
+        skip_memory=True,
+    )
+    defaults.update(kwargs)
+    agent = AIAgent(**defaults)
+    agent.api_mode = "chat_completions"
+    return agent
+
+
+class TestStreamInterruptBeforeRetry:
+    """Verify _interrupt_requested is checked before each streaming retry."""
+
+    @pytest.mark.filterwarnings(
+        "ignore::pytest.PytestUnhandledThreadExceptionWarning"
+    )
+    @patch("run_agent.AIAgent._create_request_openai_client")
+    @patch("run_agent.AIAgent._close_request_openai_client")
+    def test_interrupt_prevents_stream_retry(self, mock_close, mock_create):
+        """When _interrupt_requested is set during a transient stream error,
+        the retry loop must NOT retry — it should raise InterruptedError
+        immediately instead of opening a fresh connection."""
+        import httpx
+
+        attempt_count = [0]
+
+        def fail_once_then_interrupt(*args, **kwargs):
+            attempt_count[0] += 1
+            if attempt_count[0] == 1:
+                # First attempt: simulate normal failure, then set interrupt
+                # (as if /stop arrived while the retry loop processes the error)
+                agent._interrupt_requested = True
+                raise httpx.ConnectError("connection reset by /stop")
+            # Should never reach here — the interrupt check should fire first
+            raise httpx.ConnectError("unexpected retry — interrupt not checked!")
+
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.side_effect = fail_once_then_interrupt
+        mock_create.return_value = mock_client
+
+        agent = _make_agent()
+        agent._interrupt_requested = False
+
+        with pytest.raises(InterruptedError, match="interrupted"):
+            agent._interruptible_streaming_api_call({})
+
+        # Only 1 attempt should have been made — the interrupt should prevent retry
+        assert attempt_count[0] == 1, (
+            f"Expected 1 attempt but got {attempt_count[0]}. "
+            "The retry loop retried despite _interrupt_requested being set."
+        )
+
+    @pytest.mark.filterwarnings(
+        "ignore::pytest.PytestUnhandledThreadExceptionWarning"
+    )
+    @patch("run_agent.AIAgent._create_request_openai_client")
+    @patch("run_agent.AIAgent._close_request_openai_client")
+    def test_interrupt_before_first_attempt(self, mock_close, mock_create):
+        """If _interrupt_requested is already set when the streaming call
+        starts, it should exit immediately without making any API call."""
+        mock_client = MagicMock()
+        mock_create.return_value = mock_client
+
+        agent = _make_agent()
+        agent._interrupt_requested = True  # Pre-set before call
+
+        with pytest.raises(InterruptedError, match="interrupted"):
+            agent._interruptible_streaming_api_call({})
+
+        # No API call should have been made at all
+        assert mock_client.chat.completions.create.call_count == 0
+
+    @patch("run_agent.AIAgent._create_request_openai_client")
+    @patch("run_agent.AIAgent._close_request_openai_client")
+    def test_normal_retry_still_works_without_interrupt(self, mock_close, mock_create):
+        """Without an interrupt, transient errors should still retry normally."""
+        import httpx
+
+        attempts = [0]
+
+        def fail_twice_then_succeed(*args, **kwargs):
+            attempts[0] += 1
+            if attempts[0] <= 2:
+                raise httpx.ConnectError("transient failure")
+            # Third attempt succeeds
+            chunks = [
+                SimpleNamespace(
+                    choices=[
+                        SimpleNamespace(
+                            index=0,
+                            delta=SimpleNamespace(
+                                content="ok",
+                                tool_calls=None,
+                                reasoning_content=None,
+                                reasoning=None,
+                            ),
+                            finish_reason=None,
+                        )
+                    ],
+                    model="test/model",
+                    usage=None,
+                ),
+                SimpleNamespace(
+                    choices=[
+                        SimpleNamespace(
+                            index=0,
+                            delta=SimpleNamespace(
+                                content=None,
+                                tool_calls=None,
+                                reasoning_content=None,
+                                reasoning=None,
+                            ),
+                            finish_reason="stop",
+                        )
+                    ],
+                    model="test/model",
+                    usage=None,
+                ),
+            ]
+            stream = MagicMock()
+            stream.__iter__ = MagicMock(return_value=iter(chunks))
+            stream.response = MagicMock()
+            stream.response.headers = {}
+            return stream
+
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.side_effect = fail_twice_then_succeed
+        mock_create.return_value = mock_client
+
+        agent = _make_agent()
+        agent._interrupt_requested = False
+
+        # Should succeed on the third attempt
+        result = agent._interruptible_streaming_api_call({})
+        assert result is not None
+        assert attempts[0] == 3