diff --git a/packages/cli/src/ui/hooks/useRewind.test.ts b/packages/cli/src/ui/hooks/useRewind.test.ts new file mode 100644 index 0000000000..7694dbd7a7 --- /dev/null +++ b/packages/cli/src/ui/hooks/useRewind.test.ts @@ -0,0 +1,134 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { act } from 'react'; +import { renderHook } from '../../test-utils/render.js'; +import { useRewind } from './useRewind.js'; +import * as rewindFileOps from '../utils/rewindFileOps.js'; +import type { FileChangeStats } from '../utils/rewindFileOps.js'; +import type { + ConversationRecord, + MessageRecord, +} from '@google/gemini-cli-core'; + +// Mock the dependency +vi.mock('../utils/rewindFileOps.js', () => ({ + calculateTurnStats: vi.fn(), + calculateRewindImpact: vi.fn(), +})); + +describe('useRewindLogic', () => { + const mockUserMessage: MessageRecord = { + id: 'msg-1', + type: 'user', + content: 'Hello', + timestamp: new Date(1000).toISOString(), + }; + + const mockModelMessage: MessageRecord = { + id: 'msg-2', + type: 'gemini', + content: 'Hi there', + timestamp: new Date(1001).toISOString(), + }; + + const mockConversation: ConversationRecord = { + sessionId: 'conv-1', + projectHash: 'hash-1', + startTime: new Date(1000).toISOString(), + lastUpdated: new Date(1001).toISOString(), + messages: [mockUserMessage, mockModelMessage], + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('should initialize with no selection', () => { + const { result } = renderHook(() => useRewind(mockConversation)); + + expect(result.current.selectedMessageId).toBeNull(); + expect(result.current.confirmationStats).toBeNull(); + }); + + it('should update state when a message is selected', () => { + const mockStats: FileChangeStats = { + fileCount: 1, + addedLines: 5, + removedLines: 0, + }; + vi.mocked(rewindFileOps.calculateRewindImpact).mockReturnValue(mockStats); + + const { result } = renderHook(() => useRewind(mockConversation)); + + act(() => { + result.current.selectMessage('msg-1'); + }); + + expect(result.current.selectedMessageId).toBe('msg-1'); + expect(result.current.confirmationStats).toEqual(mockStats); + expect(rewindFileOps.calculateRewindImpact).toHaveBeenCalledWith( + mockConversation, + mockUserMessage, + ); + }); + + it('should not update state if selected message is not found', () => { + const { result } = renderHook(() => useRewind(mockConversation)); + + act(() => { + result.current.selectMessage('non-existent-id'); + }); + + expect(result.current.selectedMessageId).toBeNull(); + expect(result.current.confirmationStats).toBeNull(); + }); + + it('should clear selection correctly', () => { + const mockStats: FileChangeStats = { + fileCount: 1, + addedLines: 5, + removedLines: 0, + }; + vi.mocked(rewindFileOps.calculateRewindImpact).mockReturnValue(mockStats); + + const { result } = renderHook(() => useRewind(mockConversation)); + + // Select first + act(() => { + result.current.selectMessage('msg-1'); + }); + expect(result.current.selectedMessageId).toBe('msg-1'); + + // Then clear + act(() => { + result.current.clearSelection(); + }); + + expect(result.current.selectedMessageId).toBeNull(); + expect(result.current.confirmationStats).toBeNull(); + }); + + it('should proxy getStats call to utility function', () => { + const mockStats: FileChangeStats = { + fileCount: 2, + addedLines: 10, + removedLines: 2, + }; + vi.mocked(rewindFileOps.calculateTurnStats).mockReturnValue(mockStats); + + const { result } = renderHook(() => useRewind(mockConversation)); + + const stats = result.current.getStats(mockUserMessage); + + expect(stats).toEqual(mockStats); + expect(rewindFileOps.calculateTurnStats).toHaveBeenCalledWith( + mockConversation, + mockUserMessage, + ); + }); +}); diff --git a/packages/cli/src/ui/hooks/useRewind.ts b/packages/cli/src/ui/hooks/useRewind.ts new file mode 100644 index 0000000000..66074be694 --- /dev/null +++ b/packages/cli/src/ui/hooks/useRewind.ts @@ -0,0 +1,54 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useState, useCallback } from 'react'; +import type { + ConversationRecord, + MessageRecord, +} from '@google/gemini-cli-core'; +import { + calculateTurnStats, + calculateRewindImpact, + type FileChangeStats, +} from '../utils/rewindFileOps.js'; + +export function useRewind(conversation: ConversationRecord) { + const [selectedMessageId, setSelectedMessageId] = useState( + null, + ); + const [confirmationStats, setConfirmationStats] = + useState(null); + + const getStats = useCallback( + (userMessage: MessageRecord) => + calculateTurnStats(conversation, userMessage), + [conversation], + ); + + const selectMessage = useCallback( + (messageId: string) => { + const msg = conversation.messages.find((m) => m.id === messageId); + if (msg) { + setSelectedMessageId(messageId); + setConfirmationStats(calculateRewindImpact(conversation, msg)); + } + }, + [conversation], + ); + + const clearSelection = useCallback(() => { + setSelectedMessageId(null); + setConfirmationStats(null); + }, []); + + return { + selectedMessageId, + getStats, + confirmationStats, + selectMessage, + clearSelection, + }; +}