backend: admission queue with reservations

In-memory AdmissionController (idle/reserved/running, FIFO tickets keyed by
the session cookie, reservation and heartbeat expiry, promotion after
cleanup, EMA wait estimate) behind POST /api/v1/queue/join,
GET /api/v1/queue/status and DELETE /api/v1/queue/ticket. POST /api/v1/query
claims the session's reservation first and joins implicitly when idle.
Settings: CONFLUENCE_WEB_QUEUE_{RESERVATION_SECONDS,HEARTBEAT_SECONDS,MAX_LENGTH}.
This commit is contained in:
Artur Mukhamadiev 2026-09-15 13:24:15 +03:00
parent 3751ab26b5
commit 5aad7a0b45
9 changed files with 1103 additions and 10 deletions

View File

@ -9,6 +9,7 @@ backend/
app.py # FastAPI application factory, session cookies, security headers, endpoints app.py # FastAPI application factory, session cookies, security headers, endpoints
settings.py # Validated deployment settings and URL validation/canonicalization settings.py # Validated deployment settings and URL validation/canonicalization
runner.py # QueryRunner orchestrating deadlines, gate, disconnect cancellation, cleanup 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) transport.py # Bounded NDJSON bridge reader/writer and protocol state machine (Rev 1)
containers.py # Rootless Docker container lifecycle, tmpfs mounts, kill/remove, reconciliation containers.py # Rootless Docker container lifecycle, tmpfs mounts, kill/remove, reconciliation
confluence.py # Request-scoped Confluence client, streaming bounds, and tool dispatcher 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_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 (603600). Passed into the agent container so the supervisor and bridge enforce the same bound. | | `CONFLUENCE_WEB_MAX_DEADLINE_SECONDS` | `900.0` | Protocol maximum for one query (603600). 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_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 ## Running the Service

295
backend/admission.py Normal file
View File

@ -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,
}

View File

