From 5c30867681bcfe713fc287ddf64ff31634400258 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 19 May 2026 11:04:53 -0400 Subject: [PATCH] fix: Preserve Agent Draft Keys Across Switches --- .../SidePanel/Agents/AgentPanel.tsx | 41 +++++++--- .../__tests__/AgentPanel.drafts.spec.tsx | 78 ++++++++++++++++++- 2 files changed, 107 insertions(+), 12 deletions(-) diff --git a/client/src/components/SidePanel/Agents/AgentPanel.tsx b/client/src/components/SidePanel/Agents/AgentPanel.tsx index 1f45ed817d..9b85fe44c1 100644 --- a/client/src/components/SidePanel/Agents/AgentPanel.tsx +++ b/client/src/components/SidePanel/Agents/AgentPanel.tsx @@ -319,12 +319,9 @@ export default function AgentPanel() { [debouncedPersistAgentDraft], ); - const handleAgentChange = useCallback( - (selectedAgentId?: string | null) => { - clearDraftsForAgentIds([currentAgentIdRef.current, selectedAgentId]); - }, - [clearDraftsForAgentIds], - ); + const handleAgentChange = useCallback(() => { + clearDraftsForAgentIds([currentAgentIdRef.current]); + }, [clearDraftsForAgentIds]); const handleCreateNewAgent = useCallback(() => { clearDraftsForAgentIds([currentAgentIdRef.current, undefined]); @@ -365,8 +362,10 @@ export default function AgentPanel() { ); useEffect(() => { - const agentChanged = currentAgentIdRef.current !== current_agent_id; - const userChanged = draftUserIdRef.current !== draftUserId; + const previousAgentId = currentAgentIdRef.current; + const previousUserId = draftUserIdRef.current; + const agentChanged = previousAgentId !== current_agent_id; + const userChanged = previousUserId !== draftUserId; if (!agentChanged && !userChanged) { return; @@ -377,12 +376,25 @@ export default function AgentPanel() { } if (agentChanged && !userChanged && shouldPersistDraftRef.current) { - saveAgentDraft(currentAgentIdRef.current, getValues(), draftUserIdRef.current); + saveAgentDraft(previousAgentId, getValues(), previousUserId); + } + + const nextDraft = getAgentDraft(current_agent_id, draftUserId); + if ( + agentChanged && + !userChanged && + previousAgentId == null && + shouldPersistDraftRef.current && + !nextDraft + ) { + // AgentPanelSwitch can restore the chat agent id after a remount while a new-agent draft is visible. + setCurrentAgentId(undefined); + setHasDraft(shouldPersistDraftRef.current); + return; } currentAgentIdRef.current = current_agent_id; draftUserIdRef.current = draftUserId; - const nextDraft = getAgentDraft(current_agent_id, draftUserId); const nextHasDraft = nextDraft != null; shouldPersistDraftRef.current = nextHasDraft; setHasDraft(nextHasDraft); @@ -394,7 +406,14 @@ export default function AgentPanel() { if (userChanged) { reset(getDefaultAgentFormValues()); } - }, [current_agent_id, debouncedPersistAgentDraft, draftUserId, getValues, reset]); + }, [ + current_agent_id, + debouncedPersistAgentDraft, + draftUserId, + getValues, + reset, + setCurrentAgentId, + ]); useEffect(() => { if (hasDraft) { diff --git a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.drafts.spec.tsx b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.drafts.spec.tsx index f79ff71ac0..85d046df31 100644 --- a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.drafts.spec.tsx +++ b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.drafts.spec.tsx @@ -26,6 +26,7 @@ const PROGRAMMATIC_UPDATE_LABEL = 'Programmatic agent update'; const AVATAR_ACTION_LABEL = 'Avatar action'; const RESET_AVATAR_LABEL = 'Reset avatar'; const DELETE_AGENT_LABEL = 'Delete Agent'; +const SELECT_AGENT_WITH_DRAFT_LABEL = 'Select agent with draft'; const MOCK_USER_ID = 'user-123'; type MockButtonProps = React.ButtonHTMLAttributes & { @@ -35,6 +36,7 @@ type MockButtonProps = React.ButtonHTMLAttributes & { type MockAgentSelectProps = { hasDraft?: boolean; + onAgentChange?: (selectedAgentId?: string | null) => void; }; type MockAgentFooterProps = { @@ -146,7 +148,7 @@ jest.mock('~/common', () => { jest.mock('../AgentSelect', () => ({ __esModule: true, - default: function MockAgentSelect({ hasDraft }: MockAgentSelectProps) { + default: function MockAgentSelect({ hasDraft, onAgentChange }: MockAgentSelectProps) { const { useFormContext } = jest.requireActual( 'react-hook-form', ) as typeof import('react-hook-form'); @@ -162,6 +164,9 @@ jest.mock('../AgentSelect', () => ({ > {PROGRAMMATIC_UPDATE_LABEL} + ); }, @@ -384,6 +389,77 @@ describe('AgentPanel draft preservation', () => { expect(mockLastAgentSelectHasDraft).toBe(true); }); + it('preserves a selected agent draft when leaving another agent', async () => { + mockCurrentAgentId = 'agent-123'; + const { rerender } = render(); + + fireEvent.change(screen.getByLabelText('Draft name'), { + target: { value: 'Draft for current agent' }, + }); + await waitFor(() => { + expect(getAgentDraft('agent-123', MOCK_USER_ID)?.name).toBe('Draft for current agent'); + }); + + saveAgentDraft( + 'agent-456', + { + id: 'agent-456', + name: 'Destination agent draft', + description: '', + instructions: 'Keep the destination draft', + model: 'gpt-4o', + model_parameters: {} as AgentForm['model_parameters'], + provider: { label: 'OpenAI', value: 'openAI' }, + tools: [], + tool_options: {}, + category: 'general', + execute_code: false, + file_search: false, + web_search: false, + } as AgentForm, + MOCK_USER_ID, + ); + + fireEvent.click(screen.getByRole('button', { name: SELECT_AGENT_WITH_DRAFT_LABEL })); + + expect(getAgentDraft('agent-123', MOCK_USER_ID)).toBeUndefined(); + expect(getAgentDraft('agent-456', MOCK_USER_ID)?.name).toBe('Destination agent draft'); + + mockCurrentAgentId = 'agent-456'; + rerender(); + + await waitFor(() => { + expect(screen.getByLabelText('Draft name')).toHaveValue('Destination agent draft'); + }); + }); + + it('keeps remounted new-agent drafts keyed as new when a chat agent id is restored', async () => { + const { rerender } = render(); + + fireEvent.change(screen.getByLabelText('Draft name'), { + target: { value: 'New agent draft' }, + }); + await waitFor(() => { + expect(getAgentDraft(undefined, MOCK_USER_ID)?.name).toBe('New agent draft'); + }); + + mockCurrentAgentId = 'agent-123'; + rerender(); + + await waitFor(() => { + expect(screen.getByLabelText('Draft name')).toHaveValue('New agent draft'); + }); + + fireEvent.change(screen.getByLabelText('Draft name'), { + target: { value: 'New agent draft continued' }, + }); + + await waitFor(() => { + expect(getAgentDraft(undefined, MOCK_USER_ID)?.name).toBe('New agent draft continued'); + }); + expect(getAgentDraft('agent-123', MOCK_USER_ID)).toBeUndefined(); + }); + it('does not carry drafts across user changes while mounted', async () => { const { rerender } = render();