"""Same-origin Studio and strict Decision Gateway.""" import asyncio import json import logging import os import time from contextlib import asynccontextmanager from pathlib import Path from fastapi import FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from starlette.concurrency import run_in_threadpool from archive_cache import NOTICE as ARCHIVE_NOTICE from archive_cache import ArchiveCache from contract import MODEL, to_records from direct_gateway import DirectGateway, DirectGatewayError from engine import Busy, Engine, Unavailable from relay import Relay, RelayError from systemone_api import ( arrange_models, public_models, resolve_canonical_model, resolve_model, public_response, sdk_response, ) from tetris_arena import ( ArenaBusyError, ArenaUpstreamError, ArenaValidationError, HTTPDecisionAdapter, RaceManager, encode_sse, ) ROOT = Path(__file__).resolve().parent logger = logging.getLogger("decision.studio") def unique_object(pairs): result = {} for key, value in pairs: if key in result: raise ValueError("Duplicate JSON key: " + key) result[key] = value return result async def read_json(request, limit=256 * 1024, *, with_size=False): if request.headers.get("content-type", "").split(";", 1)[0].strip().lower() != "application/json": raise HTTPException(415, "Use Content-Type: application/json") body = bytearray() async for chunk in request.stream(): body.extend(chunk) if len(body) > limit: raise HTTPException(413, "Request is too large. No field is truncated.") try: document = json.loads(body, object_pairs_hook=unique_object, parse_constant=lambda x: (_ for _ in ()).throw(ValueError("Nonfinite JSON value"))) return (document, len(body)) if with_size else document except (ValueError, TypeError, UnicodeError, RecursionError) as exc: raise HTTPException(422, str(exc)) from None async def read_input(request): payload = await read_json(request) try: records = to_records(payload) except (ValueError, TypeError, UnicodeError, RecursionError) as exc: raise HTTPException(422, str(exc)) from None return payload, records def create_app( mode=None, *, native_engine=None, relay=None, registry=None, relays=None, direct_gateway=None, tetris_adapter=None, tetris_manager=None, archive=None, ): mode = mode or os.getenv("DECISION_BACKEND", "native") if mode not in {"native", "pull_queue", "direct", "cache"}: raise ValueError("DECISION_BACKEND must be native, pull_queue, direct, or cache") # cache: an archive that answers only from recorded results and calls no model. archive = (archive or ArchiveCache.load(ROOT)) if mode == "cache" else None from model_registry import model_registry if relay is not None and registry is None: registry = [{'id': MODEL, 'label': 'Kai', 'version': '1.0', 'manifest_sha256': relay.manifest}] registry = arrange_models(model_registry(registry)) studio_generation = os.getenv("DECISION_STUDIO_GENERATION", "").strip() if studio_generation not in {"", "1.0", "2.0"}: raise ValueError("DECISION_STUDIO_GENERATION must be 1.0 or 2.0 when set") studio_registry = { wire: item for wire, item in registry.items() if not studio_generation or item['version'] == studio_generation } if not studio_registry: raise ValueError("The Studio generation must include a configured model") default_model = next(iter(studio_registry)) canonical_models = tuple(item['repo_id'] for item in registry.values()) expected_artifacts = {item['repo_id']: item for item in registry.values()} owns_direct_gateway = mode == "direct" and direct_gateway is None if mode == "direct": if any('revision' not in item for item in registry.values()): raise ValueError("Direct serving mode requires a pinned Hub revision for every model") direct_gateway = direct_gateway or DirectGateway.from_environment(expected_artifacts) if getattr(direct_gateway, "models", None) != frozenset(canonical_models): raise ValueError( "The direct Gateway must match every configured canonical model" ) gateway_pins = getattr(direct_gateway, "expected_artifacts", None) if gateway_pins is not None and any( gateway_pins.get(model) != { "revision": item["revision"], "manifest_sha256": item["manifest_sha256"], **({"content_sha256": item["content_sha256"]} if "content_sha256" in item else {}), } for model, item in expected_artifacts.items() ): raise ValueError("The direct Gateway artifact pins must match the registry") elif direct_gateway is not None: raise ValueError("A direct Gateway may only be supplied in direct mode") owns_tetris_manager = tetris_manager is None owns_tetris_adapter = owns_tetris_manager and tetris_adapter is None @asynccontextmanager async def lifespan(_api): yield try: if owns_tetris_manager: try: await tetris_manager.aclose() finally: if owns_tetris_adapter: await tetris_adapter.close() finally: if owns_direct_gateway: await direct_gateway.aclose() api = FastAPI( title="Decision Studio", version="0.6.0", docs_url="/api/docs", redoc_url=None, lifespan=lifespan, ) api.state.engine = native_engine or Engine() if mode == "native" and (set(registry) != {MODEL} or api.state.engine.model != MODEL or api.state.engine.manifest != registry[MODEL]['manifest_sha256']): raise ValueError("Native mode requires only Decision 1.0 Kai with its matching manifest") if mode == "pull_queue": relays = relays or ({MODEL: relay} if relay else { key: Relay(os.getenv("DECISION_WORKER_TOKEN", ""), item['manifest_sha256'], model=key, complete_input_tokens=item['complete_input_tokens']) for key, item in registry.items()}) if set(relays) != set(registry) or any(r.model != key or r.manifest != registry[key]['manifest_sha256'] or r.complete_input_tokens != registry[key]['complete_input_tokens'] for key, r in relays.items()): raise ValueError("Every queue must match its registry ID and native manifest") else: relays = {} if relays and len({r.token for r in relays.values()}) != 1: raise ValueError("This gateway uses one shared private worker token") api.state.relays = relays api.state.relay = relays.get(MODEL) # backward-compatible default for local tooling api.state.registry = registry api.state.direct = direct_gateway api.state.archive = archive def select(model): if not isinstance(model, str) or model not in registry: raise HTTPException(422, "This model is not available in this Studio") return relays.get(model) async def direct_readiness(models): """Probe distinct configured Decision models without exposing their origins.""" selected = tuple(dict.fromkeys(models)) observations = await asyncio.gather( *(direct_gateway.probe(model) for model in selected), return_exceptions=True, ) return { model: isinstance(observed, dict) and observed.get('loaded') is True for model, observed in zip(selected, observations, strict=True) } async def tetris_model_readiness(competitor_ids=None): if mode not in {"direct", "pull_queue"}: return {} competitors = ( api.state.tetris.catalog() if hasattr(api.state.tetris, 'catalog') else () ) selected = { competitor.id: competitor.request_model for competitor in competitors if competitor.family == 'decision' and competitor.ready and (competitor_ids is None or competitor.id in competitor_ids) } if mode == "direct": by_model = await direct_readiness(selected.values()) else: wire_by_canonical = { item['repo_id']: wire for wire, item in registry.items() } wires = tuple(dict.fromkeys( wire_by_canonical[model] for model in selected.values() if model in wire_by_canonical )) observations = await asyncio.gather( *(relays[wire].status() for wire in wires), return_exceptions=True, ) by_model = { registry[wire]['repo_id']: isinstance(observed, dict) and observed.get('loaded') is True for wire, observed in zip(wires, observations, strict=True) } return { competitor_id: by_model.get(model, False) for competitor_id, model in selected.items() } async def input_for(request, *, public_contract=None): payload, request_bytes = await read_json( request, limit=2 * 1024 * 1024, with_size=True, ) return prepare_input(payload, request_bytes, public_contract=public_contract) def prepare_input(payload, request_bytes, *, public_contract=None): if request_bytes > 256 * 1024 and ( not isinstance(payload, dict) or 'states' not in payload ): raise HTTPException(413, "Request is too large. No field is truncated.") if public_contract is not None: expected = {'model', 'questions', 'states' if public_contract == 'batch' else 'state'} if not isinstance(payload, dict) or set(payload) != expected: raise HTTPException(422, f"Provide exactly {', '.join(sorted(expected))}.") requested_model = payload.get('model') if isinstance(payload, dict) else None model = ( resolve_canonical_model(requested_model, registry) if public_contract is not None else resolve_model( requested_model if mode == "direct" else requested_model or default_model, registry, ) ) selected = select(model) public_model = registry[model]['repo_id'] canonical_payload = ( dict(payload, model=public_model) if isinstance(payload, dict) else payload ) if isinstance(payload, dict): payload = dict(payload, model=model) try: records = to_records(payload, model=model) except (ValueError, TypeError, UnicodeError, RecursionError) as exc: if ( mode == "direct" and "Expanded context/question input exceeds" in str(exc) ): raise HTTPException(413, str(exc)) from None raise HTTPException(422, str(exc)) from None return payload, records, selected, public_model, canonical_payload, model async def queued_result(selected, payload, public_model, public_contract, is_disconnected=None): job = await selected.submit(payload, keep_until_read=True) deadline = time.monotonic() + 45 try: while time.monotonic() < deadline: remaining = deadline - time.monotonic() current = await selected.wait(job["id"], max(0.0, min(0.25, remaining))) if current["status"] == "succeeded": if current['result'].get('model') == public_model: return current['result'] if public_contract is not None: return public_response(current['result'], public_model=public_model, batch=public_contract == 'batch') return sdk_response(current['result']) if current["status"] not in Relay.ACTIVE: raise HTTPException(503, current.get("detail", "The request did not finish")) if is_disconnected is not None and await is_disconnected(): raise asyncio.CancelledError() raise HTTPException(504, "Request exceeded the synchronous wait. Use the Studio asynchronous job API.") finally: await selected.release(job["id"]) async def arena_turn(payload): """Run one Tetris turn as a strict /v1/systemone request on the selected model's queue.""" try: payload, _, selected, public_model, _, _ = prepare_input( payload, len(json.dumps(payload).encode()), public_contract='single', ) return await queued_result(selected, payload, public_model, 'single') except (HTTPException, RelayError) as exc: raise ArenaUpstreamError("The selected model could not complete this turn.") from exc if tetris_manager is None: tetris_adapter = tetris_adapter or HTTPDecisionAdapter.from_environment( # The archive reads no upstream URL or key from the environment. {} if archive else None, allowed_local_models=canonical_models, include_cloud=studio_generation != "2.0" and not archive, local_endpoints=( direct_gateway.local_systemone_endpoints() if mode == "direct" else None ), direct_gateway=direct_gateway if mode == "direct" else None, decision_queue=( arena_turn if mode == "pull_queue" else archive.arena_turn if archive else None ), ) tetris_manager = RaceManager.from_environment( tetris_adapter, archive.race_environment(os.environ) if archive else None, ) api.state.tetris = tetris_manager @api.exception_handler(RelayError) async def relay_error(request, exc): return JSONResponse({"detail": exc.message}, status_code=exc.code, headers={"Retry-After": "2"} if exc.code in {429, 503} else None) @api.exception_handler(DirectGatewayError) async def direct_error(request, exc): headers = {"Retry-After": exc.retry_after} if exc.retry_after else None return JSONResponse({"detail": exc.message}, status_code=exc.code, headers=headers) @api.get("/api/status") async def status(model: str = default_model): model = resolve_model(model, registry) selected = select(model) if archive: return { "configured": True, "loaded": True, "phase": "ready", "model": model, "manifest_sha256": registry[model]['manifest_sha256'], "revision": registry[model].get('revision'), "complete_input_tokens": registry[model]['complete_input_tokens'], "context_batch": True, "backend": "cache", "live_inference": False, "running": 0, "queued": 0, "readiness": "recorded results only; this archive calls no model", } if mode == "direct": observed = await direct_gateway.probe(registry[model]['repo_id']) loaded = observed['loaded'] artifact = observed['artifact'] or {} return { "configured": True, "loaded": loaded, "phase": "ready" if loaded else "unavailable", "model": model, "manifest_sha256": registry[model]['manifest_sha256'], "revision": registry[model]['revision'], "loaded_revision": artifact.get('revision'), "loaded_manifest_sha256": artifact.get('manifest_sha256'), "loaded_content_sha256": artifact.get('content_sha256'), "complete_input_tokens": registry[model]['complete_input_tokens'], "context_batch": loaded, "backend": "direct", "live_inference": loaded, "running": False, "queued": 0, } return await selected.status() if selected else api.state.engine.status() @api.get("/api/ready") async def ready(): """Return HTTP 503 when any configured model is unavailable.""" if archive: loaded = {model: True for model in canonical_models} elif mode == 'direct': loaded = await direct_readiness(canonical_models) elif relays: observations = await asyncio.gather( *(relay.status() for relay in relays.values()), return_exceptions=True, ) loaded = { registry[key]['repo_id']: isinstance(observed, dict) and observed.get('loaded') is True for key, observed in zip(relays, observations, strict=True) } else: observed = api.state.engine.status() loaded = { registry[default_model]['repo_id']: observed.get('loaded') is True } healthy = bool(loaded) and all(loaded.values()) return JSONResponse( { 'backend': mode, 'status': 'ready' if healthy else 'unavailable', 'models': loaded, }, status_code=200 if healthy else 503, ) @api.get("/api/examples") def examples(): return json.loads((ROOT / "examples.json").read_text()) @api.get("/api/tetris/config") async def tetris_config(): # Credentials and upstream URLs intentionally never enter this response. config = api.state.tetris.public_config() if studio_generation: visible_models = {item['repo_id'] for item in studio_registry.values()} visible_competitors = { competitor.id for competitor in api.state.tetris.catalog() if competitor.family == 'decision' and competitor.request_model in visible_models } config['competitors'] = [ competitor for competitor in config.get('competitors', ()) if (competitor['family'] != 'decision' or competitor['id'] in visible_competitors) and (studio_generation != "2.0" or competitor['family'] != 'cloud') ] readiness = await tetris_model_readiness({ competitor['id'] for competitor in config.get('competitors', ()) }) for competitor in config.get('competitors', ()): if competitor['id'] in readiness: competitor['ready'] = competitor['ready'] and readiness[competitor['id']] return config @api.post("/api/tetris/races", status_code=201) async def create_tetris_race(request: Request): payload = await read_json(request) if archive and isinstance(payload, dict) and not archive.replays_race(payload): raise HTTPException(422, archive.race_notice()) if mode in {'direct', 'pull_queue'} and isinstance(payload, dict): selected_ids = { value for side in ('left', 'right') if isinstance(value := payload.get(side), str) } readiness = await tetris_model_readiness(selected_ids) if any(not ready for ready in readiness.values()): raise HTTPException(503, "The selected Decision model is unavailable.") try: race = await api.state.tetris.create(payload) except ArenaValidationError as exc: headers = ( {"Retry-After": str(exc.retry_after)} if isinstance(exc, ArenaBusyError) else None ) raise HTTPException(exc.status_code, str(exc), headers=headers) from None return { "id": race.id, "status": race.status, "events_url": f"/api/tetris/races/{race.id}/events", "result_url": f"/api/tetris/races/{race.id}", } def get_tetris_race(race_id): try: return api.state.tetris.get(race_id) except KeyError: raise HTTPException(404, "Race not found") from None @api.get("/api/tetris/races/{race_id}") async def tetris_race(race_id: str): return get_tetris_race(race_id).snapshot() @api.get("/api/tetris/races/{race_id}/trace") async def tetris_race_trace(race_id: str): return get_tetris_race(race_id).snapshot(include_traces=True) @api.get("/api/tetris/races/{race_id}/events") async def tetris_race_events(race_id: str, request: Request, after: int = 0): race = get_tetris_race(race_id) header_cursor = request.headers.get("last-event-id") if header_cursor: try: after = max(after, int(header_cursor)) except ValueError: raise HTTPException(400, "Last-Event-ID must be an integer") from None if after < 0: raise HTTPException(400, "after must not be negative") async def stream(): async for event in race.iter_events(after): yield encode_sse(event) return StreamingResponse( stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-store", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) @api.delete("/api/tetris/races/{race_id}", status_code=204) async def cancel_tetris_race(race_id: str): try: await api.state.tetris.cancel(race_id) except KeyError: raise HTTPException(404, "Race not found") from None @api.get("/v1/models") def models(): synchronous_wait = ( getattr(direct_gateway, "timeout_seconds", 45) if mode == "direct" else 45 ) limits = { "request_bytes": 256 * 1024, "expanded_input_bytes": 16 * 1024 * 1024, "questions": None, "contexts": None, # direct capacity belongs to each live model instance, not the # retired Gateway's fixed eight-request/eight-row profile. "active_requests_per_model": None if mode == "direct" else 8, "gpu_microbatch": None if mode == "direct" else 8, "synchronous_wait_seconds": synchronous_wait, } result = { "models": public_models(studio_registry), "limits": limits, } if mode != "direct": result["default"] = registry[default_model]['repo_id'] return result @api.post("/v1/systemone") async def evaluate_systemone(request: Request): return await evaluate(request, public_contract='single') @api.post("/v1/systemone/batches") async def evaluate_decision_batches(request: Request): return await evaluate(request, public_contract='batch') @api.post("/api/evaluate") async def evaluate_studio(request: Request): return await evaluate(request, public_contract=None) async def evaluate(request: Request, *, public_contract: str | None): payload, records, selected, public_model, canonical_payload, _ = await input_for( request, public_contract=public_contract ) if archive: recorded = ( archive.turn(canonical_payload) if public_contract == 'single' else archive.answer(canonical_payload) if public_contract is None else None ) if recorded is None: raise HTTPException(410, ARCHIVE_NOTICE) return recorded if mode == "direct": return await direct_gateway.evaluate( canonical_payload, batch='states' in canonical_payload, ) if selected: return await queued_result(selected, payload, public_model, public_contract, request.is_disconnected) try: result = await run_in_threadpool(api.state.engine.evaluate, payload, records) if public_contract is not None: return public_response(result, public_model=public_model, batch=public_contract == 'batch') return sdk_response(result) except Unavailable as exc: raise HTTPException(503, str(exc)) from None except Busy as exc: raise HTTPException(429, str(exc), headers={"Retry-After": "2"}) from None except ValueError as exc: if "tokens" in str(exc) or "room for state" in str(exc): raise HTTPException(422, str(exc)) from None logger.warning("Native validation failure: %s", type(exc).__name__) raise HTTPException(503, "The configured native package could not be validated. Check the server configuration.") from None except Exception as exc: logger.warning("Inference failure: %s", type(exc).__name__) raise HTTPException(503, "The model could not finish this request. No prediction was substituted.") from None if relays: @api.post("/api/jobs", status_code=202) async def submit(request: Request): payload, _, selected, _, _, _ = await input_for(request) return await selected.submit(payload) @api.get("/api/jobs/{job_id}") async def read_job(job_id: str, model: str = default_model): return await select(resolve_model(model, registry)).read(job_id) @api.delete("/api/jobs/{job_id}") async def cancel_job(job_id: str, model: str = default_model): return await select(resolve_model(model, registry)).cancel(job_id) async def worker_body(request, required, optional=(), *, limit=512 * 1024): # All queues use the same private trust boundary, but identities are mandatory. # Authentication precedes reading user data. The token is never returned or logged. next(iter(relays.values())).authenticate(request.headers.get("authorization", "")) body = await read_json(request, limit=limit) required = set(required) | {"model", "manifest_sha256"} if not isinstance(body, dict) or set(body) - required - set(optional) or required - set(body): raise HTTPException(422, "Invalid worker request fields") selected = select(body['model']) selected.authenticate(request.headers.get("authorization", "")) if body['manifest_sha256'] != selected.manifest: raise HTTPException(409, "Worker model and native manifest do not match this queue") return body, selected @api.post("/internal/worker/heartbeat") async def heartbeat(request: Request): body, selected = await worker_body(request, {"worker_id", "phase"}, {"capabilities"}) if any(other is not selected and other.worker and other.worker['id'] == body['worker_id'] for other in relays.values()): raise HTTPException(409, "A worker identity cannot serve multiple models") return await selected.heartbeat(body["worker_id"], body["manifest_sha256"], body["phase"], body.get("capabilities")) @api.post("/internal/worker/claim") async def claim(request: Request): body, selected = await worker_body(request, {"worker_id"}, {"wait_seconds"}) return await selected.claim(body["worker_id"], body.get("wait_seconds", 20)) @api.post("/internal/worker/result") async def complete(request: Request): # Leave room for the worker/lease envelope around an 8 MiB # attested runtime response. body, selected = await worker_body(request, {"worker_id", "id", "lease_token"}, {"result", "error_code"}, limit=8 * 1024 * 1024 + 64 * 1024) return await selected.complete(body["worker_id"], body["id"], body["lease_token"], body.get("result"), body.get("error_code")) @api.middleware("http") async def response_headers(request, call_next): response = await call_next(request) response.headers["X-Content-Type-Options"] = "nosniff" response.headers["Referrer-Policy"] = "no-referrer" if request.url.path.startswith(("/api/", "/v1/", "/internal/", "/tetris/")): response.headers["Cache-Control"] = "no-store" return response api.mount("/", StaticFiles(directory=ROOT / "static", html=True), name="studio") return api app = create_app()