EasyFlow/easyflow-ui-admin/app/src/views/ai/shared/agent-agui/projection.ts

import type {
  ChatTimelineItem,
  ChatTimelineKnowledgeHit,
  ChatTimelineMessageItem,
  ChatTimelineSkillInvocationStatus,
  ChatTimelineToolStatus,
} from '@easyflow/common-ui';

import type { AguiEvent } from './client';

import { ChatTimelineBuilder } from '@easyflow/common-ui';

import { EventType } from '@ag-ui/client';

import { projectArtifactPayload } from './artifact-projection';
import { easyFlowAguiCustomEvent } from './custom-events';

export interface AguiTimelineProjectionOptions {
  finishedAt?: number;
  onInputAccepted?: (payload: Record<string, unknown>) => Promise<void> | void;
  roundId?: string;
  startedAt?: number;
}

export interface AguiTimelineProjectionState {
  startedAt?: number;
  toolArgs: Map<string, string>;
  toolNames: Map<string, string>;
}

export function createAguiTimelineProjectionState(
  startedAt?: number,
): AguiTimelineProjectionState {
  return {
    startedAt,
    toolArgs: new Map(),
    toolNames: new Map(),
  };
}

function asText(value: unknown) {
  return value === null || value === undefined ? '' : String(value);
}

function asRecord(value: unknown): Record<string, unknown> {
  return value && typeof value === 'object' && !Array.isArray(value)
    ? (value as Record<string, unknown>)
    : {};
}

function asArray(value: unknown) {
  return Array.isArray(value) ? value : [];
}

function metadata(
  options: AguiTimelineProjectionOptions,
  state: AguiTimelineProjectionState,
  messageId?: string,
): Partial<ChatTimelineMessageItem> {
  return {
    ...(messageId ? { id: messageId } : {}),
    ...(options.roundId ? { roundId: options.roundId } : {}),
    ...(state.startedAt ? { turnStartedAt: state.startedAt } : {}),
  };
}

function parseToolArgs(value: string) {
  if (!value) return undefined;
  try {
    return JSON.parse(value);
  } catch {
    return value;
  }
}

function normalizeKnowledgeItems(
  payload: Record<string, unknown>,
): ChatTimelineKnowledgeHit[] {
  const source =
    payload.items ??
    payload.hits ??
    payload.documents ??
    payload.knowledgeReferences ??
    payload.knowledgeCitations ??
    [];
  return asArray(source)
    .map((value, index) => {
      const item = asRecord(value);
      const itemMetadata = asRecord(item.metadata);
      const documentName = asText(
        item.documentName ?? item.documentTitle ?? item.title,
      );
      const sourceFileName = asText(
        item.sourceFileName ?? itemMetadata.sourceFileName,
      );
      const chunkContent = asText(
        item.chunkContent ?? item.content ?? item.text ?? item.summary,
      );
      const rawScore = item.score ?? item.similarity;
      const score =
        typeof rawScore === 'number' || typeof rawScore === 'string'
          ? rawScore
          : undefined;
      return {
        ...item,
        chunkContent,
        chunkId: asText(item.chunkId ?? itemMetadata.chunkId),
        content: asText(item.content ?? item.text ?? item.summary),
        documentId: asText(item.documentId ?? itemMetadata.documentId),
        documentName,
        id: asText(item.id ?? item.chunkId ?? index),
        knowledgeId: asText(item.knowledgeId ?? payload.knowledgeId),
        knowledgeName: asText(item.knowledgeName ?? payload.knowledgeName),
        metadata: itemMetadata,
        score,
        sourceFileName,
        sourceUri: asText(item.sourceUri ?? itemMetadata.sourceUri),
        title: documentName || sourceFileName || asText(item.source),
      } satisfies ChatTimelineKnowledgeHit;
    })
    .filter((item) => item.chunkContent || item.title || item.documentName);
}

function asyncToolStatus(
  payload: Record<string, unknown>,
): ChatTimelineToolStatus {
  const status = asText(payload.status).toUpperCase();
  if (status === 'SUCCEEDED') return 'success';
  if (['CANCELLED', 'FAILED', 'TIMEOUT'].includes(status)) return 'error';
  return 'running';
}

function statusKey(
  payload: Record<string, unknown>,
  options: AguiTimelineProjectionOptions,
  fallback: string,
) {
  const value = asText(payload.statusKey) || fallback;
  return options.roundId ? `${value}:${options.roundId}` : value;
}

