Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/gateway/test_sse_agent_cancel.py

@@ -0,0 +1,482 @@
+"""Tests for SSE client disconnect → agent task cancellation.
+
+When a streaming /v1/chat/completions client disconnects mid-stream
+(network drop, browser tab close), the agent is interrupted via
+agent.interrupt() so it stops making LLM API calls, and the asyncio
+task wrapper is cancelled.
+"""
+
+import asyncio
+import threading
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_adapter():
+    """Build a minimal APIServerAdapter with mocked internals."""
+    from gateway.platforms.api_server import APIServerAdapter
+    from gateway.config import PlatformConfig
+
+    config = PlatformConfig(enabled=True, token="test-key")
+    adapter = APIServerAdapter(config)
+    return adapter
+
+
+def _make_request():
+    """Build a mock aiohttp request."""
+    req = MagicMock()
+    req.headers = {}
+    return req
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+class TestSSEAgentCancelOnDisconnect:
+    """gateway/platforms/api_server.py — _write_sse_chat_completion()"""
+
+    def test_agent_task_cancelled_on_client_disconnect(self):
+        """When response.write raises ConnectionResetError (client dropped),
+        the agent task must be cancelled."""
+        adapter = _make_adapter()
+
+        # Agent task that runs forever (simulates a long LLM call)
+        agent_done = asyncio.Event()
+
+        async def fake_agent():
+            await agent_done.wait()
+            return {"final_response": "done"}, {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            # Constructed inside the running loop — ThreadSafeAsyncQueue
+            # captures asyncio.get_running_loop() at construction time.
+            stream_q = ThreadSafeAsyncQueue()
+            stream_q.put_nowait("hello ")  # Some data already queued
+
+            agent_task = asyncio.ensure_future(fake_agent())
+
+            # Mock response that raises ConnectionResetError on second write
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            call_count = 0
+
+            async def write_side_effect(data):
+                nonlocal call_count
+                call_count += 1
+                if call_count >= 2:
+                    raise ConnectionResetError("client disconnected")
+
+            mock_response.write = AsyncMock(side_effect=write_side_effect)
+            mock_response.prepare = AsyncMock()
+
+            with patch.object(type(adapter), '_write_sse_chat_completion',
+                              adapter._write_sse_chat_completion):
+                # Patch StreamResponse creation
+                with patch("gateway.platforms.api_server.web.StreamResponse",
+                           return_value=mock_response):
+                    await adapter._write_sse_chat_completion(
+                        _make_request(), "cmpl-123", "gpt-4", 1234567890,
+                        stream_q, agent_task,
+                    )
+
+            # The critical assertion: agent_task must be cancelled
+            assert agent_task.cancelled() or agent_task.done()
+            # Clean up
+            agent_done.set()
+
+        asyncio.run(run())
+
+    def test_agent_task_not_cancelled_on_normal_completion(self):
+        """On normal stream completion, agent task should NOT be cancelled."""
+        adapter = _make_adapter()
+
+        async def fake_agent():
+            return {"final_response": "done"}, {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+            stream_q.put_nowait("hello")
+            stream_q.put_nowait(None)  # End-of-stream sentinel
+
+            agent_task = asyncio.ensure_future(fake_agent())
+            await asyncio.sleep(0)  # Let agent complete
+
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            mock_response.write = AsyncMock()
+            mock_response.prepare = AsyncMock()
+
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=mock_response):
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-456", "gpt-4", 1234567890,
+                    stream_q, agent_task,
+                )
+
+            # Agent should have completed normally, not been cancelled
+            assert agent_task.done()
+            assert not agent_task.cancelled()
+
+        asyncio.run(run())
+
+    def test_broken_pipe_also_cancels_agent(self):
+        """BrokenPipeError (another disconnect variant) also cancels the task."""
+        adapter = _make_adapter()
+
+        async def fake_agent():
+            await asyncio.sleep(0.2)  # Never completes
+            return {}, {}
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+
+            agent_task = asyncio.ensure_future(fake_agent())
+
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            mock_response.write = AsyncMock(side_effect=BrokenPipeError("pipe broken"))
+            mock_response.prepare = AsyncMock()
+
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=mock_response):
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-789", "gpt-4", 1234567890,
+                    stream_q, agent_task,
+                )
+
+            assert agent_task.cancelled() or agent_task.done()
+
+        asyncio.run(run())
+
+    def test_already_done_task_not_cancelled_on_disconnect(self):
+        """If agent already finished before disconnect, don't try to cancel."""
+        adapter = _make_adapter()
+
+        async def fake_agent():
+            return {"final_response": "done"}, {}
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+            stream_q.put_nowait("data")
+
+            agent_task = asyncio.ensure_future(fake_agent())
+            await asyncio.sleep(0)  # Let agent complete
+
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            call_count = 0
+
+            async def write_side_effect(data):
+                nonlocal call_count
+                call_count += 1
+                if call_count >= 2:
+                    raise ConnectionResetError("late disconnect")
+
+            mock_response.write = AsyncMock(side_effect=write_side_effect)
+            mock_response.prepare = AsyncMock()
+
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=mock_response):
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-done", "gpt-4", 1234567890,
+                    stream_q, agent_task,
+                )
+
+            # Task was already done — should not be cancelled
+            assert agent_task.done()
+            assert not agent_task.cancelled()
+
+        asyncio.run(run())
+
+    def test_agent_interrupt_called_on_disconnect(self):
+        """When the client disconnects, agent.interrupt() must be called
+        so the agent thread stops making LLM API calls."""
+        adapter = _make_adapter()
+
+        agent_done = asyncio.Event()
+
+        async def fake_agent():
+            await agent_done.wait()
+            return {"final_response": "done"}, {}
+
+        # Mock agent with an interrupt method
+        mock_agent = MagicMock()
+        mock_agent.interrupt = MagicMock()
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+            stream_q.put_nowait("hello ")
+
+            agent_task = asyncio.ensure_future(fake_agent())
+            agent_ref = [mock_agent]
+
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            call_count = 0
+
+            async def write_side_effect(data):
+                nonlocal call_count
+                call_count += 1
+                if call_count >= 2:
+                    raise ConnectionResetError("client disconnected")
+
+            mock_response.write = AsyncMock(side_effect=write_side_effect)
+            mock_response.prepare = AsyncMock()
+
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=mock_response):
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-int", "gpt-4", 1234567890,
+                    stream_q, agent_task, agent_ref,
+                )
+
+            # agent.interrupt() must have been called
+            mock_agent.interrupt.assert_called_once_with("SSE client disconnected")
+            # Clean up
+            agent_done.set()
+
+        asyncio.run(run())
+
+    def test_agent_ref_none_still_cancels_task(self):
+        """When agent_ref is not provided (None), the task is still cancelled
+        on disconnect — just without the interrupt() call."""
+        adapter = _make_adapter()
+
+        async def fake_agent():
+            await asyncio.sleep(999)
+            return {}, {}
+
+        async def run():
+            from aiohttp import web
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+
+            agent_task = asyncio.ensure_future(fake_agent())
+
+            mock_response = AsyncMock(spec=web.StreamResponse)
+            mock_response.write = AsyncMock(side_effect=BrokenPipeError("gone"))
+            mock_response.prepare = AsyncMock()
+
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=mock_response):
+                # No agent_ref passed — should still handle disconnect cleanly
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-noref", "gpt-4", 1234567890,
+                    stream_q, agent_task,
+                )
+
+            assert agent_task.cancelled() or agent_task.done()
+
+        asyncio.run(run())
+
+
+def _capturing_response():
+    """Mock StreamResponse that records all written SSE bytes as text."""
+    from aiohttp import web
+
+    chunks: list = []
+    resp = AsyncMock(spec=web.StreamResponse)
+    resp.prepare = AsyncMock()
+
+    async def _write(data):
+        chunks.append(data.decode() if isinstance(data, (bytes, bytearray)) else data)
+
+    resp.write = AsyncMock(side_effect=_write)
+    return resp, chunks
+
+
+def _finish_reason(chunks: list):
+    """Extract the terminal finish_reason and its chunk from captured SSE."""
+    import json
+
+    sse = "".join(chunks)
+    finish = None
+    for line in sse.splitlines():
+        if line.startswith("data: ") and '"finish_reason"' in line:
+            obj = json.loads(line[6:])
+            if obj["choices"][0].get("finish_reason") is not None:
+                finish = obj
+    return (finish["choices"][0]["finish_reason"] if finish else None), finish, sse
+
+
+class TestSSEAgentFailureFinishReason:
+    """gateway/platforms/api_server.py — _write_sse_chat_completion()
+
+    A clean stream-queue termination (sentinel received) followed by an agent
+    failure must NOT report finish_reason: "stop". Both failure modes — an
+    ``agent_task`` that raises and a ``result`` dict flagged failed — surface
+    as finish_reason: "error", mirroring the non-streaming path. Issue #12422.
+    """
+
+    def _run(self, fake_agent, queue_items=("partial",)):
+        adapter = _make_adapter()
+
+        async def run():
+            from gateway.platforms.api_server import ThreadSafeAsyncQueue
+
+            stream_q = ThreadSafeAsyncQueue()
+            for item in queue_items:
+                stream_q.put_nowait(item)
+            stream_q.put_nowait(None)  # clean end-of-stream sentinel
+
+            agent_task = asyncio.ensure_future(fake_agent())
+            resp, chunks = _capturing_response()
+            with patch("gateway.platforms.api_server.web.StreamResponse",
+                       return_value=resp):
+                await adapter._write_sse_chat_completion(
+                    _make_request(), "cmpl-fail", "gpt-4", 1234567890,
+                    stream_q, agent_task,
+                )
+            return _finish_reason(chunks)
+
+        return asyncio.run(run())
+
+    def test_agent_task_raises_reports_error_not_stop(self):
+        async def crash():
+            raise RuntimeError("boom from agent")
+
+        reason, finish, sse = self._run(crash)
+        assert reason == "error"
+        assert "error" in finish
+        assert "data: [DONE]" in sse
+
+    def test_failed_result_dict_reports_error_not_stop(self):
+        async def failed():
+            return (
+                {"final_response": "", "failed": True, "completed": False,
+                 "error": "upstream model 500"},
+                {"input_tokens": 5, "output_tokens": 0, "total_tokens": 5},
+            )
+
+        reason, finish, _ = self._run(failed)
+        assert reason == "error"
+        assert finish.get("hermes", {}).get("failed") is True
+
+    def test_truncated_result_reports_length(self):
+        async def trunc():
+            return (
+                {"final_response": "half", "partial": True, "completed": False,
+                 "error": "output was truncated"},
+                {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8},
+            )
+
+        reason, finish, _ = self._run(trunc)
+        assert reason == "length"
+        assert finish["hermes"]["error_code"] == "output_truncated"
+
+    def test_successful_completion_reports_stop(self):
+        async def ok():
+            return (
+                {"final_response": "hi", "completed": True},
+                {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7},
+            )
+
+        reason, finish, _ = self._run(ok)
+        assert reason == "stop"
+        # No error/hermes pollution on the happy path.
+        assert "error" not in finish
+        assert "hermes" not in finish
+
+
+# ---------------------------------------------------------------------------
+# Sweeper review fix (teknium1, 2026-07-30): cover the cross-thread
+# ``put_threadsafe`` boundary that #72610 introduces via ``ThreadSafeAsyncQueue``.
+# ``run_conversation`` runs in a worker thread (``loop.run_in_executor``),
+# so its ``_on_delta`` / ``_on_tool_*`` callbacks must be able to push into
+# the queue from off the owning event loop and immediately wake the
+# consumer ``get()`` — this is the production boundary the original tests
+# only exercised with same-loop ``put_nowait``.
+# ---------------------------------------------------------------------------
+
+class TestThreadSafeAsyncQueueCrossThreadBoundary:
+    """gateway/platforms/api_server.py — ThreadSafeAsyncQueue"""
+
+    def test_worker_thread_put_threadsafe_wakes_owning_loop_get(self):
+        """A real daemon-Thread calling ``put_threadsafe`` from off-loop must
+        immediately unblock an ``await q.get()`` on the owning event loop.
+        This mirrors the ``run_conversation``/``run_in_executor`` boundary."""
+
+        loop = asyncio.new_event_loop()
+
+        async def consumer():
+            q = ThreadSafeAsyncQueue()
+
+            async def wait_for_item():
+                return await asyncio.wait_for(q.get(), timeout=2)
+
+            def worker():
+                time.sleep(0.05)
+                # No ``loop=`` kwarg on purpose: production callers
+                # (_on_delta / _on_tool_*) never pass one, so the queue
+                # must resolve its own ``_loop_ref``. Passing loop= here
+                # would make a broken _loop_ref pass this test.
+                q.put_threadsafe("from-worker")
+
+            thread = threading.Thread(target=worker, daemon=True)
+            thread.start()
+
+            got = await wait_for_item()
+            assert got == "from-worker"
+
+            thread.join(timeout=2)
+            assert not thread.is_alive()
+
+        loop.run_until_complete(consumer())
+        loop.close()
+
+    def test_twenty_concurrent_threads_no_drop(self):
+        """Twenty concurrent off-loop ``put_threadsafe`` calls — all arrive.
+        Regression test for the #65003 producer path."""
+
+        loop = asyncio.new_event_loop()
+        n = 20
+
+        async def consumer():
+            q = ThreadSafeAsyncQueue()
+            received = []
+
+            async def drain():
+                for _ in range(n):
+                    received.append(await q.get())
+
+            def worker(idx):
+                time.sleep(0.01 + idx * 0.002)
+                # No ``loop=`` kwarg — exercise the production
+                # ``_loop_ref`` resolution path (see the note above).
+                q.put_threadsafe(f"item-{idx}")
+
+            threads = [
+                threading.Thread(target=worker, args=(i,), daemon=True)
+                for i in range(n)
+            ]
+            for t in threads:
+                t.start()
+
+            await asyncio.wait_for(drain(), timeout=5)
+
+            for t in threads:
+                t.join(timeout=2)
+                assert not t.is_alive()
+
+            assert len(received) == n
+            assert set(received) == {f"item-{i}" for i in range(n)}
+
+        loop.run_until_complete(consumer())
+        loop.close()