fix(api): reserve cron fire work during drain

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)
 
 
     # ------------------------------------------------------------------

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."""