refactor GET /resources endpoint

Author: skmkqw

.DS_Store

@@ -13,7 +13,8 @@
 
 from chain.processor import ChainProcessor
 from orchestration import ContainerManager, DataStore, Guardian, Orchestrator
-from server.utils import is_local_ip
+from server.exceptions import BadRequestError, ForbiddenError, RESTServerError, RitualNodeError
+from server.utils import handle_errors, is_local_ip
 from shared import AsyncTask, JobResult
 from shared.config import ConfigChain, ConfigServer
 from shared.message import (
@@ -159,11 +160,12 @@ async def info() -> Tuple[Response, int]:
                 ),
                 200,
             )
-
+        
         @self._app.route("/resources", methods=["GET"])
         @rate_limit(
             self._rate_limit.num_requests, timedelta(seconds=self._rate_limit.period)
         )
+        @handle_errors
         async def resources() -> Tuple[Response, int]:
             """Collects container resources
 
@@ -175,10 +177,14 @@ async def resources() -> Tuple[Response, int]:
             """
             model_id = request.args.get("model_id")
 
-            return (
-                jsonify(await self._orchestrator.collect_service_resources(model_id)),
-                200,
-            )
+            try:
+                resources = await self._orchestrator.collect_service_resources(model_id)
+                return jsonify(resources), 200
+            except Exception as e:
+                raise RESTServerError(
+                    "Failed to collect resources", 
+                    {"model_id": model_id, "error": str(e)}
+                )
 
         def filter_create_job(func):  # type: ignore
             """Decorator to filter and preprocess incoming off-chain messages"""
@@ -198,20 +204,24 @@ async def wrapper() -> Any:
                     # Get the IP address of the client
                     client_ip = request.remote_addr
                     if not client_ip:
-                        return (
-                            jsonify({"error": "Could not get client IP address"}),
-                            400,
-                        )
+                        raise BadRequestError("Could not get client IP address")
 
                     # Parse message data, inject uuid and client IP
                     job_id = str(uuid4())  # Generate a unique job ID
                     log.debug(
                         "Received new off-chain raw message", msg=data, job_id=job_id
                     )
-                    parsed: OffchainMessage = from_union(
-                        OffchainMessage,
-                        {"id": job_id, "ip": client_ip, **data},
-                    )
+                    
+                    try:
+                        parsed: OffchainMessage = from_union(
+                            OffchainMessage,
+                            {"id": job_id, "ip": client_ip, **data},
+                        )
+                    except Exception as e:
+                        raise BadRequestError(
+                            "Failed to parse message", 
+                            {"error": str(e)}
+                        )
 
                     # Filter message through guardian
                     filtered = self._guardian.process_message(parsed)
@@ -225,19 +235,20 @@ async def wrapper() -> Any:
                             err=filtered.error,
                             **filtered.params,
                         )
