fix: allow custom HTTP status codes from entrypoint handlers (#284) (#296)

* fix: allow custom HTTP status codes from entrypoint handlers (#284)

* test: add tests for custom HTTP status codes from entrypoint handlers (#284)

Author: aidandaly24

src/bedrock_agentcore/runtime/app.py

@@ -18,6 +18,7 @@
 
 from starlette.applications import Starlette
 from starlette.concurrency import run_in_threadpool
+from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route, WebSocketRoute
@@ -418,6 +419,18 @@ async def _handle_invocation(self, request):
                 self.logger.info("Returning streaming response (async generator) (%.3fs)", duration)
                 return StreamingResponse(self._stream_with_error_handling(result), media_type="text/event-stream")
 
+            # If handler returned a Starlette Response directly, pass it through.
+            # This lets handlers control status codes (e.g. JSONResponse(data, status_code=404)).
+            if isinstance(result, Response):
+                status = getattr(result, "status_code", 200)
+                # Log at warning level for error responses so operators can distinguish
+                # intentional error responses from successful invocations in logs.
+                if status >= 400:
+                    self.logger.warning("Invocation returned HTTP %d (%.3fs)", status, duration)
+                else:
+                    self.logger.info("Invocation completed successfully (%.3fs)", duration)
+                return result
+
             self.logger.info("Invocation completed successfully (%.3fs)", duration)
             # Use safe serialization for consistency with streaming paths
             safe_json_string = self._safe_serialize_to_json_string(result)
@@ -427,6 +440,16 @@ async def _handle_invocation(self, request):
             duration = time.time() - start_time
             self.logger.warning("Invalid JSON in request (%.3fs): %s", duration, e)
             return JSONResponse({"error": "Invalid JSON", "details": str(e)}, status_code=400)
+        except HTTPException as e:
+            duration = time.time() - start_time
+            # Use error level for 5xx to match the generic Exception handler's severity,
+            # since server errors warrant the same urgency regardless of how they're raised.
+            # Use warning for 4xx since those are intentional client-error responses.
+            if e.status_code >= 500:
+                self.logger.error("HTTP %d (%.3fs): %s", e.status_code, duration, e.detail)
+            else:
+                self.logger.warning("HTTP %d (%.3fs): %s", e.status_code, duration, e.detail)
+            return JSONResponse({"error": e.detail}, status_code=e.status_code)
         except Exception as e:
             duration = time.time() - start_time
             self.logger.exception("Invocation failed (%.3fs)", duration)

tests/bedrock_agentcore/runtime/test_app.py

@@ -421,6 +421,151 @@ def handler(payload):
         assert response.headers.get("x-custom-mw") == "mw"
 
 
+class TestCustomStatusCodes:
+    """Test that entrypoint handlers can return custom HTTP status codes (#284)."""
+
+    def test_http_exception_returns_custom_status(self):
+        """Test raising HTTPException returns its status code instead of 500."""
+        from starlette.exceptions import HTTPException
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            raise HTTPException(status_code=400, detail="Prompt missing")
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 400
+        assert response.json() == {"error": "Prompt missing"}
+
+    def test_http_exception_422(self):
+        """Test HTTPException with 422 status code."""
+        from starlette.exceptions import HTTPException
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            raise HTTPException(status_code=422, detail="Validation failed")
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 422
+        assert response.json() == {"error": "Validation failed"}
+
+    def test_http_exception_500_still_returns_500(self):
+        """Test HTTPException with 500 is not swallowed."""
+        from starlette.exceptions import HTTPException
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            raise HTTPException(status_code=500, detail="Intentional server error")
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 500
+        assert response.json() == {"error": "Intentional server error"}
+
+    def test_return_response_passthrough(self):
+        """Test returning a Response object passes it through without wrapping."""
+        from starlette.responses import JSONResponse
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            return JSONResponse({"error": "not found"}, status_code=404)
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 404
+        assert response.json() == {"error": "not found"}
+
+    def test_return_response_200_passthrough(self):
+        """Test returning a Response with 200 also passes through."""
+        from starlette.responses import JSONResponse
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            return JSONResponse({"custom": True}, status_code=200, headers={"x-custom": "yes"})
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 200
+        assert response.json() == {"custom": True}
+        assert response.headers.get("x-custom") == "yes"
+
+    def test_normal_dict_return_still_200(self):
+        """Test that normal dict returns are unaffected."""
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            return {"message": "ok"}
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 200
+        assert response.json() == {"message": "ok"}
+
+    def test_generic_exception_still_500(self):
+        """Test that non-HTTP exceptions still return 500."""
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        def handler(payload):
+            raise ValueError("something broke")
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 500
+        assert response.json() == {"error": "something broke"}
+
+    def test_async_handler_http_exception(self):
+        """Test HTTPException works from async handlers too."""
+        from starlette.exceptions import HTTPException
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        async def handler(payload):
+            raise HTTPException(status_code=400, detail="Bad request")
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 400
+        assert response.json() == {"error": "Bad request"}
+
+    def test_async_handler_response_passthrough(self):
+        """Test returning Response from async handler."""
+        from starlette.responses import JSONResponse
+
+        app = BedrockAgentCoreApp()
+
+        @app.entrypoint
+        async def handler(payload):
+            return JSONResponse({"error": "gone"}, status_code=410)
+
+        client = TestClient(app)
+        response = client.post("/invocations", json={})
+
+        assert response.status_code == 410
+        assert response.json() == {"error": "gone"}
+
+
 class TestConcurrentInvocations:
     """Test concurrent invocation handling simplified without limits."""