deerflow-code/frontend-web/src/open-canvas/hooks/useCanvasThreadStream.ts

// Canvas-specific variant of useThreadStream (core/threads/hooks.ts).
// Copied and adapted for canvas_agent so canvas writing-mode interactions
// can be customized independently from the main lead_agent flow.

import type { AIMessage, Message } from "@langchain/langgraph-sdk";
import { useStream } from "@langchain/langgraph-sdk/react";
import { useCallback, useEffect, useRef, useState } from "react";
import { useNavigate } from "react-router-dom";
import { toast } from "sonner";

import { getCanvasAPIClient } from "@/core/api/canvas-client";
import { useUpdateSubtask } from "@/core/tasks/context";

import { CANVAS_AGENT_ID, CANVAS_THREAD_METADATA } from "../constants";
import { useCanvasGraphContext } from "../contexts/CanvasGraphContext";
import { useCanvasThreadContext } from "../contexts/CanvasThreadContext";
import type { StreamingArtifactPartial } from "../api/runs";
import type { CanvasArtifact } from "../types/artifact";
import type { CanvasThreadValues } from "../types/thread";
import { getCurrentArtifactContent } from "../utils/artifact";

function getStreamErrorMessage(error: unknown): string {
  if (typeof error === "string" && error.trim()) return error;
  if (error instanceof Error && error.message.trim()) return error.message;
  if (typeof error === "object" && error !== null) {
    const message = Reflect.get(error, "message");
    if (typeof message === "string" && message.trim()) return message;
  }
  return "Request failed.";
}

/**
 * Canvas-specific stream hook — a direct copy of the main agent's useThreadStream,
 * configured for canvas_agent. Modify sendMessage here to inject canvas-specific
 * context (writing_artifact_path, canvas_mode, highlighted text, etc.).
 */
