Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 134 additions & 0 deletions packages/cli/src/ui/hooks/useRewind.test.ts
Original file line number Diff line number Diff line change
@@ -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,
);
});
});
54 changes: 54 additions & 0 deletions packages/cli/src/ui/hooks/useRewind.ts
Original file line number Diff line number Diff line change
@@ -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<string | null>(
null,
);
const [confirmationStats, setConfirmationStats] =
useState<FileChangeStats | null>(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,
};
}
Loading