Moves AI feedback context management into provider

(#4449, #4518)

Author: sergeibbb

src/commands/aiFeedback.ts

@@ -1,5 +1,5 @@
-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';
@@ -8,7 +8,6 @@ import type { QuickPickItemOfT } from '../quickpicks/items/common';
 import { command } from '../system/-webview/command';
 import { map } from '../system/iterable';
 import { Logger } from '../system/logger';
-import { createDisposable } from '../system/unifiedDisposable';
 import { ActiveEditorCommand } from './commandBase';
 import { getCommandUri } from './commandBase.utils';
 
@@ -42,31 +41,6 @@ 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);
-}
-
 async function sendFeedback(container: Container, uri: Uri, sentiment: AIFeedbackEvent['sentiment']): Promise<void> {
 	const context = extractAIResultContext(container, uri);
 	if (!context) return;

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/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 { getMarkdownDocument } from '../../../../commands/aiFeedback';
 import { Schemes } from '../../../../constants';
 import type { AIProviders } from '../../../../constants.ai';
 import type { Container } from '../../../../container';
@@ -289,7 +288,7 @@ export function extractAIResultContext(container: Container, uri: Uri | undefine
 		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);

src/telemetry/aiFeedbackProvider.ts

@@ -12,7 +12,9 @@ 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)),
+			workspace.onDidCloseTextDocument(document => {
+				this.removeDocument(document.uri);
+			}),
 		);
 	}
 
@@ -21,16 +23,19 @@ export class AIFeedbackProvider implements Disposable {
 		this.addChangelogUri(uri);
 	}
 
-	private removeChangelogDocument(uri: Uri): void {
-		this.deleteChangelogDocument(uri.toString());
+	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;
@@ -70,6 +75,18 @@ export class AIFeedbackProvider implements Disposable {
 		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;