diff --git a/site/src/pages/AgentsPage/DiffViewer.tsx b/site/src/pages/AgentsPage/DiffViewer.tsx index 064c5aad7a..5ff92bee36 100644 --- a/site/src/pages/AgentsPage/DiffViewer.tsx +++ b/site/src/pages/AgentsPage/DiffViewer.tsx @@ -402,6 +402,24 @@ const LazyFileDiff = memo<{ /> ); }, + (prev, next) => { + if ( + prev.fileDiff !== next.fileDiff || + prev.options !== next.options || + prev.lineAnnotations !== next.lineAnnotations + ) { + return false; + } + // When neither the previous nor next props include line + // annotations, a new `renderAnnotation` reference is + // irrelevant because there is nothing to render. Skip the + // re-render so that unrelated state changes (e.g. + // activeCommentBox) don't bust memo for every file. + if (prev.renderAnnotation !== next.renderAnnotation) { + return !prev.lineAnnotations && !next.lineAnnotations; + } + return true; + }, ); // ------------------------------------------------------------------- diff --git a/site/src/pages/AgentsPage/useGitWatcher.test.ts b/site/src/pages/AgentsPage/useGitWatcher.test.ts index 5a8f3a5612..69954cdc7c 100644 --- a/site/src/pages/AgentsPage/useGitWatcher.test.ts +++ b/site/src/pages/AgentsPage/useGitWatcher.test.ts @@ -425,4 +425,194 @@ describe("useGitWatcher", () => { }); expect(result.current.repositories.has("/home/user/project-x")).toBe(true); }); + + it("stale close event after re-mount does not create duplicate connections", () => { + vi.useFakeTimers(); + + try { + const socket1 = createMockSocket(); + + const { result, rerender } = renderHook( + ({ chatId }: { chatId: string | undefined }) => + useGitWatcher({ chatId, agentStatus: "connected" }), + { initialProps: { chatId: "chat-aaa" as string | undefined } }, + ); + + act(() => socket1.simulateOpen()); + expect(mockWatchChatGit).toHaveBeenCalledTimes(1); + + // Prepare socket2 for the re-mount triggered by chatId change. + const socket2 = createMockSocket(); + rerender({ chatId: "chat-bbb" }); + + expect(socket1.close).toHaveBeenCalled(); + expect(mockWatchChatGit).toHaveBeenCalledTimes(2); + + // Simulate socket1's close event arriving late (stale). + // This must NOT clobber socketRef or schedule a reconnect. + act(() => socket1.simulateClose()); + + expect(mockWatchChatGit).toHaveBeenCalledTimes(2); + + // Advance timers to prove no reconnect was scheduled. + act(() => vi.advanceTimersByTime(60_000)); + expect(mockWatchChatGit).toHaveBeenCalledTimes(2); + + // socket2 should still work: open sets isConnected, + // messages update repositories. + act(() => socket2.simulateOpen()); + expect(result.current.isConnected).toBe(true); + + act(() => { + socket2.simulateMessage({ + type: "changes", + repositories: [{ repo_root: "/repo", branch: "main" }], + }); + }); + expect(result.current.repositories.size).toBe(1); + } finally { + vi.useRealTimers(); + } + }); + + it("preserves reference on duplicate messages", () => { + const socket = createMockSocket(); + + const { result } = renderHook(() => + useGitWatcher({ chatId: "chat-123", agentStatus: "connected" }), + ); + + act(() => socket.simulateOpen()); + + const message = { + type: "changes" as const, + repositories: [ + { + repo_root: "/repo", + branch: "main", + unified_diff: "diff1", + }, + ], + }; + + act(() => socket.simulateMessage(message)); + const ref1 = result.current.repositories; + expect(ref1.size).toBe(1); + + // Sending the exact same data should not produce a new reference. + act(() => socket.simulateMessage(message)); + expect(result.current.repositories).toBe(ref1); + }); + + it("single field change triggers update", () => { + const socket = createMockSocket(); + + const { result } = renderHook(() => + useGitWatcher({ chatId: "chat-123", agentStatus: "connected" }), + ); + + act(() => socket.simulateOpen()); + + const base = { + repo_root: "/repo", + branch: "main", + remote_origin: "git@github.com:org/repo.git", + unified_diff: "diff1", + }; + + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [base], + }); + }); + let ref = result.current.repositories; + expect(ref.get("/repo")?.branch).toBe("main"); + + // Changing only branch. + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [{ ...base, branch: "feature" }], + }); + }); + expect(result.current.repositories).not.toBe(ref); + expect(result.current.repositories.get("/repo")?.branch).toBe("feature"); + ref = result.current.repositories; + + // Changing only remote_origin. + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [ + { + ...base, + branch: "feature", + remote_origin: "https://github.com/org/repo", + }, + ], + }); + }); + expect(result.current.repositories).not.toBe(ref); + ref = result.current.repositories; + + // Changing only unified_diff. + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [ + { + ...base, + branch: "feature", + remote_origin: "https://github.com/org/repo", + unified_diff: "diff2", + }, + ], + }); + }); + expect(result.current.repositories).not.toBe(ref); + expect(result.current.repositories.get("/repo")?.unified_diff).toBe( + "diff2", + ); + }); + + it("removing unknown repo preserves reference", () => { + const socket = createMockSocket(); + + const { result } = renderHook(() => + useGitWatcher({ chatId: "chat-123", agentStatus: "connected" }), + ); + + act(() => socket.simulateOpen()); + + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [ + { + repo_root: "/repo", + branch: "main", + unified_diff: "diff1", + }, + ], + }); + }); + const ref1 = result.current.repositories; + expect(ref1.size).toBe(1); + + // Removing a repo_root that was never added should be a no-op. + act(() => { + socket.simulateMessage({ + type: "changes", + repositories: [ + { + repo_root: "/unknown", + branch: "", + removed: true, + }, + ], + }); + }); + expect(result.current.repositories).toBe(ref1); + }); }); diff --git a/site/src/pages/AgentsPage/useGitWatcher.ts b/site/src/pages/AgentsPage/useGitWatcher.ts index eda827c172..2ee117b1bd 100644 --- a/site/src/pages/AgentsPage/useGitWatcher.ts +++ b/site/src/pages/AgentsPage/useGitWatcher.ts @@ -7,6 +7,19 @@ import type { } from "api/typesGenerated"; import { useCallback, useEffect, useRef, useState } from "react"; +// Compile-time guard: ensures the bailout comparison in setRepositories +// covers every data field. If WorkspaceAgentRepoChanges gains a new +// field, this will error until the comparison is updated. +type _ComparedRepoFields = Omit< + WorkspaceAgentRepoChanges, + "repo_root" | "removed" +>; +const _repoFieldGuard: Record = { + branch: true, + remote_origin: true, + unified_diff: true, +}; + interface UseGitWatcherOptions { chatId: string | undefined; agentStatus: WorkspaceAgentStatus | undefined; @@ -65,11 +78,19 @@ export function useGitWatcher({ socketRef.current = socket; 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; + } try { const data = JSON.parse( String(event.data), @@ -111,6 +132,10 @@ export function useGitWatcher({ // 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; + } setIsConnected(false); socketRef.current = null;