Genarrative/src/services/sseStream.ts

export type SseStreamEvent = {
  eventName: string;
  data: string;
};

export type SseJsonStreamEvent = SseStreamEvent & {
  parsed: Record<string, unknown>;
};

type SseEventBoundary = {
  index: number;
  length: number;
};

type SseStreamEventHandler<TEvent extends SseStreamEvent> = (
  event: TEvent,
) => void | boolean;

function findSseEventBoundary(buffer: string): SseEventBoundary | null {
  const lfBoundary = buffer.indexOf('\n\n');
  const crlfBoundary = buffer.indexOf('\r\n\r\n');

  if (lfBoundary === -1 && crlfBoundary === -1) {
    return null;
  }

  if (lfBoundary === -1) {
    return {
      index: crlfBoundary,
      length: 4,
    };
  }

  if (crlfBoundary === -1 || lfBoundary < crlfBoundary) {
    return {
      index: lfBoundary,
      length: 2,
    };
  }

  return {
    index: crlfBoundary,
    length: 4,
  };
}

function parseSseEventBlock(eventBlock: string): SseStreamEvent | null {
  let eventName = 'message';
  const dataLines: string[] = [];

  for (const rawLine of eventBlock.split(/\r?\n/u)) {
    const line = rawLine.trim();

    if (line.startsWith('event:')) {
      eventName = line.slice(6).trim() || 'message';
      continue;
    }

    if (line.startsWith('data:')) {
      dataLines.push(line.slice(5).trim());
    }
  }

  const data = dataLines.join('\n');
  if (!data) {
    return null;
  }

  return {
    eventName,
    data,
  };
}

export function parseSseJsonObject(data: string): Record<string, unknown> | null {
  try {
    const parsed = JSON.parse(data) as unknown;
    return typeof parsed === 'object' && parsed !== null
      ? (parsed as Record<string, unknown>)
      : null;
  } catch {
    return null;
  }
}

export async function readSseStream(
  response: Response,
  onEvent: SseStreamEventHandler<SseStreamEvent>,
) {
  const streamBody = response.body;
  if (!streamBody) {
    throw new Error('streaming response body is unavailable');
  }

  const reader = streamBody.getReader();
  const decoder = new TextDecoder('utf-8');
  let buffer = '';
  let shouldContinue = true;
  let completed = false;

  const consumeBuffer = () => {
    for (;;) {
      if (!shouldContinue) {
        break;
      }

      const boundary = findSseEventBoundary(buffer);
      if (!boundary) {
        break;
      }

      const eventBlock = buffer.slice(0, boundary.index);
      buffer = buffer.slice(boundary.index + boundary.length);
      const event = parseSseEventBlock(eventBlock);
      if (!event) {
        continue;
      }

      if (onEvent(event) === false) {
        shouldContinue = false;
      }
    }
  };

  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) {
        break;
      }

      buffer += decoder.decode(value, { stream: true });
      consumeBuffer();
      if (!shouldContinue) {
        break;
      }
    }

    if (shouldContinue) {
      // 流结束后 flush 解码器，避免 UTF-8 多字节字符残留在内部缓冲里。
      buffer += decoder.decode();
      consumeBuffer();
      completed = true;
    }
  } finally {
    if (!completed && typeof reader.cancel === 'function') {
      await reader.cancel().catch(() => {});
    }
    reader.releaseLock?.();
  }
}

export function readSseJsonStream(
  response: Response,
  onEvent: SseStreamEventHandler<SseJsonStreamEvent>,
) {
  return readSseStream(response, (event) => {
    const parsed = parseSseJsonObject(event.data);
    if (!parsed) {
      return;
    }

    return onEvent({
      ...event,
      parsed,
    });
  });
}