API Gateway WebSockets for Multi-Turn Coding Agents: Connection State Without Sticky Myths

import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import {
 DynamoDBDocumentClient,
 PutCommand,
 DeleteCommand,
 UpdateCommand,
} from "@aws-sdk/lib-dynamodb";
import type {
 APIGatewayProxyWebsocketHandlerV2,
} from "aws-lambda";

const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const TABLE = process.env.SESSION_TABLE!;

export const connectHandler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
 const connectionId = event.requestContext.connectionId;
 const sessionId =
 event.queryStringParameters?.sessionId ?? crypto.randomUUID();
 const tenantId = event.queryStringParameters?.tenantId;
 if (!tenantId) return { statusCode: 401, body: "missing tenantId" };

 await ddb.send(
 new PutCommand({
 TableName: TABLE,
 Item: {
 pk: `conn#${connectionId}`,
 sk: "meta",
 sessionId,
 tenantId,
 status: "connected",
 ttl: Math.floor(Date.now() / 1000) + 86400,
 },
 })
 );
 await ddb.send(
 new PutCommand({
 TableName: TABLE,
 Item: {
 pk: `session#${sessionId}`,
 sk: "conn",
 connectionId,
 tenantId,
 ttl: Math.floor(Date.now() / 1000) + 86400,
 },
 })
 );

 return { statusCode: 200, body: "connected" };
};

export const disconnectHandler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
 const connectionId = event.requestContext.connectionId;
 await ddb.send(
 new DeleteCommand({
 TableName: TABLE,
 Key: { pk: `conn#${connectionId}`, sk: "meta" },
 })
 );
 // ✅ Keep session# row for resume; clear connectionId pointer
 // (lookup session via GSI or store reverse update)
 return { statusCode: 200, body: "disconnected" };
};

import {
 ApiGatewayManagementApiClient,
 PostToConnectionCommand,
} from "@aws-sdk/client-apigatewaymanagementapi";

function mgmtClient(domainName: string, stage: string) {
 return new ApiGatewayManagementApiClient({
 endpoint: `https://${domainName}/${stage}`,
 });
}

export async function pushToConnection(input: {
 domainName: string;
 stage: string;
 connectionId: string;
 payload: unknown;
}) {
 const client = mgmtClient(input.domainName, input.stage);
 try {
 await client.send(
 new PostToConnectionCommand({
 ConnectionId: input.connectionId,
 Data: Buffer.from(JSON.stringify(input.payload)),
 })
 );
 } catch (e: any) {
 // ✅ GoneException → connection dead; clear Dynamo mapping
 if (e.name === "GoneException" || e.$metadata?.httpStatusCode === 410) {
 await ddb.send(
 new DeleteCommand({
 TableName: TABLE,
 Key: { pk: `conn#${input.connectionId}`, sk: "meta" },
 })
 );
 return { ok: false, reason: "gone" };
 }
 throw e;
 }
 return { ok: true };
}

export const defaultHandler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
 const connectionId = event.requestContext.connectionId;
 const body = JSON.parse(event.body ?? "{}") as {
 type: string;
 sessionId: string;
 prompt?: string;
 interrupt?: boolean;
 };

 if (body.type === "agent.turn") {
 // Enqueue heavy work — do not run 90s Bedrock inside the WS Lambda if concurrency is tight
 await enqueueAgentJob({
 sessionId: body.sessionId,
 connectionId,
 prompt: body.prompt,
 domainName: event.requestContext.domainName!,
 stage: event.requestContext.stage,
 });
 await pushToConnection({
 domainName: event.requestContext.domainName!,
 stage: event.requestContext.stage,
 connectionId,
 payload: { type: "ack", sessionId: body.sessionId },
 });
 }

 if (body.type === "agent.interrupt") {
 await flagInterrupt(body.sessionId);
 }

 return { statusCode: 200, body: "ok" };
};

// ✅ Batch deltas
let buf = "";
let timer: NodeJS.Timeout | null = null;
function queueDelta(text: string, send: (s: string) => Promise<void>) {
 buf += text;
 if (!timer) {
 timer = setTimeout(async () => {
 const out = buf;
 buf = "";
 timer = null;
 await send(out);
 }, 50);
 }
}

export async function resolveConnectionId(sessionId: string): Promise<string | null> {
 const out = await ddb.send(
 new UpdateCommand({
 // actually Get — shown as helper
 TableName: TABLE,
 Key: { pk: `session#${sessionId}`, sk: "conn" },
 UpdateExpression: "SET lastSeen = :t",
 ExpressionAttributeValues: { ":t": Date.now() },
 ReturnValues: "ALL_NEW",
 })
 );
 return (out.Attributes?.connectionId as string) ?? null;
}