Client-Side Runtime Agent Resilience: Defending Multi-Turn Streaming Against Silent Disconnects, Token Drops, and In-Memory Buffer Overflows

import { getToken } from "firebase/app-check";

/**
 * Resilient Stream Consumer with Reconnection Offsets & App Check Attestation
 */
export class ResilientAgentConsumer {
  private resumeOffset = 0;
  private maxRetries = 3;

  constructor(
    private readonly endpoint: string,
    private readonly appCheckInstance: any,
    private readonly turnId: string
  ) {}

  public async executeStream(
    prompt: string,
    onRenderTick: (text: string) => void
  ): Promise<void> {
    let attempt = 0;
    let accumulatedText = "";

    while (attempt < this.maxRetries) {
      try {
        // Fetch fresh App Check token with offline-safe fallback
        let appCheckToken = "";
        try {
          appCheckToken = (await getToken(this.appCheckInstance, false)).token;
        } catch (tokenErr) {
          throw new Error(`App Check attestation failed (offline or unverified): ${tokenErr}`);
        }

        const response = await fetch(this.endpoint, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            "X-Firebase-AppCheck": appCheckToken,
            "X-Stream-Turn-Id": this.turnId,
            "X-Stream-Resume-Offset": String(this.resumeOffset),
          },
          body: JSON.stringify({ prompt, resumeFrom: this.resumeOffset }),
        });

        if (!response.ok || !response.body) {
          throw new Error(`HTTP Transport Failure: ${response.status}`);
        }

        const reader = response.body.getReader();
        const decoder = new TextDecoder("utf-8");
        let lineBuffer = "";

        while (true) {
          const { done, value } = await reader.read();
          if (done) break;

          lineBuffer += decoder.decode(value, { stream: true });
          const lines = lineBuffer.split(/\r?\n/);
          lineBuffer = lines.pop() || "";

          for (const line of lines) {
            // Filter out W3C SSE comment frames (:keep-alive) even when packet fragmented
            if (line.startsWith(":") || !line.trim()) continue;
            accumulatedText += line + "\n";
          }
          this.resumeOffset += value.byteLength;

          // Throttled UI dispatch to prevent client memory GC thrashing
          onRenderTick(accumulatedText);
        }

        // Stream completed successfully
        return;

      } catch (err) {
        attempt++;
        if (attempt >= this.maxRetries) {
          throw new Error(`Agent stream terminated after ${this.maxRetries} retries: ${err}`);
        }

        // Exponential backoff with jitter before resuming from last byte offset
        const backoffMs = Math.pow(2, attempt) * 500 + Math.random() * 200;
        await new Promise((res) => setTimeout(res, backoffMs));
      }
    }
  }
}

import { Request, Response, NextFunction } from "express";
import { getAppCheck } from "firebase-admin/app-check";

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

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

  try {
    const claims = await getAppCheck().verifyToken(appCheckToken);
    (req as any).appCheckClaims = claims;
    next();
  } catch (err) {
    return res.status(401).json({ error: "Unauthorized: Invalid App Check token" });
  }
}

/**
 * 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"); // Disable proxy buffering on Cloud Run / Envoy

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