const skillInvocationStatuses = new Set<ChatTimelineSkillInvocationStatus>([
  'CANCELLED',
  'FAILED',
  'INCOMPLETE',
  'RUNNING',
  'SUCCESS',
]);

function normalizeSkillInvocationStatus(
  value: unknown,
): ChatTimelineSkillInvocationStatus | undefined {
  const status = asText(value).trim().toUpperCase();
  return skillInvocationStatuses.has(
    status as ChatTimelineSkillInvocationStatus,
  )
    ? (status as ChatTimelineSkillInvocationStatus)
    : undefined;
}

function skillInvocationStatusKey(
  payload: Record<string, unknown>,
  options: AguiTimelineProjectionOptions,
) {
  const supplied = asText(payload.statusKey).trim();
  if (supplied) {
    return supplied;
  }
  const identity =
    asText(payload.skillId).trim() ||
    asText(payload.skillName).trim() ||
    'unknown';
  return `skill-invocation:${options.roundId || 'draft'}:${identity}`;
}

function applyCustomEvent(
  items: ChatTimelineItem[],
  event: Extract<AguiEvent, { type: EventType.CUSTOM }>,
  options: AguiTimelineProjectionOptions,
  state: AguiTimelineProjectionState,
) {
  const payload = asRecord(event.value);
  const turnMetadata = metadata(options, state);
  if (event.name === easyFlowAguiCustomEvent.inputAccepted) {
    void options.onInputAccepted?.(payload);
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.artifactPublished) {
    projectArtifactPayload(items, payload, turnMetadata);
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.toolApprovalRequired) {
    ChatTimelineBuilder.appendToolApproval(
      items,
      {
        approvalId: asText(payload.approvalId),
        expiresAt: asText(payload.expiresAt),
        input: payload.input,
        metadata: payload.metadata,
        toolCallId: asText(payload.toolCallId),
        toolDisplayName: asText(payload.toolDisplayName),
        toolName: asText(payload.toolName) || '工具调用',
        toolType: asText(payload.toolType),
      },
      turnMetadata,
    );
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.toolApprovalResolved) {
    const toolCallId = asText(payload.toolCallId);
    if (asText(payload.status).toUpperCase() === 'APPROVED') {
      ChatTimelineBuilder.upsertToolCall(items, {
        ...turnMetadata,
        approvalId: asText(payload.approvalId),
        status: 'running',
        toolCallId,
      });
    } else {
      ChatTimelineBuilder.markToolRejected(items, {
        ...turnMetadata,
        approvalId: asText(payload.approvalId),
        reason: asText(payload.reason),
        toolCallId,
      });
    }
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.toolMetadata) {
    const toolCallId = asText(payload.toolCallId);
    const toolDisplayName = asText(payload.toolDisplayName);
    if (toolCallId && toolDisplayName) {
      state.toolNames.set(toolCallId, toolDisplayName);
      ChatTimelineBuilder.upsertToolCall(items, {
        ...turnMetadata,
        status: 'running',
        toolCallId,
        toolName: toolDisplayName,
      });
    }
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.knowledgeCitations) {
    ChatTimelineBuilder.appendKnowledge(
      items,
      normalizeKnowledgeItems(payload),
      turnMetadata,
    );
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.asyncToolStatus) {
    const input = asRecord(payload.input ?? payload.toolInput);
    ChatTimelineBuilder.upsertToolCall(items, {
      ...turnMetadata,
      input: payload.input ?? payload.toolInput,
      output:
        payload.summary ??
        payload.label ??
        payload.output ??
        payload.result ??
        payload.text,
      sourceToolCallId: asText(
        payload.sourceToolCallId ?? payload.source_tool_call_id,
      ),
      status: asyncToolStatus(payload),
      statusKey: statusKey(payload, options, 'knowledge-retrieval'),
      taskId: asText(payload.taskId ?? input.taskId ?? input.task_id),
      toolCallId: asText(payload.toolCallId ?? payload.taskId ?? payload.id),
      toolName: asText(
        payload.toolDisplayName ?? payload.toolName ?? payload.name,
      ),
    });
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.knowledgeRetrievalStatus) {
    ChatTimelineBuilder.upsertKnowledgeRetrievalStatus(
      items,
      asText(payload.status).toLowerCase() === 'running' ? 'running' : 'done',
      statusKey(payload, options, 'knowledge-retrieval'),
      turnMetadata,
    );
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.skillInvocationStatus) {
    const status = normalizeSkillInvocationStatus(payload.status);
    if (!status) {
      return;
    }
    ChatTimelineBuilder.upsertSkillInvocationStatus(items, {
      ...turnMetadata,
      displayName:
        asText(payload.skillDisplayName).trim() ||
        asText(payload.skillName).trim() ||
        '技能',
      status,
      statusKey: skillInvocationStatusKey(payload, options),
    });
    return;
  }
  if (event.name === easyFlowAguiCustomEvent.runtimeContextStatus) {
    ChatTimelineBuilder.upsertMemoryCompressionStatus(items, {
      ...turnMetadata,
      compressed:
        typeof payload.compressed === 'boolean'
          ? payload.compressed
          : undefined,
      label: asText(payload.label),
      phase: asText(payload.phase),
      status: asText(payload.status),
      statusKey: statusKey(payload, options, 'memory-compression'),
    });
  }
}

