arena/api.py

"""FastAPI transport for the unified Arena service."""

from __future__ import annotations

import json
import os
from typing import AsyncIterator

from fastapi import FastAPI, HTTPException, Query, Request, Response
from pydantic import BaseModel, Field
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse, StreamingResponse

from .model_catalog import ModelProviderConfig, ProviderModel
from .run_models import RunRequest, RunView
from .run_store import InvalidRunToken, RunNotFound
from .validation_models import ValidationReport
from .service import (
    ArenaService,
    FileAccessFailed,
)
from .skill_catalog import SkillInfo
from .thread_inspector import (
    AgentRunFailed,
    FileEntry,
    MessageResponse,
    ThreadInspector,
    ThreadNotFound,
    ThreadResponse,
)


app = FastAPI(title="Backyard Demo Builder", version="0.1.0")
service = ArenaService()
thread_inspector = ThreadInspector()
sessions = thread_inspector.sessions

API_KEY = os.getenv("ARENA_API_KEY", "")


class AuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.url.path in ("/health",) or not API_KEY:
            return await call_next(request)
        auth = request.headers.get("Authorization", "")
        expected = f"Bearer {API_KEY}"
        if auth != expected:
            return JSONResponse(
                status_code=401,
                content={"detail": "Unauthorized"},
            )
        return await call_next(request)


if API_KEY:
    app.add_middleware(AuthMiddleware)


class MessageRequest(BaseModel):
    content: str = Field(min_length=1)


class RejoinRequest(BaseModel):
    token: str = Field(min_length=1)


class RunMessageRequest(BaseModel):
    content: str = Field(min_length=1)


@app.get("/health")
async def health() -> dict[str, str]:
    return service.health()


@app.get("/skills")
async def list_skills() -> list[SkillInfo]:
    return service.list_skills()


@app.post("/models/catalog")
async def list_provider_models(req: ModelProviderConfig) -> list[ProviderModel]:
    try:
        return service.list_provider_models(req)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=502, detail="Could not fetch provider models") from exc


@app.post("/models/catalog/refresh")
async def refresh_provider_models(req: ModelProviderConfig) -> list[ProviderModel]:
    try:
        return service.refresh_provider_models(req)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=502, detail="Could not refresh provider models") from exc


@app.post("/runs")
async def create_run(req: RunRequest) -> RunView:
    return service.create_run(req).run


@app.post("/runs/start")
async def start_run(req: RunRequest) -> RunView:
    return service.start_run(req).run


@app.get("/runs/{run_id}")
async def get_run(run_id: str) -> RunView:
    try:
        return service.get_run(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc


@app.get("/runs/{run_id}/events")
async def stream_run_events(run_id: str) -> StreamingResponse:
    try:
        service.get_run(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc

    async def event_generator() -> AsyncIterator[str]:
        import asyncio
        q = service.subscribe_events(run_id)
        loop = asyncio.get_event_loop()
        while True:
            event = await loop.run_in_executor(None, q.get)
            if event is None:
                break
            yield f"data: {json.dumps(event, default=str)}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/runs/{run_id}/codebase.zip")
async def get_run_codebase_zip(run_id: str) -> Response:
    try:
        payload = service.get_run_codebase_zip(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    except FileAccessFailed as exc:
        raise HTTPException(status_code=409, detail="Run has no codebase archive") from exc
    return Response(
        content=payload.data,
        media_type=payload.content_type,
        headers={"Content-Disposition": f'attachment; filename="{payload.filename}"'},
    )


@app.get("/runs/{run_id}/field_notes.md")
async def get_run_field_notes(run_id: str) -> Response:
    try:
        payload = service.get_run_field_notes(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    return Response(
        content=payload.data,
        media_type=payload.content_type,
        headers={"Content-Disposition": f'attachment; filename="{payload.filename}"'},
    )


@app.get("/runs/{run_id}/trace.json")
async def get_run_trace_json(run_id: str) -> Response:
    try:
        payload = service.get_run_trace_json(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    return Response(
        content=payload.data,
        media_type=payload.content_type,
        headers={"Content-Disposition": f'attachment; filename="{payload.filename}"'},
    )


@app.get("/runs/{run_id}/submission.md")
async def get_run_submission_markdown(run_id: str) -> Response:
    try:
        payload = service.get_run_submission_markdown(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    return Response(
        content=payload.data,
        media_type=payload.content_type,
        headers={"Content-Disposition": f'attachment; filename="{payload.filename}"'},
    )


@app.get("/runs/{run_id}/space_README.md")
async def get_run_space_readme(run_id: str) -> Response:
    try:
        payload = service.get_run_space_readme(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    return Response(
        content=payload.data,
        media_type=payload.content_type,
        headers={"Content-Disposition": f'attachment; filename="{payload.filename}"'},
    )


@app.post("/runs/rejoin")
async def rejoin_run(req: RejoinRequest) -> RunView:
    try:
        return service.rejoin_run(req.token).run
    except InvalidRunToken as exc:
        raise HTTPException(status_code=404, detail="Token not found") from exc


@app.post("/runs/{run_id}/cancel")
async def cancel_run(run_id: str) -> RunView:
    try:
        return service.cancel_run(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc


@app.post("/runs/{run_id}/messages")
async def add_run_message(run_id: str, req: RunMessageRequest) -> RunView:
    try:
        return service.add_run_message(run_id, req.content)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc


@app.post("/runs/{run_id}/validate")
async def validate_run(run_id: str) -> ValidationReport:
    try:
        return service.validate_run(run_id)
    except RunNotFound as exc:
        raise HTTPException(status_code=404, detail="Run not found") from exc
    except FileAccessFailed as exc:
        raise HTTPException(status_code=409, detail="Run has no codebase archive") from exc


@app.post("/threads")
async def create_thread() -> ThreadResponse:
    try:
        return thread_inspector.create_thread()
    except Exception as exc:
        raise HTTPException(status_code=500, detail="Could not create agent session") from exc


@app.post("/threads/{thread_id}/messages")
async def send_message(thread_id: str, req: MessageRequest) -> MessageResponse:
    try:
        return thread_inspector.send_message(thread_id, req.content)
    except ThreadNotFound as exc:
        raise HTTPException(status_code=404, detail="Thread not found") from exc
    except AgentRunFailed as exc:
        raise HTTPException(status_code=502, detail="Agent run failed") from exc


@app.get("/threads/{thread_id}/files")
async def list_files(thread_id: str, path: str = Query("/")) -> list[FileEntry]:
    try:
        return thread_inspector.list_files(thread_id, path)
    except ThreadNotFound as exc:
        raise HTTPException(status_code=404, detail="Thread not found") from exc
    except FileAccessFailed as exc:
        raise HTTPException(status_code=400, detail="Could not list files") from exc


@app.get("/threads/{thread_id}/files/content")
async def read_file(thread_id: str, path: str) -> dict[str, str]:
    try:
        return thread_inspector.read_file(thread_id, path)
    except ThreadNotFound as exc:
        raise HTTPException(status_code=404, detail="Thread not found") from exc
    except FileAccessFailed as exc:
        raise HTTPException(status_code=400, detail="Could not read file") from exc


@app.delete("/threads/{thread_id}")
async def delete_thread(thread_id: str) -> dict[str, bool]:
    try:
        return thread_inspector.delete_thread(thread_id)
    except ThreadNotFound as exc:
        raise HTTPException(status_code=404, detail="Thread not found") from exc