200 Tokens a Second Locked Up My Browser: Client-Side Defense for Agent Streams

// ❌ ANTI-PATTERN: Re-rendering 50-file diffs on every token chunk
onChunk((chunk) => {
  setText((prev) => prev + chunk);
});

import { getLimitedUseToken, type AppCheck } from "firebase/app-check";

const MAX_ATTEMPTS = 3;
const IDLE_TIMEOUT_MS = 10_000; // no byte for 10 s means two 4 s heartbeats never arrived
const EVENT_BOUNDARY = /(?:\r\n|\r(?!\n)|\n){2}/; // blank line; a CR followed by LF is one CRLF

class NonRetryableError extends Error {}

/**
 * Resilient Stream Consumer with Last-Event-ID Resume & App Check Attestation
 */
export class ResilientAgentConsumer {
  constructor(
    private readonly endpoint: string,
    private readonly appCheck: AppCheck
  ) {}

  public async executeStream(
    prompt: string,
    onRenderTick: (text: string) => void,
    turnId: string = crypto.randomUUID(), // pass the same ID to resubmit a failed turn
    signal?: AbortSignal
  ): Promise<void> {
    let lastEventId = ""; // ID of the last complete event, sent back on reconnect
    let text = "";
    let frame = 0;
    const render = () => { frame = 0; onRenderTick(text); };

    for (let attempt = 1; ; attempt++) {
      const controller = new AbortController();
      const abort = () => controller.abort();
      let idleTimer = setTimeout(abort, IDLE_TIMEOUT_MS);

      try {
        // Single-use token: the gateway consumes it, so every attempt fetches a new one
        const { token } = await getLimitedUseToken(this.appCheck);
        const headers: Record<string, string> = {
          "Content-Type": "application/json",
          "X-Firebase-AppCheck": token,
          "X-Stream-Turn-Id": turnId,
        };
        if (lastEventId) headers["Last-Event-ID"] = lastEventId;

        const response = await fetch(this.endpoint, {
          method: "POST",
          headers,
          body: JSON.stringify({ prompt }),
          signal: signal ? AbortSignal.any([signal, controller.signal]) : controller.signal,
        });

        if (!response.ok || !response.body) {
          // A 401, 403 or other 4xx fails the same way again; 408, 429 and 5xx may not
          const retryable = response.status === 408 || response.status === 429 || response.status >= 500;
          const message = `HTTP ${response.status}`;
          throw retryable ? new Error(message) : new NonRetryableError(message);
        }

        // TextDecoderStream holds a split multi-byte character until its last byte arrives
        const reader = response.body.pipeThrough(new TextDecoderStream()).getReader();
        let buffer = "";

        while (true) {
          const { done, value } = await reader.read();
          // A clean EOF without a done event is still a dropped stream
          if (done) throw new Error("Stream ended before the done event");

          clearTimeout(idleTimer); // any byte, heartbeat included, proves the socket is alive
          idleTimer = setTimeout(abort, IDLE_TIMEOUT_MS);
          buffer += value;

          let match: RegExpExecArray | null;
          while ((match = EVENT_BOUNDARY.exec(buffer)) !== null) {
            const event = parseSseEvent(buffer.slice(0, match.index));
            buffer = buffer.slice(match.index + match[0].length);

            if (event.type === "token") {
              text += JSON.parse(event.data);
              frame ||= requestAnimationFrame(render); // at most one render per display frame
            } else if (event.type === "done") {
              cancelAnimationFrame(frame);
              render();
              return;
            } else if (event.type === "error") {
              throw new NonRetryableError(event.data); // the server ended the turn with an error
            }
            // Count an event as received only once it is handled, so a failed parse is not skipped on resume
            if (event.id !== undefined) lastEventId = event.id;
          }
        }
      } catch (err) {
        if (err instanceof NonRetryableError || signal?.aborted || attempt >= MAX_ATTEMPTS) {
          throw new Error(`Agent stream failed after ${attempt} attempt(s): ${err}`);
        }
      } finally {
        clearTimeout(idleTimer);
        abort(); // release the old socket before a retry opens a new one
      }

      // Exponential backoff with jitter, then resume after lastEventId
      const backoffMs = 2 ** attempt * 500 + Math.random() * 200;
      await new Promise((res) => setTimeout(res, backoffMs));
    }
  }
}

