Dropped Tokens: Fixing Multi-Turn Agent Streams That Die Mid-Flight

import type { Response } from 'express';
import { Timestamp, type DocumentReference, type Firestore } from '@google-cloud/firestore';

export interface StreamEvent {
  seq: number;
  type: 'token' | 'tool_start' | 'tool_end' | 'done' | 'error';
  data: string; // JSON-encoded payload, stored and replayed byte for byte
  timestamp: number;
}

export class LeaseLostError extends Error {}

export class AgentSessionStreamGateway {
  private heartbeatTimer?: NodeJS.Timeout;
  private flushTimer?: NodeJS.Timeout;
  private currentSeq: number = 0;
  private replayed = false;
  private eventBuffer: { ref: DocumentReference; event: StreamEvent }[] = [];
  private flushChain: Promise<void> = Promise.resolve();
  private readonly REPLAY_PAGE_SIZE = 200;
  private readonly TOKEN_FLUSH_MS = 500;
  private readonly JOURNAL_TTL_MS = 24 * 60 * 60 * 1000;

  constructor(
    private readonly sessionId: string,
    private readonly res: Response,
    private readonly db: Firestore,
    private readonly ownerId: string // the ID the caller wrote to leaseOwner when it took the lease
  ) {}

  /**
   * Initializes SSE response headers and begins periodic keep-alive pings.
   */
  public initHeaders(): void {
    this.res.setHeader('Content-Type', 'text/event-stream');
    this.res.setHeader('Cache-Control', 'no-cache, no-transform');
    this.res.setHeader('Connection', 'keep-alive');
    this.res.setHeader('X-Accel-Buffering', 'no');
    this.res.flushHeaders();

    // Emit an SSE comment ping every 10 seconds to prevent proxy timeouts
    this.heartbeatTimer = setInterval(() => {
      if (this.isOpen()) {
        this.res.write(': ping\n\n');
      }
    }, 10_000);
    // Stop pinging the moment the client disconnects
    this.res.on('close', () => clearInterval(this.heartbeatTimer));

    // Persist buffered tokens every 500 ms, not only at tool boundaries
    this.flushTimer = setInterval(() => this.flushInBackground(), this.TOKEN_FLUSH_MS);
  }

  /**
   * Replays every event after lastEventId, one page at a time, and continues the
   * sequence from there so a new event never reuses an ID the client already has.
   * Call it once before emit(), with 0 for a new session.
   */
  public async replayFrom(lastEventId: number, follow = false): Promise<number> {
    let cursor = Number.isSafeInteger(lastEventId) && lastEventId > 0 ? lastEventId : 0;
    let finished = false;

    while (true) {
      const page = await this.eventsRef()
        .where('seq', '>', cursor)
        .orderBy('seq', 'asc')
        .limit(this.REPLAY_PAGE_SIZE)
        .get();

      for (const doc of page.docs) {
        const event = doc.data() as StreamEvent;
        this.writeSseFrame(event);
        cursor = event.seq;
        if (event.type === 'done' || event.type === 'error') finished = true;
      }
      if (page.size === this.REPLAY_PAGE_SIZE) continue;
      // follow: another instance owns the turn, so keep reading its journal until done
      if (!follow || finished || !this.isOpen()) break;
      await new Promise((resolve) => setTimeout(resolve, this.TOKEN_FLUSH_MS));
    }

    this.currentSeq = cursor;
    this.replayed = true;
    return cursor;
  }

  /**
   * Emits tokens instantly to client socket and buffers state for batched persistence.
   */
  public emit(type: StreamEvent['type'], payload: unknown): Promise<void> {
    if (!this.replayed) throw new Error('Call replayFrom() before emit()');
    this.currentSeq += 1;
    const event: StreamEvent = {
      seq: this.currentSeq,
      type,
      data: JSON.stringify(payload ?? null),
      timestamp: Date.now(),
    };

    // 1. Flush immediately to client socket (zero latency penalty on streaming)
    this.writeSseFrame(event);

    // 2. Buffer in memory for batched commit; the auto-ID is fixed here so a retry rewrites the same doc
    this.eventBuffer.push({ ref: this.eventsRef().doc(), event });

    // 3. Checkpoint tool boundaries now; await emit('tool_start') before the tool runs
    if (type === 'token') return Promise.resolve();
    const checkpoint = this.flushBuffer();
    checkpoint.catch((error) => console.error(`Journal flush failed for session ${this.sessionId}`, error));
    return checkpoint;
  }

  /**
   * Flushes in-flight event buffer to Firestore in one transaction that also checks the lease.
   * Commits run one at a time; a failed commit goes back into the buffer unless the lease was lost.
   */
  public flushBuffer(): Promise<void> {
    const run = this.flushChain.then(() => this.commitBuffered());
    this.flushChain = run.catch(() => {}); // one failure must not block later flushes
    return run;
  }

  private async commitBuffered(): Promise<void> {
    const pending = this.eventBuffer.splice(0);
    if (pending.length === 0) return;

    try {
      // The lease check and the writes commit together, so a container that lost the lease writes nothing
      await this.db.runTransaction(async (tx) => {
        const session = await tx.get(this.db.collection('agent_sessions').doc(this.sessionId));
        if (session.get('leaseOwner') !== this.ownerId) throw new LeaseLostError(`Lost the lease on ${this.sessionId}`);
        for (const { ref, event } of pending) {
          // A TTL policy on expireAt (exempt from indexing) deletes each event after it expires
          tx.set(ref, { ...event, expireAt: Timestamp.fromMillis(event.timestamp + this.JOURNAL_TTL_MS) });
        }
      });
    } catch (error) {
      if (!(error instanceof LeaseLostError)) this.eventBuffer.unshift(...pending); // retried on the next flush
      throw error;
    }
  }

  private flushInBackground(): void {
    this.flushBuffer().catch((error) => {
      console.error(`Journal flush failed for session ${this.sessionId}; will retry`, error);
    });
  }

  private writeSseFrame(event: StreamEvent): void {
    if (!this.isOpen()) return;
    // JSON.stringify never emits a raw newline, so one data: line carries the payload
    this.res.write(`id: ${event.seq}\nevent: ${event.type}\ndata: ${event.data}\n\n`);
  }

  private isOpen(): boolean {
    // After a client disconnect, writableEnded stays false and destroyed turns true
    return !this.res.writableEnded && !this.res.destroyed;
  }

  private eventsRef() {
    return this.db.collection('agent_sessions').doc(this.sessionId).collection('events');
  }

  public async close(): Promise<void> {
    clearInterval(this.heartbeatTimer);
    clearInterval(this.flushTimer);
    try {
      // Queued behind any in-flight commit, and retried: nothing else will store the final events
      for (let attempt = 1; ; attempt++) {
        try {
          await this.flushBuffer();
          break;
        } catch (error) {
          if (attempt >= 5 || error instanceof LeaseLostError) throw error;
          await new Promise((resolve) => setTimeout(resolve, 250 * 2 ** attempt));
        }
      }
    } finally {
      if (!this.res.writableEnded) {
        this.res.end();
      }
    }
  }
}