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

An IDE agent is not a single HTTP request. It is a multi-turn conversation: plan → tool call → patch preview → user interrupt → continue. HTTP with Lambda response streaming is excellent for one-shot token streams. Multi-turn with interrupts and server-pushed tool events wants API Gateway WebSocket APIs, connectionId state in DynamoDB, and the management callback API — not sticky ALB folklore.

⚡ TL;DR: Use API Gateway WebSockets for multi-turn coding agents. Persist connectionId ↔ sessionId in DynamoDB on $connect, clean up on $disconnect, handle messages on default (or typed routes). Push tokens/events with PostToConnection. Design for reconnect: client resumes with sessionId, not the old connectionId. Pair HTTP streaming for short prompts; use WS for sessions. See also SQS FIFO agent queues and Bedrock Converse idempotent tool results.

Route selection: $connect, $disconnect, default

API Gateway WebSocket routes:

• $connect — auth, create session row, store connectionId
• $disconnect — mark offline / TTL; do not delete session history yet
• $default or custom routes (sendmessage, interrupt) — agent turns

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" };
};

Fan-out via the callback API (PostToConnection)

Lambda that generates tokens or finishes a tool must push to the client using the WebSocket management endpoint, not by “returning” from a different connection.

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 };
}

Default route handler sketch:

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" };
};

✅ Put long Bedrock work on SQS/FIFO worker Lambdas; WS Lambdas stay fast. Ordered per-session jobs map cleanly to SQS FIFO + Lambda.

Lambda concurrency and backpressure

Each WS message can invoke Lambda. A chatty IDE + streaming token pushes can amplify concurrency.

• Set reserved concurrency on the WS routers (connect/default) separate from workers.
• Coalesce token deltas (e.g. 50ms batching) before PostToConnection.
• If PostToConnection throttles, buffer in SQS per sessionId — do not block Bedrock generation on WS send.

// ✅ 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);
  }
}

❌ Sticky ALB myths

You do not need sticky sessions on an ALB for API Gateway WebSockets. API Gateway owns the connection; your Lambdas are stateless relative to TCP. State lives in DynamoDB (connectionId, sessionId, conversation transcript pointers).

ALB sticky sessions matter for ALB → target WebSocket architectures you manage yourself. With API Gateway WebSocket APIs, invest in session persistence, not stickiness.

Resume after reconnect with sessionId

Clients drop (laptops sleep). On reconnect:

1. Client opens WS with ?sessionId=...&tenantId=...
2. $connect stores new connectionId against the same sessionId
3. Worker that had the old id looks up session# → conn before each push
4. Replay missed events from an event log (DynamoDB or SQS) keyed by session

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;
}

Use idempotent tool results so a replayed turn does not double-apply patches.

Pair with HTTP response streaming

Do not force every call through WS. Do not force multi-turn IDE sessions through a single buffered HTTP request.

Checklist

• [ ] $connect writes conn# and session# mappings with TTL
• [ ] $disconnect clears connection mapping; keeps session for resume
• [ ] Long Bedrock/tool work on worker queue, not in the WS default Lambda
• [ ] PostToConnection handles 410 Gone by deleting stale conn rows
• [ ] Client resumes with sessionId; never assumes connectionId survives
• [ ] Reserved concurrency split: WS routers vs workers
• [ ] No reliance on ALB sticky myths for API Gateway WS
• [ ] Idempotent tool applies under reconnect replay

Multi-turn coding agents fail when connection identity is confused with session identity. DynamoDB + callback API + resume tokens fix that — stickiness marketing does not.