Refactor Exceptions

lib/security_headers.py

@@ -6,26 +6,8 @@
 
 import os
 
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.requests import Request
-from starlette.responses import Response
-
 
 def _hsts_enabled() -> bool:
     """Return True if ENABLE_HSTS is set and truthy."""
     val = os.environ.get("ENABLE_HSTS", "").lower()
     return val in ("1", "true", "yes")
-
-
-class SecurityHeadersMiddleware(BaseHTTPMiddleware):
-    """Add security headers to all API responses."""
-
-    async def dispatch(self, request: Request, call_next) -> Response:
-        response = await call_next(request)
-        response.headers["X-Content-Type-Options"] = "nosniff"
-        response.headers["X-Frame-Options"] = "DENY"
-        if _hsts_enabled():
-            response.headers["Strict-Transport-Security"] = (
-                "max-age=31536000; includeSubDomains"
-            )
-        return response

simulation/api/errors.py

@@ -56,3 +56,15 @@ class ApiRunCreationFailedError(Exception):
     def __init__(self, message: str):
         super().__init__(message)
         self.message = message
+
+
+class ApiValidationError(ValueError):
+    """Raised for 422 Unprocessable Entity errors."""
+
+    pass
+
+
+class ApiInvalidInputError(ValueError):
+    """Raised for 400 Bad Request errors."""
+
+    pass

simulation/api/exception_handlers.py

