gongxue-base/apps/admin/src/components/AiChat/sseReducer.ts

import type {
  AiArtifactSchema,
  AiAttachment,
  AiChartSchema,
  AiChatMessage,
  AiFormSchema,
  AiImportPreflight,
  AiModelRetryInfo,
  AiReviewSchema,
  AiSseChunk,
  AiToolRun,
} from './types';
import { mergeArtifactIntoMessage, mergeById, mergeForms } from './uiArtifacts';

export interface AiSsePayload {
  messageId?: number;
  userMessageId?: number;
  assistantMessageId?: number;
  delta?: string;
  content?: string;
  reasoningContent?: string | null;
  toolCallId?: string;
  toolName?: string;
  skillKey?: string | null;
  status?: string;
  summary?: string | null;
  durationMs?: number | null;
  attachment?: AiAttachment;
  form?: AiFormSchema;
  artifact?: AiArtifactSchema;
  review?: AiReviewSchema;
  chart?: AiChartSchema;
  preflight?: AiImportPreflight;
  wizard?: unknown;
  retry?: AiModelRetryInfo;
  message?:
    | string
    | {
        id?: number;
        content?: string;
        reasoningContent?: string | null;
        status?: string;
        toolRuns?: AiToolRun[];
        attachments?: AiAttachment[];
        replyToMessageId?: number | null;
        metadata?: Record<string, unknown> | null;
      };
  error?: string;
}

export function emptyAssistant(): AiChatMessage {
  return {
    role: 'assistant',
    content: '',
    reasoningContent: '',
    toolRuns: [],
    attachments: [],
    forms: [],
    uiArtifacts: [],
  };
}

export function parseSsePayload(chunk?: AiSseChunk): {
  event: string;
  payload: AiSsePayload;
} {
  if (!chunk) return { event: '', payload: {} };
  const event = chunk.event?.trim() || 'message';
  if (!chunk.data || chunk.data === '[DONE]') return { event, payload: {} };
  try {
    const parsed: unknown = JSON.parse(chunk.data);
    return {
      event,
      payload: parsed && typeof parsed === 'object' ? (parsed as AiSsePayload) : {},
    };
  } catch {
    return { event, payload: { delta: chunk.data } };
  }
}

function upsertToolRun(
  toolRuns: AiToolRun[],
  payload: AiSsePayload,
  fallbackStatus: AiToolRun['status'],
): AiToolRun[] {
  const toolCallId = payload.toolCallId || `${payload.toolName || 'tool'}-${toolRuns.length}`;
  const next: AiToolRun = {
    toolCallId,
    toolName: payload.toolName || '查询工具',
    skillKey: payload.skillKey,
    status: (payload.status as AiToolRun['status']) || fallbackStatus,
    summary: payload.summary,
    resultSummary: fallbackStatus === 'running' ? undefined : payload.summary,
    argumentsSummary: fallbackStatus === 'running' ? payload.summary : undefined,
    durationMs: payload.durationMs,
  };
  const index = toolRuns.findIndex((item) => item.toolCallId === toolCallId);
  if (index === -1) return [...toolRuns, next];
  return toolRuns.map((item, itemIndex) => (itemIndex === index ? { ...item, ...next } : item));
}

function normalizeToolRuns(toolRuns: AiToolRun[] | undefined, fallback: AiToolRun[]): AiToolRun[] {
  if (!toolRuns) return fallback;
  return toolRuns.map((tool) => ({
    ...tool,
    status: tool.status === 'error' ? 'failed' : tool.status,
    summary: tool.resultSummary ?? tool.argumentsSummary ?? tool.summary,
  }));
}

