Refactor AI CoAuthor Attribution#311796
Conversation
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <copilot@github.com>
There was a problem hiding this comment.
Pull request overview
Refactors AI co-author attribution tracking to persist AI contributions per-URI in workspace storage, so Git’s addAICoAuthor behavior survives document close and window reloads.
Changes:
- Replace per-open-document tracking with a persisted
ResourceMap<AiContributionLevel>keyed by URI and backed byIStorageService. - Debounce storage writes and flush pending state on workspace save and on dispose; add a command to clear all contributions.
- Update contribution wiring and tests to use
ObservableWorkspace+ storage-backed persistence semantics.
Show a summary per file
| File | Description |
|---|---|
| src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts | Implements per-URI contribution recording, persistence, debounced saving, and updated commands. |
| src/vs/workbench/contrib/editTelemetry/browser/editTelemetryContribution.ts | Updates feature instantiation to pass ObservableWorkspace instead of annotated documents. |
| src/vs/workbench/contrib/editTelemetry/test/browser/aiContributionFeature.test.ts | Updates and extends tests for persistence across close/reload, clearing persisted-only entries, and dispose flush behavior. |
Copilot's findings
- Files reviewed: 3/3 changed files
- Comments generated: 1
Screenshot ChangesBase: Changed (1)
|
| Before | After |
|---|---|
 |
 |
There was a problem hiding this comment.
Copilot's findings
Comments suppressed due to low confidence (1)
src/vs/workbench/contrib/editTelemetry/browser/aiContributionFeature.ts:128
- Same concern here:
_clearAiContributionsrevives untrusted URI components viaURI.revive. UseURI.from(..., true)(or otherwise validate) to avoid accepting invalid/malicious URIs and to keep persisted keys stable.
for (const r of resources) {
if (this._contributions.delete(URI.revive(r))) {
changed = true;
}
}
- Files reviewed: 3/3 changed files
- Comments generated: 3
| const obj: Record<string, AiContributionLevel> = {}; | ||
| 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); |
There was a problem hiding this comment.
When building the persisted snapshot, using a normal object ({}) with dynamic string keys can be vulnerable to prototype-pollution edge cases (e.g. a crafted URI string of __proto__) and can also behave oddly with such keys. Use a null-prototype dictionary (Object.create(null)) or persist as an array of [uriString, level] tuples instead.
There was a problem hiding this comment.
Fixed in e608489 — switched to Object.create(null) for the persisted snapshot object.
| 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); |
There was a problem hiding this comment.
Test name still refers to “AI ranges”, but the feature now persists per-URI contribution levels (not tracked ranges). Renaming this test (and any related comments) would keep the test intent aligned with the new behavior.
There was a problem hiding this comment.
Fixed in e608489 — renamed to "persisted AI contribution levels survive a workspace reload".
| 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.revive(r)); | ||
| if (recorded === undefined) { | ||
| continue; | ||
| } | ||
| if (level === 'all' || recorded === 'chatAndAgent') { | ||
| return true; |
There was a problem hiding this comment.
resources are command arguments coming from extensions (untrusted URI components). URI.revive explicitly does no validation and is intended for URI#toJSON data; using it here can accept malformed URIs and create surprising keys/collisions (and interacts badly with the object-key persistence below). Prefer URI.from(resource, /*strict*/ true) (or validate with isUriComponents + URI.from) when reviving command args.
This issue also appears on line 124 of the same file.
There was a problem hiding this comment.
Fixed in e608489 — replaced URI.revive(r) with URI.from(r, true) on both lines 105 and 125.
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>
|
Just as a heads up, I was blocked by some firewall rules while working on your feedback. Expand below for details. Warning Firewall rules blocked me from connecting to one or more addresses (expand for details)I tried to connect to the following addresses, but was blocked by firewall rules:
If you need me to access, download, or install something from one of these locations, you can either:
|
Persistence and Attribution Tracking:
aiContributionFeature.tsnow records AI contributions per-URI in aResourceMap, persisting this state in workspace storage so that attributions survive document closure and window reloads.API and Command Changes:
Dependency and Construction Updates:
ObservableWorkspaceandIStorageServiceinstead of annotated documents, updating all instantiations and tests accordingly. [1] [2]Testing Improvements:
Code Cleanup: