EasyFlow/easyflow-ui-admin/app/src/views/ai/agents/composables/useAgentTryoutStream.ts

import type {
  ChatDocumentAttachment,
  ChatImageAttachment,
  ChatTimelineItem as ChatTimelineItemType,
  ChatTimelineMessageItem,
} from '@easyflow/common-ui';

import type { AguiEvent } from '../../shared/agent-agui/client';
import type {
  AgentInfo,
  AgentKnowledgeBinding,
  AgentSkillBinding,
  AgentToolBinding,
} from '../types';

import { ref } from 'vue';

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

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

import {
  createAguiRunId,
  EasyFlowAguiClient,
} from '../../shared/agent-agui/client';
import {
  isRetryableAguiTransportError,
  resumeAguiRunUntilTerminal,
} from '../../shared/agent-agui/reconnect';
import { clearAgentDraftSession } from '../api';
import { useAgentTryoutRawRounds } from './useAgentTryoutRawRounds';

function resolveDraftSessionId(agent: AgentInfo) {
  return `agent-draft-${agent.id || agent.localId || 'unsaved'}`;
}

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 draftAgentTransport(agent: AgentInfo) {
  return {
    avatar: agent.avatar,
    categoryId: agent.categoryId,
    description: agent.description,
    executionConfigJson: agent.executionConfigJson,
    generationConfigJson: agent.generationConfigJson,
    id: agent.id,
    memoryConfigJson: agent.memoryConfigJson,
    modelConfigJson: agent.modelConfigJson,
    modelId: agent.modelId,
    name: agent.name,
    promptConfigJson: agent.promptConfigJson,
    publishStatus: agent.publishStatus,
    status: agent.status,
    visibilityScope: agent.visibilityScope,
  };
}

function draftToolBindingTransport(binding: AgentToolBinding) {
  return {
    enabled: binding.enabled,
    hitlConfigJson: binding.hitlConfigJson,
    hitlEnabled: binding.hitlEnabled,
    id: binding.id,
    optionsJson: binding.optionsJson,
    sortNo: binding.sortNo,
    targetId: binding.targetId,
    toolName: binding.toolName,
    toolType: binding.toolType,
  };
}

function draftKnowledgeBindingTransport(binding: AgentKnowledgeBinding) {
  return {
    enabled: binding.enabled,
    id: binding.id,
    knowledgeId: binding.knowledgeId,
    optionsJson: binding.optionsJson,
    retrievalMode: binding.retrievalMode,
    sortNo: binding.sortNo,
  };
}

function draftSkillBindingTransport(binding: AgentSkillBinding) {
  return {
    skillId: binding.skillId,
    sortNo: binding.sortNo,
  };
}

interface DraftRuntimeContext {
  agent: AgentInfo;
  knowledgeBindings: AgentKnowledgeBinding[];
  skillBindings: AgentSkillBinding[];
  toolBindings: AgentToolBinding[];
}

interface ActiveDraftRun {
  cancelling: boolean;
  inputAcceptedNotified: boolean;
  lastCursor: number;
  roundId: string;
  runId: string;
  sessionId: string;
  startedAt: number;
  stopped: boolean;
}

