Refactors AI feedback context handling for changelog docs

Centralizes management of AI feedback context and changelog document URIs
into a dedicated provider, improving modularity and cleanup. Updates
relevant usages to access changelog feedback context through the new
provider, removing redundant context management logic from command files.

Simplifies resource cleanup on document close and ensures consistent
context availability for AI-related features.
(#4449)

Author: sergeibbb

src/commands/aiFeedback.ts

@@ -49,7 +49,7 @@ 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 {

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,6 +1,5 @@
 import type { Disposable, QuickInputButton } from 'vscode';
 import { env, ThemeIcon, Uri, window } from 'vscode';
-import { getChangelogFeedbackContext } from '../../../../commands/generateChangelog';
 import { Schemes } from '../../../../constants';
 import type { AIProviders } from '../../../../constants.ai';
 import type { Container } from '../../../../container';
@@ -282,7 +281,7 @@ 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;
@@ -299,7 +298,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.getChangelogFeedback(uri.toString());
 		} catch {
 			return undefined;
 		}

src/telemetry/aiFeedbackProvider.ts

@@ -0,0 +1,67 @@
+import type { Disposable, Uri } from 'vscode';
+import { workspace } from 'vscode';
+import type { AIResultContext } from '../plus/ai/aiProviderService';
+import { setContext } from '../system/-webview/context';
+import type { Deferrable } from '../system/function/debounce';
+import { debounce } from '../system/function/debounce';
+
+export class AIFeedbackProvider implements Disposable {
+	constructor() {
+		// Listen for document close events to clean up contexts
+		this._disposables.push(
+			workspace.onDidCloseTextDocument(document => this.removeChangelogDocument(document.uri)),
+		);
+	}
+
+	public addChangelogDocument(uri: Uri, context: AIResultContext): void {
+		this.setChangelogFeedback(uri.toString(), context);
+		this.addChangelogUri(uri);
+	}
+
+	private removeChangelogDocument(uri: Uri): void {
+		this.deleteChangelogFeedback(uri.toString());
+		this.removeChangelogUri(uri);
+	}
+
+	private readonly _disposables: Disposable[] = [];
+	dispose(): void {
+		this._disposables.forEach(d => void d.dispose());
+		this._changelogFeedbacks.clear();
+		this._changelogUris.clear();
+		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 _changelogFeedbacks = new Map<string, AIResultContext>();
+	getChangelogFeedback(documentUri: string): AIResultContext | undefined {
+		return this._changelogFeedbacks.get(documentUri);
+	}
+	private setChangelogFeedback(documentUri: string, context: AIResultContext): void {
+		this._changelogFeedbacks.set(documentUri, context);
+	}
+	private deleteChangelogFeedback(documentUri: string): void {
+		this._changelogFeedbacks.delete(documentUri);
+	}
+}