diff --git a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingCheckpointTimelineImpl.ts b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingCheckpointTimelineImpl.ts index a904ab1304d55..378963dd20278 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingCheckpointTimelineImpl.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingCheckpointTimelineImpl.ts @@ -123,16 +123,23 @@ export class ChatEditingCheckpointTimelineImpl implements IChatEditingCheckpoint { equalsFn: (a, b) => arraysEqual(a, b, objectsEqual) }, reader => { const currentEpoch = this._currentEpoch.read(reader); - const checkpoints = this._checkpoints.read(reader); + const operations = this._operations.read(reader); + const firstUnappliedOperationIdx = operations.findIndex(op => op.epoch >= currentEpoch); + if (firstUnappliedOperationIdx === -1) { + return []; + } + const lastAppliedOperation = firstUnappliedOperationIdx > 0 ? operations[firstUnappliedOperationIdx - 1].epoch : 0; + const checkpoints = this._checkpoints.read(reader); const disablement = new Map(); - // Go through the checkpoints and disable any that are after our current epoch. + + // Go through the checkpoints and disable any until the one that contains the last applied operation. // Subtle: the request will first make a checkpoint with an 'undefined' undo // stop, and in this loop we'll "automatically" disable the entire request when // we reach that checkpoint. for (let i = checkpoints.length - 1; i >= 0; i--) { const { undoStopId, requestId, epoch } = checkpoints[i]; - if (epoch < currentEpoch) { + if (epoch < lastAppliedOperation) { break; } @@ -165,19 +172,23 @@ export class ChatEditingCheckpointTimelineImpl implements IChatEditingCheckpoint return existing.checkpointId; } + const { checkpoints, operations } = this._getVisibleOperationsAndCheckpoints(); const checkpointId = generateUuid(); - const checkpoint: ICheckpoint = { + const epoch = this.incrementEpoch(); + + checkpoints.push({ checkpointId, requestId, undoStopId, - epoch: this.incrementEpoch(), + epoch, label, description - }; + }); transaction(tx => { - this._checkpoints.set([...existingCheckpoints, checkpoint], tx); - this._currentEpoch.set(checkpoint.epoch + 1, tx); + this._checkpoints.set(checkpoints, tx); + this._operations.set(operations, tx); + this._currentEpoch.set(epoch + 1, tx); }); return checkpointId; @@ -203,26 +214,36 @@ export class ChatEditingCheckpointTimelineImpl implements IChatEditingCheckpoint throw new Error(`Checkpoint ${checkpointId} not found`); } - return this._navigateToEpoch(targetCheckpoint.epoch + 1); + if (targetCheckpoint.undoStopId === undefined) { + // If we're navigating to the start of a request, we want to restore the file + // to whatever baseline we captured, _not_ the result state from the prior request + // because there may have been user changes in the meantime. But we still want + // to set the epoch marking that checkpoint as having been undone (the second + // arg below) so that disablement works and so it's discarded if appropriate later. + return this._navigateToEpoch(targetCheckpoint.epoch + 1, targetCheckpoint.epoch); + } else { + return this._navigateToEpoch(targetCheckpoint.epoch + 1); + } + } public getContentURIAtStop(requestId: string, fileURI: URI, stopId: string | undefined): URI { return ChatEditingSnapshotTextModelContentProvider.getSnapshotFileURI(this.chatSessionId, requestId, stopId, fileURI.path); } - private async _navigateToEpoch(targetEpoch: number): Promise { + private async _navigateToEpoch(restoreToEpoch: number, navigateToEpoch = restoreToEpoch): Promise { const currentEpoch = this._currentEpoch.get(); - if (currentEpoch === targetEpoch) { + if (currentEpoch === restoreToEpoch) { return; // Already at target epoch } - const urisToRestore = await this._applyFileSystemOperations(currentEpoch, targetEpoch); + const urisToRestore = await this._applyFileSystemOperations(currentEpoch, restoreToEpoch); // Reconstruct content for files affected by operations in the range - await this._reconstructAllFileContents(targetEpoch, urisToRestore); + await this._reconstructAllFileContents(restoreToEpoch, urisToRestore); // Update current epoch - this._currentEpoch.set(targetEpoch, undefined); + this._currentEpoch.set(navigateToEpoch, undefined); } private _getCheckpoint(checkpointId: string): ICheckpoint | undefined { @@ -234,26 +255,31 @@ export class ChatEditingCheckpointTimelineImpl implements IChatEditingCheckpoint } public recordFileOperation(operation: FileOperation): void { - const currentEpoch = this._currentEpoch.get(); - const currentCheckpoints = this._checkpoints.get(); - - const operations = this._operations.get(); - const insertAt = findLastIdx(operations, op => op.epoch < currentEpoch); - operations[insertAt + 1] = operation; - operations.length = insertAt + 2; // Truncate any operations beyond this point + const { currentEpoch, checkpoints, operations } = this._getVisibleOperationsAndCheckpoints(); + if (operation.epoch < currentEpoch) { + throw new Error(`Cannot record operation at epoch ${operation.epoch} when current epoch is ${currentEpoch}`); + } - // If we undid some operations and are dropping them out of history, also remove - // any associated checkpoints. - const newCheckpoints = currentCheckpoints.filter(c => c.epoch < currentEpoch); + operations.push(operation); transaction(tx => { - if (newCheckpoints.length !== currentCheckpoints.length) { - this._checkpoints.set(newCheckpoints, tx); - } - this._currentEpoch.set(operation.epoch + 1, tx); + this._checkpoints.set(checkpoints, tx); this._operations.set(operations, tx); + this._currentEpoch.set(operation.epoch + 1, tx); }); } + private _getVisibleOperationsAndCheckpoints() { + const currentEpoch = this._currentEpoch.get(); + const checkpoints = this._checkpoints.get(); + const operations = this._operations.get(); + + return { + currentEpoch, + checkpoints: checkpoints.filter(c => c.epoch < currentEpoch), + operations: operations.filter(op => op.epoch < currentEpoch) + }; + } + public recordFileBaseline(baseline: IFileBaseline): void { const key = this._getBaselineKey(baseline.uri, baseline.requestId); this._fileBaselines.set(key, baseline); diff --git a/src/vs/workbench/contrib/chat/test/browser/chatEditingCheckpointTimeline.test.ts b/src/vs/workbench/contrib/chat/test/browser/chatEditingCheckpointTimeline.test.ts index 25fc09cf6236e..28c9069aaff0a 100644 --- a/src/vs/workbench/contrib/chat/test/browser/chatEditingCheckpointTimeline.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/chatEditingCheckpointTimeline.test.ts @@ -6,7 +6,7 @@ import assert from 'assert'; import { DisposableStore } from '../../../../../base/common/lifecycle.js'; import { ResourceMap } from '../../../../../base/common/map.js'; -import { autorun, transaction } from '../../../../../base/common/observable.js'; +import { transaction } from '../../../../../base/common/observable.js'; import { URI } from '../../../../../base/common/uri.js'; import { upcastPartial } from '../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; @@ -422,25 +422,56 @@ suite('ChatEditingCheckpointTimeline', function () { }); test('requestDisablement tracks disabled requests', async function () { - const disabledRequests: string[] = []; - store.add(autorun(reader => { - const disabled = timeline.requestDisablement.read(reader); - disabledRequests.length = 0; - disabledRequests.push(...disabled.map(d => d.requestId)); - })); + const uri = URI.parse('file:///test.txt'); - // Create checkpoints for multiple requests timeline.createCheckpoint('req1', undefined, 'Start req1'); + timeline.recordFileOperation(createFileCreateOperation(uri, 'req1', timeline.incrementEpoch(), 'a')); + timeline.createCheckpoint('req1', 'stop1', 'Stop req1'); + timeline.recordFileOperation(createTextEditOperation(uri, 'req1', timeline.incrementEpoch(), [{ range: new Range(1, 1, 1, 2), text: 'b' }])); + timeline.createCheckpoint('req2', undefined, 'Start req2'); + timeline.recordFileOperation(createTextEditOperation(uri, 'req2', timeline.incrementEpoch(), [{ range: new Range(1, 1, 1, 2), text: 'c' }])); - // Navigate to req1 stop1 - const req1StopId = timeline.getCheckpointIdForRequest('req1', 'stop1')!; - await timeline.navigateToCheckpoint(req1StopId); + // Undo sequence: + assert.deepStrictEqual(timeline.requestDisablement.get(), []); + + await timeline.undoToLastCheckpoint(); + assert.strictEqual(fileContents.get(uri), 'b'); + assert.deepStrictEqual(timeline.requestDisablement.get(), [ + { requestId: 'req2', afterUndoStop: undefined }, + ]); + + await timeline.undoToLastCheckpoint(); + assert.strictEqual(fileContents.get(uri), 'a'); + assert.deepStrictEqual(timeline.requestDisablement.get(), [ + { requestId: 'req2', afterUndoStop: undefined }, + { requestId: 'req1', afterUndoStop: 'stop1' }, + ]); - // req2 should be disabled as we're before it. req1 is also disabled because - // we're in the past relative to the current tip - assert.ok(disabledRequests.includes('req2')); + await timeline.undoToLastCheckpoint(); + assert.strictEqual(fileContents.get(uri), undefined); + assert.deepStrictEqual(timeline.requestDisablement.get(), [ + { requestId: 'req2', afterUndoStop: undefined }, + { requestId: 'req1', afterUndoStop: undefined }, + ]); + + // Redo sequence: + await timeline.redoToNextCheckpoint(); + assert.strictEqual(fileContents.get(uri), 'a'); + assert.deepStrictEqual(timeline.requestDisablement.get(), [ + { requestId: 'req2', afterUndoStop: undefined }, + { requestId: 'req1', afterUndoStop: 'stop1' }, + ]); + + await timeline.redoToNextCheckpoint(); + assert.strictEqual(fileContents.get(uri), 'b'); + assert.deepStrictEqual(timeline.requestDisablement.get(), [ + { requestId: 'req2', afterUndoStop: undefined }, + ]); + + await timeline.redoToNextCheckpoint(); + assert.strictEqual(fileContents.get(uri), 'c'); }); test('persistence - save and restore state', function () { @@ -597,7 +628,7 @@ suite('ChatEditingCheckpointTimeline', function () { // Verify we're at the start of req1, which has epoch 2 (0 = initial, 1 = baseline, 2 = start checkpoint) const state = timeline.getStateForPersistence(); - assert.strictEqual(state.currentEpoch, 3); // Should be at the "Start req1" checkpoint epoch + assert.strictEqual(state.currentEpoch, 2); // Should be at the "Start req1" checkpoint epoch }); test('operations use incrementing epochs', function () {