export function useCanvasThreadStream() {
  const navigate = useNavigate();
  const { threadId, setThreadId, modelName } = useCanvasThreadContext();
  const { artifact, setArtifact, setIsStreaming, setStreamingArtifact } = useCanvasGraphContext();
  const updateSubtask = useUpdateSubtask();

  const [optimisticMessages, setOptimisticMessages] = useState<Message[]>([]);
  const sendInFlightRef = useRef(false);
  const prevMsgCountRef = useRef(0);

  const [onStreamThreadId, setOnStreamThreadId] = useState<string | null>(
    () => threadId,
  );
  const threadIdRef = useRef<string | null>(threadId ?? null);

  useEffect(() => {
    const normalized = threadId ?? null;
    setOnStreamThreadId(normalized);
    threadIdRef.current = normalized;
    if (!normalized) {
      sendInFlightRef.current = false;
      prevMsgCountRef.current = 0;
      setOptimisticMessages([]);
    }
  }, [threadId]);

  const thread = useStream<CanvasThreadValues>({
    client: getCanvasAPIClient(),
    assistantId: CANVAS_AGENT_ID,
    threadId: onStreamThreadId,
    fetchStateHistory: { limit: 1 },
    onCreated(meta) {
      setOnStreamThreadId(meta.thread_id);
      threadIdRef.current = meta.thread_id;
      setThreadId(meta.thread_id);
      navigate(`/page/canvas/${meta.thread_id}`, { replace: true });
      void getCanvasAPIClient()
        .threads.update(meta.thread_id, { metadata: CANVAS_THREAD_METADATA })
        .catch(() => {});
    },
    onUpdateEvent(data) {
      const updates: Array<Partial<CanvasThreadValues> | null> = Object.values(
        data || {},
      );
      for (const update of updates) {
        if (update?.artifact) {
          setArtifact(update.artifact as CanvasArtifact);
        }
      }
    },
    onCustomEvent(event: unknown) {
      if (
        typeof event === "object" &&
        event !== null &&
        "type" in event &&
        (event as { type: unknown }).type === "task_running"
      ) {
        const e = event as {
          type: "task_running";
          task_id: string;
          message: AIMessage;
        };
        updateSubtask({ id: e.task_id, latestMessage: e.message });
        return;
      }
      if (
        typeof event === "object" &&
        event !== null &&
        "type" in event &&
        (event as { type: unknown }).type === "llm_retry" &&
        "message" in event &&
        typeof (event as { message: unknown }).message === "string"
      ) {
        // Model retry events are self-healing. Keep them off the user-facing
        // toast surface so a short model stall does not look like a failure.
      }
    },
    onError(error) {
      setOptimisticMessages([]);
      toast.error(getStreamErrorMessage(error));
    },
  });

  // Keep CanvasGraphContext.isStreaming in sync so the artifact panel and
  // quick-action bar can react to streaming state.
  useEffect(() => {
    setIsStreaming(thread.isLoading);
  }, [thread.isLoading, setIsStreaming]);

  // Detect update_artifact tool calls from streaming messages to populate
  // streamingArtifact so the right panel shows content before the final
  // state update fires via onUpdateEvent.
  useEffect(() => {
    if (!thread.isLoading) {
      setStreamingArtifact(null);
      return;
    }
    const msgs = thread.messages;
    for (let i = msgs.length - 1; i >= 0; i--) {
      const msg = msgs[i];
      if (msg.type !== "ai") continue;
      const toolCalls = (msg as AIMessage).tool_calls;
      if (!toolCalls?.length) continue;
      const call = toolCalls.find((tc) => tc.name === "update_artifact");
      if (!call?.args) continue;
      const args = call.args as Record<string, unknown>;
      const partial: StreamingArtifactPartial = {
        type: args.type as "text" | "code" | undefined,
        title: typeof args.title === "string" ? args.title : undefined,
        content: typeof args.content === "string" ? args.content : undefined,
        language: typeof args.language === "string" ? args.language : undefined,
      };
      if (partial.title !== undefined || partial.content !== undefined) {
        setStreamingArtifact(partial);
        return;
      }
    }
  }, [thread.messages, thread.isLoading, setStreamingArtifact]);

  // Clear optimistic messages once the server stream begins delivering real ones.
  useEffect(() => {
    if (
      optimisticMessages.length > 0 &&
      thread.messages.length > prevMsgCountRef.current
    ) {
      setOptimisticMessages([]);
    }
  }, [thread.messages.length, optimisticMessages.length]);

  /**
   * Canvas writing-mode sendMessage — copied from ChatPage's handleSubmit flow.
   * Modify here to inject writing_artifact_path, canvas_mode, highlighted
   * text/code, or any other canvas-specific context the agent needs.
   */
  const sendMessage = useCallback(
    async (text: string, extraContext?: Record<string, unknown>) => {
      const trimmed = text.trim();
      if (!trimmed || sendInFlightRef.current) return;
      sendInFlightRef.current = true;

      prevMsgCountRef.current = thread.messages.length;
      setOptimisticMessages([
        {
          type: "human",
          id: `opt-human-${Date.now()}`,
          content: [{ type: "text", text: trimmed }],
          additional_kwargs: {},
        },
      ]);

      // Inject current artifact content into the human message so the LLM can see
      // what to edit. create_agent only passes state.messages to the LLM — state.artifact
      // is NOT automatically visible. Skip injection when the caller already embedded
      // the artifact (e.g. quick actions build the full message themselves).
      const currentContent = artifact ? getCurrentArtifactContent(artifact) : null;
      const artifactContext = currentContent
        ? {
            canvas_artifact_title: currentContent.title,
            canvas_artifact_type: currentContent.type,
          }
        : {};

      const messageText =
        currentContent && !trimmed.includes("<artifact")
          ? (() => {
              const body =
                currentContent.type === "code"
                  ? currentContent.code
                  : currentContent.fullMarkdown;
              const tag =
                currentContent.type === "code"
                  ? `language="${currentContent.language}"`
                  : "";
              return `${trimmed}\n\n<artifact ${tag}>\n${body}\n</artifact>`;
            })()
          : trimmed;

      try {
        await thread.submit(
          {
            messages: [
              { type: "human", content: [{ type: "text", text: messageText }] },
            ],
          },
          {
            threadId: threadIdRef.current ?? undefined,
            streamSubgraphs: true,
            streamResumable: true,
            config: {
              recursion_limit: 1000,
              ...(modelName ? { configurable: { model_name: modelName } } : {}),
            },
            context: {
              ...extraContext,
              ...artifactContext,
              thread_id: threadIdRef.current ?? "",
              writing_mode: true,
            },
          },
        );
      } catch {
        setOptimisticMessages([]);
      } finally {
        sendInFlightRef.current = false;
      }
    },
    [thread, modelName, artifact],
  );

  const mergedThread =
    optimisticMessages.length > 0
      ? ({
          ...thread,
          messages: [...thread.messages, ...optimisticMessages],
        } as typeof thread)
      : thread;

  return { thread: mergedThread, sendMessage, threadId };
}