From c98f749f7e2a41e0d7c1798e33e4575a81a6c49b Mon Sep 17 00:00:00 2001
From: Peng Lyu <penn.lv@gmail.com>
Date: Wed, 29 May 2024 10:27:54 -0700
Subject: [PATCH] Fix #213525. Move nb chat variable to service (#213835)

---
 .../chat/notebook.chat.contribution.ts        | 34 +++++++++++++++++++
 .../controller/chat/notebookChatController.ts | 21 ++++++------
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebook.chat.contribution.ts b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebook.chat.contribution.ts
index 995340fb19188..f35264aa8cc80 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebook.chat.contribution.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebook.chat.contribution.ts
@@ -3,4 +3,38 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { Disposable } from 'vs/base/common/lifecycle';
+import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from 'vs/workbench/common/contributions';
+import { IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
 import 'vs/workbench/contrib/notebook/browser/controller/chat/cellChatActions';
+import { NotebookChatController } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController';
+import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
+
+class NotebookChatVariables extends Disposable implements IWorkbenchContribution {
+
+	static readonly ID = 'workbench.contrib.notebookChatVariables';
+
+	constructor(
+		@IChatVariablesService private readonly _chatVariableService: IChatVariablesService,
+		@INotebookEditorService private readonly _notebookEditorService: INotebookEditorService
+	) {
+		super();
+
+		this._register(this._chatVariableService.registerVariable(
+			{ id: '_notebookChatInput', name: '_notebookChatInput', description: '', hidden: true },
+			async (_message, _arg, model) => {
+				const editors = this._notebookEditorService.listNotebookEditors();
+				for (const editor of editors) {
+					const chatController = editor.getContribution(NotebookChatController.id) as NotebookChatController | undefined;
+					if (chatController?.hasSession(model)) {
+						return chatController.getSessionInputUri();
+					}
+				}
+
+				return undefined;
+			}
+		));
+	}
+}
+
+registerWorkbenchContribution2(NotebookChatVariables.ID, NotebookChatVariables, WorkbenchPhase.BlockRestore);
diff --git a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
index e5c83779451a3..4662ab11501b1 100644
--- a/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
+++ b/src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts
@@ -31,9 +31,8 @@ import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/c
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
 import { ChatAgentLocation } from 'vs/workbench/contrib/chat/common/chatAgents';
-import { ChatModel } from 'vs/workbench/contrib/chat/common/chatModel';
+import { ChatModel, IChatModel } from 'vs/workbench/contrib/chat/common/chatModel';
 import { IChatService } from 'vs/workbench/contrib/chat/common/chatService';
-import { IChatVariablesService } from 'vs/workbench/contrib/chat/common/chatVariables';
 import { countWords } from 'vs/workbench/contrib/chat/common/chatWordCounter';
 import { ProgressingEditsOptions } from 'vs/workbench/contrib/inlineChat/browser/inlineChatStrategies';
 import { InlineChatWidget } from 'vs/workbench/contrib/inlineChat/browser/inlineChatWidget';
@@ -279,8 +278,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		@ILanguageService private readonly _languageService: ILanguageService,
 		@INotebookExecutionStateService private _executionStateService: INotebookExecutionStateService,
 		@IStorageService private readonly _storageService: IStorageService,
-		@IChatService private readonly _chatService: IChatService,
-		@IChatVariablesService private readonly _chatVariableService: IChatVariablesService,
+		@IChatService private readonly _chatService: IChatService
 	) {
 		super();
 		this._ctxHasActiveRequest = CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST.bindTo(this._contextKeyService);
@@ -301,13 +299,6 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			this._historyCandidate = '';
 			this._storageService.store(NotebookChatController._storageKey, JSON.stringify(NotebookChatController._promptHistory), StorageScope.PROFILE, StorageTarget.USER);
 		};
-
-		this._register(this._chatVariableService.registerVariable(
-			{ id: '_notebookChatInput', name: '_notebookChatInput', description: '', hidden: true },
-			async (_message, _arg, model) => {
-				return this._widget?.parentEditor.getModel()?.uri;
-			}
-		));
 	}
 
 	private _registerFocusTracker() {
@@ -541,6 +532,14 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		this._widget.updateNotebookEditorFocusNSelections();
 	}
 
+	hasSession(chatModel: IChatModel) {
+		return this._model.value === chatModel;
+	}
+
+	getSessionInputUri() {
+		return this._widget?.parentEditor.getModel()?.uri;
+	}
+
 	async acceptInput() {
 		assertType(this._widget);
 		await this._sessionCtor;