Lambda Response Streaming: Keep AI Coding Gateways Under Client Timeouts

0 views

Your AI coding gateway sits behind API Gateway HTTP API or an ALB. The model needs 40–90 seconds to stream a multi-file patch. The client (IDE extension, Cursor-like agent, browser SSE) has a 29–60s idle timeout. BUFFERED invoke mode holds the entire payload until Lambda returns — so the gateway kills the connection while Bedrock is still generating. The unfair advantage is not “enable streaming” as a checkbox. It is RESPONSE_STREAM with early first-byte, cancel propagation, and metering that still works when you sample logs.

⚡ TL;DR: Use Lambda response streaming (InvokeMode: RESPONSE_STREAM) for AI gateways that must beat client idle timeouts. Stream via awslambda.streamifyResponse in Node, flush headers/first token fast, and treat the stream as at-least-once from the client’s view (reconnect + resume tokens). Pair with Lambda Invoke Modes: BUFFERED Versus RESPONSE_STREAM and Streaming RAG Pipelines. Never stream secrets or raw tool dumps; gate with schema contracts from AI Coding Agent Tool Schemas.

Why BUFFERED kills coding gateways

API Gateway HTTP API max integration timeout is 30s for many setups; ALB idle timeout defaults to 60s. A coding agent that waits for a full Converse response will hit that wall on any non-trivial prompt. BUFFERED mode also inflates memory: you buffer the entire model output in the Lambda process before flush.

// ❌ BUFFERED: client sees nothing until full completion
import {
  BedrockRuntimeClient,
  ConverseCommand,
} from "@aws-sdk/client-bedrock-runtime";

const bedrock = new BedrockRuntimeClient({});

export const handler = async (event: { prompt: string }) => {
  const out = await bedrock.send(
    new ConverseCommand({
      modelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
      messages: [{ role: "user", content: [{ text: event.prompt }] }],
    })
  );
  // Gateway often already timed out waiting for this return
  return {
    statusCode: 200,
    body: JSON.stringify({ text: out.output?.message?.content?.[0]?.text }),
  };
};

✅ Fix: stream tokens as they arrive so time-to-first-byte (TTFB) is seconds, not the full generation window.

Wire RESPONSE_STREAM in Node.js

Lambda response streaming requires a Function URL (or supported invoke path) with InvokeMode: RESPONSE_STREAM, and Node’s awslambda.streamifyResponse wrapper. Prefer Function URLs with IAM auth or JWT at the edge for coding gateways — do not leave them public.

// handler.ts — Node 20 Lambda with response streaming
import {
  BedrockRuntimeClient,
  ConverseStreamCommand,
} from "@aws-sdk/client-bedrock-runtime";

declare const awslambda: {
  streamifyResponse: (
    fn: (
      event: unknown,
      responseStream: NodeJS.WritableStream,
      context: unknown
    ) => Promise<void>
  ) => unknown;
  HttpResponseStream: {
    from: (
      stream: NodeJS.WritableStream,
      metadata: { statusCode: number; headers: Record<string, string> }
    ) => NodeJS.WritableStream;
  };
};

const bedrock = new BedrockRuntimeClient({});

export const handler = awslambda.streamifyResponse(
  async (event: any, responseStream, _context) => {
    const httpStream = awslambda.HttpResponseStream.from(responseStream, {
      statusCode: 200,
      headers: {
        "Content-Type": "text/event-stream; charset=utf-8",
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",
      },
    });

    const write = (obj: unknown) => {
      httpStream.write(`data: ${JSON.stringify(obj)}\n\n`);
    };

    try {
      // ✅ Flush a heartbeat immediately — beats idle timeouts
      write({ type: "session.start", ts: Date.now() });

      const body = typeof event.body === "string" ? JSON.parse(event.body) : event;
      const stream = await bedrock.send(
        new ConverseStreamCommand({
          modelId: body.modelId ?? "anthropic.claude-3-5-sonnet-20241022-v2:0",
          messages: [
            { role: "user", content: [{ text: String(body.prompt ?? "") }] },
          ],
          inferenceConfig: { maxTokens: 4096, temperature: 0.2 },
        })
      );

      let inputTokens = 0;
      let outputTokens = 0;

      for await (const evt of stream.stream ?? []) {
        if (evt.contentBlockDelta?.delta?.text) {
          write({ type: "delta", text: evt.contentBlockDelta.delta.text });
        }
        if (evt.metadata?.usage) {
          inputTokens = evt.metadata.usage.inputTokens ?? inputTokens;
          outputTokens = evt.metadata.usage.outputTokens ?? outputTokens;
        }
        if (evt.messageStop) {
          write({
            type: "done",
            stopReason: evt.messageStop.stopReason,
            usage: { inputTokens, outputTokens },
          });
        }
      }
    } catch (err: any) {
      // ✅ Surface errors on the stream — don't hang the client
      write({ type: "error", message: err?.name ?? "StreamError" });
    } finally {
      httpStream.end();
    }
  }
);

