docs(event_handler): finish API Gateway WebSocket resolver docs

Adds the overview and workflow mermaid diagrams, accessing the event
and Lambda context, and the sending-messages recipe: a connection
store capturing connection_id and callback_url at $connect, an
in-resolver broadcast, and a separate function pushing a result to
the requesting client later — with GoneException handling,
ManageConnections IAM, and a custom domain caveat.

Author: Iamrodos

docs/core/event_handler/api_gateway_websocket.md

@@ -6,6 +6,32 @@ status: new
 
 Event Handler for Amazon API Gateway WebSocket APIs.
 
+```mermaid
+stateDiagram-v2
+    direction LR
+    EventSource: WebSocket client
+    Gateway: Amazon API Gateway
+    LambdaInit: Lambda invocation
+    EventHandler: Event Handler
+    EventHandlerResolver: Route event based on route key
+    YourLogic: Run your registered route handler
+    EventHandlerResolverBuilder: Normalize response (statusCode/body)
+    LambdaResponse: Lambda response
+
+    EventSource --> Gateway: $connect, messages, $disconnect
+    Gateway --> LambdaInit: Route selection
+
+    LambdaInit --> EventHandler
+
+    state EventHandler {
+        [*] --> EventHandlerResolver: app.resolve(event, context)
+        EventHandlerResolver --> YourLogic
+        YourLogic --> EventHandlerResolverBuilder
+    }
+
+    EventHandler --> LambdaResponse
+```
+
 ## Key Features
 
 * Route WebSocket events by route key with dedicated `$connect`, `$disconnect`, and `$default` decorators
@@ -185,6 +211,109 @@ Register handlers for specific exception types with `@app.exception_handler`; it
 ???+ note "Unhandled exceptions never reach the client"
     An exception with no registered handler is logged and becomes a bare `{"statusCode": 500}`. If it propagated to the Lambda runtime instead, the runtime's error response — exception type, message, and stack trace — would be delivered to the connected client.
 
