RAG/frontend/src/utils/sse.ts

/**
 * SSE 流式读取工具 —— 从 utils.js 原封不动搬移
 *
 * 统一处理 Flux<String> / ServerSentEvent / SseEmitter 三种 SSE 接口。
 * 此文件使用原生 fetch + ReadableStream API，零框架依赖。
 */
import type { SSECallbacks } from '@/types/sse'
import { getToken } from '@/utils/token'

/** 构建带认证的请求头 */
function authHeaders(extra?: Record<string, string>): Record<string, string> {
  const token = getToken()
  const base: Record<string, string> = extra ? { ...extra } : {}
  // 仅当调用方未显式提供 Authorization 时才注入管理后台 token，避免覆盖调用方传入的鉴权头（如测试面板的 SDK Token）
  if (token && !base['Authorization']) base['Authorization'] = `Bearer ${token}`
  return base
}

/** 尝试从 OpenAI Chat Completions chunk 中提取 delta.content；返回 undefined 表示非 OpenAI 格式（回退纯文本） */
function extractOpenAIDelta(text: string): string | undefined {
  let obj: any
  try { obj = JSON.parse(text) } catch { return undefined }
  // OpenAI 错误 chunk（{"error":{"message":...}}）：提取错误信息作为内容，避免原始 JSON 泄漏到对话框
  if (obj && obj.error) {
    const msg = obj.error.message || obj.error.type
    return typeof msg === 'string' && msg ? msg : '服务异常'
  }
  if (obj && Array.isArray(obj.choices)) {
    const delta = obj.choices[0]?.delta
    if (delta && typeof delta.content === 'string') return delta.content
    return '' // OpenAI 形状但无 content（role/finish_reason 空 chunk）→ 跳过
  }
  return undefined // 是 JSON 但非 OpenAI 形状 → 回退纯文本
}

/**
 * 通用 SSE 流式读取 —— 统一处理 Flux<String> / ServerSentEvent / SseEmitter 三种 SSE 接口
 *
 * @param url 请求地址
 * @param onChunk 每收到一段文本的回调
 * @param onDone 流结束的回调
 * @param headers 额外请求头
 * @param signal AbortSignal 用于取消请求（组件卸载时必须传入以释放网络资源）
 */
export async function readSSEStream(
  url: string,
  onChunk: (text: string) => void,
  onDone?: () => void,
  headers?: Record<string, string>,
  signal?: AbortSignal
): Promise<void> {
  const res = await fetch(url, { headers: authHeaders(headers), signal, credentials: 'include' })
  if (!res.ok) throw new Error('HTTP ' + res.status)
  const reader = res.body!.getReader()
  const decoder = new TextDecoder()
  let buffer = ''
  // SSE 规范：同一事件内的多行 data: 字段用 \n 拼接，事件间用空行分隔
  let eventDataLines: string[] = []
  let currentEvent = 'message'

  const flushEvent = () => {
    if (eventDataLines.length === 0) return
    const text = eventDataLines.join('\n')
    eventDataLines = []
    const ev = currentEvent
    currentEvent = 'message'
    if (text === '[DONE]') return
    // 跳过 status / faq 等系统事件，不显示在对话框中
    if (ev === 'status') return
    // OpenAI Chat Completions 格式：提取 delta.content；无 content 的空 chunk 跳过
    const extracted = extractOpenAIDelta(text)
    if (extracted !== undefined) {
      if (extracted) onChunk(extracted)
      return
    }
    // 空事件视为 LLM 流式输出的换行符（Spring 将 "\n" 编码为单条空 data: 事件）
    onChunk(text || '\n')
  }

  while (true) {
    const { done, value } = await reader.read()
    if (done) break
    buffer += decoder.decode(value, { stream: true })
    const lines = buffer.split('\n')
    buffer = lines.pop() || ''
    for (let line of lines) {
      // 兼容 \r\n 行结束符
      if (line.endsWith('\r')) line = line.slice(0, -1)
      if (line === '') {
        // 空行 = SSE 事件边界
        flushEvent()
      } else if (line.startsWith('event:')) {
        // 记录事件类型，用于 flushEvent 时过滤系统事件
        currentEvent = line.slice(6).trim()
      } else if (line.startsWith('data:')) {
        // 累积 data 字段；仅剥离 data: 后的一个可选空格，保留 markdown 列表缩进
        let data = line.slice(5)
        if (data.startsWith(' ')) data = data.slice(1)
        eventDataLines.push(data)
      } else if (!line.startsWith(':')) {
        // Flux<String> 模式（非标准 SSE），先把已累积的 SSE 事件 flush 再处理
        flushEvent()
        if (line.trim()) onChunk(line)
      }
    }
  }
  // 流结束，flush 末尾未以空行收尾的事件
  flushEvent()
  if (onDone) onDone()
}

