feat(event_handler): add middleware support to API Gateway WebSocket resolver

Reuses the shared middleware framework (MiddlewareFrame,
BaseMiddlewareHandler) with a WebSocket terminal adapter: global
middlewares registered via use() run first, then route-level
middlewares=[...], then the handler. Return values — including
short-circuits that skip next_middleware — go through the same
response normalization as handler returns, and middleware exceptions
follow the existing exception-handling path. Global middlewares also
run around the 400 for unmatched route keys, matching the REST
resolver's not-found behaviour.

Includes middleware tests (ordering, short-circuit, isolation across
routes, exception paths, class-based middleware) and docs: Middleware
plus Authentication patterns (Lambda authorizer first, then
authenticate/inject_user middleware with a bring-your-own connection
store).

Author: Iamrodos

aws_lambda_powertools/event_handler/api_gateway_websocket/base.py

@@ -43,6 +43,10 @@ def on_default(
     ) -> Callable:
         raise NotImplementedError
 
+    @abstractmethod
+    def use(self, middlewares: list[Callable]) -> None:
+        raise NotImplementedError
+
     def append_context(self, **additional_context) -> None:
         """
         Appends context information available under any route.

aws_lambda_powertools/event_handler/api_gateway_websocket/router.py

@@ -55,6 +55,7 @@ def __init__(self):
         self.context = {}  # early init as customers might add context before event resolution
         self._route_registry = RouteRegistry()
         self._exception_handlers: dict[type[Exception], Callable] = {}
+        self._router_middlewares: list[Callable] = []
 
     def route(
         self,
@@ -189,6 +190,33 @@ def on_default(
         """
         return self.route(route_key=ROUTE_KEY_DEFAULT, middlewares=middlewares)
 
+    def use(self, middlewares: list[Callable]) -> None:
+        """
+        Add one or more global middlewares that run before route-specific middlewares.
+
+        Middlewares run in insertion order: global middlewares first, then route-level
+        `middlewares=[...]`, then the route handler.
+
+        Parameters
+        ----------
+        middlewares : list[Callable]
+            List of middlewares to run on every event, including unmatched route keys
+
+        Examples
+        --------
+        >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+        >>> from aws_lambda_powertools.event_handler.middlewares import NextMiddleware
+        >>>
+        >>> app = APIGatewayWebSocketResolver()
+        >>>
+        >>> def log_route(app: APIGatewayWebSocketResolver, next_middleware: NextMiddleware):
+        >>>     print(f"dispatching {app.current_event.request_context.route_key}")
+        >>>     return next_middleware(app)
+        >>>
+        >>> app.use(middlewares=[log_route])
+        """
+        self._router_middlewares.extend(middlewares)
+
     def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) -> Callable:
         """
         Register a handler for one or more exception types.

aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py

@@ -6,7 +6,9 @@
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Any
 
+from aws_lambda_powertools.event_handler.api_gateway import MiddlewareFrame
 from aws_lambda_powertools.event_handler.api_gateway_websocket.router import Router
+from aws_lambda_powertools.event_handler.api_gateway_websocket.types import WebSocketRoute
 from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager
 from aws_lambda_powertools.shared.json_encoder import Encoder
 from aws_lambda_powertools.utilities.data_classes.api_gateway_websocket_event import APIGatewayWebSocketEvent
@@ -20,6 +22,20 @@
 logger = logging.getLogger(__name__)
 
 
+def _registered_route_adapter(app: APIGatewayWebSocketResolver, next_middleware: Callable[..., Any]) -> Any:
+    """Terminal middleware frame: calls the registered route handler (no arguments) and returns its raw value.
+
+    Response normalization happens once, after the whole middleware chain returns, so middlewares
+    may short-circuit with any of the same shapes route handlers can return.
+    """
+    return next_middleware()
+
+
+def _route_not_found() -> tuple[None, int]:
+    """Stand-in handler for unmatched route keys, so global middlewares still run around the 400."""
+    return None, HTTPStatus.BAD_REQUEST.value
+
+
 class APIGatewayWebSocketResolver(Router):
     """
     API Gateway WebSocket API resolver.
