diff --git a/tests/gateway/test_api_server_runs.py b/tests/gateway/test_api_server_runs.py index c273bc4b6f..18c6e243a9 100644 --- a/tests/gateway/test_api_server_runs.py +++ b/tests/gateway/test_api_server_runs.py @@ -28,6 +28,7 @@ from gateway.platforms.api_server import ( cors_middleware, security_headers_middleware, ) +from gateway.platforms.api_server_runs import _RunStream from tools import approval as approval_mod from tools import approval_gateway_wait @@ -598,7 +599,7 @@ class TestRunEvents: @pytest.mark.asyncio async def test_tool_completed_event_includes_redacted_bounded_result_preview(self, adapter): loop = asyncio.get_running_loop() - adapter._run_streams["run_tool"] = asyncio.Queue() + adapter._run_streams["run_tool"] = _RunStream() callback = adapter._make_run_event_callback("run_tool", loop) callback( @@ -610,7 +611,8 @@ class TestRunEvents: "output": "x" * 600, }, ) - event = await adapter._run_streams["run_tool"].get() + await asyncio.sleep(0) # the callback hops onto the loop via call_soon_threadsafe + _, event = adapter._run_streams["run_tool"].backlog[-1] assert event["error"] is True assert "BLOCKED: approval required" in event["preview"] @@ -862,9 +864,9 @@ class TestSteerRun: app = _create_runs_app(adapter) agent = MagicMock() agent.steer.return_value = True - queue = asyncio.Queue() + stream = _RunStream() adapter._active_run_agents["run_123"] = agent - adapter._run_streams["run_123"] = queue + adapter._run_streams["run_123"] = stream adapter._set_run_status("run_123", "running") _claim_run(adapter, "run_123") @@ -880,7 +882,7 @@ class TestSteerRun: } agent.steer.assert_called_once_with("tighten the ending") assert adapter._run_statuses["run_123"]["last_event"] == "run.steered" - event = queue.get_nowait() + _, event = stream.backlog[-1] assert event["event"] == "run.steered" assert event["run_id"] == "run_123" assert event["accepted"] is True diff --git a/tests/gateway/test_api_server_sse_keepalive.py b/tests/gateway/test_api_server_sse_keepalive.py index 6eb3776bde..e2aa68ebf4 100644 --- a/tests/gateway/test_api_server_sse_keepalive.py +++ b/tests/gateway/test_api_server_sse_keepalive.py @@ -68,8 +68,10 @@ async def test_idle_openai_stream_writes_keepalive_before_remote_client_deadline async def test_idle_run_events_stream_uses_shared_keepalive_cadence(monkeypatch, adapter): """``GET /v1/runs/{id}/events`` follows the same keepalive constant as the OpenAI routes.""" monkeypatch.setattr(api_server, "CHAT_COMPLETIONS_SSE_KEEPALIVE_SECONDS", 0.2) - queue: asyncio.Queue = asyncio.Queue() - adapter._run_streams["run_idle"] = queue + from gateway.platforms.api_server_runs import _RunStream + + stream = _RunStream() + adapter._run_streams["run_idle"] = stream adapter._set_run_status("run_idle", "running") monkeypatch.setattr(adapter, "_request_owns_run", lambda request, run_id: True) @@ -79,7 +81,7 @@ async def test_idle_run_events_stream_uses_shared_keepalive_cadence(monkeypatch, async def _close_after_idle(): await asyncio.sleep(0.5) - await queue.put(None) + stream.put_nowait(None) async with TestClient(TestServer(app)) as cli: closer = asyncio.create_task(_close_after_idle())