function applyMessagePayload(
  message: AiChatMessage,
  nested: AiSsePayload['message'],
  payload: AiSsePayload,
): void {
  if (typeof nested !== 'object' || nested === null) return;
  message.forms = mergeForms(
    message.forms,
    (nested.metadata?.a2uiForm as AiFormSchema | undefined) ?? payload.form,
  );
  message.reviews = mergeById<AiReviewSchema>(
    message.reviews,
    (nested.metadata?.a2uiReview as AiReviewSchema | undefined) ?? payload.review,
  );
  message.charts = mergeById<AiChartSchema>(
    message.charts,
    (nested.metadata?.a2uiChart as AiChartSchema | AiChartSchema[] | undefined) ?? payload.chart,
  );
  const artifacts = nested.metadata?.uiArtifacts;
  if (Array.isArray(artifacts)) {
    for (const artifact of artifacts) {
      if (artifact && typeof artifact === 'object' && typeof artifact.id === 'string') {
        mergeArtifactIntoMessage(message, artifact as AiArtifactSchema);
      }
    }
  }
  message.replyToMessageId = nested.replyToMessageId ?? message.replyToMessageId;
  message.metadata = nested.metadata ?? message.metadata;
}

export function reduceAiSseMessage(
  originMessage: AiChatMessage | undefined,
  chunk?: AiSseChunk,
): AiChatMessage {
  const message = originMessage ? { ...originMessage } : emptyAssistant();
  const { event, payload } = parseSsePayload(chunk);

  if (event === 'message.created') {
    const nested = typeof payload.message === 'object' ? payload.message : undefined;
    message.id = nested?.id ?? payload.assistantMessageId ?? payload.messageId ?? message.id;
    message.content = nested?.content ?? message.content;
    message.reasoningContent = nested?.reasoningContent ?? message.reasoningContent;
    message.toolRuns = normalizeToolRuns(nested?.toolRuns, message.toolRuns);
    message.attachments = nested?.attachments ?? message.attachments;
    applyMessagePayload(message, nested, payload);
  } else if (event === 'reasoning.delta') {
    message.retrying = null;
    message.reasoningContent += payload.delta ?? payload.reasoningContent ?? '';
  } else if (event === 'content.delta') {
    message.retrying = null;
    message.content += payload.delta ?? payload.content ?? '';
  } else if (event === 'model.retrying' && payload.retry) {
    message.retrying = payload.retry;
  } else if (event === 'ui.form' && payload.form) {
    message.forms = mergeForms(message.forms, payload.form);
  } else if (event === 'ui.review' && payload.review) {
    message.reviews = mergeById<AiReviewSchema>(message.reviews, payload.review);
  } else if (event === 'ui.chart' && payload.chart) {
    message.charts = mergeById<AiChartSchema>(message.charts, payload.chart);
  } else if (event === 'ui.artifact' && payload.artifact) {
    mergeArtifactIntoMessage(message, payload.artifact);
  } else if (event === 'ui.import_preflight' && payload.preflight) {
    message.metadata = { ...message.metadata, a2uiImportPreflight: payload.preflight };
  } else if (event === 'ui.import_wizard' && payload.wizard) {
    message.metadata = { ...message.metadata, a2uiImportWizard: payload.wizard };
  } else if (event === 'tool.started') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'running');
  } else if (event === 'tool.completed') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'success');
  } else if (event === 'tool.failed') {
    message.toolRuns = upsertToolRun(message.toolRuns, payload, 'failed');
  } else if (event === 'attachment.processed' && payload.attachment) {
    if (!message.attachments.some((item) => item.id === payload.attachment?.id)) {
      message.attachments = [...message.attachments, payload.attachment];
    }
  } else if (event === 'message.completed') {
    const nested = typeof payload.message === 'object' ? payload.message : undefined;
    message.id = nested?.id ?? payload.messageId ?? message.id;
    message.content = nested?.content ?? payload.content ?? message.content;
    message.reasoningContent =
      nested?.reasoningContent ?? payload.reasoningContent ?? message.reasoningContent;
    message.toolRuns = normalizeToolRuns(nested?.toolRuns, message.toolRuns);
    message.attachments = nested?.attachments ?? message.attachments;
    applyMessagePayload(message, nested, payload);
    message.retrying = null;
  } else if (event === 'message.cancelled') {
    message.id = payload.messageId ?? message.id;
    message.cancelled = true;
    message.retrying = null;
  } else if (event === 'error') {
    message.retrying = null;
    message.error =
      (typeof payload.message === 'string' ? payload.message : undefined) ||
      payload.error ||
      'AI 回答生成失败';
  }
  return message;
}