fix: resolve logical issues when generating multi-page data

app/components/chat/Artifact.tsx

@@ -2,7 +2,7 @@ import { useStore } from '@nanostores/react';
 import classNames from 'classnames';
 import { AnimatePresence, motion } from 'framer-motion';
 import { computed } from 'nanostores';
-import { memo, useEffect, useRef, useState } from 'react';
+import { memo, useEffect, useMemo, useRef, useState } from 'react';
 import { type BundledLanguage, type BundledTheme, createHighlighter, type HighlighterGeneric } from 'shiki';
 import type { ActionState } from '~/lib/runtime/action-runner';
 import { webBuilderStore } from '~/lib/stores/web-builder';
@@ -22,18 +22,26 @@ if (import.meta.hot && import.meta.hot.data) {
 
 interface ArtifactProps {
   messageId: string;
+  pageName: string;
 }
 
-export const Artifact = memo(({ messageId }: ArtifactProps) => {
+export const Artifact = memo(({ messageId, pageName }: ArtifactProps) => {
   const userToggledActions = useRef(false);
   const [showActions, setShowActions] = useState(false);
   const [allActionFinished, setAllActionFinished] = useState(false);
 
   const artifacts = useStore(webBuilderStore.chatStore.artifacts);
-  const artifact = artifacts[messageId];
+  const artifact = useMemo(() => {
+    const artifactsByPageName = artifacts.get(messageId);
+    if (!artifactsByPageName) {
+      return undefined;
+    }
+
+    return artifactsByPageName.get(pageName);
+  }, [artifacts, messageId, pageName]);
 
   const actions = useStore(
-    computed(artifact.runner.actions, (actions) => {
+    computed(artifact?.runner.actions!, (actions) => {
       return Object.values(actions);
     }),
   );
@@ -44,11 +52,14 @@ export const Artifact = memo(({ messageId }: ArtifactProps) => {
   };
 
   useEffect(() => {
-    const actionsMap = artifact.runner.actions.get();
+    const actionsMap = artifact?.runner.actions.get();
+    if (!actionsMap) {
+      return;
+    }
 
     Object.entries(actionsMap).forEach(([actionId, action]) => {
       if (action.status === 'running' || action.status === 'pending') {
-        artifact.runner.actions.setKey(actionId, {
+        artifact?.runner.actions.setKey(actionId, {
           ...action,
           status: 'aborted',
         });
@@ -61,7 +72,7 @@ export const Artifact = memo(({ messageId }: ArtifactProps) => {
       setShowActions(true);
     }
 
-    if (actions.length !== 0 && artifact.type === 'bundled') {
+    if (actions.length !== 0 && artifact?.type === 'bundled') {
       const finished = !actions.find((action) => action.status !== 'complete');
 
       if (allActionFinished !== finished) {
@@ -80,7 +91,7 @@ export const Artifact = memo(({ messageId }: ArtifactProps) => {
             webBuilderStore.showWorkbench.set(!showWorkbench);
           }}
         >
-          {artifact.type == 'bundled' && (
+          {artifact?.type == 'bundled' && (
             <>
               <div className="p-4">
                 {allActionFinished ? (
@@ -101,7 +112,7 @@ export const Artifact = memo(({ messageId }: ArtifactProps) => {
         </button>
         <div className="bg-upage-elements-artifacts-borderColor w-[1px]" />
         <AnimatePresence>
-          {actions.length && artifact.type !== 'bundled' && (
+          {actions.length && artifact?.type !== 'bundled' && (
             <motion.button
               initial={{ width: 0 }}
               animate={{ width: 'auto' }}
@@ -118,7 +129,7 @@ export const Artifact = memo(({ messageId }: ArtifactProps) => {
         </AnimatePresence>
       </div>
       <AnimatePresence>
-        {artifact.type !== 'bundled' && showActions && actions.length > 0 && (
+        {artifact?.type !== 'bundled' && showActions && actions.length > 0 && (
           <motion.div
             className="actions"
             initial={{ height: 0 }}
@@ -152,6 +163,7 @@ function openArtifactInWebBuilder(pageName: string, rootDomId: string) {
     webBuilderStore.currentView.set('code');
   }
   webBuilderStore.setSelectedPage(pageName);
+  webBuilderStore.setActiveSectionByPageName(pageName);
   webBuilderStore.editorStore.scrollToElement(rootDomId);
 }
 

app/components/chat/Markdown.tsx

@@ -25,12 +25,16 @@ export const Markdown = memo(({ children, html = false, limitedMarkdown = false
       div: ({ className, children, node, ...props }) => {
         if (className?.includes('__uPageArtifact__')) {
           const messageId = node?.properties.dataMessageId as string;
+          const pageName = node?.properties.dataPageName as string;
 
           if (!messageId) {
             logger.error(`Invalid message id ${messageId}`);
           }
+          if (!pageName) {
+            logger.error(`Invalid page name ${pageName}`);
+          }
 
-          return <Artifact messageId={messageId} />;
+          return <Artifact messageId={messageId} pageName={pageName} />;
         }
 
         if (className?.includes('__uPageThought__')) {

app/lib/runtime/message-parser.ts

@@ -34,6 +34,7 @@ export interface ParserCallbacks {
 
 interface ElementFactoryProps {
   messageId: string;
+  pageName: string;
 }
 
 type ElementFactory = (props: ElementFactoryProps) => string;
@@ -207,7 +208,7 @@ export class StreamingMessageParser {
 
               const artifactFactory = this._options.artifactElement ?? createArtifactElement;
 
-              output += artifactFactory({ messageId });
+              output += artifactFactory({ messageId, pageName: artifactName });
 
               i = openTagEnd + 1;
             } else {

app/lib/stores/chat.ts

@@ -18,8 +18,9 @@ export interface ArtifactState {
 }
 
 export type ArtifactUpdateState = Pick<ArtifactState, 'title' | 'closed'>;
-
-type Artifacts = MapStore<Record<string, ArtifactState>>;
+type ArtifactsByPageName = Map<string, ArtifactState>;
+type ArtifactsByMessageId = Map<string, ArtifactsByPageName>;
+type Artifacts = MapStore<ArtifactsByMessageId>;
 
 export class ChatStore {
   private globalExecutionQueue = Promise.resolve();
@@ -31,8 +32,8 @@ export class ChatStore {
   currentDescription: WritableAtom<string | undefined> =
     import.meta.hot?.data?.currentDescription ?? atom<string | undefined>(undefined);
 
-  artifacts: Artifacts = import.meta.hot?.data?.artifacts ?? map({});
-  artifactIdList: string[] = [];
+  artifacts: Artifacts = import.meta.hot?.data?.artifacts ?? map(new Map());
+  artifactIdList: { messageId: string; pageName: string }[] = [];
   actionAlert: WritableAtom<ActionAlert | undefined> =
     import.meta.hot?.data?.actionAlert ?? atom<ActionAlert | undefined>(undefined);
 
@@ -60,7 +61,12 @@ export class ChatStore {
   }
 
   get firstArtifact(): ArtifactState | undefined {
-    return this.getArtifact(this.artifactIdList[0]);
+    if (this.artifactIdList.length === 0) {
+      return undefined;
+    }
+
+    const { messageId, pageName } = this.artifactIdList[0];
+    return this.getArtifact(messageId, pageName);
   }
 
   get description() {
@@ -76,31 +82,30 @@ export class ChatStore {
   }
 
   abortAllActions() {
-    // TODO: what do we wanna do and how do we wanna recover from this?
     const artifacts = this.artifacts.get();
 
-    Object.values(artifacts).forEach((artifact) => {
+    artifacts.values().forEach((artifactByPageNames) => {
+      artifactByPageNames.values().forEach((artifact) => {
         const actions = artifact.runner.actions.get();
-
         Object.values(actions).forEach((action) => {
           if (action.status === 'running' || action.status === 'pending') {
             action.abort();
           }
         });
       });
+    });
   }
 
   addArtifact({ messageId, name, title, id }: ArtifactCallbackData) {
-    const artifact = this.getArtifact(messageId);
+    const artifact = this.getArtifact(messageId, name);
     if (artifact) {
       return;
     }
 
-    if (!this.artifactIdList.includes(messageId)) {
-      this.artifactIdList.push(messageId);
+    if (!this.artifactIdList.includes({ messageId, pageName: name })) {
+      this.artifactIdList.push({ messageId, pageName: name });
     }
-
-    this.artifacts.setKey(messageId, {
+    const newArtifact = {
       id,
       name,
       title,
@@ -112,21 +117,56 @@ export class ChatStore {
 
         this.actionAlert.set(alert);
       }),
-    });
+    };
+
+    const artifactsByMessageId = this.artifacts.get();
+    let artifactsByPageName = artifactsByMessageId.get(messageId);
+    if (!artifactsByPageName) {
+      artifactsByPageName = new Map();
+      artifactsByMessageId.set(messageId, artifactsByPageName);
     }
 
-  updateArtifact({ messageId }: ArtifactCallbackData, state: Partial<ArtifactUpdateState>) {
-    const artifact = this.getArtifact(messageId);
+    artifactsByPageName.set(name, newArtifact);
+
+    this.artifacts.set(artifactsByMessageId);
+  }
+
+  updateArtifact({ messageId, name }: ArtifactCallbackData, state: Partial<ArtifactUpdateState>) {
+    const artifact = this.getArtifact(messageId, name);
     if (!artifact) {
       return;
     }
 
-    this.artifacts.setKey(messageId, { ...artifact, ...state });
+    const artifactsByMessageId = this.artifacts.get();
+    const artifactsByPageName = artifactsByMessageId.get(messageId);
+    if (!artifactsByPageName) {
+      return;
+    }
+    artifactsByPageName.set(name, { ...artifact, ...state });
+    artifactsByMessageId.set(messageId, artifactsByPageName);
+
+    this.artifacts.set(artifactsByMessageId);
   }
 
-  private getArtifact(id: string) {
+  private getArtifact(messageId: string, pageName: string) {
     const artifacts = this.artifacts.get();
-    return artifacts[id];
+    const artifactsByPageName = artifacts.get(messageId);
+    if (!artifactsByPageName) {
+      return undefined;
+    }
+
+    return artifactsByPageName.get(pageName);
+  }
+
+  private getArtifactByArtifactId(messageId: string, artifactId: string) {
+    const artifacts = this.artifacts.get();
+
+    const artifactsByPageName = artifacts.get(messageId);
+    if (!artifactsByPageName) {
+      return undefined;
+    }
+
+    return artifactsByPageName.values().find((artifact) => artifact.id === artifactId);
   }
 
   setReloadedMessages(messages: string[]) {
@@ -138,8 +178,8 @@ export class ChatStore {
   }
 
   private async _addAction(data: ActionCallbackData) {
-    const { messageId } = data;
-    const artifact = this.getArtifact(messageId);
+    const { messageId, artifactId } = data;
+    const artifact = this.getArtifactByArtifactId(messageId, artifactId);
 
     if (!artifact) {
       unreachable('Artifact not found');
@@ -157,9 +197,9 @@ export class ChatStore {
   }
 
   async _runAction(data: ActionCallbackData, isRunning: boolean = false) {
-    const { messageId } = data;
+    const { messageId, artifactId } = data;
 
-    const artifact = this.getArtifact(messageId);
+    const artifact = this.getArtifactByArtifactId(messageId, artifactId);
     if (!artifact) {
       unreachable('Artifact not found');
     }

app/lib/stores/web-builder.ts

@@ -77,6 +77,7 @@ export class WebBuilderStore {
       // 找到第一个页面并选中
       for (const [pageName] of Object.entries(pages)) {
         this.setSelectedPage(pageName);
+        this.setActiveSectionByPageName(pageName);
         break;
       }
     }
@@ -84,14 +85,14 @@ export class WebBuilderStore {
 
   setSelectedPage(pageName: string | undefined) {
     this.pagesStore.setActivePage(pageName);
+  }
 
-    if (pageName) {
+  setActiveSectionByPageName(pageName: string) {
     const page = this.pagesStore.getPage(pageName);
     if (page) {
       this.setActiveSection(page.actionIds[page.actionIds.length - 1]);
     }
   }
-  }
 
   setActiveSection(sectionId: string | undefined) {
     this.pagesStore.setActiveSection(sectionId);

app/types/artifact.ts

@@ -5,7 +5,7 @@
 export interface UPageArtifactData {
   // artifact id，唯一
   id: string;
-  // 页面名称，最终渲染为页面文件名，如 `index.html`，不包含后缀。唯一
+  // 页面名称，如 `index`，最终渲染为页面文件名，唯一
   name: string;
   // 页面标题，最终渲染为页面标题
   title: string;