From 6b50ace2bd517e6f511fdccc6df3d33da72c406a Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 21 Mar 2026 19:56:23 +0700 Subject: [PATCH] Refactor agent input draft persistence --- ...ent-input-area.draft-persist-guard.test.ts | 48 ------ .../agent-input-area.draft-persist-guard.ts | 20 --- .../app/src/components/agent-input-area.tsx | 161 +++--------------- .../app/src/hooks/use-agent-input-draft.ts | 161 ++++++++++++++++++ .../src/hooks/use-draft-agent-create-flow.ts | 19 +-- packages/app/src/panels/agent-panel.tsx | 13 ++ .../src/screens/agent/draft-agent-screen.tsx | 20 ++- .../workspace/workspace-draft-agent-tab.tsx | 19 ++- 8 files changed, 224 insertions(+), 237 deletions(-) delete mode 100644 packages/app/src/components/agent-input-area.draft-persist-guard.test.ts delete mode 100644 packages/app/src/components/agent-input-area.draft-persist-guard.ts create mode 100644 packages/app/src/hooks/use-agent-input-draft.ts diff --git a/packages/app/src/components/agent-input-area.draft-persist-guard.test.ts b/packages/app/src/components/agent-input-area.draft-persist-guard.test.ts deleted file mode 100644 index c27a16255..000000000 --- a/packages/app/src/components/agent-input-area.draft-persist-guard.test.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { shouldSkipDraftPersist } from "./agent-input-area.draft-persist-guard"; - -describe("shouldSkipDraftPersist", () => { - it("blocks persist while hydrate for current uncontrolled generation is incomplete", () => { - expect( - shouldSkipDraftPersist({ - isControlled: false, - currentGeneration: 2, - hydratedGeneration: 1, - isCurrentGeneration: true, - }) - ).toBe(true); - }); - - it("allows persist after hydrate completes for current generation", () => { - expect( - shouldSkipDraftPersist({ - isControlled: false, - currentGeneration: 3, - hydratedGeneration: 3, - isCurrentGeneration: true, - }) - ).toBe(false); - }); - - it("blocks persist for stale generations", () => { - expect( - shouldSkipDraftPersist({ - isControlled: false, - currentGeneration: 4, - hydratedGeneration: 4, - isCurrentGeneration: false, - }) - ).toBe(true); - }); - - it("does not block controlled draft persistence", () => { - expect( - shouldSkipDraftPersist({ - isControlled: true, - currentGeneration: 0, - hydratedGeneration: 0, - isCurrentGeneration: true, - }) - ).toBe(false); - }); -}); diff --git a/packages/app/src/components/agent-input-area.draft-persist-guard.ts b/packages/app/src/components/agent-input-area.draft-persist-guard.ts deleted file mode 100644 index ba2b60374..000000000 --- a/packages/app/src/components/agent-input-area.draft-persist-guard.ts +++ /dev/null @@ -1,20 +0,0 @@ -export function shouldSkipDraftPersist(input: { - isControlled: boolean; - currentGeneration: number; - hydratedGeneration: number; - isCurrentGeneration: boolean; -}): boolean { - if (input.isControlled) { - return false; - } - - if (input.currentGeneration <= 0) { - return true; - } - - if (!input.isCurrentGeneration) { - return true; - } - - return input.hydratedGeneration !== input.currentGeneration; -} diff --git a/packages/app/src/components/agent-input-area.tsx b/packages/app/src/components/agent-input-area.tsx index 1aeece82b..6934c835e 100644 --- a/packages/app/src/components/agent-input-area.tsx +++ b/packages/app/src/components/agent-input-area.tsx @@ -10,8 +10,6 @@ import { generateMessageId, type StreamItem } from '@/types/stream' import { AgentStatusBar, DraftAgentStatusBar, type DraftAgentStatusBarProps } from './agent-status-bar' import { useImageAttachmentPicker } from '@/hooks/use-image-attachment-picker' import { useSessionStore } from '@/stores/session-store' -import { useDraftStore } from '@/stores/draft-store' -import { buildDraftStoreKey } from '@/stores/draft-keys' import { MessageInput, type MessagePayload, @@ -35,7 +33,6 @@ import { persistAttachmentFromFileUri, } from '@/attachments/service' import { resolveStatusControlMode } from '@/components/agent-input-area.status-controls' -import { shouldSkipDraftPersist } from '@/components/agent-input-area.draft-persist-guard' import { markScrollInvestigationRender } from '@/utils/scroll-jank-investigation' import { useKeyboardShiftStyle } from '@/hooks/use-keyboard-shift-style' import { useKeyboardActionHandler } from '@/hooks/use-keyboard-action-handler' @@ -47,17 +44,21 @@ type QueuedMessage = { images?: ImageAttachment[] } +type ImageListUpdater = ImageAttachment[] | ((prev: ImageAttachment[]) => ImageAttachment[]) + interface AgentInputAreaProps { agentId: string serverId: string - draftId?: string onSubmitMessage?: (payload: MessagePayload) => Promise /** Externally controlled loading state. When true, disables the submit button. */ isSubmitLoading?: boolean /** When true, blurs the input immediately when submitting. */ blurOnSubmit?: boolean - value?: string - onChangeText?: (text: string) => void + value: string + onChangeText: (text: string) => void + images: ImageAttachment[] + onChangeImages: (updater: ImageListUpdater) => void + clearDraft: (lifecycle: 'sent' | 'abandoned') => void /** When true, auto-focuses the text input on web. */ autoFocus?: boolean /** Callback to expose the addImages function to parent components */ @@ -80,12 +81,14 @@ const MOBILE_MESSAGE_PLACEHOLDER = 'Message, @files, /commands' export function AgentInputArea({ agentId, serverId, - draftId, onSubmitMessage, isSubmitLoading = false, blurOnSubmit = false, value, onChangeText, + images, + onChangeImages, + clearDraft, autoFocus = false, onAddImages, commandDraftConfig, @@ -114,14 +117,6 @@ export function AgentInputArea({ const agent = useSessionStore((state) => state.sessions[serverId]?.agents?.get(agentId)) - const draftStoreKey = buildDraftStoreKey({ serverId, agentId, draftId }) - const getDraftInput = useDraftStore((state) => state.getDraftInput) - const hydrateDraftInput = useDraftStore((state) => state.hydrateDraftInput) - const saveDraftInput = useDraftStore((state) => state.saveDraftInput) - const clearDraftInput = useDraftStore((state) => state.clearDraftInput) - const beginDraftGeneration = useDraftStore((state) => state.beginDraftGeneration) - const isDraftGenerationCurrent = useDraftStore((state) => state.isDraftGenerationCurrent) - const queuedMessagesRaw = useSessionStore((state) => state.sessions[serverId]?.queuedMessages?.get(agentId) ) @@ -131,7 +126,6 @@ export function AgentInputArea({ const setAgentStreamTail = useSessionStore((state) => state.setAgentStreamTail) const setAgentStreamHead = useSessionStore((state) => state.setAgentStreamHead) - const [internalInput, setInternalInput] = useState('') const isDesktopWebBreakpoint = Platform.OS === 'web' && UnistylesRuntime.breakpoint !== 'xs' && @@ -139,17 +133,16 @@ export function AgentInputArea({ const messagePlaceholder = isDesktopWebBreakpoint ? DESKTOP_MESSAGE_PLACEHOLDER : MOBILE_MESSAGE_PLACEHOLDER - const userInput = value ?? internalInput - const setUserInput = onChangeText ?? setInternalInput + const userInput = value + const setUserInput = onChangeText + const selectedImages = images + const setSelectedImages = onChangeImages const [cursorIndex, setCursorIndex] = useState(0) const [isProcessing, setIsProcessing] = useState(false) - const [selectedImages, setSelectedImages] = useState([]) const [isCancellingAgent, setIsCancellingAgent] = useState(false) const [sendError, setSendError] = useState(null) const [isMessageInputFocused, setIsMessageInputFocused] = useState(false) const messageInputRef = useRef(null) - const draftGenerationRef = useRef(0) - const hydratedGenerationRef = useRef(0) const keyboardHandlerIdRef = useRef( `message-input:${serverId}:${agentId}:${Math.random().toString(36).slice(2)}` ) @@ -187,7 +180,7 @@ export function AgentInputArea({ // Expose addImages function to parent for drag-and-drop support const addImages = useCallback((images: ImageAttachment[]) => { setSelectedImages((prev) => [...prev, ...images]) - }, []) + }, [setSelectedImages]) useEffect(() => { onAddImages?.(addImages) @@ -326,10 +319,7 @@ export function AgentInputArea({ return next }) - const isControlled = value !== undefined - if (!isControlled) { - setUserInput('') - } + setUserInput('') setSelectedImages([]) } @@ -354,28 +344,20 @@ export function AgentInputArea({ // Save values so we can restore on error. const savedImages = imageAttachments if (!onSubmitMessageRef.current) { - if (value !== undefined) { - onChangeText?.('') - } else { - setUserInput('') - } + setUserInput('') + setSelectedImages([]) } - setSelectedImages([]) setSendError(null) setIsProcessing(true) try { await submitMessage(trimmedMessage, imageAttachments) - clearDraftInput({ draftKey: draftStoreKey, lifecycle: 'sent' }) + clearDraft('sent') } catch (error) { console.error('[AgentInput] Failed to send message:', error) // Restore input so the user never loses their message if (!onSubmitMessageRef.current) { - if (value !== undefined) { - onChangeText?.(trimmedMessage) - } else { - setUserInput(trimmedMessage) - } + setUserInput(trimmedMessage) } if (savedImages) { setSelectedImages(savedImages) @@ -434,109 +416,6 @@ export function AgentInputArea({ } }, [isAgentRunning, isConnected]) - // Hydrate draft only when switching agents (uncontrolled mode only) - const isControlled = value !== undefined - useEffect(() => { - // Skip draft hydration for controlled inputs - parent manages state - if (isControlled) { - return - } - const generation = beginDraftGeneration(draftStoreKey) - draftGenerationRef.current = generation - hydratedGenerationRef.current = 0 - setUserInput('') - setSelectedImages([]) - let cancelled = false - - void (async () => { - const draft = await hydrateDraftInput(draftStoreKey) - if (cancelled) { - return - } - if (!isDraftGenerationCurrent({ draftKey: draftStoreKey, generation })) { - return - } - if (!draft) { - hydratedGenerationRef.current = generation - return - } - - setUserInput(draft.text) - setSelectedImages(draft.images) - hydratedGenerationRef.current = generation - })() - - return () => { - cancelled = true - } - }, [ - beginDraftGeneration, - draftStoreKey, - hydrateDraftInput, - isControlled, - isDraftGenerationCurrent, - setUserInput, - ]) - - // Persist drafts into the shared session store with change detection to avoid redundant work - useEffect(() => { - const currentGeneration = draftGenerationRef.current - const isCurrentGeneration = - currentGeneration > 0 - ? isDraftGenerationCurrent({ draftKey: draftStoreKey, generation: currentGeneration }) - : true - if ( - shouldSkipDraftPersist({ - isControlled, - currentGeneration, - hydratedGeneration: hydratedGenerationRef.current, - isCurrentGeneration, - }) - ) { - return - } - - const existing = getDraftInput(draftStoreKey) - const isSameText = existing?.text === userInput - const existingImages: ImageAttachment[] = existing?.images ?? [] - const isSameImages = - existingImages.length === selectedImages.length && - existingImages.every((img, idx) => { - return ( - img.id === selectedImages[idx]?.id && - img.mimeType === selectedImages[idx]?.mimeType && - img.storageType === selectedImages[idx]?.storageType && - img.storageKey === selectedImages[idx]?.storageKey - ) - }) - - if (isSameText && isSameImages) { - return - } - - const hasContent = userInput.trim().length > 0 || selectedImages.length > 0 - if (!hasContent) { - if (existing) { - clearDraftInput({ draftKey: draftStoreKey, lifecycle: 'abandoned' }) - } - return - } - - saveDraftInput({ - draftKey: draftStoreKey, - draft: { text: userInput, images: selectedImages }, - }) - }, [ - clearDraftInput, - draftStoreKey, - getDraftInput, - isControlled, - isDraftGenerationCurrent, - saveDraftInput, - selectedImages, - userInput, - ]) - const handleKeyboardAction = useCallback((action: KeyboardActionDefinition): boolean => { if (!isScreenFocused) { return false diff --git a/packages/app/src/hooks/use-agent-input-draft.ts b/packages/app/src/hooks/use-agent-input-draft.ts new file mode 100644 index 000000000..5ba90855a --- /dev/null +++ b/packages/app/src/hooks/use-agent-input-draft.ts @@ -0,0 +1,161 @@ +import { useCallback, useEffect, useRef, useState } from 'react' +import type { AttachmentMetadata } from '@/attachments/types' +import { useDraftStore } from '@/stores/draft-store' + +type ImageUpdater = + | AttachmentMetadata[] + | ((prev: AttachmentMetadata[]) => AttachmentMetadata[]) + +interface AgentInputDraft { + text: string + setText: (text: string) => void + images: AttachmentMetadata[] + setImages: (updater: ImageUpdater) => void + clear: (lifecycle: 'sent' | 'abandoned') => void + isHydrated: boolean +} + +function hasDraftContent(input: { text: string; images: AttachmentMetadata[] }): boolean { + return input.text.trim().length > 0 || input.images.length > 0 +} + +function areImagesEqual(input: { + left: AttachmentMetadata[] + right: AttachmentMetadata[] +}): boolean { + if (input.left.length !== input.right.length) { + return false + } + + return input.left.every((image, index) => { + const other = input.right[index] + return ( + image.id === other?.id && + image.mimeType === other?.mimeType && + image.storageType === other?.storageType && + image.storageKey === other?.storageKey + ) + }) +} + +export function useAgentInputDraft(draftKey: string): AgentInputDraft { + const [text, setText] = useState('') + const [images, setImagesState] = useState([]) + const [isHydrated, setIsHydrated] = useState(false) + const draftGenerationRef = useRef(0) + const hydratedGenerationRef = useRef(0) + + const setImages = useCallback((updater: ImageUpdater) => { + setImagesState((previousImages) => { + if (typeof updater === 'function') { + return updater(previousImages) + } + return updater + }) + }, []) + + const clear = useCallback( + (lifecycle: 'sent' | 'abandoned') => { + const store = useDraftStore.getState() + store.clearDraftInput({ draftKey, lifecycle }) + + const generation = store.beginDraftGeneration(draftKey) + draftGenerationRef.current = generation + hydratedGenerationRef.current = generation + + setText('') + setImagesState([]) + setIsHydrated(true) + }, + [draftKey] + ) + + useEffect(() => { + const store = useDraftStore.getState() + const generation = store.beginDraftGeneration(draftKey) + draftGenerationRef.current = generation + hydratedGenerationRef.current = 0 + + setText('') + setImagesState([]) + setIsHydrated(false) + + let cancelled = false + + void (async () => { + const draft = await store.hydrateDraftInput(draftKey) + if (cancelled) { + return + } + if (!useDraftStore.getState().isDraftGenerationCurrent({ draftKey, generation })) { + return + } + + if (draft) { + setText(draft.text) + setImagesState(draft.images) + } + + hydratedGenerationRef.current = generation + setIsHydrated(true) + })() + + return () => { + cancelled = true + } + }, [draftKey]) + + useEffect(() => { + const currentGeneration = draftGenerationRef.current + if (currentGeneration <= 0) { + return + } + + const store = useDraftStore.getState() + const isCurrentGeneration = store.isDraftGenerationCurrent({ + draftKey, + generation: currentGeneration, + }) + if (!isCurrentGeneration) { + return + } + if (hydratedGenerationRef.current !== currentGeneration) { + return + } + + const existing = store.getDraftInput(draftKey) + const isSameDraft = + existing?.text === text && + areImagesEqual({ + left: existing?.images ?? [], + right: images, + }) + if (isSameDraft) { + return + } + + if (!hasDraftContent({ text, images })) { + if (existing) { + store.clearDraftInput({ draftKey, lifecycle: 'abandoned' }) + } + return + } + + store.saveDraftInput({ + draftKey, + draft: { + text, + images, + }, + }) + }, [draftKey, images, text]) + + return { + text, + setText, + images, + setImages, + clear, + isHydrated, + } +} diff --git a/packages/app/src/hooks/use-draft-agent-create-flow.ts b/packages/app/src/hooks/use-draft-agent-create-flow.ts index d2eb633af..3ce5aa40a 100644 --- a/packages/app/src/hooks/use-draft-agent-create-flow.ts +++ b/packages/app/src/hooks/use-draft-agent-create-flow.ts @@ -12,11 +12,10 @@ type CreateAttempt = { } type DraftAgentMachineState = - | { tag: 'draft'; promptText: string; errorMessage: string } + | { tag: 'draft'; errorMessage: string } | { tag: 'creating'; attempt: CreateAttempt } type DraftAgentMachineEvent = - | { type: 'DRAFT_SET_PROMPT'; text: string } | { type: 'DRAFT_SET_ERROR'; message: string } | { type: 'SUBMIT'; attempt: CreateAttempt } | { type: 'CREATE_FAILED'; message: string } @@ -27,12 +26,6 @@ function assertNever(value: never): never { function reducer(state: DraftAgentMachineState, event: DraftAgentMachineEvent): DraftAgentMachineState { switch (event.type) { - case 'DRAFT_SET_PROMPT': { - if (state.tag !== 'draft') { - return state - } - return { ...state, promptText: event.text } - } case 'DRAFT_SET_ERROR': { if (state.tag !== 'draft') { return state @@ -46,7 +39,7 @@ function reducer(state: DraftAgentMachineState, event: DraftAgentMachineEvent): if (state.tag !== 'creating') { return state } - return { tag: 'draft', promptText: state.attempt.text, errorMessage: event.message } + return { tag: 'draft', errorMessage: event.message } } default: return assertNever(event) @@ -94,7 +87,6 @@ export function useDraftAgentCreateFlow({ }: UseDraftAgentCreateFlowOptions) { const [machine, dispatch] = useReducer(reducer, { tag: 'draft', - promptText: '', errorMessage: '', } as DraftAgentMachineState) @@ -103,7 +95,6 @@ export function useDraftAgentCreateFlow({ const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle) const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear) - const promptValue = machine.tag === 'draft' ? machine.promptText : '' const formErrorMessage = machine.tag === 'draft' ? machine.errorMessage : '' const isSubmitting = machine.tag === 'creating' @@ -132,10 +123,6 @@ export function useDraftAgentCreateFlow({ return buildDraftAgent(machine.attempt) }, [buildDraftAgent, machine]) - const setPromptText = useCallback((text: string) => { - dispatch({ type: 'DRAFT_SET_PROMPT', text }) - }, []) - const handleCreateFromInput = useCallback( async ({ text, images }: SubmitContext) => { if (isSubmitting) { @@ -226,12 +213,10 @@ export function useDraftAgentCreateFlow({ return { machine, - promptValue, formErrorMessage, isSubmitting, optimisticStreamItems, draftAgent, - setPromptText, handleCreateFromInput, } } diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index ccaea5a86..0a45a4a55 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -20,6 +20,7 @@ import { } from "@/hooks/use-agent-screen-state-machine"; import { useArchiveAgent } from "@/hooks/use-archive-agent"; import { useDelayedHistoryRefreshToast } from "@/hooks/use-delayed-history-refresh-toast"; +import { useAgentInputDraft } from "@/hooks/use-agent-input-draft"; import { useKeyboardShiftStyle } from "@/hooks/use-keyboard-shift-style"; import { usePaneContext } from "@/panels/pane-context"; import type { PanelDescriptor, PanelRegistration } from "@/panels/panel-registry"; @@ -42,6 +43,7 @@ import { import { mergePendingCreateImages } from "@/utils/pending-create-images"; import { deriveSidebarStateBucket } from "@/utils/sidebar-agent-state"; import { useCreateFlowStore } from "@/stores/create-flow-store"; +import { buildDraftStoreKey } from "@/stores/draft-keys"; import { useSessionStore, type Agent } from "@/stores/session-store"; import type { PendingPermission } from "@/types/shared"; import type { StreamItem } from "@/types/stream"; @@ -234,6 +236,12 @@ function AgentPanelBody({ routeKey: string; reason: "initial-entry" | "resume"; } | null>(null); + const agentInputDraft = useAgentInputDraft( + buildDraftStoreKey({ + serverId, + agentId: agentId ?? "__pending__", + }) + ); const handleFilesDropped = useCallback((files: ImageAttachment[]) => { addImagesRef.current?.(files); @@ -851,6 +859,11 @@ function AgentPanelBody({ ( initialWorktreeMode @@ -766,12 +774,10 @@ function DraftAgentScreenContent({ ]) const { - promptValue, formErrorMessage, isSubmitting, optimisticStreamItems, draftAgent, - setPromptText, handleCreateFromInput, } = useDraftAgentCreateFlow({ draftId: draftIdRef.current, @@ -1184,12 +1190,14 @@ function DraftAgentScreenContent({ onSubmitMessage={handleCreateFromInput} isSubmitLoading={isSubmitting} blurOnSubmit={true} - value={promptValue} - onChangeText={setPromptText} + value={draftInput.text} + onChangeText={draftInput.setText} + images={draftInput.images} + onChangeImages={draftInput.setImages} + clearDraft={draftInput.clear} autoFocus={!isSubmitting} onAddImages={handleAddImagesCallback} commandDraftConfig={draftCommandConfig} - draftId={draftIdRef.current} statusControls={{ providerDefinitions, selectedProvider, diff --git a/packages/app/src/screens/workspace/workspace-draft-agent-tab.tsx b/packages/app/src/screens/workspace/workspace-draft-agent-tab.tsx index 3669d68e3..4ad98745f 100644 --- a/packages/app/src/screens/workspace/workspace-draft-agent-tab.tsx +++ b/packages/app/src/screens/workspace/workspace-draft-agent-tab.tsx @@ -6,8 +6,10 @@ import { FileDropZone } from '@/components/file-drop-zone' import { AgentStreamView } from '@/components/agent-stream-view' import type { ImageAttachment } from '@/components/message-input' import { useAgentFormState } from '@/hooks/use-agent-form-state' +import { useAgentInputDraft } from '@/hooks/use-agent-input-draft' import { useDraftAgentCreateFlow } from '@/hooks/use-draft-agent-create-flow' import { useHostRuntimeClient, useHostRuntimeIsConnected } from '@/runtime/host-runtime' +import { buildDraftStoreKey } from '@/stores/draft-keys' import type { Agent } from '@/stores/session-store' import { encodeImages } from '@/utils/encode-images' import type { AgentCapabilityFlags, AgentSessionConfig } from '@server/server/agent/agent-sdk-types' @@ -43,6 +45,13 @@ export function WorkspaceDraftAgentTab({ const client = useHostRuntimeClient(serverId) const isConnected = useHostRuntimeIsConnected(serverId) const addImagesRef = useRef<((images: ImageAttachment[]) => void) | null>(null) + const draftInput = useAgentInputDraft( + buildDraftStoreKey({ + serverId, + agentId: tabId, + draftId, + }) + ) const { selectedProvider, @@ -81,12 +90,10 @@ export function WorkspaceDraftAgentTab({ }, [setWorkingDir, workingDir, workspaceId]) const { - promptValue, formErrorMessage, isSubmitting, optimisticStreamItems, draftAgent, - setPromptText, handleCreateFromInput, } = useDraftAgentCreateFlow({ draftId, @@ -233,12 +240,14 @@ export function WorkspaceDraftAgentTab({ onSubmitMessage={handleCreateFromInput} isSubmitLoading={isSubmitting} blurOnSubmit={true} - value={promptValue} - onChangeText={setPromptText} + value={draftInput.text} + onChangeText={draftInput.setText} + images={draftInput.images} + onChangeImages={draftInput.setImages} + clearDraft={draftInput.clear} autoFocus={!isSubmitting} onAddImages={handleAddImagesCallback} commandDraftConfig={draftCommandConfig} - draftId={draftId} statusControls={{ providerDefinitions, selectedProvider,