API Gateway WebSocket PostToConnection With AWS SDK v3

// connections.ts: $connect and $disconnect route handlers that keep a DynamoDB table of live connection IDs.
// Table "WsConnections": partition key connectionId (string), TTL attribute expiresAt (epoch seconds).
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DynamoDBDocumentClient, PutCommand, DeleteCommand } from "@aws-sdk/lib-dynamodb";
import type { APIGatewayProxyWebsocketEventV2, APIGatewayProxyResultV2 } from "aws-lambda";

const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const TABLE = process.env.CONNECTIONS_TABLE ?? "WsConnections";
const MAX_CONNECTION_SECONDS = 2 * 60 * 60; // API Gateway closes WebSocket connections after 2 hours

export async function onConnect(event: APIGatewayProxyWebsocketEventV2): Promise<APIGatewayProxyResultV2> {
  const { connectionId, domainName, stage } = event.requestContext;
  const now = Math.floor(Date.now() / 1000);
  await ddb.send(
    new PutCommand({
      TableName: TABLE,
      Item: {
        connectionId,
        callbackUrl: `https://${domainName}/${stage}`,
        connectedAt: now,
        expiresAt: now + MAX_CONNECTION_SECONDS + 300, // TTL sweeps rows whose $disconnect never arrived
      },
    }),
  );
  return { statusCode: 200 };
}

export async function onDisconnect(event: APIGatewayProxyWebsocketEventV2): Promise<APIGatewayProxyResultV2> {
  await ddb.send(new DeleteCommand({ TableName: TABLE, Key: { connectionId: event.requestContext.connectionId } }));
  return { statusCode: 200 };
}

// ws-send.ts: send messages to WebSocket clients with PostToConnectionCommand, remove stale
// connection IDs on GoneException (HTTP 410), and fan out to many clients with a concurrency limit.
import {
  ApiGatewayManagementApiClient,
  PostToConnectionCommand,
  GoneException,
  PayloadTooLargeException,
  LimitExceededException,
} from "@aws-sdk/client-apigatewaymanagementapi";
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DynamoDBDocumentClient, DeleteCommand, paginateScan } from "@aws-sdk/lib-dynamodb";

const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const TABLE = process.env.CONNECTIONS_TABLE ?? "WsConnections";
const MAX_PAYLOAD_BYTES = 128 * 1024; // API Gateway WebSocket message payload quota

// The endpoint is the connection URL with https:// instead of wss://, including the stage:
// https://{api-id}.execute-api.{region}.amazonaws.com/{stage}. With a custom domain, use https://{domain}
// plus the base path you mapped, if any.
const clients = new Map<string, ApiGatewayManagementApiClient>();
function clientFor(endpoint: string): ApiGatewayManagementApiClient {
  let client = clients.get(endpoint);
  if (!client) {
    client = new ApiGatewayManagementApiClient({ endpoint, region: process.env.AWS_REGION });
    clients.set(endpoint, client);
  }
  return client;
}

export type SendResult = "sent" | "gone" | "too-large" | "failed";

/** Send one JSON message to one connection. A 410 means the client is gone: delete its row. */
export async function sendToConnection(endpoint: string, connectionId: string, message: unknown): Promise<SendResult> {
  const data = new TextEncoder().encode(JSON.stringify(message));
  if (data.byteLength > MAX_PAYLOAD_BYTES) return "too-large";
  try {
    await clientFor(endpoint).send(new PostToConnectionCommand({ ConnectionId: connectionId, Data: data }));
    return "sent";
  } catch (err) {
    if (err instanceof GoneException) {
      await ddb.send(new DeleteCommand({ TableName: TABLE, Key: { connectionId } }));
      return "gone";
    }
    if (err instanceof PayloadTooLargeException) return "too-large";
    if (err instanceof LimitExceededException) throw err; // let the caller slow down or retry later
    console.error(`postToConnection ${connectionId}:`, err);
    return "failed";
  }
}

/** Run tasks with at most `limit` in flight. */
async function withConcurrency<T>(items: T[], limit: number, task: (item: T) => Promise<void>): Promise<void> {
  let next = 0;
  const workers = Array.from({ length: Math.min(limit, items.length) }, async () => {
    while (next < items.length) {
      const item = items[next++];
      await task(item);
    }
  });
  await Promise.all(workers);
}

/** Broadcast to every stored connection. Scans the table, so keep it for small audiences. */
export async function broadcast(endpoint: string, message: unknown, concurrency = 25): Promise<Record<SendResult, number>> {
  const counts: Record<SendResult, number> = { sent: 0, gone: 0, "too-large": 0, failed: 0 };
  const ids: string[] = [];
  for await (const page of paginateScan({ client: ddb }, { TableName: TABLE, ProjectionExpression: "connectionId" })) {
    for (const item of page.Items ?? []) ids.push(String(item.connectionId));
  }
  await withConcurrency(ids, concurrency, async (id) => {
    counts[await sendToConnection(endpoint, id, message)]++;
  });
  return counts;
}

