dify/e2e/features/agent-v2/support/service-api-sse.ts

import type { EventSourceMessage } from 'eventsource-parser'
import { EventSourceParserStream } from 'eventsource-parser/stream'

const SERVICE_API_SSE_MAX_BUFFER_SIZE = 1_048_576
const SERVICE_API_SSE_SUMMARY_LIMIT = 2_000
const SUCCESS_TERMINAL_EVENTS = new Set(['message_end', 'workflow_finished'])

export const SERVICE_API_STREAM_TIMEOUT_MS = 120_000
export const SERVICE_API_RUNTIME_STEP_TIMEOUT_MS = SERVICE_API_STREAM_TIMEOUT_MS + 10_000

export type ServiceApiSseEvent = {
  data: unknown
  event?: string
  id?: string
}

export type ServiceApiSseResult = {
  answer: string
  events: ServiceApiSseEvent[]
}

const isRecord = (value: unknown): value is Record<string, unknown> => (
  Boolean(value) && typeof value === 'object' && !Array.isArray(value)
)

const truncate = (value: string, limit = SERVICE_API_SSE_SUMMARY_LIMIT) => (
  value.length > limit ? `${value.slice(0, limit)}...` : value
)

const parseEventData = (event: EventSourceMessage) => {
  try {
    return JSON.parse(event.data) as unknown
  }
  catch (error) {
    throw new Error(
      `Agent v2 Service API SSE event contained invalid JSON: ${truncate(JSON.stringify(event.data))}`,
      { cause: error },
    )
  }
}

const getEventName = (event: EventSourceMessage, data: unknown) => {
  if (event.event)
    return event.event
  if (isRecord(data) && typeof data.event === 'string')
    return data.event
  return undefined
}

const summarizeEvents = (events: ServiceApiSseEvent[]) => truncate(JSON.stringify(events.map((event) => {
  if (!isRecord(event.data))
    return { event: event.event, data: truncate(String(event.data), 200) }

  const data = event.data
  return {
    event: event.event,
    ...(['code', 'message', 'conversation_id', 'message_id', 'task_id'] as const).reduce<Record<string, unknown>>(
      (summary, key) => {
        if (key in data)
          summary[key] = data[key]
        return summary
      },
      {},
    ),
  }
})))

const createBackendError = (event: ServiceApiSseEvent, events: ServiceApiSseEvent[]) => {
  const data = event.data
  const details = isRecord(data)
    ? (['code', 'message', 'conversation_id', 'message_id', 'task_id'] as const)
        .flatMap(key => key in data ? [`${key}=${JSON.stringify(data[key])}`] : [])
        .join(' ')
    : `data=${JSON.stringify(data)}`

  return new Error(
    `Agent v2 Service API SSE error${details ? `: ${details}` : ''}. Received events: ${summarizeEvents(events)}`,
  )
}

export async function consumeServiceApiSse(
  body: ReadableStream<BufferSource> | null,
): Promise<ServiceApiSseResult> {
  if (!body)
    throw new Error('Agent v2 Service API SSE response did not expose a readable body.')

  const events: ServiceApiSseEvent[] = []
  const answers: string[] = []
  const stream = body
    .pipeThrough(new TextDecoderStream())
    .pipeThrough(new EventSourceParserStream({
      maxBufferSize: SERVICE_API_SSE_MAX_BUFFER_SIZE,
      onError: 'terminate',
    }))

  try {
    for await (const message of stream) {
      const data = parseEventData(message)
      const eventName = getEventName(message, data)
      const event: ServiceApiSseEvent = {
        data,
        ...(eventName ? { event: eventName } : {}),
        ...(message.id ? { id: message.id } : {}),
      }
      events.push(event)

      if (isRecord(data) && typeof data.answer === 'string')
        answers.push(data.answer)

      if (eventName === 'error')
        throw createBackendError(event, events)

      if (eventName && SUCCESS_TERMINAL_EVENTS.has(eventName)) {
        return {
          answer: answers.join(''),
          events,
        }
      }
    }
  }
  catch (error) {
    if (error instanceof Error && error.message.startsWith('Agent v2 Service API SSE'))
      throw error

    const message = error instanceof Error ? error.message : String(error)
    throw new Error(
      `Agent v2 Service API SSE parsing failed: ${message}. Received events: ${summarizeEvents(events)}`,
      { cause: error },
    )
  }

  throw new Error(
    `Agent v2 Service API SSE stream closed before a terminal event. Received events: ${summarizeEvents(events)}`,
  )
}