UHD_Media_Tool/apps/web/app/components/agent/use-agent-stream.ts

import { type AgentRunEvent, type AgentUsageTotals, EventType } from "@umt/contracts";
import { useEffect, useMemo, useRef, useState } from "react";
import { agentMockEnabled, fetchRunEvents } from "~/api/agent.ts";
import { sse } from "~/lib/sse.ts";

// Agent 运行流:先按 seq 回放 agent_run_event,再跟随 SSE(agent.* 事件只作"有新事件"的提示,
// 收到后按 lastSeq 增量补读;thinking/compaction 不上全局总线,补读时一并拿到)。SSE 不可用时退回轮询。

export type ToolCallView = {
  toolCallId: string;
  capabilityId: string;
  inputSummary: string | null;
  outputSummary: string | null;
  status: "running" | "succeeded" | "failed" | "denied";
  reason: string | null;
  startedAt: string;
  endedAt: string | null;
  recovery: boolean;
};

export type StreamItem =
  | { kind: "text"; key: string; text: string; at: string }
  | { kind: "thinking"; key: string; text: string; at: string }
  | { kind: "tool"; key: string; call: ToolCallView }
  | { kind: "compaction"; key: string; at: string; reason: string; status: string | null; done: boolean }
  | {
      kind: "end";
      key: string;
      at: string;
      ok: boolean;
      message: string;
      reason: string | null;
      budget: string | null;
    };

export type AgentStreamState = {
  events: AgentRunEvent[];
  items: StreamItem[];
  usage: AgentUsageTotals | null;
  finished: boolean;
  error: string | null;
  connected: "sse" | "poll" | "mock" | "idle";
};

const SSE_TYPES: string[] = [EventType.agentTextDelta, EventType.agentToolCall, EventType.agentRunStatus];

/** 事件 → 展示项:相邻文本增量合并;工具调用按 toolCallId 折叠成一张卡 */
export function foldEvents(events: AgentRunEvent[]): StreamItem[] {
  const items: StreamItem[] = [];
  const tools = new Map<string, ToolCallView>();
  for (const e of events) {
    const last = items[items.length - 1];
    switch (e.type) {
      case "text_delta":
        if (last?.kind === "text") last.text += e.delta;
        else items.push({ kind: "text", key: `t${e.seq}`, text: e.delta, at: e.at });
        break;
      case "thinking_delta":
        if (last?.kind === "thinking") last.text += e.delta;
        else items.push({ kind: "thinking", key: `th${e.seq}`, text: e.delta, at: e.at });
        break;
      case "tool_call_start": {
        const call: ToolCallView = {
          toolCallId: e.toolCallId,
          capabilityId: e.capabilityId,
          inputSummary: e.inputSummary,
          outputSummary: null,
          status: "running",
          reason: null,
          startedAt: e.at,
          endedAt: null,
          recovery: e.recovery,
        };
        tools.set(e.toolCallId, call);
        items.push({ kind: "tool", key: `c${e.toolCallId}`, call });
        break;
      }
      case "tool_call_end": {
        const call = tools.get(e.toolCallId);
        if (call) {
          call.outputSummary = e.outputSummary;
          call.status = e.isError ? "failed" : "succeeded";
          call.endedAt = e.at;
        } else {
          items.push({
            kind: "tool",
            key: `c${e.toolCallId}`,
            call: {
              toolCallId: e.toolCallId,
              capabilityId: e.capabilityId,
              inputSummary: null,
              outputSummary: e.outputSummary,
              status: e.isError ? "failed" : "succeeded",
              reason: null,
              startedAt: e.at,
              endedAt: e.at,
              recovery: e.recovery,
            },
          });
        }
        break;
      }
      case "tool_call_rejected": {
        const call = tools.get(e.toolCallId);
        if (call) {
          call.status = "denied";
          call.reason = e.reason;
          call.endedAt = e.at;
        } else {
          items.push({
            kind: "tool",
            key: `c${e.toolCallId}`,
            call: {
              toolCallId: e.toolCallId,
              capabilityId: e.capabilityId ?? e.toolName,
              inputSummary: null,
              outputSummary: null,
              status: "denied",
              reason: e.reason,
              startedAt: e.at,
              endedAt: e.at,
              recovery: false,
            },
          });
        }
        break;
      }
      case "compaction":
        if (e.phase === "start")
          items.push({ kind: "compaction", key: `k${e.seq}`, at: e.at, reason: e.reason, status: null, done: false });
        else {
          const open = [...items].reverse().find((i) => i.kind === "compaction" && !i.done);
          if (open && open.kind === "compaction") {
            open.done = true;
            open.status = e.status ?? "completed";
          } else
            items.push({
              kind: "compaction",
              key: `k${e.seq}`,
              at: e.at,
              reason: e.reason,
              status: e.status ?? "completed",
              done: true,
            });
        }
        break;
      case "run_completed":
        items.push({
          kind: "end",
          key: `e${e.seq}`,
          at: e.at,
          ok: true,
          message: e.finalText,
          reason: null,
          budget: null,
        });
        break;
      case "run_failed":
        items.push({
          kind: "end",
          key: `e${e.seq}`,
          at: e.at,
          ok: false,
          message: e.message,
          reason: e.reason,
          budget: e.budget ?? null,
        });
        break;
      default:
        break;
    }
  }
  return items;
}

