Streaming over an API Gateway WebSocket API is not the same shape as streaming over HTTP. The process generating tokens does not hold the socket, cannot write to it, and has to address the client by an opaque id through a separate REST call.
The connection is not where the tokens are
With Server-Sent Events or a chunked HTTP response, one process owns the response stream and writes into it as tokens arrive. With API Gateway WebSockets, API Gateway owns the socket. Your integration is invoked per inbound message, returns, and dies. Anything the backend wants to send afterwards goes out through the @connections management API — a normal, signed, request/response HTTPS call, addressed to a connection id.
Everything awkward about this design follows from that one fact. The connection id has to be stored somewhere durable so that a different process can find it. Ordering is your problem, because two concurrent posts to the same connection are two independent HTTP requests. And a disconnect is discovered by a call failing, not by a socket closing under you.
Routing is by routeSelectionExpression, an attribute set on the API itself, which names a JSON property in the inbound message — commonly "$request.body.action". Three route keys are predefined: $connect, called as the connection is established, $disconnect, called when either side goes away, and $default, which catches messages that do not match a route or cannot be evaluated as JSON.
$connect and the connection id
Authorise on $connect, not later. It is the only point in the lifecycle where you have a normal request context to authorise against — query string, headers, an authorizer — and returning a non-2xx status from it refuses the connection outright. After that, all you have on each frame is the connection id.
export const handler = async (event) => {
const { connectionId, requestTimeEpoch } = event.requestContext;
const userId = event.requestContext.authorizer?.principalId;
if (!userId) return { statusCode: 401 };
await ddb.send(new PutItemCommand({
TableName: "ws-connections",
Item: {
connectionId: { S: connectionId },
userId: { S: userId },
// 2h max lifetime plus slack; DynamoDB TTL sweeps the leftovers
expiresAt: { N: String(Math.floor(requestTimeEpoch / 1000) + 8000) },
},
}));
return { statusCode: 200 };
};
The TTL attribute is not optional housekeeping. Because a connection can end without $disconnect ever running — a lost network, an abrupt close — the table accumulates rows that name sockets which no longer exist, and every one of them is a future failed API call. A TTL comfortably beyond the two-hour maximum connection lifetime bounds that growth without you writing a sweeper.
Pushing tokens with @connections
The producing process — a second Lambda, a Step Functions task, a container consuming a model’s stream — posts each chunk to the management endpoint, which is the API’s execute-api hostname with the stage appended:
import {
ApiGatewayManagementApiClient,
PostToConnectionCommand,
} from "@aws-sdk/client-apigatewaymanagementapi";
const mgmt = new ApiGatewayManagementApiClient({
endpoint: "https://abc123.execute-api.us-east-1.amazonaws.com/prod",
});
async function relay(connectionId, stream) {
let buffer = "";
let lastFlush = Date.now();
for await (const chunk of stream) {
buffer += chunk.delta ?? "";
// One POST per token is one signed HTTPS request per token.
if (buffer.length > 200 || Date.now() - lastFlush > 100) {
await mgmt.send(new PostToConnectionCommand({
ConnectionId: connectionId,
Data: JSON.stringify({ type: "delta", text: buffer }),
}));
buffer = "";
lastFlush = Date.now();
}
}
if (buffer) {
await mgmt.send(new PostToConnectionCommand({
ConnectionId: connectionId,
Data: JSON.stringify({ type: "delta", text: buffer }),
}));
}
await mgmt.send(new PostToConnectionCommand({
ConnectionId: connectionId,
Data: JSON.stringify({ type: "done" }),
}));
}
The buffering is the part to keep. A model emitting 60 tokens a second into an unbuffered relay is 60 signed HTTPS requests a second per reader, each one billed and each one counting against the same account-level API Gateway throttle as your ordinary traffic — see throttling limits on AWS API Gateway for how those interact. Flushing on either a character threshold or a 100 ms timer keeps the interface feeling live while cutting the call count by an order of magnitude.
The caller’s IAM role needs execute-api:ManageConnections, which is a different action from the execute-api:Invoke that governs calling the API:
{
"Effect": "Allow",
"Action": "execute-api:ManageConnections",
"Resource": "arn:aws:execute-api:us-east-1:111122223333:abc123/prod/POST/@connections/*"
}
Frame, message and duration limits
AWS documents four numbers that shape this design, and all four are listed as not adjustable at the time of writing: a WebSocket frame size of 32 KB, a message payload of 128 KB, a maximum connection duration of 2 hours, and an idle connection timeout of 10 minutes. A message larger than the frame size must be split across frames.
The status codes API Gateway returns on close tell you which limit you hit, and they are worth handling explicitly in the client. Code 1001 covers both the 10-minute idle timeout and the 2-hour lifetime ceiling — the same code for two very different situations, so the client should reconnect rather than diagnose. 1009 means a message was too big to process. 1003 means binary was sent; binary media types are not supported on WebSocket APIs, which rules out shipping raw audio frames over this transport. 1008 is returned when a client sends too many requests.
These quotas and the 500-new-connections-per-second account limit are from the Amazon API Gateway quotas page at the time of writing. The connection-rate limit is adjustable; the frame, message and duration limits are documented as not.
The 10-minute idle timeout is the one that surprises people, because “idle” means no traffic in either direction. A user reading a long answer sends nothing and, once generation finishes, receives nothing, so a chat session drops in the middle of the reader’s attention span. A client-side ping every few minutes on a custom route resets it. The 2-hour ceiling cannot be reset by anything; design the client to reconnect and re-identify rather than to treat one socket as a session.
GoneException and the reaping problem
When you post to a connection that has closed, the management API returns HTTP 410 and the SDK raises GoneException. This is the normal way to learn about a disconnect, not an error condition, and it should be handled at the call site by deleting the stored row and abandoning the stream:
try {
await mgmt.send(new PostToConnectionCommand({ ConnectionId, Data }));
} catch (err) {
if (err.name === "GoneException") {
await ddb.send(new DeleteItemCommand({
TableName: "ws-connections",
Key: { connectionId: { S: ConnectionId } },
}));
throw new ClientGoneError(); // stop consuming the model stream
}
throw err;
}
Rethrowing rather than swallowing matters more than it looks. If the relay keeps consuming the model’s stream after the reader has gone, you are paying for output tokens that nobody will ever see, for the full length of the generation. Cancelling the upstream request on GoneException is the difference between a disconnect costing nothing and a disconnect costing a full completion — and under a flaky mobile network, disconnects are not rare.
The relay above assumes one stream shape. It is not: providers differ in event framing, in where the text lives on a delta, in how they signal completion, and in how they report a mid-stream error — so a relay written against one provider needs a second parser the day you add a fallback. Multigrid normalises streaming into one event shape across providers, which means the WebSocket side stays a single codepath while the model behind it changes.
Top comments (0)