From f00534e76a72ecad9d9eb38fca9cbbae715921bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E5=BE=B7=E5=AE=87?= <kvtodev@outlook.com>
Date: Thu, 16 Jul 2026 13:18:09 +0800
Subject: [PATCH] simplify draft state

---
 .../components/RichInputArea/attachment.tsx   |  3 +-
 .../RichInputArea/attachmentContext.tsx       |  1 -
 .../home/components/RichInputArea/index.tsx   | 69 +++++++++++--------
 .../RichInputArea/richTextToPrompt.tsx        | 29 ++++----
 .../src/view/home/index.tsx                   | 24 +++----
 .../src/view/home/useHomeDraftStore.ts        | 32 ++-------
 6 files changed, 72 insertions(+), 86 deletions(-)

diff --git a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachment.tsx b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachment.tsx
index ddb5f34f2..41bffc9a0 100644
--- a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachment.tsx
+++ b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachment.tsx
@@ -13,7 +13,7 @@ export function AttachmentToken({
   nodeKey: NodeKey;
 }) {
   const [editor] = useLexicalComposerContext();
-  const { attachments, removeAttachment } = useContext(AttachmentContext);
+  const { attachments } = useContext(AttachmentContext);
   const attachment = attachments.get(attachmentId);
   if (!attachment) {
     return null;
@@ -29,7 +29,6 @@ export function AttachmentToken({
           editor.update(() => {
             $getNodeByKey(nodeKey)?.remove();
           });
-          removeAttachment(attachmentId);
         }}
       >
         <X size={12} aria-hidden="true" />
diff --git a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachmentContext.tsx b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachmentContext.tsx
index 3b4d2c0b8..04c3330a1 100644
--- a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachmentContext.tsx
+++ b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/attachmentContext.tsx
@@ -4,7 +4,6 @@ import type { HomeAttachmentDraft } from '../../useHomeDraftStore';
 
 const AttachmentContext = createContext({
   attachments: new Map<string, HomeAttachmentDraft>(),
-  removeAttachment: (_attachmentId: string): void => undefined,
 });
 
 export default AttachmentContext;
diff --git a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/index.tsx b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/index.tsx
index 5b43eea18..7d20f67e6 100644
--- a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/index.tsx
+++ b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/index.tsx
@@ -21,7 +21,7 @@ import {
   PASTE_COMMAND,
 } from 'lexical';
 import { Upload } from 'lucide-react';
-import React, { useEffect, useMemo, useRef } from 'react';
+import React, { useEffect, useRef } from 'react';
 
 import type { HomeAttachmentDraft, RichText } from '../../useHomeDraftStore';
 import AttachmentContext from './attachmentContext';
@@ -32,14 +32,14 @@ import {
 } from './attachmentNode';
 
 type RichInputAreaProps = {
-  attachments: readonly HomeAttachmentDraft[];
   value: readonly RichText[];
   placeholder: string;
   onChange: (value: RichText[]) => void;
-  onAttachmentsAdd: (attachments: HomeAttachmentDraft[]) => void;
-  onAttachmentRemove: (attachmentId: string) => void;
   children?: React.ReactNode;
 };
+type AttachmentRegistry = {
+  current: Map<string, HomeAttachmentDraft>;
+};
 const INSERT_ATTACHMENTS_COMMAND: LexicalCommand<HomeAttachmentDraft[]> =
   createCommand('INSERT_HOME_ATTACHMENTS_COMMAND');
 const INSERT_CLIPBOARD_TEXT_COMMAND: LexicalCommand<string> = createCommand(
@@ -113,8 +113,10 @@ function selectEditableEndWhenNeeded() {
 
 function EditorPlugins({
   onChange,
-  onAttachmentsAdd,
-}: Pick<RichInputAreaProps, 'onChange' | 'onAttachmentsAdd'>) {
+  attachmentByIdRef,
+}: Pick<RichInputAreaProps, 'onChange'> & {
+  attachmentByIdRef: AttachmentRegistry;
+}) {
   const [editor] = useLexicalComposerContext();
 
   useEffect(
@@ -122,6 +124,9 @@ function EditorPlugins({
       editor.registerCommand(
         INSERT_ATTACHMENTS_COMMAND,
         (nextAttachments) => {
+          nextAttachments.forEach((attachment) => {
+            attachmentByIdRef.current.set(attachment.id, attachment);
+          });
           selectEditableEndWhenNeeded();
           const currentSelection = $getSelection();
           if ($isRangeSelection(currentSelection)) {
@@ -135,7 +140,7 @@ function EditorPlugins({
         },
         COMMAND_PRIORITY_EDITOR,
       ),
-    [editor],
+    [attachmentByIdRef, editor],
   );
 
   useEffect(
@@ -172,13 +177,12 @@ function EditorPlugins({
             }
             const nextAttachments = createAttachmentDrafts([file]);
             editor.dispatchCommand(INSERT_ATTACHMENTS_COMMAND, nextAttachments);
-            onAttachmentsAdd(nextAttachments);
           });
           return true;
         },
         COMMAND_PRIORITY_HIGH,
       ),
-    [editor, onAttachmentsAdd],
+    [editor],
   );
 
   return (
@@ -186,7 +190,7 @@ function EditorPlugins({
       onChange={(editorState) => {
         editorState.read(() => {
           const parts: RichText[] = [];
-          collectRichText($getRoot(), parts);
+          collectRichText($getRoot(), parts, attachmentByIdRef.current);
           onChange(parts);
         });
       }}
@@ -194,23 +198,30 @@ function EditorPlugins({
   );
 }
 
-function collectRichText(node: LexicalNode, parts: RichText[]) {
+function collectRichText(
+  node: LexicalNode,
+  parts: RichText[],
+  attachmentById: ReadonlyMap<string, HomeAttachmentDraft>,
+) {
   if ($isTextNode(node)) {
     parts.push({ type: 'text', text: node.getTextContent() });
     return;
   }
   if ($isAttachmentNode(node)) {
-    parts.push({ type: 'attachment', attachmentId: node.__attachmentId });
+    const attachment = attachmentById.get(node.__attachmentId);
+    if (attachment) {
+      parts.push({ type: 'attachment', attachment });
+    }
     return;
   }
   if ($isElementNode(node)) {
-    node.getChildren().forEach((child) => collectRichText(child, parts));
+    node
+      .getChildren()
+      .forEach((child) => collectRichText(child, parts, attachmentById));
   }
 }
 
-export function UploadButton({
-  onAttachmentsAdd,
-}: Pick<RichInputAreaProps, 'onAttachmentsAdd'>) {
+export function UploadButton() {
   const [editor] = useLexicalComposerContext();
   const inputRef = useRef<HTMLInputElement | null>(null);
   return (
@@ -226,7 +237,6 @@ export function UploadButton({
           if (files.length === 0) return;
           const nextAttachments = createAttachmentDrafts(files);
           editor.dispatchCommand(INSERT_ATTACHMENTS_COMMAND, nextAttachments);
-          onAttachmentsAdd(nextAttachments);
         }}
       />
       <button
@@ -241,19 +251,22 @@ export function UploadButton({
   );
 }
 
+function richTextToAttachmentMap(parts: readonly RichText[]) {
+  const attachmentById = new Map<string, HomeAttachmentDraft>();
+  for (const part of parts) {
+    if (part.type === 'attachment') {
+      attachmentById.set(part.attachment.id, part.attachment);
+    }
+  }
+  return attachmentById;
+}
+
 export default function RichInputArea(props: RichInputAreaProps) {
-  const attachmentMap = useMemo(
-    () =>
-      new Map(
-        props.attachments.map((attachment) => [attachment.id, attachment]),
-      ),
-    [props.attachments],
-  );
+  const attachmentByIdRef = useRef(richTextToAttachmentMap(props.value));
   return (
     <AttachmentContext.Provider
       value={{
-        attachments: attachmentMap,
-        removeAttachment: props.onAttachmentRemove,
+        attachments: attachmentByIdRef.current,
       }}
     >
       <LexicalComposer
@@ -266,7 +279,7 @@ export default function RichInputArea(props: RichInputAreaProps) {
               paragraph.append(
                 part.type === 'text'
                   ? $createTextNode(part.text)
-                  : $createAttachmentNode(part.attachmentId),
+                  : $createAttachmentNode(part.attachment.id),
               );
             }
             $getRoot().append(paragraph);
@@ -295,7 +308,7 @@ export default function RichInputArea(props: RichInputAreaProps) {
         </div>
         <EditorPlugins
           onChange={props.onChange}
-          onAttachmentsAdd={props.onAttachmentsAdd}
+          attachmentByIdRef={attachmentByIdRef}
         />
       </LexicalComposer>
     </AttachmentContext.Provider>
diff --git a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/richTextToPrompt.tsx b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/richTextToPrompt.tsx
index 1af271084..387ef90db 100644
--- a/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/richTextToPrompt.tsx
+++ b/apps/ai-game-creator-shell/src/view/home/components/RichInputArea/richTextToPrompt.tsx
@@ -1,7 +1,4 @@
-import type {
-  HomeAttachmentDraft,
-  RichText,
-} from '../../useHomeDraftStore';
+import type { HomeAttachmentDraft, RichText } from '../../useHomeDraftStore';
 
 function ordinalSuffix(index: number) {
   if (index % 100 >= 11 && index % 100 <= 13) return 'th';
@@ -10,22 +7,26 @@ function ordinalSuffix(index: number) {
     'th'
   );
 }
-export function richTextToPrompt(
-  parts: readonly RichText[],
-  attachments: readonly HomeAttachmentDraft[],
-) {
-  const attachmentById = new Map(
-    attachments.map((attachment) => [attachment.id, attachment]),
-  );
+export function richTextToPrompt(parts: readonly RichText[]) {
   let attachmentIndex = 0;
   return parts
     .map((part) => {
       if (part.type === 'text') return part.text;
-      const attachment = attachmentById.get(part.attachmentId);
-      if (!attachment) return '';
       attachmentIndex += 1;
-      return `{${attachmentIndex}${ordinalSuffix(attachmentIndex)} attachment: ${attachment.file.name}}`;
+      return `{${attachmentIndex}${ordinalSuffix(attachmentIndex)} attachment: ${part.attachment.file.name}}`;
     })
     .join('')
     .trim();
 }
+
+export function richTextToAttachments(parts: readonly RichText[]) {
+  const attachments: HomeAttachmentDraft[] = [];
+  const seenAttachmentIds = new Set<string>();
+  for (const part of parts) {
+    if (part.type !== 'attachment') continue;
+    if (seenAttachmentIds.has(part.attachment.id)) continue;
+    seenAttachmentIds.add(part.attachment.id);
+    attachments.push(part.attachment);
+  }
+  return attachments;
+}
diff --git a/apps/ai-game-creator-shell/src/view/home/index.tsx b/apps/ai-game-creator-shell/src/view/home/index.tsx
index 569631ad8..8b77e2097 100644
--- a/apps/ai-game-creator-shell/src/view/home/index.tsx
+++ b/apps/ai-game-creator-shell/src/view/home/index.tsx
@@ -11,7 +11,10 @@ import { useState } from 'react';
 
 import BRAND_ICON from '../../../../../packages/shared/src/icons/taonier-product-ip.png'
 import RichInputArea, { UploadButton } from './components/RichInputArea';
-import { richTextToPrompt } from './components/RichInputArea/richTextToPrompt';
+import {
+  richTextToAttachments,
+  richTextToPrompt,
+} from './components/RichInputArea/richTextToPrompt';
 import {
   type HomeAgentMode,
   type HomeDraft,
@@ -62,19 +65,10 @@ export default function HomeView({
 }: HomeViewProps) {
   const homeAgentMode = useLauncherHomeDraftStore((state) => state.mode);
   const homeRichText = useLauncherHomeDraftStore((state) => state.richText);
-  const homeAttachments = useLauncherHomeDraftStore(
-    (state) => state.attachments,
-  );
   const setHomeAgentMode = useLauncherHomeDraftStore((state) => state.setMode);
   const setHomeRichText = useLauncherHomeDraftStore(
     (state) => state.setRichText,
   );
-  const addHomeAttachments = useLauncherHomeDraftStore(
-    (state) => state.addAttachments,
-  );
-  const removeHomeAttachment = useLauncherHomeDraftStore(
-    (state) => state.removeAttachment,
-  );
   const [homeCreationBusy, setHomeCreationBusy] = useState(false);
   const { resources: showcaseResources, status: showcaseStatus } =
     useHomeShowcase();
@@ -88,7 +82,8 @@ export default function HomeView({
 
   async function handleHomeSubmit(event: FormEvent<HTMLFormElement>) {
     event.preventDefault();
-    const prompt = richTextToPrompt(homeRichText, homeAttachments);
+    const referencedAttachments = richTextToAttachments(homeRichText);
+    const prompt = richTextToPrompt(homeRichText);
     if (!prompt) {
       onStatusChange(
         homeAgentModeItems.find((item) => item.mode === homeAgentMode)
@@ -103,7 +98,7 @@ export default function HomeView({
         await onCreateDraft({
           mode: homeAgentMode,
           prompt,
-          attachments: homeAttachments,
+          attachments: referencedAttachments,
         }),
       );
     } catch (error) {
@@ -166,15 +161,12 @@ export default function HomeView({
           onSubmit={handleHomeSubmit}
         >
           <RichInputArea
-            attachments={homeAttachments}
             value={homeRichText}
             placeholder={activeHomeMode.placeholder}
             onChange={setHomeRichText}
-            onAttachmentsAdd={addHomeAttachments}
-            onAttachmentRemove={removeHomeAttachment}
           >
             <div className="grid grid-cols-[1fr_auto] items-center gap-2.5 text-[12px] text-(--platform-text-soft)">
-              <UploadButton onAttachmentsAdd={addHomeAttachments} />
+              <UploadButton />
               <button
                 className="grid size-6 cursor-pointer place-items-center rounded-full border-0 bg-(image:--platform-button-primary-fill) p-0 text-(--platform-button-primary-text) shadow-(--platform-profile-action-shadow) disabled:cursor-not-allowed disabled:opacity-55"
                 type="submit"
diff --git a/apps/ai-game-creator-shell/src/view/home/useHomeDraftStore.ts b/apps/ai-game-creator-shell/src/view/home/useHomeDraftStore.ts
index 05480f284..d816df036 100644
--- a/apps/ai-game-creator-shell/src/view/home/useHomeDraftStore.ts
+++ b/apps/ai-game-creator-shell/src/view/home/useHomeDraftStore.ts
@@ -1,5 +1,10 @@
 import { create } from 'zustand';
 
+export type HomeAttachmentDraft = {
+  id: string;
+  file: File;
+};
+
 export type RichText =
   | {
       type: 'text';
@@ -7,16 +12,11 @@ export type RichText =
     }
   | {
       type: 'attachment';
-      attachmentId: string;
+      attachment: HomeAttachmentDraft;
     };
 
 export type HomeAgentMode = 'game' | 'art' | 'doc';
 
-export type HomeAttachmentDraft = {
-  id: string;
-  file: File;
-};
-
 export type HomeDraft = {
   mode: HomeAgentMode;
   prompt: string;
@@ -26,21 +26,17 @@ export type HomeDraft = {
 type UseHomeDraftStore = {
   mode: HomeAgentMode;
   richText: RichText[];
-  attachments: HomeAttachmentDraft[];
   setMode: (mode: HomeAgentMode) => void;
   setRichText: (richText: RichText[]) => void;
-  addAttachments: (attachments: HomeAttachmentDraft[]) => void;
-  removeAttachment: (attachmentId: string) => void;
   reset: () => void;
 };
 
 const initialHomeDraft: Pick<
   UseHomeDraftStore,
-  'mode' | 'richText' | 'attachments'
+  'mode' | 'richText'
 > = {
   mode: 'game',
   richText: [],
-  attachments: [],
 };
 
 // Keep non-serializable File objects available while the Home view is unmounted.
@@ -48,19 +44,5 @@ export const useLauncherHomeDraftStore = create<UseHomeDraftStore>((set) => ({
   ...initialHomeDraft,
   setMode: (mode) => set({ mode }),
   setRichText: (richText) => set({ richText }),
-  addAttachments: (attachments) =>
-    set((state) => ({
-      attachments: [...state.attachments, ...attachments],
-    })),
-  removeAttachment: (attachmentId) =>
-    set((state) => ({
-      attachments: state.attachments.filter(
-        (attachment) => attachment.id !== attachmentId,
-      ),
-      richText: state.richText.filter(
-        (part) =>
-          part.type !== 'attachment' || part.attachmentId !== attachmentId,
-      ),
-    })),
   reset: () => set(initialHomeDraft),
 }));