diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 76f69e73e814..6287132243ad 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -35,6 +35,7 @@ import asyncio import hashlib import hmac import json +from contextlib import contextmanager from contextvars import ContextVar from functools import wraps import logging @@ -703,6 +704,29 @@ def _admit_api_agent_request(handler): return _wrapped +def _release_pending_api_work(adapter, reservation: dict[str, bool]) -> None: + """Release a pending-work reservation exactly once.""" + if reservation["active"]: + reservation["active"] = False + adapter._pending_agent_requests = max(0, adapter._pending_agent_requests - 1) + + +@contextmanager +def _reserve_pending_api_work(adapter): + """Keep externally-triggered background work visible across awaits. + + A handler can detach the reservation to an asyncio task; its done callback + then owns release so shutdown cannot miss the handoff to background work. + """ + reservation = {"active": True, "detached": False} + adapter._pending_agent_requests += 1 + try: + yield reservation + finally: + if not reservation["detached"]: + _release_pending_api_work(adapter, reservation) + + if AIOHTTP_AVAILABLE: @web.middleware async def body_limit_middleware(request, handler): @@ -3927,30 +3951,35 @@ class APIServerAdapter(BasePlatformAdapter): if draining is not None: return draining - try: - body = await request.json() - except Exception: - body = {} - job_id = (body or {}).get("job_id") - if not job_id: - return web.json_response({"error": "missing job_id"}, status=400) + with _reserve_pending_api_work(self) as reservation: + try: + body = await request.json() + except Exception: + body = {} + job_id = (body or {}).get("job_id") + if not job_id: + return web.json_response({"error": "missing job_id"}, status=400) - from cron.scheduler_provider import resolve_cron_scheduler - provider = resolve_cron_scheduler() + from cron.scheduler_provider import resolve_cron_scheduler + provider = resolve_cron_scheduler() - loop = asyncio.get_running_loop() - # Fire in the background (202 immediately). fire_due claims via the - # store CAS, so a retry while this is in flight is de-duped. - task = asyncio.create_task( - asyncio.to_thread(provider.fire_due, job_id, adapters=None, loop=loop) - ) - try: - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) - except (TypeError, AttributeError): - pass + loop = asyncio.get_running_loop() + # Fire in the background (202 immediately). fire_due claims via the + # store CAS, so a retry while this is in flight is de-duped. + task = asyncio.create_task( + asyncio.to_thread(provider.fire_due, job_id, adapters=None, loop=loop) + ) + reservation["detached"] = True + task.add_done_callback( + lambda _task: _release_pending_api_work(self, reservation) + ) + try: + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + except (TypeError, AttributeError): + pass - return web.json_response({"status": "accepted", "job_id": job_id}, status=202) + return web.json_response({"status": "accepted", "job_id": job_id}, status=202) # ------------------------------------------------------------------ diff --git a/tests/gateway/test_cron_fire_webhook.py b/tests/gateway/test_cron_fire_webhook.py index fad9f4de1021..ae745f98c0be 100644 --- a/tests/gateway/test_cron_fire_webhook.py +++ b/tests/gateway/test_cron_fire_webhook.py @@ -8,6 +8,7 @@ test_chronos_verify.py. """ import asyncio +import threading from types import SimpleNamespace from unittest.mock import patch @@ -135,6 +136,61 @@ async def test_valid_token_refuses_during_gateway_drain(adapter, monkeypatch): assert spy.fired == [] +@pytest.mark.asyncio +async def test_valid_fire_reservation_blocks_drain_before_body_and_task(adapter, monkeypatch): + runner = SimpleNamespace(_draining=False, _external_drain_active=False) + body_started = asyncio.Event() + release_body = asyncio.Event() + fired = threading.Event() + release_fire = threading.Event() + + class BlockingProvider: + def fire_due(self, job_id, *, adapters=None, loop=None): + fired.set() + release_fire.wait(timeout=2) + return True + + original_json = web.Request.json + + async def delayed_json(request): + body_started.set() + await release_body.wait() + return await original_json(request) + + monkeypatch.setattr("cron.scheduler_provider.resolve_cron_scheduler", BlockingProvider) + monkeypatch.setattr( + "plugins.cron_providers.chronos.verify.get_fire_verifier", + lambda: (lambda **kw: {"purpose": "cron_fire"}), + ) + app = _create_app(adapter) + with patch("gateway.run._gateway_runner_ref", lambda: runner), patch.object( + web.Request, "json", delayed_json + ): + async with TestClient(TestServer(app)) as cli: + request_task = asyncio.create_task( + cli.post( + "/api/cron/fire", + headers={"Authorization": "Bearer good"}, + json={"job_id": "abc123"}, + ) + ) + await body_started.wait() + assert adapter.active_agent_work_count() == 1 + + release_body.set() + response = await request_task + assert response.status == 202 + await asyncio.to_thread(fired.wait, 2) + assert adapter.active_agent_work_count() == 1 + release_fire.set() + for _ in range(50): + if adapter.active_agent_work_count() == 0: + break + await asyncio.sleep(0.01) + + assert adapter.active_agent_work_count() == 0 + + @pytest.mark.asyncio async def test_missing_job_id_400(adapter, monkeypatch): """Valid token but no job_id → 400, no fire."""