clawdforge/clawdforge/server.py

"""FastAPI app exposing /run, /files, /admin/tokens, /healthz."""
import os
import time
from pathlib import Path
from typing import Annotated

from fastapi import FastAPI, Header, HTTPException, Request, UploadFile, File, Form
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from .auth import Auth
from .config import load
from .runner import Runner
from .store import Store


cfg = load()
store = Store(cfg.db_path)
auth = Auth(
    store=store,
    global_cidrs=list(cfg.allow_cidrs),
    admin_token=cfg.admin_bootstrap_token,
)
runner = Runner(
    claude_bin=cfg.claude_bin,
    default_model=cfg.default_model,
    default_timeout=cfg.default_timeout_secs,
    runs_dir=cfg.runs_dir,
)

app = FastAPI(title="clawdforge", version="0.1.0")


# ---------- schemas ----------------------------------------------------------


class RunRequest(BaseModel):
    prompt: str = Field(min_length=1)
    model: str | None = None
    system: str | None = None
    files: list[str] | None = None
    timeout_secs: int | None = Field(default=None, ge=5, le=600)


class TokenCreateRequest(BaseModel):
    name: str = Field(min_length=1, max_length=64, pattern=r"^[a-z0-9][a-z0-9_-]*$")
    ip_cidrs: list[str] = Field(default_factory=list)


# ---------- endpoints --------------------------------------------------------


@app.get("/healthz")
def healthz(request: Request):
    auth.require_global_ip(request)
    # Smoke-check that claude binary exists and responds (--version is fast, no auth needed)
    import shutil, subprocess
    found = shutil.which(cfg.claude_bin) is not None
    version = None
    if found:
        try:
            r = subprocess.run([cfg.claude_bin, "--version"], capture_output=True, text=True, timeout=5)
            version = (r.stdout or r.stderr or "").strip().splitlines()[0] if (r.stdout or r.stderr) else None
        except Exception as e:
            version = f"err: {e}"
    return {"ok": True, "claude_present": found, "claude_version": version}


@app.post("/run")
def run(
    request: Request,
    body: RunRequest,
    authorization: Annotated[str | None, Header()] = None,
):
    rec = auth.require_app(request, authorization)

    # Resolve any file_tokens to actual paths owned by this app
    file_paths: list[str] = []
    if body.files:
        for ftoken in body.files:
            path = store.resolve_file(ftoken, rec["name"])
            if not path:
                raise HTTPException(404, f"unknown or expired file token: {ftoken}")
            file_paths.append(path)

    started_at = int(time.time())
    res = runner.run(
        prompt=body.prompt,
        model=body.model,
        system=body.system,
        files=file_paths or None,
        timeout_secs=body.timeout_secs,
    )

    result_chars = len(str(res.result)) if res.result is not None else 0
    store.log_run(
        app_name=rec["name"],
        started_at=started_at,
        duration_ms=res.duration_ms,
        model=body.model or cfg.default_model,
        prompt_chars=len(body.prompt),
        result_chars=result_chars,
        ok=res.ok,
        error=res.error,
        file_count=len(file_paths),
    )

    if not res.ok:
        return JSONResponse(
            status_code=502,
            content={
                "ok": False,
                "error": res.error,
                "stderr": res.raw_stderr[-4000:],
                "duration_ms": res.duration_ms,
                "stop_reason": res.stop_reason,
            },
        )

    return {
        "ok": True,
        "result": res.result,
        "duration_ms": res.duration_ms,
        "stop_reason": res.stop_reason,
    }


@app.post("/files")
def upload_file(
    request: Request,
    file: Annotated[UploadFile, File()],
    ttl_secs: Annotated[int, Form()] = 3600,
    authorization: Annotated[str | None, Header()] = None,
):
    rec = auth.require_app(request, authorization)
    if ttl_secs < 60 or ttl_secs > 86400:
        raise HTTPException(400, "ttl_secs out of range (60..86400)")

    safe_name = (file.filename or "upload").replace("/", "_").replace("..", "_")
    target_dir = Path(cfg.runs_dir) / "uploads" / rec["name"]
    target_dir.mkdir(parents=True, exist_ok=True)

    import secrets as _s
    suffix = _s.token_hex(8)
    target = target_dir / f"{int(time.time())}-{suffix}-{safe_name}"
    with target.open("wb") as fh:
        while True:
            chunk = file.file.read(1024 * 1024)
            if not chunk:
                break
            fh.write(chunk)

    file_token = store.register_file(rec["name"], str(target), ttl_secs)
    return {"file_token": file_token, "ttl_secs": ttl_secs, "size": target.stat().st_size}


@app.post("/admin/tokens")
def create_token(
    request: Request,
    body: TokenCreateRequest,
    authorization: Annotated[str | None, Header()] = None,
):
    auth.require_admin(request, authorization)
    try:
        token = store.create_token(body.name, body.ip_cidrs)
    except Exception as e:
        raise HTTPException(409, f"token create failed: {e}")
    return {"name": body.name, "token": token, "ip_cidrs": body.ip_cidrs}


@app.get("/admin/tokens")
def list_tokens(
    request: Request,
    authorization: Annotated[str | None, Header()] = None,
):
    auth.require_admin(request, authorization)
    return {"tokens": store.list_tokens()}


@app.delete("/admin/tokens/{name}")
def revoke_token(
    name: str,
    request: Request,
    authorization: Annotated[str | None, Header()] = None,
):
    auth.require_admin(request, authorization)
    if not store.revoke_token(name):
        raise HTTPException(404, "no such token")
    return {"ok": True}


@app.on_event("startup")
def _startup_gc():
    # Best-effort GC of expired staged files at boot
    try:
        for p in store.gc_expired_files():
            try:
                os.remove(p)
            except FileNotFoundError:
                pass
    except Exception:
        pass