Track all extension hosts in runtime startup service (#5939)

src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts

@@ -1165,11 +1165,15 @@ export class MainThreadLanguageRuntime
 		this._proxy = extHostContext.getProxy(ExtHostPositronContext.ExtHostLanguageRuntime);
 		this._id = MainThreadLanguageRuntime.MAX_ID++;
 
-		this._runtimeStartupService.onDidChangeRuntimeStartupPhase((phase) => {
-			if (phase === RuntimeStartupPhase.Discovering) {
-				this._proxy.$discoverLanguageRuntimes();
-			}
-		});
+		this._runtimeStartupService.registerMainThreadLanguageRuntime(this._id);
+
+		this._disposables.add(
+			this._runtimeStartupService.onDidChangeRuntimeStartupPhase((phase) => {
+				if (phase === RuntimeStartupPhase.Discovering) {
+					this._proxy.$discoverLanguageRuntimes();
+				}
+			})
+		);
 
 		this._disposables.add(this._runtimeSessionService.registerSessionManager(this));
 	}
@@ -1245,7 +1249,7 @@ export class MainThreadLanguageRuntime
 
 	// Signals that language runtime discovery is complete.
 	$completeLanguageRuntimeDiscovery(): void {
-		this._runtimeStartupService.completeDiscovery();
+		this._runtimeStartupService.completeDiscovery(this._id);
 	}
 
 	$unregisterLanguageRuntime(handle: number): void {
@@ -1275,6 +1279,8 @@ export class MainThreadLanguageRuntime
 				session.emitExit(exit);
 			}
 		});
+
+		this._runtimeStartupService.unregisterMainThreadLanguageRuntime(this._id);
 		this._disposables.dispose();
 	}
 

src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts

@@ -84,6 +84,12 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup
 	// (metadata.languageId) of the runtime.
 	private readonly _mostRecentlyStartedRuntimesByLanguageId = new Map<string, ILanguageRuntimeMetadata>();
 
+	// A map of each extension host and its runtime discovery completion state.
+	// This is keyed by the the extension host's mainThreadLanguageRuntime's id
+	// This map is used to determine if runtime discovery has been completed
+	// across all extension hosts.
+	private readonly _discoveryCompleteByExtHostId = new Map<number, boolean>();
+
 	// The current startup phase; an observeable value.
 	private _startupPhase: ISettableObservable<RuntimeStartupPhase>;
 
@@ -344,11 +350,65 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup
 	}
 
 	/**
-	 * Completes the language runtime discovery phase. If no runtimes were
-	 * started or will be started, automatically start one.
+	 * Signals that the runtime discovery phase is completed only after all
+	 * extension hosts have completed runtime discovery.
+	 *
+	 * If no runtimes were started or will be started, automatically start one.
+	 */
+	public completeDiscovery(id: number): void {
+		// Update the extension host's runtime discovery state to 'Complete'
+		this._discoveryCompleteByExtHostId.set(id, true);
+		this._logService.debug(`[Runtime startup] Discovery completed for extension host with id: ${id}.`);
+
+		// Determine if all extension hosts have completed discovery
+		let discoveryCompletedByAllExtensionHosts = true;
+		for (const disoveryCompleted of this._discoveryCompleteByExtHostId.values()) {
+			if (!disoveryCompleted) {
+				discoveryCompletedByAllExtensionHosts = false;
+				break;
+			}
+		}
+
+		// The 'Discovery' phase is considered complete only after all extension hosts
+		// have signaled they have completed their own runtime discovery
+		if (discoveryCompletedByAllExtensionHosts) {
+			this._startupPhase.set(RuntimeStartupPhase.Complete, undefined);
+			// Reset the discovery state for each ext host so we are ready
+			// for possible re-discovery of runtimes
+			this._discoveryCompleteByExtHostId.forEach((_, extHostId, m) => {
+				m.set(extHostId, false);
+			});
+		}
+	}
+
+	/**
+	 * Used to register an instance of a MainThreadLanguageRuntime.
+	 *
+	 * This is required because there can be multiple extension hosts
+	 * and the startup service needs to know of all of them to track
+	 * the startup phase across all extension hosts.
+	 *
+	 * @param id The id of the MainThreadLanguageRuntime instance being registered.
+	 */
+	public registerMainThreadLanguageRuntime(id: number): void {
+		// Add the mainThreadLanguageRuntime instance id to the set of mainThreadLanguageRuntimes.
+		this._discoveryCompleteByExtHostId.set(id, false);
+		this._logService.debug(`[Runtime startup] Registered extension host with id: ${id}.`);
+	}
+
+	/**
+	 * Used to un-registers an instance of a MainThreadLanguageRuntime.
+	 *
+	 * This is required because there can be multiple extension hosts
+	 * and the startup service needs to know of all of them to track
+	 * the startup phase across all extension hosts.
+	 *
+	 * @param id The id of the MainThreadLanguageRuntime instance being un-registered.
 	 */
-	completeDiscovery(): void {
-		this._startupPhase.set(RuntimeStartupPhase.Complete, undefined);
+	public unregisterMainThreadLanguageRuntime(id: number): void {
+		// Remove the mainThreadLanguageRuntime instance id to the set of mainThreadLanguageRuntimes.
+		this._discoveryCompleteByExtHostId.delete(id);
+		this._logService.debug(`[Runtime startup] Unregistered extension host with id: ${id}.`);
 	}
 
 	/**

src/vs/workbench/services/runtimeStartup/common/runtimeStartupService.ts

@@ -100,8 +100,31 @@ export interface IRuntimeStartupService {
 	getAffiliatedRuntimeMetadata(languageId: string): ILanguageRuntimeMetadata | undefined;
 
 	/**
-	 * Signal that discovery of language runtimes is complete. Called from the
-	 * extension host.
+	 * Signal that discovery of language runtimes is completed for an extension host.
+	 *
+	 * @param id the id of the MainThreadLanguageRuntime instance for the extension host
+	 */
+	completeDiscovery(id: number): void;
+
+	/**
+	 * Used to register an instance of a MainThreadLanguageRuntime.
+	 *
+	 * This is required because there can be multiple extension hosts
+	 * and the startup service needs to know of all of them to track
+	 * the startup phase across all extension hosts.
+	 *
+	 * @param id The id of the MainThreadLanguageRuntime instance for the extension host.
+	 */
+	registerMainThreadLanguageRuntime(id: number): void;
+
+	/**
+	 * Used to un-register an instance of a MainThreadLanguageRuntime.
+	 *
+	 * This is required because there can be multiple extension hosts
+	 * and the startup service needs to know of all of them to track
+	 * the startup phase across all extension hosts.
+	 *
+	 * @param id The id of the MainThreadLanguageRuntime instance for the extension host.
 	 */
-	completeDiscovery(): void;
+	unregisterMainThreadLanguageRuntime(id: number): void;
 }