state-hub/api/main.py

import hashlib
import os
import time
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.services.write_idempotency import WriteIdempotencyMiddleware
from api.routers import decisions, extension_points, intake, ops_runs, progress, state, suggestions, tasks, technical_debt, topics, workstreams, workstream_dependencies
from api.routers import domains, repos, contributions, sbom, policy, domain_goals, repo_goals, messages, memo_notifications, capability_requests, tpsc, services
from api.routers import token_events
from api.routers import message_feed
from api.routers import interface_changes
from api.routers import flows
from api.routers import recently_on_scope
from api.routers import consistency_sweep
from api.routers import reconciliation
from api.routers import execution
from api.routers import fabric
from api.routers import legacy_meter
from api.routers import review_contracts
from api.routers import identifier_migrations
from api.routers import repository_renames
from api.routers import work_record_projection


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

    async def dispatch(self, request: Request, call_next):
        started = time.perf_counter()
        response = await call_next(request)
        if request.method != "GET":
            response.headers["X-StateHub-Elapsed-Ms"] = f"{(time.perf_counter() - started) * 1000:.1f}"
            return response
        if "application/json" not in response.headers.get("content-type", ""):
            response.headers["X-StateHub-Elapsed-Ms"] = f"{(time.perf_counter() - started) * 1000:.1f}"
            return response

        body_parts = []
        async for chunk in response.body_iterator:
            body_parts.append(chunk)
        body = b"".join(body_parts)
        elapsed_ms = f"{(time.perf_counter() - started) * 1000:.1f}"

        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",
                    "X-StateHub-Elapsed-Ms": elapsed_ms,
                    "X-StateHub-Response-Bytes": "0",
                },
            )

        headers = {k: v for k, v in response.headers.items() if k.lower() != "content-length"}
        headers["ETag"] = etag
        headers["X-StateHub-Elapsed-Ms"] = elapsed_ms
        headers["X-StateHub-Response-Bytes"] = str(len(body))
        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,
)

_default_dashboard_origins = [
    *(f"http://localhost:{port}" for port in range(3000, 3006)),
    *(f"http://127.0.0.1:{port}" for port in range(3000, 3006)),
    *(f"http://[::1]:{port}" for port in range(3000, 3006)),
]
_cors_env = os.getenv("CORS_ORIGINS", ",".join(_default_dashboard_origins))
_cors_origins = [o.strip() for o in _cors_env.split(",") if o.strip()]

app.add_middleware(WriteIdempotencyMiddleware)
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", "Idempotency-Key", "X-StateHub-Source-Agent", "X-StateHub-Source-Host"],
    expose_headers=["ETag", "X-StateHub-Elapsed-Ms", "X-StateHub-Response-Bytes", "X-StateHub-Cache", "X-StateHub-Idempotency-Replay"],
)

app.include_router(domains.router)
app.include_router(recently_on_scope.hourly_router)
app.include_router(recently_on_scope.router)
app.include_router(consistency_sweep.router)
app.include_router(repos.router)
app.include_router(repository_renames.router)
app.include_router(repository_renames.operation_router)
app.include_router(work_record_projection.router)
app.include_router(topics.router)
app.include_router(workstreams.router)
app.include_router(workstreams.workplan_router)
app.include_router(workstream_dependencies.router)
app.include_router(workstream_dependencies.workplan_router)
app.include_router(tasks.router)
app.include_router(decisions.router)
app.include_router(intake.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)
# Before messages.router so /messages/feed is not captured by /messages/{id}.
app.include_router(message_feed.router)
app.include_router(messages.router)
app.include_router(memo_notifications.router)
app.include_router(capability_requests.router)
app.include_router(suggestions.router)
app.include_router(tpsc.router)
app.include_router(services.router)
app.include_router(token_events.router)
app.include_router(interface_changes.router)
app.include_router(flows.router)
app.include_router(reconciliation.router)
app.include_router(execution.router)
app.include_router(fabric.router)
app.include_router(legacy_meter.router)
app.include_router(review_contracts.router)
app.include_router(identifier_migrations.router)
app.include_router(state.router)
app.include_router(ops_runs.router)
app.include_router(policy.router)


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