simplify draft state

2026-07-16 13:18:09 +08:00
parent 0b8fe81362
commit f00534e76a
6 changed files with 72 additions and 86 deletions
@@ -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" />
@@ -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;
@@ -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>
@@ -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;
+}
@@ -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"
@@ -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),
 }));