记一次类 ChatGPT 的网页模块开发

import axios, {
  type AxiosRequestConfig,
  type CancelTokenSource,
  type AxiosResponse,
  type AxiosError,
} from 'axios';

import { axios as axiosInstance } from '@/lib/axios';

// 基础事件类型定义
export interface SSEEventBase {
  type: string;
  target: SimpleAxiosSSE;
  data?: string;
  lastEventId?: string | null;
  retry?: number | null;
}

export interface SSEOpenEvent extends SSEEventBase {
  type: 'open';
}

export interface SSEMessageEvent extends SSEEventBase {
  type: 'message';
  data: string;
}

export interface SSEErrorEvent extends SSEEventBase {
  type: 'error';
  error: Error;
  message: string;
}

export interface SSECloseEvent extends SSEEventBase {
  type: 'close';
}

// 联合类型
export type SSEEvent = SSEOpenEvent | SSEMessageEvent | SSEErrorEvent | SSECloseEvent;

// 事件监听器类型
export type SSEEventListener = (event: SSEEvent) => void;
export type SSEOpenEventListener = (event: SSEOpenEvent) => void;
export type SSEMessageEventListener = (event: SSEMessageEvent) => void;
export type SSEErrorEventListener = (event: SSEErrorEvent) => void;
export type SSECloseEventListener = (event: SSECloseEvent) => void;

// 配置选项接口
export interface SimpleAxiosSSEOptions
  extends Omit<AxiosRequestConfig, 'url' | 'responseType' | 'onDownloadProgress' | 'method'> {
  // 只支持 GET 和 POST 方法
  method?: 'GET' | 'POST';
  autoConnect?: boolean;
  // 重连相关选项
  autoReconnect?: boolean;
  reconnectDelay?: number;
  maxRetries?: number;
  useLastEventId?: boolean;

  // 自定义 lastEventId 的 header 名称
  lastEventIdHeaderName?: string;
}

// 连接状态枚举
export enum ReadyState {
  CONNECTING = 0,
  OPEN = 1,
  CLOSED = 2,
}

const DefaultLastEventIdHeaderName = 'Last-Event-ID';

/**
 * 基于 Axios onDownloadProgress 的简单 SSE 实现
 * 提供类似原生 EventSource 的 API
 */
export class SimpleAxiosSSE {
  // 静态常量
  static readonly CONNECTING = ReadyState.CONNECTING;
  static readonly OPEN = ReadyState.OPEN;
  static readonly CLOSED = ReadyState.CLOSED;

  // 基本属性
  public readonly url: string;
  private readonly options: AxiosRequestConfig;

  // 事件监听器
  private listeners: Record<string, SSEEventListener[]> = {};

  // 连接状态
  public readyState: ReadyState = ReadyState.CONNECTING;

  // 数据缓存和处理
  private buffer = '';
  private lastProcessedIndex = 0;

  // Axios 请求取消器
  private cancelSource: CancelTokenSource | null = null;

  // 重连相关属性
  private autoReconnect: boolean;
  private reconnectDelay: number;
  private maxRetries: number;
  private useLastEventId: boolean;
  private retryCount = 0;
  private reconnectTimer: NodeJS.Timeout | null = null;
  private lastEventId: string | null = null;
  private serverRetryDelay: number | null = null;
  private manualClose = false; // 标记是否为手动关闭

  private lastEventIdHeaderName: string = DefaultLastEventIdHeaderName;

  /**
   * 获取当前重试次数
   */
  getRetryCount(): number {
    return this.retryCount;
  }

  /**
   * 获取最后的事件ID
   */
  getLastEventId(): string | null {
    return this.lastEventId;
  }

  /**
   * 手动设置最后的事件ID
   */
  setLastEventId(eventId: string | null): void {
    this.lastEventId = eventId;
  }

  // 是否已关闭
  private get closed() {
    return this.readyState === ReadyState.CLOSED;
  }

  // 事件处理器
  public onmessage: SSEMessageEventListener | null = null;
  public onopen: SSEOpenEventListener | null = null;
  public onerror: SSEErrorEventListener | null = null;
  public onclose: SSECloseEventListener | null = null;

  public AxiosResponse: AxiosResponse | null = null;