-                        return (
-                            jsonify(
-                                {"error": filtered.error, "params": filtered.params}
-                            ),
-                            405,
+                        raise ForbiddenError(
+                            filtered.error,
+                            filtered.params
                         )
 
                     # Call actual endpoint function
                     return await func(message=filtered)
 
+                except RitualNodeError:
+                    # Pass through our custom exceptions
+                    raise
                 except Exception as e:
-                    log.error(f"Error in endpoint preprocessing: {e}")
-                    return jsonify({"error": f"Internal server error: {str(e)}"}), 500
+                    log.error(f"Error in endpoint preprocessing: {e}", exc_info=True)
+                    raise RESTServerError(f"Internal server error", {"message": str(e)})
 
             return wrapper
 
@@ -246,6 +257,7 @@ async def wrapper() -> Any:
         @rate_limit(
             self._rate_limit.num_requests, timedelta(seconds=self._rate_limit.period)
         )
+        @handle_errors
         async def create_job(message: OffchainMessage) -> Tuple[Response, int]:
             """Creates new off-chain job (direct compute request or subscription)
 
@@ -277,7 +289,7 @@ async def create_job(message: OffchainMessage) -> Tuple[Response, int]:
                     # Should only reach this point if chain is enabled (else, filtered
                     # out upstream)
                     if self._processor is None:
-                        raise RuntimeError("Chain not enabled")
+                        raise RESTServerError("Chain not enabled")
 
                     # Submit delegated subscription message to processor
                     create_task(self._processor.track(message))
@@ -286,6 +298,8 @@ async def create_job(message: OffchainMessage) -> Tuple[Response, int]:
                     # fetched via REST so it would be misleading. They are tracked
                     # on-chain instead
                     return_obj = {}
+                else:
+                    raise BadRequestError("Unsupported message type", {"type": message.type})
 
                 # Return created message ID
                 log.debug(
@@ -297,16 +311,14 @@ async def create_job(message: OffchainMessage) -> Tuple[Response, int]:
                     id=str(message.id),
                 )
                 return jsonify(return_obj), 200
+            except RitualNodeError:
+                # Pass through our custom exceptions
+                raise
             except Exception as e:
-                # Return error
-                log.error(
-                    "Processed REST response",
-                    endpoint=request.path,
-                    method=request.method,
-                    status=500,
-                    err=str(e),
+                raise RESTServerError(
+                    "Could not enqueue job", 
+                    {"error": str(e)}
                 )
-                return jsonify({"error": f"Could not enqueue job: {str(e)}"}), 500
 
         @self._app.route("/api/jobs/stream", methods=["POST"])
         @filter_create_job  # type: ignore
@@ -357,6 +369,7 @@ async def generator() -> Any:
         @rate_limit(
             self._rate_limit.num_requests, timedelta(seconds=self._rate_limit.period)
         )
+        @handle_errors
         async def create_job_batch() -> Tuple[Response, int]:
             """Creates off-chain jobs in batch (direct compute requests / subscriptions)
 
@@ -365,27 +378,36 @@ async def create_job_batch() -> Tuple[Response, int]:
             """
             try:
                 # Collect JSON body
-                data = await request.get_json(force=True)
+                try:
+                    data = await request.get_json(force=True)
+                except Exception:
+                    raise BadRequestError("Invalid JSON body")
 
                 # Get the IP address of the client
                 client_ip = request.remote_addr
                 if not client_ip:
-                    return jsonify({"error": "Could not get client IP address"}), 400
+                    raise BadRequestError("Could not get client IP address")
 
                 log.debug("Received new off-chain raw message batch", msg=data)
 
                 # If data is not an array, return error
                 if not isinstance(data, list):
-                    return jsonify({"error": "Expected a list"}), 400
+                    raise BadRequestError("Expected a list of job requests")
 
                 # Inject uuid and client IP to each message
-                parsed: list[OffchainMessage] = [
-                    from_union(
-                        OffchainMessage,
-                        {"id": str(uuid4()), "ip": client_ip, **item},
+                try:
+                    parsed: list[OffchainMessage] = [
+                        from_union(
+                            OffchainMessage,
+                            {"id": str(uuid4()), "ip": client_ip, **item},
+                        )
+                        for item in data
+                    ]
+                except Exception as e:
+                    raise BadRequestError(
+                        "Failed to parse message batch", 
+                        {"error": str(e)}
                     )
-                    for item in data
-                ]
 
                 # Filter messages through guardian
                 filtered = cast(
@@ -419,7 +441,8 @@ async def create_job_batch() -> Tuple[Response, int]:
                             # Should only reach this point if chain is enabled (else,
                             # filtered out upstream)
                             if self._processor is None:
-                                raise RuntimeError("Chain not enabled")
+                                results.append({"error": "Chain not enabled"})
+                                continue
 
                             # Submit filtered delegated subscription message to processor
                             create_task(
@@ -429,7 +452,9 @@ async def create_job_batch() -> Tuple[Response, int]:
                             )
                             results.append({})
                         else:
-                            results.append({"error": "Could not parse message"})
+                            results.append({"error": "Unsupported message type"})
+                    else:
+                        results.append({"error": "Invalid message format"})
 
                 # Return created message IDs or errors
                 log.debug(
@@ -440,26 +465,31 @@ async def create_job_batch() -> Tuple[Response, int]:
                     results=results,
                 )
                 return jsonify(results), 200
+                
+            except RitualNodeError:
+                # Pass through our custom exceptions
+                raise
             except Exception as e:
-                # Return error
-                log.error(
-                    "Processed REST response",
-                    endpoint=request.path,
-                    method=request.method,
-                    status=500,
-                    err=str(e),
+                raise RESTServerError(
+                    "Failed to process batch request", 
+                    {"error": str(e)}
                 )
-                return jsonify({"error": f"Could not enqueue job:  {str(e)}"}), 500
 
         @self._app.route("/api/jobs", methods=["GET"])
         @rate_limit(
             self._rate_limit.num_requests, timedelta(seconds=self._rate_limit.period)
         )
+        @handle_errors
         async def get_job() -> Tuple[Response, int]:
+            """Get job status by ID or list all jobs for a client.
+            
+            Returns:
+                Response: Job data or list of job IDs
+            """
             # Get the IP address of the client
             client_ip = request.remote_addr
             if not client_ip:
-                return jsonify({"error": "Could not get client IP address"}), 400
+                raise BadRequestError("Could not get client IP address")
 
             # Get job ID from query
             job_ids = request.args.getlist("id")
@@ -485,6 +515,7 @@ async def get_job() -> Tuple[Response, int]:
                 return jsonify(data), 200
 
         @self._app.route("/api/status", methods=["PUT"])
+        @handle_errors
         async def store_job_status() -> Tuple[Response, int]:
             """Stores job status in data store"""
 
@@ -494,56 +525,55 @@ async def store_job_status() -> Tuple[Response, int]:
                     "Unauthorized attempt to store job status",
                     remote_addr=request.remote_addr,
                 )
-                return jsonify({"error": "Unauthorized"}), 403
+                raise ForbiddenError("Unauthorized access", 
+                                    {"remote_addr": str(request.remote_addr)})
 
-            try:
-                # Collect JSON body
-                data = await request.get_json(force=True)
+            # Collect JSON body
+            data = await request.get_json(force=True)
 
-                # Get the IP address of the client
-                client_ip = request.remote_addr
-                if not client_ip:
-                    return jsonify({"error": "Could not get client IP address"}), 400
-
-                log.debug("Received new result", result=data)
-
-                # Create off-chain message with client IP
-                parsed: OffchainMessage = from_union(
-                    OffchainMessage,
-                    {
-                        "id": data["id"],
-                        "ip": client_ip,
-                        "containers": data["containers"],
-                        "data": {},
-                    },
-                )
+            # Get the IP address of the client
+            client_ip = request.remote_addr
+            if not client_ip:
+                raise BadRequestError("Could not get client IP address")
+
+            log.debug("Received new result", result=data)
+
+            # Validate required fields
+            if "id" not in data:
+                raise BadRequestError("Missing required field: id")
+            if "status" not in data:
+                raise BadRequestError("Missing required field: status")
+            if "containers" not in data:
+                raise BadRequestError("Missing required field: containers")
+
+            # Create off-chain message with client IP
+            parsed: OffchainMessage = from_union(
+                OffchainMessage,
+                {
+                    "id": data["id"],
+                    "ip": client_ip,
+                    "containers": data["containers"],
+                    "data": {},
+                },
+            )
 
-                # Store job status
-                match data["status"]:
-                    case "success":
-                        self._store.set_success(parsed, [])
-                        for container in data["containers"]:
-                            self._store.track_container_status(container, "success")
-                    case "failed":
-                        self._store.set_failed(parsed, [])
-                        for container in data["containers"]:
-                            self._store.track_container_status(container, "failed")
-                    case "running":
-                        self._store.set_running(parsed)
-                    case _:
-                        return jsonify({"error": "Status is invalid"}), 400
-
-                return jsonify(), 200
-            except Exception as e:
-                # Return error
-                log.error(
-                    "Processed REST response",
-                    endpoint=request.path,
-                    method=request.method,
-                    status=500,
-                    err=e,
-                )
-                return jsonify({"error": "Could not store job status"}), 500
+            # Store job status
+            match data["status"]:
+                case "success":
+                    self._store.set_success(parsed, [])
+                    for container in data["containers"]:
+                        self._store.track_container_status(container, "success")
+                case "failed":
+                    self._store.set_failed(parsed, [])
+                    for container in data["containers"]:
+                        self._store.track_container_status(container, "failed")
+                case "running":
+                    self._store.set_running(parsed)
+                case _:
+                    raise BadRequestError("Invalid status value", 
+                                        {"status": data["status"]})
+
+            return jsonify(), 200
 
     async def run_forever(self: RESTServer) -> None:
         """Main RESTServer lifecycle loop. Uses production hypercorn server"""