@amaster.ai/components-templates 1.4.1 → 1.4.2

package/components/ai-assistant/template/hooks/useConversation.ts

@@ -0,0 +1,856 @@
+import { useState, useCallback, useMemo, useRef, useEffect } from "react";
+import { produce } from "immer";
+import type {
+  MessagesItem,
+  Conversation,
+  Role,
+  TextMessage,
+  ThoughtMessage,
+  ToolMessage,
+  ErrorMessage,
+  UIRenderMessage,
+} from "../types";
+import type {
+  Message1,
+  SendStreamingMessageResponse,
+  SendStreamingMessageSuccessResponse,
+} from "@a2a-js/sdk";
+import { client } from "@/lib/client";
+import { getText } from "../i18n";
+
+interface HistoryState {
+  next: string | null;
+  hasMore: boolean;
+}
+
+export interface UpdateMessageInput {
+  taskId: string;
+  messageId: string;
+  content?: string;
+  partial?: boolean;
+  status?: string;
+  response?: any;
+  metadata?: Record<string, any>;
+}
+
+type MessageHandler = (
+  conv: Conversation,
+  part: any,
+  id: string,
+  role: Role,
+  status?: any,
+) => Conversation;
+
+const handlers: Record<string, MessageHandler> = {
+  "text-content": (conv, part, id, role, status) => {
+    if (!id) return conv;
+    const existing = conv.messages.find((m) => m.messageId === id);
+    const newText = part.text || "";
+    if (existing && "content" in existing) {
+      return {
+        ...conv,
+        messages: conv.messages.map((m) =>
+          m.messageId === id
+            ? { ...m, content: (m as TextMessage).content + newText }
+            : m,
+        ),
+        lastUpdated: new Date().toISOString(),
+      };
+    }
+
+    return {
+      ...conv,
+      messages: [
+        ...conv.messages,
+        {
+          messageId: id,
+          role,
+          kind: "text-content",
+          content: newText,
+          timestamp: status.timestamp || new Date().toISOString(),
+        } as TextMessage,
+      ],
+      lastUpdated: new Date().toISOString(),
+    };
+  },
+
+  thought: (conv, part, id, role, status) => {
+    if (!id) return conv;
+    const newChunk = part?.data?.description || "";
+    const existing = conv.messages.find((m) => m.messageId === id);
+    if (existing && "thought" in existing) {
+      return {
+        ...conv,
+        messages: conv.messages.map((m) =>
+          m.messageId === id
+            ? { ...m, thought: (m as ThoughtMessage).thought + newChunk }
+            : m,
+        ),
+        lastUpdated: new Date().toISOString(),
+      };
+    }
+
+    return {
+      ...conv,
+      messages: [
+        ...conv.messages,
+        {
+          messageId: id,
+          role,
+          kind: "thought",
+          thought: newChunk,
+          timestamp: status?.timestamp || new Date().toISOString(),
+        } as ThoughtMessage,
+      ],
+      lastUpdated: new Date().toISOString(),
+    };
+  },
+
+  tool: (conv, part, _id, role, status) => {
+    const partData = part?.data || {};
+    const tool = partData?.tool;
+    const name = tool?.displayName || tool?.name || "";
+    const request = partData?.request;
+    const id = request?.callId;
+    if (!id || !name) return conv;
+
+    const toolStatus = partData?.status;
+    const existing = conv.messages.find((m) => m.messageId === id);
+
+    if (existing && "toolStatus" in existing) {
+      return {
+        ...conv,
+        messages: conv.messages.map((m) =>
+          m.messageId === id
+            ? {
+                ...m,
+                toolStatus: toolStatus || (m as ToolMessage).toolStatus,
+              }
+            : m,
+        ),
+        lastUpdated: new Date().toISOString(),
+      };
+    }
+
+    return {
+      ...conv,
+      messages: [
+        ...conv.messages,
+        {
+          messageId: id,
+          role,
+          kind: "tool",
+          toolName: name,
+          toolDescription: request?.args?.query || "",
+          toolStatus: toolStatus || "pending",
+          timestamp: status?.timestamp || new Date().toISOString(),
+        } as ToolMessage,
+      ],
+      lastUpdated: new Date().toISOString(),
+    };
+  },
+
+  error: (conv, part, id, role, status) => {
+    const errorMsg = part?.text || part?.data?.error || "发生错误";
+    return {
+      ...conv,
+      messages: [
+        ...conv.messages,
+        {
+          messageId: id || `error-${Date.now()}`,
+          role,
+          kind: "error",
+          content: errorMsg,
+          timestamp: new Date().toISOString(),
+        } as ErrorMessage,
+      ],
+      status: "error",
+      lastUpdated: new Date().toISOString(),
+    };
+  },
+
+  "ui-render": (conv, part, id, role, status) => {
+    if (!id) return conv;
+    const existing = conv.messages.find((m) => m.messageId === id);
+    const partData = part?.data || {};
+    const spec = partData.spec || { root: "", elements: {} };
+
+    if (existing && existing.kind === "ui-render") {
+      return {
+        ...conv,
+        messages: conv.messages.map((m) =>
+          m.messageId === id ? { ...(m as UIRenderMessage), spec } : m,
+        ),
+        lastUpdated: new Date().toISOString(),
+      };
+    }
+
+    return {
+      ...conv,
+      messages: [
+        ...conv.messages,
+        {
+          messageId: id,
+          role,
+          kind: "ui-render",
+          spec,
+          timestamp: status?.timestamp || new Date().toISOString(),
+        } as UIRenderMessage,
+      ],
+      lastUpdated: new Date().toISOString(),
+    };
+  },
+};
+
+const getKind = (part: any, metadata?: any) => {
+  let kind = "unknown";
+  if (part?.data?.tool) {
+    kind = "tool";
+  } else if (part?.data?.type === "ui") {
+    kind = "ui-render";
+  } else if (part?.kind === "text") {
+    kind = "text-content";
+  } else if (
+    part?.data?.type === "though" ||
+    metadata?.coderAgent?.kind === "thought"
+  ) {
+    kind = "thought";
+  } else if (metadata?.error) {
+    kind = "error";
+  }
+  return kind;
+};
+
+const generateId = () =>
+  `msg-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
+
+type ConversationsState = {
+  byId: Record<string, Conversation>;
+  order: string[];
+};
+
+export function useConversation() {
+  const [state, setState] = useState<ConversationsState>({
+    byId: {},
+    order: [],
+  });
+  const [isLoading, setIsLoading] = useState(false);
+  const [isLoadingHistory, setIsLoadingHistory] = useState(false);
+  const [historyState, setHistoryState] = useState<HistoryState>({
+    next: null,
+    hasMore: false,
+  });
+  const loadingRef = useRef(false);
+  const abortControllerRef = useRef<AbortController | null>(null);
+  const forceStopRef = useRef(false);
+  const lastMessageTypeRef = useRef<string>("unknown");
+  const lastMessageIdRef = useRef<string>("");
+  const lastMessageRoleRef = useRef<Role | string>("");
+  const currentStreamingTaskIdRef = useRef<string | null>(null);
+  const [starting, setStarting] = useState(true);
+
+  useEffect(() => {
+    loadingRef.current = isLoading;
+  }, [isLoading]);
+
+  const abort = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+    }
+    forceStopRef.current = true;
+    setIsLoading(false);
+  }, []);
+
+  const handleEnd = useCallback(() => {
+    abort();
+    lastMessageTypeRef.current = "";
+    lastMessageIdRef.current = "";
+    lastMessageRoleRef.current = "";
+    currentStreamingTaskIdRef.current = null;
+  }, [abort]);
+
+  // 填充数据到 conversation 中
+  const fillData = ({
+    kind,
+    taskId,
+    part,
+    status,
+    handler,
+    isHistory,
+    historyId,
+    messageId,
+  }: {
+    kind: string;
+    taskId: string;
+    part: any;
+    status: any;
+    handler: MessageHandler;
+    isHistory?: boolean;
+    historyId?: string;
+    messageId?: string;
+  }) => {
+    setState(
+      produce((draft) => {
+        let id = messageId || lastMessageIdRef.current;
+        const role = status?.message?.role === "agent" ? "assistant" : "user";
+
+        if (
+          !messageId &&
+          (kind === "ui-render" ||
+            (kind !== "unknown" && kind !== lastMessageTypeRef.current))
+        ) {
+          lastMessageTypeRef.current = kind;
+          id = generateId();
+          lastMessageIdRef.current = id;
+        }
+
+        const prevConv = draft.byId[taskId] ?? {
+          taskId,
+          status: isHistory ? "completed" : "submitted",
+          messages: [],
+          role,
+          historyId,
+          lastUpdated: new Date().toISOString(),
+        };
+
+        const nextConv = handler(prevConv, part, id, role, status);
+
+        draft.byId[taskId] = {
+          ...nextConv,
+          status: status?.state || prevConv.status,
+          lastUpdated: new Date().toISOString(),
+        };
+
+        const index = draft.order.indexOf(taskId);
+        if (index !== -1) {
+          draft.order.splice(index, 1);
+        }
+
+        if (!isHistory) {
+          draft.order.push(taskId);
+        }
+      }),
+    );
+  };
+
+  // 处理实时数据
+  const processLiveData = useCallback(
+    (data: SendStreamingMessageResponse) => {
+      if (!(data as any)?.result) return;
+      const result = (data as SendStreamingMessageSuccessResponse).result;
+
+      const { taskId, status, metadata, final } = result as any;
+
+      if (final) {
+        handleEnd();
+        return;
+      }
+
+      if (!taskId) return;
+
+      currentStreamingTaskIdRef.current = taskId;
+
+      const parts = status?.message?.parts || [];
+      const firstPart = parts?.[0];
+      const kind = getKind(firstPart, metadata);
+
+      const handler = handlers[kind];
+      if (!handler) {
+        return;
+      }
+      removeLoadingPlaceholder();
+
+      fillData({
+        kind,
+        taskId,
+        part: firstPart,
+        status,
+        handler,
+      });
+    },
+    [handleEnd],
+  );
+
+  const processHistoryData = useCallback(
+    (data: Message1, taskId: string) => {
+      const { messageId, historyId } = (data || {}) as any;
+      if (!messageId) return;
+
+      const parts = data.parts || [];
+      const firstPart = parts?.[0];
+
+      const kind = getKind(firstPart);
+
+      const handler = handlers[kind];
+      if (!handler) {
+        return;
+      }
+      fillData({
+        kind,
+        taskId,
+        part: firstPart,
+        handler,
+        isHistory: true,
+        historyId,
+        messageId,
+        status: { message: data },
+      });
+    },
+    [handleEnd],
+  );
+
+  const conversations = useMemo(() => {
+    return state.order
+      .map((id) => state.byId[id])
+      .filter(Boolean) as Conversation[];
+  }, [state]);
+
+  const getConversation = useCallback(
+    (taskId: string) => {
+      return state.byId[taskId];
+    },
+    [state],
+  );
+
+  const clearConversation = useCallback((taskId: string) => {
+    setState(
+      produce((draft) => {
+        if (!draft.byId[taskId]) return;
+        delete draft.byId[taskId];
+        draft.order = draft.order.filter((id) => id !== taskId);
+      }),
+    );
+  }, []);
+
+  const removeMessage = useCallback((taskId: string, messageId: string) => {
+    setState(
+      produce((draft) => {
+        const conv = draft.byId[taskId];
+        if (!conv) return;
+
+        const nextMessages = conv.messages.filter(
+          (m) => m.messageId !== messageId,
+        );
+
+        draft.byId[taskId] = {
+          ...conv,
+          messages: nextMessages,
+          lastUpdated: new Date().toISOString(),
+        };
+
+        const index = draft.order.indexOf(taskId);
+        if (index !== -1) {
+          draft.order.splice(index, 1);
+        }
+        draft.order.push(taskId);
+      }),
+    );
+  }, []);
+
+  const updateMessage = useCallback((input: UpdateMessageInput) => {
+    const {
+      taskId,
+      messageId,
+      content,
+      partial = false,
+      status,
+      response,
+      metadata = {},
+    } = input;
+
+    setState(
+      produce((draft) => {
+        const conv = draft.byId[taskId];
+        if (!conv) return;
+
+        const msgIndex = conv.messages.findIndex(
+          (m) => m.messageId === messageId,
+        );
+        if (msgIndex === -1) {
+          return;
+        }
+
+        const oldMsg = conv.messages[msgIndex];
+        let updatedMsg: MessagesItem = { ...oldMsg, ...metadata };
+
+        if (content !== undefined) {
+          if ("content" in updatedMsg) {
+            updatedMsg.content = partial
+              ? (updatedMsg as TextMessage).content + content
+              : content;
+          } else if ("thought" in updatedMsg) {
+            updatedMsg.thought = partial
+              ? (updatedMsg as ThoughtMessage).thought + content
+              : content;
+          } else if (oldMsg.kind === "text-content") {
+            (updatedMsg as any).content = partial
+              ? ((oldMsg as any).content || "") + content
+              : content;
+          }
+        }
+
+        if ("toolStatus" in updatedMsg || status) {
+          (updatedMsg as any).toolStatus =
+            status || (updatedMsg as any).toolStatus;
+        }
+        if (response !== undefined) {
+          (updatedMsg as any).response = response;
+        }
+
+        const nextMessages = [...conv.messages];
+        nextMessages[msgIndex] = updatedMsg;
+
+        draft.byId[taskId] = {
+          ...conv,
+          messages: nextMessages,
+          lastUpdated: new Date().toISOString(),
+        };
+
+        const index = draft.order.indexOf(taskId);
+        if (index !== -1) {
+          draft.order.splice(index, 1);
+        }
+        draft.order.push(taskId);
+      }),
+    );
+  }, []);
+
+  const addConversation = useCallback((conversation: Conversation) => {
+    setState(
+      produce((draft) => {
+        const id = conversation.taskId;
+        draft.byId[id] = conversation;
+        const index = draft.order.indexOf(id);
+        if (index !== -1) {
+          draft.order.splice(index, 1);
+        }
+        draft.order.push(id);
+      }),
+    );
+  }, []);
+
+  const addLoadingPlaceholder = useCallback(() => {
+    addConversation({
+      taskId: "loading-placeholder",
+      status: "working",
+      messages: [
+        {
+          messageId: "loading-placeholder",
+          role: "assistant",
+          kind: "text-content",
+          content: getText().thinking,
+          timestamp: new Date().toISOString(),
+        },
+      ],
+      lastUpdated: new Date().toISOString(),
+    });
+  }, [addConversation]);
+
+  const removeLoadingPlaceholder = useCallback(() => {
+    clearConversation("loading-placeholder");
+  }, [clearConversation]);
+
+  useEffect(() => {
+    if (!isLoading) {
+      removeLoadingPlaceholder();
+    }
+  }, [isLoading, removeLoadingPlaceholder]);
+
+  const onSendError = useCallback(() => {
+    addConversation({
+      taskId: generateId(),
+      status: "error",
+      messages: [
+        {
+          messageId: `error-${Date.now()}`,
+          role: "assistant",
+          kind: "error",
+          content: getText().errorMessage,
+          timestamp: new Date().toISOString(),
+        } as ErrorMessage,
+      ],
+      lastUpdated: new Date().toISOString(),
+    });
+    setIsLoading(false);
+  }, [addConversation]);
+
+  const loadHistory = useCallback(async (limit = 10, next?: string) => {
+    setIsLoadingHistory(true);
+    let scrollToId = "";
+    try {
+      const copilot = client.copilot as any;
+      const result = await copilot.getHistory?.(limit, next);
+      let messages = result?.messages || [];
+      const length = messages.length;
+      scrollToId = length > 0 ? messages?.[length - 1]?.messageId : "";
+
+      let taskId = "";
+      let lastRole = "";
+      let taskIds: string[] = [];
+      messages.forEach((msg: any) => {
+        let role = msg.role === "agent" ? "assistant" : "user";
+        if (role !== lastRole) {
+          taskId = `history-conv-${generateId()}`;
+          taskIds.push(taskId);
+          lastRole = role;
+        }
+        processHistoryData(msg, taskId);
+      });
+      setState(
+        produce((draft) => {
+          draft.order = [...taskIds, ...draft.order];
+        }),
+      );
+      setHistoryState({
+        next: result?.next || null,
+        hasMore: result?.hasMore || false,
+      });
+    } catch (error) {
+      console.error("[Chat] Failed to load history:", error);
+    } finally {
+      setIsLoadingHistory(false);
+      scrollToId &&
+        next &&
+        setTimeout(() => {
+          document
+            .querySelector(`[data-conversation-message-id="${scrollToId}"]`)
+            ?.scrollIntoView({ behavior: "instant", block: "center" });
+        }, 20);
+    }
+  }, []);
+
+  const loadMoreHistory = useCallback(async () => {
+    if (!historyState.hasMore || isLoadingHistory || !historyState.next) return;
+    await loadHistory(20, historyState.next);
+  }, [historyState.hasMore, historyState.next, isLoadingHistory]);
+
+  const checkActiveTask = useCallback(async () => {
+    try {
+      const result = await client.copilot.getChatStatus();
+      const taskId = result.taskId;
+      if (taskId && result.working) {
+        currentStreamingTaskIdRef.current = taskId;
+
+        const stream = await client.copilot.chat([], { taskId });
+        for await (const chunk of stream) {
+          if (forceStopRef.current) break;
+          processLiveData(chunk);
+        }
+      }
+    } catch (error) {
+      console.error("[Chat] Failed to check active task:", error);
+    }
+  }, [processLiveData]);
+
+  const cancelChat = useCallback(
+    async (taskId?: string) => {
+      abort();
+      const targetTaskId = taskId || currentStreamingTaskIdRef.current;
+      if (!targetTaskId) return;
+
+      try {
+        await client.copilot.cancelChat(targetTaskId);
+        currentStreamingTaskIdRef.current = null;
+      } catch (error) {
+        console.error("[Chat] Failed to cancel chat:", error);
+      }
+    },
+    [abort],
+  );
+
+  const resetConversation = useCallback(
+    (greeting?: string) => {
+      abort();
+      setState({ byId: {}, order: [] });
+      currentStreamingTaskIdRef.current = null;
+
+      if (greeting) {
+        const taskId = `conv-${Date.now()}`;
+        setState({
+          byId: {
+            [taskId]: {
+              taskId,
+              status: "submitted",
+              messages: [
+                {
+                  messageId: `greeting-${Date.now()}`,
+                  role: "assistant",
+                  kind: "text-content",
+                  content: greeting,
+                  timestamp: new Date().toISOString(),
+                } as TextMessage,
+              ],
+              lastUpdated: new Date().toISOString(),
+            },
+          },
+          order: [taskId],
+        });
+      }
+    },
+    [abort],
+  );
+
+  const startNewConversation = useCallback(async () => {
+    addConversation({
+      taskId: generateId(),
+      status: "submitted",
+      messages: [],
+      lastUpdated: new Date().toISOString(),
+      system: { level: "newConversation" },
+    });
+    try {
+      const copilot = client.copilot as any;
+      await copilot.newConversation?.();
+    } catch (error) {
+      console.error("[Chat] Failed to create new conversation:", error);
+    }
+  }, []);
+
+  const subscribeChatEvents = useCallback(
+    async ({
+      userContent,
+      taskId,
+    }: {
+      userContent?: string;
+      taskId?: string;
+    }) => {
+      let hasResponse = false;
+
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+      }
+      if (typeof AbortController !== "undefined") {
+        abortControllerRef.current = new AbortController();
+      }
+      const controller = abortControllerRef.current;
+      forceStopRef.current = false;
+
+      // mock
+      if (userContent === "MAGA") {
+        const { generateMockUIStream } = await import("../mock/mock-data");
+        const cancelMock = generateMockUIStream(
+          taskId || generateId(),
+          (chunk) => {
+            if (forceStopRef.current) return;
+            hasResponse = true;
+            processLiveData(chunk as any);
+            if (chunk.result.final) {
+              setIsLoading(false);
+            }
+          },
+        );
+        abortControllerRef.current?.signal.addEventListener(
+          "abort",
+          cancelMock,
+        );
+        return;
+      }
+
+      // real chat
+      try {
+        const stream = client.copilot.chat(
+          userContent ? [{ role: "user", content: userContent }] : [],
+          taskId ? { taskId } : undefined,
+        );
+
+        for await (const chunk of stream) {
+          if (controller?.signal.aborted || forceStopRef.current) break;
+          hasResponse = true;
+
+          processLiveData(chunk);
+
+          if ((chunk as any)?.isFinal || (chunk as any)?.final) {
+            break;
+          }
+        }
+      } catch (err: any) {
+        if (err.name === "AbortError" || forceStopRef.current) {
+          console.error("[Chat] Request aborted", err);
+        } else {
+          console.error("[Chat] Stream error:", err);
+          onSendError();
+        }
+      } finally {
+        abort();
+        if (!hasResponse) {
+          onSendError();
+        }
+      }
+    },
+    [],
+  );
+
+  const sendMessage = useCallback(
+    async (userContent: string) => {
+      if (loadingRef.current) {
+        return;
+      }
+      if (!userContent.trim()) return;
+      const userMsgId = generateId();
+
+      addConversation({
+        taskId: generateId(),
+        status: "submitted",
+        messages: [
+          {
+            messageId: userMsgId,
+            role: "user",
+            kind: "text-content",
+            content: userContent,
+            timestamp: new Date().toISOString(),
+          },
+        ],
+        lastUpdated: new Date().toISOString(),
+      });
+
+      addLoadingPlaceholder();
+
+      setIsLoading(true);
+      subscribeChatEvents({ userContent });
+    },
+    [
+      addConversation,
+      addLoadingPlaceholder,
+      processLiveData,
+      onSendError,
+      abort,
+    ],
+  );
+
+  useEffect(() => {
+    if (starting) {
+      const init = async () => {
+        console.debug("[Chat] Initializing conversation...");
+        try {
+          await loadHistory();
+          await checkActiveTask();
+          setStarting(false);
+        } catch (error) {
+          console.error("[Chat] Initialization error:", error);
+          setStarting(false);
+        }
+      };
+      init();
+    }
+  }, [starting, loadHistory, checkActiveTask]);
+
+  return {
+    starting,
+    isLoading,
+    isLoadingHistory,
+    historyState,
+    conversations,
+    sendMessage,
+    abort,
+    cancelChat,
+    resetConversation,
+    startNewConversation,
+    loadHistory,
+    loadMoreHistory,
+    getConversation,
+    clearConversation,
+    removeMessage,
+    updateMessage,
+  };
+}