Edit: Improve Rejection Tracker. Handle file deletions and formatters (…

…#4097)

vscode/src/non-stop/FixupController.ts

@@ -32,6 +32,7 @@ import { FixupFileObserver } from './FixupFileObserver'
 import { FixupScheduler } from './FixupScheduler'
 import { FixupTask, type FixupTaskID, type FixupTelemetryMetadata } from './FixupTask'
 import { type Diff, computeDiff } from './diff'
+import { trackRejection } from './rejection-tracker'
 import type { FixupActor, FixupFileCollection, FixupIdleTaskRunner, FixupTextChanged } from './roles'
 import { CodyTaskState, getMinimumDistanceToRangeBoundary } from './utils'
 
@@ -533,60 +534,19 @@ export class FixupController
             })
         }
 
-        /**
-         * Tracks when a user clicks "Undo" in the Edit codelens.
-         * This is important as VS Code doesn't let us easily differentiate between
-         * document changes made by specific commands.
-         *
-         * This logic ensures we can still mark as task as rejected if a user clicks "Undo".
-         */
-        const commandUndoListener = this.undoCommandEvent.event(id => {
-            if (id !== task.id) {
-                return
-            }
-
-            // Immediately dispose of the rejectionListener, otherwise this will also run
-            // and mark the "Undo" change here as an "acccepted" change made by the user.
-            rejectionListener.dispose()
-            commandUndoListener.dispose()
-
-            // If a user manually clicked "Undo", we can be confident that they reject the fixup.
-            logAcceptance('rejected')
-        })
-        let undoCount = 0
-        /**
-         * Tracks the rejection of a Fixup task via the users' next action.
-         * As in, if the user immediately undos the change via the system undo command,
-         * or if they persist to make new edits to the file.
-         *
-         * Will listen for changes to the text document and tracks whether the Edit changes were undone or redone.
-         * When a change is made, it logs telemetry about whether the change was rejected or accepted.
-         */
-        const rejectionListener = vscode.workspace.onDidChangeTextDocument(event => {
-            if (event.document.uri !== document.uri || event.contentChanges.length === 0) {
-                // Irrelevant change, ignore
-                return
-            }
-
-            if (event.reason === vscode.TextDocumentChangeReason.Undo) {
-                // Set state, but don't fire telemetry yet as the user could still "Redo".
-                undoCount += 1
-                return
-            }
-
-            if (event.reason === vscode.TextDocumentChangeReason.Redo) {
-                // User re-did the change, so reset state
-                undoCount = Math.max(0, undoCount - 1)
-                return
+        trackRejection(
+            document,
+            vscode.workspace,
+            {
+                onAccepted: () => logAcceptance('accepted'),
+                onRejected: () => logAcceptance('rejected'),
+            },
+            {
+                id: task.id,
+                intent: task.intent,
+                undoEvent: this.undoCommandEvent,
             }
-
-            // User has made a change, we can now fire our stored state as to if the change was undone or not
-            logAcceptance(undoCount > 0 ? 'rejected' : 'accepted')
-
-            // We no longer need to track this change, so dispose of our listeners
-            rejectionListener.dispose()
-            commandUndoListener.dispose()
-        })
+        )
     }
 
     private async streamTask(task: FixupTask, state: 'streaming' | 'complete'): Promise<void> {

vscode/src/non-stop/rejection-tracker.test.ts

@@ -0,0 +1,285 @@
+import { testFileUri } from '@sourcegraph/cody-shared'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import * as vscode from 'vscode'
+import { document } from '../completions/test-helpers'
+import { range } from '../testutils/textDocument'
+import { trackRejection } from './rejection-tracker'
+
+const singleCharacterChange: vscode.TextDocumentContentChangeEvent = {
+    range: range(0, 0, 0, 0),
+    text: 'B',
+    rangeLength: 0,
+    rangeOffset: 0,
+}
+
+const multiCharacterChange: vscode.TextDocumentContentChangeEvent = {
+    ...singleCharacterChange,
+    text: 'Beans', // multi-character change, could be from a command or formatter
+}
+
+const whitespaceChange: vscode.TextDocumentContentChangeEvent = {
+    ...singleCharacterChange,
+    text: ' ', // whitespace only change
+}
+
+describe('trackRejection', () => {
+    const onAcceptedSpy = vi.fn()
+    const onRejectedSpy = vi.fn()
+
+    // Mock workspace APIs to trigger document changes
+    let onDidChangeTextDocument: (event: vscode.TextDocumentChangeEvent) => void
+    let onDidDeleteFiles: (event: vscode.FileDeleteEvent) => void
+    let onDidCloseTextDocument: (event: vscode.TextDocument) => void
+
+    const mockWorkspace: Pick<
+        typeof vscode.workspace,
+        | 'onDidChangeTextDocument'
+        | 'onDidDeleteFiles'
+        | 'onDidCloseTextDocument'
+        | 'onDidSaveTextDocument'
+    > = {
+        onDidChangeTextDocument(listener) {
+            onDidChangeTextDocument = listener
+            return { dispose: () => {} }
+        },
+        onDidDeleteFiles(listener) {
+            onDidDeleteFiles = listener
+            return { dispose: () => {} }
+        },
+        onDidCloseTextDocument(listener) {
+            onDidCloseTextDocument = listener
+            return { dispose: () => {} }
+        },
+        onDidSaveTextDocument(listener) {
+            return { dispose: () => {} }
+        },
+    }
+    const mockDocument = document('Hello, world!')
+    const mockTask = {
+        id: 'task-id',
+        intent: 'edit',
+        undoEvent: new vscode.EventEmitter(),
+    }
+
+    afterEach(() => {
+        vi.resetAllMocks()
+    })
+
+    describe('file rejections', () => {
+        it('should call onRejected when document is deleted', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+
+            onDidDeleteFiles({ files: [mockDocument.uri] })
+
+            // File deleted, mark task as rejected
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).toHaveBeenCalled()
+        })
+
+        it('should not onRejected when a different document is deleted', () => {
+            const otherDoc = document('World, Hello!', undefined, testFileUri('other.ts').toString())
+
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+
+            onDidDeleteFiles({ files: [otherDoc.uri] })
+
+            // Other file deleted, do nothing
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+        })
+
+        it('for the test intent should call onRejected when document is closed without being saved', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                { ...mockTask, intent: 'test' } as any
+            )
+
+            Object.assign(mockDocument, { isUntitled: true })
+            onDidCloseTextDocument(mockDocument)
+
+            // File closed without being saved, mark task as rejected
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).toHaveBeenCalled()
+        })
+
+        it('for the test intent should not call onRejected when document is closed after being saved', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                { ...mockTask, intent: 'test' } as any
+            )
+
+            onDidCloseTextDocument(mockDocument)
+
+            // File closed without being saved, mark task as rejected
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).toHaveBeenCalled()
+        })
+    })
+
+    describe('in-file rejections', () => {
+        it('should call onRejected when a task is undone and another change is made', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: vscode.TextDocumentChangeReason.Undo,
+            })
+
+            // Not called yet, we've only undone and redo is available
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: undefined,
+            })
+
+            // Redo no longer available, accepted is called
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).toHaveBeenCalled()
+        })
+
+        it('should call onRejected when the task is finally rejected', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [multiCharacterChange],
+                reason: undefined,
+            })
+
+            // Mulit character change so we don't accept or reject
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: vscode.TextDocumentChangeReason.Undo,
+            })
+
+            // We undid the multi-character change, not this task, do nothing
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: vscode.TextDocumentChangeReason.Undo,
+            })
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: undefined,
+            })
+
+            // Finally we have undone the task and made another edit
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).toHaveBeenCalled()
+        })
+    })
+
+    describe('in-file acceptance', () => {
+        it('should call onAccepted when document is changed with valid characters', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: undefined,
+            })
+
+            // Single character change, accept task
+            expect(onAcceptedSpy).toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+        })
+
+        it('should not call onAccepted immediately for multi-character changes, but after the next change', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [multiCharacterChange],
+                reason: undefined,
+            })
+
+            // Multi character change, so we do nothing as it could be from a
+            // command or formatter
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: undefined,
+            })
+
+            // Finally we make a single character change, we can accept
+            expect(onAcceptedSpy).toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+        })
+
+        it('should not call onAccepted immediately for white-space only changes, but after the next change', () => {
+            trackRejection(
+                mockDocument,
+                mockWorkspace,
+                { onAccepted: onAcceptedSpy, onRejected: onRejectedSpy },
+                mockTask as any
+            )
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [whitespaceChange],
+                reason: undefined,
+            })
+
+            // Single character change, but only whitespace, not enough of a signal
+            // for acceptance, do nothing.
+            expect(onAcceptedSpy).not.toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+
+            onDidChangeTextDocument({
+                document: mockDocument,
+                contentChanges: [singleCharacterChange],
+                reason: undefined,
+            })
+
+            // Finally a single character change, accept
+            expect(onAcceptedSpy).toHaveBeenCalled()
+            expect(onRejectedSpy).not.toHaveBeenCalled()
+        })
+    })
+})