From 9ac74295a87e037ae3a33342146982009ee8eccc Mon Sep 17 00:00:00 2001 From: cwebster-99 Date: Tue, 21 Apr 2026 17:08:29 -0500 Subject: [PATCH 1/5] Address gaps in ai co author feature Co-authored-by: Copilot --- .../browser/aiContributionFeature.ts | 463 +++++++++++++++--- .../browser/editTelemetryContribution.ts | 2 +- .../browser/aiContributionFeature.test.ts | 202 +++++++- 3 files changed, 590 insertions(+), 77 deletions(-) diff --git a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts index 6dcef778f840a..905c6dfe68e6a 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts @@ -3,113 +3,448 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js'; +import { RunOnceScheduler } from '../../../../base/common/async.js'; +import { Disposable, toDisposable } from '../../../../base/common/lifecycle.js'; import { ResourceMap } from '../../../../base/common/map.js'; -import { autorun } from '../../../../base/common/observable.js'; +import { autorun, mapObservableArrayCached, runOnChange } from '../../../../base/common/observable.js'; import { URI, UriComponents } from '../../../../base/common/uri.js'; +import { StringReplacement } from '../../../../editor/common/core/edits/stringEdit.js'; import { CommandsRegistry } from '../../../../platform/commands/common/commands.js'; -import { AnnotatedDocument, IAnnotatedDocuments } from './helpers/annotatedDocuments.js'; -import { createDocWithJustReason } from './helpers/documentWithAnnotatedEdits.js'; -import { DocumentEditSourceTracker } from './telemetry/editTracker.js'; +import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; +import { EditSourceBase } from './helpers/documentWithAnnotatedEdits.js'; +import { ObservableWorkspace } from './helpers/observableWorkspace.js'; export type AiContributionLevel = 'chatAndAgent' | 'all'; -interface TrackerEntry { - readonly trackerStore: DisposableStore; - readonly tracker: DocumentEditSourceTracker; -} +const STORAGE_KEY = 'aiEdits.contributions'; +const SAVE_DEBOUNCE_MS = 250; /** - * Tracks AI-generated edits across open documents using the edit telemetry pipeline. + * Tracks which URIs contain *surviving* AI-generated content for git co-author + * trailer attribution. + * + * Unlike a simple "AI ever touched this URI" flag, this tracker shrinks the + * recorded AI ranges as edits replace, delete, or split them, and removes the + * URI entirely once no AI-authored content remains. That way users who revert + * an AI suggestion before committing do not get a misleading + * `Co-authored-by: Copilot` trailer. + * + * State is persisted per workspace, keyed by URI, together with the document + * length at the time of snapshot. On window reload, persisted state is only + * restored for documents whose current length still matches; otherwise the + * entry is dropped (we cannot meaningfully rebase ranges across an offline + * change). */ export class AiContributionFeature extends Disposable { - private readonly _trackers = new ResourceMap(); - private readonly _documentsByUri = new ResourceMap(); + /** In-memory tracker for currently-loaded documents. */ + private readonly _live = new ResourceMap(); + + /** + * Persisted snapshots for URIs that are not currently loaded (or that + * still match their loaded content). Used to answer queries for + * commits made on closed files and to seed live trackers on document + * load. + */ + private readonly _persisted = new ResourceMap(); + + private readonly _saveScheduler: RunOnceScheduler; + private _dirty = false; + private _disposing = false; constructor( - annotatedDocuments: IAnnotatedDocuments, + workspace: ObservableWorkspace, + @IStorageService private readonly _storageService: IStorageService, ) { super(); - this._register(autorun(reader => { - const docs = annotatedDocuments.documents.read(reader); - const activeUris = new ResourceMap(); + this._loadFromStorage(); - for (const doc of docs) { - const uri = doc.document.uri; - activeUris.set(uri, true); - this._documentsByUri.set(uri, doc); + this._saveScheduler = this._register(new RunOnceScheduler(() => this._saveToStorage(), SAVE_DEBOUNCE_MS)); + this._register(this._storageService.onWillSaveState(() => { + if (this._dirty) { + this._saveScheduler.cancel(); + this._saveToStorage(); + } + })); - if (!this._trackers.has(uri)) { - this._trackers.set(uri, this._createTrackerEntry(doc)); + // Subscribe to every loaded document. We deliberately do NOT gate on editor + // visibility, because the trailer should still be added when the agent edits + // a file that the user never opens, and it should survive closing the file. + const trackedDocs = mapObservableArrayCached(this, workspace.documents, (doc, store) => { + const initialLength = doc.value.get().value.length; + const persisted = this._persisted.get(doc.uri); + let ranges: AiSurvivingRanges; + if (persisted && persisted.contentLength === initialLength) { + ranges = AiSurvivingRanges.fromSerialized(persisted.ranges); + } else { + if (persisted) { + // Document content changed between sessions; we cannot rebase + // AI ranges across an offline edit, so drop the stale entry. + this._persisted.delete(doc.uri); + this._markDirty(); } + ranges = new AiSurvivingRanges(); } + this._live.set(doc.uri, ranges); - for (const [uri, entry] of this._trackers) { - if (!activeUris.has(uri)) { - entry.trackerStore.dispose(); - this._trackers.delete(uri); - this._documentsByUri.delete(uri); + store.add(runOnChange(doc.value, (_val, _prev, edits) => { + let changed = false; + for (const e of edits) { + const source = EditSourceBase.create(e.reason); + const level: AiContributionLevel | undefined = source.category === 'ai' + ? (source.feature === 'chat' ? 'chatAndAgent' : 'all') + : undefined; + if (ranges.apply(e.replacements, level)) { + changed = true; + } } - } - })); + if (changed) { + this._snapshot(doc.uri, ranges, doc.value.get().value.length); + } + })); - this._register(CommandsRegistry.registerCommand('_aiEdits.hasAiContributions', (_accessor, resources: UriComponents[], level: AiContributionLevel) => { - return this._hasAiContributions(resources, level); - })); + store.add(toDisposable(() => { + // Snapshot one last time when the document leaves the workspace + // so closed files keep their attribution across reloads. + this._snapshot(doc.uri, ranges, doc.value.get().value.length); + this._live.delete(doc.uri); + })); + }); - this._register(CommandsRegistry.registerCommand('_aiEdits.clearAiContributions', (_accessor, resources: UriComponents[]) => { - this._clearAiContributions(resources); - })); + // Force the cached array to be evaluated so the per-document subscriptions are wired up. + this._register(autorun(reader => { trackedDocs.read(reader); })); - this._register(CommandsRegistry.registerCommand('_aiEdits.clearAllAiContributions', () => { - this._clearAiContributions(); - })); + this._register(CommandsRegistry.registerCommand('_aiEdits.hasAiContributions', + (_acc, resources: UriComponents[], level: AiContributionLevel) => this._hasAiContributions(resources, level))); + this._register(CommandsRegistry.registerCommand('_aiEdits.clearAiContributions', + (_acc, resources: UriComponents[]) => this._clearAiContributions(resources))); + this._register(CommandsRegistry.registerCommand('_aiEdits.clearAllAiContributions', + () => this._clearAiContributions())); } - override dispose(): void { - for (const [, entry] of this._trackers) { - entry.trackerStore.dispose(); - } + public override dispose(): void { + // Cancel the debounced save first, then run super.dispose() which, via the + // per-document `toDisposable` callbacks, takes one final snapshot of every + // live document. Only after those snapshots have updated `_persisted` do + // we flush to storage. Flushing before super.dispose() would lose the + // final snapshots, because the scheduler they try to schedule has just + // been disposed along with the rest of this feature. + this._disposing = true; + this._saveScheduler.cancel(); super.dispose(); + if (this._dirty) { + this._saveToStorage(); + } } - private _createTrackerEntry(doc: AnnotatedDocument): TrackerEntry { - const trackerStore = new DisposableStore(); - const docWithJustReason = createDocWithJustReason(doc.documentWithAnnotations, trackerStore); - const tracker = trackerStore.add(new DocumentEditSourceTracker(docWithJustReason, undefined)); - return { trackerStore, tracker }; + private _snapshot(uri: URI, ranges: AiSurvivingRanges, contentLength: number): void { + if (ranges.isEmpty()) { + if (this._persisted.delete(uri)) { + this._markDirty(); + } + return; + } + this._persisted.set(uri, { contentLength, ranges: ranges.serialize() }); + this._markDirty(); + } + + private _markDirty(): void { + this._dirty = true; + if (!this._disposing) { + this._saveScheduler.schedule(); + } } private _hasAiContributions(resources: UriComponents[], level: AiContributionLevel): boolean { - for (const resource of resources) { - const entry = this._trackers.get(URI.revive(resource)); - if (entry) { - for (const edit of entry.tracker.getTrackedRanges()) { - if (edit.source.category === 'ai' && (level === 'all' || edit.source.feature === 'chat')) { - return true; - } + for (const r of resources) { + const uri = URI.revive(r); + const live = this._live.get(uri); + if (live) { + if (live.hasLevel(level)) { + return true; } + continue; + } + const persisted = this._persisted.get(uri); + if (persisted && hasLevel(persisted.ranges, level)) { + return true; } } return false; } private _clearAiContributions(resources?: UriComponents[]): void { - const uris = resources ? resources.map(r => URI.revive(r)) : [...this._trackers.keys()]; - for (const uri of uris) { - const entry = this._trackers.get(uri); - if (entry) { - entry.trackerStore.dispose(); - const doc = this._documentsByUri.get(uri); - if (doc) { - this._trackers.set(uri, this._createTrackerEntry(doc)); + let changed = false; + if (!resources) { + if (this._persisted.size > 0) { + this._persisted.clear(); + changed = true; + } + for (const ranges of this._live.values()) { + if (!ranges.isEmpty()) { + ranges.clear(); + changed = true; + } + } + } else { + for (const r of resources) { + const uri = URI.revive(r); + if (this._persisted.delete(uri)) { + changed = true; + } + const live = this._live.get(uri); + if (live && !live.isEmpty()) { + live.clear(); + changed = true; + } + } + } + if (changed) { + this._markDirty(); + } + } + + private _loadFromStorage(): void { + const raw = this._storageService.get(STORAGE_KEY, StorageScope.WORKSPACE); + if (!raw) { + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch { + return; + } + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + return; + } + for (const [uriString, value] of Object.entries(parsed as Record)) { + const entry = parsePersistedEntry(value); + if (!entry) { + continue; + } + let uri: URI; + try { + uri = URI.parse(uriString); + } catch { + continue; + } + this._persisted.set(uri, entry); + } + } + + private _saveToStorage(): void { + // Only clear the dirty flag after a successful write, so that if the + // storage call throws the onWillSaveState / dispose paths still see + // pending state and can retry. + try { + if (this._persisted.size === 0) { + this._storageService.remove(STORAGE_KEY, StorageScope.WORKSPACE); + } else { + const obj: Record = {}; + for (const [uri, entry] of this._persisted) { + obj[uri.toString()] = entry; + } + // StorageTarget.MACHINE: do not sync via Settings Sync. AI range + // offsets are tied to on-disk document content, which differs + // per machine, so syncing would produce stale attributions. + this._storageService.store(STORAGE_KEY, JSON.stringify(obj), StorageScope.WORKSPACE, StorageTarget.MACHINE); + } + this._dirty = false; + } catch { + // Keep _dirty true so a later save attempt can retry. + } + } +} + +interface ISerializedRange { + readonly start: number; + readonly length: number; + readonly level: AiContributionLevel; +} + +interface IPersistedRangesEntry { + readonly contentLength: number; + readonly ranges: readonly ISerializedRange[]; +} + +function hasLevel(ranges: readonly ISerializedRange[], level: AiContributionLevel): boolean { + if (ranges.length === 0) { + return false; + } + if (level === 'all') { + return true; + } + for (const r of ranges) { + if (r.level === 'chatAndAgent') { + return true; + } + } + return false; +} + +function parsePersistedEntry(value: unknown): IPersistedRangesEntry | undefined { + if (!value || typeof value !== 'object') { + return undefined; + } + const v = value as { contentLength?: unknown; ranges?: unknown }; + if (typeof v.contentLength !== 'number' || !Number.isFinite(v.contentLength) || v.contentLength < 0) { + return undefined; + } + if (!Array.isArray(v.ranges)) { + return undefined; + } + const ranges: ISerializedRange[] = []; + for (const r of v.ranges) { + if (!r || typeof r !== 'object') { + continue; + } + const { start, length, level } = r as { start?: unknown; length?: unknown; level?: unknown }; + if (typeof start !== 'number' || typeof length !== 'number' || length <= 0 || start < 0) { + continue; + } + if (level !== 'all' && level !== 'chatAndAgent') { + continue; + } + if (start + length > v.contentLength) { + continue; + } + ranges.push({ start, length, level }); + } + if (ranges.length === 0) { + return undefined; + } + ranges.sort((a, b) => a.start - b.start); + return { contentLength: v.contentLength, ranges: mergeTouching(ranges) }; +} + +/** + * Maintains a sorted, non-overlapping list of AI-authored offset ranges and + * keeps them in sync with arbitrary text edits. + */ +class AiSurvivingRanges { + + public static fromSerialized(ranges: readonly ISerializedRange[]): AiSurvivingRanges { + const copy = ranges.map(r => ({ start: r.start, length: r.length, level: r.level })); + copy.sort((a, b) => a.start - b.start); + return new AiSurvivingRanges(mergeTouching(copy)); + } + + private _ranges: ISerializedRange[]; + + constructor(ranges: ISerializedRange[] = []) { + this._ranges = ranges; + } + + public isEmpty(): boolean { + return this._ranges.length === 0; + } + + public hasLevel(level: AiContributionLevel): boolean { + return hasLevel(this._ranges, level); + } + + public clear(): void { + this._ranges = []; + } + + public serialize(): ISerializedRange[] { + return this._ranges.map(r => ({ start: r.start, length: r.length, level: r.level })); + } + + /** + * Apply a batch of replacements (in coordinates of the document state + * before the batch). Returns `true` iff the surviving ranges changed. + * + * @param level if defined, the inserted text is recorded as AI-authored + * at the given level; otherwise the edit is treated as + * non-AI (it can still trim or split existing AI ranges). + */ + public apply(replacements: readonly StringReplacement[], level: AiContributionLevel | undefined): boolean { + if (replacements.length === 0) { + return false; + } + // Process in reverse order so that each replacement's "before" coordinates + // remain valid against the (still-unmodified) lower part of the document. + let changed = false; + for (let i = replacements.length - 1; i >= 0; i--) { + const r = replacements[i]; + if (this._applyOne(r.replaceRange.start, r.replaceRange.endExclusive, r.newText.length, level)) { + changed = true; + } + } + return changed; + } + + private _applyOne(start: number, endExclusive: number, newLen: number, level: AiContributionLevel | undefined): boolean { + const delta = newLen - (endExclusive - start); + const out: ISerializedRange[] = []; + let touched = false; + + for (const r of this._ranges) { + const rEnd = r.start + r.length; + if (rEnd <= start) { + // Entirely before the edit - unaffected. + out.push(r); + } else if (r.start >= endExclusive) { + // Entirely after the edit - shifted by delta. + if (delta !== 0) { + out.push({ start: r.start + delta, length: r.length, level: r.level }); + touched = true; } else { - this._trackers.delete(uri); - this._documentsByUri.delete(uri); + out.push(r); + } + } else { + // Overlaps the deleted range. Keep the parts outside it. + touched = true; + if (r.start < start) { + out.push({ start: r.start, length: start - r.start, level: r.level }); + } + if (rEnd > endExclusive) { + out.push({ start: endExclusive + delta, length: rEnd - endExclusive, level: r.level }); } } } + + if (level !== undefined && newLen > 0) { + insertSorted(out, { start, length: newLen, level }); + touched = true; + } + + if (!touched) { + return false; + } + this._ranges = mergeTouching(out); + return true; + } +} + +function insertSorted(list: ISerializedRange[], r: ISerializedRange): void { + let i = 0; + while (i < list.length && list[i].start < r.start) { + i++; + } + list.splice(i, 0, r); +} + +function mergeTouching(ranges: ISerializedRange[]): ISerializedRange[] { + if (ranges.length <= 1) { + return ranges; + } + const out: ISerializedRange[] = []; + let cur: ISerializedRange = ranges[0]; + for (let i = 1; i < ranges.length; i++) { + const next = ranges[i]; + const curEnd = cur.start + cur.length; + if (next.start <= curEnd && next.level === cur.level) { + const end = Math.max(curEnd, next.start + next.length); + cur = { start: cur.start, length: end - cur.start, level: cur.level }; + } else { + out.push(cur); + cur = next; + } } + out.push(cur); + return out; } diff --git a/src/vs/workbench/contrib/editTelemetry/browser/editTelemetryContribution.ts b/src/vs/workbench/contrib/editTelemetry/browser/editTelemetryContribution.ts index 6f2bbf8fc54aa..d1f2fdafc9413 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/editTelemetryContribution.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/editTelemetryContribution.ts @@ -63,7 +63,7 @@ export class EditTelemetryContribution extends Disposable { if (addAICoAuthor.read(r) === 'off') { return; } - r.store.add(instantiationService.createInstance(AiContributionFeature, annotatedDocuments.read(r))); + r.store.add(instantiationService.createInstance(AiContributionFeature, workspace.read(r))); })); } } diff --git a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts index f6217cf441770..e5448886f0c44 100644 --- a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts +++ b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts @@ -8,18 +8,16 @@ import { DisposableStore } from '../../../../../base/common/lifecycle.js'; import { URI } from '../../../../../base/common/uri.js'; import { ServiceCollection } from '../../../../../platform/instantiation/common/serviceCollection.js'; import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; -import { AnnotatedDocuments, UriVisibilityProvider } from '../../browser/helpers/annotatedDocuments.js'; import { StringEditWithReason } from '../../browser/helpers/observableWorkspace.js'; import { AiContributionFeature } from '../../browser/aiContributionFeature.js'; import { EditSources } from '../../../../../editor/common/textModelEditSource.js'; -import { DiffService } from '../../browser/helpers/documentWithAnnotatedEdits.js'; -import { computeStringDiff } from '../../../../../editor/common/services/editorWebWorker.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; -import { ILogService, NullLogService } from '../../../../../platform/log/common/log.js'; import { MutableObservableWorkspace } from './editTelemetry.test.js'; import { CommandsRegistry } from '../../../../../platform/commands/common/commands.js'; import { timeout } from '../../../../../base/common/async.js'; import { runWithFakedTimers } from '../../../../../base/test/common/timeTravelScheduler.js'; +import { IStorageService } from '../../../../../platform/storage/common/storage.js'; +import { TestStorageService } from '../../../../test/common/workbenchTestServices.js'; suite('AiContributionFeature', () => { ensureNoDisposablesAreLeakedInTestSuite(); @@ -49,16 +47,16 @@ suite('AiContributionFeature', () => { correlationId: undefined, }); - function setup(): void { + let storage: TestStorageService; + + function setup(sharedStorage?: TestStorageService): void { disposables = new DisposableStore(); const instantiationService = disposables.add(new TestInstantiationService(new ServiceCollection(), false, undefined, true)); - instantiationService.stubInstance(DiffService, { computeDiff: async (original, modified) => computeStringDiff(original, modified, { maxComputationTimeMs: 500 }, 'advanced') }); - instantiationService.stubInstance(UriVisibilityProvider, { isVisible: () => true }); - instantiationService.stub(ILogService, new NullLogService()); + storage = sharedStorage ?? disposables.add(new TestStorageService()); + instantiationService.stub(IStorageService, storage); workspace = new MutableObservableWorkspace(); - const annotatedDocuments = disposables.add(new AnnotatedDocuments(workspace, instantiationService)); - disposables.add(instantiationService.createInstance(AiContributionFeature, annotatedDocuments)); + disposables.add(instantiationService.createInstance(AiContributionFeature, workspace)); } function hasAiContributions(uris: URI[], level: 'chatAndAgent' | 'all'): boolean { @@ -176,7 +174,7 @@ suite('AiContributionFeature', () => { disposables.dispose(); })); - test('cleans up tracker when document is closed', () => runWithFakedTimers({}, async () => { + test('contributions persist after document is closed', () => runWithFakedTimers({}, async () => { setup(); const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); await timeout(1500); @@ -189,7 +187,10 @@ suite('AiContributionFeature', () => { d.dispose(); await timeout(1500); - assert.strictEqual(hasAiContributions([fileA], 'all'), false); + // Contributions are tracked per-URI and must survive document close, + // otherwise commits made after the file was closed would lose the trailer. + assert.strictEqual(hasAiContributions([fileA], 'all'), true); + assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), true); disposables.dispose(); })); @@ -214,4 +215,181 @@ suite('AiContributionFeature', () => { assert.strictEqual(hasAiContributions([fileB], 'all'), false); disposables.dispose(); })); + + test('user edit that overwrites the AI range removes the contribution', () => runWithFakedTimers({}, async () => { + setup(); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + await timeout(1500); + assert.strictEqual(hasAiContributions([d.uri], 'all'), true); + + // User selects and replaces the AI-authored text - no AI content remains. + d.applyEdit(StringEditWithReason.replace(d.findRange('world'), 'manual', userEdit)); + await timeout(1500); + + assert.strictEqual(hasAiContributions([d.uri], 'all'), false); + assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), false); + disposables.dispose(); + })); + + test('removing chat range while inline AI range survives downgrades level', () => runWithFakedTimers({}, async () => { + setup(); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'aaa bbb' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('aaa'), 'CHAT', chatEdit)); + d.applyEdit(StringEditWithReason.replace(d.findRange('bbb'), 'INLINE', inlineCompletionEdit)); + await timeout(1500); + assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), true); + assert.strictEqual(hasAiContributions([d.uri], 'all'), true); + + // User overwrites the chat-authored region only. + d.applyEdit(StringEditWithReason.replace(d.findRange('CHAT'), 'user', userEdit)); + await timeout(1500); + + assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), false); + assert.strictEqual(hasAiContributions([d.uri], 'all'), true); + disposables.dispose(); + })); + + test('user edit inside the AI range only shrinks attribution', () => runWithFakedTimers({}, async () => { + setup(); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'xx' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('xx'), 'aibody', chatEdit)); + await timeout(1500); + + // Delete one character in the middle of the AI text - some AI content still remains. + d.applyEdit(StringEditWithReason.replace(d.findRange('b'), '', userEdit)); + await timeout(1500); + + assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), true); + disposables.dispose(); + })); + + test('persisted AI ranges survive a workspace reload', () => runWithFakedTimers({}, async () => { + const reloadStore = new DisposableStore(); + const sharedStorage = reloadStore.add(new TestStorageService()); + + setup(sharedStorage); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + await timeout(1500); + + // Simulate a window reload: dispose everything (which flushes pending writes), + // then bring up a fresh feature against the same backing storage. + disposables.dispose(); + + setup(sharedStorage); + await timeout(1500); + + assert.strictEqual(hasAiContributions([fileA], 'all'), true); + assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), true); + disposables.dispose(); + reloadStore.dispose(); + })); + + test('persisted entry is dropped when document length changed offline', () => runWithFakedTimers({}, async () => { + const reloadStore = new DisposableStore(); + const sharedStorage = reloadStore.add(new TestStorageService()); + + setup(sharedStorage); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + await timeout(1500); + disposables.dispose(); + + // Reopen the document with a different length than what was persisted - we + // cannot rebase ranges across an offline change, so the entry must be dropped. + setup(sharedStorage); + const d2 = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'changed externally' }, undefined)); + await timeout(1500); + + assert.strictEqual(hasAiContributions([d2.uri], 'all'), false); + assert.strictEqual(hasAiContributions([d2.uri], 'chatAndAgent'), false); + + // After the stale entry is dropped, closing the document must not resurrect + // it: subsequent queries (against the now-closed URI) must still be false. + d2.dispose(); + await timeout(1500); + assert.strictEqual(hasAiContributions([fileA], 'all'), false); + assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), false); + + disposables.dispose(); + reloadStore.dispose(); + })); + + test('clear removes contributions for a closed (persisted-only) file', () => runWithFakedTimers({}, async () => { + setup(); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + await timeout(1500); + + // Close the document so the entry lives only in the persisted map. + d.dispose(); + await timeout(1500); + assert.strictEqual(hasAiContributions([fileA], 'all'), true); + + clearAiContributions([fileA]); + + assert.strictEqual(hasAiContributions([fileA], 'all'), false); + assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), false); + disposables.dispose(); + })); + + test('dispose flushes pending writes even before the debounce fires', () => runWithFakedTimers({}, async () => { + const reloadStore = new DisposableStore(); + const sharedStorage = reloadStore.add(new TestStorageService()); + + setup(sharedStorage); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + + // Apply an AI edit and immediately tear down, without waiting for the + // save scheduler's debounce window to elapse. The dispose path must + // still flush the pending snapshot, otherwise the next window would + // see no attribution for a file that the agent just edited. + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + disposables.dispose(); + + setup(sharedStorage); + await timeout(1500); + + assert.strictEqual(hasAiContributions([fileA], 'all'), true); + assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), true); + disposables.dispose(); + reloadStore.dispose(); + })); + + test('dispose snapshots live documents that had no pending save', () => runWithFakedTimers({}, async () => { + const reloadStore = new DisposableStore(); + const sharedStorage = reloadStore.add(new TestStorageService()); + + setup(sharedStorage); + const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); + await timeout(1500); + + d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); + await timeout(1500); // let the debounce flush normally + + // Now tear down with no new dirty state: the per-document `toDisposable` + // callbacks will re-snapshot, which must not be lost to the already- + // cancelled scheduler. + disposables.dispose(); + + setup(sharedStorage); + await timeout(1500); + + assert.strictEqual(hasAiContributions([fileA], 'all'), true); + disposables.dispose(); + reloadStore.dispose(); + })); }); From 58ee477c4c7ef7faac3bbc4385e1d49f872fbedf Mon Sep 17 00:00:00 2001 From: cwebster-99 Date: Tue, 21 Apr 2026 17:23:32 -0500 Subject: [PATCH 2/5] clean up comments Co-authored-by: Copilot --- .../browser/aiContributionFeature.ts | 82 ++++++++----------- 1 file changed, 35 insertions(+), 47 deletions(-) diff --git a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts index 905c6dfe68e6a..115931764c80a 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts @@ -20,20 +20,17 @@ const STORAGE_KEY = 'aiEdits.contributions'; const SAVE_DEBOUNCE_MS = 250; /** - * Tracks which URIs contain *surviving* AI-generated content for git co-author - * trailer attribution. + * Tracks which URIs contain *surviving* AI-generated content for git + * co-author trailer attribution. * - * Unlike a simple "AI ever touched this URI" flag, this tracker shrinks the - * recorded AI ranges as edits replace, delete, or split them, and removes the - * URI entirely once no AI-authored content remains. That way users who revert - * an AI suggestion before committing do not get a misleading - * `Co-authored-by: Copilot` trailer. + * Recorded AI ranges are shrunk, split, or removed as later edits overwrite + * them, so reverting an AI suggestion before committing does not produce a + * misleading `Co-authored-by: Copilot` trailer. * - * State is persisted per workspace, keyed by URI, together with the document - * length at the time of snapshot. On window reload, persisted state is only - * restored for documents whose current length still matches; otherwise the - * entry is dropped (we cannot meaningfully rebase ranges across an offline - * change). + * State is persisted per workspace, keyed by URI plus document length at + * snapshot time. On reload, an entry is restored only if the current + * document length still matches; otherwise it is dropped (offline edits + * cannot be rebased). */ export class AiContributionFeature extends Disposable { @@ -41,10 +38,8 @@ export class AiContributionFeature extends Disposable { private readonly _live = new ResourceMap(); /** - * Persisted snapshots for URIs that are not currently loaded (or that - * still match their loaded content). Used to answer queries for - * commits made on closed files and to seed live trackers on document - * load. + * Latest snapshot per URI. Used to answer queries for closed files and + * to seed live trackers on document load. */ private readonly _persisted = new ResourceMap(); @@ -68,9 +63,8 @@ export class AiContributionFeature extends Disposable { } })); - // Subscribe to every loaded document. We deliberately do NOT gate on editor - // visibility, because the trailer should still be added when the agent edits - // a file that the user never opens, and it should survive closing the file. + // Track every loaded document, regardless of editor visibility: the + // trailer must apply to agent-only edits and survive closing the file. const trackedDocs = mapObservableArrayCached(this, workspace.documents, (doc, store) => { const initialLength = doc.value.get().value.length; const persisted = this._persisted.get(doc.uri); @@ -79,8 +73,7 @@ export class AiContributionFeature extends Disposable { ranges = AiSurvivingRanges.fromSerialized(persisted.ranges); } else { if (persisted) { - // Document content changed between sessions; we cannot rebase - // AI ranges across an offline edit, so drop the stale entry. + // Length mismatch: cannot rebase ranges across an offline edit. this._persisted.delete(doc.uri); this._markDirty(); } @@ -105,14 +98,13 @@ export class AiContributionFeature extends Disposable { })); store.add(toDisposable(() => { - // Snapshot one last time when the document leaves the workspace - // so closed files keep their attribution across reloads. + // Final snapshot so closed files keep their attribution. this._snapshot(doc.uri, ranges, doc.value.get().value.length); this._live.delete(doc.uri); })); }); - // Force the cached array to be evaluated so the per-document subscriptions are wired up. + // Force the cached array so per-document subscriptions get wired up. this._register(autorun(reader => { trackedDocs.read(reader); })); this._register(CommandsRegistry.registerCommand('_aiEdits.hasAiContributions', @@ -124,12 +116,10 @@ export class AiContributionFeature extends Disposable { } public override dispose(): void { - // Cancel the debounced save first, then run super.dispose() which, via the - // per-document `toDisposable` callbacks, takes one final snapshot of every - // live document. Only after those snapshots have updated `_persisted` do - // we flush to storage. Flushing before super.dispose() would lose the - // final snapshots, because the scheduler they try to schedule has just - // been disposed along with the rest of this feature. + // Order matters: `super.dispose()` runs the per-document `toDisposable` + // callbacks, each of which takes a final snapshot into `_persisted`. + // `_disposing` suppresses re-scheduling from `_markDirty` while that + // happens, then we flush once at the end. this._disposing = true; this._saveScheduler.cancel(); super.dispose(); @@ -235,9 +225,8 @@ export class AiContributionFeature extends Disposable { } private _saveToStorage(): void { - // Only clear the dirty flag after a successful write, so that if the - // storage call throws the onWillSaveState / dispose paths still see - // pending state and can retry. + // Only clear `_dirty` after a successful write so retries can flush + // pending state on the next save attempt. try { if (this._persisted.size === 0) { this._storageService.remove(STORAGE_KEY, StorageScope.WORKSPACE); @@ -246,14 +235,13 @@ export class AiContributionFeature extends Disposable { for (const [uri, entry] of this._persisted) { obj[uri.toString()] = entry; } - // StorageTarget.MACHINE: do not sync via Settings Sync. AI range - // offsets are tied to on-disk document content, which differs - // per machine, so syncing would produce stale attributions. + // MACHINE: ranges are tied to on-disk content, which differs per + // machine — syncing would produce stale attributions. this._storageService.store(STORAGE_KEY, JSON.stringify(obj), StorageScope.WORKSPACE, StorageTarget.MACHINE); } this._dirty = false; } catch { - // Keep _dirty true so a later save attempt can retry. + // Keep `_dirty` so the next save can retry. } } } @@ -354,19 +342,19 @@ class AiSurvivingRanges { } /** - * Apply a batch of replacements (in coordinates of the document state - * before the batch). Returns `true` iff the surviving ranges changed. + * Apply a batch of replacements (in pre-batch coordinates). Returns + * `true` iff the surviving ranges changed. * - * @param level if defined, the inserted text is recorded as AI-authored - * at the given level; otherwise the edit is treated as - * non-AI (it can still trim or split existing AI ranges). + * @param level if defined, the inserted text is recorded at the given + * level; otherwise the edit is non-AI (and may still trim + * or split existing AI ranges). */ public apply(replacements: readonly StringReplacement[], level: AiContributionLevel | undefined): boolean { if (replacements.length === 0) { return false; } - // Process in reverse order so that each replacement's "before" coordinates - // remain valid against the (still-unmodified) lower part of the document. + // Reverse order keeps each replacement's pre-batch coordinates valid + // against the still-unmodified lower part of the document. let changed = false; for (let i = replacements.length - 1; i >= 0; i--) { const r = replacements[i]; @@ -385,10 +373,10 @@ class AiSurvivingRanges { for (const r of this._ranges) { const rEnd = r.start + r.length; if (rEnd <= start) { - // Entirely before the edit - unaffected. + // Before the edit. out.push(r); } else if (r.start >= endExclusive) { - // Entirely after the edit - shifted by delta. + // After the edit - shift by delta. if (delta !== 0) { out.push({ start: r.start + delta, length: r.length, level: r.level }); touched = true; @@ -396,7 +384,7 @@ class AiSurvivingRanges { out.push(r); } } else { - // Overlaps the deleted range. Keep the parts outside it. + // Overlaps - keep the parts outside the deleted span. touched = true; if (r.start < start) { out.push({ start: r.start, length: start - r.start, level: r.level }); From 6d13fed3a066be9971225d5effdd228bc8b9d05b Mon Sep 17 00:00:00 2001 From: cwebster-99 Date: Tue, 21 Apr 2026 17:34:48 -0500 Subject: [PATCH 3/5] remove survival code Co-authored-by: Copilot --- .../browser/aiContributionFeature.ts | 321 ++---------------- .../browser/aiContributionFeature.test.ts | 109 ------ 2 files changed, 33 insertions(+), 397 deletions(-) diff --git a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts index 115931764c80a..954ad3260020a 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts @@ -4,11 +4,10 @@ *--------------------------------------------------------------------------------------------*/ import { RunOnceScheduler } from '../../../../base/common/async.js'; -import { Disposable, toDisposable } from '../../../../base/common/lifecycle.js'; +import { Disposable } from '../../../../base/common/lifecycle.js'; import { ResourceMap } from '../../../../base/common/map.js'; import { autorun, mapObservableArrayCached, runOnChange } from '../../../../base/common/observable.js'; import { URI, UriComponents } from '../../../../base/common/uri.js'; -import { StringReplacement } from '../../../../editor/common/core/edits/stringEdit.js'; import { CommandsRegistry } from '../../../../platform/commands/common/commands.js'; import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; import { EditSourceBase } from './helpers/documentWithAnnotatedEdits.js'; @@ -20,32 +19,21 @@ const STORAGE_KEY = 'aiEdits.contributions'; const SAVE_DEBOUNCE_MS = 250; /** - * Tracks which URIs contain *surviving* AI-generated content for git - * co-author trailer attribution. + * Tracks which URIs have received AI-generated edits, so the git extension + * can add a `Co-authored-by: Copilot` trailer to commits that touch them. * - * Recorded AI ranges are shrunk, split, or removed as later edits overwrite - * them, so reverting an AI suggestion before committing does not produce a - * misleading `Co-authored-by: Copilot` trailer. - * - * State is persisted per workspace, keyed by URI plus document length at - * snapshot time. On reload, an entry is restored only if the current - * document length still matches; otherwise it is dropped (offline edits - * cannot be rebased). + * Attribution is recorded per URI: once an AI edit lands in a file, the + * file is considered AI-contributed until explicitly cleared. State is + * persisted per workspace so the trailer is still applied after a window + * reload, after the file is closed, or for files that the agent edited + * without the user ever opening them in an editor. */ export class AiContributionFeature extends Disposable { - /** In-memory tracker for currently-loaded documents. */ - private readonly _live = new ResourceMap(); - - /** - * Latest snapshot per URI. Used to answer queries for closed files and - * to seed live trackers on document load. - */ - private readonly _persisted = new ResourceMap(); + private readonly _contributions = new ResourceMap(); private readonly _saveScheduler: RunOnceScheduler; private _dirty = false; - private _disposing = false; constructor( workspace: ObservableWorkspace, @@ -66,41 +54,15 @@ export class AiContributionFeature extends Disposable { // Track every loaded document, regardless of editor visibility: the // trailer must apply to agent-only edits and survive closing the file. const trackedDocs = mapObservableArrayCached(this, workspace.documents, (doc, store) => { - const initialLength = doc.value.get().value.length; - const persisted = this._persisted.get(doc.uri); - let ranges: AiSurvivingRanges; - if (persisted && persisted.contentLength === initialLength) { - ranges = AiSurvivingRanges.fromSerialized(persisted.ranges); - } else { - if (persisted) { - // Length mismatch: cannot rebase ranges across an offline edit. - this._persisted.delete(doc.uri); - this._markDirty(); - } - ranges = new AiSurvivingRanges(); - } - this._live.set(doc.uri, ranges); - store.add(runOnChange(doc.value, (_val, _prev, edits) => { - let changed = false; for (const e of edits) { const source = EditSourceBase.create(e.reason); - const level: AiContributionLevel | undefined = source.category === 'ai' - ? (source.feature === 'chat' ? 'chatAndAgent' : 'all') - : undefined; - if (ranges.apply(e.replacements, level)) { - changed = true; + if (source.category !== 'ai') { + continue; } + const level: AiContributionLevel = source.feature === 'chat' ? 'chatAndAgent' : 'all'; + this._record(doc.uri, level); } - if (changed) { - this._snapshot(doc.uri, ranges, doc.value.get().value.length); - } - })); - - store.add(toDisposable(() => { - // Final snapshot so closed files keep their attribution. - this._snapshot(doc.uri, ranges, doc.value.get().value.length); - this._live.delete(doc.uri); })); }); @@ -116,11 +78,6 @@ export class AiContributionFeature extends Disposable { } public override dispose(): void { - // Order matters: `super.dispose()` runs the per-document `toDisposable` - // callbacks, each of which takes a final snapshot into `_persisted`. - // `_disposing` suppresses re-scheduling from `_markDirty` while that - // happens, then we flush once at the end. - this._disposing = true; this._saveScheduler.cancel(); super.dispose(); if (this._dirty) { @@ -128,36 +85,28 @@ export class AiContributionFeature extends Disposable { } } - private _snapshot(uri: URI, ranges: AiSurvivingRanges, contentLength: number): void { - if (ranges.isEmpty()) { - if (this._persisted.delete(uri)) { - this._markDirty(); - } + private _record(uri: URI, level: AiContributionLevel): void { + const existing = this._contributions.get(uri); + // `chatAndAgent` is the stronger attribution; never downgrade to `all`. + if (existing === 'chatAndAgent' || existing === level) { return; } - this._persisted.set(uri, { contentLength, ranges: ranges.serialize() }); + this._contributions.set(uri, level); this._markDirty(); } private _markDirty(): void { this._dirty = true; - if (!this._disposing) { - this._saveScheduler.schedule(); - } + this._saveScheduler.schedule(); } private _hasAiContributions(resources: UriComponents[], level: AiContributionLevel): boolean { for (const r of resources) { - const uri = URI.revive(r); - const live = this._live.get(uri); - if (live) { - if (live.hasLevel(level)) { - return true; - } + const recorded = this._contributions.get(URI.revive(r)); + if (recorded === undefined) { continue; } - const persisted = this._persisted.get(uri); - if (persisted && hasLevel(persisted.ranges, level)) { + if (level === 'all' || recorded === 'chatAndAgent') { return true; } } @@ -167,25 +116,13 @@ export class AiContributionFeature extends Disposable { private _clearAiContributions(resources?: UriComponents[]): void { let changed = false; if (!resources) { - if (this._persisted.size > 0) { - this._persisted.clear(); + if (this._contributions.size > 0) { + this._contributions.clear(); changed = true; } - for (const ranges of this._live.values()) { - if (!ranges.isEmpty()) { - ranges.clear(); - changed = true; - } - } } else { for (const r of resources) { - const uri = URI.revive(r); - if (this._persisted.delete(uri)) { - changed = true; - } - const live = this._live.get(uri); - if (live && !live.isEmpty()) { - live.clear(); + if (this._contributions.delete(URI.revive(r))) { changed = true; } } @@ -210,8 +147,7 @@ export class AiContributionFeature extends Disposable { return; } for (const [uriString, value] of Object.entries(parsed as Record)) { - const entry = parsePersistedEntry(value); - if (!entry) { + if (value !== 'chatAndAgent' && value !== 'all') { continue; } let uri: URI; @@ -220,7 +156,7 @@ export class AiContributionFeature extends Disposable { } catch { continue; } - this._persisted.set(uri, entry); + this._contributions.set(uri, value); } } @@ -228,15 +164,15 @@ export class AiContributionFeature extends Disposable { // Only clear `_dirty` after a successful write so retries can flush // pending state on the next save attempt. try { - if (this._persisted.size === 0) { + if (this._contributions.size === 0) { this._storageService.remove(STORAGE_KEY, StorageScope.WORKSPACE); } else { - const obj: Record = {}; - for (const [uri, entry] of this._persisted) { - obj[uri.toString()] = entry; + const obj: Record = {}; + for (const [uri, level] of this._contributions) { + obj[uri.toString()] = level; } - // MACHINE: ranges are tied to on-disk content, which differs per - // machine — syncing would produce stale attributions. + // MACHINE: attribution is tied to on-disk content of this + // workspace, which differs per machine. this._storageService.store(STORAGE_KEY, JSON.stringify(obj), StorageScope.WORKSPACE, StorageTarget.MACHINE); } this._dirty = false; @@ -245,194 +181,3 @@ export class AiContributionFeature extends Disposable { } } } - -interface ISerializedRange { - readonly start: number; - readonly length: number; - readonly level: AiContributionLevel; -} - -interface IPersistedRangesEntry { - readonly contentLength: number; - readonly ranges: readonly ISerializedRange[]; -} - -function hasLevel(ranges: readonly ISerializedRange[], level: AiContributionLevel): boolean { - if (ranges.length === 0) { - return false; - } - if (level === 'all') { - return true; - } - for (const r of ranges) { - if (r.level === 'chatAndAgent') { - return true; - } - } - return false; -} - -function parsePersistedEntry(value: unknown): IPersistedRangesEntry | undefined { - if (!value || typeof value !== 'object') { - return undefined; - } - const v = value as { contentLength?: unknown; ranges?: unknown }; - if (typeof v.contentLength !== 'number' || !Number.isFinite(v.contentLength) || v.contentLength < 0) { - return undefined; - } - if (!Array.isArray(v.ranges)) { - return undefined; - } - const ranges: ISerializedRange[] = []; - for (const r of v.ranges) { - if (!r || typeof r !== 'object') { - continue; - } - const { start, length, level } = r as { start?: unknown; length?: unknown; level?: unknown }; - if (typeof start !== 'number' || typeof length !== 'number' || length <= 0 || start < 0) { - continue; - } - if (level !== 'all' && level !== 'chatAndAgent') { - continue; - } - if (start + length > v.contentLength) { - continue; - } - ranges.push({ start, length, level }); - } - if (ranges.length === 0) { - return undefined; - } - ranges.sort((a, b) => a.start - b.start); - return { contentLength: v.contentLength, ranges: mergeTouching(ranges) }; -} - -/** - * Maintains a sorted, non-overlapping list of AI-authored offset ranges and - * keeps them in sync with arbitrary text edits. - */ -class AiSurvivingRanges { - - public static fromSerialized(ranges: readonly ISerializedRange[]): AiSurvivingRanges { - const copy = ranges.map(r => ({ start: r.start, length: r.length, level: r.level })); - copy.sort((a, b) => a.start - b.start); - return new AiSurvivingRanges(mergeTouching(copy)); - } - - private _ranges: ISerializedRange[]; - - constructor(ranges: ISerializedRange[] = []) { - this._ranges = ranges; - } - - public isEmpty(): boolean { - return this._ranges.length === 0; - } - - public hasLevel(level: AiContributionLevel): boolean { - return hasLevel(this._ranges, level); - } - - public clear(): void { - this._ranges = []; - } - - public serialize(): ISerializedRange[] { - return this._ranges.map(r => ({ start: r.start, length: r.length, level: r.level })); - } - - /** - * Apply a batch of replacements (in pre-batch coordinates). Returns - * `true` iff the surviving ranges changed. - * - * @param level if defined, the inserted text is recorded at the given - * level; otherwise the edit is non-AI (and may still trim - * or split existing AI ranges). - */ - public apply(replacements: readonly StringReplacement[], level: AiContributionLevel | undefined): boolean { - if (replacements.length === 0) { - return false; - } - // Reverse order keeps each replacement's pre-batch coordinates valid - // against the still-unmodified lower part of the document. - let changed = false; - for (let i = replacements.length - 1; i >= 0; i--) { - const r = replacements[i]; - if (this._applyOne(r.replaceRange.start, r.replaceRange.endExclusive, r.newText.length, level)) { - changed = true; - } - } - return changed; - } - - private _applyOne(start: number, endExclusive: number, newLen: number, level: AiContributionLevel | undefined): boolean { - const delta = newLen - (endExclusive - start); - const out: ISerializedRange[] = []; - let touched = false; - - for (const r of this._ranges) { - const rEnd = r.start + r.length; - if (rEnd <= start) { - // Before the edit. - out.push(r); - } else if (r.start >= endExclusive) { - // After the edit - shift by delta. - if (delta !== 0) { - out.push({ start: r.start + delta, length: r.length, level: r.level }); - touched = true; - } else { - out.push(r); - } - } else { - // Overlaps - keep the parts outside the deleted span. - touched = true; - if (r.start < start) { - out.push({ start: r.start, length: start - r.start, level: r.level }); - } - if (rEnd > endExclusive) { - out.push({ start: endExclusive + delta, length: rEnd - endExclusive, level: r.level }); - } - } - } - - if (level !== undefined && newLen > 0) { - insertSorted(out, { start, length: newLen, level }); - touched = true; - } - - if (!touched) { - return false; - } - this._ranges = mergeTouching(out); - return true; - } -} - -function insertSorted(list: ISerializedRange[], r: ISerializedRange): void { - let i = 0; - while (i < list.length && list[i].start < r.start) { - i++; - } - list.splice(i, 0, r); -} - -function mergeTouching(ranges: ISerializedRange[]): ISerializedRange[] { - if (ranges.length <= 1) { - return ranges; - } - const out: ISerializedRange[] = []; - let cur: ISerializedRange = ranges[0]; - for (let i = 1; i < ranges.length; i++) { - const next = ranges[i]; - const curEnd = cur.start + cur.length; - if (next.start <= curEnd && next.level === cur.level) { - const end = Math.max(curEnd, next.start + next.length); - cur = { start: cur.start, length: end - cur.start, level: cur.level }; - } else { - out.push(cur); - cur = next; - } - } - out.push(cur); - return out; -} diff --git a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts index e5448886f0c44..d0a4e96084109 100644 --- a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts +++ b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts @@ -216,60 +216,6 @@ suite('AiContributionFeature', () => { disposables.dispose(); })); - test('user edit that overwrites the AI range removes the contribution', () => runWithFakedTimers({}, async () => { - setup(); - const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); - await timeout(1500); - - d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); - await timeout(1500); - assert.strictEqual(hasAiContributions([d.uri], 'all'), true); - - // User selects and replaces the AI-authored text - no AI content remains. - d.applyEdit(StringEditWithReason.replace(d.findRange('world'), 'manual', userEdit)); - await timeout(1500); - - assert.strictEqual(hasAiContributions([d.uri], 'all'), false); - assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), false); - disposables.dispose(); - })); - - test('removing chat range while inline AI range survives downgrades level', () => runWithFakedTimers({}, async () => { - setup(); - const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'aaa bbb' }, undefined)); - await timeout(1500); - - d.applyEdit(StringEditWithReason.replace(d.findRange('aaa'), 'CHAT', chatEdit)); - d.applyEdit(StringEditWithReason.replace(d.findRange('bbb'), 'INLINE', inlineCompletionEdit)); - await timeout(1500); - assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), true); - assert.strictEqual(hasAiContributions([d.uri], 'all'), true); - - // User overwrites the chat-authored region only. - d.applyEdit(StringEditWithReason.replace(d.findRange('CHAT'), 'user', userEdit)); - await timeout(1500); - - assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), false); - assert.strictEqual(hasAiContributions([d.uri], 'all'), true); - disposables.dispose(); - })); - - test('user edit inside the AI range only shrinks attribution', () => runWithFakedTimers({}, async () => { - setup(); - const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'xx' }, undefined)); - await timeout(1500); - - d.applyEdit(StringEditWithReason.replace(d.findRange('xx'), 'aibody', chatEdit)); - await timeout(1500); - - // Delete one character in the middle of the AI text - some AI content still remains. - d.applyEdit(StringEditWithReason.replace(d.findRange('b'), '', userEdit)); - await timeout(1500); - - assert.strictEqual(hasAiContributions([d.uri], 'chatAndAgent'), true); - disposables.dispose(); - })); - test('persisted AI ranges survive a workspace reload', () => runWithFakedTimers({}, async () => { const reloadStore = new DisposableStore(); const sharedStorage = reloadStore.add(new TestStorageService()); @@ -294,38 +240,6 @@ suite('AiContributionFeature', () => { reloadStore.dispose(); })); - test('persisted entry is dropped when document length changed offline', () => runWithFakedTimers({}, async () => { - const reloadStore = new DisposableStore(); - const sharedStorage = reloadStore.add(new TestStorageService()); - - setup(sharedStorage); - const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); - await timeout(1500); - - d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); - await timeout(1500); - disposables.dispose(); - - // Reopen the document with a different length than what was persisted - we - // cannot rebase ranges across an offline change, so the entry must be dropped. - setup(sharedStorage); - const d2 = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'changed externally' }, undefined)); - await timeout(1500); - - assert.strictEqual(hasAiContributions([d2.uri], 'all'), false); - assert.strictEqual(hasAiContributions([d2.uri], 'chatAndAgent'), false); - - // After the stale entry is dropped, closing the document must not resurrect - // it: subsequent queries (against the now-closed URI) must still be false. - d2.dispose(); - await timeout(1500); - assert.strictEqual(hasAiContributions([fileA], 'all'), false); - assert.strictEqual(hasAiContributions([fileA], 'chatAndAgent'), false); - - disposables.dispose(); - reloadStore.dispose(); - })); - test('clear removes contributions for a closed (persisted-only) file', () => runWithFakedTimers({}, async () => { setup(); const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); @@ -369,27 +283,4 @@ suite('AiContributionFeature', () => { reloadStore.dispose(); })); - test('dispose snapshots live documents that had no pending save', () => runWithFakedTimers({}, async () => { - const reloadStore = new DisposableStore(); - const sharedStorage = reloadStore.add(new TestStorageService()); - - setup(sharedStorage); - const d = disposables.add(workspace.createDocument({ uri: fileA, initialValue: 'hello' }, undefined)); - await timeout(1500); - - d.applyEdit(StringEditWithReason.replace(d.findRange('hello'), 'world', chatEdit)); - await timeout(1500); // let the debounce flush normally - - // Now tear down with no new dirty state: the per-document `toDisposable` - // callbacks will re-snapshot, which must not be lost to the already- - // cancelled scheduler. - disposables.dispose(); - - setup(sharedStorage); - await timeout(1500); - - assert.strictEqual(hasAiContributions([fileA], 'all'), true); - disposables.dispose(); - reloadStore.dispose(); - })); }); From 4deee89b526618966d67ce7102e5114f52f82d31 Mon Sep 17 00:00:00 2001 From: cwebster-99 Date: Tue, 21 Apr 2026 17:43:07 -0500 Subject: [PATCH 4/5] address code review --- .../editTelemetry/test/browser/aiContributionFeature.test.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts index d0a4e96084109..4248ed5a21955 100644 --- a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts +++ b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts @@ -47,12 +47,10 @@ suite('AiContributionFeature', () => { correlationId: undefined, }); - let storage: TestStorageService; - function setup(sharedStorage?: TestStorageService): void { disposables = new DisposableStore(); const instantiationService = disposables.add(new TestInstantiationService(new ServiceCollection(), false, undefined, true)); - storage = sharedStorage ?? disposables.add(new TestStorageService()); + const storage = sharedStorage ?? disposables.add(new TestStorageService()); instantiationService.stub(IStorageService, storage); workspace = new MutableObservableWorkspace(); From e608489096381d57523c5a8294c2d07b6855754a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 21 Apr 2026 23:11:22 +0000 Subject: [PATCH 5/5] Address review: Object.create(null), URI.from strict, rename test Agent-Logs-Url: https://github.com/microsoft/vscode/sessions/05a5e7f2-2cda-4e62-9224-07c4a8763cc4 Co-authored-by: cwebster-99 <60238438+cwebster-99@users.noreply.github.com> --- .../contrib/editTelemetry/browser/aiContributionFeature.ts | 6 +++--- .../test/browser/aiContributionFeature.test.ts | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts index 954ad3260020a..dbeff51512c96 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts @@ -102,7 +102,7 @@ export class AiContributionFeature extends Disposable { private _hasAiContributions(resources: UriComponents[], level: AiContributionLevel): boolean { for (const r of resources) { - const recorded = this._contributions.get(URI.revive(r)); + const recorded = this._contributions.get(URI.from(r, true)); if (recorded === undefined) { continue; } @@ -122,7 +122,7 @@ export class AiContributionFeature extends Disposable { } } else { for (const r of resources) { - if (this._contributions.delete(URI.revive(r))) { + if (this._contributions.delete(URI.from(r, true))) { changed = true; } } @@ -167,7 +167,7 @@ export class AiContributionFeature extends Disposable { if (this._contributions.size === 0) { this._storageService.remove(STORAGE_KEY, StorageScope.WORKSPACE); } else { - const obj: Record = {}; + const obj: Record = Object.create(null); for (const [uri, level] of this._contributions) { obj[uri.toString()] = level; } diff --git a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts index 4248ed5a21955..de2bc9c908e8d 100644 --- a/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts +++ b/src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts @@ -214,7 +214,7 @@ suite('AiContributionFeature', () => { disposables.dispose(); })); - test('persisted AI ranges survive a workspace reload', () => runWithFakedTimers({}, async () => { + test('persisted AI contribution levels survive a workspace reload', () => runWithFakedTimers({}, async () => { const reloadStore = new DisposableStore(); const sharedStorage = reloadStore.add(new TestStorageService());