export function applyAguiEventToTimeline(
  items: ChatTimelineItem[],
  event: AguiEvent,
  options: AguiTimelineProjectionOptions = {},
  state = createAguiTimelineProjectionState(),
) {
  state.startedAt ??= options.startedAt ?? Date.now();
  switch (event.type) {
    case EventType.CUSTOM: {
      applyCustomEvent(items, event, options, state);
      return;
    }
    case EventType.MESSAGES_SNAPSHOT: {
      const assistantMessage = [...event.messages]
        .reverse()
        .find((message) => message.role === 'assistant');
      if (assistantMessage?.content !== undefined) {
        ChatTimelineBuilder.replaceMessageContent(
          items,
          assistantMessage.content,
          metadata(options, state, assistantMessage.id),
        );
      }
      return;
    }
    case EventType.REASONING_MESSAGE_CONTENT: {
      ChatTimelineBuilder.appendThinkingDelta(
        items,
        event.delta,
        metadata(options, state, event.messageId),
      );
      return;
    }
    case EventType.RUN_ERROR: {
      if (event.code === 'RUN_CANCELLED') {
        ChatTimelineBuilder.finalize(
          items,
          {
            ...metadata(options, state),
            turnFinishedAt: options.finishedAt ?? Date.now(),
            turnSucceeded: false,
          },
          {
            runningSkillStatus: 'cancelled',
          },
        );
        return;
      }
      ChatTimelineBuilder.appendError(
        items,
        event.message || '请求失败',
        metadata(options, state),
      );
      ChatTimelineBuilder.finalize(items, {
        ...metadata(options, state),
        turnFinishedAt: options.finishedAt ?? Date.now(),
        turnSucceeded: false,
      });
      return;
    }
    case EventType.RUN_FINISHED: {
      ChatTimelineBuilder.finalize(items, {
        ...metadata(options, state),
        roundCompleted: true,
        turnFinishedAt: options.finishedAt ?? Date.now(),
        turnSucceeded: true,
      });
      return;
    }
    case EventType.RUN_STARTED: {
      ChatTimelineBuilder.ensureAssistantTurn(
        items,
        metadata(
          options,
          state,
          options.roundId ? `turn-${options.roundId}` : undefined,
        ),
      );
      return;
    }
    case EventType.TEXT_MESSAGE_CONTENT: {
      ChatTimelineBuilder.appendMessageDelta(
        items,
        event.delta,
        metadata(options, state, event.messageId),
      );
      return;
    }
    case EventType.TOOL_CALL_ARGS: {
      const args = `${state.toolArgs.get(event.toolCallId) || ''}${event.delta}`;
      state.toolArgs.set(event.toolCallId, args);
      ChatTimelineBuilder.upsertToolCall(items, {
        ...metadata(options, state),
        input: parseToolArgs(args),
        status: 'running',
        toolCallId: event.toolCallId,
        toolName: state.toolNames.get(event.toolCallId),
      });
      return;
    }
    case EventType.TOOL_CALL_RESULT: {
      ChatTimelineBuilder.upsertToolCall(items, {
        ...metadata(options, state),
        output: event.content,
        status: 'success',
        toolCallId: event.toolCallId,
        toolName: state.toolNames.get(event.toolCallId),
      });
      return;
    }
    case EventType.TOOL_CALL_START: {
      state.toolNames.set(event.toolCallId, event.toolCallName);
      ChatTimelineBuilder.upsertToolCall(items, {
        ...metadata(options, state),
        status: 'running',
        toolCallId: event.toolCallId,
        toolName: event.toolCallName,
      });
      break;
    }
    default: {
      break;
    }
  }
}