fix(prompt-input): preserve draft edits during async submission

packages/elements/__tests__/prompt-input.test.tsx

@@ -21,6 +21,7 @@ import {
   PromptInputActionMenuTrigger,
   PromptInputBody,
   PromptInputButton,
+  PromptInputProvider,
   PromptInputSelect,
   PromptInputSelectContent,
   PromptInputSelectItem,
@@ -1054,6 +1055,98 @@ describe("promptInputProvider", () => {
 
     expect(screen.getByTestId("count")).toHaveTextContent("1");
   });
+
+  const startPendingSubmission = async () => {
+    setupPromptInputTests();
+    const user = userEvent.setup();
+
+    let resolveSubmit!: () => void;
+    let rejectSubmit!: (reason?: unknown) => void;
+    // oxlint-disable-next-line eslint-plugin-promise(avoid-new)
+    const submission = new Promise<void>((resolve, reject) => {
+      resolveSubmit = resolve;
+      rejectSubmit = reject;
+    });
+    const onSubmit = vi.fn(() => submission);
+
+    render(
+      <PromptInputProvider>
+        <PromptInput onSubmit={onSubmit}>
+          <PromptInputBody>
+            <PromptInputTextarea />
+            <PromptInputSubmit />
+          </PromptInputBody>
+        </PromptInput>
+      </PromptInputProvider>
+    );
+
+    const textarea = screen.getByRole("textbox");
+
+    await user.type(textarea, "first");
+    await user.keyboard("{Enter}");
+
+    return {
+      onSubmit,
+      rejectSubmit,
+      resolveSubmit,
+      submission,
+      textarea,
+      user,
+    };
+  };
+
+  it("preserves text edited during an async submission", async () => {
+    const { onSubmit, user, textarea, submission, resolveSubmit } =
+      await startPendingSubmission();
+
+    await vi.waitFor(() => {
+      expect(onSubmit).toHaveBeenCalledOnce();
+    });
+
+    await user.clear(textarea);
+    await user.type(textarea, "second");
+
+    await act(async () => {
+      resolveSubmit();
+      await submission;
+    });
+
+    expect(textarea).toHaveValue("second");
+  });
+
+  it("clears unchanged text after successful async submission", async () => {
+    const { onSubmit, textarea, submission, resolveSubmit } =
+      await startPendingSubmission();
+
+    await vi.waitFor(() => {
+      expect(onSubmit).toHaveBeenCalledOnce();
+    });
+
+    expect(textarea).toHaveValue("first");
+
+    await act(async () => {
+      resolveSubmit();
+      await submission;
+    });
+
+    expect(textarea).toHaveValue("");
+  });
+
+  it("preserves text after failed async submission", async () => {
+    const { onSubmit, textarea, submission, rejectSubmit } =
+      await startPendingSubmission();
+
+    await vi.waitFor(() => {
+      expect(onSubmit).toHaveBeenCalledOnce();
+    });
+
+    await act(async () => {
+      rejectSubmit(new Error("Submission failed"));
+      await expect(submission).rejects.toThrow("Submission failed");
+    });
+
+    expect(textarea).toHaveValue("first");
+  });
 });
 
 describe("file validation", () => {

packages/elements/src/prompt-input.tsx

@@ -1,5 +1,20 @@
 "use client";
 
+import type { ChatStatus, FileUIPart, SourceDocumentUIPart } from "ai";
+import type {
+  ChangeEvent,
+  ChangeEventHandler,
+  ClipboardEventHandler,
+  ComponentProps,
+  FormEvent,
+  FormEventHandler,
+  HTMLAttributes,
+  KeyboardEventHandler,
+  PropsWithChildren,
+  ReactNode,
+  RefObject,
+} from "react";
+
 import {
   Command,
   CommandEmpty,
@@ -40,7 +55,6 @@ import {
   TooltipTrigger,
 } from "@repo/shadcn-ui/components/ui/tooltip";
 import { cn } from "@repo/shadcn-ui/lib/utils";
-import type { ChatStatus, FileUIPart, SourceDocumentUIPart } from "ai";
 import {
   CornerDownLeftIcon,
   ImageIcon,
@@ -50,19 +64,6 @@ import {
   XIcon,
 } from "lucide-react";
 import { nanoid } from "nanoid";
-import type {
-  ChangeEvent,
-  ChangeEventHandler,
-  ClipboardEventHandler,
-  ComponentProps,
-  FormEvent,
-  FormEventHandler,
-  HTMLAttributes,
-  KeyboardEventHandler,
-  PropsWithChildren,
-  ReactNode,
-  RefObject,
-} from "react";
 import {
   Children,
   createContext,
@@ -191,6 +192,7 @@ export interface TextInputContext {
   value: string;
   setInput: (v: string) => void;
   clear: () => void;
+  clearIfValueMatches: (submittedText: string) => void;
 }
 
 export interface PromptInputControllerProps {
@@ -253,6 +255,10 @@ export const PromptInputProvider = ({
   const [textInput, setTextInput] = useState(initialTextInput);
   const clearInput = useCallback(() => setTextInput(""), []);
 
+  const clearInputIfValueMatches = useCallback((submittedText: string) => {
+    setTextInput((current) => (current === submittedText ? "" : current));
+  }, []);
+
   // ----- attachments state (global when wrapped)
   const [attachmentFiles, setAttachmentFiles] = useState<
     (FileUIPart & { id: string })[]
@@ -349,11 +355,18 @@ export const PromptInputProvider = ({
       attachments,
       textInput: {
         clear: clearInput,
+        clearIfValueMatches: clearInputIfValueMatches,
         setInput: setTextInput,
         value: textInput,
       },
     }),
-    [textInput, clearInput, attachments, __registerFileInput]
+    [
+      textInput,
+      clearInput,
+      clearInputIfValueMatches,
+      attachments,
+      __registerFileInput,
+    ]
   );
 
   return (
@@ -883,7 +896,7 @@ export const PromptInput = ({
             await result;
             clear();
             if (usingProvider) {
-              controller.textInput.clear();
+              controller.textInput.clearIfValueMatches(text);
             }
           } catch {
             // Don't clear on error - user may want to retry
@@ -892,7 +905,7 @@ export const PromptInput = ({
           // Sync function completed without throwing, clear inputs
           clear();
           if (usingProvider) {
-            controller.textInput.clear();
+            controller.textInput.clearIfValueMatches(text);
           }
         }
       } catch {