diff --git a/src/main/git/status.test.ts b/src/main/git/status.test.ts new file mode 100644 index 00000000000..0796e16e045 --- /dev/null +++ b/src/main/git/status.test.ts @@ -0,0 +1,78 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { execFileAsyncMock, readFileMock, rmMock } = vi.hoisted(() => ({ + execFileAsyncMock: vi.fn(), + readFileMock: vi.fn(), + rmMock: vi.fn() +})) + +vi.mock('util', async () => { + const actual = await vi.importActual('util') + return { + ...actual, + promisify: vi.fn(() => execFileAsyncMock) + } +}) + +vi.mock('fs/promises', () => ({ + readFile: readFileMock, + rm: rmMock +})) + +import { discardChanges } from './status' + +describe('discardChanges', () => { + beforeEach(() => { + execFileAsyncMock.mockReset() + readFileMock.mockReset() + rmMock.mockReset() + }) + + it('restores tracked files from HEAD', async () => { + execFileAsyncMock.mockResolvedValueOnce({ stdout: 'src/file.ts\n' }) + execFileAsyncMock.mockResolvedValueOnce({ stdout: '' }) + + await discardChanges('/repo', 'src/file.ts') + + expect(execFileAsyncMock).toHaveBeenNthCalledWith( + 1, + 'git', + ['ls-files', '--error-unmatch', '--', 'src/file.ts'], + { + cwd: '/repo', + encoding: 'utf-8' + } + ) + expect(execFileAsyncMock).toHaveBeenNthCalledWith( + 2, + 'git', + ['restore', '--worktree', '--source=HEAD', '--', 'src/file.ts'], + { + cwd: '/repo', + encoding: 'utf-8' + } + ) + expect(rmMock).not.toHaveBeenCalled() + }) + + it('removes untracked files from disk', async () => { + execFileAsyncMock.mockRejectedValueOnce(new Error('not tracked')) + + await discardChanges('/repo', 'src/new-file.ts') + + expect(execFileAsyncMock).toHaveBeenCalledTimes(1) + expect(rmMock).toHaveBeenCalledWith('/repo/src/new-file.ts', { + force: true, + recursive: true + }) + }) + + it('rejects paths that traverse outside the worktree', async () => { + await expect(discardChanges('/repo', '../../etc/passwd')).rejects.toThrow( + 'resolves outside the worktree' + ) + + expect(execFileAsyncMock).not.toHaveBeenCalled() + expect(rmMock).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/git/status.ts b/src/main/git/status.ts index 87ccf2ae5cd..ca34989e4e4 100644 --- a/src/main/git/status.ts +++ b/src/main/git/status.ts @@ -1,7 +1,7 @@ import { execFile } from 'child_process' -import { readFile } from 'fs/promises' +import { readFile, rm } from 'fs/promises' import { promisify } from 'util' -import { join } from 'path' +import { join, resolve } from 'path' import type { GitStatusEntry, GitFileStatus, GitDiffResult } from '../../shared/types' const execFileAsync = promisify(execFile) @@ -163,8 +163,27 @@ export async function unstageFile(worktreePath: string, filePath: string): Promi * Discard working tree changes for a file. */ export async function discardChanges(worktreePath: string, filePath: string): Promise { - await execFileAsync('git', ['checkout', '--', filePath], { - cwd: worktreePath, - encoding: 'utf-8' - }) + const resolvedWorktree = resolve(worktreePath) + const resolvedTarget = resolve(worktreePath, filePath) + if (!resolvedTarget.startsWith(`${resolvedWorktree}/`)) { + throw new Error(`Path "${filePath}" resolves outside the worktree`) + } + + let tracked = false + try { + await execFileAsync('git', ['ls-files', '--error-unmatch', '--', filePath], { + cwd: worktreePath, + encoding: 'utf-8' + }) + tracked = true + } catch { + // File is not tracked by git + } + + await (tracked + ? execFileAsync('git', ['restore', '--worktree', '--source=HEAD', '--', filePath], { + cwd: worktreePath, + encoding: 'utf-8' + }) + : rm(resolvedTarget, { force: true, recursive: true })) } diff --git a/src/renderer/src/components/right-sidebar/SourceControl.tsx b/src/renderer/src/components/right-sidebar/SourceControl.tsx index 51bfb68a936..7e9c3d09b0c 100644 --- a/src/renderer/src/components/right-sidebar/SourceControl.tsx +++ b/src/renderer/src/components/right-sidebar/SourceControl.tsx @@ -15,6 +15,7 @@ import { useAppStore } from '@/store' import { detectLanguage } from '@/lib/language-detect' import { cn } from '@/lib/utils' import type { GitStatusEntry, GitStagingArea } from '../../../../shared/types' +import { getSourceControlActions } from './source-control-actions' const STATUS_ICONS: Record> = { modified: FileEdit, @@ -252,6 +253,7 @@ export default function SourceControl(): React.JSX.Element { const dirPath = entry.path.includes('/') ? entry.path.slice(0, entry.path.lastIndexOf('/')) : '' + const actions = getSourceControlActions(area) return (
- {area === 'unstaged' || area === 'untracked' ? ( - <> - {area === 'unstaged' && ( - { - e.stopPropagation() - void handleDiscard(entry.path) - }} - /> - )} - { - e.stopPropagation() - void handleStage(entry.path) - }} - /> - - ) : ( + {actions.includes('discard') && ( + { + e.stopPropagation() + if (area === 'untracked') { + if ( + !window.confirm( + `Delete untracked file "${entry.path}"? This cannot be undone.` + ) + ) { + return + } + } + void handleDiscard(entry.path) + }} + /> + )} + {actions.includes('stage') && ( + { + e.stopPropagation() + void handleStage(entry.path) + }} + /> + )} + {actions.includes('unstage') && ( { + it('shows discard and stage actions for untracked files', () => { + expect(getSourceControlActions('untracked')).toEqual(['discard', 'stage']) + }) + + it('shows discard and stage actions for unstaged files', () => { + expect(getSourceControlActions('unstaged')).toEqual(['discard', 'stage']) + }) + + it('shows unstage action for staged files', () => { + expect(getSourceControlActions('staged')).toEqual(['unstage']) + }) +}) diff --git a/src/renderer/src/components/right-sidebar/source-control-actions.ts b/src/renderer/src/components/right-sidebar/source-control-actions.ts new file mode 100644 index 00000000000..d79bd53595e --- /dev/null +++ b/src/renderer/src/components/right-sidebar/source-control-actions.ts @@ -0,0 +1,15 @@ +import type { GitStagingArea } from '../../../../shared/types' + +export type SourceControlAction = 'discard' | 'stage' | 'unstage' + +export function getSourceControlActions(area: GitStagingArea): SourceControlAction[] { + switch (area) { + case 'staged': + return ['unstage'] + case 'unstaged': + case 'untracked': + return ['discard', 'stage'] + default: + return [] + } +}