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); }} > 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(), - 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; +}; const INSERT_ATTACHMENTS_COMMAND: LexicalCommand = createCommand('INSERT_HOME_ATTACHMENTS_COMMAND'); const INSERT_CLIPBOARD_TEXT_COMMAND: LexicalCommand = createCommand( @@ -113,8 +113,10 @@ function selectEditableEndWhenNeeded() { function EditorPlugins({ onChange, - onAttachmentsAdd, -}: Pick) { + attachmentByIdRef, +}: Pick & { + 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, +) { 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) { +export function UploadButton() { const [editor] = useLexicalComposerContext(); const inputRef = useRef(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); }} /> (); + 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 ( 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(); + 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) { 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} > - + 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((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), }));