diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 69f33d0b6a..a8a1d0bac4 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -37,6 +37,7 @@ import { import { isMobileViewport } from "#/utils/mobile"; import { pageTitle } from "#/utils/page"; import { rewriteLocalhostURL } from "#/utils/portForward"; +import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; import { AgentChatPageLoadingView, AgentChatPageNotFoundView, @@ -283,6 +284,56 @@ function resolveCompactionThreshold( return config.compression_threshold; } +// Compile-time guard: ensures the workspace watcher bailout comparison +// covers every WorkspaceAgent field the UI reads. If WorkspaceAgent +// gains a new field, this will error until the field is either added +// to the comparison or explicitly excluded here. +type _UncoveredAgentFields = Omit< + TypesGen.WorkspaceAgent, + | "id" + | "status" + | "name" + | "expanded_directory" + // Fields below are intentionally not compared. They change + // frequently (stats, metadata) or are objects/arrays that would + // require deep comparison, and the UI does not read them. + | "parent_id" + | "created_at" + | "updated_at" + | "first_connected_at" + | "last_connected_at" + | "disconnected_at" + | "started_at" + | "ready_at" + | "lifecycle_state" + | "resource_id" + | "instance_id" + | "architecture" + | "environment_variables" + | "operating_system" + | "logs_length" + | "logs_overflowed" + | "directory" + | "version" + | "api_version" + | "apps" + | "latency" + | "connection_timeout_seconds" + | "troubleshooting_url" + | "subsystems" + | "health" + | "display_apps" + | "log_sources" + | "scripts" + | "startup_script_behavior" +>; +// If this errors, a new field was added to WorkspaceAgent. +// Decide: does the UI read it? If yes, add it to the first +// section of the Omit above and to the bailout comparison +// in the workspace watcher message handler. If no, add it +// to the excluded section of the Omit. +const _agentFieldGuard: Record = {}; + const AgentChatPage: FC = () => { const { agentId } = useParams<{ agentId: string }>(); const { @@ -387,35 +438,54 @@ const AgentChatPage: FC = () => { if (!workspaceId) { return; } - const socket = watchWorkspace(workspaceId); - socket.addEventListener("message", (event) => { - if (event.parseError) { - return; - } - if (event.parsedMessage.type === "data") { - const next = event.parsedMessage.data as TypesGen.Workspace; - queryClient.setQueryData( - workspaceByIdKey(workspaceId), - (prev) => { - // Return the same reference when nothing the UI - // reads has changed. This prevents react-query - // from notifying subscribers and avoids a full - // AgentChatPage re-render on every heartbeat. - if ( - prev && - prev.latest_build.status === next.latest_build.status && - prev.latest_build.resources === next.latest_build.resources && - prev.name === next.name && - prev.owner_name === next.owner_name - ) { - return prev; - } - return next; - }, - ); - } + return createReconnectingWebSocket({ + connect() { + const socket = watchWorkspace(workspaceId); + socket.addEventListener("message", (event) => { + if (event.parseError) { + return; + } + if (event.parsedMessage.type === "data") { + const next = event.parsedMessage.data as TypesGen.Workspace; + queryClient.setQueryData( + workspaceByIdKey(workspaceId), + (prev) => { + // Return the same reference when nothing the UI + // reads has changed. This prevents react-query + // from notifying subscribers and avoids a full + // AgentChatPage re-render on every heartbeat. + const prevAgent = getWorkspaceAgent(prev, undefined); + const nextAgent = getWorkspaceAgent(next, undefined); + if ( + prev && + prev.latest_build.status === next.latest_build.status && + prev.name === next.name && + prev.owner_name === next.owner_name && + prevAgent?.id === nextAgent?.id && + prevAgent?.status === nextAgent?.status && + prevAgent?.name === nextAgent?.name && + prevAgent?.expanded_directory === + nextAgent?.expanded_directory + ) { + return prev; + } + return next; + }, + ); + } + }); + return socket; + }, + onOpen() { + // Refetch workspace data on reconnection to cover + // events missed while disconnected. Also fires on the + // initial connection (harmless, may deduplicate with + // the in-flight useQuery fetch). + void queryClient.invalidateQueries({ + queryKey: workspaceByIdKey(workspaceId), + }); + }, }); - return () => socket.close(); }, [workspaceId, queryClient]); const sshConfigQuery = useQuery(deploymentSSHConfig()); const workspace = workspaceQuery.data; diff --git a/site/src/pages/AgentsPage/hooks/useGitWatcher.test.ts b/site/src/pages/AgentsPage/hooks/useGitWatcher.test.ts index 119ea08ab2..5cafc47b05 100644 --- a/site/src/pages/AgentsPage/hooks/useGitWatcher.test.ts +++ b/site/src/pages/AgentsPage/hooks/useGitWatcher.test.ts @@ -169,7 +169,8 @@ describe("useGitWatcher", () => { expect(result.current.isConnected).toBe(false); // Simulate the browser firing the close event after - // socket.close() — the disposedRef guard must prevent + // socket.close(). The dispose guard inside + // createReconnectingWebSocket must prevent // the reconnect handler from scheduling a new attempt. mockWatchChatGit.mockClear(); act(() => socket.simulateClose()); @@ -403,7 +404,8 @@ describe("useGitWatcher", () => { expect(socket.close).toHaveBeenCalledTimes(1); // Simulate the browser firing the close event after - // socket.close() — the disposedRef guard must prevent + // socket.close(). The dispose guard inside + // createReconnectingWebSocket must prevent // the reconnect handler from scheduling a new attempt. mockWatchChatGit.mockClear(); act(() => socket.simulateClose()); diff --git a/site/src/pages/AgentsPage/hooks/useGitWatcher.ts b/site/src/pages/AgentsPage/hooks/useGitWatcher.ts index 1615ed2ed7..0db57e0629 100644 --- a/site/src/pages/AgentsPage/hooks/useGitWatcher.ts +++ b/site/src/pages/AgentsPage/hooks/useGitWatcher.ts @@ -6,6 +6,7 @@ import type { WorkspaceAgentRepoChanges, WorkspaceAgentStatus, } from "#/api/typesGenerated"; +import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; // Compile-time guard: ensures the bailout comparison in setRepositories // covers every data field. If WorkspaceAgentRepoChanges gains a new @@ -34,8 +35,6 @@ interface UseGitWatcherResult { refresh: () => boolean; } -const MAX_BACKOFF_MS = 30_000; - export function useGitWatcher({ chatId, agentStatus, @@ -46,10 +45,6 @@ export function useGitWatcher({ const [isConnected, setIsConnected] = useState(false); const socketRef = useRef(null); - const reconnectAttemptRef = useRef(0); - const reconnectTimerRef = useRef | null>(null); - // Track whether we've been disposed to avoid reconnecting after unmount. - const disposedRef = useRef(false); const sendMessage = (msg: WorkspaceAgentGitClientMessage): boolean => { const socket = socketRef.current; @@ -69,107 +64,82 @@ export function useGitWatcher({ return; } - disposedRef.current = false; + const activeChatId = chatId; - function connect() { - if (disposedRef.current) { - return; - } + const dispose = createReconnectingWebSocket({ + connect() { + const socket = watchChatGit(activeChatId); + socketRef.current = socket; - const socket = watchChatGit(chatId!); - socketRef.current = socket; + socket.addEventListener("message", (event) => { + // Ignore messages from superseded connections. + if (socketRef.current !== socket) { + return; + } + let data: WorkspaceAgentGitServerMessage; + try { + data = JSON.parse( + String((event as MessageEvent).data), + ) as WorkspaceAgentGitServerMessage; + } catch { + // Ignore unparsable messages. + return; + } - socket.addEventListener("open", () => { - // Ignore open events from superseded connections. - if (socketRef.current !== socket) { - return; - } - setIsConnected(true); - reconnectAttemptRef.current = 0; - }); - - socket.addEventListener("message", (event) => { - // Ignore messages from superseded connections. - if (socketRef.current !== socket) { - return; - } - let data: WorkspaceAgentGitServerMessage; - try { - data = JSON.parse( - String(event.data), - ) as WorkspaceAgentGitServerMessage; - } catch { - // Ignore unparsable messages. - return; - } - - if (data.type === "changes" && data.repositories) { - setRepositories((prev) => { - let changed = false; - const next = new Map(prev); - for (const repo of data.repositories!) { - if (repo.removed) { - if (next.has(repo.repo_root)) { - next.delete(repo.repo_root); - changed = true; - } - } else { - const existing = next.get(repo.repo_root); - if ( - !existing || - existing.branch !== repo.branch || - existing.remote_origin !== repo.remote_origin || - existing.unified_diff !== repo.unified_diff - ) { - next.set(repo.repo_root, repo); - changed = true; + if (data.type === "changes" && data.repositories) { + setRepositories((prev) => { + let changed = false; + const next = new Map(prev); + for (const repo of data.repositories!) { + if (repo.removed) { + if (next.has(repo.repo_root)) { + next.delete(repo.repo_root); + changed = true; + } + } else { + const existing = next.get(repo.repo_root); + if ( + !existing || + existing.branch !== repo.branch || + existing.remote_origin !== repo.remote_origin || + existing.unified_diff !== repo.unified_diff + ) { + next.set(repo.repo_root, repo); + changed = true; + } } } - } - return changed ? next : prev; - }); - } else if (data.type === "error") { - console.warn("[useGitWatcher] server error:", data.message); - } - }); + return changed ? next : prev; + }); + } else if (data.type === "error") { + console.warn("[useGitWatcher] server error:", data.message); + } + }); - // Note: WebSocket "error" events are always followed by a "close" - // event, so reconnection is handled here. - socket.addEventListener("close", () => { - // Ignore close events from superseded connections. - if (socketRef.current !== socket) { - return; - } + return socket; + }, + + onOpen() { + setIsConnected(true); + }, + + onDisconnect() { setIsConnected(false); socketRef.current = null; + }, - if (disposedRef.current) { - return; - } - - // Reconnect with exponential backoff. - const attempt = reconnectAttemptRef.current; - const delay = Math.min(1000 * 2 ** attempt, MAX_BACKOFF_MS); - reconnectAttemptRef.current = attempt + 1; - reconnectTimerRef.current = setTimeout(connect, delay); - }); - } - - connect(); + // 30s cap instead of the utility default 10s. The git + // endpoint may be slow to respond after a workspace wakes. + maxMs: 30_000, + }); return () => { - disposedRef.current = true; - if (reconnectTimerRef.current !== null) { - clearTimeout(reconnectTimerRef.current); - reconnectTimerRef.current = null; - } - if (socketRef.current) { - socketRef.current.close(); - socketRef.current = null; - } + // dispose() suppresses onDisconnect, so reset state + // explicitly. + dispose(); setIsConnected(false); setRepositories(new Map()); - reconnectAttemptRef.current = 0; + socketRef.current = null; }; }, [chatId, agentStatus]);