state-hub/api/main.py

import hashlib
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response as StarletteResponse

from api.database import engine
from api.events import shutdown_publisher
from api.routers import decisions, extension_points, progress, state, tasks, technical_debt, topics, workstreams, workstream_dependencies
from api.routers import domains, repos, contributions, sbom, policy, domain_goals, repo_goals, messages, capability_requests, tpsc
from api.routers import token_events
from api.routers import interface_changes
from api.routers import flows


class ETagMiddleware(BaseHTTPMiddleware):
    """Add ETag + conditional-GET (304) support to all JSON GET responses."""

    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        if request.method != "GET":
            return response
        if "application/json" not in response.headers.get("content-type", ""):
            return response

        body_parts = []
        async for chunk in response.body_iterator:
            body_parts.append(chunk)
        body = b"".join(body_parts)

        etag = '"' + hashlib.md5(body, usedforsecurity=False).hexdigest() + '"'
        if request.headers.get("if-none-match") == etag:
            return StarletteResponse(
                status_code=304,
                headers={"ETag": etag, "Cache-Control": "no-cache"},
            )

        headers = {k: v for k, v in response.headers.items() if k.lower() != "content-length"}
        headers["ETag"] = etag
        if not any(k.lower() == "cache-control" for k in headers):
            headers["Cache-Control"] = "no-cache"
        return StarletteResponse(
            content=body,
            status_code=response.status_code,
            headers=headers,
            media_type=response.media_type,
        )


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield
    await shutdown_publisher()
    await engine.dispose()


app = FastAPI(
    title="Custodian State Hub",
    description="Local-first state API for the Custodian agent system.",
    version="0.6.0",
    lifespan=lifespan,
)

_cors_env = os.getenv("CORS_ORIGINS", "http://localhost:3000,http://127.0.0.1:3000")
_cors_origins = [o.strip() for o in _cors_env.split(",") if o.strip()]

app.add_middleware(ETagMiddleware)
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_methods=["GET", "POST", "PATCH", "DELETE", "PUT"],
    allow_headers=["Content-Type", "If-None-Match"],
    expose_headers=["ETag"],
)

app.include_router(domains.router)
app.include_router(repos.router)
app.include_router(topics.router)
app.include_router(workstreams.router)
app.include_router(workstream_dependencies.router)
app.include_router(tasks.router)
app.include_router(decisions.router)
app.include_router(extension_points.router)
app.include_router(technical_debt.router)
app.include_router(progress.router)
app.include_router(domain_goals.router)
app.include_router(repo_goals.router)
app.include_router(contributions.router)
app.include_router(sbom.router)
app.include_router(messages.router)
app.include_router(capability_requests.router)
app.include_router(tpsc.router)
app.include_router(token_events.router)
app.include_router(interface_changes.router)
app.include_router(flows.router)
app.include_router(state.router)
app.include_router(policy.router)


@app.get("/", include_in_schema=False)
async def root():
    return {"service": "state-hub", "docs": "/docs"}