feat: add WebSocket support with @app.websocket decorator (#41)

* feat: add WebSocket support with @app.websocket decorator

Add bidirectional streaming support for agent invocations:
- Add @app.websocket decorator for registering WebSocket handlers
- Register /ws endpoint in BedrockAgentCoreApp constructor
- Implement _handle_websocket method with proper error handling
- Follow existing decorator patterns (entrypoint, ping) for consistency
- Support WebSocketDisconnect for graceful connection handling
- Integrate with existing RequestContext for session management

test: add unit and integration tests for WebSocket support

Add comprehensive test coverage for WebSocket decorator:

Unit tests (10 tests):
- WebSocket route initialization and registration
- Decorator functionality and handler storage
- Basic send/receive communication
- Context integration with session IDs
- Exception handling and error cases
- Multiple message handling
- Graceful disconnect handling
- Custom request headers via context
- Streaming data functionality

Integration test:
- End-to-end WebSocket echo server
- Streaming multiple messages
- Session ID propagation through headers
- Real WebSocket client connection testing

All tests pass successfully.

* fix: remove exception details from WebSocket close reason

Don't send exception messages to clients when closing WebSocket connections
with code 1011. This prevents leaking internal error details to clients.

Changes:
- Remove reason parameter from websocket.close() calls
- Exception details are still logged server-side
- Tests continue to pass as they only check for disconnect exceptions

---------

Co-authored-by: Abishek Kumar <[email protected]>

Author: 2 people

.github/workflows/integration-testing.yml

@@ -130,7 +130,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -e .
-          pip install --no-cache-dir pytest requests strands-agents uvicorn httpx starlette
+          pip install --no-cache-dir pytest requests strands-agents uvicorn httpx starlette websockets
 
       - name: Run integration tests
         env:

pyproject.toml

@@ -143,6 +143,7 @@ dev = [
     "pytest-asyncio>=0.24.0",
     "pytest-cov>=6.0.0",
     "ruff>=0.12.0",
+    "websockets>=14.1",
     "wheel>=0.45.1",
     "strands-agents>=1.18.0",
 ]

src/bedrock_agentcore/runtime/app.py

@@ -17,8 +17,9 @@
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.responses import JSONResponse, Response, StreamingResponse
-from starlette.routing import Route
+from starlette.routing import Route, WebSocketRoute
 from starlette.types import Lifespan
+from starlette.websockets import WebSocket, WebSocketDisconnect
 
 from .context import BedrockAgentCoreContext, RequestContext
 from .models import (
@@ -90,6 +91,7 @@ def __init__(
         """
         self.handlers: Dict[str, Callable] = {}
         self._ping_handler: Optional[Callable] = None
+        self._websocket_handler: Optional[Callable] = None
         self._active_tasks: Dict[int, Dict[str, Any]] = {}
         self._task_counter_lock: threading.Lock = threading.Lock()
         self._forced_ping_status: Optional[PingStatus] = None
@@ -98,6 +100,7 @@ def __init__(
         routes = [
             Route("/invocations", self._handle_invocation, methods=["POST"]),
             Route("/ping", self._handle_ping, methods=["GET"]),
+            WebSocketRoute("/ws", self._handle_websocket),
         ]
         super().__init__(routes=routes, lifespan=lifespan, middleware=middleware)
         self.debug = debug  # Set after super().__init__ to avoid override
@@ -135,6 +138,24 @@ def ping(self, func: Callable) -> Callable:
         self._ping_handler = func
         return func
 
+    def websocket(self, func: Callable) -> Callable:
+        """Decorator to register a WebSocket handler at /ws endpoint.
+
+        Args:
+            func: The function to register as WebSocket handler
+
+        Returns:
+            The decorated function
+
+        Example:
+            @app.websocket
+            async def handler(websocket, context):
+                await websocket.accept()
+                # ... handle messages ...
+        """
+        self._websocket_handler = func
+        return func
+
     def async_task(self, func: Callable) -> Callable:
         """Decorator to track async tasks for ping status.
 
@@ -390,6 +411,29 @@ def _handle_ping(self, request):
             self.logger.exception("Ping endpoint failed")
             return JSONResponse({"status": PingStatus.HEALTHY.value, "time_of_last_update": int(time.time())})
 
+    async def _handle_websocket(self, websocket: WebSocket):
+        """Handle WebSocket connections."""
+        request_context = self._build_request_context(websocket)
+
+        try:
+            handler = self._websocket_handler
+            if not handler:
+                self.logger.error("No WebSocket handler defined")
+                await websocket.close(code=1011)
+                return
+
+            self.logger.debug("WebSocket connection established")
+            await handler(websocket, request_context)
+
+        except WebSocketDisconnect:
+            self.logger.debug("WebSocket disconnected")
+        except Exception:
+            self.logger.exception("WebSocket handler failed")
+            try:
+                await websocket.close(code=1011)
+            except Exception:
+                pass
+
     def run(self, port: int = 8080, host: Optional[str] = None, **kwargs):
         """Start the Bedrock AgentCore server.
 

tests/bedrock_agentcore/runtime/test_app.py

@@ -1949,3 +1949,202 @@ def __init__(self):
         assert result["X-Amzn-Bedrock-AgentCore-Runtime-Custom-Empty"] == ""
         assert result["X-Amzn-Bedrock-AgentCore-Runtime-Custom-Valid"] == "valid-value"
         assert len(result) == 3
+
+
+class TestWebSocketSupport:
+    """Test WebSocket decorator and handler functionality."""
+
+    def test_websocket_initialization(self):
+        """Test that WebSocket route is registered during initialization."""
+        app = BedrockAgentCoreApp()
+        routes = app.routes
+        route_paths = [route.path for route in routes]  # type: ignore
+
+        assert "/ws" in route_paths
+
+    def test_websocket_decorator(self):
+        """Test @app.websocket decorator registers handler."""
+        app = BedrockAgentCoreApp()
+
+        @app.websocket
+        async def test_handler(websocket, context):
+            await websocket.accept()
+
+        assert app._websocket_handler is not None
+        assert app._websocket_handler == test_handler
+
+    def test_websocket_no_handler_defined(self):
+        """Test WebSocket endpoint when no handler is defined."""
+        from starlette.websockets import WebSocketDisconnect
+
+        app = BedrockAgentCoreApp()
+        client = TestClient(app)
+
+        with pytest.raises((WebSocketDisconnect, RuntimeError)):
+            with client.websocket_connect("/ws"):
+                pass
+
+    def test_websocket_basic_communication(self):
+        """Test basic WebSocket send/receive."""
+        app = BedrockAgentCoreApp()
+
+        @app.websocket
+        async def handler(websocket, context):
+            await websocket.accept()
+            data = await websocket.receive_json()
+            await websocket.send_json({"echo": data})
+            await websocket.close()
+
+        client = TestClient(app)
+
+        with client.websocket_connect("/ws") as websocket:
+            websocket.send_json({"message": "Hello"})
+            response = websocket.receive_json()
+            assert response == {"echo": {"message": "Hello"}}
+
+    def test_websocket_with_context(self):
+        """Test WebSocket handler receives context with session ID."""
+        app = BedrockAgentCoreApp()
+
+        received_context = None
+
+        @app.websocket
+        async def handler(websocket, context):
+            nonlocal received_context
+            received_context = context
+            await websocket.accept()
+            await websocket.send_json({"session_id": context.session_id})
+            await websocket.close()
+
+        client = TestClient(app)
+
+        with client.websocket_connect(
+            "/ws", headers={"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "ws-session-123"}
+        ) as websocket:
+            response = websocket.receive_json()
+            assert response["session_id"] == "ws-session-123"
+            assert received_context is not None
+            assert received_context.session_id == "ws-session-123"
+
+    def test_websocket_handler_exception(self):
+        """Test WebSocket handler exceptions are caught and logged."""
+        from starlette.websockets import WebSocketDisconnect
+
+        app = BedrockAgentCoreApp()
+
+        @app.websocket
+        async def handler(websocket, context):
+            await websocket.accept()
+            raise ValueError("Test WebSocket error")
+
+        client = TestClient(app)
+
+        with pytest.raises((WebSocketDisconnect, ValueError, RuntimeError)):
+            with client.websocket_connect("/ws") as websocket:
+                websocket.receive_json()
+
+    def test_websocket_multiple_messages(self):
+        """Test WebSocket can handle multiple messages."""
+        app = BedrockAgentCoreApp()
+
+        @app.websocket
+        async def handler(websocket, context):
+            await websocket.accept()
+            for _ in range(3):
+                data = await websocket.receive_json()
+                await websocket.send_json({"received": data})
+            await websocket.close()
+
+        client = TestClient(app)
+
+        with client.websocket_connect("/ws") as websocket:
+            for i in range(3):
+                websocket.send_json({"count": i})
+                response = websocket.receive_json()
+                assert response == {"received": {"count": i}}
+
+    def test_websocket_disconnect_handling(self):
+        """Test WebSocket gracefully handles client disconnect."""
+        from starlette.websockets import WebSocketDisconnect
+
+        app = BedrockAgentCoreApp()
+
+        disconnect_handled = False
+
+        @app.websocket
+        async def handler(websocket, context):
+            nonlocal disconnect_handled
+            await websocket.accept()
+            try:
+                while True:
+                    await websocket.receive_json()
+            except WebSocketDisconnect:
+                disconnect_handled = True
+                raise
+
+        client = TestClient(app)
+
+        with client.websocket_connect("/ws") as websocket:
+            websocket.send_json({"message": "test"})
+
+        # Disconnect should be handled gracefully
+        assert disconnect_handled
+
+    def test_websocket_with_request_headers(self):
+        """Test WebSocket handler receives custom request headers via context."""
+        app = BedrockAgentCoreApp()
+
+        received_headers = None
+
+        @app.websocket
+        async def handler(websocket, context):
+            nonlocal received_headers
+            received_headers = context.request_headers
+            await websocket.accept()
+            await websocket.send_json({"has_headers": context.request_headers is not None})
+            await websocket.close()
+
+        client = TestClient(app)
+
+        headers = {
+            "Authorization": "Bearer ws-token",
+            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-ClientId": "ws-client-123",
+        }
+
+        with client.websocket_connect("/ws", headers=headers) as websocket:
+            response = websocket.receive_json()
+            assert response["has_headers"] is True
+
+        assert received_headers is not None
+        # Find authorization header (case-insensitive)
+        auth_key = next((k for k in received_headers.keys() if k.lower() == "authorization"), None)
+        assert auth_key is not None
+        assert received_headers[auth_key] == "Bearer ws-token"
+
+    def test_websocket_streaming_data(self):
+        """Test WebSocket can stream multiple data chunks."""
+        app = BedrockAgentCoreApp()
+
+        @app.websocket
+        async def handler(websocket, context):
+            await websocket.accept()
+            # Stream data
+            for i in range(5):
+                await websocket.send_json({"chunk": i, "data": f"chunk_{i}"})
+            await websocket.send_json({"done": True})
+            await websocket.close()
+
+        client = TestClient(app)
+
+        with client.websocket_connect("/ws") as websocket:
+            chunks = []
+            for _ in range(5):
+                chunk = websocket.receive_json()
+                chunks.append(chunk)
+
+            final = websocket.receive_json()
+
+            assert len(chunks) == 5
+            assert chunks[0] == {"chunk": 0, "data": "chunk_0"}
+            assert chunks[4] == {"chunk": 4, "data": "chunk_4"}
+            assert final == {"done": True}