+### Accessing the event and Lambda context
+
+Inside handlers and middlewares, `app.current_event` is the current `APIGatewayWebSocketEvent` ([Event Source Data Classes](../../utilities/data_classes.md){target="_blank"} utility) and `app.lambda_context` is the Lambda context.
+
+Use `app.append_context` / `app.context` to share data between middlewares and handlers within a single invocation — the context is cleared after each `resolve`.
+
+=== "accessing_websocket_event_and_context.py"
+
+    ```python hl_lines="9 13 14"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py"
+    ```
+
+    1. Route, connection, and identity details live in `request_context`.
+    2. `json_body` parses the message body; `body` and `decoded_body` are also available.
+    3. The standard Lambda context object.
+
+### Sending messages to connected clients
+
+Returned bodies only reply to the **calling** client, on routes with a route response. To push a message to a client at any time — including after the invocation that received the request has long finished — call the API Gateway Management API's `PostToConnection` with the connection ID, against the endpoint the `callback_url` property gives you.
+
+Capture `connection_id` and `callback_url` together at `$connect` and persist them: that record is everything **any** process needs to push to that client later.
+
+The example shows both sending situations. Inside the resolver, `broadcast` pushes to every stored connection from the same invocation, building the client from the current event.
+
+Outside it, `submitReport` acknowledges a long-running request immediately, and a **separate Lambda function** — for example the final state of a Step Functions workflow — pushes the finished result to the one client that asked, using the stored `callback_url`. Progress updates are that same call made mid-work.
+
+A stale connection ID raises `GoneException` — the client is no longer connected, so treat it as a signal to remove that connection from your store. The posting function also needs the `execute-api:ManageConnections` IAM permission on `arn:aws:execute-api:{region}:{account}:{api-id}/{stage}/POST/@connections/*`.
+
+=== "working_with_post_to_connection.py"
+
+    ```python hl_lines="2 16 28 35 41"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py"
+    ```
+
+    1. The only thing the WebSocket resolver and the pushing function share.
+    2. Captured at `$connect`: the only moment you get for free, and all anyone needs to post to this client later.
+    3. Acknowledge immediately; the result is pushed by another function when it is ready. For the ack to reach the client, configure a route response on `submitReport`.
+    4. Sending from inside the resolver: the current event provides the endpoint directly.
+
+=== "working_with_post_to_connection_report_ready.py"
+
+    ```python hl_lines="18 23"
+    --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py"
+    ```
+
+    1. The stored `callback_url` rebuilds the Management API client in any execution environment.
+    2. The client disconnected before the result was ready — remove it from the store.
+
+=== "my_connection_store.py"
+
+    ```python
+    --8<-- "examples/event_handler_api_gateway_websocket/src/my_connection_store.py"
+    ```
+
+???+ note "Custom domain names"
+    `callback_url` is built from the event's domain name. When clients connect through a custom domain whose API mapping path differs from the stage name, construct the endpoint yourself instead: `https://{api_id}.execute-api.{region}.amazonaws.com/{stage}`.
+
+## Event Handler workflow
+
+### Connection lifecycle
+
+Connection accepted, message exchanged, connection closed.
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client as WebSocket client
+    participant Gateway as Amazon API Gateway
+    participant Lambda as Lambda (Event Handler)
+    Client->>Gateway: Upgrade request (wss://)
+    Gateway->>Lambda: $connect event
+    Lambda-->>Gateway: {"statusCode": 200}
+    Gateway-->>Client: 101 Switching Protocols
+    Client->>Gateway: {"action": "orderUpdate", ...}
+    Gateway->>Lambda: orderUpdate event (route selection expression)
+    Lambda-->>Gateway: {"statusCode": 200, "body": "..."}
+    Gateway-->>Client: body (only with a route response configured)
+    Client->>Gateway: Close frame
+    Gateway->>Lambda: $disconnect event (best effort)
+    Lambda-->>Gateway: {"statusCode": 200}
+```
+</center>
+
+### Rejected connection
+
+Connection rejected by the `$connect` handler.
+
+<center>
+```mermaid
+sequenceDiagram
+    participant Client as WebSocket client
+    participant Gateway as Amazon API Gateway
+    participant Lambda as Lambda (Event Handler)
+    Client->>Gateway: Upgrade request (wss://)
+    Gateway->>Lambda: $connect event
+    Lambda-->>Gateway: {"statusCode": 401}
+    Gateway-->>Client: HTTP 401 (upgrade refused)
+    Note over Gateway,Lambda: $disconnect may still be delivered for this connection
+    Gateway->>Lambda: $disconnect event
+    Lambda-->>Gateway: {"statusCode": 200}
+```
+</center>
+
 ## Testing your code
 
 Test your handlers by passing a WebSocket event payload to `app.resolve()` — no mocks of API Gateway needed.

examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py

@@ -0,0 +1,19 @@
+from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+app = APIGatewayWebSocketResolver()
+
+
+@app.on_default()
+def default():
+    request_context = app.current_event.request_context  # (1)!
+    return {
+        "connectionId": request_context.connection_id,
+        "routeKey": request_context.route_key,
+        "message": app.current_event.json_body,  # (2)!
+        "remainingTimeMs": app.lambda_context.get_remaining_time_in_millis(),  # (3)!
+    }
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)

examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json

@@ -46,4 +46,4 @@
       "apiId": "asasasas"
   },
   "isBase64Encoded": false
-}
+}

examples/event_handler_api_gateway_websocket/src/my_connection_store.py

@@ -0,0 +1,21 @@
+# Illustrative connection store. In production, back these functions with a durable store
+# such as a DynamoDB table: the WebSocket resolver and the functions that push results run
+# in different Lambda execution environments and cannot share process memory.
+
+_connections: dict[str, str] = {}
+
+
+def save(connection_id: str, callback_url: str) -> None:
+    _connections[connection_id] = callback_url
+
+
+def get_callback_url(connection_id: str) -> str:
+    return _connections[connection_id]
+
+
+def all_connection_ids() -> list[str]:
+    return list(_connections)
+
+
+def delete(connection_id: str) -> None:
+    _connections.pop(connection_id, None)

examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py

@@ -0,0 +1,46 @@
+import boto3
+import my_connection_store  # (1)!
+
+from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+app = APIGatewayWebSocketResolver()
+
+
+def start_report_generation(report_id: str, connection_id: str) -> None: ...  # e.g. start a Step Functions execution
+
+
+@app.on_connect()
+def connect():
+    request_context = app.current_event.request_context
+    my_connection_store.save(request_context.connection_id, request_context.callback_url)  # (2)!
+
+
+@app.on_disconnect()
+def disconnect():
+    my_connection_store.delete(app.current_event.request_context.connection_id)
+
+
+@app.route("submitReport")
+def submit_report():
+    request = app.current_event.json_body
+    start_report_generation(request["reportId"], app.current_event.request_context.connection_id)
+    return {"status": "accepted"}  # (3)!
+
+
+@app.route("broadcast")
+def broadcast():
+    client = boto3.client(
+        "apigatewaymanagementapi",
+        endpoint_url=app.current_event.request_context.callback_url,  # (4)!
+    )
+    message: str = app.current_event.json_body["message"]
+    for connection_id in my_connection_store.all_connection_ids():
+        try:
+            client.post_to_connection(ConnectionId=connection_id, Data=message.encode())
+        except client.exceptions.GoneException:
+            my_connection_store.delete(connection_id)
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)

examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py

@@ -0,0 +1,25 @@
+# A separate Lambda function, in a different execution environment from the WebSocket
+# resolver — invoked when the work completes, e.g. as the final state of a Step Functions
+# workflow. The connection store is the only thing the two functions share.
+
+import json
+
+import boto3
+import my_connection_store
+
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger()
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> None:
+    connection_id = event["connectionId"]
+    callback_url = my_connection_store.get_callback_url(connection_id)  # (1)!
+
+    client = boto3.client("apigatewaymanagementapi", endpoint_url=callback_url)
+    try:
+        client.post_to_connection(ConnectionId=connection_id, Data=json.dumps(event["report"]).encode())
+    except client.exceptions.GoneException:  # (2)!
+        logger.info("Client disconnected before the report was ready", connection_id=connection_id)
+        my_connection_store.delete(connection_id)