function parseSseEvent(block: string): { id?: string; type: string; data: string } {
  let id: string | undefined;
  let type = "message";
  const data: string[] = [];
  for (const line of block.split(/\r\n|\r|\n/)) {
    if (line.startsWith(":")) continue; // :keep-alive heartbeat
    const colon = line.indexOf(":");
    const field = colon === -1 ? line : line.slice(0, colon);
    const value = colon === -1 ? "" : line.slice(colon + 1).replace(/^ /, "");
    if (field === "id") id = value;
    else if (field === "event") type = value;
    else if (field === "data") data.push(value);
  }
  return { id, type, data: data.join("\n") };
}

import type { Request, Response, NextFunction } from "express";
import { getAppCheck, type VerifyAppCheckTokenResponse } from "firebase-admin/app-check";
import { FieldValue, type Firestore, type Transaction } from "firebase-admin/firestore";

declare global {
  namespace Express {
    interface Request { appCheckClaims?: VerifyAppCheckTokenResponse }
  }
}

/**
 * Cloud Run Middleware: App Check Token Verification
 */
export async function verifyAppCheckMiddleware(
  req: Request, 
  res: Response, 
  next: NextFunction
): Promise<void> {
  const appCheckToken = req.header("X-Firebase-AppCheck");

  if (!appCheckToken) {
    res.status(401).json({ error: "Unauthorized: Missing App Check token" });
    return;
  }

  let claims: VerifyAppCheckTokenResponse;
  try {
    // Replay protection: this route runs tools, so each token is accepted once
    claims = await getAppCheck().verifyToken(appCheckToken, { consume: true });
  } catch (err) {
    console.warn("App Check verification failed", err); // also surfaces a misconfigured Admin SDK
    // A bad or expired token fails the same way again; a failed call to the App Check backend may not
    const code = (err as { code?: string }).code ?? "";
    const badToken = code === "app-check/invalid-argument" || code === "app-check/app-check-token-expired";
    res.status(badToken ? 401 : 503).json({ error: badToken ? "Unauthorized: Invalid App Check token" : "App Check unavailable" });
    return;
  }
  if (claims.alreadyConsumed) {
    res.status(401).json({ error: "Unauthorized: App Check token already used" });
    return;
  }
  req.appCheckClaims = claims;
  next();
}

/**
 * Configures Cloud Run Streaming Headers & Heartbeat Keep-Alives
 */
export function setupStreamingHeaders(res: Response): NodeJS.Timeout {
  res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
  res.setHeader("Cache-Control", "no-cache, no-transform");
  res.setHeader("Connection", "keep-alive");
  res.setHeader("X-Accel-Buffering", "no"); // For nginx-style proxies; Cloud Run streams without it
  res.flushHeaders(); // Send headers now, not with the first heartbeat 4 s later

  // Emit SSE keep-alive heartbeat comment frame every 4 seconds
  const heartbeatTimer = setInterval(() => {
    if (!res.writableEnded) {
      res.write(":keep-alive\n\n");
    }
  }, 4000);

  res.on("close", () => clearInterval(heartbeatTimer));
  res.on("finish", () => clearInterval(heartbeatTimer));

  return heartbeatTimer;
}

/**
 * Runs one side-effecting step of a turn at most once per (turnId, step).
 * The step's writes and the record of the step commit in one Firestore transaction.
 * write() must only stage writes on tx: a retried transaction calls it again.
 */
export async function runStepOnce(
  db: Firestore,
  turnId: string,
  step: number,
  write: (tx: Transaction) => void
): Promise<boolean> {
  const turnRef = db.collection("turns").doc(turnId);
  return db.runTransaction(async (tx) => {
    const turn = await tx.get(turnRef);
    if (turn.get(`steps.${step}`) !== undefined) return false; // an earlier attempt committed this step
    write(tx);
    tx.set(turnRef, { steps: { [step]: FieldValue.serverTimestamp() } }, { merge: true });
    return true;
  });
}