feat: improve citation display and skill persistence

frontend-web/src/components/workspace/chats/chat-box.tsx

@@ -40,9 +40,11 @@ const ChatBox: React.FC<{ children: React.ReactNode; threadId: string }> = ({
   const { thread } = useThread();
   const pathname = usePathname();
   const threadIdRef = useRef(threadId);
+  const artifactsCountRef = useRef(0);
   const layoutRef = useRef<GroupImperativeHandle>(null);
   const [sideTab, setSideTab] = useState("document");
   const [referencesOpen, setReferencesOpen] = useState(true);
+  const [activeReferenceBatchId, setActiveReferenceBatchId] = useState<string | null>(null);
   const [settings] = useLocalSettings();
   const { skills } = useSkills();
 
@@ -61,24 +63,32 @@ const ChatBox: React.FC<{ children: React.ReactNode; threadId: string }> = ({
     () => (settings.display.ragModeEnabled ? buildSkillDisplayMap(skills) : new Map()),
     [settings.display.ragModeEnabled, skills],
   );
-  const latestHumanIndex = useMemo(
-    () => findLatestHumanMessageIndex(thread.values.messages ?? []),
-    [thread.values.messages],
-  );
-  const displaySources = useMemo(() => {
+  const referenceBatches = useMemo(() => {
     if (!settings.display.ragModeEnabled) {
       return [];
     }
     const messages = thread.values.messages ?? [];
-    return collectRuntimeDisplaySources(messages, skillDisplayMap, {
-      afterMessageId: latestHumanIndex >= 0 ? messages[latestHumanIndex]?.id : undefined,
-      includeModes: ["citation"],
-    });
-  }, [settings.display.ragModeEnabled, skillDisplayMap, latestHumanIndex, thread.values.messages]);
-  const referenceSources = useMemo(
-    () => displaySources.filter((source) => source.mode === "citation"),
-    [displaySources],
-  );
+    return messages
+      .filter((message) => message.type === "ai" && message.id)
+      .map((message) => ({
+        id: message.id!,
+        sources: collectRuntimeDisplaySources(messages, skillDisplayMap, {
+          afterMessageId: findPreviousHumanMessageId(messages, message.id),
+          beforeMessageId: message.id,
+          includeBeforeMessage: true,
+          includeModes: ["citation"],
+        }).filter((source) => source.mode === "citation"),
+      }))
+      .filter((batch) => batch.sources.length > 0);
+  }, [settings.display.ragModeEnabled, skillDisplayMap, thread.values.messages]);
+  const latestReferenceBatch =
+    referenceBatches.length > 0
+      ? referenceBatches[referenceBatches.length - 1]!
+      : null;
+  const activeReferenceBatch =
+    referenceBatches.find((batch) => batch.id === activeReferenceBatchId) ??
+    latestReferenceBatch;
+  const referenceSources = activeReferenceBatch?.sources ?? [];
   const hasCitationReferences = referenceSources.length > 0;
   useEffect(() => {
     if (hasCitationReferences) {
@@ -86,7 +96,22 @@ const ChatBox: React.FC<{ children: React.ReactNode; threadId: string }> = ({
     }
   }, [hasCitationReferences]);
   useEffect(() => {
-    const handleShowReference = () => {
+    if (!latestReferenceBatch) {
+      setActiveReferenceBatchId(null);
+      return;
+    }
+    setActiveReferenceBatchId((current) =>
+      current && referenceBatches.some((batch) => batch.id === current)
+        ? current
+        : latestReferenceBatch.id,
+    );
+  }, [latestReferenceBatch?.id, referenceBatches]);
+  useEffect(() => {
+    const handleShowReference = (event: Event) => {
+      const detail = (event as CustomEvent<{ batchId?: string }>).detail;
+      if (detail?.batchId) {
+        setActiveReferenceBatchId(detail.batchId);
+      }
       setReferencesOpen(true);
       setSideTab("references");
     };
@@ -98,12 +123,22 @@ const ChatBox: React.FC<{ children: React.ReactNode; threadId: string }> = ({
   useEffect(() => {
     if (threadIdRef.current !== threadId) {
       threadIdRef.current = threadId;
+      artifactsCountRef.current = 0;
       deselect();
       setReferencesOpen(true);
     }
 
     // Update artifacts from the current thread
-    setArtifacts(thread.values.artifacts);
+    const nextArtifacts = thread.values.artifacts ?? [];
+    const hadArtifacts = artifactsCountRef.current > 0;
+    const hasNewArtifact = nextArtifacts.length > artifactsCountRef.current;
+    artifactsCountRef.current = nextArtifacts.length;
+    setArtifacts(nextArtifacts);
+    if ((hasNewArtifact || (!hadArtifacts && nextArtifacts.length > 0)) && !selectedArtifact) {
+      selectArtifact(nextArtifacts[nextArtifacts.length - 1]!);
+      setArtifactsOpen(true);
+      setSideTab("document");
+    }
 
     // DO NOT automatically deselect the artifact when switching threads, because the artifacts auto discovering is not work now.
     // if (
@@ -142,10 +177,10 @@ const ChatBox: React.FC<{ children: React.ReactNode; threadId: string }> = ({
   const showTabs = !!selectedArtifact && hasCitationReferences;
 
   useEffect(() => {
-    if (!selectedArtifact && hasCitationReferences) {
-      setSideTab("references");
-    } else if (selectedArtifact && !hasCitationReferences) {
+    if (selectedArtifact) {
       setSideTab("document");
+    } else if (hasCitationReferences) {
+      setSideTab("references");
     }
   }, [hasCitationReferences, selectedArtifact]);
 
@@ -315,13 +350,22 @@ function ReferencePanel({
   );
 }
 
-function findLatestHumanMessageIndex(messages: Message[]) {
-  for (let index = messages.length - 1; index >= 0; index--) {
-    if (messages[index]?.type === "human") {
-      return index;
+function findPreviousHumanMessageId(
+  messages: Message[],
+  beforeMessageId?: string,
+): string | undefined {
+  const beforeIndex =
+    beforeMessageId === undefined
+      ? messages.length
+      : messages.findIndex((message) => message.id === beforeMessageId);
+  const end = beforeIndex < 0 ? messages.length : beforeIndex;
+  for (let index = end - 1; index >= 0; index--) {
+    const message = messages[index];
+    if (message?.type === "human") {
+      return message.id;
     }
   }
-  return -1;
+  return undefined;
 }
 
 export { ChatBox };