deerflow-code/frontend-web/src/core/threads/use-research-report-stream.ts

"use client";

/**
 * 聊天「写作模式」内联深度研究报告的流式状态。
 *
 * 后端 `deep_research_report` 工具把引擎事件经 LangGraph custom 通道
 * （`get_stream_writer`）以 `{type: "deep_research_report", event: {...}}`
 * 帧发到前端 `onCustomEvent`；本模块按 threadId 归约这些帧，供聊天页注入
 * 进度步骤条（`deep_research_progress` 虚拟工具调用）与 report.md 的
 * `write_file` stub 流式预览（复用沙箱 write-file 虚拟 URL 机制）。
 *
 * 与工作台 SSE 不同：custom 事件不落库、断线不重放，刷新页面即丢——
 * 报告完成后 md 已持久化在聊天线程 artifacts，由真实产物接管展示。
 */

import { useCallback, useSyncExternalStore } from "react";

/** 深度研究 basic 管线的阶段序（与后端 events.ALL_PHASES 对齐的子集）。 */
const PHASE_ORDER = [
  "initializing",
  "planning",
  "collecting",
  "curating",
  "compressing",
  "writing",
  "done",
] as const;

function phaseRank(phase: string): number {
  const index = PHASE_ORDER.indexOf(phase as (typeof PHASE_ORDER)[number]);
  return index === -1 ? 0 : index;
}

export type ResearchReportStreamState = {
  /** "active" | "success" | "error"（含被取消/断流）。 */
  status: "active" | "success" | "error";
  /** 工具调用 id（帧里的 jobId），一次写作轮一个。 */
  jobId: string;
  phase: string;
  queries: string[];
  completedQueries: number;
  sourceCount: number;
  /** report_delta/report_chunk 按 offset 去重拼装的报告全文。 */
  reportText: string;
  /** report_thinking 增量（写作思考，可选展示）。 */
  thinking: string;
  error: string | null;
};

type Listener = () => void;

type StoreEntry = {
  state: ResearchReportStreamState | null;
  listeners: Set<Listener>;
  emitTimer: ReturnType<typeof setTimeout> | null;
};

const EMPTY_ACTIVE: ResearchReportStreamState = {
  status: "active",
  jobId: "",
  phase: "initializing",
  queries: [],
  completedQueries: 0,
  sourceCount: 0,
  reportText: "",
  thinking: "",
  error: null,
};

const stores = new Map<string, StoreEntry>();

function getStore(threadId: string): StoreEntry {
  let store = stores.get(threadId);
  if (!store) {
    store = { state: null, listeners: new Set(), emitTimer: null };
    stores.set(threadId, store);
  }
  return store;
}

/**
 * token 帧的通知合并间隔（ms）。report_delta / report_thinking 是 token 级
 * 高频帧（每秒可达数十帧）；状态本身同步更新（字符串拼接很便宜），但订阅
 * 通知若逐帧发出，会让聊天页整棵树（长消息列表 + 沙箱 Markdown 预览）按
 * token 频率重渲染，页面直接卡死。工作台用 requestAnimationFrame 批处理
 * （DeepResearchWorkbench），聊天这里等效地按固定间隔合并——每秒最多
 * 10 次通知，流式观感不受影响。终态帧（tool_finished）与
 * finalize/clear 立即 flush，结束态不滞后。
 */
const EMIT_COALESCE_MS = 100;

function emitNow(store: StoreEntry) {
  if (store.emitTimer !== null) {
    clearTimeout(store.emitTimer);
    store.emitTimer = null;
  }
  for (const listener of store.listeners) listener();
}

function emitCoalesced(store: StoreEntry) {
  if (store.emitTimer !== null) return;
  store.emitTimer = setTimeout(() => {
    store.emitTimer = null;
    for (const listener of store.listeners) listener();
  }, EMIT_COALESCE_MS);
}

export function clearResearchReportStream(threadId: string) {
  const store = getStore(threadId);
  if (store.state === null) return;
  store.state = null;
  emitNow(store);
}

/** run 结束（成功/失败/取消/断流）时收尾仍处于 active 的状态。 */
export function finalizeResearchReportStream(threadId: string) {
  const store = getStore(threadId);
  if (store.state?.status !== "active") return;
  store.state = { ...store.state, status: "error", error: store.state.error ?? "已中断" };
  emitNow(store);
}

