diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index e437b51a1331..b442c0c60f9e 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -882,9 +882,13 @@ class APIServerAdapter(BasePlatformAdapter): self._run_streams: Dict[str, "asyncio.Queue[Optional[Dict]]"] = {} # Creation timestamps for orphaned-run TTL sweep self._run_streams_created: Dict[str, float] = {} + # Runs with a connected SSE consumer; their queue is actively draining. + self._run_stream_subscribers: set[str] = set() # Active run agent/task references for stop support self._active_run_agents: Dict[str, Any] = {} self._active_run_tasks: Dict[str, "asyncio.Task"] = {} + # Stop is cooperative: the executor thread may outlive the HTTP request. + self._stopping_run_ids: set[str] = set() # Pollable run status for dashboards and external control-plane UIs. self._run_statuses: Dict[str, Dict[str, Any]] = {} # Active approval session key for each run_id. The approval core @@ -4304,6 +4308,8 @@ class APIServerAdapter(BasePlatformAdapter): def _text_cb(delta: Optional[str]) -> None: if delta is None: return + if run_id not in self._run_streams: + return try: loop.call_soon_threadsafe(q.put_nowait, { "event": "message.delta", @@ -4328,6 +4334,18 @@ class APIServerAdapter(BasePlatformAdapter): async def _run_and_close(): try: self._set_run_status(run_id, "running") + if run_id in self._stopping_run_ids: + q.put_nowait({ + "event": "run.cancelled", + "run_id": run_id, + "timestamp": time.time(), + }) + self._set_run_status( + run_id, + "cancelled", + last_event="run.cancelled", + ) + return agent = self._create_agent( ephemeral_system_prompt=ephemeral_system_prompt, session_id=session_id, @@ -4412,10 +4430,21 @@ class APIServerAdapter(BasePlatformAdapter): return r, u result, usage = await asyncio.get_running_loop().run_in_executor(None, _run_sync) + if run_id in self._stopping_run_ids: + q.put_nowait({ + "event": "run.cancelled", + "run_id": run_id, + "timestamp": time.time(), + }) + self._set_run_status( + run_id, + "cancelled", + last_event="run.cancelled", + ) # Check for structured failure (non-retryable client errors like # 401/400 return failed=True instead of raising, so the except # block below never fires — issue #15561). - if isinstance(result, dict) and result.get("failed"): + elif isinstance(result, dict) and result.get("failed"): error_msg = _redact_api_error_text(result.get("error") or "agent run failed") q.put_nowait({ "event": "run.failed", @@ -4497,6 +4526,7 @@ class APIServerAdapter(BasePlatformAdapter): self._active_run_agents.pop(run_id, None) self._active_run_tasks.pop(run_id, None) self._run_approval_sessions.pop(run_id, None) + self._stopping_run_ids.discard(run_id) task = asyncio.create_task(_run_and_close()) self._active_run_tasks[run_id] = task @@ -4548,6 +4578,7 @@ class APIServerAdapter(BasePlatformAdapter): return web.json_response(_openai_error(f"Run not found: {run_id}", code="run_not_found"), status=404) q = self._run_streams[run_id] + self._run_stream_subscribers.add(run_id) response = web.StreamResponse( status=200, @@ -4575,6 +4606,7 @@ class APIServerAdapter(BasePlatformAdapter): except Exception as exc: logger.debug("[api_server] SSE stream error for run %s: %s", run_id, exc) finally: + self._run_stream_subscribers.discard(run_id) self._run_streams.pop(run_id, None) self._run_streams_created.pop(run_id, None) @@ -4683,6 +4715,7 @@ class APIServerAdapter(BasePlatformAdapter): return web.json_response(_openai_error(f"Run not found: {run_id}", code="run_not_found"), status=404) self._set_run_status(run_id, "stopping", last_event="run.stopping") + self._stopping_run_ids.add(run_id) if agent is not None: try: @@ -4690,41 +4723,29 @@ class APIServerAdapter(BasePlatformAdapter): except Exception: pass - if task is not None and not task.done(): - task.cancel() - # Bounded wait: run_conversation() executes in the default - # executor thread which task.cancel() cannot preempt — we rely on - # agent.interrupt() above to break the loop. Cap the wait so a - # slow/unresponsive interrupt can't hang this handler. - try: - await asyncio.wait_for(asyncio.shield(task), timeout=5.0) - except asyncio.TimeoutError: - logger.warning( - "[api_server] stop for run %s timed out after 5s; " - "agent may still be finishing the current step", - run_id, - ) - except (asyncio.CancelledError, Exception): - pass - return web.json_response({"run_id": run_id, "status": "stopping"}) async def _sweep_orphaned_runs(self) -> None: - """Periodically clean up expired transport state for inactive runs.""" + """Periodically expire transport buffers and terminal status records.""" while True: await asyncio.sleep(60) + self._sweep_orphaned_runs_once(time.time()) + + def _sweep_orphaned_runs_once(self, now: Optional[float] = None) -> None: + """Expire old SSE buffers without treating transport age as run age.""" + if now is None: now = time.time() - stale = [ - run_id - for run_id, created_at in list(self._run_streams_created.items()) - if now - created_at > self._RUN_STREAM_TTL - and ( - (task := self._active_run_tasks.get(run_id)) is None - or task.done() - ) - ] - for run_id in stale: - logger.debug("[api_server] sweeping orphaned run %s", run_id) + stale = [ + run_id + for run_id, created_at in list(self._run_streams_created.items()) + if now - created_at > self._RUN_STREAM_TTL + and run_id not in self._run_stream_subscribers + ] + for run_id in stale: + logger.debug("[api_server] sweeping expired run transport %s", run_id) + task = self._active_run_tasks.get(run_id) + task_done = task is None or task.done() + if task_done: try: from tools.approval import unregister_gateway_notify @@ -4733,20 +4754,24 @@ class APIServerAdapter(BasePlatformAdapter): unregister_gateway_notify(approval_session_key) except Exception: pass - self._run_streams.pop(run_id, None) - self._run_streams_created.pop(run_id, None) + # The transport TTL always bounds buffering. Live control state is + # independent and survives until the executor-backed task returns. + self._run_streams.pop(run_id, None) + self._run_streams_created.pop(run_id, None) + if task_done: self._active_run_agents.pop(run_id, None) self._active_run_tasks.pop(run_id, None) self._run_approval_sessions.pop(run_id, None) + self._stopping_run_ids.discard(run_id) - stale_statuses = [ - run_id - for run_id, status in list(self._run_statuses.items()) - if status.get("status") in {"completed", "failed", "cancelled"} - and now - float(status.get("updated_at", 0) or 0) > self._RUN_STATUS_TTL - ] - for run_id in stale_statuses: - self._run_statuses.pop(run_id, None) + stale_statuses = [ + run_id + for run_id, status in list(self._run_statuses.items()) + if status.get("status") in {"completed", "failed", "cancelled"} + and now - float(status.get("updated_at", 0) or 0) > self._RUN_STATUS_TTL + ] + for run_id in stale_statuses: + self._run_statuses.pop(run_id, None) # ------------------------------------------------------------------ # BasePlatformAdapter interface diff --git a/tests/gateway/test_api_server_runs.py b/tests/gateway/test_api_server_runs.py index 4f160416b184..957dbfac383e 100644 --- a/tests/gateway/test_api_server_runs.py +++ b/tests/gateway/test_api_server_runs.py @@ -10,6 +10,7 @@ Covers: import asyncio import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -448,9 +449,21 @@ class TestRunEvents: class TestRunLifecycleSweep: + def test_sweep_keeps_transport_with_active_subscriber(self, adapter): + run_id = "run_subscribed" + queue = asyncio.Queue() + adapter._run_streams[run_id] = queue + adapter._run_streams_created[run_id] = 0 + adapter._run_stream_subscribers.add(run_id) + + adapter._sweep_orphaned_runs_once(time.time()) + + assert adapter._run_streams[run_id] is queue + assert run_id in adapter._run_streams_created + @pytest.mark.asyncio - async def test_expired_live_run_remains_counted_approvable_and_stoppable(self, adapter): - """Stream TTL must not detach control state from a still-running task.""" + async def test_expired_live_run_drops_transport_but_keeps_control_state(self, adapter): + """Stream TTL bounds buffering without detaching a live run.""" app = _create_runs_app(adapter) adapter._max_concurrent_runs = 1 @@ -487,7 +500,8 @@ class TestRunLifecycleSweep: assert adapter._active_run_tasks[run_id] is task assert adapter._active_run_agents[run_id] is mock_agent - assert run_id in adapter._run_streams + assert run_id not in adapter._run_streams + assert run_id not in adapter._run_streams_created assert adapter._run_approval_sessions[run_id] == run_id limited = adapter._concurrency_limited_response() @@ -506,6 +520,30 @@ class TestRunLifecycleSweep: assert stop_resp.status == 200 mock_agent.interrupt.assert_called_once_with("Stop requested via API") + @pytest.mark.asyncio + async def test_expired_transport_stops_buffering_new_deltas(self, adapter): + """An unconsumed expired queue must not grow for the rest of a live run.""" + app = _create_runs_app(adapter) + + async with TestClient(TestServer(app)) as cli: + with patch.object(adapter, "_create_agent") as mock_create: + mock_agent, agent_ready, _ = _make_slow_agent() + mock_create.return_value = mock_agent + + start_resp = await cli.post("/v1/runs", json={"input": "hello"}) + run_id = (await start_resp.json())["run_id"] + assert agent_ready.wait(timeout=3.0) + expired_queue = adapter._run_streams[run_id] + stream_delta = mock_create.call_args.kwargs["stream_delta_callback"] + + adapter._run_streams_created[run_id] -= adapter._RUN_STREAM_TTL + 1 + adapter._sweep_orphaned_runs_once(time.time()) + before = expired_queue.qsize() + stream_delta("must-not-buffer") + await asyncio.sleep(0) + + assert expired_queue.qsize() == before + @pytest.mark.asyncio async def test_expired_orphan_run_state_is_reaped(self, adapter): run_id = "run_expired_orphan" @@ -542,6 +580,88 @@ class TestRunLifecycleSweep: class TestStopRun: + @pytest.mark.asyncio + async def test_stop_before_agent_creation_prevents_run_start(self, adapter): + """A stop accepted while queued must prevent agent construction.""" + app = _create_runs_app(adapter) + original_create_task = asyncio.create_task + task_started = asyncio.Event() + allow_task = asyncio.Event() + + def _delayed_create_task(coro): + async def _delayed(): + task_started.set() + await allow_task.wait() + return await coro + + return original_create_task(_delayed()) + + with patch("gateway.platforms.api_server.asyncio.create_task", side_effect=_delayed_create_task), \ + patch.object(adapter, "_create_agent") as mock_create: + async with TestClient(TestServer(app)) as cli: + resp = await cli.post("/v1/runs", json={"input": "hello"}) + run_id = (await resp.json())["run_id"] + await task_started.wait() + + stop_resp = await cli.post(f"/v1/runs/{run_id}/stop") + assert stop_resp.status == 200 + allow_task.set() + + for _ in range(20): + if run_id not in adapter._active_run_tasks: + break + await asyncio.sleep(0.05) + + mock_create.assert_not_called() + assert adapter._run_statuses[run_id]["status"] == "cancelled" + + @pytest.mark.asyncio + async def test_stop_keeps_uncooperative_executor_tracked_until_exit(self, adapter): + """Cancelling an asyncio wrapper must not hide its live executor thread.""" + app = _create_runs_app(adapter) + run_can_finish = threading.Event() + run_finished = threading.Event() + + async with TestClient(TestServer(app)) as cli: + with patch.object(adapter, "_create_agent") as mock_create: + mock_agent = MagicMock() + mock_agent.session_prompt_tokens = 0 + mock_agent.session_completion_tokens = 0 + mock_agent.session_total_tokens = 0 + started = threading.Event() + + def _run_conversation(*_args, **_kwargs): + started.set() + run_can_finish.wait(timeout=5) + run_finished.set() + return {"final_response": "late result"} + + mock_agent.run_conversation.side_effect = _run_conversation + mock_create.return_value = mock_agent + + resp = await cli.post("/v1/runs", json={"input": "hello"}) + run_id = (await resp.json())["run_id"] + assert started.wait(timeout=3) + + stop_resp = await cli.post(f"/v1/runs/{run_id}/stop") + assert stop_resp.status == 200 + await asyncio.sleep(0.1) + + assert not run_finished.is_set() + assert run_id in adapter._active_run_agents + assert run_id in adapter._active_run_tasks + assert adapter._run_statuses[run_id]["status"] == "stopping" + + run_can_finish.set() + for _ in range(40): + if run_id not in adapter._active_run_tasks: + break + await asyncio.sleep(0.05) + + assert run_id not in adapter._active_run_agents + assert run_id not in adapter._active_run_tasks + assert adapter._run_statuses[run_id]["status"] == "cancelled" + @pytest.mark.asyncio async def test_stop_running_agent(self, adapter): """Stop should interrupt the agent and cancel the task."""