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

// TODO: 1. Handle reconnection logic.
//       2. Allow specifying custom
//       3. Try to make it so that connect DOES NOT resolve until connection is actually establish.
//          There's a very slim chance of a race-condition otherwise where we might miss an event.
//          Note that this IS NOT trivial. Claude will probably NOT be able to do this for you!

/**
 * A Server-Sent Events client implementation using fetch API.
 * Supports custom headers, POST requests, and is non-blocking.
 */
export class SSEClient {
  /**
   * Creates and connects a new SSEClient instance.
   * @param url The URL to connect to
   * @param options Configuration options
   * @returns A new connected SSEClient instance
   */
  public static async from(
    url: string,
    options: SSEClientOptions,
  ): Promise<SSEClient> {
    const client = new SSEClient(url, options);
    await client.connect();
    return client;
  }

  // State.
  private url: string;
  private options: SSEClientOptions;
  private controller: AbortController;
  private connected: boolean = false;

  /**
   * Creates a new SSEClient instance.
   * @param url The URL to connect to
   * @param options Configuration options
   */
  constructor(url: string, options: SSEClientOptions) {
    this.url = url;
    this.options = {
      method: 'GET',
      headers: {
        Accept: 'text/event-stream',
        'Cache-Control': 'no-cache',
      },
      onOpen: () => {},
      onMessage: () => {},
      onError: (error) => console.error('SSEClient error:', error),
      onClose: () => {},
      ...options,
    };
    this.controller = new AbortController();
  }

  /**
   * Connects to the SSE endpoint.
   *
   * @remarks
   *
   * This currently does not await properly and has a chance of introducing a race-condition.
   * We want to wait for the connection to ESTABLISH before resolving.
   * Otherwise, there is a slight possibility of missing events during the connection phase.
   * If you try to fix this, be cautious about using a while(true) loop: JS is single threaded and you may block.
   * Try to use the Streams event features instead.
   */
  public async connect(): Promise<void> {
    if (this.connected) return;

    this.connected = true;
    const { method, headers, body } = this.options;

    const fetchOptions: RequestInit = {
      method,
      headers,
      body,
      signal: this.controller.signal,
      cache: 'no-store',
    };

    // Using an async IIFE to handle the connection without blocking
    (async () => {
      try {
        // TODO: We want to resolve only AFTER we handle connection.
        //       But this is not trivial!
        const response = await fetch(this.url, fetchOptions);

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

        if (!response.body) {
          throw new Error('Response body is null');
        }

        // Notify that the connection has been established
        if (this.options.onOpen) {
          this.options.onOpen();
        }

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

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

          if (done) {
            this.connected = false;
            if (this.options.onClose) {
              this.options.onClose();
            }
            break;
          }

          // Decode the received chunk and add it to the buffer
          const chunk = decoder.decode(value, { stream: true });
          buffer += chunk;

          // Process any complete events in the buffer
          const events = this.parseEvents(buffer);
          buffer = events.remainder;

          // Handle each complete event
          for (const event of events.parsed) {
            if (this.options.onMessage) {
              this.options.onMessage(event);
            }
          }
        }
      } catch (error) {
        this.connected = false;

        if (this.options.onError) {
          this.options.onError(error);
        }

        if (this.options.onClose) {
          this.options.onClose();
        }
      }
    })();
  }

  /**
   * Parses SSE events from a text buffer.
   * @param buffer The buffer to parse
   * @returns Parsed events and the remaining buffer
   */
  private parseEvents(buffer: string): {
    parsed: SSEvent[];
    remainder: string;
  } {
    const events: SSEvent[] = [];
    const lines = buffer.split(/\r\n|\r|\n/);

    let currentEvent: Partial<SSEvent> = {};
    let completeEventCount = 0;

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

      // Empty line signals the end of an event
      if (line === '') {
        if (currentEvent.data) {
          // Remove trailing newline if present
          currentEvent.data = currentEvent.data.replace(/\n$/, '');
          events.push(currentEvent as SSEvent);
          currentEvent = {};
          completeEventCount = i + 1;
        }
        continue;
      }

      // Parse field: value format
      const colonIndex = line.indexOf(':');
      if (colonIndex === -1) continue;

      const field = line.slice(0, colonIndex);
      // Skip initial space after colon if present
      const valueStartIndex =
        colonIndex + 1 + (line[colonIndex + 1] === ' ' ? 1 : 0);
      const value = line.slice(valueStartIndex);

      if (field === 'data') {
        currentEvent.data = currentEvent.data
          ? currentEvent.data + '\n' + value
          : value;
      } else if (field === 'event') {
        currentEvent.event = value;
      } else if (field === 'id') {
        currentEvent.id = value;
      } else if (field === 'retry') {
        const retryMs = parseInt(value, 10);
        if (!isNaN(retryMs)) {
          currentEvent.retry = retryMs;
        }
      }
    }

    const remainder = lines.slice(completeEventCount).join('\n');
    return { parsed: events, remainder };
  }

  /**
   * Closes the SSE connection.
   */
  public close(): void {
    if (!this.connected) return;

    this.controller.abort();
    this.connected = false;
  }

  /**
   * Checks if the client is currently connected.
   * @returns Whether the client is connected
   */
  public isConnected(): boolean {
    return this.connected;
  }
}

/**
 * Configuration options for the SSEClient.
 */
export interface SSEClientOptions {
  /**
   * HTTP method to use (GET or POST).
   */
  method?: 'GET' | 'POST';

  /**
   * HTTP headers to send with the request.
   */
  headers?: Record<string, string>;

  /**
   * Body to send with POST requests.
   */
  body?: string | FormData | URLSearchParams;

  /**
   * Called when the connection is established.
   */
  onOpen?: () => void;

  /**
   * Called when a message is received.
   */
  onMessage?: (event: SSEvent) => void;

  /**
   * Called when an error occurs.
   */
  onError?: (error: unknown) => void;

  /**
   * Called when the connection is closed.
   */
  onClose?: () => void;
}

/**
 * Represents a Server-Sent Event.
 */
export interface SSEvent {
  /**
   * Event data.
   */
  data: string;

  /**
   * Event type.
   */
  event?: string;

  /**
   * Event ID.
   */
  id?: string;

  /**
   * Reconnection time in milliseconds.
   */
  retry?: number;
}