@ -17,6 +17,7 @@ from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.exceptions import HTTPException as StarletteHTTPException
from backend.admission import AdmissionController
from backend.artifacts import ArtifactStore from backend.artifacts import ArtifactStore
from backend.confluence import ConfluenceClient from backend.confluence import ConfluenceClient
from backend.containers import ( from backend.containers import (
@ -27,11 +28,13 @@ from backend.containers import (
from backend.errors import ( from backend.errors import (
AppError, AppError,
ArtifactNotFoundError, ArtifactNotFoundError,
BusyError,
InvalidInputError, InvalidInputError,
OriginDeniedError, OriginDeniedError,
ConnectivityFailedError, ConnectivityFailedError,
sanitize_message, sanitize_message,
RequestTooLargeError, RequestTooLargeError,
TicketNotFoundError,
) )
from backend.model import FakeModelAdapter, ModelAdapter, OpenAIModelAdapter from backend.model import FakeModelAdapter, ModelAdapter, OpenAIModelAdapter
from backend.runner import QueryRunner from backend.runner import QueryRunner
@ -71,6 +74,12 @@ class QueryRequest(BaseModel):
credentials: QueryCredentials credentials: QueryCredentials
class QueueJoinRequest(BaseModel):
"""Body is `{}`. extra="forbid" rejects any field, including credentials/pat."""
model_config = ConfigDict(extra="forbid")
class ServerSessionStore: class ServerSessionStore:
"""In-memory store of server-issued session IDs to prevent client-forged sessions.""" """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") 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: while True:
await asyncio.sleep(60) await asyncio.sleep(60)
if not runner._gate.locked(): if not runner._gate.locked():
@ -174,6 +183,7 @@ async def _periodic_maintenance(runner, store, session_store) -> None:
try: try:
await asyncio.to_thread(store.expire_artifacts) await asyncio.to_thread(store.expire_artifacts)
session_store.prune_stale() session_store.prune_stale()
admission.expire()
except Exception: except Exception:
logger.warning("Artifact or session maintenance failed") logger.warning("Artifact or session maintenance failed")
@ -266,6 +276,7 @@ def create_app(
artifact_store: Optional[ArtifactStore] = None, artifact_store: Optional[ArtifactStore] = None,
model_adapter: Optional[ModelAdapter] = None, model_adapter: Optional[ModelAdapter] = None,
confluence_client_factory: Optional[Callable[..., ConfluenceClient]] = None, confluence_client_factory: Optional[Callable[..., ConfluenceClient]] = None,
admission_controller: Optional[AdmissionController] = None,
) -> FastAPI: ) -> FastAPI:
"""Create and configure FastAPI application with dependency injection.""" """Create and configure FastAPI application with dependency injection."""
app_settings = settings or Settings.from_env() app_settings = settings or Settings.from_env()
@ -311,6 +322,11 @@ def create_app(
) )
session_store = ServerSessionStore() 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 @contextlib.asynccontextmanager
async def lifespan(fastapi_app: FastAPI): async def lifespan(fastapi_app: FastAPI):
@ -320,7 +336,7 @@ def create_app(
if isinstance(mgr, DockerContainerManager): if isinstance(mgr, DockerContainerManager):
await mgr.verify_rootless() await mgr.verify_rootless()
await runner.reconcile() 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 yield
finally: finally:
if maintenance_task is not None: if maintenance_task is not None:
@ -393,6 +409,7 @@ def create_app(
app.add_middleware(SecurityMiddleware) app.add_middleware(SecurityMiddleware)
app.state.runner = runner app.state.runner = runner
app.state.artifact_store = store app.state.artifact_store = store
app.state.admission = admission
@app.get("/") @app.get("/")
async def root_index(request: Request): async def root_index(request: Request):
@ -437,6 +454,44 @@ def create_app(
get_or_set_session_cookie(request, resp, session_store) get_or_set_session_cookie(request, resp, session_store)
return resp 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") @app.post("/api/v1/query")
async def execute_query(req: QueryRequest, request: Request): async def execute_query(req: QueryRequest, request: Request):
check_same_origin(request) check_same_origin(request)
@ -446,7 +501,15 @@ def create_app(
else: else:
session_store.touch(session_id) # type: ignore[arg-type] session_store.touch(session_id) # type: ignore[arg-type]
# Execute query via runner # 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( result = await runner.run(
prompt=req.prompt, prompt=req.prompt,
confluence_url=req.credentials.url, confluence_url=req.credentials.url,
@ -454,6 +517,16 @@ def create_app(
session_id=session_id, session_id=session_id,
is_disconnected=request.is_disconnected, 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) resp = JSONResponse(content=result)
# Ensure session cookie is set # Ensure session cookie is set
secure = request.url.scheme == "https" secure = request.url.scheme == "https"

View File

@ -125,3 +125,13 @@ class ExecutionFailedError(AppError):
class CleanupFailedError(AppError): class CleanupFailedError(AppError):
def __init__(self, message: str = "Container cleanup could not be confirmed"): def __init__(self, message: str = "Container cleanup could not be confirmed"):
super().__init__("cleanup_failed", message, status_code=500) 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)

View File

@ -69,6 +69,11 @@ class Settings(BaseModel):
# Development mode (explicit opt-in) # Development mode (explicit opt-in)
dev_mode: bool = False 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") @field_validator("approved_confluence_origins")
@classmethod @classmethod
def validate_origins(cls, v: List[str]) -> List[str]: 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") raise ValueError("query_timeout_seconds must not exceed max_deadline_seconds")
return self 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 @classmethod
def from_env(cls) -> Settings: def from_env(cls) -> Settings:
"""Load settings from environment variables.""" """Load settings from environment variables."""
@ -210,6 +236,13 @@ class Settings(BaseModel):
cleanup_timeout_seconds=_parse_float( cleanup_timeout_seconds=_parse_float(
"CONFLUENCE_WEB_CLEANUP_TIMEOUT_SECONDS", 10.0 "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),
) )

View File

@ -51,5 +51,13 @@ CONFLUENCE_WEB_QUERY_TIMEOUT_SECONDS=600
CONFLUENCE_WEB_MAX_DEADLINE_SECONDS=900 CONFLUENCE_WEB_MAX_DEADLINE_SECONDS=900
CONFLUENCE_WEB_CLEANUP_TIMEOUT_SECONDS=10 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. # Never set in production. Substitutes the container, model and Confluence with fakes.
CONFLUENCE_WEB_DEV_MODE=false CONFLUENCE_WEB_DEV_MODE=false

View File

@ -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

View File

@ -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

View File

@ -1,4 +1,6 @@
import pytest import pytest
from pydantic import ValidationError
from backend.errors import DestinationDeniedError, InvalidInputError from backend.errors import DestinationDeniedError, InvalidInputError
from backend.settings import Settings, canonicalize_url, validate_confluence_url 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) 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(): def test_max_deadline_seconds_defaults_and_bounds():
assert Settings().max_deadline_seconds == 900.0 assert Settings().max_deadline_seconds == 900.0
Settings(max_deadline_seconds=60.0, query_timeout_seconds=60.0) Settings(max_deadline_seconds=60.0, query_timeout_seconds=60.0)