Agent-System/web/src/hooks/useRecoveredTaskPhases.ts

import { useEffect, useRef } from 'react';
import type { StreamEvent } from '../api/chatStream';
import { readTaskRunSSEOnce, subscribeTaskRunEvents } from '../api/taskRuns';
import { applyPhaseEvent } from '../engine/phaseEventMapper';
import {
  messageNeedsPhaseReplay,
  phaseReplayKey,
  type PhaseReplayCacheEntry,
} from '../lib/taskPhaseReplay';
import type { ChatMessage, PhaseEvent } from '../types/workbench';

/**
 * 刷新后进度条只存在于 task_run_events。
 * running 继续订阅；interrupted 用 replay_only 读完即关，避免没有 done 事件时挂住连接。
 */
export function useRecoveredTaskPhases(
  sessionId: string,
  messages: ChatMessage[],
  updateMessage: (id: string, patch: Partial<ChatMessage>) => void,
) {
  const started = useRef(new Set<string>());
  const cache = useRef(new Map<string, PhaseReplayCacheEntry>());
  const updateRef = useRef(updateMessage);
  updateRef.current = updateMessage;
  const messagesRef = useRef(messages);
  messagesRef.current = messages;
  const abortRef = useRef<AbortController | null>(null);

  useEffect(() => {
    started.current.clear();
    cache.current.clear();
    const ac = new AbortController();
    abortRef.current = ac;
    return () => ac.abort();
  }, [sessionId]);

  const key = phaseReplayKey(messages);

  useEffect(() => {
    const signal = abortRef.current?.signal;
    if (!signal || signal.aborted) return;
    const current = messagesRef.current;

    for (const message of current) {
      const runId = message.taskRunId;
      if (!runId) continue;
      const cached = cache.current.get(runId);
      if (cached && cached.events.length > 0 && (!message.phaseEvents || message.phaseEvents.length === 0)) {
        updateRef.current(message.id, {
          phaseEvents: cached.events,
          livePhaseLabel: cached.label,
        });
      }
    }

    for (const message of current.filter(messageNeedsPhaseReplay)) {
      const runId = message.taskRunId;
      if (!runId || started.current.has(runId)) continue;
      started.current.add(runId);
      const messageId = message.id;
      const live = message.status === 'streaming_status';
      void replayRun(runId, messageId, live, cache.current, updateRef.current, signal);
    }
  }, [key, sessionId]);
}

async function replayRun(
  runId: string,
  messageId: string,
  live: boolean,
  cache: Map<string, PhaseReplayCacheEntry>,
  updateMessage: (id: string, patch: Partial<ChatMessage>) => void,
  signal: AbortSignal,
) {
  let strip: PhaseEvent[] = cache.get(runId)?.events || [];
  const onEvent = (ev: StreamEvent) => {
    if (ev.type !== 'phase' || !ev.phase) return;
    strip = applyPhaseEvent(strip, ev.phase);
    const label = ev.phase.message;
    cache.set(runId, { events: strip, label });
    updateMessage(messageId, { phaseEvents: strip, livePhaseLabel: label });
  };
  try {
    if (live) {
      await subscribeTaskRunEvents(runId, { onEvent }, signal);
      return;
    }
    await readTaskRunSSEOnce(runId, 0, onEvent, signal, { replayOnly: true });
  } catch (err) {
    if (!signal.aborted) {
      console.warn('task phase replay failed', runId, err);
    }
  }
}