hermes-agent/tests/gateway/test_api_server_active_work_drain.py

"""Regression coverage for #63529 API-server shutdown draining.

API-server work is adapter-owned rather than tracked by
``GatewayRunner._running_agents``. The shutdown drain must account for the
same live state as the API concurrency limiter, including a ``/v1/runs`` task
that exists before its agent has been constructed, and it must refuse new API
turns once the gateway starts draining.
"""

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer

from gateway.config import Platform, PlatformConfig
from gateway.platforms.api_server import APIServerAdapter
from tests.gateway.restart_test_helpers import make_restart_runner


class _RunTask:
    def __init__(self, done: bool = False):
        self._done = done

    def done(self) -> bool:
        return self._done


def _make_api_adapter(*, inflight: int = 0, queued_ids=()):
    tasks = {run_id: _RunTask() for run_id in queued_ids}
    adapter = SimpleNamespace(
        platform=Platform.API_SERVER,
        _inflight_agent_runs=inflight,
        _active_run_tasks=tasks,
    )

    def active_agent_work_count() -> int:
        return int(getattr(adapter, "_pending_agent_requests", 0)) + int(
            adapter._inflight_agent_runs
        ) + sum(not task.done() for task in adapter._active_run_tasks.values())

    adapter.active_agent_work_count = active_agent_work_count
    return adapter


def _make_admission_app(adapter: APIServerAdapter) -> web.Application:
    app = web.Application()
    app.router.add_post("/api/sessions/{session_id}/chat", adapter._handle_session_chat)
    app.router.add_post(
        "/api/sessions/{session_id}/chat/stream", adapter._handle_session_chat_stream
    )
    app.router.add_post("/v1/chat/completions", adapter._handle_chat_completions)
    app.router.add_post("/v1/responses", adapter._handle_responses)
    app.router.add_post("/v1/runs", adapter._handle_runs)
    return app


class TestActiveApiRunCount:
    def test_zero_when_no_api_adapter(self):
        runner, _adapter = make_restart_runner()
        runner.adapters = {}
        assert runner._active_api_run_count() == 0


class TestAPIServerAdapterWorkCount:

    @pytest.mark.asyncio
    async def test_concurrency_limit_excludes_current_pending_admission(self):
        adapter = APIServerAdapter(PlatformConfig(enabled=True))
        adapter._max_concurrent_runs = 1
        app = _make_admission_app(adapter)

        async with TestClient(TestServer(app)) as client:
            with patch.object(adapter, "_run_agent", new=AsyncMock(return_value=({}, {}))):
                response = await client.post(
                    "/api/sessions/s/chat",
                    json={"message": "hello"},
                )

        assert response.status == 404


    def test_counts_live_run_task_before_agent_creation(self):
        adapter = APIServerAdapter(PlatformConfig(enabled=True))
        adapter._inflight_agent_runs = 2
        adapter._active_run_tasks = {
            "queued": _RunTask(),
            "finished": _RunTask(done=True),
        }
        adapter._active_run_agents = {}

        assert adapter.active_agent_work_count() == 3


class TestDrainWaitsForApiWork:

    @pytest.mark.asyncio
    async def test_drain_waits_for_real_queued_run_before_agent_creation(self):
        """A live /v1/runs task must block drain before it has an agent."""
        runner, _adapter = make_restart_runner()
        api = APIServerAdapter(PlatformConfig(enabled=True))
        runner.adapters = {Platform.API_SERVER: api}
        app = _make_admission_app(api)
        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())

        mock_agent = MagicMock()
        mock_agent.run_conversation.return_value = {"final_response": "done"}
        mock_agent.session_prompt_tokens = 0
        mock_agent.session_completion_tokens = 0
        mock_agent.session_total_tokens = 0

        with patch(
            "gateway.platforms.api_server.asyncio.create_task",
            side_effect=delayed_create_task,
        ), patch.object(api, "_create_agent", return_value=mock_agent):
            async with TestClient(TestServer(app)) as client:
                response = await client.post("/v1/runs", json={"input": "hello"})
                assert response.status == 202
                await task_started.wait()

                assert api._active_run_agents == {}
                assert runner._active_api_run_count() == 1
                drain_task = original_create_task(runner._drain_active_agents(2.0))
                await asyncio.sleep(0.1)
                assert not drain_task.done()

                allow_task.set()
                _snapshot, timed_out = await drain_task

        assert timed_out is False


class TestDrainAdmission:
    @pytest.mark.asyncio
    async def test_drain_refuses_every_agent_start_endpoint(self):
        adapter = APIServerAdapter(PlatformConfig(enabled=True))
        runner = SimpleNamespace(_draining=True, _external_drain_active=False)
        app = _make_admission_app(adapter)
        paths = (
            "/api/sessions/missing/chat",
            "/api/sessions/missing/chat/stream",
            "/v1/chat/completions",
            "/v1/responses",
            "/v1/runs",
        )

        with patch("gateway.run._gateway_runner_ref", lambda: runner):
            async with TestClient(TestServer(app)) as client:
                for path in paths:
                    response = await client.post(path, json={})
                    payload = await response.json()

                    assert response.status == 503
                    assert response.headers["Retry-After"] == "1"
                    assert payload["error"]["code"] == "gateway_draining"