/** 与工作台 appendReportDelta 相同的 offset 去重追加。 */
function appendWithOffset(current: string, delta: string, offset: unknown): string {
  if (!delta) return current;
  if (typeof offset !== "number" || !Number.isFinite(offset) || offset < 0) return current + delta;
  if (offset >= current.length) return current + delta;
  if (current.slice(offset, offset + delta.length) === delta) return current;
  return current.slice(0, offset) + delta;
}

/** onCustomEvent 深度研究帧的 event 载荷。 */
export type ResearchReportEventFrame = {
  type: string;
  phase?: string;
  payload?: Record<string, unknown>;
};

export function ingestResearchReportEvent(threadId: string, frame: ResearchReportEventFrame) {
  const store = getStore(threadId);
  const jobId = typeof (frame as { jobId?: unknown }).jobId === "string"
    ? String((frame as { jobId?: unknown }).jobId)
    : "";
  const previous = store.state;
  // 新一轮工具调用（jobId 变化）从零开始，旧状态自然淘汰。
  const base: ResearchReportStreamState = previous && previous.jobId === jobId
    ? previous
    : { ...EMPTY_ACTIVE, jobId };

  const payload = frame.payload ?? {};
  let next = base;
  switch (frame.type) {
    case "phase_changed":
      next = { ...base, phase: typeof payload.to === "string" ? payload.to : base.phase };
      break;
    case "queries_planned":
      next = {
        ...base,
        queries: Array.isArray(payload.queries)
          ? payload.queries.filter((q): q is string => typeof q === "string")
          : base.queries,
        completedQueries: 0,
      };
      break;
    case "query_completed":
      next = {
        ...base,
        completedQueries: base.completedQueries + 1,
        sourceCount: base.sourceCount + (typeof payload.resultCount === "number" ? payload.resultCount : 0),
      };
      break;
    case "sources_curated": {
      const selected = Array.isArray(payload.selectedIds) ? payload.selectedIds.length : null;
      next = {
        ...base,
        sourceCount: selected !== null && selected > 0 ? selected : base.sourceCount,
      };
      break;
    }
    case "report_delta":
    case "report_chunk": {
      if (typeof payload.target === "string" && payload.target !== "report") break;
      const delta = typeof payload.delta === "string" ? payload.delta : "";
      const reportText = appendWithOffset(base.reportText, delta, payload.offset);
      next = reportText === base.reportText ? base : { ...base, reportText };
      break;
    }
    case "report_thinking": {
      const delta = typeof payload.delta === "string" ? payload.delta : "";
      next = delta ? { ...base, thinking: base.thinking + delta } : base;
      break;
    }
    case "report_completed":
      next = {
        ...base,
        phase: "writing",
        sourceCount: typeof payload.sourceCount === "number" && payload.sourceCount > 0
          ? payload.sourceCount
          : base.sourceCount,
      };
      break;
    case "warning":
    case "error": {
      const message = typeof payload.message === "string" ? payload.message : "";
      if (payload.recoverable === false && message) {
        next = { ...base, error: message };
      }
      break;
    }
    case "tool_finished":
      next = {
        ...base,
        status: payload.status === "success" ? "success" : "error",
        error: typeof payload.error === "string" ? payload.error : base.error,
        phase: payload.status === "success" ? "done" : base.phase,
      };
      break;
    default:
      break;
  }

  if (frame.phase && next === base) next = { ...base, phase: frame.phase };
  if (next !== previous) {
    store.state = next;
    // 状态同步更新（peek/快照立即可见）；通知按 token 帧频率合并，终态立即。
    if (frame.type === "tool_finished") {
      emitNow(store);
    } else {
      emitCoalesced(store);
    }
  }
}

/** 订阅某线程的报告流状态；无进行中/已结束状态时返回 null。 */
export function useResearchReportStream(threadId: string | null | undefined): ResearchReportStreamState | null {
  const subscribe = useCallback(
    (listener: Listener) => {
      if (!threadId) return () => {};
      return subscribeResearchReportStream(threadId, listener);
    },
    [threadId],
  );
  const getSnapshot = useCallback(() => (threadId ? getStore(threadId).state : null), [threadId]);
  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
}

/**
 * 直接订阅某线程的通知（节流合并后触发，非逐 token 帧）；返回退订函数。
 * `useResearchReportStream` 内部也走这里，测试可直接验证合并语义。
 */
export function subscribeResearchReportStream(threadId: string, listener: Listener): () => void {
  const store = getStore(threadId);
  store.listeners.add(listener);
  return () => {
    store.listeners.delete(listener);
  };
}

