lingji-work-fe/src/utils/querySSE.ts

import { fetchEventSource } from '@microsoft/fetch-event-source';
import type { EventSourceMessage } from '@microsoft/fetch-event-source';
import { useEnterpriseStore } from '@/stores/modules/enterprise';
import { glsWithCatch } from '@/utils/stroage';
import { BASE_PYTHON_URL } from '@/api/all/chat';
import { genRandomId } from '@/utils/tools';

const DEFAULT_SSE_URL = `${BASE_PYTHON_URL}/api/agent/runs`;

const SSE_HEADERS = {
  'Content-Type': 'application/json',
  'Cache-Control': 'no-cache',
  Connection: 'keep-alive',
  Accept: 'text/event-stream',
};

interface SSEConfig {
  headers?: Record<string, string | number>;
  method?: string;
  body?: any;
  handleMessage?: (data: any) => void;
  handleError?: (err: any) => number | null | undefined | void;
  handleClose?: () => void;
  handleOpen?: (response: Response) => void;
}

/**
 * 创建服务器发送事件（SSE）连接
 * @param config SSE 配置
 * @param url 可选的自定义 URL
 * @returns 包含abort方法的对象，用于中断SSE连接
 */
export default async (config: SSEConfig, url: string = DEFAULT_SSE_URL): Promise<{ abort: () => void }> => {
  const {
    body = undefined,
    headers = {},
    method = 'post',
    handleMessage,
    handleError,
    handleOpen,
    handleClose,
  } = config;
  const store = useEnterpriseStore();

  // 创建AbortController实例用于中断请求
  const abortController = new AbortController();

  fetchEventSource(url, {
    method,
    // credentials: 'include',
    headers: {
      ...SSE_HEADERS,
      Authorization: glsWithCatch('accessToken'),
      'enterprise-id': store.enterpriseInfo?.id?.toString(),
      requestid: genRandomId(),
      ...headers,
    },
    body,
    signal: abortController.signal, // 传递signal给fetchEventSource
    openWhenHidden: true, // 用户切换到另一个页面后仍能保持SSE连接
    onmessage(event: EventSourceMessage) {
      if (event.data) {
        try {
          const parsedData = JSON.parse(event.data);
          handleMessage?.({ ...event, data: parsedData });
        } catch (error) {
          console.error('Error parsing SSE message:', error);
          handleError(new Error('Failed to parse SSE message'));
        }
      }
    },
    onerror(error: Error) {
      // 请求失败时主动关闭SSE连接
      abortController.abort();
      handleError?.(error);
    },
    onclose() {
      handleClose?.();
    },
    async onopen(response: Response) {
      handleOpen?.(response);
    },
  });

  // 返回abort方法供外部调用
  return {
    abort: () => abortController.abort(),
  };
};