AI feedback disposable provider

src/commands/aiFeedback.ts

@@ -1,18 +1,13 @@
-import type { Disposable, TextEditor, Uri } from 'vscode';
-import { window, workspace } from 'vscode';
+import type { TextEditor, Uri } from 'vscode';
+import { window } from 'vscode';
 import type { AIFeedbackEvent, AIFeedbackUnhelpfulReasons, Source } from '../constants.telemetry';
 import type { Container } from '../container';
 import type { AIResultContext } from '../plus/ai/aiProviderService';
 import { extractAIResultContext } from '../plus/ai/utils/-webview/ai.utils';
 import type { QuickPickItemOfT } from '../quickpicks/items/common';
 import { command } from '../system/-webview/command';
-import { setContext } from '../system/-webview/context';
-import { UriMap } from '../system/-webview/uriMap';
-import type { Deferrable } from '../system/function/debounce';
-import { debounce } from '../system/function/debounce';
-import { filterMap, map } from '../system/iterable';
+import { map } from '../system/iterable';
 import { Logger } from '../system/logger';
-import { createDisposable } from '../system/unifiedDisposable';
 import { ActiveEditorCommand } from './commandBase';
 import { getCommandUri } from './commandBase.utils';
 
@@ -46,57 +41,20 @@ export class AIFeedbackUnhelpfulCommand extends ActiveEditorCommand {
 
 type UnhelpfulResult = { reasons?: AIFeedbackUnhelpfulReasons[]; custom?: string };
 
-let _documentCloseTracker: Disposable | undefined;
-const _markdownDocuments = new Map<string, AIResultContext>();
-export function getMarkdownDocument(documentUri: string): AIResultContext | undefined {
-	return _markdownDocuments.get(documentUri);
-}
-export function setMarkdownDocument(documentUri: string, context: AIResultContext, container: Container): void {
-	_markdownDocuments.set(documentUri, context);
-
-	if (!_documentCloseTracker) {
-		_documentCloseTracker = workspace.onDidCloseTextDocument(document => {
-			deleteMarkdownDocument(document.uri.toString());
-		});
-		container.context.subscriptions.push(
-			createDisposable(() => {
-				_documentCloseTracker?.dispose();
-				_documentCloseTracker = undefined;
-				_markdownDocuments.clear();
-			}),
-		);
-	}
-}
-function deleteMarkdownDocument(documentUri: string): void {
-	_markdownDocuments.delete(documentUri);
-}
-
-const uriResponses = new UriMap<AIFeedbackEvent['sentiment']>();
-let _updateContextDebounced: Deferrable<() => void> | undefined;
-
 async function sendFeedback(container: Container, uri: Uri, sentiment: AIFeedbackEvent['sentiment']): Promise<void> {
-	const context = extractAIResultContext(uri);
+	const context = extractAIResultContext(container, uri);
 	if (!context) return;
 
 	try {
-		const previous = uriResponses.get(uri);
+		const previous = container.aiFeedback.getFeedbackResponse(uri);
 		if (sentiment === previous) return;
 
 		let unhelpful: UnhelpfulResult | undefined;
 		if (sentiment === 'unhelpful') {
 			unhelpful = await showUnhelpfulFeedbackPicker();
 		}
 
-		uriResponses.set(uri, sentiment);
-		_updateContextDebounced ??= debounce(() => {
-			void setContext('gitlens:tabs:ai:helpful', [
-				...filterMap(uriResponses, ([uri, sentiment]) => (sentiment === 'helpful' ? uri : undefined)),
-			]);
-			void setContext('gitlens:tabs:ai:unhelpful', [
-				...filterMap(uriResponses, ([uri, sentiment]) => (sentiment === 'unhelpful' ? uri : undefined)),
-			]);
-		}, 100);
-		_updateContextDebounced();
+		container.aiFeedback.setFeedbackResponse(uri, sentiment);
 
 		sendFeedbackEvent(container, { source: 'ai:markdown-preview' }, context, sentiment, unhelpful);
 	} catch (ex) {

src/commands/explainBase.ts

@@ -11,7 +11,6 @@ import type { AIModel } from '../plus/ai/models/model';
 import { getAIResultContext } from '../plus/ai/utils/-webview/ai.utils';
 import { getBestRepositoryOrShowPicker } from '../quickpicks/repositoryPicker';
 import { showMarkdownPreview } from '../system/-webview/markdown';
-import { setMarkdownDocument } from './aiFeedback';
 import { GlCommandBase } from './commandBase';
 import { getCommandUri } from './commandBase.utils';
 
@@ -142,7 +141,7 @@ export abstract class ExplainCommandBase extends GlCommandBase {
 		const content = `${headerContent}\n\n${result.parsed.summary}\n\n${result.parsed.body}`;
 
 		// Store the AI result context in the feedback provider for documents that cannot store it in their URI
-		setMarkdownDocument(documentUri.toString(), context, this.container);
+		this.container.aiFeedback.setMarkdownDocument(documentUri.toString(), context);
 
 		this.container.markdown.updateDocument(documentUri, content);
 	}

src/commands/generateChangelog.ts

@@ -1,18 +1,15 @@
-import type { CancellationToken, ProgressOptions, Uri } from 'vscode';
+import type { CancellationToken, ProgressOptions } from 'vscode';
 import { ProgressLocation, window, workspace } from 'vscode';
 import type { Source } from '../constants.telemetry';
 import type { Container } from '../container';
 import type { GitReference } from '../git/models/reference';
 import { getChangesForChangelog } from '../git/utils/-webview/log.utils';
 import { createRevisionRange, shortenRevision } from '../git/utils/revision.utils';
 import { showGenericErrorMessage } from '../messages';
-import type { AIGenerateChangelogChanges, AIResultContext } from '../plus/ai/aiProviderService';
+import type { AIGenerateChangelogChanges } from '../plus/ai/aiProviderService';
 import { getAIResultContext } from '../plus/ai/utils/-webview/ai.utils';
 import { showComparisonPicker } from '../quickpicks/comparisonPicker';
 import { command } from '../system/-webview/command';
-import { setContext } from '../system/-webview/context';
-import type { Deferrable } from '../system/function/debounce';
-import { debounce } from '../system/function/debounce';
 import type { Lazy } from '../system/lazy';
 import { lazy } from '../system/lazy';
 import { Logger } from '../system/logger';
@@ -25,36 +22,6 @@ export interface GenerateChangelogCommandArgs {
 	source?: Source;
 }
 
-// Storage for AI feedback context associated with changelog documents
-const changelogFeedbackContexts = new Map<string, AIResultContext>();
-export function getChangelogFeedbackContext(documentUri: string): AIResultContext | undefined {
-	return changelogFeedbackContexts.get(documentUri);
-}
-function setChangelogFeedbackContext(documentUri: string, context: AIResultContext): void {
-	changelogFeedbackContexts.set(documentUri, context);
-}
-function clearChangelogFeedbackContext(documentUri: string): void {
-	changelogFeedbackContexts.delete(documentUri);
-}
-
-// Storage for changelog document URIs
-const changelogUris = new Set<Uri>();
-let _updateChangelogContextDebounced: Deferrable<() => void> | undefined;
-function updateChangelogContext(): void {
-	_updateChangelogContextDebounced ??= debounce(() => {
-		void setContext('gitlens:tabs:ai:changelog', [...changelogUris]);
-	}, 100);
-	_updateChangelogContextDebounced();
-}
-function addChangelogUri(uri: Uri): void {
-	changelogUris.add(uri);
-	updateChangelogContext();
-}
-function removeChangelogUri(uri: Uri): void {
-	changelogUris.delete(uri);
-	updateChangelogContext();
-}
-
 @command()
 export class GenerateChangelogCommand extends GlCommandBase {
 	constructor(private readonly container: Container) {
@@ -158,17 +125,7 @@ export async function generateChangelogAndOpenMarkdownDocument(
 	const document = await workspace.openTextDocument({ language: 'markdown', content: content });
 	if (feedbackContext) {
 		// Store feedback context for this document
-		setChangelogFeedbackContext(document.uri.toString(), feedbackContext);
-		// Add to changelog URIs context even for no-results documents
-		addChangelogUri(document.uri);
-		// Clean up context when document is closed
-		const disposable = workspace.onDidCloseTextDocument(closedDoc => {
-			if (closedDoc.uri.toString() === document.uri.toString()) {
-				clearChangelogFeedbackContext(document.uri.toString());
-				removeChangelogUri(document.uri);
-				disposable.dispose();
-			}
-		});
+		container.aiFeedback.addChangelogDocument(document.uri, feedbackContext);
 	}
 	await window.showTextDocument(document);
 }

src/container.ts

@@ -57,6 +57,7 @@ import type { Storage } from './system/-webview/storage';
 import { memoize } from './system/decorators/-webview/memoize';
 import { log } from './system/decorators/log';
 import { Logger } from './system/logger';
+import { AIFeedbackProvider } from './telemetry/aiFeedbackProvider';
 import { TelemetryService } from './telemetry/telemetry';
 import { UsageTracker } from './telemetry/usageTracker';
 import { isWalkthroughSupported, WalkthroughStateProvider } from './telemetry/walkthroughStateProvider';
@@ -365,6 +366,14 @@ export class Container {
 		return this._ai;
 	}
 
+	private _aiFeedback: AIFeedbackProvider | undefined;
+	get aiFeedback(): AIFeedbackProvider {
+		if (this._aiFeedback == null) {
+			this._disposables.push((this._aiFeedback = new AIFeedbackProvider()));
+		}
+		return this._aiFeedback;
+	}
+
 	private _autolinks: AutolinksProvider | undefined;
 	get autolinks(): AutolinksProvider {
 		if (this._autolinks == null) {

src/plus/ai/utils/-webview/ai.utils.ts

@@ -1,7 +1,5 @@
 import type { Disposable, QuickInputButton } from 'vscode';
 import { env, ThemeIcon, Uri, window } from 'vscode';
-import { getMarkdownDocument } from '../../../../commands/aiFeedback';
-import { getChangelogFeedbackContext } from '../../../../commands/generateChangelog';
 import { Schemes } from '../../../../constants';
 import type { AIProviders } from '../../../../constants.ai';
 import type { Container } from '../../../../container';
@@ -284,13 +282,13 @@ export function getAIResultContext(result: AIResult): AIResultContext {
 	};
 }
 
-export function extractAIResultContext(uri: Uri | undefined): AIResultContext | undefined {
+export function extractAIResultContext(container: Container, uri: Uri | undefined): AIResultContext | undefined {
 	if (uri?.scheme === Schemes.GitLensAIMarkdown) {
 		const { authority } = uri;
 		if (!authority) return undefined;
 
 		try {
-			const context: AIResultContext | undefined = getMarkdownDocument(uri.toString());
+			const context: AIResultContext | undefined = container.aiFeedback.getMarkdownDocument(uri.toString());
 			if (context) return context;
 
 			const metadata = decodeGitLensRevisionUriAuthority<MarkdownContentMetadata>(authority);
@@ -304,7 +302,7 @@ export function extractAIResultContext(uri: Uri | undefined): AIResultContext |
 	// Check for untitled documents with stored changelog feedback context
 	if (uri?.scheme === 'untitled') {
 		try {
-			return getChangelogFeedbackContext(uri.toString());
+			return container.aiFeedback.getChangelogDocument(uri.toString());
 		} catch {
 			return undefined;
 		}

src/telemetry/aiFeedbackProvider.ts

@@ -0,0 +1,114 @@
+import type { Disposable, Uri } from 'vscode';
+import { workspace } from 'vscode';
+import type { AIFeedbackEvent } from '../constants.telemetry';
+import type { AIResultContext } from '../plus/ai/aiProviderService';
+import { setContext } from '../system/-webview/context';
+import { UriMap } from '../system/-webview/uriMap';
+import type { Deferrable } from '../system/function/debounce';
+import { debounce } from '../system/function/debounce';
+import { filterMap } from '../system/iterable';
+
+export class AIFeedbackProvider implements Disposable {
+	constructor() {
+		// Listen for document close events to clean up contexts
+		this._disposables.push(
+			workspace.onDidCloseTextDocument(document => {
+				this.removeDocument(document.uri);
+			}),
+		);
+	}
+
+	public addChangelogDocument(uri: Uri, context: AIResultContext): void {
+		this.setChangelogDocument(uri.toString(), context);
+		this.addChangelogUri(uri);
+	}
+
+	private removeDocument(uri: Uri): void {
+		const uriString = uri.toString();
+		this.deleteChangelogDocument(uriString);
+		this.removeChangelogUri(uri);
+		this.deleteMarkdownDocument(uriString);
+	}
+
+	private readonly _disposables: Disposable[] = [];
+	dispose(): void {
+		this._disposables.forEach(d => void d.dispose());
+		this._uriResponses.clear();
+		this._changelogDocuments.clear();
+		this._markdownDocuments.clear();
+		this._changelogUris.clear();
+		this._updateFeedbackContextDebounced = undefined;
+		this._updateChangelogContextDebounced = undefined;
+	}
+
+	// Storage for changelog document URIs
+	private readonly _changelogUris = new Set<Uri>();
+	private _updateChangelogContextDebounced: Deferrable<() => void> | undefined;
+	private updateChangelogContext(): void {
+		this._updateChangelogContextDebounced ??= debounce(() => {
+			void setContext('gitlens:tabs:ai:changelog', [...this._changelogUris]);
+		}, 100);
+		this._updateChangelogContextDebounced();
+	}
+	private addChangelogUri(uri: Uri): void {
+		if (!this._changelogUris.has(uri)) {
+			this._changelogUris.add(uri);
+			this.updateChangelogContext();
+		}
+	}
+	private removeChangelogUri(uri: Uri): void {
+		if (this._changelogUris.has(uri)) {
+			this._changelogUris.delete(uri);
+			this.updateChangelogContext();
+		}
+	}
+
+	// Storage for AI feedback context associated with changelog documents
+	private readonly _changelogDocuments = new Map<string, AIResultContext>();
+	getChangelogDocument(documentUri: string): AIResultContext | undefined {
+		return this._changelogDocuments.get(documentUri);
+	}
+	private setChangelogDocument(documentUri: string, context: AIResultContext): void {
+		this._changelogDocuments.set(documentUri, context);
+	}
+	private deleteChangelogDocument(documentUri: string): void {
+		this._changelogDocuments.delete(documentUri);
+	}
+
+	// Storage for AI feedback context associated with any document
+	private readonly _markdownDocuments = new Map<string, AIResultContext>();
+	getMarkdownDocument(documentUri: string): AIResultContext | undefined {
+		return this._markdownDocuments.get(documentUri);
+	}
+	setMarkdownDocument(documentUri: string, context: AIResultContext): void {
+		this._markdownDocuments.set(documentUri, context);
+	}
+	private deleteMarkdownDocument(documentUri: string): void {
+		this._markdownDocuments.delete(documentUri);
+	}
+
+	// Storage for AI feedback responses by URI
+	private readonly _uriResponses = new UriMap<AIFeedbackEvent['sentiment']>();
+	private _updateFeedbackContextDebounced: Deferrable<() => void> | undefined;
+	private updateFeedbackContext(): void {
+		this._updateFeedbackContextDebounced ??= debounce(() => {
+			void setContext('gitlens:tabs:ai:helpful', [
+				...filterMap(this._uriResponses, ([uri, sentiment]) => (sentiment === 'helpful' ? uri : undefined)),
+			]);
+			void setContext('gitlens:tabs:ai:unhelpful', [
+				...filterMap(this._uriResponses, ([uri, sentiment]) => (sentiment === 'unhelpful' ? uri : undefined)),
+			]);
+		}, 100);
+		this._updateFeedbackContextDebounced();
+	}
+	setFeedbackResponse(uri: Uri, sentiment: AIFeedbackEvent['sentiment']): void {
+		const previous = this._uriResponses.get(uri);
+		if (sentiment === previous) return;
+
+		this._uriResponses.set(uri, sentiment);
+		this.updateFeedbackContext();
+	}
+	getFeedbackResponse(uri: Uri): AIFeedbackEvent['sentiment'] | undefined {
+		return this._uriResponses.get(uri);
+	}
+}