  constructor(url: string, options: SimpleAxiosSSEOptions = {}) {
    this.url = url;
    this.options = {
      method: 'GET',
      headers: {
        Accept: 'text/event-stream',
        ...options.headers,
      },
      timeout: 6 * 60 * 1000, // 默认6分钟超时
      ...options,
    };

    // 初始化重连相关配置
    this.autoReconnect = options.autoReconnect ?? true;
    this.reconnectDelay = options.reconnectDelay ?? 3000;
    this.maxRetries = options.maxRetries ?? 3;
    this.useLastEventId = options.useLastEventId ?? true;

    this.lastEventIdHeaderName = options.lastEventIdHeaderName ?? DefaultLastEventIdHeaderName;

    if (options.autoConnect !== false) {
      this.connect();
    }
  }

  /**
   * 添加事件监听器
   */
  addEventListener(type: string, listener: SSEEventListener): void {
    if (!this.listeners[type]) {
      this.listeners[type] = [];
    }
    if (!this.listeners[type].includes(listener)) {
      this.listeners[type].push(listener);
    }
  }

  /**
   * 移除事件监听器
   */
  removeEventListener(type: string, listener: SSEEventListener): void {
    if (this.listeners[type]) {
      const index = this.listeners[type].indexOf(listener);
      if (index > -1) {
        this.listeners[type].splice(index, 1);
      }
    }
  }

  /**
   * 分发事件
   */
  private dispatchEvent(event: SSEEvent): void {
    // 调用 onXXX 处理器
    const handlerName = `on${event.type}` as keyof this;
    const handler = this[handlerName];
    if (typeof handler === 'function') {
      (handler as SSEEventListener)(event);
    }

    // 调用 addEventListener 注册的监听器
    if (this.listeners[event.type]) {
      this.listeners[event.type].forEach((listener) => {
        try {
          listener(event);
        } catch (error) {
          console.error('Error in event listener:', error);
        }
      });
    }
  }

  private clearReconnectTimer(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
  }

  /**
   * 开始连接
   */
  connect(): void {
    if (this.readyState === ReadyState.OPEN) {
      return;
    }

    // 清除重连定时器
    this.clearReconnectTimer();

    this.readyState = ReadyState.CONNECTING;
    this.cancelSource = axios.CancelToken.source();
    this.manualClose = false;

    // 准备请求头，如果有 lastEventId 则添加
    const headers = { ...this.options.headers };
    if (this.useLastEventId && this.lastEventId) {
      headers[this.lastEventIdHeaderName] = this.lastEventId;
    }

    const requestConfig: AxiosRequestConfig = {
      ...this.options,
      url: this.url,
      headers,
      responseType: 'text',
      cancelToken: this.cancelSource.token,
      onDownloadProgress: (progressEvent) => {
        this.handleDownloadProgress(progressEvent);
      },
    };

    axiosInstance
      .request({ ...requestConfig, url: this.url })
      .then((response: AxiosResponse) => {
        this.AxiosResponse = response;
        if (!this.closed) {
          this.handleComplete(response);
        }
      })
      .catch((error: AxiosError) => {
        if (!this.closed && !axios.isCancel(error)) {
          this.handleError(error);
        }
      });
  }

  /**
   * 处理下载进度，解析 SSE 数据
   */
  private handleDownloadProgress(progressEvent: ProgressEvent): void {
    if (this.closed) return;

    // 获取新接收到的数据
    const target = progressEvent.target as XMLHttpRequest;

    // 首次接收数据时触发 open 事件
    if (this.readyState === ReadyState.CONNECTING) {
      this.readyState = ReadyState.OPEN;
      // 重置重试计数，连接成功
      if (target.status >= 200 && target.status < 300) {
        this.retryCount = 0;
      }
      this.dispatchEvent({
        type: 'open',
        target: this,
      } as SSEOpenEvent);
    }

    if (!target || !target.responseText) return;

    const fullText = target.responseText;
    const newData = fullText.substring(this.lastProcessedIndex);
    this.lastProcessedIndex = fullText.length;

    if (!newData) return;

    // 将新数据添加到缓冲区
    this.buffer += newData;

    // 解析 SSE 数据块
    this.parseSSEData();
  }

  /**
   * 解析 SSE 数据格式
   */
  private parseSSEData(): void {
    // SSE 数据块以双换行符分隔
    const chunks = this.buffer.split(/\r?\n\r?\n/);

    // 保留最后一个可能不完整的块
    this.buffer = chunks.pop() || '';

    chunks.forEach((chunk) => {
      if (chunk.trim()) {
        this.parseSSEChunk(chunk.trim());
      }
    });
  }

  /**
   * 解析单个 SSE 数据块
   */
  private parseSSEChunk(chunk: string): void {
    const lines = chunk.split(/\r?\n/);
    const eventData: {
      id: string | null;
      event: string;
      data: string[];
      retry: number | null;
    } = {
      id: null,
      event: 'message',
      data: [],
      retry: null,
    };

    lines.forEach((line) => {
      if (line.startsWith(':')) {
        // 注释行，忽略
        return;
      }

      const colonIndex = line.indexOf(':');
      let field: string;
      let value: string;

      if (colonIndex === -1) {
        field = line;
        value = '';
      } else {
        field = line.substring(0, colonIndex);
        // 跳过冒号后的第一个空格（如果存在）
        value = line.substring(colonIndex + 1).replace(/^ /, '');
      }

      switch (field) {
        case 'id':
          eventData.id = value;
          break;
        case 'event':
          eventData.event = value;
          break;
        case 'data':
          eventData.data.push(value);
          break;
        case 'retry':
          eventData.retry = parseInt(value, 10);
          break;
      }
    });

    // 保存 lastEventId 用于重连
    if (eventData.id !== null && this.useLastEventId) {
      this.lastEventId = eventData.id;
    }

    // 保存服务器指定的重连延迟时间
    if (eventData.retry !== null) {
      this.serverRetryDelay = eventData.retry;
    }

    // 如果有数据，触发事件
    if (eventData.data.length > 0 || eventData.event !== 'message') {
      const event: SSEEventBase = {
        type: eventData.event,
        target: this,
        data: eventData.data.join('\n'),
        lastEventId: eventData.id,
        retry: eventData.retry,
      };

      this.dispatchEvent(event as SSEEvent);
    }
  }

  /**
   * 处理请求完成
   */
  private handleComplete(response: AxiosResponse): void {
    // 处理剩余的缓冲区数据
    if (this.buffer.trim()) {
      this.parseSSEChunk(this.buffer.trim());
    }

    // 连接正常结束，尝试重连
    this._markClosed();
  }

  /**
   * 处理错误
   */
  private handleError(error: AxiosError): void {
    const errorEvent: SSEErrorEvent = {
      type: 'error',
      target: this,
      error: error as Error,
      message: error.message,
    };
    this.dispatchEvent(errorEvent);

    // 发生错误，尝试重连
    this._markClosed();
  }

  /**
   * 标记连接关闭并处理重连逻辑
   */
  private _markClosed(): void {
    if (this.closed) return;

    // 重置连接状态
    this.readyState = ReadyState.CLOSED;
    this.buffer = '';
    this.lastProcessedIndex = 0;

    if (this.cancelSource) {
      this.cancelSource = null;
    }

    // 触发 close 事件
    this.dispatchEvent({
      type: 'close',
      target: this,
    } as SSECloseEvent);

    // 如果不是手动关闭且启用自动重连，则尝试重连
    if (!this.manualClose && this.autoReconnect) {
      this._attemptReconnect();
    }
  }

  /**
   * 尝试重连
   */
  private _attemptReconnect(): void {
    // 检查是否超过最大重试次数
    if (this.maxRetries > 0 && this.retryCount >= this.maxRetries) {
      console.warn(`SSE 重连失败：超过最大重试次数 ${this.maxRetries}`);
      return;
    }

    this.retryCount++;

    // 使用服务器指定的延迟时间，否则使用配置的延迟时间
    const delay = this.serverRetryDelay ?? this.reconnectDelay;

    console.log(`SSE 将在 ${delay}ms 后尝试第 ${this.retryCount} 次重连`);

    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect();
    }, delay);
  }

  /**
   * 关闭连接
   */
  close(): void {
    if (this.closed) return;

    this.manualClose = true;

    // 清除重连定时器
    this.clearReconnectTimer();

    this.readyState = ReadyState.CLOSED;

    if (this.cancelSource) {
      this.cancelSource.cancel('Connection closed by user');
      this.cancelSource = null;
    }

    this.dispatchEvent({
      type: 'close',
      target: this,
    } as SSECloseEvent);
  }

  /**
   * 强制重连（会重置重试计数）
   */
  reconnect(): void {
    this.retryCount = 0;
    this.manualClose = false;

    if (this.readyState !== ReadyState.CLOSED) {
      this.close();
    }

    // 延迟一下再连接，确保上次连接完全关闭
    setTimeout(() => {
      this.connect();
    }, 100);
  }
}

/**
 * 创建 SimpleAxiosSSE 实例的工厂函数
 */
export function createSimpleAxiosSSE(
  url: string,
  options: SimpleAxiosSSEOptions = {}
): SimpleAxiosSSE {
  const sse = new SimpleAxiosSSE(url, options);

  return sse;
}

export default createSimpleAxiosSSE;