diff --git a/tracecat/api/app.py b/tracecat/api/app.py index d9d94656..c0c09f8d 100644 --- a/tracecat/api/app.py +++ b/tracecat/api/app.py @@ -5,10 +5,10 @@ import polars as pl from aio_pika import Channel from aio_pika.pool import Pool -from fastapi import Depends, FastAPI, HTTPException, status +from fastapi import Depends, FastAPI, HTTPException, Request, status from fastapi.middleware.cors import CORSMiddleware from fastapi.params import Body -from fastapi.responses import StreamingResponse +from fastapi.responses import ORJSONResponse, StreamingResponse from sqlalchemy import Engine, or_ from sqlalchemy.exc import NoResultFound from sqlmodel import Session, select @@ -132,6 +132,15 @@ async def lifespan(app: FastAPI): ) +# Catch-all exception handler to prevent stack traces from leaking +@app.exception_handler(Exception) +async def custom_exception_handler(request: Request, exc: Exception): + return ORJSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={"message": "An unexpected error occurred. Please try again later."}, + ) + + @app.get("/") def root() -> dict[str, str]: return {"message": "Hello world. I am the API."} diff --git a/tracecat/runner/app.py b/tracecat/runner/app.py index 0df7bd80..d1089880 100644 --- a/tracecat/runner/app.py +++ b/tracecat/runner/app.py @@ -173,6 +173,15 @@ async def valid_workflow(workflow_id: str) -> str: return workflow_id +# Catch-all exception handler to prevent stack traces from leaking +@app.exception_handler(Exception) +async def custom_exception_handler(request: Request, exc: Exception): + return ORJSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={"message": "An unexpected error occurred. Please try again later."}, + ) + + # Endpoints @app.get("/") def root() -> dict[str, str]: