xo-stack/packages/stack/src/utils/sse-client.ts

// TODO: This probably needs some work.
//       The reason this exists is that Browser's native EventSource is rubbish:
//       1. It does not allow passing of headers.
//       2. It is browser-only (this implementation aims to work for Node also)
//       3. It does not handle disconnects/reconnects gracefully
//       4. It does not handle last-event-id (important for history).
//       5. It does not allow a custom fetch function (probably not necessary, but there's a reason I'd rather support this which I'll explain some other time).
//       This might eventually be shifted out of XO/stack package and into its own.
//       Q: "Why not just leverage an existing package?"
//       A: They're all bad - the "go to" is a Microsoft (@azure/fetch-event-soruce) that is no longer maintained.
//          It does not support NodeJS (or Workers) and we cannot depend upon forks of it.

export interface SSEOptions {
  fetch?: typeof fetch;
  headers?: Record<string, string>;
  onMessage?: (event: SSEMessageEvent) => void;
  onError?: (error: Error) => void;
  onOpen?: () => void;
  onClose?: () => void;
  maxRetries?: number;
  retryDelay?: number;
}

export interface SSEMessageEvent {
  data: string;
  type: string;
  id: string | null;
  lastEventId: string;
  origin: string;
}

export interface SSEEvent {
  data: string;
  event: string;
  id: string | null;
}

export class SSEClient {
  private url: string;
  private options: Required<Omit<SSEOptions, 'headers'>> & {
    headers: Record<string, string>;
  };
  private abortController: AbortController | null;
  private retryCount: number;
  private lastEventId: string;
  private connected: boolean;

  constructor(url: string, options: SSEOptions = {}) {
    this.url = url;
    this.options = {
      fetch: globalThis.fetch,
      headers: {},
      onMessage: () => {},
      onError: () => {},
      onOpen: () => {},
      onClose: () => {},
      maxRetries: 5,
      retryDelay: 3000,
      ...options,
    };

    this.abortController = null;
    this.retryCount = 0;
    this.lastEventId = '';
    this.connected = false;
  }

  public async connect(): Promise<void> {
    if (this.connected) {
      return;
    }

    this.abortController = new AbortController();

    try {
      const response = await this.options.fetch(this.url, {
        headers: {
          Accept: 'text/event-stream',
          'Cache-Control': 'no-cache',
          'Last-Event-ID': this.lastEventId,
          ...this.options.headers,
        },
        signal: this.abortController.signal,
      });

      if (!response.ok) {
        throw new Error(`HTTP error! status: ${response.status}`);
      }

      if (!response.body) {
        throw new Error('ReadableStream not supported');
      }

      this.connected = true;
      this.options.onOpen();

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = '';

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

        if (done) {
          this.handleDisconnect();
          break;
        }

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

        buffer = this.processEventStream(lines);
      }
    } catch (error) {
      if (error instanceof Error && error.name === 'AbortError') {
        return;
      }

      this.options.onError(
        error instanceof Error ? error : new Error(String(error)),
      );
      this.handleDisconnect();

      if (this.retryCount < this.options.maxRetries) {
        this.retryCount++;
        setTimeout(() => this.connect(), this.options.retryDelay);
      }
    }
  }

  private processEventStream(lines: string[]): string {
    let event: SSEEvent = {
      data: '',
      event: 'message',
      id: null,
    };

    for (let i = 0; i < lines.length - 1; i++) {
      const line = lines[i];

      if (line.length === 0) {
        // Empty line means dispatch the event
        if (event.data.length > 0) {
          this.dispatchEvent(event);
          event = {
            data: '',
            event: 'message',
            id: null,
          };
        }
        continue;
      }

      const [fieldName, ...fieldValue] = line.split(':');
      const value = fieldValue.join(':').trim();

      switch (fieldName) {
        case 'event':
          event.event = value;
          break;
        case 'data':
          event.data += value + '\n';
          break;
        case 'id':
          event.id = value;
          this.lastEventId = value;
          break;
        case 'retry': {
          const retryInt = parseInt(value, 10);
          if (!isNaN(retryInt)) {
            this.options.retryDelay = retryInt;
          }
          break;
        }
      }
    }

    // Return the last incomplete line as the new buffer
    return lines[lines.length - 1];
  }

  private dispatchEvent(event: SSEEvent): void {
    const messageEvent: SSEMessageEvent = {
      data: event.data.slice(0, -1), // Remove trailing newline
      type: event.event,
      id: event.id,
      lastEventId: this.lastEventId,
      origin: this.url,
    };

    this.options.onMessage(messageEvent);
  }

  private handleDisconnect(): void {
    this.connected = false;
    this.options.onClose();
  }

  public close(): void {
    if (this.abortController) {
      this.abortController.abort();
      this.handleDisconnect();
    }
  }
}