How to use AWS Websocket API Gateway Lambda endpoint callback function

Question

When the client sends a websocket message to AWS websocket API Gateway the websocket Lambda handler function receives three arguments: event, context and callback:

import type { APIGatewayProxyEvent, Context, } from "aws-lambda";

export class Handler {
  constructor() {}

  async execute(event: APIGatewayProxyEvent, context: Context, callback: any): Promise<void> {
    console.log(`${event}, ${context}, ${callback}`);

    if (!event.requestContext || !event.requestContext.connectionId) {
      callback(null, { statusCode: 400, body: "Missing RequestContext" });
    };

    callback(null, {statusCode: 200, body: 'Your message was received'});
  }
}

const Handler = new Handler(service);
async function handler(event: APIGatewayProxyEvent, context: Context, callback: any): Promise<void> {
  return await Handler.execute(event, context, callback);
}
export default handler;

The Handler checks if the event that it received contains the value for event.requestContext.connectionId. If not, it calls the callback function passing null as the first argument and a message containing the 400 statusCode and bodywith message telling that the connectionId is missing:

callback(null, { statusCode: 400, body: "Missing RequestContext" });

Unfortunately, I was not able to locate any AWS documentation describing how exactly we should be using this callback function. Instead I've used some code examples found on Internet.

I was expecting that calling the callback function passing it 400 status will make the Handler function exit and raise an exception. But it doesn't happen. The handler just keeps going executing the next lines like nothing happened and callback function was never called.

1. What is a real purpose for the callback function?
2. How do we return an error message from websocket Handler function letting the client application know about the occurred problem?
3. How do we stop the Handler from executing the rest of the code if there was an exception or invalid parameter received?
4. Should the websocket Handler return any value or should it be void - no value returned?

Please post your answer as a script example or as a snippet so we could test it.

P.S. The doc posted at https://aws.amazon.com/blogs/compute/announcing-websocket-apis-in-amazon-api-gateway/ says that:

The backend can send messages to specific users via a callback URL that is provided after the user is connected to the WebSocket API.

And there is a snippet showing it:

exports.handler = function(event, context, callback) {
  var putParams = {
    TableName: process.env.TABLE_NAME,
    Item: {
      connectionId: { S: event.requestContext.connectionId }
    }
  };

  DDB.putItem(putParams, function(err, data) {
    callback(null, {
      statusCode: err ? 500 : 200,
      body: err ? "Failed to connect: " + JSON.stringify(err) : "Connected"
    });
  });
};

Answer 1

There are a few important things to understand about two-way communication patterns using Lambda integrations for Websocket APIs:

Disambiguating "Callback"

"Callback" is an overloaded term. First, it refers to the legacy pattern for Nodejs Lambda function handlers. We should factor this out in favour of the AWS-recommended async-await handler pattern. Second, AWS has sometimes referred to "Callback URLs" as way for backends to communicate with Websocket clients using the @aws-sdk/client-apigatewaymanagementapi client¹. The docs also refer to the same thing as @connections commands.

One problem with your code is that you are mixing up the unrelated Lambda and Websocket API "callback" concepts. To avoid ambiguity, I will avoid the term "callback" altogether.

There are 2 two-way communication patterns

There are two ways to send data from backend services to connected clients:

Method 1: $default route integration response with a RouteResponse

For the $default route only, your Lambda can return a message to a client in the handler response. This will only work if you have added a RouteResponse to the route².

// default-route-handler.ts
import { APIGatewayProxyWebsocketEventV2, APIGatewayProxyResult } from 'aws-lambda';

export async function handler(event: APIGatewayProxyWebsocketEventV2): Promise<APIGatewayProxyResult> {
  return {
    statusCode: 200,
    body: JSON.stringify('hello client, I am the $default route!'),
  };
}

Method 2: @connections API

For all routes, you can use the @connections API to send a POST request. Do this with the ApiGatewayManagementApi client's PostToConnection API, which handles the heavy lifting of sending the message to the client. The Lambda just returns a status.

// another-great-handler.ts
import { APIGatewayProxyWebsocketEventV2 } from 'aws-lambda';
import { ApiGatewayManagementApiClient, PostToConnectionCommand } from '@aws-sdk/client-apigatewaymanagementapi';

interface StatusCodeResponse {
  statusCode: 200 | 500;
}

export async function handler(event: APIGatewayProxyWebsocketEventV2): Promise<StatusCodeResponse> {
  const { requestContext: { domainName, stage, connectionId, connectedAt, routeKey },} = event;

  try {
    const client = new ApiGatewayManagementApiClient({
      region: process.env.AWS_REGION,
      endpoint: 'https://' + domainName + '/' + stage,
    });

    const encoder = new TextEncoder();

    const postCmd = new PostToConnectionCommand({
      ConnectionId: connectionId,
      Data: encoder.encode(`Hello from ${routeKey}. You connected at: ${connectedAt}`),
    });

    await client.send(postCmd);
  } catch (err: unknown) {
    return { statusCode: 500 };
  }

  return { statusCode: 200 };
}

---

1. I am using the AWS JS SDK v3

2. Add a RouteResponse in the Console ($default route > "Add integration response" button), or with the create-route-response API, or with CloudFormation AWS::ApiGatewayV2::RouteResponse