@@ -0,0 +1,171 @@
+import logging
+
+from fastapi import Request
+from fastapi.encoders import jsonable_encoder
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
+
+from lib.env_utils import is_local_mode
+from simulation.api.dependencies.auth import UnauthorizedError
+from simulation.api.errors import (
+    ApiAgentFollowEdgeAlreadyExistsError,
+    ApiAgentFollowEdgeNotFoundError,
+    ApiAgentNotFoundError,
+    ApiHandleAlreadyExistsError,
+    ApiInvalidInputError,
+    ApiRunCreationFailedError,
+    ApiRunNotFoundError,
+    ApiSelfFollowNotAllowedError,
+    ApiTargetAgentNotFoundError,
+    ApiValidationError,
+)
+from simulation.api.routes._helpers import error_response
+
+logger = logging.getLogger(__name__)
+
+
+async def global_exception_handler(request: Request, exc: Exception) -> JSONResponse:
+    logger.exception(
+        "500 Internal Server Error: %s %s", request.method, request.url.path
+    )
+    return error_response(
+        status_code=500,
+        code="INTERNAL_ERROR",
+        message="Internal server error.",
+        detail=str(exc) if is_local_mode() else None,
+    )
+
+
+def validation_exception_handler(
+    _request: Request, exc: RequestValidationError
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=HTTP_422_UNPROCESSABLE_ENTITY,
+        content={
+            "error": {
+                "code": "VALIDATION_ERROR",
+                "message": "Request validation failed",
+                "detail": jsonable_encoder(exc.errors()),
+            }
+        },
+    )
+
+
+def unauthorized_handler(_request: Request, exc: UnauthorizedError) -> JSONResponse:
+    return error_response(
+        status_code=401,
+        code="UNAUTHORIZED",
+        message=exc.message,
+    )
+
+
+def run_not_found_handler(_request: Request, exc: ApiRunNotFoundError) -> JSONResponse:
+    return error_response(
+        status_code=404,
+        code="RUN_NOT_FOUND",
+        message="Run not found",
+        detail=exc.run_id,
+    )
+
+
+def run_creation_failed_handler(
+    _request: Request, exc: ApiRunCreationFailedError
+) -> JSONResponse:
+    return error_response(
+        status_code=500,
+        code="RUN_CREATION_FAILED",
+        message=exc.message,
+    )
+
+
+def handle_already_exists_handler(
+    _request: Request, exc: ApiHandleAlreadyExistsError
+) -> JSONResponse:
+    return error_response(
+        status_code=409,
+        code="HANDLE_ALREADY_EXISTS",
+        message="Agent with this handle already exists",
+        detail=exc.handle,
+    )
+
+
+def agent_not_found_handler(
+    _request: Request, exc: ApiAgentNotFoundError
+) -> JSONResponse:
+    return error_response(
+        status_code=404,
+        code="AGENT_NOT_FOUND",
+        message="Agent not found",
+        detail=exc.handle,
+    )
+
+
+def target_agent_not_found_handler(
+    _request: Request, exc: ApiTargetAgentNotFoundError
+) -> JSONResponse:
+    return error_response(
+        status_code=404,
+        code="TARGET_AGENT_NOT_FOUND",
+        message="Target agent not found",
+        detail=exc.handle,
+    )
+
+
+def follow_edge_already_exists_handler(
+    _request: Request, exc: ApiAgentFollowEdgeAlreadyExistsError
+) -> JSONResponse:
+    return error_response(
+        status_code=409,
+        code="FOLLOW_EDGE_ALREADY_EXISTS",
+        message="Follow edge already exists",
+        detail=f"{exc.follower_handle}->{exc.target_handle}",
+    )
+
+
+def follow_edge_not_found_handler(
+    _request: Request, exc: ApiAgentFollowEdgeNotFoundError
+) -> JSONResponse:
+    return error_response(
+        status_code=404,
+        code="FOLLOW_EDGE_NOT_FOUND",
+        message="Follow edge not found",
+        detail=f"{exc.follower_handle}->{exc.target_handle}",
+    )
+
+
+def self_follow_not_allowed_handler(
+    _request: Request, exc: ApiSelfFollowNotAllowedError
+) -> JSONResponse:
+    return error_response(
+        status_code=422,
+        code="SELF_FOLLOW_NOT_ALLOWED",
+        message="Agent cannot follow itself",
+        detail=exc.handle,
+    )
+
+
+def value_error_handler(_request: Request, exc: ValueError) -> JSONResponse:
+    if isinstance(exc, ApiInvalidInputError):
+        return error_response(status_code=400, code="INVALID_INPUT", message=str(exc))
+    if isinstance(exc, ApiValidationError):
+        return error_response(
+            status_code=422, code="VALIDATION_ERROR", message=str(exc)
+        )
+    return error_response(status_code=422, code="VALIDATION_ERROR", message=str(exc))
+
+
+EXCEPTION_HANDLERS: dict[type[Exception], object] = {
+    Exception: global_exception_handler,
+    ValueError: value_error_handler,
+    RequestValidationError: validation_exception_handler,
+    UnauthorizedError: unauthorized_handler,
+    ApiRunNotFoundError: run_not_found_handler,
+    ApiRunCreationFailedError: run_creation_failed_handler,
+    ApiHandleAlreadyExistsError: handle_already_exists_handler,
+    ApiAgentNotFoundError: agent_not_found_handler,
+    ApiTargetAgentNotFoundError: target_agent_not_found_handler,
+    ApiAgentFollowEdgeAlreadyExistsError: follow_edge_already_exists_handler,
+    ApiAgentFollowEdgeNotFoundError: follow_edge_not_found_handler,
+    ApiSelfFollowNotAllowedError: self_follow_not_allowed_handler,
+}

simulation/api/main.py

@@ -7,28 +7,19 @@
 import asyncio
 import logging
 import os
-import uuid
 from contextlib import asynccontextmanager
 
-from fastapi import FastAPI, Request
-from fastapi.encoders import jsonable_encoder
-from fastapi.exceptions import RequestValidationError
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
 from slowapi.errors import RateLimitExceeded
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
 
 from db.adapters.sqlite.sqlite import get_db_path, initialize_database
 from lib.env_utils import is_local_mode, parse_bool_env
 from lib.rate_limiting import limiter, rate_limit_exceeded_handler
-from lib.request_logging import log_request_start
-from lib.security_headers import SecurityHeadersMiddleware
 from simulation.api.context import build_app_context
-from simulation.api.dependencies.auth import (
-    UnauthorizedError,
-    disallow_auth_bypass_in_production,
-)
+from simulation.api.dependencies.auth import disallow_auth_bypass_in_production
+from simulation.api.exception_handlers import EXCEPTION_HANDLERS
+from simulation.api.middleware import RequestIdMiddleware, SecurityHeadersMiddleware
 from simulation.api.routes.simulation import router as simulation_router
 from simulation.local_dev.local_mode import disallow_local_mode_in_production
 from simulation.local_dev.seed_loader import seed_database_from_fixtures_if_needed
