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:
parent
3751ab26b5
commit
5aad7a0b45
@ -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 (60–3600). 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 (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_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
295
backend/admission.py
Normal 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,
|
||||||
|
}
|
||||||
@ -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"
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
@ -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),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
296
tests/backend/test_admission.py
Normal file
296
tests/backend/test_admission.py
Normal 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
|
||||||
329
tests/backend/test_queue_api.py
Normal file
329
tests/backend/test_queue_api.py
Normal 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
|
||||||
@ -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)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user