The Leaky Abstraction (Vol 2): Multi-Turn Agent State Transport and Buffer Disconnects

import { Response } from 'express';
import { Firestore } from '@google-cloud/firestore';

export interface StreamEvent {
  seq: number;
  type: 'token' | 'tool_start' | 'tool_end' | 'done' | 'error';
  payload: unknown;
  timestamp: number;
}

export class AgentSessionStreamGateway {
  private heartbeatTimer?: NodeJS.Timeout;
  private currentSeq: number = 0;
  private eventBuffer: StreamEvent[] = [];
  private readonly MAX_REPLAY_LIMIT = 200;

  constructor(
    private readonly sessionId: string,
    private readonly res: Response,
    private readonly db: Firestore
  ) {}

  /**
   * 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.res.writableEnded) {
        this.res.write(': ping\n\n');
      }
    }, 10_000);
  }

  /**
   * Replays unacknowledged events with bounded limit on client reconnect.
   */
  public async replayFrom(lastEventId: number): Promise<number> {
    const snapshot = await this.db
      .collection('agent_sessions')
      .doc(this.sessionId)
      .collection('events')
      .where('seq', '>', lastEventId)
      .orderBy('seq', 'asc')
      .limit(this.MAX_REPLAY_LIMIT)
      .get();

    for (const doc of snapshot.docs) {
      const event = doc.data() as StreamEvent;
      this.writeSseFrame(event);
      this.currentSeq = Math.max(this.currentSeq, event.seq);
    }

    return this.currentSeq;
  }

  /**
   * Emits tokens instantly to client socket and buffers state for batched persistence.
   */
  public emit(type: StreamEvent['type'], payload: unknown): void {
    this.currentSeq += 1;
    const event: StreamEvent = {
      seq: this.currentSeq,
      type,
      payload,
      timestamp: Date.now(),
    };

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

    // 2. Buffer in memory for batched commit
    this.eventBuffer.push(event);

    // 3. Flush checkpoints immediately on tool execution boundaries
    if (type !== 'token') {
      void this.flushBuffer();
    }
  }

  /**
   * Flushes in-flight event buffer to Firestore using atomic batched writes.
   */
  public async flushBuffer(): Promise<void> {
    if (this.eventBuffer.length === 0) return;

    const eventsToCommit = [...this.eventBuffer];
    this.eventBuffer = [];

    const batch = this.db.batch();
    const sessionRef = this.db.collection('agent_sessions').doc(this.sessionId);

    for (const event of eventsToCommit) {
      const docId = event.seq.toString().padStart(8, '0');
      const docRef = sessionRef.collection('events').doc(docId);
      batch.set(docRef, event);
    }

    await batch.commit();
  }

  private writeSseFrame(event: StreamEvent): void {
    if (this.res.writableEnded) return;
    this.res.write(`id: ${event.seq}\n`);
    this.res.write(`event: ${event.type}\n`);
    this.res.write(`data: ${JSON.stringify(event.payload)}\n\n`);
  }

  public async close(): Promise<void> {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
    }
    await this.flushBuffer();
    if (!this.res.writableEnded) {
      this.res.end();
    }
  }
}