RAG/frontend/src/utils/sse.ts

/** Shared SSE reader for plain text, OpenAI chunks and tool events. */
import type { SSECallbacks, SourceReference } from '@/types/sse'
import { getToken } from '@/utils/token'

/** Preserve explicit SDK authorization instead of replacing it with the admin token. */
function authHeaders(extra?: Record<string, string>): Record<string, string> {
  const token = getToken()
  const headers = { ...extra }
  if (token && !headers['Authorization']) headers['Authorization'] = `Bearer ${token}`
  return headers
}

function isSourceReference(value: unknown): value is SourceReference {
  if (!value || typeof value !== 'object') return false
  const nullableString = (v: unknown) => v === null || typeof v === 'string'
  const nullableNumber = (v: unknown) => v === null || typeof v === 'number'
  return 'documentId' in value && nullableString(value.documentId)
    && 'title' in value && nullableString(value.title)
    && 'sourceName' in value && nullableString(value.sourceName)
    && 'chunkIndex' in value && nullableNumber(value.chunkIndex)
    && 'score' in value && nullableNumber(value.score)
    && 'snippet' in value && nullableString(value.snippet)
}

/** Returns false only for legacy plain-text payloads. Metadata never enters message text. */
function dispatchOpenAI(text: string, handlers: SSECallbacks): boolean {
  let value: unknown
  try { value = JSON.parse(text) } catch { return false }
  if (!value || typeof value !== 'object') return false
  if ('error' in value && value.error && typeof value.error === 'object') {
    const error = value.error
    const message = 'message' in error ? error.message : 'type' in error ? error.type : undefined
    handlers.onMessage?.(typeof message === 'string' && message ? message : '服务异常')
    return true
  }
  if (!('choices' in value) || !Array.isArray(value.choices)) return false
  if ('sources' in value && Array.isArray(value.sources)) {
    if (!value.sources.every(isSourceReference)) throw new Error('无效的引用来源数据')
    handlers.onSources?.(value.sources)
  }
  const content = value.choices[0]?.delta?.content
  if (typeof content === 'string' && content) handlers.onMessage?.(content)
  return true
}

/** Text-only callers share the same framing, completion and cleanup semantics. */
export function readSSEStream(
  url: string,
  onChunk: (text: string) => void,
  onDone?: () => void,
  headers?: Record<string, string>,
  signal?: AbortSignal,
): Promise<void> {
  return readSSEStreamWithEvents(url, { onMessage: onChunk, onDone }, headers, signal)
}

/** Read complete SSE events; [DONE] terminates immediately without waiting for network EOF. */
export async function readSSEStreamWithEvents(
  url: string,
  handlers: SSECallbacks,
  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((await res.text()) || 'HTTP ' + res.status)
  if (!res.body) throw new Error('响应流为空')
  const reader = res.body.getReader()
  const decoder = new TextDecoder()
  let buffer = ''
  let event = 'message'
  let data: string[] = []
  let completed = false
  let eof = false

  const flushEvent = () => {
    const currentEvent = event
    event = 'message'
    if (!data.length) return
    const raw = data.join('\n')
    data = []
    if (raw === '[DONE]') {
      completed = true
      return
    }
    switch (currentEvent) {
      case 'status': return
      case 'tool_call_start': handlers.onToolCallStart?.(JSON.parse(raw)); return
      case 'tool_call_result': handlers.onToolCallResult?.(JSON.parse(raw)); return
      case 'error': handlers.onError?.(JSON.parse(raw)); return
      default:
        if (!dispatchOpenAI(raw, handlers)) handlers.onMessage?.(raw || '\n')
    }
  }

  const consumeLine = (raw: string) => {
    const line = raw.endsWith('\r') ? raw.slice(0, -1) : raw
    if (!line) {
      flushEvent()
    } else if (line.startsWith('event:')) {
      event = line.slice(6).trim()
    } else if (line === 'data' || line.startsWith('data:')) {
      const value = line === 'data' ? '' : line.slice(5)
      data.push(value.startsWith(' ') ? value.slice(1) : value)
    } else if (!line.startsWith(':') && !line.startsWith('id:') && !line.startsWith('retry:')) {
      // Legacy Flux<String> bodies may contain unframed lines.
      flushEvent()
      if (!completed && line.trim()) handlers.onMessage?.(line)
    }
  }

  try {
    while (!completed) {
      signal?.throwIfAborted()
      const result = await reader.read()
      signal?.throwIfAborted()
      eof = result.done
      buffer += eof ? decoder.decode() : decoder.decode(result.value, { stream: true })
      let start = 0
      let end: number
      while (!completed && (end = buffer.indexOf('\n', start)) !== -1) {
        consumeLine(buffer.slice(start, end))
        start = end + 1
      }
      buffer = buffer.slice(start)
      if (eof) {
        if (!completed && buffer) consumeLine(buffer)
        if (!completed) flushEvent()
        break
      }
    }
    handlers.onDone?.()
  } finally {
    // Cancellation may itself reject or stall; it must not delay DONE or mask the original error.
    if (!eof) {
      try { void reader.cancel().catch(() => {}) } catch { /* Preserve the stream outcome. */ }
    }
    try { reader.releaseLock() } catch { /* Preserve the stream outcome. */ }
  }
}