DGXAI/driftcall
0
1"""DriftCall env Space — FastAPI + OpenEnv-compliant REST surface.2 3Implements ``docs/modules/deploy_env_space.md`` and DESIGN.md §3.3 / §11.1.4 5Endpoints:6 GET /healthz → 200 text/plain "ok" (unauthenticated)7 POST /reset → 200 application/json (create / recycle session)8 POST /step → 200 application/json (advance one turn)9 GET /state → 200 application/json (read DriftCallState)10 POST /close → 200 application/json (evict session)11 12Headers (mutating endpoints): ``Authorization: Bearer <DRIFTCALL_ENV_TOKEN>``13and ``X-Session-Id: <[A-Za-z0-9_-]{1,64}>``.14 15Error modes (deploy_env_space.md §5):16 M1 401 unauthorized M7 400 bad_json17 M2 400 missing_session_id M8 400 invalid_action18 M3 404 session_not_found M9 500 internal_error19 M4 404 session_expired M10 500 io_error20 M5 429 max_sessions M11 413 payload_too_large21 M6 503 model_not_ready M12 409 reset_in_progress22 23All error bodies: ``{"error": {"code": <slug>, "message": <str>,24"request_id": <asgi-id>}}``; ``Cache-Control: no-store``; only M5 carries25``Retry-After: 30``. No stack traces ever leak across the wire.26"""27 28from __future__ import annotations29 30import asyncio31import contextlib32import dataclasses33import json34import logging35import os36import re37import time38from contextlib import asynccontextmanager39from dataclasses import dataclass, replace40from typing import TYPE_CHECKING, Any41 42from fastapi import FastAPI, Request, Response43from fastapi.responses import JSONResponse, PlainTextResponse44from starlette.middleware.base import BaseHTTPMiddleware45 46from cells.step_04_models import ActionType, DriftCallAction47from cells.step_10_env import (48 DriftCallEnv,49 EnvClosedError,50 EnvNotReadyError,51 EpisodeAlreadyTerminalError,52 InvalidActionError,53 InvalidConfigError,54 UnknownDomainError,55 UnknownToolError,56)57 58if TYPE_CHECKING:59 from collections.abc import AsyncIterator, Awaitable, Callable60 61 from starlette.types import ASGIApp62 63logger = logging.getLogger(__name__)64 65# ---------------------------------------------------------------------------66# Constants67# ---------------------------------------------------------------------------68 69_MAX_SESSIONS: int = 1070_TTL_S: float = 3600.071_SWEEP_INTERVAL_S: float = 60.072_MAX_SESSION_ID_LEN: int = 6473_SESSION_ID_RE: re.Pattern[str] = re.compile(r"^[A-Za-z0-9_-]{1,64}$")74_MAX_BODY_BYTES: int = 1 * 1024 * 1024 # 1 MiB75_RETRY_AFTER_S: str = "30"76_TOKEN_ENV_VAR: str = "DRIFTCALL_ENV_TOKEN"77 78 79# ---------------------------------------------------------------------------80# Time source (test-overridable)81# ---------------------------------------------------------------------------82 83 84def _monotonic() -> float:85 """Indirection for tests to monkeypatch."""86 87 return time.monotonic()88 89 90# ---------------------------------------------------------------------------91# Errors / envelope92# ---------------------------------------------------------------------------93 94 95@dataclass(frozen=True)96class _ApiError(Exception):97 """Internal exception → uniform error envelope (deploy_env_space.md §5)."""98 99 code: str100 message: str101 http_status: int102 retry_after: bool = False103 104 105_NO_STORE: dict[str, str] = {"Cache-Control": "no-store"}106 107 108def _error_response(err: _ApiError, request_id: str) -> JSONResponse:109 body = {110 "error": {111 "code": err.code,112 "message": err.message,113 "request_id": request_id,114 }115 }116 headers = dict(_NO_STORE)117 if err.retry_after:118 headers["Retry-After"] = _RETRY_AFTER_S119 return JSONResponse(status_code=err.http_status, content=body, headers=headers)120 121 122# ---------------------------------------------------------------------------123# Session cache124# ---------------------------------------------------------------------------125 126 127@dataclass(frozen=True)128class SessionEntry:129 """Frozen per project rule — every touch produces a new entry."""130 131 env: DriftCallEnv132 created_at: float133 last_touched: float134 reset_count: int135 lock: asyncio.Lock136 137 138class SessionCache:139 """In-memory session registry with LRU + TTL eviction."""140 141 def __init__(self, *, max_sessions: int = _MAX_SESSIONS, ttl_s: float = _TTL_S) -> None:142 self._max = max_sessions143 self._ttl = ttl_s144 self._store: dict[str, SessionEntry] = {}145 self._guard = asyncio.Lock()146 147 @property148 def size(self) -> int:149 return len(self._store)150 151 def get(self, sid: str) -> SessionEntry | None:152 return self._store.get(sid)153 154 async def acquire_lock(self, sid: str) -> asyncio.Lock:155 """Return (or lazily create) the per-session lock."""156 async with self._guard:157 entry = self._store.get(sid)158 if entry is not None:159 return entry.lock160 return asyncio.Lock()161 162 async def insert_or_replace(self, sid: str, env_factory: Callable[[], DriftCallEnv]) -> SessionEntry:163 """Insert a new env or replace an existing one (in-place reset)."""164 async with self._guard:165 now = _monotonic()166 existing = self._store.get(sid)167 if existing is not None:168 # In-place reset (§7.1 case after winner completed).169 try:170 existing.env.close()171 except Exception:172 logger.exception("env.close() raised on in-place reset for sid=%s", sid)173 env = env_factory()174 entry = SessionEntry(175 env=env,176 created_at=now,177 last_touched=now,178 reset_count=existing.reset_count + 1,179 lock=existing.lock,180 )181 self._store[sid] = entry182 return entry183 # New session — enforce cap.184 if len(self._store) >= self._max:185 # Try LRU evict only if any entry is older than the others by TTL/2.186 victim_sid = min(self._store, key=lambda k: self._store[k].last_touched)187 victim = self._store[victim_sid]188 age = now - victim.last_touched189 if age <= 0.0:190 raise _ApiError(191 code="max_sessions",192 message=f"max concurrent sessions reached ({self._max})",193 http_status=429,194 retry_after=True,195 )196 try:197 victim.env.close()198 except Exception:199 logger.exception("env.close() raised on LRU eviction for sid=%s", victim_sid)200 self._store.pop(victim_sid, None)201 env = env_factory()202 entry = SessionEntry(203 env=env,204 created_at=now,205 last_touched=now,206 reset_count=0,207 lock=asyncio.Lock(),208 )209 self._store[sid] = entry210 return entry211 212 def touch(self, sid: str) -> tuple[SessionEntry | None, bool]:213 """Update last_touched. Returns ``(entry, was_expired)``.214 215 - ``(entry, False)`` on hit216 - ``(None, True)`` if the entry was present but evicted by this call217 due to TTL expiry218 - ``(None, False)`` if there was never an entry under this sid219 """220 entry = self._store.get(sid)221 if entry is None:222 return None, False223 now = _monotonic()224 if now - entry.last_touched > self._ttl:225 try:226 entry.env.close()227 except Exception:228 logger.exception("env.close() raised on expired touch for sid=%s", sid)229 self._store.pop(sid, None)230 return None, True231 new = replace(entry, last_touched=now)232 self._store[sid] = new233 return new, False234 235 def evict(self, sid: str) -> SessionEntry | None:236 """Pop a session out of the cache. Returns the removed entry or None."""237 return self._store.pop(sid, None)238 239 def sweep(self) -> int:240 """Synchronous TTL sweep — evict every entry past TTL."""241 now = _monotonic()242 expired = [sid for sid, e in self._store.items() if now - e.last_touched > self._ttl]243 for sid in expired:244 entry = self._store.pop(sid)245 try:246 entry.env.close()247 except Exception:248 logger.exception("env.close() raised on sweep for sid=%s", sid)249 if expired:250 logger.info(251 json.dumps(252 {253 "event": "session_sweep",254 "expired_count": len(expired),255 "cache_size": len(self._store),256 }257 )258 )259 return len(expired)260 261 262# ---------------------------------------------------------------------------263# App state container264# ---------------------------------------------------------------------------265 266 267@dataclass268class _AppState:269 """Mutable (intentional) — owned by lifespan; readers go through getters."""270 271 cache: SessionCache272 models_ready: bool = False273 sweep_task: asyncio.Task[None] | None = None274 bearer_token: str = ""275 276 277def _get_state(app: FastAPI) -> _AppState:278 state: _AppState = app.state.driftcall279 return state280 281 282# ---------------------------------------------------------------------------283# Lifespan — eager-load Kokoro + Whisper before serving (M6 guard)284# ---------------------------------------------------------------------------285 286 287def _eager_load_models() -> None:288 """Force-load TTS + ASR singletons. Test patches this to avoid network."""289 from cells.step_09_audio import get_asr_engine, get_tts_engine290 291 get_tts_engine()292 get_asr_engine()293 294 295@asynccontextmanager296async def lifespan(app: FastAPI) -> AsyncIterator[None]:297 cache = SessionCache()298 token = os.environ.get(_TOKEN_ENV_VAR, "")299 if not token:300 # Fail-fast per deploy_env_space.md §3.5.301 raise RuntimeError(302 f"{_TOKEN_ENV_VAR} environment variable not set; refusing to start"303 )304 state = _AppState(cache=cache, bearer_token=token)305 app.state.driftcall = state306 307 # Eager model load (M6 guard — must complete before serving).308 try:309 await asyncio.to_thread(_eager_load_models)310 except Exception:311 logger.exception("eager model load failed")312 raise313 state.models_ready = True314 315 # Background TTL sweep.316 async def _sweep_loop() -> None:317 try:318 while True:319 await asyncio.sleep(_SWEEP_INTERVAL_S)320 cache.sweep()321 except asyncio.CancelledError:322 raise323 324 state.sweep_task = asyncio.create_task(_sweep_loop())325 try:326 yield327 finally:328 if state.sweep_task is not None:329 state.sweep_task.cancel()330 with contextlib.suppress(asyncio.CancelledError, Exception):331 await state.sweep_task332 333 334# ---------------------------------------------------------------------------335# Body-size middleware (M11)336# ---------------------------------------------------------------------------337 338 339class _BodySizeMiddleware(BaseHTTPMiddleware):340 def __init__(self, app: ASGIApp, *, max_bytes: int = _MAX_BODY_BYTES) -> None:341 super().__init__(app)342 self._max_bytes = max_bytes343 344 async def dispatch(345 self, request: Request, call_next: Callable[[Request], Awaitable[Response]]346 ) -> Response:347 cl = request.headers.get("content-length")348 if cl is not None:349 try:350 cl_int = int(cl)351 except ValueError:352 cl_int = -1353 if cl_int > self._max_bytes:354 err = _ApiError(355 code="payload_too_large",356 message="request body exceeds 1 MiB",357 http_status=413,358 )359 return _error_response(err, _request_id(request))360 return await call_next(request)361 362 363# ---------------------------------------------------------------------------364# Helpers — auth, headers, body parsing365# ---------------------------------------------------------------------------366 367 368def _request_id(request: Request) -> str:369 return str(id(request))370 371 372def _check_bearer(request: Request, state: _AppState) -> None:373 auth = request.headers.get("authorization", "")374 if not auth.startswith("Bearer "):375 raise _ApiError(376 code="unauthorized",377 message="missing or non-Bearer Authorization header",378 http_status=401,379 )380 token = auth[len("Bearer ") :].strip()381 if token != state.bearer_token or not token:382 raise _ApiError(383 code="unauthorized",384 message="invalid bearer token",385 http_status=401,386 )387 388 389def _check_session_header(request: Request) -> str:390 sid = request.headers.get("x-session-id", "")391 if not sid or not _SESSION_ID_RE.match(sid):392 raise _ApiError(393 code="missing_session_id",394 message="X-Session-Id header missing or malformed",395 http_status=400,396 )397 return sid398 399 400def _check_models_ready(state: _AppState) -> None:401 if not state.models_ready:402 raise _ApiError(403 code="model_not_ready",404 message="audio models still loading; retry shortly",405 http_status=503,406 )407 408 409async def _parse_json_body(request: Request) -> dict[str, Any]:410 raw = await request.body()411 if len(raw) > _MAX_BODY_BYTES:412 raise _ApiError(413 code="payload_too_large",414 message="request body exceeds 1 MiB",415 http_status=413,416 )417 if not raw:418 return {}419 try:420 parsed = json.loads(raw)421 except (json.JSONDecodeError, UnicodeDecodeError) as exc:422 raise _ApiError(423 code="bad_json",424 message=f"malformed JSON: {exc.__class__.__name__}",425 http_status=400,426 ) from exc427 if not isinstance(parsed, dict):428 raise _ApiError(429 code="bad_json",430 message="request body must be a JSON object",431 http_status=400,432 )433 return parsed434 435 436# ---------------------------------------------------------------------------437# Action / config validation (envelope-level — env owns deep validation)438# ---------------------------------------------------------------------------439 440 441def _build_action(raw: Any) -> DriftCallAction:442 if not isinstance(raw, dict):443 raise _ApiError(444 code="invalid_action",445 message="action must be a JSON object",446 http_status=400,447 )448 atype_raw = raw.get("action_type")449 if not isinstance(atype_raw, str):450 raise _ApiError(451 code="invalid_action",452 message="action.action_type must be a string",453 http_status=400,454 )455 try:456 atype = ActionType(atype_raw)457 except ValueError as exc:458 raise _ApiError(459 code="invalid_action",460 message=f"unknown action_type {atype_raw!r}",461 http_status=400,462 ) from exc463 464 tool_name = raw.get("tool_name")465 tool_args = raw.get("tool_args")466 message = raw.get("message")467 confidence = raw.get("confidence")468 rationale = raw.get("rationale")469 470 # Action-type contract checks (deep checks happen inside env._validate_action).471 if atype == ActionType.TOOL_CALL and (472 tool_name is None or not isinstance(tool_name, str) or tool_args is None473 ):474 raise _ApiError(475 code="invalid_action",476 message="TOOL_CALL requires tool_name (str) and tool_args (object)",477 http_status=400,478 )479 return DriftCallAction(480 action_type=atype,481 tool_name=tool_name if isinstance(tool_name, str) else None,482 tool_args=tool_args if isinstance(tool_args, dict) else None,483 message=message if isinstance(message, str) else None,484 confidence=float(confidence) if isinstance(confidence, (int, float)) and not isinstance(confidence, bool) else None,485 rationale=rationale if isinstance(rationale, str) else None,486 )487 488 489def _build_env_config(reset_body: dict[str, Any]) -> dict[str, Any]:490 raw_cfg = reset_body.get("config")491 if raw_cfg is None:492 raw_cfg = {}493 if not isinstance(raw_cfg, dict):494 raise _ApiError(495 code="invalid_action",496 message="config must be a JSON object",497 http_status=400,498 )499 return raw_cfg500 501 502# ---------------------------------------------------------------------------503# Serialization helpers504# ---------------------------------------------------------------------------505 506 507def _to_jsonable(obj: Any) -> Any:508 """Recursively convert frozen dataclasses / tuples / enums to JSON-safe form."""509 if dataclasses.is_dataclass(obj) and not isinstance(obj, type):510 return {k: _to_jsonable(v) for k, v in dataclasses.asdict(obj).items()}511 if isinstance(obj, ActionType):512 return obj.value513 if isinstance(obj, dict):514 return {k: _to_jsonable(v) for k, v in obj.items()}515 if isinstance(obj, (list, tuple)):516 return [_to_jsonable(v) for v in obj]517 return obj518 519 520# ---------------------------------------------------------------------------521# Endpoint handlers (one function per route)522# ---------------------------------------------------------------------------523 524 525async def _handle_reset(request: Request, state: _AppState) -> Response:526 _check_bearer(request, state)527 _check_models_ready(state)528 sid = _check_session_header(request)529 body = await _parse_json_body(request)530 cfg = _build_env_config(body)531 seed_raw = body.get("seed")532 if seed_raw is not None and (not isinstance(seed_raw, int) or isinstance(seed_raw, bool)):533 raise _ApiError(534 code="invalid_action",535 message="seed must be an int or null",536 http_status=400,537 )538 seed: int | None = seed_raw if isinstance(seed_raw, int) and not isinstance(seed_raw, bool) else None539 540 cache = state.cache541 # Per-session reset lock (§7.1).542 existing = cache.get(sid)543 if existing is not None and existing.lock.locked():544 raise _ApiError(545 code="reset_in_progress",546 message="concurrent /reset on same session id",547 http_status=409,548 )549 550 # Acquire lock (creates one if not present).551 lock = await cache.acquire_lock(sid)552 if lock.locked():553 raise _ApiError(554 code="reset_in_progress",555 message="concurrent /reset on same session id",556 http_status=409,557 )558 559 async with lock:560 def _factory() -> DriftCallEnv:561 try:562 return DriftCallEnv(cfg)563 except InvalidConfigError as exc:564 raise _ApiError(565 code="invalid_action",566 message=f"invalid config: {exc}",567 http_status=400,568 ) from exc569 570 try:571 entry = await cache.insert_or_replace(sid, _factory)572 except _ApiError:573 raise574 except Exception as exc:575 logger.exception("env construction failed for sid=%s", sid)576 raise _ApiError(577 code="internal_error",578 message="env construction failed",579 http_status=500,580 ) from exc581 582 try:583 obs = await asyncio.to_thread(entry.env.reset, seed)584 except InvalidConfigError as exc:585 cache.evict(sid)586 raise _ApiError(587 code="invalid_action",588 message=f"invalid config at reset: {exc}",589 http_status=400,590 ) from exc591 except OSError as exc:592 cache.evict(sid)593 raise _ApiError(594 code="io_error",595 message=f"I/O error during reset: {exc.__class__.__name__}",596 http_status=500,597 ) from exc598 except Exception as exc:599 cache.evict(sid)600 logger.exception("env.reset raised for sid=%s", sid)601 raise _ApiError(602 code="internal_error",603 message="env.reset raised",604 http_status=500,605 ) from exc606 607 body_out = {608 "observation": _to_jsonable(obs),609 "episode_id": entry.env.state().episode_id,610 "max_turns": entry.env.state().max_turns,611 }612 return JSONResponse(status_code=200, content=body_out)613 614 615async def _handle_step(request: Request, state: _AppState) -> Response:616 _check_bearer(request, state)617 _check_models_ready(state)618 sid = _check_session_header(request)619 body = await _parse_json_body(request)620 raw_action = body.get("action")621 action = _build_action(raw_action)622 623 entry, was_expired = state.cache.touch(sid)624 if entry is None:625 if was_expired:626 raise _ApiError(627 code="session_expired",628 message="session TTL expired; call /reset",629 http_status=404,630 )631 raise _ApiError(632 code="session_not_found",633 message="X-Session-Id has no live session; call /reset",634 http_status=404,635 )636 637 try:638 obs = await asyncio.to_thread(entry.env.step, action)639 except (InvalidActionError, UnknownToolError, UnknownDomainError) as exc:640 raise _ApiError(641 code="invalid_action",642 message=str(exc),643 http_status=400,644 ) from exc645 except (EnvNotReadyError, EnvClosedError, EpisodeAlreadyTerminalError) as exc:646 raise _ApiError(647 code="invalid_action",648 message=str(exc),649 http_status=400,650 ) from exc651 except OSError as exc:652 raise _ApiError(653 code="io_error",654 message=f"I/O error during step: {exc.__class__.__name__}",655 http_status=500,656 ) from exc657 except Exception as exc:658 logger.exception("env.step raised for sid=%s", sid)659 raise _ApiError(660 code="internal_error",661 message="env.step raised",662 http_status=500,663 ) from exc664 665 reward: float | None = None666 info: dict[str, Any] = {}667 if entry.env.done():668 try:669 rewards = entry.env.rewards()670 reward = float(getattr(rewards, "reward", 0.0))671 info["terminated_by"] = entry.env.episode().terminated_by672 except Exception:673 reward = None674 675 body_out = {676 "observation": _to_jsonable(obs),677 "reward": reward,678 "done": bool(entry.env.done()),679 "info": info,680 }681 return JSONResponse(status_code=200, content=body_out)682 683 684async def _handle_state(request: Request, state: _AppState) -> Response:685 _check_bearer(request, state)686 _check_models_ready(state)687 sid = _check_session_header(request)688 entry, was_expired = state.cache.touch(sid)689 if entry is None:690 if was_expired:691 raise _ApiError(692 code="session_expired",693 message="session TTL expired; call /reset",694 http_status=404,695 )696 raise _ApiError(697 code="session_not_found",698 message="X-Session-Id has no live session; call /reset",699 http_status=404,700 )701 try:702 st = entry.env.state()703 except EnvNotReadyError as exc:704 raise _ApiError(705 code="invalid_action",706 message=str(exc),707 http_status=400,708 ) from exc709 body_out = {"state": _to_jsonable(st), "turn": st.turn}710 return JSONResponse(status_code=200, content=body_out)711 712 713async def _handle_close(request: Request, state: _AppState) -> Response:714 _check_bearer(request, state)715 _check_models_ready(state)716 sid = _check_session_header(request)717 entry = state.cache.evict(sid)718 if entry is None:719 return JSONResponse(status_code=200, content={"closed": True, "final_state": None})720 final_state: Any = None721 try:722 final_state = _to_jsonable(entry.env.state())723 except EnvNotReadyError:724 final_state = None725 try:726 entry.env.close()727 except Exception:728 logger.exception("env.close raised on /close for sid=%s", sid)729 return JSONResponse(status_code=200, content={"closed": True, "final_state": final_state})730 731 732# ---------------------------------------------------------------------------733# App factory + route wiring734# ---------------------------------------------------------------------------735 736 737def create_app() -> FastAPI:738 """Construct a fresh FastAPI app. Used by tests to get an isolated instance."""739 app = FastAPI(lifespan=lifespan, title="DriftCall Env", version="0.1.0")740 app.add_middleware(_BodySizeMiddleware, max_bytes=_MAX_BODY_BYTES)741 742 @app.get("/healthz", response_class=PlainTextResponse)743 async def healthz() -> PlainTextResponse:744 return PlainTextResponse(content="ok", status_code=200)745 746 @app.post("/reset")747 async def reset_route(request: Request) -> Response:748 try:749 return await _handle_reset(request, _get_state(app))750 except _ApiError as err:751 return _error_response(err, _request_id(request))752 753 @app.post("/step")754 async def step_route(request: Request) -> Response:755 try:756 return await _handle_step(request, _get_state(app))757 except _ApiError as err:758 return _error_response(err, _request_id(request))759 760 @app.get("/state")761 async def state_route(request: Request) -> Response:762 try:763 return await _handle_state(request, _get_state(app))764 except _ApiError as err:765 return _error_response(err, _request_id(request))766 767 @app.post("/close")768 async def close_route(request: Request) -> Response:769 try:770 return await _handle_close(request, _get_state(app))771 except _ApiError as err:772 return _error_response(err, _request_id(request))773 774 return app775 776 777app = create_app()778 779 780__all__ = [781 "SessionCache",782 "SessionEntry",783 "app",784 "create_app",785 "lifespan",786]787