From 1da89a5f3dd41ecad97272557a1623c851bfc12c Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:56:14 -0700 Subject: [PATCH] fix(api): keep live runs tracked past stream ttl --- gateway/platforms/api_server.py | 21 ++++-- tests/gateway/test_api_server.py | 25 +++++-- tests/gateway/test_api_server_runs.py | 94 +++++++++++++++++++++++++++ 3 files changed, 129 insertions(+), 11 deletions(-) diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index c71d92a29b52..e437b51a1331 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -899,7 +899,8 @@ class APIServerAdapter(BasePlatformAdapter): # from a request flood (#7483). self._max_concurrent_runs: int = self._resolve_max_concurrent_runs() # Number of in-flight runs on the non-streaming chat/responses paths - # (the /v1/runs path tracks its own in-flight set via _run_streams). + # (the /v1/runs path tracks its own in-flight set via + # _active_run_tasks). self._inflight_agent_runs: int = 0 def _readiness_work_counts(self) -> tuple[int, int, int]: @@ -4001,14 +4002,18 @@ class APIServerAdapter(BasePlatformAdapter): The cap bounds total in-flight agent activity across every agent-serving endpoint: the non-streaming chat/responses paths - (tracked by ``_inflight_agent_runs``) plus the ``/v1/runs`` streaming - path (tracked by ``_run_streams``). A configured value of 0 disables - the cap entirely. + (tracked by ``_inflight_agent_runs``) plus the ``/v1/runs`` path + (tracked by live entries in ``_active_run_tasks``). Stream queues are + transport state and may disappear while their underlying run remains + active, so they must not define run concurrency. A configured value of + 0 disables the cap entirely. """ limit = self._max_concurrent_runs if limit <= 0: return None - inflight = self._inflight_agent_runs + len(self._run_streams) + inflight = self._inflight_agent_runs + sum( + not task.done() for task in self._active_run_tasks.values() + ) if inflight >= limit: return web.json_response( _openai_error( @@ -4705,7 +4710,7 @@ class APIServerAdapter(BasePlatformAdapter): return web.json_response({"run_id": run_id, "status": "stopping"}) async def _sweep_orphaned_runs(self) -> None: - """Periodically clean up run streams that were never consumed.""" + """Periodically clean up expired transport state for inactive runs.""" while True: await asyncio.sleep(60) now = time.time() @@ -4713,6 +4718,10 @@ class APIServerAdapter(BasePlatformAdapter): 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) diff --git a/tests/gateway/test_api_server.py b/tests/gateway/test_api_server.py index 1d5e9ca91e10..3892652c3be7 100644 --- a/tests/gateway/test_api_server.py +++ b/tests/gateway/test_api_server.py @@ -568,14 +568,29 @@ class TestConcurrencyCap: assert resp.headers.get("Retry-After") def test_cap_counts_both_buckets(self): - # /v1/runs (tracked by _run_streams) + chat/responses (inflight) + # /v1/runs (tracked by live tasks) + chat/responses (inflight) adapter = _make_adapter() adapter._max_concurrent_runs = 4 adapter._inflight_agent_runs = 2 - adapter._run_streams = {"r1": object(), "r2": object()} - resp = adapter._concurrency_limited_response() - assert resp is not None - assert resp.status == 429 + + async def _assert_live_tasks_are_counted_without_streams(): + blocker = asyncio.Event() + + async def _live_run(): + await blocker.wait() + + tasks = [asyncio.create_task(_live_run()) for _ in range(2)] + adapter._active_run_tasks = {f"r{i}": task for i, task in enumerate(tasks)} + adapter._run_streams = {} + try: + resp = adapter._concurrency_limited_response() + assert resp is not None + assert resp.status == 429 + finally: + blocker.set() + await asyncio.gather(*tasks) + + asyncio.run(_assert_live_tasks_are_counted_without_streams()) def test_zero_disables_cap(self): adapter = _make_adapter() diff --git a/tests/gateway/test_api_server_runs.py b/tests/gateway/test_api_server_runs.py index 16d7866f1295..4f160416b184 100644 --- a/tests/gateway/test_api_server_runs.py +++ b/tests/gateway/test_api_server_runs.py @@ -442,6 +442,100 @@ class TestRunEvents: assert resp.status == 401 +# --------------------------------------------------------------------------- +# Run lifecycle TTL sweeping +# --------------------------------------------------------------------------- + + +class TestRunLifecycleSweep: + @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.""" + app = _create_runs_app(adapter) + adapter._max_concurrent_runs = 1 + + 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"}) + assert start_resp.status == 202 + run_id = (await start_resp.json())["run_id"] + assert agent_ready.wait(timeout=3.0) + + task = adapter._active_run_tasks[run_id] + assert isinstance(task, asyncio.Task) + assert not task.done() + + pending = approval_mod._ApprovalEntry({ + "command": "bash -c long-running", + "description": "approval after stream TTL", + "pattern_keys": ["shell-c"], + }) + with approval_mod._lock: + approval_mod._gateway_queues[run_id] = [pending] + + adapter._run_streams_created[run_id] -= adapter._RUN_STREAM_TTL + 1 + # Exercise one real sweeper iteration without waiting 60 seconds. + with patch( + "gateway.platforms.api_server.asyncio.sleep", + side_effect=[None, asyncio.CancelledError()], + ): + with pytest.raises(asyncio.CancelledError): + await adapter._sweep_orphaned_runs() + + 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 adapter._run_approval_sessions[run_id] == run_id + + limited = adapter._concurrency_limited_response() + assert limited is not None + assert limited.status == 429 + + approval_resp = await cli.post( + f"/v1/runs/{run_id}/approval", + json={"choice": "once"}, + ) + assert approval_resp.status == 200 + assert pending.event.is_set() + assert pending.result == "once" + + stop_resp = await cli.post(f"/v1/runs/{run_id}/stop") + assert stop_resp.status == 200 + mock_agent.interrupt.assert_called_once_with("Stop requested via API") + + @pytest.mark.asyncio + async def test_expired_orphan_run_state_is_reaped(self, adapter): + run_id = "run_expired_orphan" + adapter._run_streams[run_id] = asyncio.Queue() + adapter._run_streams_created[run_id] = 0 + adapter._run_approval_sessions[run_id] = run_id + + pending = approval_mod._ApprovalEntry({ + "command": "bash -c orphaned", + "description": "orphaned approval", + "pattern_keys": ["shell-c"], + }) + with approval_mod._lock: + approval_mod._gateway_queues[run_id] = [pending] + + with patch( + "gateway.platforms.api_server.asyncio.sleep", + side_effect=[None, asyncio.CancelledError()], + ): + with pytest.raises(asyncio.CancelledError): + await adapter._sweep_orphaned_runs() + + assert run_id not in adapter._run_streams + assert run_id not in adapter._run_streams_created + assert run_id not in adapter._run_approval_sessions + assert pending.event.is_set() + with approval_mod._lock: + assert run_id not in approval_mod._gateway_queues + + # --------------------------------------------------------------------------- # POST /v1/runs/{run_id}/stop — interrupt a running agent # ---------------------------------------------------------------------------