Genarrative/src/services/image-editor/editorAgentSse.ts

import type { EditorAgentSseEvent } from '../../../packages/shared/src/contracts/editorAgent';
import { readSseJsonStream } from '../sseStream';

const EDITOR_AGENT_SSE_EVENT_NAMES = new Set<EditorAgentSseEvent['event']>([
  'stage',
  'message_delta',
  'tool_started',
  'tool_completed',
  'generation_result',
  'error',
  'done',
]);

function isEditorAgentSseEventName(
  value: unknown,
): value is EditorAgentSseEvent['event'] {
  return (
    typeof value === 'string' &&
    EDITOR_AGENT_SSE_EVENT_NAMES.has(value as EditorAgentSseEvent['event'])
  );
}

function readEventNameFromPayload(parsed: Record<string, unknown>) {
  return isEditorAgentSseEventName(parsed.event) ? parsed.event : null;
}

function readEventDataFromPayload(parsed: Record<string, unknown>) {
  const data = parsed.data;
  return typeof data === 'object' && data !== null
    ? (data as Record<string, unknown>)
    : null;
}

export function normalizeEditorAgentSseEvent(
  eventName: string,
  parsed: Record<string, unknown>,
): EditorAgentSseEvent | null {
  const payloadEventName = readEventNameFromPayload(parsed);
  if (payloadEventName) {
    const payloadData = readEventDataFromPayload(parsed);
    if (!payloadData) {
      return null;
    }
    return {
      event: payloadEventName,
      data: payloadData,
    } as unknown as EditorAgentSseEvent;
  }

  if (!isEditorAgentSseEventName(eventName)) {
    return null;
  }

  return {
    event: eventName,
    data: parsed,
  } as unknown as EditorAgentSseEvent;
}

export async function readEditorAgentSseEvents(
  response: Response,
  onEvent: (event: EditorAgentSseEvent) => void,
) {
  await readSseJsonStream(response, ({ eventName, parsed }) => {
    const event = normalizeEditorAgentSseEvent(eventName, parsed);
    if (event) {
      onEvent(event);
    }
  });
}