export function useAgentStream(runId: string | undefined, active: boolean): AgentStreamState {
  const [events, setEvents] = useState<AgentRunEvent[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [connected, setConnected] = useState<AgentStreamState["connected"]>("idle");
  const lastSeq = useRef(-1);
  const inflight = useRef(false);
  const pending = useRef(false);

  // biome-ignore lint/correctness/useExhaustiveDependencies: 换 run 时重置
  useEffect(() => {
    setEvents([]);
    setError(null);
    lastSeq.current = -1;
  }, [runId]);

  useEffect(() => {
    if (!runId) return;
    let stopped = false;
    const ctrl = new AbortController();

    const pull = async () => {
      if (inflight.current) {
        pending.current = true;
        return;
      }
      inflight.current = true;
      try {
        do {
          pending.current = false;
          const next = await fetchRunEvents(runId, lastSeq.current, ctrl.signal);
          if (stopped) return;
          const fresh = next.filter((e) => e.seq > lastSeq.current).sort((a, b) => a.seq - b.seq);
          if (fresh.length) {
            lastSeq.current = fresh[fresh.length - 1]!.seq;
            setEvents((prev) => [...prev, ...fresh]);
          }
          setError(null);
        } while (pending.current && !stopped);
      } catch (err) {
        if ((err as Error).name !== "AbortError" && !stopped) setError((err as Error).message);
      } finally {
        inflight.current = false;
      }
    };

    void pull();
    if (!active)
      return () => {
        stopped = true;
        ctrl.abort();
      };

    const mock = agentMockEnabled();
    // 共用页面级 SSE 连接:agent.* 事件只当"有新事件"的提示,按 lastSeq 补读
    const unsubscribe = mock
      ? null
      : sse.subscribe(SSE_TYPES, (e) => {
          const p = e.payload as { runId?: string } | null;
          if (e.type === EventType.resync || p?.runId === runId) void pull();
        });
    setConnected(mock ? "mock" : typeof EventSource !== "undefined" ? "sse" : "poll");
    // mock 模拟流式:密集轮询;真实端点:SSE 为主,2.5 秒轮询兜底
    const timer = setInterval(() => void pull(), mock ? 150 : 2500);
    return () => {
      stopped = true;
      ctrl.abort();
      clearInterval(timer);
      unsubscribe?.();
      setConnected("idle");
    };
  }, [runId, active]);

  return useMemo(() => {
    const items = foldEvents(events);
    const end = [...events].reverse().find((e) => e.type === "run_completed" || e.type === "run_failed");
    const usage = end && (end.type === "run_completed" || end.type === "run_failed") ? end.usage : null;
    return { events, items, usage, finished: !!end, error, connected };
  }, [events, error, connected]);
}