diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/Composer.tsx b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/Composer.tsx index b231c7f0f2e..bdf418b70c6 100644 --- a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/Composer.tsx +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/Composer.tsx @@ -6,7 +6,6 @@ import type { UserContent, } from "@superset/chat/protocol"; import type { - ComposerChip, ComposerMentionEntry, ComposerMentionProvider, PromptInputCommand, @@ -38,7 +37,6 @@ import { useChatComposerStore } from "../../../../stores/chatComposerStore"; import { CHAT_COLUMN_CLASSNAME, CHAT_GUTTER_CLASSNAME } from "../../constants"; import { commandChip, LEADING_COMMAND } from "../../utils/commandChip"; import { mcpContextLine } from "../../utils/mcpServerChips"; -import { elementsForChips } from "../../utils/messageElements"; import { type AgentSwitcher, ModelPicker } from "./components/ModelPicker"; import { ModePicker, type SessionMode } from "./components/ModePicker"; import { useAttachmentPreview } from "./hooks/useAttachmentPreview"; @@ -50,6 +48,7 @@ import { useDrawerEntry } from "./hooks/useDrawerEntry"; import { useFooterTier } from "./hooks/useFooterTier"; import { useMcpServerPicker } from "./hooks/useMcpServerPicker"; import { useUploadAttachments } from "./hooks/useUploadAttachments"; +import { submitMessage } from "./utils/submitMessage"; export type ComposerTarget = { name: string; @@ -151,7 +150,7 @@ export const Composer = memo(function Composer({ const cloudUtils = cloudTrpc.useUtils(); const uploadAttachments = useUploadAttachments(workspaceId); const attachmentPreview = useAttachmentPreview(); - const { storedDraft, onChange, clearDraft } = useComposerDraft( + const { storedDraft, onChange, beginSubmission } = useComposerDraft( draftKey, draftPersisted, ); @@ -352,58 +351,20 @@ export const Composer = memo(function Composer({ return [...local, ...agentCommands]; }, [availableCommands, mcpCommand, resumeCommand, t]); - const send = useCallback( - async ( - text: string, - files: File[], - mentions: ComposerChip[], - pastedTexts: string[], - steer: boolean, - ) => { - const tags = await uploadAttachments(files); - if (!tags) { - promptInputRef.current?.appendText(appendPastedText(text, pastedTexts)); - return; - } - const message = text.trim(); - const elements = elementsForChips(message, mentions); - const mcpContext = mcpContextLine( - mentions.flatMap((chip) => mcpPicker.serverNameOf(chip) ?? []), - ); - onSend( - [ - { - type: "text", - text: appendPastedText( - [message, ...tags].filter(Boolean).join("\n"), - pastedTexts, - ), - ...(elements.length > 0 ? { elements } : {}), - }, - ...(mcpContext - ? [{ type: "text" as const, text: mcpContext, hidden: true }] - : []), - ], - { steer }, - ); - clearDraft(); - }, - [onSend, uploadAttachments, clearDraft, mcpPicker.serverNameOf], - ); - const handleSubmit = useCallback( - ({ + async ({ text, files, mentions, pastedTexts, steer, + isDraftUnchanged, }: PromptInputSubmitPayload) => { if ( blocked || (text.trim() === "" && files.length === 0 && pastedTexts.length === 0) ) - return; + return false; const name = LEADING_COMMAND.exec(text.trim())?.[1]; const input = promptInputRef.current; if (text.trim() === "/mcp" && mcpPicker.command && input) { @@ -438,9 +399,40 @@ export const Composer = memo(function Composer({ toast.error(t({ message: `Unknown command /${name}` })); return false; } - return send(text, files, mentions, pastedTexts, steer); + const finishSubmission = beginSubmission( + appendPastedText(text, pastedTexts), + isDraftUnchanged, + ); + if (!finishSubmission) return false; + try { + const result = await submitMessage( + { text, files, mentions, pastedTexts, steer }, + { + uploadAttachments, + onSend, + mcpContext: mcpContextLine( + mentions.flatMap((chip) => mcpPicker.serverNameOf(chip) ?? []), + ), + }, + ); + finishSubmission(result !== false); + return result; + } catch (error) { + finishSubmission(false); + throw error; + } }, - [availableCommands, blocked, send, t, target, mcpPicker, onOpenHistory], + [ + availableCommands, + blocked, + onSend, + uploadAttachments, + t, + beginSubmission, + target, + mcpPicker, + onOpenHistory, + ], ); return ( diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.test.tsx b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.test.tsx index 23fdb48b2ff..516bab3d446 100644 --- a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.test.tsx +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.test.tsx @@ -1,10 +1,65 @@ import { afterEach, beforeEach, expect, spyOn, test } from "bun:test"; import { appendPastedText, parsePastedText } from "@superset/chat/core"; -const { act, cleanup, renderHook } = await import("@testing-library/react"); +const { act, cleanup, fireEvent, render, renderHook } = await import( + "@testing-library/react" +); const { prependToDraft, takeRecoveredDraftText, useComposerDraft } = await import("./useComposerDraft"); +function requireSubmission( + submission: ReturnType< + ReturnType["beginSubmission"] + >, +) { + if (!submission) throw new Error("Expected the submission to start"); + return submission; +} + +const { PromptInput } = await import("@superset/chat-ui/PromptInput"); +type PromptProps = import("@superset/chat-ui/PromptInput").PromptInputProps; + +function PendingDraftComposer({ + draftKey, + findChips, + collapseLargePastes, + inputRef, + onSend, +}: Pick & { + inputRef?: PromptProps["ref"]; + draftKey: string; + onSend: () => Promise; +}) { + const { storedDraft, onChange, beginSubmission } = useComposerDraft(draftKey); + return ( + { + const finish = beginSubmission( + appendPastedText(payload.text, payload.pastedTexts), + payload.isDraftUnchanged, + ); + if (!finish) return false; + try { + const result = await onSend(); + finish(result); + return result === false ? false : undefined; + } catch (error) { + finish(false); + throw error; + } + }} + /> + ); +} + beforeEach(() => window.localStorage.clear()); afterEach(cleanup); test("a composer mounted before the save delay reads text typed into the one it replaces", () => { @@ -121,3 +176,544 @@ test("text restored to a subagent's draft stays in memory and reaches its compos ); expect(reopened.result.current.storedDraft).toBe("for the subagent"); }); + +test("an accepted upload/send after unmount clears the submitted persisted draft", async () => { + const { submitMessage } = await import("../../utils/submitMessage"); + const key = "chat-v3-draft:accepted-after-unmount"; + const composer = renderHook(() => useComposerDraft(key)); + act(() => composer.result.current.onChange(" submitted text ")); + const upload = Promise.withResolvers(); + const finishSubmission = requireSubmission( + composer.result.current.beginSubmission("submitted text"), + ); + const pending = submitMessage( + { + text: "submitted text", + files: [], + mentions: [], + pastedTexts: [], + steer: false, + }, + { uploadAttachments: () => upload.promise, onSend: async () => true }, + ); + composer.unmount(); + upload.resolve([]); + finishSubmission((await pending) !== false); + expect(window.localStorage.getItem(key)).toBeNull(); + const remounted = renderHook(() => useComposerDraft(key)); + expect(remounted.result.current.storedDraft).toBeUndefined(); +}); + +test("rejected or thrown submissions preserve the unmounted draft", async () => { + const { submitMessage } = await import("../../utils/submitMessage"); + for (const rejection of [false, new Error("not accepted")]) { + const key = `chat-v3-draft:rejected-${String(rejection)}`; + const composer = renderHook(() => useComposerDraft(key)); + act(() => composer.result.current.onChange("keep this")); + const finish = requireSubmission( + composer.result.current.beginSubmission("keep this"), + ); + composer.unmount(); + try { + const result = await submitMessage( + { + text: "keep this", + files: [], + mentions: [], + pastedTexts: [], + steer: false, + }, + { + uploadAttachments: async () => [], + onSend: async () => { + if (rejection instanceof Error) throw rejection; + return rejection; + }, + }, + ); + finish(result !== false); + } catch { + finish(false); + } + expect(window.localStorage.getItem(key)).toBe("keep this"); + } +}); + +test("a raw edit and undo invalidate a pending submission", () => { + const key = "chat-v3-draft:submission-aba"; + const composer = renderHook(() => useComposerDraft(key)); + act(() => composer.result.current.onChange(" same text ")); + const finish = requireSubmission( + composer.result.current.beginSubmission("same text"), + ); + act(() => composer.result.current.onChange("same text")); + act(() => composer.result.current.onChange(" same text ")); + composer.unmount(); + finish(true); + expect(window.localStorage.getItem(key)).toBe(" same text "); +}); + +test("recovered text invalidates a pending submission", () => { + const key = "chat-v3-draft:submission-recovery"; + const composer = renderHook(() => useComposerDraft(key)); + act(() => composer.result.current.onChange("submitted")); + const finish = requireSubmission( + composer.result.current.beginSubmission("submitted"), + ); + prependToDraft(key, ["recovered queued message"]); + composer.unmount(); + finish(true); + expect(window.localStorage.getItem(key)).toBe( + "recovered queued message\n\nsubmitted", + ); +}); + +test("a different mount of the same draft can edit while a submission settles", () => { + const key = "chat-v3-draft:submission-other-mount"; + const first = renderHook(() => useComposerDraft(key)); + act(() => first.result.current.onChange("submitted")); + const finish = requireSubmission( + first.result.current.beginSubmission("submitted"), + ); + const second = renderHook(() => useComposerDraft(key)); + act(() => second.result.current.onChange("newer draft")); + first.unmount(); + finish(true); + second.unmount(); + expect(window.localStorage.getItem(key)).toBe("newer draft"); +}); + +test("a same-key remount cannot replace an in-flight submission", () => { + const key = "chat-v3-draft:submission-newer"; + const first = renderHook(() => useComposerDraft(key)); + act(() => first.result.current.onChange("submitted")); + const finish = requireSubmission( + first.result.current.beginSubmission("submitted"), + ); + first.unmount(); + const replacement = renderHook(() => useComposerDraft(key)); + try { + expect(replacement.result.current.beginSubmission("submitted")).toBeNull(); + } finally { + act(() => finish(true)); + } + expect(window.localStorage.getItem(key)).toBeNull(); + expect(replacement.result.current.storedDraft).toBeUndefined(); +}); + +test("acceptance retains an empty memory fallback when storage removal fails", () => { + const key = "chat-v3-draft:submission-remove-failure"; + const composer = renderHook(() => useComposerDraft(key)); + act(() => composer.result.current.onChange("submitted")); + const finish = requireSubmission( + composer.result.current.beginSubmission("submitted"), + ); + composer.unmount(); + const removeItem = spyOn( + window.localStorage, + "removeItem", + ).mockImplementation(() => { + throw new Error("SecurityError"); + }); + const warn = spyOn(console, "warn").mockImplementation(() => {}); + try { + finish(true); + const next = renderHook(() => useComposerDraft(key)); + expect(next.result.current.storedDraft).toBe(""); + act(() => next.result.current.onChange("newer draft")); + next.unmount(); + expect(window.localStorage.getItem(key)).toBe("newer draft"); + } finally { + removeItem.mockRestore(); + warn.mockRestore(); + } +}); + +test("a submission keeps its original draft key across navigation", () => { + const firstKey = "chat-v3-draft:submission-navigation-a"; + const secondKey = "chat-v3-draft:submission-navigation-b"; + const composer = renderHook(({ key }) => useComposerDraft(key), { + initialProps: { key: firstKey }, + }); + act(() => composer.result.current.onChange("submitted")); + const finish = requireSubmission( + composer.result.current.beginSubmission("submitted"), + ); + composer.rerender({ key: secondKey }); + act(() => composer.result.current.onChange("next chat draft")); + finish(true); + composer.unmount(); + expect(window.localStorage.getItem(firstKey)).toBeNull(); + expect(window.localStorage.getItem(secondKey)).toBe("next chat draft"); +}); + +test("acceptance clears the mounted seed before navigating back to the same draft", () => { + const key = "chat-v3-draft:submission-seeded-navigation"; + window.localStorage.setItem(key, "submitted"); + const composer = renderHook(({ key }) => useComposerDraft(key), { + initialProps: { key }, + }); + const finish = requireSubmission( + composer.result.current.beginSubmission("submitted"), + ); + act(() => finish(true)); + composer.rerender({ + key: "chat-v3-draft:submission-seeded-navigation-other", + }); + composer.rerender({ key }); + expect(composer.result.current.storedDraft).toBeUndefined(); +}); + +test("acceptance cannot restore the seed of a replacement mount", () => { + const key = "chat-v3-draft:submission-replacement-seed"; + window.localStorage.setItem(key, "submitted"); + const first = renderHook(() => useComposerDraft(key)); + const finish = requireSubmission( + first.result.current.beginSubmission("submitted"), + ); + const replacement = renderHook(({ key }) => useComposerDraft(key), { + initialProps: { key }, + }); + first.unmount(); + act(() => finish(true)); + replacement.rerender({ key: "chat-v3-draft:submission-replacement-other" }); + replacement.rerender({ key }); + expect(replacement.result.current.storedDraft).toBeUndefined(); +}); + +test("a temporary storage read failure keeps the mounted draft seed", () => { + const key = "chat-v3-draft:seed-read-failure"; + window.localStorage.setItem(key, "keep this draft"); + const composer = renderHook(() => useComposerDraft(key)); + const getItem = spyOn(window.localStorage, "getItem").mockImplementation( + () => { + throw new Error("SecurityError"); + }, + ); + try { + composer.rerender(); + expect(composer.result.current.storedDraft).toBe("keep this draft"); + } finally { + getItem.mockRestore(); + } +}); + +test("acceptance while navigated away clears the original cached seed", () => { + const key = "chat-v3-draft:submission-away-seed"; + window.localStorage.setItem(key, "submitted"); + const composer = renderHook(({ key }) => useComposerDraft(key), { + initialProps: { key }, + }); + const finish = requireSubmission( + composer.result.current.beginSubmission("submitted"), + ); + composer.rerender({ key: "chat-v3-draft:submission-away-other" }); + finish(true); + composer.rerender({ key }); + expect(composer.result.current.storedDraft).toBeUndefined(); +}); + +test("a late plugin chip stays saved when a send is accepted after close", async () => { + const key = "chat-v3-draft:late-plugin-hydration"; + const text = "Review @example.ts"; + window.localStorage.setItem(key, text); + const accepted = Promise.withResolvers(); + const onSend = () => accepted.promise; + const composer = render( + , + ); + await act(async () => {}); + fireEvent.click(composer.getByRole("button", { name: "Send message" })); + composer.rerender( + [ + { + start: 7, + end: text.length, + chip: { label: "Example file", serialized: "@example.ts" }, + }, + ]} + />, + ); + await act(async () => {}); + + expect(composer.getByText("Example file")).toBeTruthy(); + expect(composer.getByRole("textbox").textContent).toBe("Review Example file"); + composer.unmount(); + await act(async () => accepted.resolve(true)); + expect(window.localStorage.getItem(key)).toBe(text); + const reopened = render( + , + ); + await act(async () => {}); + + expect(reopened.getByRole("textbox").textContent).toBe(text); +}); + +test("draft edits, clearing and recovery do not release an in-flight send", () => { + for (const change of ["edit", "clear", "recover"]) { + const key = `chat-v3-draft:pending-${change}`; + const first = renderHook(() => useComposerDraft(key)); + act(() => first.result.current.onChange("submitted")); + const finish = requireSubmission( + first.result.current.beginSubmission("submitted"), + ); + act(() => { + if (change === "edit") first.result.current.onChange("newer draft"); + if (change === "clear") first.result.current.clearDraft(); + if (change === "recover") prependToDraft(key, ["recovered"]); + }); + first.unmount(); + const replacement = renderHook(() => useComposerDraft(key)); + try { + expect( + replacement.result.current.beginSubmission("next send"), + ).toBeNull(); + } finally { + finish(false); + } + const retry = replacement.result.current.beginSubmission("next send"); + expect(retry).toBeFunction(); + retry?.(false); + replacement.unmount(); + } +}); + +test("in-flight guards stay with their draft key across navigation", () => { + const key = "chat-v3-draft:pending-navigation"; + const composer = renderHook(({ draftKey }) => useComposerDraft(draftKey), { + initialProps: { draftKey: key }, + }); + act(() => composer.result.current.onChange("first draft")); + const finishFirst = requireSubmission( + composer.result.current.beginSubmission("first draft"), + ); + composer.rerender({ draftKey: `${key}-other` }); + act(() => composer.result.current.onChange("second draft")); + const finishOther = requireSubmission( + composer.result.current.beginSubmission("second draft"), + ); + expect(finishOther).toBeFunction(); + composer.rerender({ draftKey: key }); + try { + expect(composer.result.current.beginSubmission("first draft")).toBeNull(); + } finally { + act(() => finishFirst(true)); + act(() => finishOther?.(false)); + } + const retry = composer.result.current.beginSubmission("another message"); + expect(retry).toBeFunction(); + retry?.(false); + composer.unmount(); + expect(window.localStorage.getItem(key)).toBeNull(); + expect(window.localStorage.getItem(`${key}-other`)).toBe("second draft"); +}); + +test("a remounted chip-edited composer declines duplicate sends until acceptance", async () => { + const key = "chat-v3-draft:pending-remount-chip"; + const text = "Review @example.ts"; + window.localStorage.setItem(key, text); + const accepted = Promise.withResolvers(); + let sendCount = 0; + const onSend = () => { + sendCount += 1; + return accepted.promise; + }; + const original = render( + , + ); + await act(async () => {}); + fireEvent.click(original.getByRole("button", { name: "Send message" })); + original.rerender( + [ + { + start: 7, + end: text.length, + chip: { label: "Example file", serialized: "@example.ts" }, + }, + ]} + />, + ); + await act(async () => {}); + expect(original.getByText("Example file")).toBeTruthy(); + original.unmount(); + const replacement = render( + , + ); + await act(async () => {}); + try { + await act(async () => + fireEvent.click( + replacement.getByRole("button", { name: "Send message" }), + ), + ); + expect(sendCount).toBe(1); + } finally { + await act(async () => accepted.resolve(true)); + } + expect(window.localStorage.getItem(key)).toBe(text); + expect(replacement.getByRole("textbox").textContent).toBe(text); + await act(async () => + fireEvent.click(replacement.getByRole("button", { name: "Send message" })), + ); + expect(sendCount).toBe(2); + replacement.unmount(); + expect(window.localStorage.getItem(key)).toBeNull(); +}); + +test.each([ + false, + new Error("not accepted"), +])("a remounted composer can retry after a pending send settles with %s", async (rejection) => { + const key = `chat-v3-draft:pending-remount-rejected-${String(rejection)}`; + window.localStorage.setItem(key, "retry this draft"); + const pending = Promise.withResolvers(); + let sendCount = 0; + const onSend = () => { + sendCount += 1; + return pending.promise; + }; + const original = render( + , + ); + await act(async () => {}); + fireEvent.click(original.getByRole("button", { name: "Send message" })); + original.unmount(); + const replacement = render( + , + ); + await act(async () => {}); + try { + await act(async () => + fireEvent.click( + replacement.getByRole("button", { name: "Send message" }), + ), + ); + expect(sendCount).toBe(1); + } finally { + await act(async () => { + if (rejection instanceof Error) pending.reject(rejection); + else pending.resolve(false); + }); + } + expect(replacement.getByRole("textbox").textContent).toBe("retry this draft"); + await act(async () => + fireEvent.click(replacement.getByRole("button", { name: "Send message" })), + ); + expect(sendCount).toBe(2); + replacement.unmount(); + expect(window.localStorage.getItem(key)).toBe("retry this draft"); +}); + +test("an accepted non-persisted subagent draft clears in memory without writing storage", () => { + const key = "chat-v3-draft:accepted-subagent:child"; + const storageWrite = spyOn(window.localStorage, "setItem"); + const storageRemove = spyOn(window.localStorage, "removeItem"); + try { + const composer = renderHook(() => useComposerDraft(key, false)); + act(() => composer.result.current.onChange("subagent message")); + const finish = requireSubmission( + composer.result.current.beginSubmission("subagent message"), + ); + composer.unmount(); + finish(true); + const reopened = renderHook(() => useComposerDraft(key, false)); + expect(reopened.result.current.storedDraft).toBeUndefined(); + expect(storageWrite).not.toHaveBeenCalled(); + expect(storageRemove).not.toHaveBeenCalled(); + } finally { + storageWrite.mockRestore(); + storageRemove.mockRestore(); + } +}); + +test("a declined non-persisted subagent draft survives remount and can retry", () => { + const key = "chat-v3-draft:declined-subagent:child"; + const composer = renderHook(() => useComposerDraft(key, false)); + act(() => composer.result.current.onChange("retry subagent message")); + const finish = requireSubmission( + composer.result.current.beginSubmission("retry subagent message"), + ); + composer.unmount(); + finish(false); + const reopened = renderHook(() => useComposerDraft(key, false)); + expect(reopened.result.current.storedDraft).toBe("retry subagent message"); + const retry = requireSubmission( + reopened.result.current.beginSubmission("retry subagent message"), + ); + act(() => retry(true)); + expect(window.localStorage.getItem(key)).toBeNull(); +}); + +test("clearing an ephemeral target does not discard a parent draft waiting to save", () => { + const key = "chat-v3-draft:parent-before-subagent"; + const composer = renderHook( + ({ draftKey, persisted }) => useComposerDraft(draftKey, persisted), + { initialProps: { draftKey: key, persisted: true } }, + ); + act(() => composer.result.current.onChange("parent draft")); + composer.rerender({ draftKey: `${key}:child`, persisted: false }); + act(() => composer.result.current.clearDraft()); + composer.unmount(); + expect(window.localStorage.getItem(key)).toBe("parent draft"); + expect(window.localStorage.getItem(`${key}:child`)).toBeNull(); +}); + +test("acceptance clears the serialized pasted-card draft only after sending finishes", async () => { + const key = "chat-v3-draft:accepted-pasted-card"; + const text = appendPastedText("Review the log", ["original pasted log"]); + window.localStorage.setItem(key, text); + const accepted = Promise.withResolvers(); + const composer = render( + accepted.promise} + />, + ); + await act(async () => {}); + expect(composer.getByText("original pasted log")).toBeTruthy(); + fireEvent.click(composer.getByRole("button", { name: "Send message" })); + await act(async () => {}); + expect(window.localStorage.getItem(key)).toBe(text); + expect(composer.getByText("original pasted log")).toBeTruthy(); + await act(async () => accepted.resolve(true)); + expect(composer.queryByText("original pasted log")).toBeNull(); + expect(composer.getByRole("textbox").textContent).toBe(""); + expect(window.localStorage.getItem(key)).toBeNull(); +}); + +test("a pasted card recovered while acceptance is pending keeps the newer draft", async () => { + const key = "chat-v3-draft:pending-pasted-recovery"; + window.localStorage.setItem(key, "Review the log"); + const accepted = Promise.withResolvers(); + const inputRef: { + current: import("@superset/chat-ui/PromptInput").PromptInputHandle | null; + } = { current: null }; + const composer = render( + accepted.promise} + />, + ); + await act(async () => {}); + fireEvent.click(composer.getByRole("button", { name: "Send message" })); + await act(async () => { + inputRef.current?.appendText(appendPastedText("", ["new recovered log"])); + }); + await act(async () => accepted.resolve(true)); + expect(composer.getByText("new recovered log")).toBeTruthy(); + expect(composer.getByRole("textbox").textContent).toBe("Review the log"); + composer.unmount(); + expect(parsePastedText(window.localStorage.getItem(key) ?? "")).toEqual({ + text: "Review the log", + pastes: ["new recovered log"], + }); +}); diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.ts b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.ts index cc16790e743..8ffe8d1097c 100644 --- a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.ts +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/hooks/useComposerDraft/useComposerDraft.ts @@ -5,6 +5,9 @@ const DRAFT_DEBOUNCE_MS = 300; const unsavedDrafts = new Map(); const recoveredAfterRead = new Map(); +const pendingSubmissions = new Map(); +const inFlightSubmissions = new Map(); +const clearedDraftListeners = new Map void>>(); function persistDraft(draftKey: string, text: string): boolean { try { @@ -17,13 +20,19 @@ function persistDraft(draftKey: string, text: string): boolean { } } -function readDraft(draftKey: string): string | undefined { +function clearStoredDraft(draftKey: string, persisted: boolean) { + unsavedDrafts.set(draftKey, ""); + if (!persisted || persistDraft(draftKey, "")) unsavedDrafts.delete(draftKey); + for (const notify of clearedDraftListeners.get(draftKey) ?? []) notify(); +} + +function readDraft(draftKey: string, fallback?: string): string | undefined { const unsaved = unsavedDrafts.get(draftKey); if (unsaved !== undefined) return unsaved; try { return window.localStorage.getItem(draftKey) ?? undefined; } catch { - return undefined; + return fallback; } } @@ -43,6 +52,7 @@ export function prependToDraft( texts: readonly string[], persisted = true, ) { + pendingSubmissions.delete(draftKey); const merged = joinDrafts([...texts, readDraft(draftKey) ?? ""]); unsavedDrafts.set(draftKey, merged); if (persisted) persistDraft(draftKey, merged); @@ -74,6 +84,9 @@ export function useComposerDraft(draftKey: string, persisted = true) { }, []); const onChange = useCallback( (text: string) => { + const submission = pendingSubmissions.get(draftKey); + if (submission && submission.text !== text) + pendingSubmissions.delete(draftKey); if (!persisted) { if (text === "") unsavedDrafts.delete(draftKey); else unsavedDrafts.set(draftKey, text); @@ -93,15 +106,60 @@ export function useComposerDraft(draftKey: string, persisted = true) { return { draftKey, text: readDraft(draftKey) }; }); + useEffect(() => { + setSeed((current) => { + const text = readDraft( + draftKey, + current.draftKey === draftKey ? current.text : undefined, + ); + return current.draftKey === draftKey && current.text === text + ? current + : { draftKey, text }; + }); + const listeners = + clearedDraftListeners.get(draftKey) ?? new Set<() => void>(); + const clearSeed = () => setSeed({ draftKey, text: undefined }); + listeners.add(clearSeed); + clearedDraftListeners.set(draftKey, listeners); + return () => { + listeners.delete(clearSeed); + if (listeners.size === 0) clearedDraftListeners.delete(draftKey); + }; + }, [draftKey]); + + const beginSubmission = useCallback( + (text: string, isDraftUnchanged?: () => boolean) => { + if (inFlightSubmissions.has(draftKey)) return null; + pendingSubmissions.delete(draftKey); + const submitted = { text: readDraft(draftKey) ?? "" }; + inFlightSubmissions.set(draftKey, submitted); + if (submitted.text.trim() === text.trim()) + pendingSubmissions.set(draftKey, submitted); + return (accepted: boolean) => { + if (inFlightSubmissions.get(draftKey) !== submitted) return; + inFlightSubmissions.delete(draftKey); + if (pendingSubmissions.get(draftKey) !== submitted) return; + pendingSubmissions.delete(draftKey); + if ( + !accepted || + isDraftUnchanged?.() === false || + (readDraft(draftKey) ?? "") !== submitted.text + ) + return; + clearStoredDraft(draftKey, persisted); + }; + }, + [draftKey, persisted], + ); + const clearDraft = useCallback(() => { - unsavedDrafts.delete(draftKey); - if (persisted) { + pendingSubmissions.delete(draftKey); + if (pendingDraftKey.current === draftKey) { if (draftTimer.current) clearTimeout(draftTimer.current); draftTimer.current = null; pendingDraftKey.current = null; - persistDraft(draftKey, ""); } - setSeed({ draftKey, text: undefined }); + clearStoredDraft(draftKey, persisted); }, [draftKey, persisted]); if (seed.draftKey !== draftKey) { @@ -109,9 +167,8 @@ export function useComposerDraft(draftKey: string, persisted = true) { setSeed({ draftKey, text: readDraft(draftKey) }); } - return { - storedDraft: seed.draftKey === draftKey ? seed.text : readDraft(draftKey), - onChange, - clearDraft, - }; + const storedDraft = + seed.draftKey === draftKey ? seed.text : readDraft(draftKey); + + return { storedDraft, onChange, clearDraft, beginSubmission }; } diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/index.ts b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/index.ts new file mode 100644 index 00000000000..ee2ccccce76 --- /dev/null +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/index.ts @@ -0,0 +1 @@ +export { submitMessage } from "./submitMessage"; diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.test.ts b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.test.ts new file mode 100644 index 00000000000..bf75e7569f3 --- /dev/null +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.test.ts @@ -0,0 +1,155 @@ +import { expect, mock, test } from "bun:test"; +import { parsePastedText } from "@superset/chat/core"; +import type { UserContent } from "@superset/chat/protocol"; +import type { PromptInputSubmitPayload } from "@superset/chat-ui/PromptInput"; +import { submitMessage } from "./submitMessage"; + +const draft: PromptInputSubmitPayload = { + text: "Review the file", + files: [], + mentions: [], + pastedTexts: [], + steer: false, +}; + +test("failed attachment upload declines submission without sending", async () => { + const onSend = mock(); + const result = await submitMessage(draft, { + uploadAttachments: async () => null, + onSend, + }); + expect(result).toBe(false); + expect(onSend).not.toHaveBeenCalled(); +}); + +test("submission remains pending until the send callback accepts", async () => { + let accept!: () => void; + const pending = new Promise((resolve) => { + accept = resolve; + }); + let settled = false; + const result = submitMessage(draft, { + uploadAttachments: async () => [], + onSend: () => pending, + }).then(() => { + settled = true; + }); + await Promise.resolve(); + await Promise.resolve(); + expect(settled).toBe(false); + accept(); + await result; + expect(settled).toBe(true); +}); + +test("a declined send propagates false so the composer retains its draft", async () => { + expect( + await submitMessage(draft, { + uploadAttachments: async () => [], + onSend: async () => false, + }), + ).toBe(false); +}); + +test("rejected sends remain rejected so the composer does not clear its draft", async () => { + const failure = new Error("Transport unavailable"); + await expect( + submitMessage(draft, { + uploadAttachments: async () => [], + onSend: async () => { + throw failure; + }, + }), + ).rejects.toBe(failure); +}); + +test("accepted sends retain attachment tags, chip metadata and steering", async () => { + const onSend = mock(() => ({ clientId: "accepted-outbox-entry" })); + await submitMessage( + { + ...draft, + text: "Read example.ts", + mentions: [ + { + label: "Example", + serialized: "example.ts", + data: { elementKind: "file_mention" }, + }, + ], + steer: true, + }, + { + uploadAttachments: async () => ["[attachment: review.txt]"], + onSend, + }, + ); + expect(onSend).toHaveBeenCalledWith( + [ + { + type: "text", + text: "Read example.ts\n[attachment: review.txt]", + elements: [ + { + byteRange: { start: 5, end: 15 }, + elementKind: "file_mention", + label: "Example", + }, + ], + }, + ], + { steer: true }, + ); +}); + +test("accepted sends keep pasted-card contents and hidden MCP context", async () => { + const onSend = mock< + (content: UserContent[], options: { steer: boolean }) => Promise + >(async () => undefined); + const mcpContext = + 'MCP context: Use the configured MCP server "Fixture MCP" when relevant to this request.'; + await submitMessage( + { + ...draft, + text: " Review the log ", + pastedTexts: ["first pasted log", "second pasted log"], + }, + { + uploadAttachments: async () => ["[attachment: review.txt]"], + onSend, + mcpContext, + }, + ); + expect(onSend).toHaveBeenCalledTimes(1); + const [content, options] = onSend.mock.calls[0] ?? []; + expect(content).toHaveLength(2); + const first = content?.[0]; + if (first?.type !== "text") throw new Error("Expected text content"); + expect(parsePastedText(first.text)).toEqual({ + text: "Review the log\n[attachment: review.txt]", + pastes: ["first pasted log", "second pasted log"], + }); + expect(content?.[1]).toEqual({ + type: "text", + text: mcpContext, + hidden: true, + }); + expect(options).toEqual({ steer: false }); +}); + +test("a pasted-card-only message sends its content without a fabricated visible prompt", async () => { + const onSend = mock< + (content: UserContent[], options: { steer: boolean }) => Promise + >(async () => undefined); + await submitMessage( + { ...draft, text: "", pastedTexts: ["complete pasted message"] }, + { uploadAttachments: async () => [], onSend, mcpContext: null }, + ); + const [content] = onSend.mock.calls[0] ?? []; + expect(content).toHaveLength(1); + const first = content?.[0]; + if (first?.type !== "text") throw new Error("Expected text content"); + expect(parsePastedText(first.text)).toEqual({ + text: "", + pastes: ["complete pasted message"], + }); +}); diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.ts b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.ts new file mode 100644 index 00000000000..054b88cb9ea --- /dev/null +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/v2-workspace/$workspaceId/hooks/usePaneRegistry/components/ChatSession/components/Composer/utils/submitMessage/submitMessage.ts @@ -0,0 +1,39 @@ +import { appendPastedText } from "@superset/chat/core"; +import type { UserContent } from "@superset/chat/protocol"; +import type { PromptInputSubmitPayload } from "@superset/chat-ui/PromptInput"; +import { elementsForChips } from "../../../../utils/messageElements"; + +export async function submitMessage( + { text, files, mentions, pastedTexts, steer }: PromptInputSubmitPayload, + { + uploadAttachments, + onSend, + mcpContext, + }: { + uploadAttachments: (files: File[]) => Promise; + onSend: (content: UserContent[], options: { steer: boolean }) => unknown; + mcpContext?: string | null; + }, +) { + const tags = await uploadAttachments(files); + if (!tags) return false; + const message = text.trim(); + const elements = elementsForChips(message, mentions); + const result = await onSend( + [ + { + type: "text", + text: appendPastedText( + [message, ...tags].filter(Boolean).join("\n"), + pastedTexts, + ), + ...(elements.length > 0 ? { elements } : {}), + }, + ...(mcpContext + ? [{ type: "text" as const, text: mcpContext, hidden: true }] + : []), + ], + { steer }, + ); + if (result === false) return false; +} diff --git a/packages/chat-runtime/src/harness/acp/acpAdapter.test.ts b/packages/chat-runtime/src/harness/acp/acpAdapter.test.ts index 0c8b6f3e47a..3c94ab131e2 100644 --- a/packages/chat-runtime/src/harness/acp/acpAdapter.test.ts +++ b/packages/chat-runtime/src/harness/acp/acpAdapter.test.ts @@ -4010,6 +4010,214 @@ describe("AcpAdapter on protocol v2", () => { await adapter.dispose(); }); + it.each([ + "injected", + "startedNewTurn", + ])("acknowledges a v1 child steer as %s without a second prompt", async (outcome) => { + const agent = new FakeAcpAgent(); + agent.protocolVersion = 1; + agent.steeringOutcome = outcome; + const { adapter } = startAdapter(agent); + await flush(); + agent.notify("sess-1", announce("child-1")); + await flush(); + try { + expect( + await adapter.steerSubagent("subagent:child-1", [ + { type: "text", text: "check the tests too" }, + ]), + ).toBe(true); + expect( + agent.sent.filter((frame) => frame.method === "_session/steering"), + ).toMatchObject([ + { + params: { + sessionId: "child-1", + prompt: [{ type: "text", text: "check the tests too" }], + _meta: { steering: { idleBehavior: "promptRequired" } }, + }, + }, + ]); + expect( + agent.sent.some((frame) => frame.method === "session/prompt"), + ).toBe(false); + } finally { + await adapter.dispose(); + } + }); + + it.each([ + "promptRequired", + "unsupported", + ])("handles the v1 child steer outcome %s without duplicating accepted input", async (outcome) => { + const agent = new FakeAcpAgent(); + agent.protocolVersion = 1; + agent.steeringOutcome = outcome; + const { adapter } = startAdapter(agent); + await flush(); + agent.notify("sess-1", announce("child-1")); + await flush(); + try { + expect( + await adapter.steerSubagent("subagent:child-1", [ + { type: "text", text: "continue" }, + ]), + ).toBe(outcome === "promptRequired"); + const requests = agent.sent.filter((frame) => + ["_session/steering", "session/prompt"].includes( + String(frame.method), + ), + ); + expect(requests.map((frame) => frame.method)).toEqual( + outcome === "promptRequired" + ? ["_session/steering", "session/prompt"] + : ["_session/steering"], + ); + } finally { + await adapter.dispose(); + } + }); + + it.each([ + 1, 2, + ])("keeps the normal child prompt path for protocol %s without v1 steering", async (version) => { + const agent = new FakeAcpAgent(); + agent.protocolVersion = version; + agent.steeringOutcome = version === 2 ? "injected" : null; + const { adapter } = startAdapter(agent); + await flush(); + agent.notify("sess-1", announce("child-1")); + await flush(); + try { + expect( + await adapter.steerSubagent("subagent:child-1", [ + { type: "text", text: "continue" }, + ]), + ).toBe(true); + expect( + agent.sent.some((frame) => frame.method === "_session/steering"), + ).toBe(false); + expect( + agent.sent.filter((frame) => frame.method === "session/prompt"), + ).toHaveLength(1); + } finally { + await adapter.dispose(); + } + }); + + it("a child closed while its attachment resolves cannot receive steering", async () => { + const agent = new FakeAcpAgent(); + agent.protocolVersion = 1; + agent.steeringOutcome = "injected"; + let release!: (value: { + path: string; + mimeType: string; + data: string; + }) => void; + let attachmentStarted!: () => void; + const started = new Promise((resolve) => { + attachmentStarted = resolve; + }); + const adapter = new AcpAdapter({ + command: "fake", + createTransport: (_opts, handlers) => agent.transport(handlers), + resolveAttachment: () => + new Promise((resolve) => { + release = resolve; + attachmentStarted(); + }), + now: () => 1, + mintId: () => "fixture-id", + }); + void collect(adapter.start({ cwd: "/work" }), []); + await flush(); + agent.notify("sess-1", announce("child-1")); + await flush(); + try { + const pending = adapter.steerSubagent("subagent:child-1", [ + { + type: "attachment", + attachmentId: "fixture", + name: "fixture.png", + mimeType: "image/png", + }, + ]); + await started; + agent.notify("sess-1", { + sessionUpdate: "subagent_update", + sessionId: "child-1", + state: { state: "idle", stopReason: "end_turn" }, + }); + await flush(); + release({ + path: "/fixture/image.png", + mimeType: "image/png", + data: "aW1hZ2U=", + }); + await expect(pending).resolves.toBe(false); + expect( + agent.sent.filter((frame) => + ["_session/steering", "session/prompt"].includes( + String(frame.method), + ), + ), + ).toEqual([]); + } finally { + await adapter.dispose(); + } + }); + + it.each([ + "promptRequired", + "injected", + ])("Stop before the v1 child %s reply prevents only new input", async (outcome) => { + const agent = new FakeAcpAgent(); + agent.protocolVersion = 1; + agent.steeringOutcome = outcome; + let responseSeen!: () => void; + let releaseResponse!: () => void; + const held = new Promise((resolve) => { + responseSeen = resolve; + }); + const adapter = new AcpAdapter({ + command: "fake", + now: () => 1, + mintId: () => "fixture-id", + createTransport: (_opts, handlers) => + agent.transport({ + ...handlers, + onLine(line) { + const frame = JSON.parse(line); + if (frame.result?.outcome === outcome) { + releaseResponse = () => handlers.onLine(line); + responseSeen(); + } else handlers.onLine(line); + }, + }), + }); + void collect(adapter.start({ cwd: "/work" }), []); + await flush(); + agent.notify("sess-1", announce("child-1")); + await flush(); + try { + const pending = adapter.steerSubagent("subagent:child-1", [ + { type: "text", text: "continue" }, + ]); + await held; + expect(await adapter.stopBackgroundTask("subagent:child-1")).toBe(true); + releaseResponse(); + await expect(pending).resolves.toBe(outcome === "injected"); + expect( + agent.sent.filter((frame) => frame.method === "session/prompt"), + ).toEqual([]); + expect( + agent.sent.filter((frame) => frame.method === "session/cancel"), + ).toHaveLength(1); + } finally { + await adapter.dispose(); + } + }); + it("renders a steer echo in the child", async () => { const agent = new FakeAcpAgent(); const { adapter, events } = startAdapter(agent); diff --git a/packages/chat-runtime/src/harness/acp/acpAdapter.ts b/packages/chat-runtime/src/harness/acp/acpAdapter.ts index d4bcac6c2d8..2eb81e748eb 100644 --- a/packages/chat-runtime/src/harness/acp/acpAdapter.ts +++ b/packages/chat-runtime/src/harness/acp/acpAdapter.ts @@ -237,6 +237,7 @@ type SubagentRun = { parentSessionId: string; parentTaskId?: string; steers: Set<{ echoes: Set; shared: boolean }>; + cancelGeneration: number; }; type SubagentControls = { @@ -701,6 +702,7 @@ export class AcpAdapter implements HarnessAdapter { const run = subagentSessionId && this.subagents.get(subagentSessionId); if (subagentSessionId) { if (!run || run.item.canStop !== true) return false; + run.cancelGeneration++; this.client.notify("session/cancel", { sessionId: subagentSessionId }); this.cancelSubagentApprovals(subagentSessionId, run); return true; @@ -723,13 +725,42 @@ export class AcpAdapter implements HarnessAdapter { const run = subagentSessionId && this.subagents.get(subagentSessionId); if (!client || !subagentSessionId || !run || run.item.canSteer !== true) return false; + const cancelGeneration = run.cancelGeneration; + const prompt = await this.toAcpPrompt(content); + if ( + this.disposed || + this.dead || + this.client !== client || + this.subagents.get(subagentSessionId) !== run || + run.cancelGeneration !== cancelGeneration || + run.item.canSteer !== true + ) + return false; const steer = { echoes: new Set(), shared: run.steers.size > 0 }; for (const other of run.steers) other.shared = true; run.steers.add(steer); try { + if (this.negotiatedVersion < 2 && this.supportsSteering) { + const response = await client.request(ACP_STEERING_METHOD, { + sessionId: subagentSessionId, + prompt, + _meta: { steering: { idleBehavior: "promptRequired" } }, + }); + const outcome = + acpSteeringResponseSchema.safeParse(response).data?.outcome; + if (outcome === "injected" || outcome === "startedNewTurn") return true; + if (outcome !== "promptRequired") return false; + if ( + this.client !== client || + this.subagents.get(subagentSessionId) !== run || + run.cancelGeneration !== cancelGeneration || + run.item.canSteer !== true + ) + return false; + } await client.request("session/prompt", { sessionId: subagentSessionId, - prompt: await this.toAcpPrompt(content), + prompt, }); return true; } catch (error) { @@ -1702,6 +1733,7 @@ export class AcpAdapter implements HarnessAdapter { parentSessionId, ...(parent ? { parentTaskId: parent.item.id } : {}), steers: new Set(), + cancelGeneration: 0, }; this.subagents.set(sessionId, run); this.emitItem(item, turnId); diff --git a/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.test.tsx b/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.test.tsx new file mode 100644 index 00000000000..7a46b6fc964 --- /dev/null +++ b/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.test.tsx @@ -0,0 +1,326 @@ +import { afterAll, afterEach, expect, mock, test } from "bun:test"; +import { GlobalRegistrator } from "@happy-dom/global-registrator"; +import type { LexicalEditor } from "lexical"; +import type { ComposerBodyProps } from "./ComposerBody"; + +const alreadyRegistered = GlobalRegistrator.isRegistered; +if (!alreadyRegistered) GlobalRegistrator.register(); +( + globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean } +).IS_REACT_ACT_ENVIRONMENT = true; + +mock.module("@lingui/react/macro", () => ({ + Trans: ({ children }: { children?: unknown }) => children, + useLingui: () => ({ + t: ({ message }: { message: string }) => message, + }), +})); + +const { act, cleanup, fireEvent, render } = await import( + "@testing-library/react" +); +const { useEffect } = await import("react"); +const { LexicalComposer } = await import("@lexical/react/LexicalComposer"); +const { useLexicalComposerContext } = await import( + "@lexical/react/LexicalComposerContext" +); +const { + $createParagraphNode, + $createTextNode, + $getRoot, + $isElementNode, + KEY_ARROW_DOWN_COMMAND, + KEY_ARROW_UP_COMMAND, + KEY_ENTER_COMMAND, +} = await import("lexical"); +const { MentionChipNode } = await import("../../nodes/mentionChipNode"); +const { ComposerBody } = await import("./ComposerBody"); + +function EditorProbe({ + onEditor, +}: { + onEditor: (editor: LexicalEditor) => void; +}) { + const [editor] = useLexicalComposerContext(); + useEffect(() => onEditor(editor), [editor, onEditor]); + return null; +} + +async function setup( + onSubmit: ComposerBodyProps["onSubmit"], + { history }: Pick = {}, +) { + let editor: LexicalEditor | undefined; + const view = render( + { + throw error; + }, + }} + > + { + editor = value; + }} + /> + + , + ); + const currentEditor = editor; + if (!currentEditor) throw new Error("Editor did not mount"); + const write = async (text: string, chip = false) => { + await act(async () => { + currentEditor.update( + () => { + const paragraph = $createParagraphNode(); + paragraph.append($createTextNode(text)); + if (chip) + paragraph.append( + MentionChipNode.fromChip({ + label: "Example file", + serialized: "@example.ts", + data: { path: "example.ts" }, + }), + ); + $getRoot().clear().append(paragraph); + }, + { discrete: true }, + ); + }); + }; + const attach = (file: File) => { + const input = + view.container.querySelector("input[type=file]"); + if (!input) throw new Error("File input did not mount"); + fireEvent.change(input, { target: { files: [file] } }); + }; + await write("Review the file"); + return { + ...view, + editor: currentEditor, + write, + attach, + text: () => + currentEditor.getEditorState().read(() => $getRoot().getTextContent()), + submit: () => + fireEvent.click(view.getByRole("button", { name: "Send message" })), + }; +} + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason: Error) => void; + const promise = new Promise((onResolve, onReject) => { + resolve = onResolve; + reject = onReject; + }); + return { promise, resolve, reject }; +} + +afterEach(cleanup); +afterAll(async () => { + if (!alreadyRegistered) await GlobalRegistrator.unregister(); +}); + +test("keeps the full draft and attachment while submission is pending", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.attach(new File(["fixture"], "review.txt")); + view.submit(); + await act(async () => {}); + try { + expect(view.text()).toBe("Review the file"); + expect(view.getByText("review.txt")).toBeTruthy(); + } finally { + await act(async () => pending.resolve()); + } +}); + +test("does not submit twice while acceptance is pending", async () => { + const pending = deferred(); + const onSubmit = mock(() => pending.promise); + const view = await setup(onSubmit); + view.submit(); + await view.write("Another draft"); + view.submit(); + await act(async () => pending.resolve()); + expect(onSubmit).toHaveBeenCalledTimes(1); +}); + +test("ignores another Enter key while submission is pending", async () => { + const pending = deferred(); + const onSubmit = mock(() => pending.promise); + const view = await setup(onSubmit); + view.submit(); + await view.write("Another draft"); + await act(async () => { + view.editor.dispatchCommand( + KEY_ENTER_COMMAND, + new KeyboardEvent("keydown", { key: "Enter" }), + ); + }); + await act(async () => pending.resolve()); + expect(onSubmit).toHaveBeenCalledTimes(1); + expect(view.text()).toBe("Another draft"); +}); + +test("keeps a draft when submission returns false", async () => { + const view = await setup(() => false); + view.submit(); + await act(async () => {}); + expect(view.text()).toBe("Review the file"); +}); + +test("keeps text and files when asynchronous submission declines", async () => { + const view = await setup(async () => false); + view.attach(new File(["fixture"], "review.txt")); + view.submit(); + await act(async () => {}); + expect(view.text()).toBe("Review the file"); + expect(view.getByText("review.txt")).toBeTruthy(); +}); + +test("handles rejected submission and leaves the draft available for a manual retry", async () => { + const pending = deferred(); + const onSubmit = mock(() => pending.promise); + const view = await setup(onSubmit); + view.attach(new File(["fixture"], "review.txt")); + view.submit(); + await act(async () => pending.reject(new Error("Transport unavailable"))); + expect(view.text()).toBe("Review the file"); + expect(view.getByText("review.txt")).toBeTruthy(); + expect(onSubmit).toHaveBeenCalledTimes(1); +}); + +test("clears an accepted unchanged draft and its attachments", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.attach(new File(["fixture"], "review.txt")); + view.submit(); + await act(async () => pending.resolve()); + expect(view.text()).toBe(""); + expect(view.queryByText("review.txt")).toBeNull(); +}); + +test("preserves new Lexical chips, text and files after an earlier send succeeds", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.attach(new File(["old"], "review.txt")); + view.submit(); + await view.write("Next task ", true); + view.attach(new File(["new"], "next.txt")); + await act(async () => pending.resolve()); + expect(view.text()).toBe("Next task @example.ts"); + expect(view.getByText("Example file")).toBeTruthy(); + expect(view.getByText("next.txt")).toBeTruthy(); + expect(view.queryByText("review.txt")).toBeNull(); +}); + +test("preserves a draft edited back to its original text while pending", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.submit(); + await view.write("Different draft"); + await view.write("Review the file"); + await act(async () => pending.resolve()); + expect(view.text()).toBe("Review the file"); +}); + +test("preserves the newer draft saved by history navigation while a send is pending", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise, { + history: ["Earlier question"], + }); + view.submit(); + await view.write("Next task draft"); + await act(async () => { + view.editor.update( + () => { + $getRoot().selectStart(); + view.editor.dispatchCommand(KEY_ARROW_UP_COMMAND, null); + }, + { discrete: true }, + ); + }); + expect(view.text()).toBe("Earlier question"); + await act(async () => pending.resolve()); + await act(async () => { + view.editor.update( + () => { + $getRoot().selectEnd(); + view.editor.dispatchCommand(KEY_ARROW_DOWN_COMMAND, null); + }, + { discrete: true }, + ); + }); + expect(view.text()).toBe("Next task draft"); +}); + +test("selection changes do not keep an otherwise accepted unchanged draft", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.submit(); + await act(async () => { + view.editor.update(() => $getRoot().selectStart(), { discrete: true }); + }); + await act(async () => pending.resolve()); + expect(view.text()).toBe(""); +}); + +test("flushes a pending Lexical edit before deciding whether to clear it", async () => { + const pending = deferred(); + const view = await setup(() => pending.promise); + view.submit(); + await act(async () => { + pending.resolve(); + view.editor.update(() => { + const paragraph = $getRoot().getFirstChildOrThrow(); + if (!$isElementNode(paragraph)) throw new Error("Expected a paragraph"); + paragraph.append($createTextNode(" with more context")); + }); + }); + expect(view.text()).toBe("Review the file with more context"); +}); + +test("submission exposes the flushed content revision used for draft clearing", async () => { + const pending = deferred(); + let isDraftUnchanged: (() => boolean) | undefined; + const view = await setup((payload) => { + isDraftUnchanged = payload.isDraftUnchanged; + return pending.promise; + }); + await view.write("@example.ts"); + view.submit(); + expect(isDraftUnchanged?.()).toBe(true); + await act(async () => { + view.editor.update(() => $getRoot().selectStart(), { discrete: true }); + }); + expect(isDraftUnchanged?.()).toBe(true); + await act(async () => { + view.editor.update(() => { + const paragraph = $getRoot().getFirstChildOrThrow(); + if (!$isElementNode(paragraph)) throw new Error("Expected a paragraph"); + paragraph.clear().append( + MentionChipNode.fromChip({ + label: "Example file", + serialized: "@example.ts", + }), + ); + }); + expect(isDraftUnchanged?.()).toBe(false); + }); + await act(async () => pending.resolve()); + expect(view.text()).toBe("@example.ts"); + expect(view.getByText("Example file")).toBeTruthy(); +}); diff --git a/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.tsx b/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.tsx index 5bc1db83d86..5908acfe2ee 100644 --- a/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.tsx +++ b/packages/chat-ui/src/components/PromptInput/components/ComposerBody/ComposerBody.tsx @@ -201,17 +201,20 @@ export function ComposerBody({ }: ComposerBodyProps) { const { t } = useLingui(); const [editor] = useLexicalComposerContext(); + const draftRevision = useRef(0); const onChangeRef = useRef(onChange); onChangeRef.current = onChange; const pasted = usePastedTexts(editor, { enabled: Boolean(collapseLargePastes), - onCardsChange: (cards, editorText) => + onCardsChange: (cards, editorText) => { + draftRevision.current += 1; onChangeRef.current?.( appendPastedText( editorText, cards.map((card) => card.text), ), - ), + ); + }, }); const { cardsRef } = pasted; const draftText = useCallback( @@ -285,6 +288,8 @@ export function ComposerBody({ [draftText, editor, focusAtEnd, pasted.restore, splitPastes], ); const [attachments, setAttachments] = useState([]); + const [isSubmitting, setIsSubmitting] = useState(false); + const submitting = useRef(false); const [isEmpty, setIsEmpty] = useState(true); const [dragging, setDragging] = useState(false); const [mentionQuery, setMentionQuery] = useState(null); @@ -367,10 +372,15 @@ export function ComposerBody({ useEffect( () => - editor.registerUpdateListener(({ editorState }) => { - const text = editorState.read(() => $getRoot().getTextContent()); - onChangeRef.current?.(draftText(text)); - }), + editor.registerUpdateListener( + ({ editorState, dirtyElements, dirtyLeaves }) => { + if (dirtyElements.size > 0 || dirtyLeaves.size > 0) { + draftRevision.current += 1; + } + const text = editorState.read(() => $getRoot().getTextContent()); + onChangeRef.current?.(draftText(text)); + }, + ), [draftText, editor], ); @@ -507,7 +517,8 @@ export function ComposerBody({ if (attachment.previewUrl) URL.revokeObjectURL(attachment.previewUrl); }; - const submit = ({ steer = false }: { steer?: boolean } = {}) => { + const submit = async ({ steer = false }: { steer?: boolean } = {}) => { + if (submitting.current) return; if ( stateRef.current.status === "streaming" && !stateRef.current.submitWhileStreaming @@ -529,25 +540,56 @@ export function ComposerBody({ ) { return; } - if ( - stateRef.current.onSubmit?.({ - text, - files, - mentions, - pastedTexts, - steer, - }) === false - ) - return; - historyNavigationRef.current?.reset(); - if (!stateRef.current.clearOnSubmit) return; - pasted.clear(); - editor.update(() => $getRoot().clear()); - setAttachments((previous) => { - for (const attachment of previous) releaseAttachment(attachment); - return []; - }); - setPanel(null); + const submittedRevision = draftRevision.current; + const submittedPastes = cardsRef.current.map((card) => card.id); + const isDraftUnchanged = () => + editor.read( + () => + draftRevision.current === submittedRevision && + cardsRef.current.length === submittedPastes.length && + cardsRef.current.every( + (card, index) => card.id === submittedPastes[index], + ), + ); + const submittedAttachments = new Set( + stateRef.current.attachments.map((attachment) => attachment.id), + ); + submitting.current = true; + setIsSubmitting(true); + try { + if ( + (await stateRef.current.onSubmit?.({ + isDraftUnchanged, + text, + files, + mentions, + pastedTexts, + steer, + })) === false + ) { + return; + } + if (!stateRef.current.clearOnSubmit) return; + // A delayed acceptance must not erase a newer draft or its Lexical chips. + if (isDraftUnchanged()) { + historyNavigationRef.current?.reset(); + pasted.clear(); + editor.update(() => $getRoot().clear()); + } + setAttachments((previous) => + previous.filter((attachment) => { + if (!submittedAttachments.has(attachment.id)) return true; + releaseAttachment(attachment); + return false; + }), + ); + setPanel(null); + } catch { + // The caller owns error feedback. Keep the editor intact for manual retry. + } finally { + submitting.current = false; + setIsSubmitting(false); + } }; const submitRef = useRef(submit); submitRef.current = submit; @@ -582,7 +624,7 @@ export function ComposerBody({ if (event?.shiftKey || event?.isComposing || event?.keyCode === 229) return false; event?.preventDefault(); - submitRef.current({ + void submitRef.current({ steer: Boolean(event?.metaKey || event?.ctrlKey), }); return true; @@ -1136,10 +1178,11 @@ export function ComposerBody({ aria-label={t({ message: "Send message", })} - disabled={!canSend} - onClick={() => submit()} + disabled={!canSend || isSubmitting} + aria-busy={isSubmitting} + onClick={() => void submit()} className={ - canSend + canSend && !isSubmitting ? cn( FOOTER_BUTTON_CLASS, "cursor-pointer bg-primary text-primary-foreground hover:bg-primary/90", diff --git a/packages/chat-ui/src/components/PromptInput/types.ts b/packages/chat-ui/src/components/PromptInput/types.ts index 2fc94015f2c..a7fdf2854e1 100644 --- a/packages/chat-ui/src/components/PromptInput/types.ts +++ b/packages/chat-ui/src/components/PromptInput/types.ts @@ -91,6 +91,7 @@ export type PromptInputAttachment = { }; export type PromptInputSubmitPayload = { + isDraftUnchanged?: () => boolean; text: string; files: File[]; mentions: ComposerChip[]; @@ -155,10 +156,11 @@ export type PromptInputProps = { // Spans of a stored draft that were chips before it was serialized, so // they come back as chips. Runs on the draft as read, never on typing. findChips?: (text: string) => ComposerChipMatch[]; - /** Returning `false` keeps the draft in the editor. */ + /** Clears only after acceptance. Returning `false` or rejecting keeps the draft. */ onSubmit?: ( payload: PromptInputSubmitPayload, - ) => false | void | Promise; + // biome-ignore lint/suspicious/noConfusingVoidType: existing async handlers resolve void; false explicitly declines submission. + ) => false | void | Promise; onStop?: () => void; stopOnEscape?: boolean; header?: ReactNode;