Store AI result context in memory for markdown documents.

Thats because we cannot modify URI of document to review context after
the result is rendered.
(#4328, #4489)

Author: sergeibbb

src/commands/aiFeedback.ts

@@ -1,5 +1,5 @@
-import type { TextEditor, Uri } from 'vscode';
-import { window } from 'vscode';
+import type { Disposable, TextEditor, Uri } from 'vscode';
+import { window, workspace } from 'vscode';
 import type { AIFeedbackEvent, AIFeedbackUnhelpfulReasons, Source } from '../constants.telemetry';
 import type { Container } from '../container';
 import type { AIResultContext } from '../plus/ai/aiProviderService';
@@ -12,6 +12,7 @@ import type { Deferrable } from '../system/function/debounce';
 import { debounce } from '../system/function/debounce';
 import { filterMap, map } from '../system/iterable';
 import { Logger } from '../system/logger';
+import { createDisposable } from '../system/unifiedDisposable';
 import { ActiveEditorCommand } from './commandBase';
 import { getCommandUri } from './commandBase.utils';
 
@@ -45,6 +46,31 @@ 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;
 

src/commands/explainBase.ts

@@ -11,6 +11,7 @@ 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';
 
@@ -135,10 +136,14 @@ export abstract class ExplainCommandBase extends GlCommandBase {
 		result: AISummarizeResult,
 		metadata: MarkdownContentMetadata,
 	): void {
-		const metadataWithContext: MarkdownContentMetadata = { ...metadata, context: getAIResultContext(result) };
+		const context = getAIResultContext(result);
+		const metadataWithContext: MarkdownContentMetadata = { ...metadata, context: context };
 		const headerContent = getMarkdownHeaderContent(metadataWithContext, this.container.telemetry.enabled);
 		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.markdown.updateDocument(documentUri, content);
 	}
 

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

@@ -1,5 +1,6 @@
 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,8 +290,9 @@ export function extractAIResultContext(uri: Uri | undefined): AIResultContext |
 	if (!authority) return undefined;
 
 	try {
+		const context: AIResultContext | undefined = getMarkdownDocument(uri.toString());
 		const metadata = decodeGitLensRevisionUriAuthority<MarkdownContentMetadata>(authority);
-		return metadata.context;
+		return context ?? metadata.context;
 	} catch (ex) {
 		Logger.error(ex, 'extractResultContext');
 		return undefined;