deerflow-code/frontend-web/src/core/studio/hooks.ts

import { useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { useStream } from "@langchain/langgraph-sdk/react";

import { getStoredAuth } from "@/core/auth";
import { getAPIClient } from "@/core/api";
import {
  fetchNotebookSources,
  fetchStudioNotebooks,
  fetchStudioToken,
} from "./api";
import type { AgentThreadState } from "@/core/threads/types";
import type { SourceItem, StudioNotebook } from "./types";

export function useStudioNotebooks(enabled = true): {
  notebooks: StudioNotebook[];
  isLoading: boolean;
  isError: boolean;
  token: string | undefined;
} {
  const email = getStoredAuth()?.email ?? "";

  const tokenQuery = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: enabled && Boolean(email),
    staleTime: 5 * 60 * 1000,
    // 账号不存在（404/422）是永久错误，重试无意义，立即让 isError 落地
    retry: false,
    refetchOnWindowFocus: false,
  });

  const notebooksQuery = useQuery({
    queryKey: ["studio", "notebooks", tokenQuery.data],
    queryFn: () => fetchStudioNotebooks(tokenQuery.data!),
    enabled: enabled && Boolean(tokenQuery.data),
    // 空间列表不会高频变动；新建/删除路径已通过 invalidateQueries 主动失效
    // (见 workspace-nav-studio.tsx)，不需要短 staleTime + windowFocus 双兜底。
    staleTime: 10 * 60 * 1000,
    retry: 1,
    refetchOnWindowFocus: false,
  });

  return {
    notebooks: (notebooksQuery.data ?? []).filter((nb) => nb.is_owner),
    isLoading: tokenQuery.isLoading || notebooksQuery.isLoading,
    isError: tokenQuery.isError || notebooksQuery.isError,
    token: tokenQuery.data,
  };
}

export function useNotebookThreads(notebookId: number | null) {
  const client = getAPIClient();
  return useQuery({
    queryKey: ["notebook-threads", notebookId],
    queryFn: () =>
      client.threads.search({
        metadata: { thread_type: "notebook", note_id: notebookId },
        limit: 20,
        sortBy: "updated_at",
        sortOrder: "desc",
      }),
    enabled: notebookId !== null,
    staleTime: 30 * 1000,
    refetchOnWindowFocus: false,
  });
}

export function useNotebookChat(
  notebookId: number,
  sources: SourceItem[],
  modelName?: string,
  initialThreadId?: string,
) {
  const [threadId, setThreadId] = useState<string | undefined>(initialThreadId);

  const email = getStoredAuth()?.email ?? "";
  const { data: token } = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: Boolean(email),
    staleTime: 5 * 60 * 1000,
    retry: false,
    refetchOnWindowFocus: false,
  });

  const thread = useStream<AgentThreadState>({
    client: getAPIClient(),
    assistantId: "lead_agent",
    threadId,
    fetchStateHistory: { limit: 1 },
    onCreated(meta) {
      setThreadId(meta.thread_id);
      void getAPIClient()
        .threads.update(meta.thread_id, {
          metadata: { thread_type: "notebook", note_id: notebookId },
        })
        .catch(() => ({}));
    },
  });

  const send = async (
    text: string,
    selectedSourceIds: number[],
    options?: { fromClarification?: boolean },
  ) => {
    if (!token) return;
    // 协助卡提交前会先停轮；此时 isLoading 可能仍为 true，不能直接丢弃。
    if (thread.isLoading && !options?.fromClarification) return;
    // API 需要 document.id，而不是 source.id
    const toDocIds = (ids: number[]) =>
      ids
        .map((id) => sources.find((s) => s.id === id)?.document?.id)
        .filter((id): id is number => id !== undefined);
    const sourceIds =
      selectedSourceIds.length > 0
        ? toDocIds(selectedSourceIds)
        : sources
            .map((s) => s.document?.id)
            .filter((id): id is number => id !== undefined);

    await thread.submit(
      { messages: [{ type: "human", content: text }] },
      {
        threadId,
        streamSubgraphs: true,
        config: {
          configurable: {
            notebook_search_space_id: notebookId,
            notebook_source_ids: sourceIds,
            notebook_token: token,
            ...(modelName ? { model_name: modelName } : {}),
          },
        },
      },
    );
  };

  return {
    thread,
    messages: thread.messages,
    isStreaming: thread.isLoading,
    send,
    threadId,
    setThreadId,
  };
}

export function useNotebookSources(
  notebookId: number | null,
  overrideToken?: string,
): {
  sources: SourceItem[];
  isLoading: boolean;
  isError: boolean;
} {
  const email = getStoredAuth()?.email ?? "";

  // 若外部直接提供 user_token（如 notebook-token 接口返回的），跳过 getTokenByEmail
  const tokenQuery = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: !overrideToken && Boolean(email),
    staleTime: 5 * 60 * 1000,
    retry: false,
    refetchOnWindowFocus: false,
  });

  const resolvedToken = overrideToken ?? tokenQuery.data;

  const sourcesQuery = useQuery({
    queryKey: ["studio", "sources", notebookId, overrideToken ?? "default"],
    queryFn: () => fetchNotebookSources(notebookId!, resolvedToken!),
    enabled: notebookId !== null && Boolean(resolvedToken),
    // 笔记文档列表低频变动；增删时调用方应主动 invalidateQueries(["studio", "sources"]) 失效。
    staleTime: 5 * 60 * 1000,
    retry: 1,
    refetchOnWindowFocus: false,
  });

  return {
    sources: sourcesQuery.data ?? [],
    isLoading:
      (!overrideToken && tokenQuery.isLoading) || sourcesQuery.isLoading,
    isError: (!overrideToken && tokenQuery.isError) || sourcesQuery.isError,
  };
}