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

import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { useEffect, useMemo, useRef, useState } from "react";

import { useThread } from "@/components/workspace/messages/context";

import {
  loadArtifactContent,
  loadArtifactContentFromLatestWriteFile,
  loadArtifactContentFromToolCall,
} from "./loader";
import {
  SANDBOX_STREAM_FLUSH_MS,
  isSandboxStreamPerfEnabled,
} from "./sandbox-stream-perf";
import { isArtifactApiPath } from "./streaming-write-preview";
import {
  isWriteFileInFlight,
  isWriteFileStreamingUrl,
  latestCompletedArtifactMutationKey,
} from "./write-file-in-flight";
import { updateArtifactContent } from "./api";

function hasToolResult(
  thread: ReturnType<typeof useThread>["thread"],
  toolCallId?: string,
) {
  if (!toolCallId) {
    return false;
  }
  return thread.messages.some(
    (message) => message.type === "tool" && message.tool_call_id === toolCallId,
  );
}

function applyPendingTextEdit({
  content,
  filepath,
  thread,
}: {
  content: string | undefined;
  filepath: string;
  thread: ReturnType<typeof useThread>["thread"];
}) {
  if (!content || !thread.isLoading) {
    return null;
  }

  for (const message of [...thread.messages].reverse()) {
    if (message.type !== "ai" || !message.tool_calls) {
      continue;
    }
    for (const toolCall of [...message.tool_calls].reverse()) {
      if (toolCall.name !== "str_replace" || hasToolResult(thread, toolCall.id)) {
        continue;
      }
      const args = toolCall.args as {
        path?: unknown;
        old_str?: unknown;
        new_str?: unknown;
        replace_all?: unknown;
      };
      if (args.path !== filepath) {
        continue;
      }
      if (typeof args.old_str !== "string" || typeof args.new_str !== "string") {
        continue;
      }
      if (!content.includes(args.old_str)) {
        return { content, kind: "editing" as const };
      }
      const nextContent =
        args.replace_all === true
          ? content.split(args.old_str).join(args.new_str)
          : content.replace(args.old_str, args.new_str);
      return { content: nextContent, kind: "editing" as const };
    }
  }
  return null;
}

/**
 * Keep the latest write_file text in a ref (always current for later readers)
 * but only publish a display snapshot on a timer while the run is in flight.
 * Display therefore does not follow `thread` object identity / every SSE chunk.
 */
function useThrottledWriteFileDisplay({
  filepath,
  liveContent,
  enabled,
  streaming,
}: {
  filepath: string;
  liveContent: string | undefined;
  enabled: boolean;
  streaming: boolean;
}) {
  const latestRef = useRef(liveContent);
  latestRef.current = liveContent;
  const [displayed, setDisplayed] = useState(liveContent);

  useEffect(() => {
    if (!enabled) {
      setDisplayed(latestRef.current);
      return;
    }
    const flush = () => {
      setDisplayed((previous) =>
        Object.is(previous, latestRef.current) ? previous : latestRef.current,
      );
    };
    if (!streaming) {
      flush();
      return;
    }
    flush();
    const intervalId = window.setInterval(flush, SANDBOX_STREAM_FLUSH_MS);
    return () => window.clearInterval(intervalId);
  }, [enabled, filepath, streaming]);

  return enabled && streaming ? displayed : liveContent;
}

export function useArtifactContent({
  filepath,
  threadId,
  enabled,
}: {
  filepath: string;
  threadId: string;
  enabled?: boolean;
}) {
  const useWriteFileStreaming = useMemo(() => {
    return isWriteFileStreamingUrl(filepath);
  }, [filepath]);
  const { thread, isMock, publicShareCode } = useThread();
  const streamPerfEnabled = isSandboxStreamPerfEnabled();
  const liveWriteFileContent = useWriteFileStreaming
    ? loadArtifactContentFromToolCall({ url: filepath, thread })
    : undefined;
  const writeFileInFlight =
    useWriteFileStreaming && isWriteFileInFlight(thread, filepath);
  const completedArtifactMutationKey = useMemo(() => {
    if (useWriteFileStreaming || !isArtifactApiPath(filepath)) {
      return null;
    }
    return latestCompletedArtifactMutationKey(thread, filepath);
  }, [filepath, thread, useWriteFileStreaming]);
  const displayedWriteFileContent = useThrottledWriteFileDisplay({
    filepath,
    liveContent: liveWriteFileContent,
    enabled: streamPerfEnabled && useWriteFileStreaming,
    streaming: writeFileInFlight,
  });
  const content = useWriteFileStreaming ? displayedWriteFileContent : null;
  const toolCallFallbackContent = useMemo(() => {
    if (useWriteFileStreaming || !isArtifactApiPath(filepath)) {
      return undefined;
    }
    return loadArtifactContentFromLatestWriteFile({ filepath, thread });
  }, [filepath, thread, useWriteFileStreaming]);

  const { data, isLoading, error, refetch } = useQuery({
    queryKey: [
      "artifact",
      filepath,
      threadId,
      isMock,
      publicShareCode,
      completedArtifactMutationKey,
    ],
    queryFn: () => {
      return loadArtifactContent({ filepath, threadId, isMock, publicShareCode });
    },
    enabled: Boolean(enabled) && !useWriteFileStreaming && isArtifactApiPath(filepath),
    refetchInterval: thread.isLoading ? 1000 : false,
    // Cache artifact content for 5 minutes to avoid repeated fetches (especially for .skill ZIP extraction)
    staleTime: 5 * 60 * 1000,
  });
  const pendingEdit = useMemo(() => {
    if (useWriteFileStreaming) {
      return null;
    }
    return applyPendingTextEdit({
      content: data?.content,
      filepath,
      thread,
    });
  }, [data?.content, filepath, useWriteFileStreaming, thread]);

  return {
    content: useWriteFileStreaming
      ? content
      : (pendingEdit?.content ?? data?.content ?? toolCallFallbackContent),
    url: useWriteFileStreaming ? undefined : data?.url,
    isLoading,
    isApplyingEdit: pendingEdit?.kind === "editing",
    error,
    refetch,
  };
}

export function useUpdateArtifactContent() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: updateArtifactContent,
    onSuccess: (content, variables) => {
      // useArtifactContent 的查询键带第 5 个元素 publicShareCode，这里用前缀
      // 过滤（setQueriesData）匹配，避免精确键缺少 publicShareCode 写到幽灵缓存、
      // 导致保存后内容被旧缓存重置回原值。
      queryClient.setQueriesData(
        {
          queryKey: [
            "artifact",
            variables.filepath,
            variables.threadId,
            variables.isMock,
          ],
        },
        {
          content,
          url: undefined,
        },
      );
    },
  });
}