"""Several people asking at once, against an endpoint with a few slots. The gate is the only place that bounds how much work reaches an endpoint, so these tests are about the three answers it can give: go, wait, or busy. """ import asyncio from collections.abc import Iterator import pytest from app.config import get_settings from app.llm import gate from app.llm.errors import LLMError ENDPOINT = "http://endpoint.test/v1" @pytest.fixture(autouse=True) def _fresh_gate() -> Iterator[None]: gate.reset() yield gate.reset() def _limits( monkeypatch: pytest.MonkeyPatch, *, parallel: int, wait: float = 20.0, queued: int = 24, ) -> None: settings = get_settings() monkeypatch.setattr(settings, "llm_max_parallel", parallel) monkeypatch.setattr(settings, "llm_queue_wait_seconds", wait) monkeypatch.setattr(settings, "llm_max_queued", queued) gate.reset() async def test_only_as_many_calls_reach_the_endpoint_as_it_has_slots( monkeypatch: pytest.MonkeyPatch, ) -> None: _limits(monkeypatch, parallel=2) concurrent = 0 peak = 0 release = asyncio.Event() async def call() -> None: nonlocal concurrent, peak async with gate.slot(ENDPOINT, "chat"): concurrent += 1 peak = max(peak, concurrent) await release.wait() concurrent -= 1 tasks = [asyncio.create_task(call()) for _ in range(6)] await asyncio.sleep(0) # let everyone reach the gate assert peak == 2, "more calls were in flight than the endpoint has slots" release.set() await asyncio.gather(*tasks) assert peak == 2 async def test_a_waiting_call_gets_the_slot_the_previous_one_frees( monkeypatch: pytest.MonkeyPatch, ) -> None: _limits(monkeypatch, parallel=1) order: list[str] = [] first_in = asyncio.Event() let_go = asyncio.Event() async def first() -> None: async with gate.slot(ENDPOINT, "chat"): order.append("first in") first_in.set() await let_go.wait() order.append("first out") async def second() -> None: await first_in.wait() async with gate.slot(ENDPOINT, "chat"): order.append("second in") task_one = asyncio.create_task(first()) task_two = asyncio.create_task(second()) await first_in.wait() await asyncio.sleep(0) assert order == ["first in"], "the second call did not wait" let_go.set() await asyncio.gather(task_one, task_two) assert order == ["first in", "first out", "second in"] async def test_waiting_longer_than_allowed_is_reported_as_busy( monkeypatch: pytest.MonkeyPatch, ) -> None: """A caller that cannot be served soon is told so, rather than being held until the HTTP timeout makes it look like a broken endpoint.""" _limits(monkeypatch, parallel=1, wait=0.05) let_go = asyncio.Event() async def holder() -> None: async with gate.slot(ENDPOINT, "chat"): await let_go.wait() held = asyncio.create_task(holder()) await asyncio.sleep(0) with pytest.raises(LLMError) as raised: async with gate.slot(ENDPOINT, "chat"): pass assert raised.value.code == "llm_busy" let_go.set() await held async def test_beyond_the_queue_limit_the_answer_is_immediate( monkeypatch: pytest.MonkeyPatch, ) -> None: """Once far more work has arrived than the endpoint can absorb, the useful reply is "busy" now — not "busy" in twenty seconds.""" _limits(monkeypatch, parallel=1, wait=5.0, queued=2) let_go = asyncio.Event() async def occupy() -> None: async with gate.slot(ENDPOINT, "chat"): await let_go.wait() async def queue_up() -> None: async with gate.slot(ENDPOINT, "chat"): pass holder = asyncio.create_task(occupy()) await asyncio.sleep(0) waiters = [asyncio.create_task(queue_up()) for _ in range(2)] await asyncio.sleep(0) started = asyncio.get_running_loop().time() with pytest.raises(LLMError) as raised: async with gate.slot(ENDPOINT, "chat"): pass assert raised.value.code == "llm_busy" assert asyncio.get_running_loop().time() - started < 1.0 let_go.set() await asyncio.gather(holder, *waiters) async def test_endpoints_do_not_share_a_limit(monkeypatch: pytest.MonkeyPatch) -> None: """The chat and embedding roles usually run on different servers; a busy chat endpoint must not stop retrieval from embedding a query.""" _limits(monkeypatch, parallel=1) let_go = asyncio.Event() async def occupy() -> None: async with gate.slot(ENDPOINT, "chat"): await let_go.wait() holder = asyncio.create_task(occupy()) await asyncio.sleep(0) async with gate.slot("http://other.test/v1", "embedding"): pass # reached its own slot while the first endpoint is full let_go.set() await holder async def test_a_failed_call_gives_its_slot_back( monkeypatch: pytest.MonkeyPatch, ) -> None: _limits(monkeypatch, parallel=1, wait=0.05) with pytest.raises(RuntimeError): async with gate.slot(ENDPOINT, "chat"): raise RuntimeError("endpoint blew up") # The slot is free again, so the next caller is served rather than queued. async with gate.slot(ENDPOINT, "chat"): pass async def test_the_mode_can_tell_whether_a_turn_will_have_to_wait( monkeypatch: pytest.MonkeyPatch, ) -> None: _limits(monkeypatch, parallel=1) assert gate.endpoint_busy(ENDPOINT) is False let_go = asyncio.Event() async def occupy() -> None: async with gate.slot(ENDPOINT, "chat"): await let_go.wait() holder = asyncio.create_task(occupy()) await asyncio.sleep(0) assert gate.endpoint_busy(ENDPOINT) is True let_go.set() await holder assert gate.endpoint_busy(ENDPOINT) is False