From 7aa17880357c82d1fc657da0a28011b7b5d98957 Mon Sep 17 00:00:00 2001
From: zyc <1439655764@qq.com>
Date: Fri, 13 Mar 2026 15:47:11 +0800
Subject: [PATCH] fix video bug

---
 web/src/components/PromptInput.tsx | 63 ++++++++++++++++++++++++++++--
 web/src/store/generation.ts        | 42 +++++++++++++-------
 2 files changed, 88 insertions(+), 17 deletions(-)

diff --git a/web/src/components/PromptInput.tsx b/web/src/components/PromptInput.tsx
index 4dc6f82..e4c0dc5 100644
--- a/web/src/components/PromptInput.tsx
+++ b/web/src/components/PromptInput.tsx
@@ -29,13 +29,19 @@ export function PromptInput() {
     editorRef.current?.focus();
   }, []);
 
-  // Sync editor when editorHtml resets (e.g. after submit)
+  // Sync editor when editorHtml changes (e.g. after submit or reEdit)
   useEffect(() => {
     const el = editorRef.current;
     if (!el) return;
-    if (editorHtml === '' && el.innerHTML !== '') {
-      el.innerHTML = '';
+    if (el.innerHTML !== editorHtml) {
+      el.innerHTML = editorHtml;
+      // If the HTML is plain text but we have references, rebuild mention spans
+      // This handles the case where editorHtml comes from backend (plain text only)
+      if (editorHtml && !editorHtml.includes('data-ref-id') && references.length > 0) {
+        rebuildMentionSpans(el);
+      }
     }
+  // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [editorHtml]);
 
   // Handle @ button from toolbar
@@ -47,6 +53,57 @@ export function PromptInput() {
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [insertAtTrigger]);
 
+  // Rebuild mention spans from plain text @label patterns
+  const rebuildMentionSpans = useCallback((el: HTMLElement) => {
+    const walker = document.createTreeWalker(el, NodeFilter.SHOW_TEXT);
+    const replacements: { node: Text; matches: { start: number; end: number; ref: UploadedFile }[] }[] = [];
+
+    let textNode: Text | null;
+    while ((textNode = walker.nextNode() as Text | null)) {
+      const text = textNode.textContent || '';
+      const matches: { start: number; end: number; ref: UploadedFile }[] = [];
+      for (const ref of references) {
+        const pattern = `@${ref.label}`;
+        let idx = text.indexOf(pattern);
+        while (idx !== -1) {
+          matches.push({ start: idx, end: idx + pattern.length, ref });
+          idx = text.indexOf(pattern, idx + pattern.length);
+        }
+      }
+      if (matches.length > 0) {
+        matches.sort((a, b) => a.start - b.start);
+        replacements.push({ node: textNode, matches });
+      }
+    }
+
+    for (const { node, matches } of replacements) {
+      const text = node.textContent || '';
+      const frag = document.createDocumentFragment();
+      let lastIdx = 0;
+      for (const m of matches) {
+        if (m.start > lastIdx) {
+          frag.appendChild(document.createTextNode(text.slice(lastIdx, m.start)));
+        }
+        const span = document.createElement('span');
+        span.className = styles.mention;
+        span.contentEditable = 'false';
+        span.dataset.refId = m.ref.id;
+        span.dataset.refType = m.ref.type;
+        span.textContent = `@${m.ref.label}`;
+        frag.appendChild(span);
+        lastIdx = m.end;
+      }
+      if (lastIdx < text.length) {
+        frag.appendChild(document.createTextNode(text.slice(lastIdx)));
+      }
+      node.parentNode?.replaceChild(frag, node);
+    }
+
+    if (replacements.length > 0) {
+      setEditorHtml(el.innerHTML);
+    }
+  }, [references, setEditorHtml]);
+
   const openMentionPopup = useCallback(() => {
     const el = editorRef.current;
     if (!el) return;
diff --git a/web/src/store/generation.ts b/web/src/store/generation.ts
index 8448938..81469c1 100644
--- a/web/src/store/generation.ts
+++ b/web/src/store/generation.ts
@@ -314,20 +314,34 @@ export const useGenerationStore = create<GenerationState>((set, get) => ({
       inputStore.switchMode(task.mode);
     }
 
-    const references: UploadedFile[] = task.references.map((r) => ({
-      id: r.id,
-      type: r.type,
-      previewUrl: r.previewUrl,
-      label: r.label,
-    }));
-
-    useInputBarStore.setState({
-      prompt: task.prompt,
-      editorHtml: task.editorHtml || task.prompt,
-      aspectRatio: task.aspectRatio,
-      duration: task.duration,
-      references: task.mode === 'universal' ? references : [],
-    });
+    if (task.mode === 'universal') {
+      const references: UploadedFile[] = task.references.map((r) => ({
+        id: r.id,
+        type: r.type,
+        previewUrl: r.previewUrl,
+        label: r.label,
+        tosUrl: r.previewUrl,
+      }));
+      useInputBarStore.setState({
+        prompt: task.prompt,
+        editorHtml: task.editorHtml || task.prompt,
+        aspectRatio: task.aspectRatio,
+        duration: task.duration,
+        references,
+      });
+    } else {
+      // Keyframe mode: restore firstFrame and lastFrame
+      const firstRef = task.references.find((r) => r.label === '首帧');
+      const lastRef = task.references.find((r) => r.label === '尾帧');
+      useInputBarStore.setState({
+        prompt: task.prompt,
+        editorHtml: task.editorHtml || task.prompt,
+        aspectRatio: task.aspectRatio,
+        duration: task.duration,
+        firstFrame: firstRef ? { id: firstRef.id, type: firstRef.type, previewUrl: firstRef.previewUrl, label: '首帧', tosUrl: firstRef.previewUrl } : null,
+        lastFrame: lastRef ? { id: lastRef.id, type: lastRef.type, previewUrl: lastRef.previewUrl, label: '尾帧', tosUrl: lastRef.previewUrl } : null,
+      });
+    }
   },
 
   regenerate: (id) => {