diff --git a/site/src/pages/AgentsPage/AgentChatInput.tsx b/site/src/pages/AgentsPage/AgentChatInput.tsx index 1ea786c582..a67775dd8a 100644 --- a/site/src/pages/AgentsPage/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/AgentChatInput.tsx @@ -263,6 +263,18 @@ export const AgentChatInput = memo( [onContentChange], ); + // Re-focus the editor after a send completes (isLoading goes + // from true → false) so the user can immediately type again. + // Uses the "store previous value in state" pattern recommended + // by React for responding to prop changes during render. + const [prevIsLoading, setPrevIsLoading] = useState(isLoading); + if (prevIsLoading !== isLoading) { + setPrevIsLoading(isLoading); + if (prevIsLoading && !isLoading) { + internalRef.current?.focus(); + } + } + const canSend = !isDisabled && !isLoading && hasModelOptions && hasContent; const handleSubmit = useCallback(() => { diff --git a/site/src/pages/AgentsPage/AgentDetail.test.ts b/site/src/pages/AgentsPage/AgentDetail.test.ts index 521e285464..4c6e54425d 100644 --- a/site/src/pages/AgentsPage/AgentDetail.test.ts +++ b/site/src/pages/AgentsPage/AgentDetail.test.ts @@ -1,4 +1,5 @@ import { act, renderHook } from "@testing-library/react"; + import { beforeEach, describe, expect, it, vi } from "vitest"; import { draftInputStorageKeyPrefix, @@ -86,6 +87,33 @@ describe("useConversationEditingState", () => { unmount(); }); + it("calls focus on the input ref after a successful send", async () => { + const { result, onSend, unmount } = renderEditing(); + + // Attach a mock ChatMessageInputRef to the chatInputRef + const mockFocus = vi.fn(); + const mockClear = vi.fn(); + const mockInputRef = { + focus: mockFocus, + clear: mockClear, + insertText: vi.fn(), + getValue: vi.fn().mockReturnValue(""), + }; + // The hook exposes chatInputRef – assign the mock to it. + result.current.chatInputRef.current = mockInputRef; + + await act(async () => { + result.current.handleSendFromInput("hello"); + await vi.waitFor(() => { + expect(onSend).toHaveBeenCalledWith("hello", undefined); + }); + }); + + expect(mockClear).toHaveBeenCalled(); + expect(mockFocus).toHaveBeenCalled(); + unmount(); + }); + it("clears the draft from localStorage on successful send", async () => { localStorage.setItem(expectedKey, "draft to clear");