From 065e6e7cb455c8fd7aeb807693773ca5217e8999 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Thu, 28 Nov 2024 12:40:53 +0200
Subject: [PATCH] add `onDidEndSession` disposable to the session's disposables

---
 .../services/runtimeSession/common/runtimeSession.ts   | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
index 9fed7c92f46..2af31f2a661 100644
--- a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
+++ b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
@@ -369,14 +369,20 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 
 	private async doShutdownRuntimeSession(
 		session: ILanguageRuntimeSession, exitReason: RuntimeExitReason): Promise<void> {
+
+		const sessionDisposables = this._sessionDisposables.get(session.sessionId);
+		if (!sessionDisposables) {
+			throw new Error(`No disposables found for session ${session.sessionId}`);
+		}
+
 		// We wait for `onDidEndSession()` rather than `RuntimeState.Exited`, because the former
 		// generates some Console output that must finish before starting up a new runtime:
 		let disposable: IDisposable | undefined;
 		const promise = new Promise<void>(resolve => {
-			disposable = session.onDidEndSession((exit) => {
+			disposable = sessionDisposables.add(session.onDidEndSession((exit) => {
 				resolve();
 				disposable?.dispose();
-			});
+			}));
 		});
 
 		const timeout = new Promise<void>((_, reject) => {