fix: Preserve Agent Draft Keys Across Switches

This commit is contained in:
Danny Avila 2026-05-19 11:04:53 -04:00
parent a1579229e8
commit 5c30867681
2 changed files with 107 additions and 12 deletions

View file

@ -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) {

View file

@ -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<HTMLButtonElement> & {
@ -35,6 +36,7 @@ type MockButtonProps = React.ButtonHTMLAttributes<HTMLButtonElement> & {
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}
</button>
<button type="button" onClick={() => onAgentChange?.('agent-456')}>
{SELECT_AGENT_WITH_DRAFT_LABEL}
</button>
</>
);
},
@ -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(<Harness />);
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(<Harness />);
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(<Harness />);
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(<Harness />);
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(<Harness />);