@@ -89,40 +80,9 @@ async def lifespan(app: FastAPI):
 app.state.limiter = limiter
 app.add_exception_handler(RateLimitExceeded, rate_limit_exceeded_handler)  # type: ignore[reportArgumentType]
 
+for exc_class, handler_func in EXCEPTION_HANDLERS.items():
+    app.add_exception_handler(exc_class, handler_func)  # type: ignore[reportArgumentType]
 
-def _unauthorized_handler(request: Request, exc: UnauthorizedError) -> JSONResponse:
-    """Return 401 with standard error shape for auth failures."""
-    return JSONResponse(
-        status_code=401,
-        content={
-            "error": {
-                "code": "UNAUTHORIZED",
-                "message": exc.message,
-                "detail": None,
-            }
-        },
-    )
-
-
-app.add_exception_handler(UnauthorizedError, _unauthorized_handler)  # type: ignore[reportArgumentType]
-
-
-class RequestIdMiddleware(BaseHTTPMiddleware):
-    """Assigns request_id and logs request start in structured format."""
-
-    async def dispatch(self, request: Request, call_next):
-        request_id = request.headers.get("X-Request-ID") or uuid.uuid4().hex
-        request.state.request_id = request_id
-        log_request_start(
-            request_id=request_id,
-            method=request.method,
-            path=request.url.path,
-        )
-        return await call_next(request)
-
-
-app.add_middleware(SecurityHeadersMiddleware)
-app.add_middleware(RequestIdMiddleware)
 _allowed_origins_raw: str = os.environ.get("ALLOWED_ORIGINS", DEFAULT_ALLOWED_ORIGINS)
 _allowed_origins: list[str] = [
     origin.strip() for origin in _allowed_origins_raw.split(",") if origin.strip()
@@ -133,26 +93,11 @@ async def dispatch(self, request: Request, call_next):
     allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
     allow_headers=["*"],
 )
+app.add_middleware(SecurityHeadersMiddleware)
+app.add_middleware(RequestIdMiddleware)
 app.include_router(simulation_router, prefix="/v1")
 
 
-@app.exception_handler(RequestValidationError)
-async def validation_exception_handler(
-    request: Request, exc: RequestValidationError
-) -> JSONResponse:
-    """Return 422 with stable error shape matching other API errors."""
-    return JSONResponse(
-        status_code=HTTP_422_UNPROCESSABLE_ENTITY,
-        content={
-            "error": {
-                "code": "VALIDATION_ERROR",
-                "message": "Request validation failed",
-                "detail": jsonable_encoder(exc.errors()),
-            }
-        },
-    )
-
-
 @app.get("/health")
 def health():
     """Health check endpoint. Returns 200 when the service is up."""

simulation/api/middleware.py

@@ -0,0 +1,62 @@
+import uuid
+
+from fastapi import Request
+from starlette.datastructures import MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from lib.request_logging import log_request_start
+from lib.security_headers import _hsts_enabled
+
+
+class SecurityHeadersMiddleware:
+    """Pure ASGI middleware that adds security headers to every HTTP response."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        async def send_with_headers(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableHeaders(scope=message)
+                headers["X-Content-Type-Options"] = "nosniff"
+                headers["X-Frame-Options"] = "DENY"
+                if _hsts_enabled():
+                    headers["Strict-Transport-Security"] = (
+                        "max-age=31536000; includeSubDomains"
+                    )
+            await send(message)
+
+        await self.app(scope, receive, send_with_headers)
+
+
+class RequestIdMiddleware:
+    """Pure ASGI middleware that assigns a request ID and logs request start."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        request = Request(scope)
+        request_id = request.headers.get("X-Request-ID") or uuid.uuid4().hex
+        request.state.request_id = request_id
+        log_request_start(
+            request_id=request_id,
+            method=request.method,
+            path=request.url.path,
+        )
+
+        async def send_with_request_id(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableHeaders(scope=message)
+                headers["X-Request-ID"] = request_id
+            await send(message)
+
+        await self.app(scope, receive, send_with_request_id)