Fix issues around comm and session lifecycles

extensions/positron-supervisor/src/KallichoreSession.ts

@@ -632,8 +632,14 @@ export class KallichoreSession implements JupyterLanguageRuntimeSession {
 	}
 
 	removeClient(id: string): void {
-		const commOpen = new CommCloseCommand(id);
-		this.sendCommand(commOpen);
+		// Ignore this if the session is already exited; an exited session has
+		// no clients
+		if (this._runtimeState === positron.RuntimeState.Exited) {
+			this.log(`Ignoring request to close comm ${id}; kernel has already exited`, vscode.LogLevel.Debug);
+			return;
+		}
+		const commClose = new CommCloseCommand(id);
+		this.sendCommand(commClose);
 	}
 
 	/**

src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts

@@ -9,7 +9,7 @@ import * as extHostProtocol from './extHost.positron.protocol.js';
 import { Emitter } from '../../../../base/common/event.js';
 import { DisposableStore, IDisposable } from '../../../../base/common/lifecycle.js';
 import { Disposable, LanguageRuntimeMessageType } from '../extHostTypes.js';
-import { RuntimeClientType } from './extHostTypes.positron.js';
+import { RuntimeClientState, RuntimeClientType, RuntimeExitReason } from './extHostTypes.positron.js';
 import { ExtHostRuntimeClientInstance } from './extHostClientInstance.js';
 import { ExtensionIdentifier, IExtensionDescription } from '../../../../platform/extensions/common/extensions.js';
 import { URI } from '../../../../base/common/uri.js';
@@ -351,7 +351,19 @@ export class ExtHostLanguageRuntime implements extHostProtocol.ExtHostLanguageRu
 
 		// Hook up the session end (exit) handler
 		session.onDidEndSession(exit => {
+			// Notify the main thread that the session has ended
 			this._proxy.$emitLanguageRuntimeExit(handle, exit);
+
+			// If the session isn't exiting in order to restart, then we need
+			// to clean up its resources.
+			if (exit.reason !== RuntimeExitReason.Restart) {
+				session.dispose();
+			}
+
+			// Note that we don't remove the session from the list of sessions;
+			// that would invalidate the handles of all subsequent sessions
+			// since we store them in an array. The session remains in an inert
+			// state.
 		});
 
 		// Register the runtime
@@ -887,6 +899,9 @@ export class ExtHostLanguageRuntime implements extHostProtocol.ExtHostLanguageRu
 		// Dispose the client instance when the runtime exits
 		this._runtimeSessions[handle].onDidChangeRuntimeState(state => {
 			if (state === RuntimeState.Exited) {
+				// Mark the client instance as already closed so disposal
+				// doesn't try to close it
+				clientInstance.setClientState(RuntimeClientState.Closed);
 				clientInstance.dispose();
 			}
 		});