Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/bedrock_agentcore/runtime/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this could be part of the if, elif statement above.

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)
Expand All @@ -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)
Expand Down
145 changes: 145 additions & 0 deletions tests/bedrock_agentcore/runtime/test_app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
Loading