@@ -70,6 +86,7 @@ class APIGatewayWebSocketResolver(Router):
     def __init__(self):
         super().__init__()
         self.exception_handler_manager = ExceptionHandlerManager()
+        self.processed_stack_frames: list[str] = []  # used by MiddlewareFrame for debug traces
 
     def __call__(
         self,
@@ -110,6 +127,7 @@ def resolve(
         >>>     return app.resolve(event, context)
         """
         try:
+            self.processed_stack_frames.clear()
             self._setup_context(event, context)
             return self._resolve_route()
         except Exception as exc:
@@ -149,7 +167,7 @@ def exception_handler(self, exc_class: type[Exception] | list[type[Exception]])
         return self.exception_handler_manager.exception_handler(exc_class=exc_class)
 
     def _resolve_route(self) -> dict[str, Any]:
-        """Dispatch the current event to its route handler and normalize the response."""
+        """Dispatch the current event through the middleware chain to its route handler and normalize the response."""
         route_key = self.current_event.request_context.route_key
         route = self._route_registry.find_route(route_key)
         if route is None:
@@ -158,10 +176,26 @@ def _resolve_route(self) -> dict[str, Any]:
                 stacklevel=2,
                 category=PowertoolsUserWarning,
             )
-            return {"statusCode": HTTPStatus.BAD_REQUEST.value}
+            # global middlewares still run around the 400, as REST does for not-found routes
+            route = WebSocketRoute(func=_route_not_found, middlewares=[])
 
         logger.debug(f"Dispatching route key `{route_key}` to `{route['func'].__name__}`")
-        return self._to_response(route["func"]())
+        middleware_stack = self._build_middleware_stack(route)
+        return self._to_response(middleware_stack(self))
+
+    def _build_middleware_stack(self, route: WebSocketRoute) -> Callable:
+        """Wrap the route handler in middleware frames: global (`use`) first, then route-level, then
+        the terminal adapter that calls the handler. Frames are wrapped in reverse so middlewares
+        run in the order they were added."""
+        stack: Callable = route["func"]
+        all_middlewares = [*self._router_middlewares, *route["middlewares"], _registered_route_adapter]
+        for middleware in reversed(all_middlewares):
+            stack = MiddlewareFrame(current_middleware=middleware, next_middleware=stack)
+        return stack
+
+    def _push_processed_stack_frame(self, frame: str) -> None:
+        """Record a processed middleware frame; MiddlewareFrame calls this for debug traces."""
+        self.processed_stack_frames.append(frame)
 
     def _handle_exception(self, exc: Exception) -> dict[str, Any]:
         """Resolve an exception to a response via registered handlers, or a bare 500.

docs/core/event_handler/api_gateway_websocket.md

@@ -99,6 +99,56 @@ The status code has different effects depending on the route:
 
 ## Advanced
 
+### Middleware
+
+The resolver reuses the same middleware framework as the REST resolver. A middleware is a callable receiving the resolver instance and the next handler in the chain; it runs code before and/or after the route handler, and its return value goes through the same [response normalization](#response-format) as handler returns.
+
+Execution order is: global middlewares (`app.use`), then route-level `middlewares=[...]`, then the route handler. Return without calling `next_middleware(app)` to short-circuit the chain; exceptions raised in middlewares follow the same [exception handling](#exception-handling) flow as handlers.
+
+=== "working_with_middleware.py"
+
+    ```python hl_lines="9 11 17 21 24"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_middleware.py"
+    ```
+
+    1. Middlewares receive the resolver instance and the next handler in the chain.
+    2. Call `next_middleware(app)` to continue the chain; its return value is the route response.
+    3. Short-circuit: the handler never runs and this value becomes the response.
+    4. Global middlewares run on every event — including unmatched route keys — before route-level ones.
+    5. Route-level middlewares run for this route only.
+
+### Authentication patterns
+
+`$connect` is the only route where credentials exist: headers and cookies are present on the WebSocket handshake only, and later messages carry nothing but the connection ID. There are two ways to bridge that gap.
+
+#### Lambda authorizer
+
+Prefer a Lambda authorizer on `$connect` when you can run a separate authorizer function. API Gateway persists the authorizer's output with the connection and injects it into `request_context.authorizer` on **every** invocation for that connection — a managed connection-to-identity store, no code needed in your handlers beyond reading it.
+
+=== "working_with_lambda_authorizer.py"
+
+    ```python hl_lines="9 14-15"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py"
+    ```
+
+    1. Injected by API Gateway on every route for this connection — message routes, `$default`, and `$disconnect` included.
+    2. Authorizer context values arrive stringified (a numeric `1` arrives as `"1"`).
+
+#### Authenticating with middleware
+
+Without a Lambda authorizer, authenticate on `$connect` with a middleware and persist identity against the connection ID in your own store; a second middleware resolves it on message routes and shares it via `app.append_context`.
+
+=== "working_with_middleware_authentication.py"
+
+    ```python hl_lines="6 13-14 16 25 29 39"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py"
+    ```
+
+    1. Bring your own connection store. The in-memory dictionary keeps this example short — real invocations for one connection can hit different Lambda environments, so use an external store such as DynamoDB with a TTL.
+    2. Headers only exist on `$connect`.
+    3. Returning without calling `next_middleware` short-circuits the chain — the handler never runs and the connection is rejected.
+    4. Handlers and later middlewares read it from `app.context`.
+
 ### Exception handling
 
 Register handlers for specific exception types with `@app.exception_handler`; it also accepts a list of types, and lookup respects inheritance. The handler receives the exception, and its return value goes through the same [response normalization](#response-format).

examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py

@@ -0,0 +1,20 @@
+from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+app = APIGatewayWebSocketResolver()
+
+
+@app.route("orderUpdate")
+def order_update():
+    identity = app.current_event.request_context.authorizer  # (1)!
+    order = app.current_event.json_body
+    return {
+        "orderId": order["orderId"],
+        "status": "received",
+        "processedFor": identity["principalId"],
+        "tenant": identity.get("tenantId"),  # (2)!
+    }
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)

examples/event_handler_api_gateway_websocket/src/working_with_middleware.py

@@ -0,0 +1,31 @@
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger()
+app = APIGatewayWebSocketResolver()
+
+
+def log_route(app, next_middleware):  # (1)!
+    logger.info("Dispatching", route_key=app.current_event.request_context.route_key)
+    return next_middleware(app)  # (2)!
+
+
+def enforce_order_limit(app, next_middleware):
+    order = app.current_event.json_body
+    if order.get("quantity", 0) > 100:
+        return {"error": "quantity over limit"}, 400  # (3)!
+    return next_middleware(app)
+
+
+app.use(middlewares=[log_route])  # (4)!
+
+
+@app.route("orderUpdate", middlewares=[enforce_order_limit])  # (5)!
+def order_update():
+    order = app.current_event.json_body
+    return {"orderId": order["orderId"], "status": "received"}
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)

examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py

@@ -0,0 +1,46 @@
+from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+app = APIGatewayWebSocketResolver()
+
+sessions: dict[str, str] = {}  # (1)!
+
+
+def validate_token(token: str | None) -> str | None:
+    return "alice" if token == "Bearer valid-token" else None
+
+
+def authenticate(app, next_middleware):
+    user = validate_token(app.current_event.headers.get("Authorization"))  # (2)!
+    if user is None:
+        return None, 401  # (3)!
+    sessions[app.current_event.request_context.connection_id] = user
+    return next_middleware(app)
+
+
+def inject_user(app, next_middleware):
+    user = sessions.get(app.current_event.request_context.connection_id)
+    if user is None:
+        return {"error": "unauthenticated"}, 401
+    app.append_context(user=user)  # (4)!
+    return next_middleware(app)
+
+
+@app.on_connect(middlewares=[authenticate])
+def connect():
+    return None
+
+
+@app.on_disconnect()
+def disconnect():
+    sessions.pop(app.current_event.request_context.connection_id, None)
+
+
+@app.route("orderUpdate", middlewares=[inject_user])
+def order_update():
+    order = app.current_event.json_body
+    return {"orderId": order["orderId"], "status": "received", "processedFor": app.context["user"]}
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)