Infra sketch (CDK-ish):

// ✅ Function URL with RESPONSE_STREAM
const fn = new lambda.Function(this, "CodingGateway", {
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: "handler.handler",
  timeout: cdk.Duration.minutes(5),
  memorySize: 1024,
  code: lambda.Code.fromAsset("dist"),
});

const url = fn.addFunctionUrl({
  authType: lambda.FunctionUrlAuthType.AWS_IAM,
  invokeMode: lambda.InvokeMode.RESPONSE_STREAM,
  cors: {
    allowedOrigins: ["https://app.cheatcoders.net"],
    allowedMethods: [lambda.HttpMethod.POST],
    allowedHeaders: ["content-type", "authorization"],
  },
});

❌ Common failure: creating a Function URL but leaving InvokeMode at BUFFERED — you still buffer everything.

Backpressure, cancel, and reconnect

Streaming does not mean fire-and-forget. IDE clients disconnect; users hit stop; ALB still has idle timeouts if you stop writing. Design for:

  1. Heartbeat frames every 5–10s while waiting on tool rounds.
  2. Cancel via client close → abort Bedrock stream (AbortController with SDK middleware where available).
  3. Resume tokens for multi-turn coding sessions so reconnect does not re-run expensive retrieval.
// ✅ Heartbeat while tools run (pseudo — wrap tool awaits)
async function withHeartbeats(
  write: (o: unknown) => void,
  work: () => Promise<void>,
  everyMs = 8000
) {
  const t = setInterval(() => write({ type: "ping", ts: Date.now() }), everyMs);
  try {
    await work();
  } finally {
    clearInterval(t);
  }
}

For long async tools (CDK synth, integration tests), do not hold the HTTP stream open for minutes — offload like Bedrock Agents Plus EventBridge and stream a correlation id + status channel instead.

Metering and security on the wire

Every stream must emit usage for billing and abuse detection. Attach tenantId / workspaceId to the session start frame and emit EMF (see today’s EMF post) on done. Never stream:

  • Raw AWS credentials from tool results
  • Full .env dumps from “read file” tools
  • Unredacted customer source beyond the allowed path scope
// ✅ Redact before write
function safeDelta(text: string): string {
  return text
    .replace(/AKIA[0-9A-Z]{16}/g, "[REDACTED_AKIA]")
    .replace(/(?<=Bearer\s+)[A-Za-z0-9._\-]+/g, "[REDACTED_TOKEN]");
}

Also warm the function path — cold starts still delay TTFB. See Python Lambda Warm Strategies for the same discipline on Node: lean imports, provisioned concurrency for the gateway alias.

Observability for streams (not just final status)

BUFFERED Lambdas make Duration and Max Memory Used enough. Streams need TTFB, bytes flushed, and abort rates. Emit custom metrics when you write the first delta and when the client disconnects mid-stream.

// ✅ Cheap counters alongside SSE frames
let firstByteMs: number | null = null;
const started = Date.now();

function onFirstDelta() {
  if (firstByteMs == null) {
    firstByteMs = Date.now() - started;
    console.log(
      JSON.stringify({
        _aws: {
          Timestamp: Date.now(),
          CloudWatchMetrics: [
            {
              Namespace: "CheatCoders/LLM",
              Dimensions: [["Operation", "Environment"]],
              Metrics: [{ Name: "StreamTTFBMs", Unit: "Milliseconds" }],
            },
          ],
        },
        Operation: "ConverseStream",
        Environment: process.env.STAGE ?? "dev",
        StreamTTFBMs: firstByteMs,
      })
    );
  }
}

Alarm on rising StreamClientAbort more than on raw Duration — aborts often mean the IDE gave up while you were still billing Bedrock. Tie abort handling to cancel so you stop paying for orphan generations.

Client contract: SSE events your gateway should document

Document a minimal event taxonomy so every IDE plugin and internal agent speaks the same dialect: session.start, delta, tool.start / tool.end, ping, done, error. Version it (schemaVersion: 1) in session.start so you can add fields without breaking old clients. If you must change semantics, bump the version and keep dual-publish for one release — the same discipline as additive EventBridge schemas.

Checklist

  • [ ] Function URL (or supported path) uses InvokeMode: RESPONSE_STREAM
  • [ ] Handler wrapped with awslambda.streamifyResponse + early first frame
  • [ ] SSE/NDJSON frames include session.start, delta, done, error, ping
  • [ ] Heartbeats keep idle timeouts alive during tool waits
  • [ ] Client disconnect aborts model stream (no orphan Bedrock spend)
  • [ ] Usage + tenant dims emitted for cost; secrets redacted on the wire
  • [ ] Long tools offloaded async — stream does not hold for minutes
  • [ ] Auth on Function URL (IAM/JWT); CORS locked to app origins

Discover more from CheatCoders

Subscribe to get the latest posts sent to your email.

Comments

No comments yet. Why don’t you start the discussion?

Leave a comment

No account needed. Name and email are optional.