edits: fix 'restore checkpoint' leaving the request (#274542)

- For restoring to a request, we actually want to restore the request's
  initial checkpoint but put our pointer before it, so that further
  requests overwrite that request checkpoint.
- We previously didn't discard undone checkpoints when new ones were applied
- Refined the behavior of disablement to more explicitly provide the
  right UX (show up to and including the last checkpoint that contains
  any applied operation)
This commit is contained in:
Connor Peet
2025-11-02 00:05:54 -07:00
committed by GitHub
parent 10650218a3
commit ceb31bb8f5
2 changed files with 100 additions and 43 deletions
@@ -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<string, string | undefined>();
// 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<void> {
private async _navigateToEpoch(restoreToEpoch: number, navigateToEpoch = restoreToEpoch): Promise<void> {
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 { currentEpoch, checkpoints, operations } = this._getVisibleOperationsAndCheckpoints();
if (operation.epoch < currentEpoch) {
throw new Error(`Cannot record operation at epoch ${operation.epoch} when current epoch is ${currentEpoch}`);
}
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
// 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);
@@ -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(), []);
// 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), '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' },
]);
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 () {