export function useAgentTryoutStream() {
  const timelineItems = ref<ChatTimelineItemType[]>([]);
  const loading = ref(false);
  let rawRounds: ReturnType<typeof useAgentTryoutRawRounds> | undefined;
  let activeRoundId = '';
  let activeSessionId = '';
  let activeRun: ActiveDraftRun | undefined;
  let resumeOnAccepted: (() => Promise<void> | void) | undefined;
  const aguiClient = new EasyFlowAguiClient();

  function errorMessageOf(error: unknown) {
    if (error instanceof Error) {
      return `${error.name} ${error.message}`.trim();
    }
    if (typeof error === 'string') {
      return error;
    }
    const value = asRecord(error);
    const nested = asRecord(value.cause);
    return [value.name, value.message, value.error, nested.name, nested.message]
      .map((item) => asText(item).trim())
      .filter(Boolean)
      .join(' ');
  }

  function finishStoppedRun(roundId: string) {
    if (roundId) {
      rawRounds?.failRound(roundId);
      rebuildTimeline();
    }
    loading.value = false;
  }

  function rebuildTimeline() {
    timelineItems.value = rawRounds?.buildTimelineItems() || [];
  }

  function syncDraftContext(
    payload: DraftRuntimeContext,
    restore = false,
    requestedSessionId?: string,
    onAccepted?: () => Promise<void> | void,
  ) {
    resumeOnAccepted = onAccepted;
    const sessionId =
      requestedSessionId || resolveDraftSessionId(payload.agent);
    const sessionChanged = activeSessionId !== sessionId;
    activeSessionId = sessionId;
    if (!rawRounds || sessionChanged) {
      rawRounds?.flush();
      rawRounds = useAgentTryoutRawRounds({
        mode: 'draft',
        sessionId,
      });
      activeRoundId = '';
    }
    if (restore && sessionChanged && !loading.value) {
      rebuildTimeline();
      void resumeDraftRun();
    }
  }

  function notifyInputAccepted(run: ActiveDraftRun) {
    if (run.inputAcceptedNotified || !resumeOnAccepted) {
      return;
    }
    run.inputAcceptedNotified = true;
    void resumeOnAccepted();
  }

  function projectRunEvent(
    run: ActiveDraftRun,
    rounds: NonNullable<typeof rawRounds>,
    event: AguiEvent,
  ) {
    if (activeRun !== run) return;
    const runtimeEvent = rounds.recordEvent(run.roundId, event);
    if (!runtimeEvent) return;
    rounds.projectEvent(timelineItems.value, run.roundId, runtimeEvent, () =>
      notifyInputAccepted(run),
    );
  }

  function resumeActiveRun(
    run: ActiveDraftRun,
    rounds: NonNullable<typeof rawRounds>,
  ) {
    return resumeAguiRunUntilTerminal({
      client: aguiClient,
      initialCursor: run.lastCursor,
      onCursor: (cursor) => {
        run.lastCursor = Math.max(run.lastCursor, cursor);
        rounds.updateCursor(run.roundId, run.lastCursor);
      },
      onEvent: (event) => projectRunEvent(run, rounds, event),
      runId: run.runId,
      shouldContinue: () => activeRun === run && !run.stopped,
      startedAt: run.startedAt,
    });
  }

  async function resumeDraftRun() {
    const rounds = rawRounds;
    const recovery = rounds?.recoverableRun();
    if (!rounds || !recovery || activeRun) {
      return;
    }
    const run: ActiveDraftRun = {
      cancelling: false,
      inputAcceptedNotified: false,
      lastCursor: recovery.lastCursor,
      roundId: recovery.roundId,
      runId: recovery.runId,
      sessionId: activeSessionId,
      startedAt: recovery.startedAt,
      stopped: false,
    };
    activeRun = run;
    activeRoundId = run.roundId;
    if (run.lastCursor === 0) {
      rounds.resetRuntimeEvents(run.roundId);
    }
    rebuildTimeline();
    loading.value = true;
    try {
      if (recovery.inputAccepted) {
        notifyInputAccepted(run);
      }
      const terminalReceived = await resumeActiveRun(run, rounds);
      if (terminalReceived && activeRun === run && !run.stopped) {
        finishAssistant();
        markRoundCompleted(run.roundId);
      }
    } catch (error) {
      if (activeRun === run && !run.stopped) {
        const runError = {
          message: errorMessageOf(error) || '试运行恢复失败，请稍后再试',
          runId: run.runId,
          threadId: run.sessionId,
          type: EventType.RUN_ERROR,
        } as AguiEvent;
        rounds.recordEvent(run.roundId, runError);
        rounds.projectEvent(timelineItems.value, run.roundId, runError);
        finishAssistant();
        rounds.flush();
      }
    } finally {
      if (activeRun === run && !run.cancelling) {
        activeRun = undefined;
        loading.value = false;
      }
    }
  }

  function markRoundCompleted(roundId: string) {
    if (!roundId) {
      return;
    }
    rawRounds?.completeRound(roundId);
    rebuildTimeline();
  }

  function finishAssistant() {
    ChatTimelineBuilder.finalize(timelineItems.value);
  }

  function markToolApproving(payload: {
    approvalId?: string;
    toolCallId?: string;
  }) {
    ChatTimelineBuilder.markToolApproving(timelineItems.value, payload);
  }

  function markToolRejected(payload: {
    approvalId?: string;
    reason?: string;
    toolCallId?: string;
  }) {
    ChatTimelineBuilder.markToolRejected(timelineItems.value, payload);
  }

  async function runDraft(payload: {
    agent: AgentInfo;
    documents?: ChatDocumentAttachment[];
    documentUploadIds?: string[];
    images?: ChatImageAttachment[];
    imageUploadIds?: string[];
    knowledgeBindings: AgentKnowledgeBinding[];
    onAccepted?: () => Promise<void> | void;
    prompt: string;
    sessionId?: string;
    skillBindings: AgentSkillBinding[];
    toolBindings: AgentToolBinding[];
  }) {
    syncDraftContext(payload, false, payload.sessionId);
    if (!rawRounds) {
      return;
    }
    const rounds = rawRounds;
    activeRoundId = rounds.createRound(
      payload.prompt,
      payload.images,
      payload.documents,
    );
    const runId = createAguiRunId();
    rounds.bindRun(activeRoundId, runId);
    const startedAt = Date.now();
    const run: ActiveDraftRun = {
      cancelling: false,
      inputAcceptedNotified: false,
      lastCursor: 0,
      roundId: activeRoundId,
      runId,
      sessionId: activeSessionId,
      startedAt,
      stopped: false,
    };
    resumeOnAccepted = payload.onAccepted;
    activeRun = run;
    rebuildTimeline();
    loading.value = true;
    try {
      let terminalReceived = true;
      try {
        await aguiClient.run({
          forwardedProps: {
            easyflow: {
              draft: {
                agent: draftAgentTransport(payload.agent),
                knowledgeBindings: payload.knowledgeBindings.map((binding) =>
                  draftKnowledgeBindingTransport(binding),
                ),
                skillBindings: payload.skillBindings.map((binding) =>
                  draftSkillBindingTransport(binding),
                ),
                toolBindings: payload.toolBindings.map((binding) =>
                  draftToolBindingTransport(binding),
                ),
              },
              input: {
                documentUploadIds: payload.documentUploadIds,
                imageUploadIds: payload.imageUploadIds,
              },
            },
          },
          onCursor(cursor) {
            run.lastCursor = Math.max(run.lastCursor, cursor);
            rounds.updateCursor(run.roundId, run.lastCursor);
          },
          onEvent: (event) => projectRunEvent(run, rounds, event),
          runId,
          threadId: run.sessionId,
          url: '/api/v1/agent/agui/run/draft',
          userMessage: {
            content: payload.prompt,
            id: `user-${run.roundId}`,
            role: 'user',
          },
        });
      } catch (error) {
        if (!isRetryableAguiTransportError(error, run.startedAt)) {
          throw error;
        }
        terminalReceived = await resumeActiveRun(run, rounds);
      }
      if (terminalReceived && activeRun === run && !run.stopped) {
        finishAssistant();
        markRoundCompleted(run.roundId);
      }
    } catch (error) {
      if (activeRun === run && !run.stopped) {
        const runError = {
          message: errorMessageOf(error) || '试运行失败，请稍后再试',
          runId: run.runId,
          threadId: run.sessionId,
          type: EventType.RUN_ERROR,
        } as AguiEvent;
        rounds.recordEvent(run.roundId, runError);
        rounds.projectEvent(timelineItems.value, run.roundId, runError);
        finishAssistant();
        rounds.flush();
      }
    } finally {
      if (activeRun === run && !run.cancelling) {
        activeRun = undefined;
        loading.value = false;
      }
    }
  }

  async function sendDraft(payload: {
    agent: AgentInfo;
    documents?: ChatDocumentAttachment[];
    documentUploadIds?: string[];
    images?: ChatImageAttachment[];
    imageUploadIds?: string[];
    knowledgeBindings: AgentKnowledgeBinding[];
    onAccepted?: () => Promise<void> | void;
    prompt: string;
    sessionId?: string;
    skillBindings: AgentSkillBinding[];
    toolBindings: AgentToolBinding[];
  }) {
    await runDraft(payload);
  }

  async function regenerateDraft(item: ChatTimelineMessageItem) {
    // 有状态 runtime 的历史重生成需要后端 session fork/rollback，第一版先禁用入口。
    void item;
  }

  function selectVariant(
    item: ChatTimelineMessageItem,
    direction: 'next' | 'previous',
  ) {
    if (!item.roundId || !rawRounds?.canSwitch(item, direction)) {
      return;
    }
    rawRounds.selectVariant(item.roundId, direction);
    rebuildTimeline();
  }

  function copyMessageText(item: ChatTimelineMessageItem) {
    return rawRounds?.copyText(item) || '';
  }

  async function clearDraftSession() {
    if (loading.value) {
      throw new Error('试运行进行中，暂时无法清理会话');
    }
    const sessionId = activeSessionId;
    rawRounds?.clear();
    timelineItems.value = [];
    activeRoundId = '';
    if (sessionId) {
      await clearAgentDraftSession(sessionId);
    }
  }

  function recoverAfterCancelFailure(run: ActiveDraftRun) {
    run.cancelling = false;
    run.stopped = false;
    aguiClient.detach();
    if (activeRun === run) {
      activeRun = undefined;
    }
    loading.value = false;
    rebuildTimeline();
    void resumeDraftRun();
  }

  async function stop() {
    if (!loading.value) {
      return;
    }
    const stoppedRoundId = activeRun?.roundId || activeRoundId;
    const run = activeRun;
    if (run) {
      if (run.cancelling) {
        return;
      }
      run.cancelling = true;
      run.stopped = true;
      try {
        await aguiClient.cancel(run.runId);
      } catch (error) {
        recoverAfterCancelFailure(run);
        throw error;
      }
      run.cancelling = false;
    }
    activeRun = undefined;
    if (run && rawRounds?.currentVariant(run.roundId)?.status !== 'running') {
      loading.value = false;
      rebuildTimeline();
      return;
    }
    if (run) {
      // 取消接口成功后重新订阅服务端日志，以 RUN_FINISHED/RUN_ERROR 决定最终展示状态。
      loading.value = false;
      rebuildTimeline();
      await resumeDraftRun();
      return;
    }
    finishStoppedRun(stoppedRoundId);
  }

  function dispose() {
    if (loading.value) {
      activeRun = undefined;
      aguiClient.detach();
      rawRounds?.flush();
      return;
    }
    rawRounds?.flush();
  }

  return {
    loading,
    clearDraftSession,
    markToolApproving,
    markToolRejected,
    copyMessageText,
    dispose,
    regenerateDraft,
    selectVariant,
    syncDraftContext,
    timelineItems,
    sendDraft,
    stop,
  };
}