diff --git a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts index 6dcef778f840a..dbeff51512c96 100644 --- a/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts +++ b/src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts @@ -3,113 +3,181 @@ * 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 } 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 { 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 have received AI-generated edits, so the git extension + * can add a `Co-authored-by: Copilot` trailer to commits that touch them. + * + * 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 { - private readonly _trackers = new ResourceMap(); - private readonly _documentsByUri = new ResourceMap(); + private readonly _contributions = new ResourceMap(); + + private readonly _saveScheduler: RunOnceScheduler; + private _dirty = 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(); - - for (const doc of docs) { - const uri = doc.document.uri; - activeUris.set(uri, true); - this._documentsByUri.set(uri, doc); + this._loadFromStorage(); - if (!this._trackers.has(uri)) { - this._trackers.set(uri, this._createTrackerEntry(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(); } - - for (const [uri, entry] of this._trackers) { - if (!activeUris.has(uri)) { - entry.trackerStore.dispose(); - this._trackers.delete(uri); - this._documentsByUri.delete(uri); - } - } - })); - - this._register(CommandsRegistry.registerCommand('_aiEdits.hasAiContributions', (_accessor, resources: UriComponents[], level: AiContributionLevel) => { - return this._hasAiContributions(resources, level); })); - this._register(CommandsRegistry.registerCommand('_aiEdits.clearAiContributions', (_accessor, resources: UriComponents[]) => { - this._clearAiContributions(resources); - })); + // 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) => { + store.add(runOnChange(doc.value, (_val, _prev, edits) => { + for (const e of edits) { + const source = EditSourceBase.create(e.reason); + if (source.category !== 'ai') { + continue; + } + const level: AiContributionLevel = source.feature === 'chat' ? 'chatAndAgent' : 'all'; + this._record(doc.uri, level); + } + })); + }); + + // Force the cached array so per-document subscriptions get wired up. + this._register(autorun(reader => { trackedDocs.read(reader); })); + + 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())); + } - this._register(CommandsRegistry.registerCommand('_aiEdits.clearAllAiContributions', () => { - this._clearAiContributions(); - })); + public override dispose(): void { + this._saveScheduler.cancel(); + super.dispose(); + if (this._dirty) { + this._saveToStorage(); + } } - override dispose(): void { - for (const [, entry] of this._trackers) { - entry.trackerStore.dispose(); + 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; } - super.dispose(); + this._contributions.set(uri, level); + this._markDirty(); } - 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 _markDirty(): void { + this._dirty = true; + 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 recorded = this._contributions.get(URI.from(r, true)); + if (recorded === undefined) { + continue; + } + if (level === 'all' || recorded === 'chatAndAgent') { + 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)); - } else { - this._trackers.delete(uri); - this._documentsByUri.delete(uri); + let changed = false; + if (!resources) { + if (this._contributions.size > 0) { + this._contributions.clear(); + changed = true; + } + } else { + for (const r of resources) { + if (this._contributions.delete(URI.from(r, true))) { + 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)) { + if (value !== 'chatAndAgent' && value !== 'all') { + continue; + } + let uri: URI; + try { + uri = URI.parse(uriString); + } catch { + continue; + } + this._contributions.set(uri, value); + } + } + + private _saveToStorage(): void { + // Only clear `_dirty` after a successful write so retries can flush + // pending state on the next save attempt. + try { + if (this._contributions.size === 0) { + this._storageService.remove(STORAGE_KEY, StorageScope.WORKSPACE); + } else { + const obj: Record = Object.create(null); + for (const [uri, level] of this._contributions) { + obj[uri.toString()] = level; } + // 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; + } catch { + // Keep `_dirty` so the next save can retry. } } } 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..de2bc9c908e8d 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,14 @@ suite('AiContributionFeature', () => { correlationId: undefined, }); - function setup(): void { + 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()); + const 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 +172,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 +185,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 +213,72 @@ suite('AiContributionFeature', () => { assert.strictEqual(hasAiContributions([fileB], 'all'), false); disposables.dispose(); })); + + test('persisted AI contribution levels 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('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(); + })); + });