diff --git a/packages/elements/__tests__/prompt-input.test.tsx b/packages/elements/__tests__/prompt-input.test.tsx index d961a3ba..34dc3d88 100644 --- a/packages/elements/__tests__/prompt-input.test.tsx +++ b/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((resolve, reject) => { + resolveSubmit = resolve; + rejectSubmit = reject; + }); + const onSubmit = vi.fn(() => submission); + + render( + + + + + + + + + ); + + 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", () => { diff --git a/packages/elements/src/prompt-input.tsx b/packages/elements/src/prompt-input.tsx index 412c846d..cc93f843 100644 --- a/packages/elements/src/prompt-input.tsx +++ b/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 {