diff --git a/backend/README.md b/backend/README.md index 2e876cc..9c1344b 100644 --- a/backend/README.md +++ b/backend/README.md @@ -9,6 +9,7 @@ backend/ app.py # FastAPI application factory, session cookies, security headers, endpoints settings.py # Validated deployment settings and URL validation/canonicalization runner.py # QueryRunner orchestrating deadlines, gate, disconnect cancellation, cleanup + admission.py # In-memory admission queue (FIFO tickets, reservation, estimate) transport.py # Bounded NDJSON bridge reader/writer and protocol state machine (Rev 1) containers.py # Rootless Docker container lifecycle, tmpfs mounts, kill/remove, reconciliation confluence.py # Request-scoped Confluence client, streaming bounds, and tool dispatcher @@ -47,6 +48,9 @@ backend/ | `CONFLUENCE_WEB_QUERY_TIMEOUT_SECONDS` | `180.0` | Total query execution deadline; must not exceed `CONFLUENCE_WEB_MAX_DEADLINE_SECONDS`. | | `CONFLUENCE_WEB_MAX_DEADLINE_SECONDS` | `900.0` | Protocol maximum for one query (60–3600). Passed into the agent container so the supervisor and bridge enforce the same bound. | | `CONFLUENCE_WEB_CLEANUP_TIMEOUT_SECONDS` | `10.0` | Dedicated cleanup timeout. | +| `CONFLUENCE_WEB_QUEUE_RESERVATION_SECONDS` | `45` | Admission queue reservation window after promotion; allowed 30-60. See [QUEUE_SPECIFICATION.md](../docs/QUEUE_SPECIFICATION.md). | +| `CONFLUENCE_WEB_QUEUE_HEARTBEAT_SECONDS` | `15` | Admission queue heartbeat timeout for queued (not yet reserved) tickets; allowed 5-60. | +| `CONFLUENCE_WEB_QUEUE_MAX_LENGTH` | `20` | Maximum queued tickets, excluding the reserved and running sessions; allowed 1-100. | ## Running the Service diff --git a/backend/admission.py b/backend/admission.py new file mode 100644 index 0000000..d44f97d --- /dev/null +++ b/backend/admission.py @@ -0,0 +1,295 @@ +"""In-memory admission queue for the single-runner query gate. + +Implements the state machine and estimate from docs/QUEUE_SPECIFICATION.md: +one runner slot that is idle, reserved for a single session, or running for +a single session, with a FIFO queue of waiting tickets behind it. All public +methods are synchronous and perform no suspension, so a single call is +atomic under asyncio's cooperative scheduling -- the same reasoning +backend.runner.QueryRunner relies on for its own gate. +""" + +from __future__ import annotations + +import secrets +import time +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional + +from backend.errors import QueueFullError, TicketNotFoundError + +# Estimate smoothing factor and history window (spec section 4/5 constants). +ESTIMATE_ALPHA = 0.3 +ESTIMATE_HISTORY = 20 + +# Runs cancelled by the client before this many seconds do not update the estimate. +CANCELLED_RUN_FLOOR_SECONDS = 5.0 + +TicketState = str # "queued" | "ready" | "running" + + +@dataclass +class Ticket: + ticket_id: str + session_id: str + joined_at: float + last_seen_at: float + state: TicketState + reserved_at: Optional[float] = None + expires_at: Optional[float] = None + + +def _generate_ticket_id() -> str: + # Opaque, case-sensitive, [A-Za-z0-9_-] per CONTRACTS section 1. + return "q_" + secrets.token_urlsafe(12) + + +class AdmissionController: + """FIFO admission queue in front of the single query runner slot.""" + + def __init__( + self, + reservation_seconds: float, + heartbeat_seconds: float, + max_queue_length: int, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self.reservation_seconds = reservation_seconds + self.heartbeat_seconds = heartbeat_seconds + self.max_queue_length = max_queue_length + self.clock = clock + + self._tickets: Dict[str, Ticket] = {} + self._queue: List[str] = [] # session_ids in "queued" state, FIFO + self._reserved_session: Optional[str] = None + self._running_session: Optional[str] = None + self._running_started_at: Optional[float] = None + + self._run_durations: List[float] = [] + self._mean_duration: Optional[float] = None + + # -- Public API ----------------------------------------------------- + + def join(self, session_id: str) -> dict: + """Idempotent join. Returns the (possibly pre-existing) ticket view.""" + self._expire() + existing = self._tickets.get(session_id) + if existing is not None: + return self._view(existing) + + now = self.clock() + if self._reserved_session is None and self._running_session is None: + # Runner is idle: book the place directly. + ticket = Ticket( + ticket_id=_generate_ticket_id(), + session_id=session_id, + joined_at=now, + last_seen_at=now, + state="ready", + reserved_at=now, + expires_at=now + self.reservation_seconds, + ) + self._tickets[session_id] = ticket + self._reserved_session = session_id + return self._view(ticket) + + if len(self._queue) >= self.max_queue_length: + raise QueueFullError("The queue is full; please try again later") + + ticket = Ticket( + ticket_id=_generate_ticket_id(), + session_id=session_id, + joined_at=now, + last_seen_at=now, + state="queued", + ) + self._tickets[session_id] = ticket + self._queue.append(session_id) + return self._view(ticket) + + def status(self, session_id: str) -> dict: + """Heartbeat poll. Refreshes last_seen_at; raises TicketNotFoundError if absent.""" + self._expire() + ticket = self._tickets.get(session_id) + if ticket is None: + raise TicketNotFoundError("No active queue ticket for this session") + ticket.last_seen_at = self.clock() + return self._view(ticket) + + def leave(self, session_id: str) -> None: + """Drop the session's ticket, if any. Promotes the next ticket if it held the reservation.""" + self._expire() + ticket = self._tickets.pop(session_id, None) + if ticket is None: + return + if session_id in self._queue: + self._queue.remove(session_id) + if self._reserved_session == session_id: + self._reserved_session = None + self._promote() + + def claim(self, session_id: str) -> bool: + """Attempt to move a query request's session into the running slot. + + Returns True if the session may proceed (claimed or implicitly joined), + False if it must be refused with 409 busy. Does not disturb the queue + on refusal. + """ + now = self.clock() + + # The reserved session's own claim must see 409 once its window has + # passed, even when this very call is what discovers the expiry -- + # it must not fall through to idle/implicit-join for itself. + if self._reserved_session == session_id: + ticket = self._tickets.get(session_id) + expired = ticket is None or ( + ticket.expires_at is not None and now >= ticket.expires_at + ) + if expired: + self._tickets.pop(session_id, None) + self._reserved_session = None + self._promote() + return False + ticket.state = "running" + ticket.reserved_at = None + ticket.expires_at = None + self._reserved_session = None + self._running_session = session_id + self._running_started_at = now + return True + + # Not the current holder: evaluate lazy expiries (a stale reservation + # held by someone else, or stale queued heartbeats) before deciding. + self._expire() + + if self._running_session is not None: + return False + if self._reserved_session is not None: + return False + + # Idle: implicit join, reserve, and claim in one step. + ticket = Ticket( + ticket_id=_generate_ticket_id(), + session_id=session_id, + joined_at=now, + last_seen_at=now, + state="running", + ) + self._tickets[session_id] = ticket + self._running_session = session_id + self._running_started_at = now + return True + + def run_end(self, session_id: str, duration_seconds: float, cancelled: bool = False) -> None: + """Record run completion (success, failure, timeout, cancellation, disconnect). + + Must be called after the runner's own cleanup has finished so the next + holder's reservation window never overlaps residual cleanup. + """ + self._tickets.pop(session_id, None) + if self._running_session == session_id: + self._running_session = None + self._running_started_at = None + + if not (cancelled and duration_seconds < CANCELLED_RUN_FLOOR_SECONDS): + self._record_duration(duration_seconds) + + self._promote() + + def expire(self) -> None: + """Public hook for periodic maintenance; evaluates lazy expiries eagerly.""" + self._expire() + + # -- Internal --------------------------------------------------------- + + def _record_duration(self, duration_seconds: float) -> None: + self._run_durations.append(max(0.0, duration_seconds)) + if len(self._run_durations) > ESTIMATE_HISTORY: + self._run_durations = self._run_durations[-ESTIMATE_HISTORY:] + mean = self._run_durations[0] + for d in self._run_durations[1:]: + mean = ESTIMATE_ALPHA * d + (1 - ESTIMATE_ALPHA) * mean + self._mean_duration = mean + + def _promote(self) -> None: + if self._reserved_session is not None or self._running_session is not None: + return + if not self._queue: + return + session_id = self._queue.pop(0) + ticket = self._tickets.get(session_id) + if ticket is None: + # Should not happen; keep state consistent and try the next one. + self._promote() + return + now = self.clock() + ticket.state = "ready" + ticket.reserved_at = now + ticket.expires_at = now + self.reservation_seconds + self._reserved_session = session_id + + def _expire(self) -> None: + now = self.clock() + + # Reservation expiry. + if self._reserved_session is not None: + ticket = self._tickets.get(self._reserved_session) + if ticket is None or (ticket.expires_at is not None and now >= ticket.expires_at): + if ticket is not None: + self._tickets.pop(self._reserved_session, None) + self._reserved_session = None + self._promote() + + # Heartbeat expiry for queued tickets only. + stale = [ + sid + for sid in self._queue + if (ticket := self._tickets.get(sid)) is not None + and now - ticket.last_seen_at > self.heartbeat_seconds + ] + for sid in stale: + self._tickets.pop(sid, None) + if sid in self._queue: + self._queue.remove(sid) + + def _position(self, session_id: str, ticket: Ticket) -> int: + if ticket.state == "running": + return 0 + front: List[str] = [] + if self._reserved_session is not None: + front.append(self._reserved_session) + front.extend(self._queue) + try: + return front.index(session_id) + 1 + except ValueError: + return 0 + + def _eta_seconds(self, ahead: int) -> Optional[float]: + if self._mean_duration is None: + return None + if self._running_session is not None: + started_at = self._running_started_at if self._running_started_at is not None else self.clock() + elapsed = self.clock() - started_at + return max(0.0, self._mean_duration * (ahead + 1) - elapsed) + return max(0.0, self._mean_duration * (ahead + 1)) + + def _view(self, ticket: Ticket) -> dict: + position = self._position(ticket.session_id, ticket) + ahead = max(0, position - 1) + eta = self._eta_seconds(ahead) + runner = "running" if self._running_session is not None else "reserved" + + reservation_expires_in: Optional[float] = None + if ticket.state == "ready" and ticket.expires_at is not None: + reservation_expires_in = max(0.0, ticket.expires_at - self.clock()) + + return { + "ticket_id": ticket.ticket_id, + "status": ticket.state, + "position": position, + "ahead": ahead, + "eta_seconds": round(eta) if eta is not None else None, + "reservation_expires_in_seconds": ( + round(reservation_expires_in) if reservation_expires_in is not None else None + ), + "runner": runner, + } diff --git a/backend/app.py b/backend/app.py index fca91f1..74f259d 100644 --- a/backend/app.py +++ b/backend/app.py @@ -17,6 +17,7 @@ from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, ConfigDict from starlette.exceptions import HTTPException as StarletteHTTPException +from backend.admission import AdmissionController from backend.artifacts import ArtifactStore from backend.confluence import ConfluenceClient from backend.containers import ( @@ -27,11 +28,13 @@ from backend.containers import ( from backend.errors import ( AppError, ArtifactNotFoundError, + BusyError, InvalidInputError, OriginDeniedError, ConnectivityFailedError, sanitize_message, RequestTooLargeError, + TicketNotFoundError, ) from backend.model import FakeModelAdapter, ModelAdapter, OpenAIModelAdapter from backend.runner import QueryRunner @@ -71,6 +74,12 @@ class QueryRequest(BaseModel): credentials: QueryCredentials +class QueueJoinRequest(BaseModel): + """Body is `{}`. extra="forbid" rejects any field, including credentials/pat.""" + + model_config = ConfigDict(extra="forbid") + + class ServerSessionStore: """In-memory store of server-issued session IDs to prevent client-forged sessions.""" @@ -163,7 +172,7 @@ def check_same_origin(request: Request) -> None: raise OriginDeniedError("Origin does not match request destination") -async def _periodic_maintenance(runner, store, session_store) -> None: +async def _periodic_maintenance(runner, store, session_store, admission) -> None: while True: await asyncio.sleep(60) if not runner._gate.locked(): @@ -174,6 +183,7 @@ async def _periodic_maintenance(runner, store, session_store) -> None: try: await asyncio.to_thread(store.expire_artifacts) session_store.prune_stale() + admission.expire() except Exception: logger.warning("Artifact or session maintenance failed") @@ -266,6 +276,7 @@ def create_app( artifact_store: Optional[ArtifactStore] = None, model_adapter: Optional[ModelAdapter] = None, confluence_client_factory: Optional[Callable[..., ConfluenceClient]] = None, + admission_controller: Optional[AdmissionController] = None, ) -> FastAPI: """Create and configure FastAPI application with dependency injection.""" app_settings = settings or Settings.from_env() @@ -311,6 +322,11 @@ def create_app( ) session_store = ServerSessionStore() + admission = admission_controller or AdmissionController( + reservation_seconds=app_settings.queue_reservation_seconds, + heartbeat_seconds=app_settings.queue_heartbeat_seconds, + max_queue_length=app_settings.queue_max_length, + ) @contextlib.asynccontextmanager async def lifespan(fastapi_app: FastAPI): @@ -320,7 +336,7 @@ def create_app( if isinstance(mgr, DockerContainerManager): await mgr.verify_rootless() await runner.reconcile() - maintenance_task = asyncio.create_task(_periodic_maintenance(runner, store, session_store)) + maintenance_task = asyncio.create_task(_periodic_maintenance(runner, store, session_store, admission)) yield finally: if maintenance_task is not None: @@ -393,6 +409,7 @@ def create_app( app.add_middleware(SecurityMiddleware) app.state.runner = runner app.state.artifact_store = store + app.state.admission = admission @app.get("/") async def root_index(request: Request): @@ -437,6 +454,44 @@ def create_app( get_or_set_session_cookie(request, resp, session_store) return resp + @app.post("/api/v1/queue/join") + async def queue_join(req: QueueJoinRequest, request: Request): + check_same_origin(request) + session_id = request.cookies.get(SESSION_COOKIE_NAME) + if not session_store.is_valid(session_id): + session_id = session_store.create_session() + else: + session_store.touch(session_id) # type: ignore[arg-type] + + view = admission.join(session_id) # type: ignore[arg-type] + resp = JSONResponse(content=view) + secure = request.url.scheme == "https" + resp.set_cookie( + key=SESSION_COOKIE_NAME, + value=session_id, + httponly=True, + samesite="strict", + path="/", + secure=secure, + ) + return resp + + @app.get("/api/v1/queue/status") + async def queue_status(request: Request): + session_id = request.cookies.get(SESSION_COOKIE_NAME) + if not session_store.is_valid(session_id): + raise TicketNotFoundError("No active queue ticket for this session") + view = admission.status(session_id) # type: ignore[arg-type] + return JSONResponse(content=view) + + @app.delete("/api/v1/queue/ticket") + async def queue_leave(request: Request): + check_same_origin(request) + session_id = request.cookies.get(SESSION_COOKIE_NAME) + if session_store.is_valid(session_id): + admission.leave(session_id) # type: ignore[arg-type] + return Response(status_code=204) + @app.post("/api/v1/query") async def execute_query(req: QueryRequest, request: Request): check_same_origin(request) @@ -446,14 +501,32 @@ def create_app( else: session_store.touch(session_id) # type: ignore[arg-type] - # Execute query via runner - result = await runner.run( - prompt=req.prompt, - confluence_url=req.credentials.url, - confluence_pat=req.credentials.pat, - session_id=session_id, - is_disconnected=request.is_disconnected, - ) + # Admission: claim the reservation (or implicit-join while idle) before + # any work starts. The runner's own gate remains a last line of defence. + if not admission.claim(session_id): + raise BusyError("Another session holds the reservation") + + run_started = admission.clock() + result: Optional[Dict[str, Any]] = None + cancelled = False + try: + result = await runner.run( + prompt=req.prompt, + confluence_url=req.credentials.url, + confluence_pat=req.credentials.pat, + session_id=session_id, + is_disconnected=request.is_disconnected, + ) + except asyncio.CancelledError: + cancelled = True + raise + finally: + # Run end (success, failure, timeout, cancellation, or disconnect): + # record duration and promote the next ticket after the runner's + # own cleanup budget above has already completed. + duration = result["duration_seconds"] if result is not None else (admission.clock() - run_started) + admission.run_end(session_id, duration_seconds=duration, cancelled=cancelled) + resp = JSONResponse(content=result) # Ensure session cookie is set secure = request.url.scheme == "https" diff --git a/backend/errors.py b/backend/errors.py index 7bfc549..92cadf3 100644 --- a/backend/errors.py +++ b/backend/errors.py @@ -125,3 +125,13 @@ class ExecutionFailedError(AppError): class CleanupFailedError(AppError): def __init__(self, message: str = "Container cleanup could not be confirmed"): super().__init__("cleanup_failed", message, status_code=500) + + +class QueueFullError(AppError): + def __init__(self, message: str = "The queue is full"): + super().__init__("queue_full", message, status_code=503) + + +class TicketNotFoundError(AppError): + def __init__(self, message: str = "Ticket not found"): + super().__init__("ticket_not_found", message, status_code=404) diff --git a/backend/settings.py b/backend/settings.py index cf59433..f100e71 100644 --- a/backend/settings.py +++ b/backend/settings.py @@ -69,6 +69,11 @@ class Settings(BaseModel): # Development mode (explicit opt-in) dev_mode: bool = False + # Admission queue (docs/QUEUE_SPECIFICATION.md section 4) + queue_reservation_seconds: float = 45.0 + queue_heartbeat_seconds: float = 15.0 + queue_max_length: int = 20 + @field_validator("approved_confluence_origins") @classmethod def validate_origins(cls, v: List[str]) -> List[str]: @@ -123,6 +128,27 @@ class Settings(BaseModel): raise ValueError("query_timeout_seconds must not exceed max_deadline_seconds") return self + @field_validator("queue_reservation_seconds") + @classmethod + def validate_queue_reservation_seconds(cls, value): + if not (30.0 <= value <= 60.0): + raise ValueError("queue_reservation_seconds must be between 30 and 60") + return value + + @field_validator("queue_heartbeat_seconds") + @classmethod + def validate_queue_heartbeat_seconds(cls, value): + if not (5.0 <= value <= 60.0): + raise ValueError("queue_heartbeat_seconds must be between 5 and 60") + return value + + @field_validator("queue_max_length") + @classmethod + def validate_queue_max_length(cls, value): + if not (1 <= value <= 100): + raise ValueError("queue_max_length must be between 1 and 100") + return value + @classmethod def from_env(cls) -> Settings: """Load settings from environment variables.""" @@ -210,6 +236,13 @@ class Settings(BaseModel): cleanup_timeout_seconds=_parse_float( "CONFLUENCE_WEB_CLEANUP_TIMEOUT_SECONDS", 10.0 ), + queue_reservation_seconds=_parse_float( + "CONFLUENCE_WEB_QUEUE_RESERVATION_SECONDS", 45.0 + ), + queue_heartbeat_seconds=_parse_float( + "CONFLUENCE_WEB_QUEUE_HEARTBEAT_SECONDS", 15.0 + ), + queue_max_length=_parse_int("CONFLUENCE_WEB_QUEUE_MAX_LENGTH", 20), ) diff --git a/deploy/confluence-web.env.example b/deploy/confluence-web.env.example index 3c8ee52..7473644 100644 --- a/deploy/confluence-web.env.example +++ b/deploy/confluence-web.env.example @@ -51,5 +51,13 @@ CONFLUENCE_WEB_QUERY_TIMEOUT_SECONDS=600 CONFLUENCE_WEB_MAX_DEADLINE_SECONDS=900 CONFLUENCE_WEB_CLEANUP_TIMEOUT_SECONDS=10 +# --- Admission queue (docs/QUEUE_SPECIFICATION.md); defaults shown, allowed ranges in comments --- +# Reservation window after promotion, seconds (30-60). +#CONFLUENCE_WEB_QUEUE_RESERVATION_SECONDS=45 +# Heartbeat timeout for queued (not yet reserved) tickets, seconds (5-60). +#CONFLUENCE_WEB_QUEUE_HEARTBEAT_SECONDS=15 +# Maximum queued tickets, excluding the reserved and running sessions (1-100). +#CONFLUENCE_WEB_QUEUE_MAX_LENGTH=20 + # Never set in production. Substitutes the container, model and Confluence with fakes. CONFLUENCE_WEB_DEV_MODE=false diff --git a/tests/backend/test_admission.py b/tests/backend/test_admission.py new file mode 100644 index 0000000..43515d7 --- /dev/null +++ b/tests/backend/test_admission.py @@ -0,0 +1,296 @@ +"""Unit tests for backend.admission.AdmissionController (docs/QUEUE_SPECIFICATION.md).""" + +from __future__ import annotations + +import pytest + +from backend.admission import AdmissionController +from backend.errors import QueueFullError, TicketNotFoundError + + +class FakeClock: + """Injectable monotonic clock for deterministic queue tests.""" + + def __init__(self, start: float = 0.0): + self._now = start + + def __call__(self) -> float: + return self._now + + def advance(self, seconds: float) -> None: + self._now += seconds + + +def make_controller(reservation_seconds=45.0, heartbeat_seconds=15.0, max_queue_length=20, clock=None): + clock = clock or FakeClock(0.0) + return AdmissionController( + reservation_seconds=reservation_seconds, + heartbeat_seconds=heartbeat_seconds, + max_queue_length=max_queue_length, + clock=clock, + ), clock + + +def run_cycle(ctl: AdmissionController, session_id: str, duration: float, cancelled: bool = False) -> None: + """Join (idle), claim, and end a run for session_id. Assumes the runner is idle.""" + view = ctl.join(session_id) + assert view["status"] == "ready" + assert ctl.claim(session_id) is True + ctl.run_end(session_id, duration_seconds=duration, cancelled=cancelled) + + +# -- Join / idle / concurrent queueing -------------------------------------- + + +def test_join_while_idle_reserves_directly(): + ctl, clock = make_controller(reservation_seconds=45.0) + view = ctl.join("s1") + assert view["status"] == "ready" + assert view["position"] == 1 + assert view["ahead"] == 0 + assert view["runner"] == "reserved" + assert view["reservation_expires_in_seconds"] == 45 + assert view["eta_seconds"] is None + assert view["ticket_id"].startswith("q_") + + +def test_concurrent_second_join_is_queued(): + ctl, clock = make_controller() + ctl.join("s1") + view2 = ctl.join("s2") + assert view2["status"] == "queued" + assert view2["position"] == 2 + assert view2["ahead"] == 1 + assert view2["runner"] == "reserved" + assert view2["reservation_expires_in_seconds"] is None + + +def test_join_is_idempotent_per_session(): + ctl, clock = make_controller() + ctl.join("s1") + view_a = ctl.join("s2") + view_b = ctl.join("s2") + assert view_a["ticket_id"] == view_b["ticket_id"] + assert view_b["status"] == "queued" + + +# -- FIFO order and promotion ------------------------------------------------ + + +def test_fifo_order_and_promotion_on_completion(): + ctl, clock = make_controller() + ctl.join("a") + ctl.join("b") + ctl.join("c") + + assert ctl.claim("a") is True + # "a" is now running, not reserved; it does not occupy a front-list slot. + assert ctl.status("b")["position"] == 1 + assert ctl.status("c")["position"] == 2 + + ctl.run_end("a", duration_seconds=10.0, cancelled=False) + + status_b = ctl.status("b") + assert status_b["status"] == "ready" + assert status_b["position"] == 1 + assert status_b["ahead"] == 0 + + status_c = ctl.status("c") + assert status_c["status"] == "queued" + assert status_c["position"] == 2 + assert status_c["ahead"] == 1 + + with pytest.raises(TicketNotFoundError): + ctl.status("a") + + +def test_promotion_on_failure_and_timeout_paths(): + # run_end is called identically for success, failure, and timeout by app.py; + # the controller has no notion of outcome beyond "cancelled". + ctl, clock = make_controller() + ctl.join("a") + ctl.join("b") + ctl.claim("a") + ctl.run_end("a", duration_seconds=5.0, cancelled=False) + assert ctl.status("b")["status"] == "ready" + + +def test_promotion_on_cancellation(): + ctl, clock = make_controller() + ctl.join("a") + ctl.join("b") + ctl.claim("a") + ctl.run_end("a", duration_seconds=1.0, cancelled=True) + assert ctl.status("b")["status"] == "ready" + + +def test_promotion_on_leave(): + ctl, clock = make_controller() + ctl.join("a") + ctl.join("b") + ctl.leave("a") + assert ctl.status("b")["status"] == "ready" + + +def test_leave_is_idempotent(): + ctl, clock = make_controller() + ctl.leave("never-joined") # no error + ctl.join("a") + ctl.leave("a") + with pytest.raises(TicketNotFoundError): + ctl.status("a") + ctl.leave("a") # still no error + + +# -- Reservation and heartbeat expiry ---------------------------------------- + + +def test_reservation_expiry_drops_holder_and_promotes_next(): + clock = FakeClock(0.0) + ctl, _ = make_controller(reservation_seconds=30.0, heartbeat_seconds=15.0, clock=clock) + ctl.join("holder") + ctl.join("waiter") + + clock.advance(31.0) + + # The dropped session's query attempt gets refused. + assert ctl.claim("holder") is False + with pytest.raises(TicketNotFoundError): + ctl.status("holder") + + status_w = ctl.status("waiter") + assert status_w["status"] == "ready" + + +def test_heartbeat_expiry_drops_only_queued_tickets(): + clock = FakeClock(0.0) + ctl, _ = make_controller(reservation_seconds=100.0, heartbeat_seconds=10.0, clock=clock) + ctl.join("ready_holder") + ctl.join("queued_one") + + clock.advance(11.0) + + # Triggers lazy expiry; the ready ticket's reservation window (100s) has not elapsed. + view = ctl.status("ready_holder") + assert view["status"] == "ready" + + with pytest.raises(TicketNotFoundError): + ctl.status("queued_one") + + +def test_heartbeat_poll_keeps_queued_ticket_alive(): + clock = FakeClock(0.0) + ctl, _ = make_controller(reservation_seconds=100.0, heartbeat_seconds=10.0, clock=clock) + ctl.join("ready_holder") + ctl.join("queued_one") + + clock.advance(6.0) + ctl.status("queued_one") # refresh heartbeat + clock.advance(6.0) # 6s since refresh, under the 10s window + + view = ctl.status("queued_one") + assert view["status"] == "queued" + + +# -- Capacity ----------------------------------------------------------------- + + +def test_queue_full_at_max_length_excludes_reserved(): + ctl, clock = make_controller(max_queue_length=2) + ctl.join("reserved_session") # does not count against max_length + ctl.join("q1") + ctl.join("q2") + with pytest.raises(QueueFullError): + ctl.join("q3") + + +# -- Estimate ------------------------------------------------------------------- + + +def test_eta_null_before_history(): + ctl, clock = make_controller() + view = ctl.join("s1") + assert view["eta_seconds"] is None + + +def test_eta_ema_update_over_multiple_runs(): + ctl, clock = make_controller() + + run_cycle(ctl, "w1", 10.0) + probe = ctl.join("probe1") + assert probe["eta_seconds"] == 10 + ctl.leave("probe1") + + run_cycle(ctl, "w2", 20.0) + probe = ctl.join("probe2") + # EMA: 0.3*20 + 0.7*10 = 13 + assert probe["eta_seconds"] == 13 + ctl.leave("probe2") + + run_cycle(ctl, "w3", 30.0) + probe = ctl.join("probe3") + # EMA: 0.3*30 + 0.7*13 = 18.1 -> rounds to 18 + assert probe["eta_seconds"] == 18 + + +def test_ahead_counting_excludes_running_session(): + ctl, clock = make_controller() + ctl.join("running_session") + ctl.claim("running_session") + + view = ctl.join("q1") + assert view["status"] == "queued" + assert view["position"] == 1 + assert view["ahead"] == 0 + assert view["runner"] == "running" + + +def test_eta_clamps_at_zero_when_elapsed_exceeds_estimate(): + clock = FakeClock(0.0) + ctl, _ = make_controller(clock=clock) + + run_cycle(ctl, "w1", 10.0) # seeds mean = 10 + + ctl.join("w2") + ctl.claim("w2") + clock.advance(50.0) # far beyond the 10s estimate + + view = ctl.status("w2") + assert view["eta_seconds"] == 0 + + +def test_cancelled_runs_under_five_seconds_do_not_update_estimate(): + ctl, clock = make_controller() + run_cycle(ctl, "w1", 3.0, cancelled=True) + view = ctl.join("probe") + assert view["eta_seconds"] is None + + +def test_cancelled_runs_at_or_over_five_seconds_update_estimate(): + ctl, clock = make_controller() + run_cycle(ctl, "w1", 7.0, cancelled=True) + view = ctl.join("probe") + assert view["eta_seconds"] == 7 + + +# -- Claim admission table ---------------------------------------------------- + + +def test_claim_refused_for_other_session_while_reserved(): + ctl, clock = make_controller() + ctl.join("holder") + assert ctl.claim("intruder") is False + # Queue/ticket state untouched by the refusal. + assert ctl.status("holder")["status"] == "ready" + + +def test_claim_refused_while_running(): + ctl, clock = make_controller() + ctl.join("holder") + ctl.claim("holder") + assert ctl.claim("someone_else") is False + + +def test_implicit_join_admits_while_idle(): + ctl, clock = make_controller() + assert ctl.claim("walk_up") is True diff --git a/tests/backend/test_queue_api.py b/tests/backend/test_queue_api.py new file mode 100644 index 0000000..cda71e8 --- /dev/null +++ b/tests/backend/test_queue_api.py @@ -0,0 +1,329 @@ +"""API tests for the admission queue endpoints and changed /api/v1/query admission. + +Covers docs/QUEUE_SPECIFICATION.md section 11 "API tests": response shapes, +status codes, no-store header, cookie issuance on join, credential fields +rejected on queue endpoints, implicit join on query while idle, 409 for a +non-holder. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import httpx +import pytest +from httpx import ASGITransport + +from backend.admission import AdmissionController +from backend.app import create_app, SESSION_COOKIE_NAME +from backend.artifacts import ArtifactStore +from backend.containers import FakeContainerManager +from backend.dev.fake_peer import ScriptedContainerPeer +from backend.model import FakeModelAdapter +from backend.settings import Settings + +from tests.backend.conftest import make_test_confluence_client_factory + + +class FakeClock: + def __init__(self, start: float = 0.0): + self._now = start + + def __call__(self) -> float: + return self._now + + def advance(self, seconds: float) -> None: + self._now += seconds + + +def build_app( + tmp_path: Path, + scenario: str = "standard", + admission_controller=None, + query_timeout_seconds: float = 30.0, + cleanup_timeout_seconds: float = 5.0, + queue_max_length: int = 20, +): + settings = Settings( + approved_confluence_origins=["https://approved.example.com"], + query_timeout_seconds=query_timeout_seconds, + cleanup_timeout_seconds=cleanup_timeout_seconds, + queue_max_length=queue_max_length, + ) + store = ArtifactStore(tmp_path / "artifacts") + container_mgr = FakeContainerManager(lambda: ScriptedContainerPeer(scenario=scenario)) + model_adapter = FakeModelAdapter() + conf_factory = make_test_confluence_client_factory() + + app = create_app( + settings=settings, + container_manager=container_mgr, + artifact_store=store, + model_adapter=model_adapter, + confluence_client_factory=conf_factory, + admission_controller=admission_controller, + ) + return app, store, container_mgr + + +QUERY_PAYLOAD = { + "prompt": "How do I deploy service X?", + "credentials": {"url": "https://approved.example.com", "pat": "valid-pat"}, +} +ORIGIN = {"Origin": "http://testserver"} + + +@pytest.mark.asyncio +async def test_queue_join_response_shape_and_cookie(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert resp.status_code == 200 + assert resp.headers.get("Cache-Control") == "no-store" + assert SESSION_COOKIE_NAME in resp.cookies + + data = resp.json() + assert set(data.keys()) == { + "ticket_id", + "status", + "position", + "ahead", + "eta_seconds", + "reservation_expires_in_seconds", + "runner", + } + assert data["status"] == "ready" + assert data["position"] == 1 + assert data["ahead"] == 0 + assert data["eta_seconds"] is None + assert data["reservation_expires_in_seconds"] == 45 + assert data["runner"] == "reserved" + + +@pytest.mark.asyncio +async def test_queue_join_idempotent(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + first = (await client.post("/api/v1/queue/join", json={}, headers=ORIGIN)).json() + second = (await client.post("/api/v1/queue/join", json={}, headers=ORIGIN)).json() + assert first["ticket_id"] == second["ticket_id"] + + +@pytest.mark.asyncio +async def test_queue_join_rejects_credential_fields(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.post( + "/api/v1/queue/join", + json={"credentials": {"url": "https://approved.example.com", "pat": "x"}}, + headers=ORIGIN, + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "invalid_input" + + resp = await client.post("/api/v1/queue/join", json={"pat": "secret-token"}, headers=ORIGIN) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "invalid_input" + assert "secret-token" not in resp.text + + +@pytest.mark.asyncio +async def test_queue_join_origin_enforcement(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.post("/api/v1/queue/join", json={}) + assert resp.status_code == 403 + assert resp.json()["error"]["code"] == "origin_denied" + + +@pytest.mark.asyncio +async def test_queue_status_404_without_ticket_then_ready(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + # Bootstrap a session cookie without ever joining. + await client.get("/") + resp = await client.get("/api/v1/queue/status") + assert resp.status_code == 404 + assert resp.json()["error"]["code"] == "ticket_not_found" + assert resp.headers.get("Cache-Control") == "no-store" + + join_data = (await client.post("/api/v1/queue/join", json={}, headers=ORIGIN)).json() + + status_resp = await client.get("/api/v1/queue/status") + assert status_resp.status_code == 200 + status_data = status_resp.json() + assert status_data["ticket_id"] == join_data["ticket_id"] + assert status_data["status"] == "ready" + + +@pytest.mark.asyncio +async def test_queue_leave_204_always_and_drops_ticket(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + # No ticket yet: still 204. + resp = await client.delete("/api/v1/queue/ticket", headers=ORIGIN) + assert resp.status_code == 204 + + await client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + resp = await client.delete("/api/v1/queue/ticket", headers=ORIGIN) + assert resp.status_code == 204 + + status_resp = await client.get("/api/v1/queue/status") + assert status_resp.status_code == 404 + + # Idempotent second leave. + resp = await client.delete("/api/v1/queue/ticket", headers=ORIGIN) + assert resp.status_code == 204 + + +@pytest.mark.asyncio +async def test_queue_leave_origin_enforcement(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.delete("/api/v1/queue/ticket") + assert resp.status_code == 403 + assert resp.json()["error"]["code"] == "origin_denied" + + +@pytest.mark.asyncio +async def test_queue_full_503(tmp_path: Path): + app, _, _ = build_app(tmp_path, queue_max_length=1) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as reserved_client, \ + httpx.AsyncClient(transport=transport, base_url="http://testserver") as queued_client, \ + httpx.AsyncClient(transport=transport, base_url="http://testserver") as overflow_client: + r1 = await reserved_client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert r1.json()["status"] == "ready" + + r2 = await queued_client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert r2.json()["status"] == "queued" + + r3 = await overflow_client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert r3.status_code == 503 + assert r3.json()["error"]["code"] == "queue_full" + + +@pytest.mark.asyncio +async def test_query_implicit_join_while_idle_then_idle_again_after(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert resp.status_code == 200 + + # Promotion/idle recovery happens after cleanup; a fresh session should + # immediately be admitted as "ready" (idle) and see a recorded estimate. + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as probe: + probe_resp = await probe.post("/api/v1/queue/join", json={}, headers=ORIGIN) + data = probe_resp.json() + assert data["status"] == "ready" + assert data["eta_seconds"] is not None + assert data["eta_seconds"] >= 0 + + +@pytest.mark.asyncio +async def test_query_409_busy_for_non_holder_session(tmp_path: Path): + app, _, _ = build_app(tmp_path) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as holder, \ + httpx.AsyncClient(transport=transport, base_url="http://testserver") as intruder: + join_resp = await holder.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert join_resp.json()["status"] == "ready" + + # A different session, holding no reservation, is refused without + # disturbing the queue. + busy_resp = await intruder.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert busy_resp.status_code == 409 + assert busy_resp.json()["error"]["code"] == "busy" + + status_after = await holder.get("/api/v1/queue/status") + assert status_after.json()["status"] == "ready" + + # The holder itself can now claim and run. + ok_resp = await holder.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert ok_resp.status_code == 200 + + +@pytest.mark.asyncio +async def test_query_busy_while_running_and_promotes_after_cleanup(tmp_path: Path): + app, _, _ = build_app( + tmp_path, scenario="timeout", query_timeout_seconds=0.2, cleanup_timeout_seconds=2.0 + ) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as running_client, \ + httpx.AsyncClient(transport=transport, base_url="http://testserver") as other_client, \ + httpx.AsyncClient(transport=transport, base_url="http://testserver") as probe_client: + task = asyncio.create_task( + running_client.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + ) + await asyncio.sleep(0.05) + + busy_resp = await other_client.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert busy_resp.status_code == 409 + assert busy_resp.json()["error"]["code"] == "busy" + + # Runner times out, cleanup runs, and the runner returns to idle; + # the timed-out run's duration is still recorded (failure counts). + timeout_resp = await asyncio.wait_for(task, timeout=10.0) + assert timeout_resp.status_code == 504 + + probe_resp = await probe_client.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert probe_resp.json()["status"] == "ready" + + +@pytest.mark.asyncio +async def test_reservation_expiry_query_after_window_is_busy(tmp_path: Path): + # The query request itself is what discovers the expiry (spec section 3, + # "Claim"): the reserved session sending the query after the window gets + # 409, and the ticket is dropped as a side effect of that request. + clock = FakeClock(0.0) + admission = AdmissionController( + reservation_seconds=1.0, heartbeat_seconds=30.0, max_queue_length=5, clock=clock + ) + app, _, _ = build_app(tmp_path, admission_controller=admission) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as holder: + join_resp = await holder.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert join_resp.json()["status"] == "ready" + + clock.advance(2.0) + + query_resp = await holder.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert query_resp.status_code == 409 + assert query_resp.json()["error"]["code"] == "busy" + + +@pytest.mark.asyncio +async def test_reservation_expiry_status_poll_is_ticket_not_found(tmp_path: Path): + # Spec section 8: "Reservation expires -> ticket dropped, next promoted; + # next status is 404; rejoin once." Polling after the window discovers + # the expiry itself and reports ticket_not_found rather than "ready". + clock = FakeClock(0.0) + admission = AdmissionController( + reservation_seconds=1.0, heartbeat_seconds=30.0, max_queue_length=5, clock=clock + ) + app, _, _ = build_app(tmp_path, admission_controller=admission) + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as holder: + join_resp = await holder.post("/api/v1/queue/join", json={}, headers=ORIGIN) + assert join_resp.json()["status"] == "ready" + + clock.advance(2.0) + + status_resp = await holder.get("/api/v1/queue/status") + assert status_resp.status_code == 404 + assert status_resp.json()["error"]["code"] == "ticket_not_found" + + # Once discovered, the runner is idle again (queue was empty), so a + # fresh query from any session -- including this one, now indistinct + # from a walk-up -- is admitted by implicit join. + query_resp = await holder.post("/api/v1/query", json=QUERY_PAYLOAD, headers=ORIGIN) + assert query_resp.status_code == 200 diff --git a/tests/backend/test_settings.py b/tests/backend/test_settings.py index 93c9ae9..3f3a918 100644 --- a/tests/backend/test_settings.py +++ b/tests/backend/test_settings.py @@ -1,4 +1,6 @@ import pytest +from pydantic import ValidationError + from backend.errors import DestinationDeniedError, InvalidInputError from backend.settings import Settings, canonicalize_url, validate_confluence_url @@ -69,6 +71,49 @@ def test_validate_confluence_url(): validate_confluence_url("http://approved.example.com", approved) +# -- Admission queue settings validators (docs/QUEUE_SPECIFICATION.md section 4) -- + + +def test_queue_reservation_seconds_defaults_and_bounds(): + assert Settings().queue_reservation_seconds == 45.0 + Settings(queue_reservation_seconds=30.0) + Settings(queue_reservation_seconds=60.0) + with pytest.raises(ValidationError): + Settings(queue_reservation_seconds=29.9) + with pytest.raises(ValidationError): + Settings(queue_reservation_seconds=60.1) + + +def test_queue_heartbeat_seconds_defaults_and_bounds(): + assert Settings().queue_heartbeat_seconds == 15.0 + Settings(queue_heartbeat_seconds=5.0) + Settings(queue_heartbeat_seconds=60.0) + with pytest.raises(ValidationError): + Settings(queue_heartbeat_seconds=4.9) + with pytest.raises(ValidationError): + Settings(queue_heartbeat_seconds=60.1) + + +def test_queue_max_length_defaults_and_bounds(): + assert Settings().queue_max_length == 20 + Settings(queue_max_length=1) + Settings(queue_max_length=100) + with pytest.raises(ValidationError): + Settings(queue_max_length=0) + with pytest.raises(ValidationError): + Settings(queue_max_length=101) + + +def test_queue_settings_from_env(monkeypatch): + monkeypatch.setenv("CONFLUENCE_WEB_QUEUE_RESERVATION_SECONDS", "50") + monkeypatch.setenv("CONFLUENCE_WEB_QUEUE_HEARTBEAT_SECONDS", "20") + monkeypatch.setenv("CONFLUENCE_WEB_QUEUE_MAX_LENGTH", "5") + settings = Settings.from_env() + assert settings.queue_reservation_seconds == 50.0 + assert settings.queue_heartbeat_seconds == 20.0 + assert settings.queue_max_length == 5 + + def test_max_deadline_seconds_defaults_and_bounds(): assert Settings().max_deadline_seconds == 900.0 Settings(max_deadline_seconds=60.0, query_timeout_seconds=60.0)