// chat-route.ts: handler for a custom "sendmessage" route. Clients send {"action":"sendmessage","text":"..."}.
import type { APIGatewayProxyWebsocketEventV2, APIGatewayProxyResultV2 } from "aws-lambda";
import { broadcast, sendToConnection } from "./ws-send.js";

export async function onSendMessage(event: APIGatewayProxyWebsocketEventV2): Promise<APIGatewayProxyResultV2> {
  const { connectionId, domainName, stage } = event.requestContext;
  const endpoint = `https://${domainName}/${stage}`; // default execute-api domain; see the custom domain note
  let text = "";
  try {
    text = String(JSON.parse(event.body ?? "{}").text ?? "").slice(0, 2000);
  } catch {
    await sendToConnection(endpoint, connectionId, { type: "error", error: "body must be JSON" });
    return { statusCode: 400 };
  }
  const counts = await broadcast(endpoint, { type: "message", from: connectionId, text, at: new Date().toISOString() });
  console.log(JSON.stringify({ route: "sendmessage", ...counts }));
  return { statusCode: 200 };
}

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "PostToWebSocketClients",
      "Effect": "Allow",
      "Action": "execute-api:ManageConnections",
      "Resource": "arn:aws:execute-api:us-east-1:123456789012:abc123/prod/POST/@connections/*"
    },
    {
      "Sid": "ConnectionsTable",
      "Effect": "Allow",
      "Action": ["dynamodb:PutItem", "dynamodb:DeleteItem", "dynamodb:Scan"],
      "Resource": "arn:aws:dynamodb:us-east-1:123456789012:table/WsConnections"
    }
  ]
}

// ws-send.test.ts: exercises ws-send.ts with aws-sdk-client-mock (run with: npx tsx test/ws-send.test.ts)
import assert from "node:assert/strict";
import { mockClient } from "aws-sdk-client-mock";
import { ApiGatewayManagementApiClient, PostToConnectionCommand, GoneException, PayloadTooLargeException } from "@aws-sdk/client-apigatewaymanagementapi";
import { DynamoDBDocumentClient, ScanCommand, DeleteCommand, PutCommand } from "@aws-sdk/lib-dynamodb";
import { broadcast, sendToConnection } from "../src/ws-send.ts";
import { onConnect } from "../src/connections.ts";

const api = mockClient(ApiGatewayManagementApiClient);
const ddb = mockClient(DynamoDBDocumentClient);
const endpoint = "https://abc123.execute-api.us-east-1.amazonaws.com/prod";

ddb.on(ScanCommand).resolvesOnce({ Items: [{ connectionId: "A" }, { connectionId: "B" }], LastEvaluatedKey: { connectionId: "B" } })
  .resolvesOnce({ Items: [{ connectionId: "C" }] });
ddb.on(DeleteCommand).resolves({});
ddb.on(PutCommand).resolves({});
api.on(PostToConnectionCommand).resolves({});
api.on(PostToConnectionCommand, { ConnectionId: "B" }).rejects(new GoneException({ message: "gone", $metadata: { httpStatusCode: 410 } }));
api.on(PostToConnectionCommand, { ConnectionId: "C" }).rejects(new PayloadTooLargeException({ message: "big", $metadata: { httpStatusCode: 413 } }));

const counts = await broadcast(endpoint, { hello: "world" }, 2);
assert.deepEqual(counts, { sent: 1, gone: 1, "too-large": 1, failed: 0 });
assert.equal(ddb.commandCalls(DeleteCommand).length, 1);
assert.deepEqual(ddb.commandCalls(DeleteCommand)[0].args[0].input.Key, { connectionId: "B" });
const sent = api.commandCalls(PostToConnectionCommand)[0].args[0].input;
assert.equal(new TextDecoder().decode(sent.Data as Uint8Array), '{"hello":"world"}');

assert.equal(await sendToConnection(endpoint, "A", { big: "x".repeat(140 * 1024) }), "too-large");

const res = await onConnect({ requestContext: { connectionId: "Z", domainName: "abc123.execute-api.us-east-1.amazonaws.com", stage: "prod" } } as never);
assert.deepEqual(res, { statusCode: 200 });
assert.equal(ddb.commandCalls(PutCommand)[0].args[0].input.Item?.callbackUrl, endpoint);
console.log("ws-send tests passed", counts);