diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 47dbf08288..d31c599eda 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -23,7 +23,10 @@ import globalAxios, { type AxiosInstance, isAxiosError } from "axios"; import type dayjs from "dayjs"; import userAgentParser from "ua-parser-js"; import { delay } from "../utils/delay"; -import { OneWayWebSocket } from "../utils/OneWayWebSocket"; +import { + OneWayWebSocket, + type OneWayWebSocketApi, +} from "../utils/OneWayWebSocket"; import { type FieldError, isApiError } from "./errors"; import type { DeleteExternalAuthByIDResponse, @@ -142,7 +145,7 @@ export const watchWorkspace = ( export const watchChat = ( chatId: string, afterMessageId?: number, -): OneWayWebSocket => { +): OneWayWebSocketApi => { const params = new URLSearchParams(); if (afterMessageId !== undefined && afterMessageId > 0) { params.set("after_id", afterMessageId.toString()); diff --git a/site/src/pages/AgentsPage/AgentDetail.tsx b/site/src/pages/AgentsPage/AgentDetail.tsx index c05b2423b7..90eb51f098 100644 --- a/site/src/pages/AgentsPage/AgentDetail.tsx +++ b/site/src/pages/AgentsPage/AgentDetail.tsx @@ -40,7 +40,11 @@ import { pageTitle } from "utils/page"; import { rewriteLocalhostURL } from "utils/portForward"; import type { AgentsOutletContext } from "./AgentsPage"; import type { ChatMessageInputRef } from "./components/AgentChatInput"; -import { useChatStore } from "./components/AgentDetail/ChatContext"; +import { + selectChatStatus, + useChatSelector, + useChatStore, +} from "./components/AgentDetail/ChatContext"; import { getParentChatID, getWorkspaceAgent, @@ -66,6 +70,7 @@ import { } from "./utils/modelOptions"; import { parsePullRequestUrl } from "./utils/pullRequest"; import { + type ChatDetailError, formatUsageLimitMessage, isUsageLimitData, } from "./utils/usageLimitMessage"; @@ -95,17 +100,22 @@ export function useConversationEditingState(deps: { ? `${draftInputStorageKeyPrefix}${chatID}` : null; const [editorInitialValue, setEditorInitialValue] = useState(() => { - if (!draftStorageKey) { + if (typeof window === "undefined" || !draftStorageKey) { return ""; } return localStorage.getItem(draftStorageKey) ?? ""; }); - // Sync the ref with the editor value so callers that read - // inputValueRef.current see the persisted draft. Uses a layout - // effect so the value is available before paint. + // Sync the ref with the initial draft value so callers that + // read inputValueRef.current see the persisted draft. Uses a + // layout effect so the value is available before paint. + const initialSyncDone = useRef(false); useLayoutEffect(() => { - inputValueRef.current = editorInitialValue; + if (!initialSyncDone.current && editorInitialValue) { + initialSyncDone.current = true; + (inputValueRef as React.MutableRefObject).current = + editorInitialValue; + } }, [editorInitialValue, inputValueRef]); // -- History editing state -- @@ -117,14 +127,6 @@ export function useConversationEditingState(deps: { readonly ChatMessagePart[] >([]); - // -- Queue editing state -- - const [editingQueuedMessageID, setEditingQueuedMessageID] = useState< - number | null - >(null); - const [draftBeforeQueueEdit, setDraftBeforeQueueEdit] = useState< - string | null - >(null); - const handleEditUserMessage = ( messageId: number, text: string, @@ -151,6 +153,14 @@ export function useConversationEditingState(deps: { } }; + // -- Queue editing state -- + const [editingQueuedMessageID, setEditingQueuedMessageID] = useState< + number | null + >(null); + const [draftBeforeQueueEdit, setDraftBeforeQueueEdit] = useState< + string | null + >(null); + const handleStartQueueEdit = ( id: number, text: string, @@ -230,6 +240,30 @@ export function useConversationEditingState(deps: { }; } +const getPersistedDetailError = ({ + chatStatus, + chatRecord, + cachedError, +}: { + chatStatus: TypesGen.ChatStatus | null; + chatRecord: TypesGen.Chat | undefined; + cachedError: ChatDetailError | undefined; +}): ChatDetailError | undefined => { + if (cachedError?.kind === "usage-limit") { + return cachedError; + } + if (chatStatus === "error") { + if (cachedError) { + return cachedError; + } + const lastError = chatRecord?.last_error?.trim(); + if (lastError) { + return { kind: "generic", message: lastError }; + } + } + return undefined; +}; + /** * Resolves the effective compaction threshold for a model configuration, * preferring the user's override when set. @@ -239,11 +273,17 @@ function resolveCompactionThreshold( userThresholds: readonly TypesGen.UserChatCompactionThreshold[] | undefined, modelConfigs: readonly TypesGen.ChatModelConfig[], ): number | undefined { - if (!modelConfigID) return undefined; - const config = modelConfigs.find((c) => c.id === modelConfigID); - if (!config) return undefined; + if (!modelConfigID) { + return undefined; + } + const config = modelConfigs.find( + (modelConfig) => modelConfig.id === modelConfigID, + ); + if (!config) { + return undefined; + } const userOverride = userThresholds?.find( - (t) => t.model_config_id === modelConfigID, + (threshold) => threshold.model_config_id === modelConfigID, ); if (userOverride) { return userOverride.threshold_percent; @@ -475,6 +515,13 @@ const AgentDetail: FC = () => { setChatErrorReason, clearChatErrorReason, }); + const liveChatStatus = + useChatSelector(store, selectChatStatus) ?? chatRecord?.status ?? null; + const persistedError = getPersistedDetailError({ + chatStatus: liveChatStatus, + chatRecord, + cachedError: agentId ? chatErrorReasons[agentId] : undefined, + }); // Git watcher: runs regardless of sidebar visibility, but only // connects when the workspace agent is in the "connected" state @@ -555,15 +602,19 @@ const AgentDetail: FC = () => { error.response?.status === 409 && isUsageLimitData(error.response.data) ) { - setChatErrorReason(agentId, { + const reason: ChatDetailError = { kind: "usage-limit", message: formatUsageLimitMessage(error.response.data), - }); + }; + store.setStreamError(reason); + setChatErrorReason(agentId, reason); } else if (isApiError(error)) { - setChatErrorReason(agentId, { + const reason: ChatDetailError = { kind: "generic", message: error.message || "An unexpected error occurred.", - }); + }; + store.setStreamError(reason); + setChatErrorReason(agentId, reason); } }; @@ -868,8 +919,7 @@ const AgentDetail: FC = () => { agentId={agentId} chatTitle={chatTitle} parentChat={parentChat} - chatErrorReasons={chatErrorReasons} - chatRecord={chatRecord} + persistedError={persistedError} isArchived={isArchived} hasWorkspace={Boolean(workspaceId)} store={store} diff --git a/site/src/pages/AgentsPage/AgentEmbedPage.tsx b/site/src/pages/AgentsPage/AgentEmbedPage.tsx index 5340126e26..1506317d71 100644 --- a/site/src/pages/AgentsPage/AgentEmbedPage.tsx +++ b/site/src/pages/AgentsPage/AgentEmbedPage.tsx @@ -13,7 +13,10 @@ import { bootstrapChatEmbedSession, EmbedContext, } from "./components/EmbedContext"; -import type { ChatDetailError } from "./utils/usageLimitMessage"; +import { + type ChatDetailError, + chatDetailErrorsEqual, +} from "./utils/usageLimitMessage"; type BootstrapMessage = { type: "coder:vscode-auth-bootstrap"; @@ -72,18 +75,18 @@ const AgentEmbedPage: FC = () => { if (!chatId || !trimmedMessage) { return; } + const nextReason: ChatDetailError = { + ...reason, + message: trimmedMessage, + }; setChatErrorReasons((current) => { const existing = current[chatId]; - if ( - existing && - existing.kind === reason.kind && - existing.message === trimmedMessage - ) { + if (chatDetailErrorsEqual(existing, nextReason)) { return current; } return { ...current, - [chatId]: { kind: reason.kind, message: trimmedMessage }, + [chatId]: nextReason, }; }); }; diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index f33768eb31..b755038d76 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -39,7 +39,10 @@ import { shouldNavigateAfterArchive, } from "./utils/agentWorkspaceUtils"; import { getModelOptionsFromCatalog } from "./utils/modelOptions"; -import type { ChatDetailError } from "./utils/usageLimitMessage"; +import { + type ChatDetailError, + chatDetailErrorsEqual, +} from "./utils/usageLimitMessage"; // Type guard for SSE events from the chat list watch endpoint. // Shallow-compare two ChatDiffStatus objects by their meaningful @@ -205,18 +208,18 @@ const AgentsPage: FC = () => { if (!chatId || !trimmedMessage) { return; } + const nextReason: ChatDetailError = { + ...reason, + message: trimmedMessage, + }; setChatErrorReasons((current) => { const existing = current[chatId]; - if ( - existing && - existing.kind === reason.kind && - existing.message === trimmedMessage - ) { + if (chatDetailErrorsEqual(existing, nextReason)) { return current; } return { ...current, - [chatId]: { kind: reason.kind, message: trimmedMessage }, + [chatId]: nextReason, }; }); }; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx index ebfc94e412..e30e8daef0 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.test.tsx @@ -39,6 +39,7 @@ import { selectChatStatus, selectOrderedMessageIDs, selectQueuedMessages, + selectReconnectState, selectRetryState, selectStreamError, selectStreamState, @@ -58,22 +59,37 @@ type ErrorListener = (payload: Event) => void; type OpenListener = (payload: Event) => void; type CloseListener = (payload: CloseEvent) => void; -interface MockSocket { - addEventListener(event: "message", callback: MessageListener): void; - addEventListener(event: "error", callback: ErrorListener): void; - addEventListener(event: "open", callback: OpenListener): void; - addEventListener(event: "close", callback: CloseListener): void; - removeEventListener(event: "message", callback: MessageListener): void; - removeEventListener(event: "error", callback: ErrorListener): void; - removeEventListener(event: "open", callback: OpenListener): void; - removeEventListener(event: "close", callback: CloseListener): void; - close: () => void; +type WatchChatSocket = ReturnType; + +type MockSocketHelpers = { emitOpen: () => void; emitData: (event: TypesGen.ChatStreamEvent) => void; emitDataBatch: (events: readonly TypesGen.ChatStreamEvent[]) => void; emitError: () => void; emitClose: () => void; -} +}; + +type MockSocket = WatchChatSocket & MockSocketHelpers; + +const mockWatchChatReturn = (socket: MockSocket): void => { + vi.mocked(watchChat).mockReturnValue(socket); +}; + +const mockWatchChatReturnOnce = (socket: MockSocket): void => { + vi.mocked(watchChat).mockReturnValueOnce(socket); +}; + +const mockWatchChatWithFreshSockets = ( + watchMock = vi.mocked(watchChat), +): MockSocket[] => { + const sockets: MockSocket[] = []; + watchMock.mockImplementation(() => { + const socket = createMockSocket(); + sockets.push(socket); + return socket; + }); + return sockets; +}; const createMockSocket = (): MockSocket => { const messageListeners = new Set(); @@ -81,7 +97,7 @@ const createMockSocket = (): MockSocket => { const openListeners = new Set(); const closeListeners = new Set(); - const addEventListener = ( + const addEventListener = (( event: "message" | "error" | "open" | "close", callback: MessageListener | ErrorListener | OpenListener | CloseListener, ): void => { @@ -98,9 +114,9 @@ const createMockSocket = (): MockSocket => { return; } errorListeners.add(callback as ErrorListener); - }; + }) as WatchChatSocket["addEventListener"]; - const removeEventListener = ( + const removeEventListener = (( event: "message" | "error" | "open" | "close", callback: MessageListener | ErrorListener | OpenListener | CloseListener, ): void => { @@ -117,9 +133,10 @@ const createMockSocket = (): MockSocket => { return; } errorListeners.delete(callback as ErrorListener); - }; + }) as WatchChatSocket["removeEventListener"]; return { + url: "ws://example.test/api/experimental/chats/mock-stream", addEventListener, removeEventListener, close: vi.fn(), @@ -225,6 +242,8 @@ const immediateAnimationFrame = (): void => { }; afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); vi.restoreAllMocks(); vi.mocked(watchChat).mockReset(); }); @@ -236,7 +255,7 @@ describe("useChatStore", () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -317,7 +336,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "user", "hello"); const newMessage = makeMessage(chatID, 2, "assistant", "done"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -392,7 +411,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "assistant", "old"); const updatedMessage = makeMessage(chatID, 1, "assistant", "updated"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -466,7 +485,7 @@ describe("useChatStore", () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const setChatErrorReason = vi.fn(); @@ -560,7 +579,7 @@ describe("useChatStore", () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -635,7 +654,7 @@ describe("useChatStore", () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -726,7 +745,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "user", "hello"); const queuedMessage = makeQueuedMessage(chatID, 10, "queued"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -800,7 +819,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "user", "hello"); const queuedMessage = makeQueuedMessage(chatID, 10, "queued"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -871,7 +890,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "user", "hello"); const queuedMessage = makeQueuedMessage(chatID, 10, "queued"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -953,11 +972,11 @@ describe("useChatStore", () => { const mockSocket2 = createMockSocket(); // Use a fallback so that extra effect re-runs (caused by // dependency changes during rerender) get a valid socket. - vi.mocked(watchChat).mockReturnValue(mockSocket2 as never); + mockWatchChatReturn(mockSocket2); vi.mocked(watchChat) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never); + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1042,7 +1061,7 @@ describe("useChatStore", () => { const existingMessage = makeMessage(chatID, 1, "user", "hello"); const queuedMessage = makeQueuedMessage(chatID, 10, "queued"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1098,7 +1117,7 @@ describe("useChatStore", () => { const chatID = "chat-1"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1200,7 +1219,7 @@ describe("useChatStore", () => { const chatID = "chat-raf"; const existingMessage = makeMessage(chatID, 1, "user", "hello"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1294,11 +1313,11 @@ describe("useChatStore", () => { const mockSocket1 = createMockSocket(); const mockSocket2 = createMockSocket(); // Use a fallback so that extra effect re-runs get a valid socket. - vi.mocked(watchChat).mockReturnValue(mockSocket2 as never); + mockWatchChatReturn(mockSocket2); vi.mocked(watchChat) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never); + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1385,11 +1404,11 @@ describe("useChatStore", () => { const mockSocket1 = createMockSocket(); const mockSocket2 = createMockSocket(); // Use a fallback so that extra effect re-runs get a valid socket. - vi.mocked(watchChat).mockReturnValue(mockSocket2 as never); + mockWatchChatReturn(mockSocket2); vi.mocked(watchChat) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never) - .mockReturnValueOnce(mockSocket1 as never); + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1) + .mockReturnValueOnce(mockSocket1); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1459,7 +1478,7 @@ describe("useChatStore", () => { const chatID = "chat-status-guard"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1529,7 +1548,7 @@ describe("useChatStore", () => { const chatID = "chat-error"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1570,18 +1589,33 @@ describe("useChatStore", () => { mockSocket.emitData({ type: "error", chat_id: chatID, - error: { message: "Rate limit exceeded", retryable: true }, + error: { + message: "Rate limit exceeded", + kind: "rate_limit", + provider: "anthropic", + retryable: true, + status_code: 429, + }, }); }); await waitFor(() => { expect(result.current.chatStatus).toBe("error"); }); - expect(result.current.streamError).toBe("Rate limit exceeded"); + expect(result.current.streamError).toEqual({ + kind: "rate_limit", + message: "Rate limit exceeded", + provider: "anthropic", + retryable: true, + statusCode: 429, + }); expect(result.current.retryState).toBeNull(); expect(setChatErrorReason).toHaveBeenCalledWith(chatID, { - kind: "generic", + kind: "rate_limit", message: "Rate limit exceeded", + provider: "anthropic", + retryable: true, + statusCode: 429, }); }); @@ -1590,7 +1624,7 @@ describe("useChatStore", () => { const chatID = "chat-error-empty"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1634,7 +1668,13 @@ describe("useChatStore", () => { }); await waitFor(() => { - expect(result.current.streamError).toBe("Chat processing failed."); + expect(result.current.streamError).toEqual({ + kind: "generic", + message: "Chat processing failed.", + provider: undefined, + retryable: false, + statusCode: undefined, + }); }); }); @@ -1643,7 +1683,7 @@ describe("useChatStore", () => { const chatID = "chat-retry"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1685,17 +1725,22 @@ describe("useChatStore", () => { retry: { attempt: 2, error: "upstream timeout", + kind: "timeout", + provider: "anthropic", delay_ms: 5000, retrying_at: "2025-01-01T00:01:00.000Z", }, }); }); - await waitFor(() => { - expect(result.current.retryState).toEqual({ - attempt: 2, - error: "upstream timeout", - }); + await act(async () => {}); + expect(result.current.retryState).toEqual({ + attempt: 2, + error: "upstream timeout", + kind: "timeout", + provider: "anthropic", + delayMs: 5000, + retryingAt: "2025-01-01T00:01:00.000Z", }); }); @@ -1704,7 +1749,7 @@ describe("useChatStore", () => { const chatID = "chat-retry-clear"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1748,6 +1793,8 @@ describe("useChatStore", () => { retry: { attempt: 1, error: "rate limited", + kind: "rate_limit", + provider: "anthropic", delay_ms: 3000, retrying_at: "2025-01-01T00:00:30.000Z", }, @@ -1755,7 +1802,14 @@ describe("useChatStore", () => { }); await waitFor(() => { - expect(result.current.retryState).not.toBeNull(); + expect(result.current.retryState).toEqual({ + attempt: 1, + error: "rate limited", + kind: "rate_limit", + provider: "anthropic", + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", + }); }); // Transition to running — should clear retry state. @@ -1779,7 +1833,7 @@ describe("useChatStore", () => { const chatID = "chat-main"; const subagentChatID = "chat-subagent-1"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1836,12 +1890,12 @@ describe("useChatStore", () => { expect(result.current.chatStatus).toBe("running"); }); - it("sets streamError on WebSocket disconnect and reconnects", async () => { + it("sets reconnectState on WebSocket disconnect and clears it after reconnect", async () => { immediateAnimationFrame(); const chatID = "chat-disconnect"; const mockSocket1 = createMockSocket(); - vi.mocked(watchChat).mockReturnValueOnce(mockSocket1 as never); + mockWatchChatReturnOnce(mockSocket1); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1866,7 +1920,8 @@ describe("useChatStore", () => { clearChatErrorReason, }); return { - streamError: useChatSelector(store, selectStreamError), + chatStatus: useChatSelector(store, selectChatStatus), + reconnectState: useChatSelector(store, selectReconnectState), }; }, { wrapper }, @@ -1882,15 +1937,20 @@ describe("useChatStore", () => { }); await waitFor(() => { - expect(result.current.streamError).toBe( - "Chat stream disconnected. Reconnecting\u2026", + expect(result.current.reconnectState).toMatchObject({ + attempt: 1, + delayMs: 1000, + }); + expect(result.current.reconnectState?.retryingAt).toEqual( + expect.any(String), ); + expect(result.current.chatStatus).toBe("running"); }); // The reconnect timer fires after 1s. Since we're not // using fake timers, waitFor will naturally wait. const mockSocket2 = createMockSocket(); - vi.mocked(watchChat).mockReturnValueOnce(mockSocket2 as never); + mockWatchChatReturnOnce(mockSocket2); await waitFor( () => { @@ -1904,16 +1964,96 @@ describe("useChatStore", () => { mockSocket2.emitOpen(); }); + await waitFor(() => { + expect(result.current.reconnectState).toBeNull(); + expect(result.current.chatStatus).toBe("running"); + }); + }); + + it("clears stale streamError when a reconnected socket opens", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + immediateAnimationFrame(); + + const chatID = "chat-reconnect-clear-error"; + const watchMock = vi.mocked(watchChat); + const sockets = mockWatchChatWithFreshSockets(watchMock); + + 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 { + store, + streamError: useChatSelector(store, selectStreamError), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchMock).toHaveBeenCalledWith(chatID, undefined); + }); + + const socket1 = sockets[0]!; + act(() => { + socket1.emitOpen(); + result.current.store.setStreamError({ + kind: "generic", + message: "Stale transport failure.", + }); + }); + + await waitFor(() => { + expect(result.current.streamError).toEqual({ + kind: "generic", + message: "Stale transport failure.", + }); + }); + + act(() => { + socket1.emitClose(); + }); + + await act(async () => { + vi.advanceTimersByTime(1_500); + }); + + expect(watchMock).toHaveBeenCalledTimes(2); + const socket2 = sockets[1]!; + + act(() => { + socket2.emitOpen(); + }); + await waitFor(() => { expect(result.current.streamError).toBeNull(); }); }); - it("does not overwrite existing streamError on WebSocket disconnect", async () => { + + it("keeps terminal streamError when a WebSocket disconnect follows it", async () => { immediateAnimationFrame(); const chatID = "chat-disconnect-existing"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1939,6 +2079,7 @@ describe("useChatStore", () => { }); return { streamError: useChatSelector(store, selectStreamError), + reconnectState: useChatSelector(store, selectReconnectState), }; }, { wrapper }, @@ -1953,25 +2094,86 @@ describe("useChatStore", () => { mockSocket.emitData({ type: "error", chat_id: chatID, - error: { message: "Rate limit exceeded", retryable: true }, + error: { message: "Rate limit exceeded", retryable: false }, }); }); await waitFor(() => { - expect(result.current.streamError).toBe("Rate limit exceeded"); + expect(result.current.streamError).toEqual({ + kind: "generic", + message: "Rate limit exceeded", + provider: undefined, + retryable: false, + statusCode: undefined, + }); + expect(result.current.reconnectState).toBeNull(); }); - // WebSocket disconnect overwrites with reconnecting message - // since the reconnect logic always shows the disconnect - // notice on first disconnect. + // WebSocket disconnect should not overwrite the terminal error + // or surface reconnect state once the turn has already failed. act(() => { mockSocket.emitError(); }); await waitFor(() => { - expect(result.current.streamError).toBe( - "Chat stream disconnected. Reconnecting\u2026", - ); + expect(result.current.streamError).toEqual({ + kind: "generic", + message: "Rate limit exceeded", + provider: undefined, + retryable: false, + statusCode: undefined, + }); + expect(result.current.reconnectState).toBeNull(); + }); + }); + + it("does not surface reconnectState for completed chats", async () => { + immediateAnimationFrame(); + + const chatID = "chat-disconnect-completed"; + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [], + chatRecord: { ...makeChat(chatID), status: "completed" }, + chatMessagesData: { + messages: [], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }); + return { + chatStatus: useChatSelector(store, selectChatStatus), + reconnectState: useChatSelector(store, selectReconnectState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.chatStatus).toBe("completed"); + expect(watchChat).toHaveBeenCalledWith(chatID, undefined); + }); + + act(() => { + mockSocket.emitError(); + }); + + await waitFor(() => { + expect(result.current.reconnectState).toBeNull(); + expect(result.current.chatStatus).toBe("completed"); }); }); it("uses exponential backoff on consecutive disconnects", async () => { @@ -1979,9 +2181,7 @@ describe("useChatStore", () => { const chatID = "chat-backoff"; const watchMock = vi.mocked(watchChat); - - // Return fresh sockets on each call. - watchMock.mockImplementation(() => createMockSocket() as never); + const sockets = mockWatchChatWithFreshSockets(watchMock); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -2011,7 +2211,7 @@ describe("useChatStore", () => { }); // Get the first socket and disconnect it. - const socket1 = watchMock.mock.results[0].value as MockSocket; + const socket1 = sockets[0]!; act(() => socket1.emitClose()); // First reconnect after 1s. @@ -2020,7 +2220,7 @@ describe("useChatStore", () => { }); // Second disconnect — reconnect after 2s. - const socket2 = watchMock.mock.results[1].value as MockSocket; + const socket2 = sockets[1]!; act(() => socket2.emitClose()); await waitFor(() => expect(watchMock).toHaveBeenCalledTimes(3), { @@ -2034,7 +2234,7 @@ describe("useChatStore", () => { const chatID = "chat-catchup"; const msg = makeMessage(chatID, 42, "assistant", "hello"); const watchMock = vi.mocked(watchChat); - watchMock.mockImplementation(() => createMockSocket() as never); + const sockets = mockWatchChatWithFreshSockets(watchMock); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -2065,7 +2265,7 @@ describe("useChatStore", () => { }); // Disconnect and reconnect. - const socket1 = watchMock.mock.results[0].value as MockSocket; + const socket1 = sockets[0]!; act(() => socket1.emitClose()); // Second connect should also use the last message ID. @@ -2097,7 +2297,7 @@ describe("useChatStore", () => { // Return a fresh MockSocket for each connection attempt // so we can control the first and second sockets // independently. - watchMock.mockImplementation(() => createMockSocket() as never); + const sockets = mockWatchChatWithFreshSockets(watchMock); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -2133,7 +2333,7 @@ describe("useChatStore", () => { expect(watchMock).toHaveBeenCalledWith(chatID, 1); }); - const socket1 = watchMock.mock.results[0].value as MockSocket; + const socket1 = sockets[0]!; // Simulate the first socket opening successfully. act(() => socket1.emitOpen()); @@ -2176,10 +2376,10 @@ describe("useChatStore", () => { // A second socket should now exist. expect(watchMock).toHaveBeenCalledTimes(2); - const socket2 = watchMock.mock.results[1].value as MockSocket; + const socket2 = sockets[1]!; // Simulate the reconnected socket opening. This is - // where onOpen fires clearStreamState(). + // where onOpen fires resetTransportReplayState(). act(() => socket2.emitOpen()); // Replay the same parts the server would send on the @@ -2203,8 +2403,9 @@ describe("useChatStore", () => { }); }); - // Without clearStreamState() in onOpen the replayed - // parts would append to the stale accumulator, producing + // Without resetTransportReplayState() in onOpen the + // replayed parts would append to the stale accumulator, + // producing // "Hello worldHello world". The fix ensures a clean // slate so we get the correct single copy. await waitFor(() => { @@ -2221,7 +2422,7 @@ describe("useChatStore", () => { const chatID = "chat-clear-error"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -2279,7 +2480,7 @@ describe("useChatStore", () => { const msg3 = makeMessage(chatID, 3, "user", "third"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper: FC = ({ children }) => ( @@ -2350,7 +2551,7 @@ describe("useChatStore", () => { const promotedMsg = makeMessage(chatID, 3, "user", "follow-up"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper: FC = ({ children }) => ( @@ -2459,7 +2660,7 @@ describe("useChatStore", () => { const msg2 = makeMessage(chatID, 2, "assistant", "hi"); const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = createTestQueryClient(); const wrapper: FC = ({ children }) => ( @@ -2563,7 +2764,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-sidebar-status"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2628,7 +2829,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-sidebar-message"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2701,7 +2902,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-sidebar-error"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2766,7 +2967,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-active"; const otherChatID = "chat-other"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2839,7 +3040,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-no-regress-msg"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2911,7 +3112,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-no-regress-status"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { @@ -2977,7 +3178,7 @@ describe("updateSidebarChat via stream events", () => { const chatID = "chat-no-regress-error"; const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + mockWatchChatReturn(mockSocket); const queryClient = new QueryClient({ defaultOptions: { diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts index 2e9e586bad..18bd3f7a86 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/ChatContext.ts @@ -1,14 +1,18 @@ import { watchChat } from "api/api"; import { chatMessagesKey, updateInfiniteChatsCache } from "api/queries/chats"; import type * as TypesGen from "api/typesGenerated"; +import { asNumber, asString } from "components/ai-elements/runtimeTypeUtils"; import { useEffect, useRef, useState, useSyncExternalStore } from "react"; import { type InfiniteData, useQueryClient } from "react-query"; import type { OneWayMessageEvent } from "utils/OneWayWebSocket"; import { createReconnectingWebSocket } from "utils/reconnectingWebSocket"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import { + type ChatDetailError, + chatDetailErrorsEqual, +} from "../../utils/usageLimitMessage"; import { applyMessagePartToStreamState } from "./streamState"; -import type { StreamState } from "./types"; +import type { ReconnectState, RetryState, StreamState } from "./types"; const isChatStreamEvent = (data: unknown): data is TypesGen.ChatStreamEvent => typeof data === "object" && @@ -31,6 +35,30 @@ const toChatStreamEvents = (data: unknown): TypesGen.ChatStreamEvent[] => { return []; }; +const normalizeChatDetailError = ( + error: TypesGen.ChatStreamError | Record | undefined, +): ChatDetailError => ({ + message: asString(error?.message).trim() || "Chat processing failed.", + kind: asString(error?.kind).trim() || "generic", + provider: asString(error?.provider).trim() || undefined, + retryable: + typeof error?.retryable === "boolean" ? error.retryable : undefined, + statusCode: asNumber(error?.status_code), +}); + +const normalizeRetryState = (retry: TypesGen.ChatStreamRetry): RetryState => { + const delayMs = asNumber(retry.delay_ms); + const retryingAt = asString(retry.retrying_at).trim() || undefined; + return { + attempt: Math.max(1, asNumber(retry.attempt) ?? 1), + error: asString(retry.error).trim() || "Retrying request shortly.", + kind: asString(retry.kind).trim() || "generic", + provider: asString(retry.provider).trim() || undefined, + ...(delayMs !== undefined ? { delayMs } : {}), + ...(retryingAt ? { retryingAt } : {}), + }; +}; + const byMessageCreatedAt = ( left: TypesGen.ChatMessage, right: TypesGen.ChatMessage, @@ -115,13 +143,60 @@ const chatQueuedMessagesEqualByID = ( return true; }; +const retryStatesEqual = ( + left: RetryState | null, + right: RetryState | null, +): boolean => { + if (left === right) { + return true; + } + if (!left || !right) { + return false; + } + return ( + left.attempt === right.attempt && + left.error === right.error && + left.kind === right.kind && + left.provider === right.provider && + left.delayMs === right.delayMs && + left.retryingAt === right.retryingAt + ); +}; + +const reconnectStatesEqual = ( + left: ReconnectState | null, + right: ReconnectState | null, +): boolean => { + if (left === right) { + return true; + } + if (!left || !right) { + return false; + } + return ( + left.attempt === right.attempt && + left.delayMs === right.delayMs && + left.retryingAt === right.retryingAt + ); +}; + +const isActiveChatStatus = (status: TypesGen.ChatStatus | null): boolean => + status === "running" || status === "pending"; + +const shouldSurfaceReconnectState = (state: ChatStoreState): boolean => + state.streamError === null && + (state.streamState !== null || + state.retryState !== null || + isActiveChatStatus(state.chatStatus)); + type ChatStoreState = { messagesByID: Map; orderedMessageIDs: readonly number[]; streamState: StreamState | null; chatStatus: TypesGen.ChatStatus | null; - streamError: string | null; - retryState: { attempt: number; error: string } | null; + streamError: ChatDetailError | null; + retryState: RetryState | null; + reconnectState: ReconnectState | null; queuedMessages: readonly TypesGen.ChatQueuedMessage[]; subagentStatusOverrides: Map; }; @@ -144,11 +219,14 @@ type ChatStore = { queuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined, ) => void; setChatStatus: (status: TypesGen.ChatStatus | null) => void; - setStreamError: (reason: string | null) => void; + setStreamError: (reason: ChatDetailError | null) => void; clearStreamError: () => void; - setRetryState: (state: { attempt: number; error: string } | null) => void; + setRetryState: (state: RetryState | null) => void; clearRetryState: () => void; + setReconnectState: (state: ReconnectState | null) => void; + clearReconnectState: () => void; clearStreamState: () => void; + resetTransportReplayState: () => void; setSubagentStatusOverride: ( chatID: string, status: TypesGen.ChatStatus, @@ -163,6 +241,7 @@ const createInitialState = (): ChatStoreState => ({ chatStatus: null, streamError: null, retryState: null, + reconnectState: null, queuedMessages: [], subagentStatusOverrides: new Map(), }); @@ -379,13 +458,15 @@ export const createChatStore = (): ChatStore => { })); }, setStreamError: (reason) => { - if (state.streamError === reason) { - return; - } - setState((current) => ({ - ...current, - streamError: reason, - })); + setState((current) => { + if (chatDetailErrorsEqual(current.streamError, reason)) { + return current; + } + return { + ...current, + streamError: reason, + }; + }); }, clearStreamError: () => { if (state.streamError === null) { @@ -397,13 +478,15 @@ export const createChatStore = (): ChatStore => { })); }, setRetryState: (retryState) => { - if (state.retryState === retryState) { - return; - } - setState((current) => ({ - ...current, - retryState, - })); + setState((current) => { + if (retryStatesEqual(current.retryState, retryState)) { + return current; + } + return { + ...current, + retryState, + }; + }); }, clearRetryState: () => { if (state.retryState === null) { @@ -414,6 +497,26 @@ export const createChatStore = (): ChatStore => { retryState: null, })); }, + setReconnectState: (reconnectState) => { + setState((current) => { + if (reconnectStatesEqual(current.reconnectState, reconnectState)) { + return current; + } + return { + ...current, + reconnectState, + }; + }); + }, + clearReconnectState: () => { + if (state.reconnectState === null) { + return; + } + setState((current) => ({ + ...current, + reconnectState: null, + })); + }, clearStreamState: () => { if (state.streamState === null) { return; @@ -423,6 +526,21 @@ export const createChatStore = (): ChatStore => { streamState: null, })); }, + resetTransportReplayState: () => { + if ( + state.reconnectState === null && + state.streamState === null && + state.streamError === null + ) { + return; + } + setState((current) => ({ + ...current, + reconnectState: null, + streamState: null, + streamError: null, + })); + }, setSubagentStatusOverride: (chatID, status) => { if (state.subagentStatusOverrides.get(chatID) === status) { return; @@ -441,6 +559,7 @@ export const createChatStore = (): ChatStore => { state.streamState === null && state.streamError === null && state.retryState === null && + state.reconnectState === null && state.subagentStatusOverrides.size === 0 ) { return; @@ -450,6 +569,7 @@ export const createChatStore = (): ChatStore => { streamState: null, streamError: null, retryState: null, + reconnectState: null, subagentStatusOverrides: new Map(), })); }, @@ -479,6 +599,39 @@ export const selectQueuedMessages = (state: ChatStoreState) => export const selectSubagentStatusOverrides = (state: ChatStoreState) => state.subagentStatusOverrides; export const selectRetryState = (state: ChatStoreState) => state.retryState; +export const selectReconnectState = (state: ChatStoreState) => + state.reconnectState; + +const selectLatestDurableMessage = ( + state: ChatStoreState, +): TypesGen.ChatMessage | undefined => { + const latestMessageID = + state.orderedMessageIDs[state.orderedMessageIDs.length - 1]; + return latestMessageID === undefined + ? undefined + : state.messagesByID.get(latestMessageID); +}; + +export const selectIsAwaitingFirstStreamChunk = ( + state: ChatStoreState, +): boolean => { + const latestMessage = selectLatestDurableMessage(state); + const latestMessageNeedsAssistantResponse = + !latestMessage || latestMessage.role !== "assistant"; + return ( + state.streamState === null && + isActiveChatStatus(state.chatStatus) && + latestMessageNeedsAssistantResponse + ); +}; + +export const useChatSelector = ( + store: ChatStore, + selector: (state: ChatStoreState) => T, +): T => { + const getSnapshot = () => selector(store.getSnapshot()); + return useSyncExternalStore(store.subscribe, getSnapshot, getSnapshot); +}; export const useChatStore = ( options: UseChatStoreOptions, @@ -511,7 +664,7 @@ export const useChatStore = ( // after a refetch producing new objects) vs. just getting a new // array reference because an unrelated field like queued_messages // was updated in the query cache. Element-level reference - // comparison works because useMemo(flatMap) preserves message + // comparison works because the flattening step preserves message // object references when only non-message fields change in the // page, while a genuine refetch returns new objects from the // server. @@ -761,7 +914,10 @@ export const useChatStore = ( return; } if (payload.parseError || !payload.parsedMessage) { - store.setStreamError("Failed to parse chat stream update."); + store.setStreamError({ + kind: "generic", + message: "Failed to parse chat stream update.", + }); return; } if (payload.parsedMessage.type !== "data") { @@ -833,6 +989,7 @@ export const useChatStore = ( if (!nextStatus) { continue; } + if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { store.setSubagentStatusOverride( streamEvent.chat_id, @@ -840,6 +997,7 @@ export const useChatStore = ( ); continue; } + store.setChatStatus(nextStatus); if (nextStatus === "pending" || nextStatus === "waiting") { store.clearStreamState(); @@ -862,15 +1020,11 @@ export const useChatStore = ( if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { continue; } - const reason = - streamEvent.error?.message.trim() || "Chat processing failed."; + const reason = normalizeChatDetailError(streamEvent.error); store.setChatStatus("error"); store.setStreamError(reason); store.clearRetryState(); - setChatErrorReasonRef.current(chatID, { - kind: "generic", - message: reason, - }); + setChatErrorReasonRef.current(chatID, reason); updateSidebarChat((chat) => chat.status === "error" ? chat : { ...chat, status: "error" }, ); @@ -883,10 +1037,7 @@ export const useChatStore = ( const retry = streamEvent.retry; if (retry) { store.clearStreamState(); - store.setRetryState({ - attempt: retry.attempt, - error: retry.error, - }); + store.setRetryState(normalizeRetryState(retry)); } continue; } @@ -919,26 +1070,20 @@ export const useChatStore = ( return socket; }, onOpen() { - // Connection succeeded — clear any previous disconnect - // error and stale stream state. Clearing stream state - // is critical for reconnections: the server replays - // all buffered message_part events, so we must start - // from a clean slate to avoid duplicating text. - store.clearStreamError(); - store.clearStreamState(); + // Connection succeeded. Before the socket replays any + // buffered message_part events, drop transport-scoped + // state from the previous socket attempt so stale + // partial output or failures do not leak into the new + // stream. + store.resetTransportReplayState(); }, - onDisconnect(attempt) { - // Show the error only on the first disconnect (not - // while we are already retrying). - if (attempt === 0) { - store.setStreamError("Chat stream disconnected. Reconnecting\u2026"); - } - // Clear "running" status on disconnect so the UI - // doesn't show a stale spinner. The reconnected - // stream will deliver the authoritative status. - const currentStatus = store.getSnapshot().chatStatus; - if (currentStatus === "running") { - store.setChatStatus(null); + onDisconnect(reconnectState) { + // Only surface reconnecting when the disconnect + // interrupted active response work. Idle watcher + // reconnects stay silent. + const snapshot = store.getSnapshot(); + if (shouldSurfaceReconnectState(snapshot)) { + store.setReconnectState(reconnectState); } }, }); @@ -960,11 +1105,3 @@ export const useChatStore = ( }, }; }; - -export const useChatSelector = ( - store: ChatStore, - selector: (state: ChatStoreState) => T, -): T => { - const getSnapshot = () => selector(store.getSnapshot()); - return useSyncExternalStore(store.subscribe, getSnapshot, getSnapshot); -}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ChatStatusCallout.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ChatStatusCallout.tsx new file mode 100644 index 0000000000..bc7596533f --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/ChatStatusCallout.tsx @@ -0,0 +1,233 @@ +import { Alert, AlertDescription, AlertTitle } from "components/Alert/Alert"; +import { Response, Shimmer } from "components/ai-elements"; +import { Button } from "components/Button/Button"; +import { Pill } from "components/Pill/Pill"; +import { ExternalLinkIcon } from "lucide-react"; +import { type FC, useEffect, useState } from "react"; +import { getProviderStatusURL } from "./chatStatusHelpers"; +import type { LiveStatusModel } from "./liveStatusModel"; + +const RESPONSE_STARTUP_GRACE_MS = 15_000; +const DELAYED_STARTUP_TEXT = "Response startup is taking longer than expected"; +const THINKING_TEXT = "Thinking..."; + +type RetryOrFailedStatus = Extract< + LiveStatusModel, + { phase: "retrying" } | { phase: "failed" } +>; +type ReconnectingStatus = Extract; + +const StatusPlaceholder: FC<{ + text: string; + shimmer?: boolean; +}> = ({ text, shimmer = false }) => { + return ( +
+ {/* Reserve the final response height without exposing a selectable copy. */} + + {text} + +
+ {shimmer ? ( + + {text} + + ) : ( + + {text} + + )} +
+
+ ); +}; + +const StartingPlaceholder: FC = () => { + const [isDelayed, setIsDelayed] = useState(false); + + useEffect(() => { + const timeout = window.setTimeout(() => { + setIsDelayed(true); + }, RESPONSE_STARTUP_GRACE_MS); + return () => window.clearTimeout(timeout); + }, []); + + return ( + + ); +}; + +/** + * Syncs with the system clock to produce a live countdown from an + * ISO-8601 deadline. Polls at 100ms so the displayed second flips + * within 100ms of the real transition. Returns 0 when no deadline is + * provided or the deadline has passed. + */ +const useDeadlineCountdown = (deadline: string | undefined): number => { + const [secondsLeft, setSecondsLeft] = useState(0); + + useEffect(() => { + if (!deadline) { + setSecondsLeft(0); + return; + } + + const targetMs = new Date(deadline).getTime(); + if (!Number.isFinite(targetMs)) { + setSecondsLeft(0); + return; + } + + const update = () => { + const remaining = Math.max(0, targetMs - Date.now()); + setSecondsLeft(Math.ceil(remaining / 1000)); + }; + + update(); + const interval = setInterval(update, 100); + return () => clearInterval(interval); + }, [deadline]); + + return secondsLeft; +}; + +/** + * Leaf component that owns the countdown interval so ticking seconds only + * re-render this span, not the parent Alert (which contains a Radix Slot + * that infinite-loops on rapid re-renders). + */ +const StatusCountdown: FC<{ + deadline: string; + label: string; +}> = ({ deadline, label }) => { + const seconds = useDeadlineCountdown(deadline); + if (seconds <= 0) { + return null; + } + return ( + + {label} {seconds}s + + ); +}; + +const StatusAlert: FC<{ status: RetryOrFailedStatus }> = ({ status }) => { + const statusURL = getProviderStatusURL(status.kind, status.provider); + const pillType = + status.phase === "failed" + ? "error" + : status.kind === "generic" + ? "inactive" + : "warning"; + const severity = + status.phase === "failed" + ? "error" + : status.kind === "generic" + ? "info" + : "warning"; + const hasMetadata = + status.phase === "retrying" || + status.provider !== undefined || + (status.phase === "failed" && status.statusCode !== undefined) || + (status.phase === "failed" && status.retryable !== undefined); + + return ( + + + Status + + + + ) + } + > +
+
+ {status.title} + + {status.kind} + +
+ {status.message} + {hasMetadata && ( +
+ {status.phase === "retrying" && status.retryingAt && ( + + )} + {status.phase === "retrying" && ( + Attempt {status.attempt} + )} + {status.provider && Provider {status.provider}} + {status.phase === "failed" && status.statusCode !== undefined && ( + HTTP {status.statusCode} + )} + {status.phase === "failed" && status.retryable !== undefined && ( + {status.retryable ? "Retryable" : "Not retryable"} + )} +
+ )} +
+ + ); +}; + +const ReconnectingAlert: FC<{ status: ReconnectingStatus }> = ({ status }) => { + return ( + +
+ {status.title} + {status.message} +
+ + Attempt {status.attempt} +
+
+ + ); +}; + +export const ChatStatusCallout: FC<{ + status: LiveStatusModel; + startingResetKey?: string; +}> = ({ status, startingResetKey }) => { + switch (status.phase) { + case "idle": + case "streaming": + return null; + case "starting": + return ; + case "retrying": + return ( + <> + + + + ); + case "reconnecting": + return ( + <> + + + + ); + case "failed": + return ; + } +}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.stories.tsx index 9986a63714..9cd8424657 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.stories.tsx @@ -62,15 +62,7 @@ const mockTextAttachmentFetch = () => { const defaultArgs: Omit< React.ComponentProps, "parsedMessages" -> = { - isEmpty: false, - hasStreamOutput: false, - streamState: null, - streamTools: [], - subagentTitles: new Map(), - subagentStatusOverrides: new Map(), - isAwaitingFirstStreamChunk: false, -}; +> = {}; const meta: Meta = { title: "pages/AgentsPage/AgentDetail/ConversationTimeline", @@ -424,47 +416,6 @@ export const UserMessageWithImagesAndFileRefs: Story = { }, }; -/** Usage-limit errors render as an info alert with analytics access. */ -export const UsageLimitExceeded: Story = { - args: { - ...defaultArgs, - parsedMessages: [], - detailError: { - kind: "usage-limit", - message: - "You've used $50.00 of your $50.00 spend limit. Your limit resets on July 1, 2025.", - }, - - subagentTitles: new Map(), - subagentStatusOverrides: new Map(), - }, - play: async ({ canvasElement }) => { - const canvas = within(canvasElement); - expect(canvas.getByText(/spend limit/i)).toBeVisible(); - const link = canvas.getByRole("link", { name: /view usage/i }); - expect(link).toBeVisible(); - expect(link).toHaveAttribute("href", "/agents/analytics"); - }, -}; - -/** Non-usage errors must not show the usage CTA. */ -export const GenericErrorDoesNotShowUsageAction: Story = { - args: { - ...defaultArgs, - parsedMessages: [], - detailError: { kind: "generic", message: "Provider request failed." }, - subagentTitles: new Map(), - subagentStatusOverrides: new Map(), - }, - play: async ({ canvasElement }) => { - const canvas = within(canvasElement); - expect(canvas.getByText(/provider request failed/i)).toBeVisible(); - expect( - canvas.queryByRole("link", { name: /view usage/i }), - ).not.toBeInTheDocument(); - }, -}; - /** File references render inline with text, matching the chat input style. */ export const UserMessageWithInlineFileRef: Story = { args: { diff --git a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx index 3d470ff083..3574ec4285 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/ConversationTimeline.tsx @@ -3,16 +3,15 @@ import { FileTextIcon, PencilIcon } from "lucide-react"; import { type FC, Fragment, + memo, type ReactNode, useEffect, useLayoutEffect, useRef, useState, } from "react"; -import { Link } from "react-router"; import type { UrlTransform } from "streamdown"; import { cn } from "utils/cn"; -import { Alert } from "#/components/Alert/Alert"; import { ConversationItem, Message, @@ -22,7 +21,6 @@ import { Tool, } from "#/components/ai-elements"; import { WebSearchSources } from "#/components/ai-elements/tool"; -import { Button } from "#/components/Button/Button"; import { FileReferenceChip } from "#/components/ChatMessageInput/FileReferenceNode"; import { Spinner } from "#/components/Spinner/Spinner"; import { @@ -35,10 +33,11 @@ import { fetchTextAttachmentContent, formatTextAttachmentPreview, } from "../../utils/fetchTextAttachment"; -import type { ChatDetailError } from "../../utils/usageLimitMessage"; import { ImageThumbnail } from "../AgentChatInput"; import { ImageLightbox } from "../ImageLightbox"; import { TextPreviewDialog } from "../TextPreviewDialog"; +import { ChatStatusCallout } from "./ChatStatusCallout"; +import type { LiveStatusModel } from "./liveStatusModel"; import { useSmoothStreamingText } from "./SmoothText"; import type { MergedTool, @@ -382,7 +381,7 @@ function renderBlockList({ return { elements, renderedToolIDs }; } -interface ChatMessageItemProps { +const ChatMessageItem = memo<{ message: TypesGen.ChatMessage; parsed: ParsedMessageContent; onEditUserMessage?: ( @@ -399,212 +398,313 @@ interface ChatMessageItemProps { fadeFromBottom?: boolean; urlTransform?: UrlTransform; mcpServers?: readonly TypesGen.MCPServerConfig[]; -} +}>( + ({ + message, + parsed, + onEditUserMessage, + editingMessageId, + savingMessageId, + isAfterEditingMessage = false, + fadeFromBottom = false, + urlTransform, + mcpServers, + }) => { + const isUser = message.role === "user"; + const isSavingMessage = savingMessageId === message.id; + const [previewImage, setPreviewImage] = useState(null); + const [previewText, setPreviewText] = useState(null); + const toolByID = new Map(parsed.tools.map((tool) => [tool.id, tool])); -const ChatMessageItem: FC = ({ - message, - parsed, - onEditUserMessage, - editingMessageId, - savingMessageId, - isAfterEditingMessage = false, - fadeFromBottom = false, + if ( + parsed.toolResults.length > 0 && + parsed.toolCalls.length === 0 && + parsed.markdown === "" && + parsed.reasoning === "" + ) { + return null; + } + + // Hide messages that consist entirely of provider-executed + // tool results. The parser skips these parts, so the parsed + // output is empty and would show a "no renderable content" + // fallback. + const parts = message.content ?? []; + if ( + parts.length > 0 && + parts.every((p) => p.type === "tool-result" && p.provider_executed) + ) { + return null; + } + + const hasRenderableContent = + parsed.blocks.length > 0 || + parsed.tools.length > 0 || + parsed.sources.length > 0; + // Pre-compute the inline content for user messages so we + // avoid a filter + map inside the JSX return path. + const userInlineContent = isUser + ? parsed.blocks.filter( + ( + b, + ): b is + | Extract + | Extract => + b.type === "response" || b.type === "file-reference", + ) + : []; + const userFileBlocks = isUser + ? parsed.blocks.filter( + (b): b is Extract => b.type === "file", + ) + : []; + const hasUserMessageBody = + userInlineContent.length > 0 || Boolean(parsed.markdown?.trim()); + const hasFileBlocks = userFileBlocks.length > 0; + + const conversationItemProps: { role: "user" | "assistant" } = { + role: isUser ? "user" : "assistant", + }; + const { elements: orderedBlocks, renderedToolIDs } = renderBlockList({ + blocks: parsed.blocks, + toolByID, + keyPrefix: String(message.id), + onImageClick: setPreviewImage, + onTextFileClick: (content) => setPreviewText(content), + urlTransform, + mcpServers, + }); + const remainingTools = parsed.tools.filter( + (tool) => !renderedToolIDs.has(tool.id), + ); + + return ( +
+ + {isUser ? ( + + +
+ {(hasUserMessageBody || hasFileBlocks) && ( +
+ {hasUserMessageBody && ( + + {userInlineContent.length > 0 + ? userInlineContent.map((block, i) => + block.type === "response" ? ( + {block.text} + ) : ( + + ), + ) + : parsed.markdown || ""} + + )} + {isSavingMessage && ( + + )} + {onEditUserMessage && !isSavingMessage && ( + + + + + + Edit message + + + )} +
+ )} + {(() => { + if (userFileBlocks.length === 0) return null; + return ( +
+ {userFileBlocks.map((block, i) => + renderFileBlock({ + block, + key: `user-file-${block.file_id ?? i}`, + onImageClick: setPreviewImage, + onTextFileClick: setPreviewText, + }), + )} +
+ ); + })()} + {fadeFromBottom && ( +
+ )} +
+ + + ) : ( + + +
+ {orderedBlocks} + {remainingTools.map((tool) => ( + + ))} + {!hasRenderableContent && ( +
+ Message has no renderable content. +
+ )} +
+ + + )} + + {previewImage && ( + setPreviewImage(null)} + /> + )} + {previewText !== null && ( + setPreviewText(null)} + /> + )} +
+ ); + }, +); + +const hasTransientLiveStatus = (liveStatus: LiveStatusModel): boolean => + liveStatus.phase === "starting" || + liveStatus.phase === "retrying" || + liveStatus.phase === "reconnecting"; + +export const StreamingOutput: FC<{ + streamState: StreamState | null; + streamTools: readonly MergedTool[]; + subagentTitles?: Map; + subagentStatusOverrides?: Map; + liveStatus: LiveStatusModel; + startingResetKey?: string; + urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; +}> = ({ + streamState, + streamTools, + subagentTitles, + subagentStatusOverrides, + liveStatus, + startingResetKey, urlTransform, mcpServers, }) => { - const isUser = message.role === "user"; - const isSavingMessage = savingMessageId === message.id; - const [previewImage, setPreviewImage] = useState(null); - const [previewText, setPreviewText] = useState(null); - const toolByID = new Map(parsed.tools.map((tool) => [tool.id, tool])); - - if ( - parsed.toolResults.length > 0 && - parsed.toolCalls.length === 0 && - parsed.markdown === "" && - parsed.reasoning === "" - ) { + if (liveStatus.phase === "idle") { return null; } - // Hide messages that consist entirely of provider-executed - // tool results. The parser skips these parts, so the parsed - // output is empty and would show a "no renderable content" - // fallback. - const parts = message.content ?? []; - if ( - parts.length > 0 && - parts.every((p) => p.type === "tool-result" && p.provider_executed) - ) { + const isStreaming = liveStatus.phase === "streaming"; + const shouldShowBlocks = + liveStatus.phase === "streaming" || liveStatus.hasAccumulatedOutput; + const shouldShowStatusCallout = hasTransientLiveStatus(liveStatus); + if (!shouldShowBlocks && !shouldShowStatusCallout) { return null; } - const hasRenderableContent = - parsed.blocks.length > 0 || - parsed.tools.length > 0 || - parsed.sources.length > 0; - // Pre-compute the inline content for user messages so we - // avoid a filter + map inside the JSX return path. - const userInlineContent = isUser - ? parsed.blocks.filter( - ( - b, - ): b is - | Extract - | Extract => - b.type === "response" || b.type === "file-reference", - ) - : []; - - const userFileBlocks = isUser - ? parsed.blocks.filter( - (b): b is Extract => b.type === "file", - ) - : []; - - const hasUserMessageBody = - userInlineContent.length > 0 || Boolean(parsed.markdown?.trim()); - const hasFileBlocks = userFileBlocks.length > 0; - - const conversationItemProps: { role: "user" | "assistant" } = { - role: isUser ? "user" : "assistant", - }; + const conversationItemProps = { role: "assistant" as const }; + const toolByID = new Map(streamTools.map((tool) => [tool.id, tool])); + const blocks = shouldShowBlocks ? (streamState?.blocks ?? []) : []; const { elements: orderedBlocks, renderedToolIDs } = renderBlockList({ - blocks: parsed.blocks, + blocks, toolByID, - keyPrefix: String(message.id), - onImageClick: setPreviewImage, - onTextFileClick: (content) => setPreviewText(content), + keyPrefix: "stream", + isStreaming, + subagentTitles, + subagentStatusOverrides, urlTransform, mcpServers, }); - const remainingTools = parsed.tools.filter( - (tool) => !renderedToolIDs.has(tool.id), - ); + const remainingTools = shouldShowBlocks + ? streamTools.filter((tool) => !renderedToolIDs.has(tool.id)) + : []; return ( -
- - {isUser ? ( - - -
- {(hasUserMessageBody || hasFileBlocks) && ( -
- {hasUserMessageBody && ( - - {userInlineContent.length > 0 - ? userInlineContent.map((block, i) => - block.type === "response" ? ( - {block.text} - ) : ( - - ), - ) - : parsed.markdown || ""} - - )} - {isSavingMessage && ( - - )} - {onEditUserMessage && !isSavingMessage && ( - - - - - Edit message - - )} -
- )} - {(() => { - if (userFileBlocks.length === 0) return null; - return ( -
- {userFileBlocks.map((block, i) => - renderFileBlock({ - block, - key: `user-file-${block.file_id ?? i}`, - onImageClick: setPreviewImage, - onTextFileClick: setPreviewText, - }), - )} -
- ); - })()} - {fadeFromBottom && ( -
- )} -
- - - ) : ( - - -
+ + + +
+ {shouldShowBlocks && ( + <> {orderedBlocks} {remainingTools.map((tool) => ( = ({ result={tool.result} status={tool.status} isError={tool.isError} + subagentTitles={isStreaming ? subagentTitles : undefined} + subagentStatusOverrides={ + isStreaming ? subagentStatusOverrides : undefined + } mcpServerConfigId={tool.mcpServerConfigId} mcpServers={mcpServers} /> ))} - {!hasRenderableContent && ( -
- Message has no renderable content. -
- )} -
- - - )} - - {previewImage && ( - setPreviewImage(null)} - /> - )} - {previewText !== null && ( - setPreviewText(null)} - /> - )} -
- ); -}; - -export const StreamingOutput: FC<{ - streamState: StreamState | null; - streamTools: readonly MergedTool[]; - subagentTitles?: Map; - subagentStatusOverrides?: Map; - showInitialPlaceholder?: boolean; - retryState?: { attempt: number; error: string } | null; - urlTransform?: UrlTransform; - mcpServers?: readonly TypesGen.MCPServerConfig[]; -}> = ({ - streamState, - streamTools, - subagentTitles, - subagentStatusOverrides, - showInitialPlaceholder = false, - retryState, - urlTransform, - mcpServers, -}) => { - const conversationItemProps = { role: "assistant" as const }; - const toolByID = new Map(streamTools.map((tool) => [tool.id, tool])); - const blocks = streamState?.blocks ?? []; - const { elements: orderedBlocks, renderedToolIDs } = renderBlockList({ - blocks, - toolByID, - keyPrefix: "stream", - isStreaming: true, - subagentTitles, - subagentStatusOverrides, - urlTransform, - mcpServers, - }); - const remainingTools = streamTools.filter( - (tool) => !renderedToolIDs.has(tool.id), - ); - - return ( - - - -
- {orderedBlocks} - {showInitialPlaceholder || - (streamState && - orderedBlocks.length === 0 && - streamTools.length === 0) ? ( -
- - {`Thinking...${retryState ? ` attempt ${retryState.attempt}` : ""}`} - -
- - Thinking... - - {retryState && ( - - attempt {retryState.attempt} - - )} -
-
- ) : null} - {remainingTools.map((tool) => ( - + )} + {shouldShowStatusCallout && ( + - ))} + )}
@@ -1013,16 +1023,7 @@ const StickyUserMessage: FC<{ }; interface ConversationTimelineProps { - isEmpty: boolean; parsedMessages: readonly ParsedMessageEntry[]; - hasStreamOutput: boolean; - streamState: StreamState | null; - streamTools: readonly MergedTool[]; - subagentTitles: Map; - subagentStatusOverrides: Map; - retryState?: { attempt: number; error: string } | null; - isAwaitingFirstStreamChunk: boolean; - detailError?: ChatDetailError | null; onEditUserMessage?: ( messageId: number, text: string, @@ -1035,25 +1036,16 @@ interface ConversationTimelineProps { } export const ConversationTimeline: FC = ({ - isEmpty, parsedMessages, - hasStreamOutput, - streamState, - streamTools, - subagentTitles, - subagentStatusOverrides, - retryState, - isAwaitingFirstStreamChunk, - detailError, onEditUserMessage, editingMessageId, savingMessageId, urlTransform, mcpServers, }) => { - const shouldRenderStreamAfterMessages = - hasStreamOutput && parsedMessages.length > 0; - const isUsageLimitError = detailError?.kind === "usage-limit"; + if (parsedMessages.length === 0) { + return null; + } // Build a set of message IDs that appear after the message // currently being edited so they can be visually faded. @@ -1072,76 +1064,29 @@ export const ConversationTimeline: FC = ({ } return ( -
- {isEmpty && !hasStreamOutput ? ( -
-

Start a conversation with your agent.

-
- ) : ( -
- {parsedMessages.map(({ message, parsed }) => - message.role === "user" ? ( - - ) : ( - - ), - )} - {shouldRenderStreamAfterMessages && ( - - )} - {hasStreamOutput && parsedMessages.length === 0 && ( - - )} -
- )} - {detailError && ( - - View Usage - - ) - } - > - {detailError.message} - +
+ {parsedMessages.map(({ message, parsed }) => + message.role === "user" ? ( + + ) : ( + + ), )}
); diff --git a/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.stories.tsx b/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.stories.tsx new file mode 100644 index 0000000000..09d9efd4d8 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.stories.tsx @@ -0,0 +1,207 @@ +import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, waitFor, within } from "storybook/test"; +import { LiveStreamTailContent } from "./LiveStreamTail"; +import { + buildLiveStatus, + buildReconnectState, + buildStreamRenderState, + FIXTURE_NOW, + textResponseStreamParts, +} from "./storyFixtures"; + +const retryThenResumedStream = buildStreamRenderState(textResponseStreamParts); + +const defaultArgs: React.ComponentProps = { + isTranscriptEmpty: true, + streamState: null, + streamTools: [], + liveStatus: buildLiveStatus(), + subagentTitles: new Map(), + subagentStatusOverrides: new Map(), +}; + +const meta: Meta = { + title: "pages/AgentsPage/AgentDetail/LiveStreamTail", + component: LiveStreamTailContent, + decorators: [ + (Story) => ( +
+ +
+ ), + ], + beforeEach: () => { + const real = Date.now; + Date.now = () => FIXTURE_NOW; + return () => { + Date.now = real; + }; + }, +}; +export default meta; +type Story = StoryObj; + +/** Empty transcripts show the standard prompt when there is no live tail. */ +export const EmptyConversationPrompt: Story = { + args: defaultArgs, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByText(/start a conversation with your agent/i), + ).toBeVisible(); + }, +}; + +/** Usage-limit failures replace the idle prompt with the analytics CTA. */ +export const UsageLimitExceeded: Story = { + args: { + ...defaultArgs, + liveStatus: buildLiveStatus({ + persistedError: { + kind: "usage-limit", + message: + "You've used $50.00 of your $50.00 spend limit. Your limit resets on July 1, 2025.", + }, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText(/spend limit/i)).toBeVisible(); + const link = canvas.getByRole("link", { name: /view usage/i }); + expect(link).toBeVisible(); + expect(link).toHaveAttribute("href", "/agents/analytics"); + }, +}; + +/** Provider failures keep the footer-level terminal callout and status link. */ +export const TerminalOverloadedError: Story = { + args: { + ...defaultArgs, + liveStatus: buildLiveStatus({ + persistedError: { + kind: "overloaded", + message: "Anthropic is currently overloaded. Please try again shortly.", + provider: "anthropic", + retryable: true, + statusCode: 529, + }, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /service overloaded/i }), + ).toBeVisible(); + expect(canvas.getByText("overloaded")).toBeVisible(); + expect(canvas.getByText(/http 529/i)).toBeVisible(); + expect(canvas.getByRole("link", { name: /status/i })).toBeVisible(); + }, +}; + +/** Generic failures do not show usage or provider CTAs. */ +export const GenericErrorDoesNotShowUsageAction: Story = { + args: { + ...defaultArgs, + liveStatus: buildLiveStatus({ + persistedError: { + kind: "generic", + message: "Provider request failed.", + }, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /request failed/i }), + ).toBeVisible(); + expect(canvas.getByText(/provider request failed/i)).toBeVisible(); + expect( + canvas.queryByText(/start a conversation with your agent/i), + ).not.toBeInTheDocument(); + expect( + canvas.queryByRole("link", { name: /view usage/i }), + ).not.toBeInTheDocument(); + expect( + canvas.queryByRole("link", { name: /status/i }), + ).not.toBeInTheDocument(); + }, +}; + +/** Reconnecting keeps already-streamed content visible without a terminal footer. */ +export const ReconnectingKeepsPartialOutputVisible: Story = { + args: { + ...defaultArgs, + isTranscriptEmpty: false, + streamState: retryThenResumedStream.streamState, + streamTools: retryThenResumedStream.streamTools, + liveStatus: buildLiveStatus({ + streamState: retryThenResumedStream.streamState, + reconnectState: buildReconnectState({ + attempt: 2, + delayMs: 2000, + }), + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText(/storybook streamed answer/i)).toBeVisible(); + expect( + canvas.getByRole("heading", { name: /reconnecting/i }), + ).toBeVisible(); + expect(canvas.getByText(/chat stream disconnected/i)).toBeVisible(); + expect( + canvas.queryByRole("heading", { name: /request failed/i }), + ).not.toBeInTheDocument(); + }, +}; + +/** Persisted errors yield to live streaming while the live tail is active. */ +export const PersistedGenericErrorDoesNotOverrideStreaming: Story = { + args: { + ...defaultArgs, + isTranscriptEmpty: false, + streamState: retryThenResumedStream.streamState, + streamTools: retryThenResumedStream.streamTools, + liveStatus: buildLiveStatus({ + streamState: retryThenResumedStream.streamState, + persistedError: { + kind: "generic", + message: "Stale persisted error.", + }, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect(canvas.getByText(/storybook streamed answer/i)).toBeVisible(); + }); + expect( + canvas.queryByRole("heading", { name: /request failed/i }), + ).not.toBeInTheDocument(); + }, +}; + +/** Terminal failures keep partial output visible above the footer callout. */ +export const FailedStreamKeepsPartialOutputVisible: Story = { + args: { + ...defaultArgs, + isTranscriptEmpty: false, + streamState: retryThenResumedStream.streamState, + streamTools: retryThenResumedStream.streamTools, + liveStatus: buildLiveStatus({ + streamState: retryThenResumedStream.streamState, + streamError: { + kind: "generic", + message: "Provider request failed.", + }, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText(/storybook streamed answer/i)).toBeVisible(); + expect( + canvas.getByRole("heading", { name: /request failed/i }), + ).toBeVisible(); + expect(canvas.getByText(/provider request failed/i)).toBeVisible(); + }, +}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.tsx b/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.tsx new file mode 100644 index 0000000000..c5cc1f5512 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/LiveStreamTail.tsx @@ -0,0 +1,162 @@ +import type * as TypesGen from "api/typesGenerated"; +import { Link } from "react-router"; +import type { UrlTransform } from "streamdown"; +import { Alert } from "#/components/Alert/Alert"; +import { Button } from "#/components/Button/Button"; +import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import { + selectIsAwaitingFirstStreamChunk, + selectReconnectState, + selectRetryState, + selectStreamError, + selectStreamState, + selectSubagentStatusOverrides, + useChatSelector, + type useChatStore, +} from "./ChatContext"; +import { ChatStatusCallout } from "./ChatStatusCallout"; +import { StreamingOutput } from "./ConversationTimeline"; +import { deriveLiveStatus, type LiveStatusModel } from "./liveStatusModel"; +import { buildStreamTools } from "./streamState"; +import type { MergedTool, StreamState } from "./types"; + +const shouldRenderStreamingSection = (liveStatus: LiveStatusModel): boolean => + liveStatus.phase === "streaming" || + liveStatus.phase === "starting" || + liveStatus.phase === "retrying" || + liveStatus.phase === "reconnecting" || + liveStatus.hasAccumulatedOutput; + +type ChatStoreHandle = ReturnType["store"]; + +interface LiveStreamTailContentProps { + isTranscriptEmpty: boolean; + streamState: StreamState | null; + streamTools: readonly MergedTool[]; + liveStatus: LiveStatusModel; + startingResetKey?: string; + subagentTitles: Map; + subagentStatusOverrides: Map; + urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; +} + +export const LiveStreamTailContent = ({ + isTranscriptEmpty, + streamState, + streamTools, + liveStatus, + startingResetKey, + subagentTitles, + subagentStatusOverrides, + urlTransform, + mcpServers, +}: LiveStreamTailContentProps) => { + const shouldRenderStreamSection = shouldRenderStreamingSection(liveStatus); + const terminalStatus = liveStatus.phase === "failed" ? liveStatus : null; + const usageLimitStatus = + terminalStatus?.kind === "usage-limit" ? terminalStatus : null; + const shouldRenderEmptyState = + isTranscriptEmpty && liveStatus.phase === "idle"; + + if ( + !shouldRenderEmptyState && + !shouldRenderStreamSection && + !terminalStatus + ) { + return null; + } + + return ( +
+ {shouldRenderEmptyState && ( +
+

Start a conversation with your agent.

+
+ )} + {shouldRenderStreamSection && ( + + )} + {usageLimitStatus ? ( + + View Usage + + } + > + {usageLimitStatus.message} + + ) : terminalStatus ? ( + + ) : null} +
+ ); +}; + +interface LiveStreamTailProps { + store: ChatStoreHandle; + persistedError: ChatDetailError | undefined; + isTranscriptEmpty: boolean; + startingResetKey?: string; + subagentTitles: Map; + urlTransform?: UrlTransform; + mcpServers?: readonly TypesGen.MCPServerConfig[]; +} + +export const LiveStreamTail = ({ + store, + persistedError, + isTranscriptEmpty, + startingResetKey, + subagentTitles, + urlTransform, + mcpServers, +}: LiveStreamTailProps) => { + const streamState = useChatSelector(store, selectStreamState); + const streamError = useChatSelector(store, selectStreamError); + const retryState = useChatSelector(store, selectRetryState); + const reconnectState = useChatSelector(store, selectReconnectState); + const isAwaitingFirstStreamChunk = useChatSelector( + store, + selectIsAwaitingFirstStreamChunk, + ); + const subagentStatusOverrides = useChatSelector( + store, + selectSubagentStatusOverrides, + ); + const streamTools = buildStreamTools(streamState); + const liveStatus = deriveLiveStatus({ + streamState, + retryState, + reconnectState, + streamError, + persistedError: persistedError ?? null, + isAwaitingFirstStreamChunk, + }); + + return ( + + ); +}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/StreamingOutput.stories.tsx b/site/src/pages/AgentsPage/components/AgentDetail/StreamingOutput.stories.tsx index b0de061f18..0115048d75 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/StreamingOutput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetail/StreamingOutput.stories.tsx @@ -1,5 +1,12 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, screen, waitFor, within } from "storybook/test"; import { StreamingOutput } from "./ConversationTimeline"; +import { + buildLiveStatus, + buildReconnectState, + buildRetryState, + FIXTURE_NOW, +} from "./storyFixtures"; // StreamingOutput renders inside a ConversationItem > Message > MessageContent // chain, but it's self-contained enough to render standalone. @@ -14,6 +21,13 @@ const meta: Meta = {
), ], + beforeEach: () => { + const real = Date.now; + Date.now = () => FIXTURE_NOW; + return () => { + Date.now = real; + }; + }, }; export default meta; type Story = StoryObj; @@ -23,73 +37,174 @@ export const ThinkingPlaceholder: Story = { args: { streamState: null, streamTools: [], - showInitialPlaceholder: true, + liveStatus: buildLiveStatus({ isAwaitingFirstStreamChunk: true }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const matches = canvas.getAllByText("Thinking..."); + expect(matches.length).toBeGreaterThanOrEqual(1); + expect( + canvas.queryByRole("heading", { name: /retrying request/i }), + ).not.toBeInTheDocument(); }, }; -/** First retry attempt. */ -export const RetryAttempt1: Story = { +/** Transport reconnects render a non-terminal reconnecting callout. */ +export const ReconnectingAfterDisconnect: Story = { args: { streamState: null, streamTools: [], - showInitialPlaceholder: true, - retryState: { attempt: 1, error: "service unavailable" }, + liveStatus: buildLiveStatus({ + reconnectState: buildReconnectState({ + attempt: 2, + delayMs: 2000, + }), + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /reconnecting/i }), + ).toBeVisible(); + expect(canvas.getByText(/chat stream disconnected/i)).toBeVisible(); + expect(canvas.getByText(/attempt 2/i)).toBeVisible(); + await waitFor(() => { + expect(canvasElement.textContent).toMatch(/reconnecting in \d+s/i); + }); + expect(canvas.queryByText("generic")).not.toBeInTheDocument(); + const thinkingMatches = canvas.getAllByText(/thinking\.\.\./i); + expect(thinkingMatches.length).toBeGreaterThanOrEqual(1); }, }; -/** Third retry attempt. */ -export const RetryAttempt3: Story = { +/** Generic retry reasons show the mux-style retry callout. */ +export const RetryWithVisibleReason: Story = { args: { streamState: null, streamTools: [], - showInitialPlaceholder: true, - retryState: { attempt: 3, error: "rate limit exceeded" }, + liveStatus: buildLiveStatus({ + retryState: buildRetryState(), + isAwaitingFirstStreamChunk: true, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /retrying request/i }), + ).toBeVisible(); + expect(canvas.getByText(/transient upstream failure/i)).toBeVisible(); + expect(canvas.getByText("generic")).toBeVisible(); + expect(canvas.getByText(/attempt 1/i)).toBeVisible(); }, }; -/** Higher attempt number to see how it looks. */ -export const RetryHighAttempt: Story = { +/** Rate-limited retries expose the normalized kind and delay metadata. */ +export const RetryRateLimited: Story = { args: { streamState: null, streamTools: [], - showInitialPlaceholder: true, - retryState: { attempt: 12, error: "overloaded" }, + liveStatus: buildLiveStatus({ + retryState: buildRetryState({ + attempt: 3, + error: "Anthropic asked us to back off briefly before retrying.", + kind: "rate_limit", + delayMs: 3000, + }), + isAwaitingFirstStreamChunk: true, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /rate limited/i }), + ).toBeVisible(); + expect(canvas.getByText("rate_limit")).toBeVisible(); + await waitFor(() => { + expect(canvasElement.textContent).toMatch(/retrying in \d+s/i); + }); + expect( + canvas.queryByRole("link", { name: /status/i }), + ).not.toBeInTheDocument(); }, }; -/** Active streaming with partial text content. */ -export const StreamingWithText: Story = { +/** Invalid retry timestamps hide the countdown instead of rendering NaN. */ +export const RetryInvalidTimestamp: Story = { args: { - streamState: { - blocks: [ - { - type: "response" as const, - text: "Here is a partial response that is still being generated...", - }, - ], - toolCalls: {}, - toolResults: {}, - sources: [], - }, + streamState: null, streamTools: [], + liveStatus: buildLiveStatus({ + retryState: buildRetryState({ + attempt: 3, + error: "Anthropic asked us to back off briefly before retrying.", + kind: "rate_limit", + delayMs: 3000, + retryingAt: "not-a-date", + }), + isAwaitingFirstStreamChunk: true, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /rate limited/i }), + ).toBeVisible(); + expect(canvas.getByText("rate_limit")).toBeVisible(); + expect(canvas.getByText(/attempt 3/i)).toBeVisible(); + await waitFor(() => { + expect(canvas.queryByText(/retrying in nan/i)).not.toBeInTheDocument(); + expect(canvas.queryByText(/retrying in \d+s/i)).not.toBeInTheDocument(); + }); }, }; -/** Content arrived after retries (no retry indicator shown). */ -export const StreamingAfterRetry: Story = { +/** Overloaded retries expose provider status links while retrying. */ +export const RetryOverloaded: Story = { args: { - streamState: { - blocks: [ - { - type: "response" as const, - text: "Successfully connected after retry. Here is your answer...", - }, - ], - toolCalls: {}, - toolResults: {}, - sources: [], - }, + streamState: null, streamTools: [], - retryState: null, + liveStatus: buildLiveStatus({ + retryState: buildRetryState({ + kind: "overloaded", + provider: "anthropic", + error: "Anthropic is currently overloaded. Retrying your request.", + }), + isAwaitingFirstStreamChunk: true, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /service overloaded/i }), + ).toBeVisible(); + expect(canvas.getByText("overloaded")).toBeVisible(); + const statusLink = screen.getByRole("link", { name: /status/i }); + expect(statusLink).toBeVisible(); + expect(statusLink).toHaveAttribute("href", "https://status.anthropic.com"); + }, +}; + +/** Timeout retries render the timeout-specific heading without a status CTA. */ +export const RetryTimeout: Story = { + args: { + streamState: null, + streamTools: [], + liveStatus: buildLiveStatus({ + retryState: buildRetryState({ + kind: "timeout", + error: "The provider took too long to respond. Retrying now.", + }), + isAwaitingFirstStreamChunk: true, + }), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByRole("heading", { name: /request time(?:out|d out)/i }), + ).toBeVisible(); + expect(canvas.getByText("timeout")).toBeVisible(); + expect( + canvas.queryByRole("link", { name: /status/i }), + ).not.toBeInTheDocument(); }, }; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/chatStatusHelpers.ts b/site/src/pages/AgentsPage/components/AgentDetail/chatStatusHelpers.ts new file mode 100644 index 0000000000..eb57273295 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/chatStatusHelpers.ts @@ -0,0 +1,29 @@ +const PROVIDER_STATUS_URLS: Record = { + anthropic: "https://status.anthropic.com", +}; + +export const getErrorTitle = ( + kind: string, + mode: "retry" | "error", +): string => { + switch (kind) { + case "overloaded": + return "Service overloaded"; + case "rate_limit": + return "Rate limited"; + case "timeout": + return "Request timeout"; + default: + return mode === "retry" ? "Retrying request" : "Request failed"; + } +}; + +export const getProviderStatusURL = ( + kind: string, + provider?: string, +): string | undefined => { + if (!provider || kind !== "overloaded") { + return undefined; + } + return PROVIDER_STATUS_URLS[provider.toLowerCase()]; +}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/chatStore.test.ts b/site/src/pages/AgentsPage/components/AgentDetail/chatStore.test.ts index 56fcc035df..d6582997f5 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/chatStore.test.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/chatStore.test.ts @@ -198,8 +198,14 @@ describe("setStreamError / clearStreamError", () => { it("stores and clears a stream error", () => { const store = createChatStore(); - store.setStreamError("connection lost"); - expect(store.getSnapshot().streamError).toBe("connection lost"); + store.setStreamError({ + kind: "generic", + message: "connection lost", + }); + expect(store.getSnapshot().streamError).toEqual({ + kind: "generic", + message: "connection lost", + }); store.clearStreamError(); expect(store.getSnapshot().streamError).toBeNull(); @@ -207,13 +213,19 @@ describe("setStreamError / clearStreamError", () => { it("does not notify when setting the same error", () => { const store = createChatStore(); - store.setStreamError("oops"); + store.setStreamError({ + kind: "generic", + message: "oops", + }); let notified = false; store.subscribe(() => { notified = true; }); - store.setStreamError("oops"); + store.setStreamError({ + kind: "generic", + message: "oops", + }); expect(notified).toBe(false); }); @@ -239,10 +251,21 @@ describe("setRetryState / clearRetryState", () => { it("stores and clears retry state", () => { const store = createChatStore(); - store.setRetryState({ attempt: 1, error: "rate limited" }); + store.setRetryState({ + attempt: 1, + error: "rate limited", + kind: "rate_limit", + provider: "anthropic", + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", + }); expect(store.getSnapshot().retryState).toEqual({ attempt: 1, error: "rate limited", + kind: "rate_limit", + provider: "anthropic", + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", }); store.clearRetryState(); @@ -262,6 +285,42 @@ describe("setRetryState / clearRetryState", () => { }); }); +// --------------------------------------------------------------------------- +// setReconnectState / clearReconnectState +// --------------------------------------------------------------------------- + +describe("setReconnectState / clearReconnectState", () => { + it("stores and clears reconnect state", () => { + const store = createChatStore(); + + store.setReconnectState({ + attempt: 2, + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", + }); + expect(store.getSnapshot().reconnectState).toEqual({ + attempt: 2, + delayMs: 3000, + retryingAt: "2025-01-01T00:00:30.000Z", + }); + + store.clearReconnectState(); + expect(store.getSnapshot().reconnectState).toBeNull(); + }); + + it("clearReconnectState is a no-op when already null", () => { + const store = createChatStore(); + + let notified = false; + store.subscribe(() => { + notified = true; + }); + store.clearReconnectState(); + + expect(notified).toBe(false); + }); +}); + // --------------------------------------------------------------------------- // setSubagentStatusOverride // --------------------------------------------------------------------------- @@ -427,11 +486,26 @@ describe("applyMessagePart / applyMessageParts", () => { // --------------------------------------------------------------------------- describe("resetTransientState", () => { - it("clears streamState, streamError, retryState, and subagentOverrides", () => { + it("clears streamState, streamError, retryState, reconnectState, and subagentOverrides", () => { const store = createChatStore(); store.applyMessagePart({ type: "text", text: "stream" }); - store.setStreamError("oops"); - store.setRetryState({ attempt: 2, error: "rate limit" }); + store.setStreamError({ + kind: "generic", + message: "oops", + }); + store.setRetryState({ + attempt: 2, + error: "rate limit", + kind: "rate_limit", + provider: "anthropic", + delayMs: 5000, + retryingAt: "2025-01-01T00:01:00.000Z", + }); + store.setReconnectState({ + attempt: 1, + delayMs: 1000, + retryingAt: "2025-01-01T00:00:01.000Z", + }); store.setSubagentStatusOverride("sub-1", "error"); store.resetTransientState(); @@ -440,6 +514,7 @@ describe("resetTransientState", () => { expect(state.streamState).toBeNull(); expect(state.streamError).toBeNull(); expect(state.retryState).toBeNull(); + expect(state.reconnectState).toBeNull(); expect(state.subagentStatusOverrides.size).toBe(0); }); @@ -447,7 +522,10 @@ describe("resetTransientState", () => { const store = createChatStore(); store.replaceMessages([makeMessage(1, "user", "hello")]); store.setQueuedMessages([makeQueuedMessage(10, "queued")]); - store.setStreamError("oops"); + store.setStreamError({ + kind: "generic", + message: "oops", + }); store.resetTransientState(); diff --git a/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.test.ts b/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.test.ts new file mode 100644 index 0000000000..937ee38481 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.test.ts @@ -0,0 +1,195 @@ +import { describe, expect, it } from "vitest"; +import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import { deriveLiveStatus } from "./liveStatusModel"; +import type { ReconnectState, RetryState, StreamState } from "./types"; + +const makeStreamState = ( + overrides: Partial = {}, +): StreamState => ({ + blocks: [], + toolCalls: {}, + toolResults: {}, + sources: [], + ...overrides, +}); + +const makeRetryState = (overrides: Partial = {}): RetryState => ({ + attempt: 2, + error: "Retrying request shortly.", + kind: "generic", + provider: "anthropic", + delayMs: 2000, + retryingAt: "2026-03-10T00:00:02.000Z", + ...overrides, +}); + +const makeReconnectState = ( + overrides: Partial = {}, +): ReconnectState => ({ + attempt: 1, + delayMs: 1000, + retryingAt: "2026-03-10T00:00:01.000Z", + ...overrides, +}); + +const makeStreamError = ( + overrides: Partial = {}, +): ChatDetailError => ({ + kind: "generic", + message: "Chat processing failed.", + provider: "anthropic", + retryable: false, + statusCode: 500, + ...overrides, +}); + +const derive = ( + overrides: Partial[0]> = {}, +) => + deriveLiveStatus({ + streamState: null, + retryState: null, + reconnectState: null, + streamError: null, + persistedError: null, + isAwaitingFirstStreamChunk: false, + ...overrides, + }); + +describe("deriveLiveStatus", () => { + const retryingStatus = { + phase: "retrying", + hasAccumulatedOutput: false, + title: "Retrying request", + kind: "generic", + message: "Retrying request shortly.", + attempt: 2, + provider: "anthropic", + delayMs: 2000, + retryingAt: "2026-03-10T00:00:02.000Z", + }; + const reconnectingStatus = { + phase: "reconnecting", + hasAccumulatedOutput: false, + title: "Reconnecting", + message: "Chat stream disconnected. Reconnecting…", + attempt: 1, + delayMs: 1000, + retryingAt: "2026-03-10T00:00:01.000Z", + }; + const failedStatus = { + phase: "failed", + hasAccumulatedOutput: false, + title: "Request failed", + kind: "generic", + message: "Chat processing failed.", + provider: "anthropic", + retryable: false, + statusCode: 500, + }; + + it.each([ + ["idle", undefined, { phase: "idle", hasAccumulatedOutput: false }], + [ + "starting", + { isAwaitingFirstStreamChunk: true }, + { phase: "starting", hasAccumulatedOutput: false }, + ], + ["retrying", { retryState: makeRetryState() }, retryingStatus], + [ + "reconnecting", + { reconnectState: makeReconnectState() }, + reconnectingStatus, + ], + ["failed", { streamError: makeStreamError() }, failedStatus], + [ + "streaming", + { streamState: makeStreamState() }, + { phase: "streaming", hasAccumulatedOutput: false }, + ], + ])("returns %s", (_phase, overrides, expected) => { + expect(derive(overrides)).toEqual(expected); + }); + + it("uses the persisted error as the idle fallback", () => { + expect(derive({ persistedError: makeStreamError() })).toEqual(failedStatus); + }); + + it("keeps live stream state ahead of the persisted error fallback", () => { + expect( + derive({ + streamState: makeStreamState(), + persistedError: makeStreamError({ kind: "timeout" }), + }), + ).toEqual({ phase: "streaming", hasAccumulatedOutput: false }); + }); + + it("tracks accumulated output on failed streams", () => { + expect( + derive({ + streamState: makeStreamState({ + blocks: [{ type: "response", text: "Partial response" }], + }), + streamError: makeStreamError(), + }), + ).toEqual({ + ...failedStatus, + hasAccumulatedOutput: true, + }); + }); + + it("tracks accumulated output while reconnecting", () => { + expect( + derive({ + streamState: makeStreamState({ + blocks: [{ type: "response", text: "Partial response" }], + }), + reconnectState: makeReconnectState(), + }), + ).toEqual({ + ...reconnectingStatus, + hasAccumulatedOutput: true, + }); + }); + + it("prioritizes retrying over failed and reconnecting", () => { + expect( + derive({ + retryState: makeRetryState({ kind: "rate_limit" }), + reconnectState: makeReconnectState({ attempt: 3 }), + streamError: makeStreamError({ kind: "timeout" }), + persistedError: makeStreamError({ kind: "generic" }), + isAwaitingFirstStreamChunk: true, + }), + ).toMatchObject({ phase: "retrying", kind: "rate_limit" }); + }); + + it("prioritizes failed over reconnecting and starting", () => { + expect( + derive({ + reconnectState: makeReconnectState(), + streamError: makeStreamError({ kind: "timeout" }), + persistedError: makeStreamError({ kind: "generic" }), + isAwaitingFirstStreamChunk: true, + }), + ).toMatchObject({ phase: "failed", kind: "timeout" }); + }); + + it("prioritizes reconnecting over starting", () => { + expect( + derive({ + reconnectState: makeReconnectState(), + isAwaitingFirstStreamChunk: true, + }), + ).toEqual(reconnectingStatus); + }); + + it("prioritizes starting over streaming", () => { + expect( + derive({ + streamState: makeStreamState(), + isAwaitingFirstStreamChunk: true, + }), + ).toEqual({ phase: "starting", hasAccumulatedOutput: false }); + }); +}); diff --git a/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.ts b/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.ts new file mode 100644 index 0000000000..dae109a725 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/liveStatusModel.ts @@ -0,0 +1,126 @@ +import type { ChatDetailError } from "../../utils/usageLimitMessage"; +import { getErrorTitle } from "./chatStatusHelpers"; +import type { ReconnectState, RetryState, StreamState } from "./types"; + +type LiveStatusBase = { + hasAccumulatedOutput: boolean; +}; + +const RECONNECTING_TITLE = "Reconnecting"; +const RECONNECTING_MESSAGE = "Chat stream disconnected. Reconnecting…"; + +export type LiveStatusModel = + | ({ phase: "idle" } & LiveStatusBase) + | ({ phase: "starting" } & LiveStatusBase) + | ({ phase: "streaming" } & LiveStatusBase) + | ({ + phase: "retrying"; + title: string; + kind: string; + message: string; + attempt: number; + provider?: string; + delayMs?: number; + retryingAt?: string; + } & LiveStatusBase) + | ({ + phase: "reconnecting"; + title: string; + message: string; + attempt: number; + delayMs: number; + retryingAt: string; + } & LiveStatusBase) + | ({ + phase: "failed"; + title: string; + kind: string; + message: string; + provider?: string; + retryable?: boolean; + statusCode?: number; + } & LiveStatusBase); + +export type DeriveLiveStatusParams = { + streamState: StreamState | null; + retryState: RetryState | null; + reconnectState: ReconnectState | null; + streamError: ChatDetailError | null; + persistedError: ChatDetailError | null; + isAwaitingFirstStreamChunk: boolean; +}; + +const getHasAccumulatedOutput = (streamState: StreamState | null): boolean => + Boolean(streamState && streamState.blocks.length > 0); + +const toReconnectingLiveStatus = ( + reconnectState: ReconnectState, + options: { hasAccumulatedOutput?: boolean } = {}, +): Extract => ({ + phase: "reconnecting", + hasAccumulatedOutput: options.hasAccumulatedOutput ?? false, + title: RECONNECTING_TITLE, + message: RECONNECTING_MESSAGE, + ...reconnectState, +}); + +const toFailedLiveStatus = ( + error: ChatDetailError, + options: { hasAccumulatedOutput?: boolean } = {}, +): Extract => ({ + phase: "failed", + hasAccumulatedOutput: options.hasAccumulatedOutput ?? false, + title: getErrorTitle(error.kind, "error"), + kind: error.kind, + message: error.message, + provider: error.provider, + retryable: error.retryable, + statusCode: error.statusCode, +}); + +export const deriveLiveStatus = ({ + streamState, + retryState, + reconnectState, + streamError, + persistedError, + isAwaitingFirstStreamChunk, +}: DeriveLiveStatusParams): LiveStatusModel => { + const hasAccumulatedOutput = getHasAccumulatedOutput(streamState); + + if (retryState) { + return { + phase: "retrying", + hasAccumulatedOutput, + title: getErrorTitle(retryState.kind, "retry"), + kind: retryState.kind, + message: retryState.error, + attempt: retryState.attempt, + provider: retryState.provider, + delayMs: retryState.delayMs, + retryingAt: retryState.retryingAt, + }; + } + + if (streamError) { + return toFailedLiveStatus(streamError, { hasAccumulatedOutput }); + } + + if (reconnectState) { + return toReconnectingLiveStatus(reconnectState, { hasAccumulatedOutput }); + } + + if (isAwaitingFirstStreamChunk) { + return { phase: "starting", hasAccumulatedOutput }; + } + + if (streamState !== null) { + return { phase: "streaming", hasAccumulatedOutput }; + } + + if (persistedError) { + return toFailedLiveStatus(persistedError, { hasAccumulatedOutput }); + } + + return { phase: "idle", hasAccumulatedOutput }; +}; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/storyFixtures.ts b/site/src/pages/AgentsPage/components/AgentDetail/storyFixtures.ts new file mode 100644 index 0000000000..b448098ae2 --- /dev/null +++ b/site/src/pages/AgentsPage/components/AgentDetail/storyFixtures.ts @@ -0,0 +1,90 @@ +import type * as TypesGen from "api/typesGenerated"; +import { + type DeriveLiveStatusParams, + deriveLiveStatus, + type LiveStatusModel, +} from "./liveStatusModel"; +import { applyMessagePartToStreamState, buildStreamTools } from "./streamState"; +import type { + MergedTool, + ReconnectState, + RetryState, + StreamState, +} from "./types"; + +type StoryStreamRenderState = { + streamState: StreamState | null; + streamTools: readonly MergedTool[]; + liveStatus: LiveStatusModel; +}; + +const DEFAULT_LIVE_STATUS_PARAMS: DeriveLiveStatusParams = { + streamState: null, + retryState: null, + reconnectState: null, + streamError: null, + persistedError: null, + isAwaitingFirstStreamChunk: false, +}; + +export const buildLiveStatus = ( + overrides: Partial = {}, +): LiveStatusModel => + deriveLiveStatus({ + ...DEFAULT_LIVE_STATUS_PARAMS, + ...overrides, + }); + +export const buildStreamRenderState = ( + parts: readonly TypesGen.ChatMessagePart[], +): StoryStreamRenderState => { + let streamState: StreamState | null = null; + for (const part of parts) { + streamState = applyMessagePartToStreamState(streamState, part); + } + + return { + streamState, + streamTools: buildStreamTools(streamState), + liveStatus: buildLiveStatus({ streamState }), + }; +}; + +/** + * Pinned clock for stories that render countdown timers. Stories + * should mock `Date.now` to return this value so the countdowns + * are deterministic across Chromatic snapshots. + * + * Set to midnight UTC on the same day as the fixture deadlines, + * giving reconnect a 1s countdown and retry a 2s countdown. + */ +export const FIXTURE_NOW = new Date("2026-03-10T00:00:00.000Z").getTime(); + +export const buildReconnectState = ( + overrides: Partial = {}, +): ReconnectState => ({ + attempt: 1, + delayMs: 1000, + retryingAt: "2026-03-10T00:00:01.000Z", + ...overrides, +}); + +export const buildRetryState = ( + overrides: Partial = {}, +): RetryState => ({ + attempt: 1, + error: + "Anthropic is retrying your request after a transient upstream failure.", + kind: "generic", + provider: "anthropic", + delayMs: 2000, + retryingAt: "2026-03-10T00:00:02.000Z", + ...overrides, +}); + +export const textResponseStreamParts = [ + { + type: "text", + text: "Storybook streamed answer.", + }, +] satisfies readonly TypesGen.ChatMessagePart[]; diff --git a/site/src/pages/AgentsPage/components/AgentDetail/types.ts b/site/src/pages/AgentsPage/components/AgentDetail/types.ts index 7b1b154fdb..c258458973 100644 --- a/site/src/pages/AgentsPage/components/AgentDetail/types.ts +++ b/site/src/pages/AgentsPage/components/AgentDetail/types.ts @@ -1,4 +1,5 @@ import type * as TypesGen from "api/typesGenerated"; +import type { ReconnectSchedule } from "utils/reconnectingWebSocket"; export type ParsedToolCall = { id: string; @@ -60,6 +61,17 @@ export type ParsedMessageEntry = { parsed: ParsedMessageContent; }; +export type ReconnectState = ReconnectSchedule; + +export type RetryState = { + attempt: number; + error: string; + kind: string; + provider?: string; + delayMs?: number; + retryingAt?: string; +}; + type StreamToolCall = { id: string; name: string; diff --git a/site/src/pages/AgentsPage/components/AgentDetailContent.tsx b/site/src/pages/AgentsPage/components/AgentDetailContent.tsx index 51fa99da89..d4ca29d0ba 100644 --- a/site/src/pages/AgentsPage/components/AgentDetailContent.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetailContent.tsx @@ -17,21 +17,16 @@ import { selectMessagesByID, selectOrderedMessageIDs, selectQueuedMessages, - selectRetryState, - selectStreamError, - selectStreamState, - selectSubagentStatusOverrides, useChatSelector, type useChatStore, } from "./AgentDetail/ChatContext"; import { ConversationTimeline } from "./AgentDetail/ConversationTimeline"; import { getLatestContextUsage } from "./AgentDetail/chatHelpers"; +import { LiveStreamTail } from "./AgentDetail/LiveStreamTail"; import { buildSubagentTitles, parseMessagesWithMergedTools, } from "./AgentDetail/messageParsing"; -import { buildStreamTools } from "./AgentDetail/streamState"; -import type { ParsedMessageEntry } from "./AgentDetail/types"; import { useOnRenderProfiler } from "./AgentDetail/useOnRenderProfiler"; type ChatStoreHandle = ReturnType["store"]; @@ -41,8 +36,9 @@ const isChatMessage = ( ): message is TypesGen.ChatMessage => Boolean(message); interface AgentDetailTimelineProps { + chatID?: string; store: ChatStoreHandle; - persistedErrorReason: ChatDetailError | undefined; + persistedError: ChatDetailError | undefined; onEditUserMessage?: ( messageId: number, text: string, @@ -54,12 +50,10 @@ interface AgentDetailTimelineProps { mcpServers?: readonly TypesGen.MCPServerConfig[]; } -// Reads only message-related store state (stable during streaming). -// Computes parsedMessages once and passes them to ConversationTimeline -// via a memo boundary so that streaming ticks don't re-parse history. -const MessageListProvider: FC = ({ +export const AgentDetailTimeline: FC = ({ + chatID, store, - persistedErrorReason, + persistedError, onEditUserMessage, editingMessageId, savingMessageId, @@ -68,123 +62,39 @@ const MessageListProvider: FC = ({ }) => { const messagesByID = useChatSelector(store, selectMessagesByID); const orderedMessageIDs = useChatSelector(store, selectOrderedMessageIDs); - const chatStatus = useChatSelector(store, selectChatStatus); - const streamError = useChatSelector(store, selectStreamError); - const subagentStatusOverrides = useChatSelector( - store, - selectSubagentStatusOverrides, - ); - const retryState = useChatSelector(store, selectRetryState); const messages = orderedMessageIDs .map((messageID) => messagesByID.get(messageID)) .filter(isChatMessage); const parsedMessages = parseMessagesWithMergedTools(messages); const subagentTitles = buildSubagentTitles(parsedMessages); - const detailError: ChatDetailError | undefined = - (persistedErrorReason?.kind === "usage-limit" || chatStatus === "error" - ? persistedErrorReason - : undefined) ?? - (streamError - ? { kind: "generic" as const, message: streamError } - : undefined); - const latestMessage = messages[messages.length - 1]; - const latestMessageNeedsAssistantResponse = - !latestMessage || latestMessage.role !== "assistant"; - - return ( - - ); -}; - -// Reads stream-specific store state (changes every token). Isolated -// so that streamState changes don't invalidate parsedMessages above. -const StreamingBridge: FC<{ - store: ChatStoreHandle; - isEmpty: boolean; - parsedMessages: ParsedMessageEntry[]; - subagentTitles: Map; - subagentStatusOverrides: Map; - retryState: { attempt: number; error: string } | null; - detailError: ChatDetailError | undefined; - latestMessageNeedsAssistantResponse: boolean; - chatStatus: TypesGen.ChatStatus | null; - onEditUserMessage?: ( - messageId: number, - text: string, - fileBlocks?: readonly TypesGen.ChatMessagePart[], - ) => void; - editingMessageId?: number | null; - savingMessageId?: number | null; - urlTransform?: UrlTransform; - mcpServers?: readonly TypesGen.MCPServerConfig[]; -}> = ({ - store, - isEmpty, - parsedMessages, - subagentTitles, - subagentStatusOverrides, - retryState, - detailError, - latestMessageNeedsAssistantResponse, - chatStatus, - onEditUserMessage, - editingMessageId, - savingMessageId, - urlTransform, - mcpServers, -}) => { - const streamState = useChatSelector(store, selectStreamState); - const streamTools = buildStreamTools(streamState); const onRenderProfiler = useOnRenderProfiler(); - const isAwaitingFirstStreamChunk = - !streamState && - (chatStatus === "running" || chatStatus === "pending") && - latestMessageNeedsAssistantResponse; - const hasStreamOutput = Boolean(streamState) || isAwaitingFirstStreamChunk; return ( - +
+ + +
); }; -export const AgentDetailTimeline: FC = (props) => { - return ; -}; - interface AgentDetailInputProps { store: ChatStoreHandle; compressionThreshold: number | undefined; @@ -218,6 +128,7 @@ interface AgentDetailInputProps { // File parts from the message being edited, converted to // File objects and pre-populated into attachments. editingFileBlocks?: readonly TypesGen.ChatMessagePart[]; + // MCP server picker state. mcpServers?: readonly TypesGen.MCPServerConfig[]; selectedMCPServerIds?: readonly string[]; onMCPSelectionChange?: (ids: string[]) => void; diff --git a/site/src/pages/AgentsPage/components/AgentDetailView.stories.tsx b/site/src/pages/AgentsPage/components/AgentDetailView.stories.tsx index add36f96f0..dcd5a29a1f 100644 --- a/site/src/pages/AgentsPage/components/AgentDetailView.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetailView.stories.tsx @@ -8,6 +8,7 @@ import type { ComponentProps, FC } from "react"; import { expect, fn, spyOn, userEvent, waitFor, within } from "storybook/test"; import { reactRouterParameters } from "storybook-addon-remix-react-router"; import type { ModelSelectorOption } from "#/components/ai-elements"; +import type { ChatDetailError } from "../utils/usageLimitMessage"; import { createChatStore } from "./AgentDetail/ChatContext"; import { AgentDetailLoadingView, @@ -97,11 +98,8 @@ const StoryAgentDetailView: FC = ({ editing, ...overrides }) => { const props = { agentId: AGENT_ID, chatTitle: "Help me refactor", - chatErrorReasons: {} as ComponentProps< - typeof AgentDetailView - >["chatErrorReasons"], + persistedError: undefined as ChatDetailError | undefined, parentChat: undefined as TypesGen.Chat | undefined, - chatRecord: buildChat(), isArchived: false, hasWorkspace: true, store: createChatStore(), @@ -187,13 +185,7 @@ export const Default: Story = { /** Archived agent displays the read-only banner below the top bar. */ export const Archived: Story = { - render: () => ( - - ), + render: () => , }; /** Shows the parent chat link in the top bar when a parent exists. */ @@ -209,8 +201,12 @@ export const WithParentChat: Story = { export const WithError: Story = { render: () => ( ), diff --git a/site/src/pages/AgentsPage/components/AgentDetailView.tsx b/site/src/pages/AgentsPage/components/AgentDetailView.tsx index 31b4c45143..3db8c57597 100644 --- a/site/src/pages/AgentsPage/components/AgentDetailView.tsx +++ b/site/src/pages/AgentsPage/components/AgentDetailView.tsx @@ -9,11 +9,7 @@ import type { ModelSelectorOption } from "#/components/ai-elements"; import { Button } from "#/components/Button/Button"; import type { ChatDetailError } from "../utils/usageLimitMessage"; import { AgentChatInput, type ChatMessageInputRef } from "./AgentChatInput"; -import { - selectChatStatus, - useChatSelector, - type useChatStore, -} from "./AgentDetail/ChatContext"; +import type { useChatStore } from "./AgentDetail/ChatContext"; import { AgentDetailTopBar } from "./AgentDetail/TopBar"; import { AgentDetailInput, AgentDetailTimeline } from "./AgentDetailContent"; import { @@ -55,8 +51,7 @@ interface AgentDetailViewProps { agentId: string; chatTitle: string | undefined; parentChat: TypesGen.Chat | undefined; - chatErrorReasons: Record; - chatRecord: TypesGen.Chat | undefined; + persistedError: ChatDetailError | undefined; isArchived: boolean; hasWorkspace: boolean; @@ -139,8 +134,7 @@ export const AgentDetailView: FC = ({ agentId, chatTitle, parentChat, - chatErrorReasons, - chatRecord, + persistedError, isArchived, hasWorkspace, store, @@ -192,7 +186,6 @@ export const AgentDetailView: FC = ({ null, ); const visualExpanded = dragVisualExpanded ?? isRightPanelExpanded; - const chatStatus = useChatSelector(store, selectChatStatus); // Compute local diff stats from git watcher unified diffs. @@ -269,13 +262,9 @@ export const AgentDetailView: FC = ({ >
{ }); }); +describe("chatDetailErrorsEqual", () => { + it("compares matching errors by value", () => { + const left: ChatDetailError = { + kind: "rate_limit", + message: "Slow down.", + provider: "anthropic", + retryable: true, + statusCode: 429, + }; + + expect(chatDetailErrorsEqual(left, { ...left })).toBe(true); + }); + + it("treats missing and mismatched errors as different", () => { + const error: ChatDetailError = { + kind: "generic", + message: "Provider request failed.", + }; + + expect(chatDetailErrorsEqual(error, null)).toBe(false); + expect(chatDetailErrorsEqual(error, { ...error, statusCode: 500 })).toBe( + false, + ); + }); +}); + describe("isUsageLimitData", () => { it("accepts a fully populated valid payload", () => { const error: ChatDetailError = { diff --git a/site/src/pages/AgentsPage/utils/usageLimitMessage.ts b/site/src/pages/AgentsPage/utils/usageLimitMessage.ts index 3775ccf08d..00193a6eb3 100644 --- a/site/src/pages/AgentsPage/utils/usageLimitMessage.ts +++ b/site/src/pages/AgentsPage/utils/usageLimitMessage.ts @@ -13,11 +13,43 @@ interface UsageLimitData { /** * Typed classification for errors surfaced in the agent detail view. * - "usage-limit": the user hit a spending cap (409 + valid usage data). - * - "generic": any other error (stream failures, last_error, etc.). + * - other kinds come from normalized stream/provider failures such as + * "generic", "overloaded", "rate_limit", or "timeout". */ export type ChatDetailError = { message: string; - kind: "generic" | "usage-limit"; + kind: + | "usage-limit" + | "generic" + | "overloaded" + | "rate_limit" + | "timeout" + | (string & {}); + provider?: string; + retryable?: boolean; + statusCode?: number; +}; + +/** + * Compare two chat-detail errors by their user-visible fields. + */ +export const chatDetailErrorsEqual = ( + left: ChatDetailError | null | undefined, + right: ChatDetailError | null | undefined, +): boolean => { + if (left === right) { + return true; + } + if (!left || !right) { + return false; + } + return ( + left.kind === right.kind && + left.message === right.message && + left.provider === right.provider && + left.retryable === right.retryable && + left.statusCode === right.statusCode + ); }; /** diff --git a/site/src/utils/OneWayWebSocket.ts b/site/src/utils/OneWayWebSocket.ts index e0d677637b..fb2dcd3a3a 100644 --- a/site/src/utils/OneWayWebSocket.ts +++ b/site/src/utils/OneWayWebSocket.ts @@ -46,7 +46,7 @@ type OneWayEventCallback = ( payload: OneWayEventPayloadMap[TEvent], ) => void; -interface OneWayWebSocketApi { +export interface OneWayWebSocketApi { get url(): string; addEventListener: ( diff --git a/site/src/utils/reconnectingWebSocket.test.ts b/site/src/utils/reconnectingWebSocket.test.ts index 0f1e2b03cc..47444c0461 100644 --- a/site/src/utils/reconnectingWebSocket.test.ts +++ b/site/src/utils/reconnectingWebSocket.test.ts @@ -1,9 +1,12 @@ -import { createReconnectingWebSocket } from "./reconnectingWebSocket"; +import { + createReconnectingWebSocket, + type ReconnectSchedule, +} from "./reconnectingWebSocket"; /** - * Minimal mock that satisfies the {@link Closable} interface used by - * the reconnection utility. Each instance records every - * `addEventListener` call and exposes helpers to fire those events. + * Minimal mock that satisfies the {@link Closable} interface used by the + * reconnection utility. Each instance records every `addEventListener` + * call and exposes helpers to fire those events. */ function createMockSocket() { const listeners: Record void>> = {}; @@ -27,8 +30,19 @@ function createMockSocket() { return socket; } +const expectReconnectSchedule = ( + event: { reconnect: ReconnectSchedule; now: number }, + expected: { attempt: number; delayMs: number }, +) => { + expect(event.reconnect).toMatchObject(expected); + expect(Date.parse(event.reconnect.retryingAt) - event.now).toBe( + expected.delayMs, + ); +}; + beforeEach(() => { vi.useFakeTimers(); + vi.setSystemTime(new Date("2025-01-01T00:00:00.000Z")); }); afterEach(() => { @@ -69,7 +83,11 @@ describe("createReconnectingWebSocket", () => { activeSocket = createMockSocket(); return activeSocket; }); - const onDisconnect = vi.fn(); + const disconnects: Array<{ reconnect: ReconnectSchedule; now: number }> = + []; + const onDisconnect = vi.fn((reconnect: ReconnectSchedule) => { + disconnects.push({ reconnect, now: Date.now() }); + }); createReconnectingWebSocket({ connect, @@ -84,7 +102,7 @@ describe("createReconnectingWebSocket", () => { // First disconnect — should schedule reconnect after 1000ms. activeSocket.emit("close"); expect(onDisconnect).toHaveBeenCalledTimes(1); - expect(onDisconnect).toHaveBeenLastCalledWith(0); + expectReconnectSchedule(disconnects[0]!, { attempt: 1, delayMs: 1000 }); vi.advanceTimersByTime(999); expect(connect).toHaveBeenCalledTimes(1); @@ -94,7 +112,7 @@ describe("createReconnectingWebSocket", () => { // Second disconnect — delay should be 2000ms. activeSocket.emit("close"); expect(onDisconnect).toHaveBeenCalledTimes(2); - expect(onDisconnect).toHaveBeenLastCalledWith(1); + expectReconnectSchedule(disconnects[1]!, { attempt: 2, delayMs: 2000 }); vi.advanceTimersByTime(1999); expect(connect).toHaveBeenCalledTimes(2); @@ -103,6 +121,7 @@ describe("createReconnectingWebSocket", () => { // Third disconnect — delay should be 4000ms. activeSocket.emit("close"); + expectReconnectSchedule(disconnects[2]!, { attempt: 3, delayMs: 4000 }); vi.advanceTimersByTime(3999); expect(connect).toHaveBeenCalledTimes(3); vi.advanceTimersByTime(1); @@ -123,8 +142,8 @@ describe("createReconnectingWebSocket", () => { factor: 2, }); - // Disconnect enough times that the uncapped delay would - // exceed maxMs: 1000, 2000, 4000, 8000 → capped at 5000. + // Disconnect enough times that the uncapped delay would exceed + // maxMs: 1000, 2000, 4000, 8000 → capped at 5000. for (let i = 0; i < 3; i++) { activeSocket.emit("close"); vi.runOnlyPendingTimers(); @@ -163,8 +182,7 @@ describe("createReconnectingWebSocket", () => { activeSocket.emit("open"); activeSocket.emit("close"); - // Next reconnect should use the base delay (1000ms), not - // 4000ms. + // Next reconnect should use the base delay (1000ms), not 4000ms. vi.advanceTimersByTime(999); expect(connect).toHaveBeenCalledTimes(3); vi.advanceTimersByTime(1); @@ -206,14 +224,14 @@ describe("createReconnectingWebSocket", () => { createReconnectingWebSocket({ connect }); - const firstSocket = sockets[0]; + const firstSocket = sockets[0]!; firstSocket.emit("close"); vi.runOnlyPendingTimers(); - // The connect function creates a new socket. The old socket - // was already "closed" by the browser, but on a fresh - // reconnection the utility closes the previous one if it's - // still the active reference. + // The connect function creates a new socket. The old socket was + // already "closed" by the browser, but on a fresh reconnection + // the utility closes the previous one if it's still the active + // reference. expect(connect).toHaveBeenCalledTimes(2); }); @@ -265,9 +283,9 @@ describe("createReconnectingWebSocket", () => { dispose(); dispose(); - // close is idempotent on real WebSockets, so calling it - // multiple times is harmless. The important thing is that - // no reconnection is scheduled after the first dispose. + // close is idempotent on real WebSockets, so calling it multiple + // times is harmless. The important thing is that no reconnection is + // scheduled after the first dispose. expect(connect).toHaveBeenCalledTimes(1); }); diff --git a/site/src/utils/reconnectingWebSocket.ts b/site/src/utils/reconnectingWebSocket.ts index 294832ab43..47ed0c976c 100644 --- a/site/src/utils/reconnectingWebSocket.ts +++ b/site/src/utils/reconnectingWebSocket.ts @@ -17,8 +17,10 @@ * onOpen() { * console.log("connected"); * }, - * onDisconnect() { - * console.log("disconnected, will reconnect automatically"); + * onDisconnect(reconnect) { + * console.log( + * `disconnected, reconnecting in ${reconnect.delayMs}ms`, + * ); * }, * }); * @@ -37,9 +39,20 @@ const RECONNECT_MAX_MS = 10_000; const RECONNECT_FACTOR = 2; /** - * A minimal WebSocket-like interface that the reconnection utility - * can manage. Both native `WebSocket` and `OneWayWebSocket` satisfy - * this contract. + * Metadata for the reconnect attempt that was just scheduled. + * `attempt` is 1-based and user-facing: `1` means the first retry after + * the connection dropped. + */ +export type ReconnectSchedule = { + attempt: number; + delayMs: number; + retryingAt: string; +}; + +/** + * A minimal WebSocket-like interface that the reconnection utility can + * manage. Both native `WebSocket` and `OneWayWebSocket` satisfy this + * contract. */ interface Closable { addEventListener(event: string, handler: (...args: unknown[]) => void): void; @@ -63,24 +76,18 @@ interface ReconnectingWebSocketOptions { connect: () => TSocket; /** - * Called when a connection succeeds (the socket fires `open`). - * The backoff counter is reset before this callback runs. + * Called when a connection succeeds (the socket fires `open`). The + * backoff counter is reset before this callback runs. */ onOpen?: (socket: TSocket) => void; /** - * Called on the first disconnect after a successful connection or - * on a connection failure. Fires at most once per socket instance - * (browsers fire both `error` and `close`; only the first is - * forwarded). A reconnection is scheduled automatically after - * this callback returns. - * - * @param attempt - The zero-based reconnection attempt counter - * *before* it is incremented for the upcoming retry. A value of - * `0` means this is the first disconnect since the last - * successful connection. + * Called on the first disconnect after a successful connection or on a + * connection failure. Fires at most once per socket instance (browsers + * fire both `error` and `close`; only the first is forwarded). The + * callback receives the reconnect attempt that was just scheduled. */ - onDisconnect?: (attempt: number) => void; + onDisconnect?: (reconnect: ReconnectSchedule) => void; /** Base delay in milliseconds. Defaults to {@link RECONNECT_BASE_MS}. */ baseMs?: number; @@ -92,22 +99,40 @@ interface ReconnectingWebSocketOptions { factor?: number; } +const getReconnectSchedule = ({ + attempt, + baseMs, + maxMs, + factor, +}: { + attempt: number; + baseMs: number; + maxMs: number; + factor: number; +}): ReconnectSchedule => { + const delayMs = Math.min(baseMs * factor ** (attempt - 1), maxMs); + return { + attempt, + delayMs, + retryingAt: new Date(Date.now() + delayMs).toISOString(), + }; +}; + /** * Creates a self-reconnecting WebSocket connection with capped * exponential backoff. * * The returned function disposes of the connection: it closes the - * active socket (if any), cancels any pending reconnection timer, - * and prevents further reconnection attempts. It is safe to call - * the dispose function more than once. + * active socket (if any), cancels any pending reconnection timer, and + * prevents further reconnection attempts. It is safe to call the + * dispose function more than once. * * Backoff delay formula: * ``` - * delay = min(baseMs * factor ^ attempt, maxMs) + * delay = min(baseMs * factor ^ (attempt - 1), maxMs) * ``` * - * The attempt counter resets to `0` whenever a connection - * successfully opens. + * The reconnect attempt counter resets after a successful `open`. * * @returns A dispose function that tears down the connection. */ @@ -124,25 +149,23 @@ export function createReconnectingWebSocket( } = options; let disposed = false; - let reconnectAttempt = 0; + let lastReconnectAttempt = 0; let reconnectTimer: ReturnType | null = null; let activeSocket: TSocket | null = null; - // Schedule a reconnect with capped exponential backoff. - // Does nothing if the connection has been disposed. - const scheduleReconnect = () => { + const scheduleReconnect = (reconnect: ReconnectSchedule) => { if (disposed) { return; } if (reconnectTimer !== null) { clearTimeout(reconnectTimer); } - const delay = Math.min(baseMs * factor ** reconnectAttempt, maxMs); - reconnectAttempt += 1; - reconnectTimer = setTimeout(connect, delay); + lastReconnectAttempt = reconnect.attempt; + reconnectTimer = setTimeout(connect, reconnect.delayMs); }; function connect() { + reconnectTimer = null; if (disposed) { return; } @@ -155,20 +178,26 @@ export function createReconnectingWebSocket( const handleOpen = () => { // Connection succeeded — reset backoff. - reconnectAttempt = 0; + lastReconnectAttempt = 0; onOpen?.(socket); }; const handleDisconnect = () => { - // Guard against duplicate calls: browsers fire both - // "error" and "close" on a failed WebSocket, so we - // only process the first event per socket instance. + // Guard against duplicate calls: browsers fire both "error" + // and "close" on a failed WebSocket, so we only process the + // first event per socket instance. if (activeSocket !== socket || disposed) { return; } activeSocket = null; - onDisconnect?.(reconnectAttempt); - scheduleReconnect(); + const reconnect = getReconnectSchedule({ + attempt: lastReconnectAttempt + 1, + baseMs, + maxMs, + factor, + }); + onDisconnect?.(reconnect); + scheduleReconnect(reconnect); }; socket.addEventListener("open", handleOpen);