/** 非订阅式快照读取（测试 / 调试用）。 */
export function peekResearchReportStream(threadId: string): ResearchReportStreamState | null {
  return getStore(threadId).state;
}

// ── 聊天页注入消息构造 ────────────────────────────────────────────────────────

export const RESEARCH_REPORT_OUTPUT_PATH = "/mnt/user-data/outputs/report.md";

export type ResearchProgressToolCall = {
  id: string;
  name: "deep_research_progress";
  args: {
    label: string;
    detail: string;
    status: "complete" | "active";
    thinking?: string;
  };
};

/** 结构与工作台 buildWritingProgressToolCalls 相同；步骤固定为 basic 管线四段。 */
export function buildResearchProgressToolCalls(state: ResearchReportStreamState): ResearchProgressToolCall[] {
  const rank = phaseRank(state.phase);
  const finished = state.status !== "active";
  const effectiveRank = finished && state.status === "success" ? PHASE_ORDER.length - 1 : rank;
  const stages = [
    {
      id: "plan",
      label: "拆解检索问题",
      boundary: phaseRank("collecting"),
      detail: state.queries.length > 0
        ? `已拆出 ${state.queries.length} 个检索子问题`
        : "正在围绕课题拆解检索子问题",
    },
    {
      id: "collect",
      label: "检索参考资料",
      boundary: phaseRank("compressing"),
      detail: state.queries.length > 0
        ? `已检索 ${Math.min(state.completedQueries, state.queries.length)}/${state.queries.length} 个子问题 · 累计 ${state.sourceCount} 条资料`
        : state.sourceCount > 0
          ? `已检索到 ${state.sourceCount} 条资料`
          : "正在逐个子问题检索资料",
    },
    {
      id: "write",
      label: "撰写研究报告",
      boundary: PHASE_ORDER.length,
      detail: state.reportText
        ? "正在按「通用分析结构」逐字撰写，右侧沙箱同步预览"
        : "正在整理资料并开始撰写报告",
    },
  ];
  const activeIndex = stages.findIndex((stage) => effectiveRank < stage.boundary);
  const visibleLast = activeIndex === -1 ? stages.length - 1 : activeIndex;

  return stages.slice(0, visibleLast + 1).map((stage, index) => {
    const complete = effectiveRank >= stage.boundary;
    const failed = finished && state.status === "error" && index === visibleLast && !complete;
    return {
      id: `research-progress-${state.jobId}-${stage.id}`,
      name: "deep_research_progress" as const,
      args: {
        label: failed ? `${stage.label}（已中断）` : stage.label,
        detail: failed && state.error ? `${stage.detail}；${state.error}` : stage.detail,
        status: complete ? ("complete" as const) : ("active" as const),
        ...(stage.id === "write" && state.thinking.trim() ? { thinking: state.thinking } : {}),
      },
    };
  });
}

export type ResearchFeedMessage = {
  id: string;
  type: "ai";
  content: string;
  tool_calls: {
    id: string;
    name: string;
    args: Record<string, unknown>;
  }[];
};

/**
 * 组装聊天页要注入的两条虚拟 ai 消息：进度步骤条 + report.md 的 write_file
 * stub（沙箱 write-file 虚拟 URL 预览的数据源）。仅在报告尚无正文或仍在
 * 流式时输出 stub —— 完成后由真实 artifacts 接管。
 */
export function buildResearchFeedMessages(state: ResearchReportStreamState): ResearchFeedMessage[] {
  const messages: ResearchFeedMessage[] = [
    {
      id: `research-progress-${state.jobId}`,
      type: "ai",
      content: "",
      tool_calls: buildResearchProgressToolCalls(state),
    },
  ];
  if (state.reportText.trim()) {
    messages.push({
      id: `research-report-${state.jobId}`,
      type: "ai",
      content: "",
      tool_calls: [
        {
          id: `research-report-${state.jobId}-write`,
          name: "write_file",
          args: { path: RESEARCH_REPORT_OUTPUT_PATH, content: state.reportText },
        },
      ],
    });
  }
  return messages;
}

/** write-file 虚拟 URL（ArtifactFileDetail 用它从 stub 消息读流式内容）。 */
export function buildResearchReportVirtualUrl(state: ResearchReportStreamState): string {
  return new URL(
    `write-file:${RESEARCH_REPORT_OUTPUT_PATH}?message_id=${encodeURIComponent(`research-report-${state.jobId}`)}&tool_call_id=${encodeURIComponent(`research-report-${state.jobId}-write`)}&language=markdown`,
  ).toString();
}