add SkyvernHTTPException base exception to automate known error handling in agent service

skyvern/exceptions.py

@@ -1,9 +1,18 @@
+from fastapi import status
+
+
 class SkyvernException(Exception):
     def __init__(self, message: str | None = None):
         self.message = message
         super().__init__(message)
 
 
+class SkyvernHTTPException(SkyvernException):
+    def __init__(self, message: str | None = None, status_code: int = status.HTTP_400_BAD_REQUEST):
+        self.status_code = status_code
+        super().__init__(message)
+
+
 class InvalidOpenAIResponseFormat(SkyvernException):
     def __init__(self, message: str | None = None):
         super().__init__(f"Invalid response format: {message}")
@@ -22,9 +31,9 @@ def __init__(self, proxy_location: str | None = None):
         super().__init__(f"Unknown proxy location: {proxy_location}")
 
 
-class TaskNotFound(SkyvernException):
+class TaskNotFound(SkyvernHTTPException):
     def __init__(self, task_id: str | None = None):
-        super().__init__(f"Task {task_id} not found")
+        super().__init__(f"Task {task_id} not found", status_code=status.HTTP_404_NOT_FOUND)
 
 
 class ScriptNotFound(SkyvernException):
@@ -97,9 +106,9 @@ def __init__(self, block_type: str) -> None:
         super().__init__(f"Unknown block type {block_type}")
 
 
-class WorkflowNotFound(SkyvernException):
+class WorkflowNotFound(SkyvernHTTPException):
     def __init__(self, workflow_id: str) -> None:
-        super().__init__(f"Workflow {workflow_id} not found")
+        super().__init__(f"Workflow {workflow_id} not found", status_code=status.HTTP_404_NOT_FOUND)
 
 
 class WorkflowRunNotFound(SkyvernException):
@@ -119,9 +128,9 @@ def __init__(self, parameter_key: str, workflow_id: str, workflow_run_id: str) -
         )
 
 
-class WorkflowParameterNotFound(SkyvernException):
+class WorkflowParameterNotFound(SkyvernHTTPException):
     def __init__(self, workflow_parameter_id: str) -> None:
-        super().__init__(f"Workflow parameter {workflow_parameter_id} not found")
+        super().__init__(f"Workflow parameter {workflow_parameter_id} not found", status_code=status.HTTP_404_NOT_FOUND)
 
 
 class FailedToNavigateToUrl(SkyvernException):
@@ -163,14 +172,17 @@ def __init__(self) -> None:
         super().__init__("BrowserState is missing the main page")
 
 
-class OrganizationNotFound(SkyvernException):
+class OrganizationNotFound(SkyvernHTTPException):
     def __init__(self, organization_id: str) -> None:
-        super().__init__(f"Organization {organization_id} not found")
+        super().__init__(f"Organization {organization_id} not found", status_code=status.HTTP_404_NOT_FOUND)
 
 
-class StepNotFound(SkyvernException):
+class StepNotFound(SkyvernHTTPException):
     def __init__(self, organization_id: str, task_id: str, step_id: str | None = None) -> None:
-        super().__init__(f"Step {step_id or 'latest'} not found. organization_id={organization_id} task_id={task_id}")
+        super().__init__(
+            f"Step {step_id or 'latest'} not found. organization_id={organization_id} task_id={task_id}",
+            status_code=status.HTTP_404_NOT_FOUND,
+        )
 
 
 class FailedToTakeScreenshot(SkyvernException):

skyvern/forge/api_app.py

@@ -10,6 +10,7 @@
 from starlette_context.middleware import RawContextMiddleware
 from starlette_context.plugins.base import Plugin
 
+from skyvern.exceptions import SkyvernHTTPException
 from skyvern.forge import app as forge_app
 from skyvern.forge.sdk.core import skyvern_context
 from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
@@ -73,6 +74,10 @@ def start_scheduler() -> None:
 
         LOG.info("Server startup complete. Skyvern is now online")
 
+    @app.exception_handler(SkyvernHTTPException)
+    async def handle_skyvern_http_exception(request: Request, exc: SkyvernHTTPException) -> JSONResponse:
+        return JSONResponse(status_code=exc.status_code, content={"detail": exc.message})
+
     @app.exception_handler(Exception)
     async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse:
         LOG.exception("Unexpected error in agent server.", exc_info=exc)