/**
 * 增强版 SSE 流式读取（支持事件类型分发）
 * 解析 SSE 标准的 event: 字段，将不同类型事件分发到对应回调。
 *
 * @param url 请求地址
 * @param handlers 回调对象：
 *   - onMessage(chunk): 普通文本内容（event: message 或无 event 的 data）
 *   - onToolCallStart(data): 工具调用开始（event: tool_call_start）
 *   - onToolCallResult(data): 工具调用结果（event: tool_call_result）
 *   - onError(data): 错误事件（event: error）
 *   - onDone(): 流结束
 * @param headers 额外请求头
 * @param signal AbortSignal 用于取消请求（组件卸载时必须传入以释放网络资源）
 */
export async function readSSEStreamWithEvents(
  url: string,
  handlers: SSECallbacks,
  headers?: Record<string, string>,
  signal?: AbortSignal
): Promise<void> {
  const { onMessage, onToolCallStart, onToolCallResult, onError, onDone } = handlers
  const res = await fetch(url, { headers: authHeaders(headers), signal, credentials: 'include' })
  if (!res.ok) throw new Error('HTTP ' + res.status)
  const reader = res.body!.getReader()
  const decoder = new TextDecoder()
  let buffer = ''
  let currentEvent = 'message'
  // SSE 规范：同一事件内的多行 data: 字段用 \n 拼接
  let eventDataLines: string[] = []

  const flushEvent = () => {
    if (eventDataLines.length === 0) return
    const ev = currentEvent
    currentEvent = 'message'
    const raw = eventDataLines.join('\n')
    eventDataLines = []
    if (raw === '[DONE]') return
    // 空事件视为换行符（仅 message 事件）；JSON 事件（工具调用等）空内容不应出现
    const text = raw || '\n'

    switch (ev) {
      case 'tool_call_start':
        if (onToolCallStart) onToolCallStart(JSON.parse(text))
        break
      case 'tool_call_result':
        if (onToolCallResult) onToolCallResult(JSON.parse(text))
        break
      case 'error':
        if (onError) onError(JSON.parse(text))
        break
      case 'status':
        // 系统状态事件（generating / faq_hit 等），不显示在对话框中
        break
      case 'message':
      default: {
        // OpenAI Chat Completions 格式：提取 delta.content；无 content 的空 chunk 跳过
        const extracted = extractOpenAIDelta(text)
        if (extracted !== undefined) {
          if (extracted && onMessage) onMessage(extracted)
        } else if (onMessage) {
          onMessage(text)
        }
        break
      }
    }
  }

  while (true) {
    const { done, value } = await reader.read()
    if (done) break
    buffer += decoder.decode(value, { stream: true })
    const lines = buffer.split('\n')
    buffer = lines.pop() || ''

    for (let line of lines) {
      if (line.endsWith('\r')) line = line.slice(0, -1)
      if (line === '') {
        // 空行 = SSE 事件边界
        flushEvent()
      } else if (line.startsWith('event:')) {
        currentEvent = line.slice(6).trim()
      } else if (line.startsWith('data:')) {
        let data = line.slice(5)
        if (data.startsWith(' ')) data = data.slice(1)
        eventDataLines.push(data)
      } else if (!line.startsWith(':')) {
        // Flux<String> 模式（非标准 SSE）
        flushEvent()
        if (line.trim() && onMessage) onMessage(line)
      }
    }
  }
  flushEvent()
  if (onDone) onDone()
}