diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx index 782cfed194..053d0ef965 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx @@ -1828,6 +1828,92 @@ describe("useChatStore", () => { expect(result.current.retryState).toBeNull(); }); + it("clears retryState when message_part arrives after retry", async () => { + immediateAnimationFrame(); + + const chatID = "chat-retry-message-part"; + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + retryState: useChatSelector(store, selectRetryState), + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, undefined); + }); + + act(() => { + mockSocket.emitData({ + type: "retry", + chat_id: chatID, + retry: { + attempt: 1, + error: "rate limited", + kind: "rate_limit", + provider: "anthropic", + delay_ms: 3000, + retrying_at: "2025-01-01T00:00:30.000Z", + }, + }); + }); + + await waitFor(() => { + expect(result.current.retryState).toEqual({ + attempt: 1, + error: "rate limited", + kind: "rate_limit", + provider: "anthropic", + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", + }); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "retry recovered" }, + }, + }); + }); + + await waitFor(() => { + expect(result.current.retryState).toBeNull(); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "retry recovered" }, + ]); + }); + }); + it("routes status events for other chatIDs to subagent overrides", async () => { immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts index 3776fc277f..0dfbe31316 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts @@ -946,6 +946,7 @@ export const useChatStore = ( } const part = streamEvent.message_part?.part; if (part) { + store.clearRetryState(); cancelScheduledStreamReset(); partsBuf.push(part); } @@ -962,6 +963,7 @@ export const useChatStore = ( if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { continue; } + store.clearRetryState(); pendingMessages.push(message); if ( message.id !== undefined && @@ -997,6 +999,7 @@ export const useChatStore = ( continue; } + store.clearRetryState(); store.setChatStatus(nextStatus); if (nextStatus === "pending" || nextStatus === "waiting") { store.clearStreamState();