debug: fix error restarting debug session

src/vs/workbench/api/browser/mainThreadDebugService.ts

@@ -26,7 +26,7 @@ export class MainThreadDebugService implements MainThreadDebugServiceShape, IDeb
 	private _debugAdaptersHandleCounter = 1;
 	private readonly _debugConfigurationProviders: Map<number, IDebugConfigurationProvider>;
 	private readonly _debugAdapterDescriptorFactories: Map<number, IDebugAdapterDescriptorFactory>;
-	private readonly _sessions: Set<DebugSessionUUID>;
+	private readonly _extHostKnownSessions: Set<DebugSessionUUID>;
 
 	constructor(
 		extHostContext: IExtHostContext,
@@ -52,16 +52,22 @@ export class MainThreadDebugService implements MainThreadDebugServiceShape, IDeb
 			}
 			store.add(session.onDidCustomEvent(event => this._proxy.$acceptDebugSessionCustomEvent(this.getSessionDto(session), event)));
 		}));
-		this._toDispose.add(debugService.onDidEndSession(session => {
+		this._toDispose.add(debugService.onDidEndSession(({ session, restart }) => {
 			this._proxy.$acceptDebugSessionTerminated(this.getSessionDto(session));
-			this._sessions.delete(session.getId());
+			this._extHostKnownSessions.delete(session.getId());
+
+			// keep the session listeners around since we still will get events after they restart
+			if (!restart) {
+				sessionListeners.deleteAndDispose(session);
+			}
+
+			// any restarted session will create a new DA, so always throw the old one away.
 			for (const [handle, value] of this._debugAdapters) {
 				if (value.session === session) {
 					this._debugAdapters.delete(handle);
 					// break;
 				}
 			}
-			sessionListeners.deleteAndDispose(session);
 		}));
 		this._toDispose.add(debugService.getViewModel().onDidFocusSession(session => {
 			this._proxy.$acceptDebugSessionActiveChanged(this.getSessionDto(session));
@@ -75,7 +81,7 @@ export class MainThreadDebugService implements MainThreadDebugServiceShape, IDeb
 		this._debugAdapters = new Map();
 		this._debugConfigurationProviders = new Map();
 		this._debugAdapterDescriptorFactories = new Map();
-		this._sessions = new Set();
+		this._extHostKnownSessions = new Set();
 
 		this._toDispose.add(this.debugService.getViewModel().onDidFocusThread(({ thread, explicit, session }) => {
 			if (session) {
@@ -352,7 +358,7 @@ export class MainThreadDebugService implements MainThreadDebugServiceShape, IDeb
 
 	public $sessionCached(sessionID: string) {
 		// remember that the EH has cached the session and we do not have to send it again
-		this._sessions.add(sessionID);
+		this._extHostKnownSessions.add(sessionID);
 	}
 
 
@@ -362,7 +368,7 @@ export class MainThreadDebugService implements MainThreadDebugServiceShape, IDeb
 	getSessionDto(session: IDebugSession | undefined): IDebugSessionDto | undefined {
 		if (session) {
 			const sessionID = <DebugSessionUUID>session.getId();
-			if (this._sessions.has(sessionID)) {
+			if (this._extHostKnownSessions.has(sessionID)) {
 				return sessionID;
 			} else {
 				// this._sessions.add(sessionID); 	// #69534: see $sessionCached above

src/vs/workbench/contrib/audioCues/browser/audioCueDebuggerContribution.ts

@@ -41,7 +41,7 @@ export class AudioCueLineDebuggerContribution
 				)
 			);
 
-			store.add(debugService.onDidEndSession(session => {
+			store.add(debugService.onDidEndSession(({ session }) => {
 				sessionDisposables.get(session)?.dispose();
 				sessionDisposables.delete(session);
 			}));

src/vs/workbench/contrib/debug/browser/debugMemory.ts

@@ -31,7 +31,7 @@ export class DebugMemoryFileSystemProvider implements IFileSystemProvider {
 		| FileSystemProviderCapabilities.FileOpenReadWriteClose;
 
 	constructor(private readonly debugService: IDebugService) {
-		debugService.onDidEndSession(session => {
+		debugService.onDidEndSession(({ session }) => {
 			for (const [fd, memory] of this.fdMemory) {
 				if (memory.session === session) {
 					this.close(fd);

src/vs/workbench/contrib/debug/browser/debugService.ts

@@ -62,7 +62,8 @@ export class DebugService implements IDebugService {
 	private readonly _onDidChangeState: Emitter<State>;
 	private readonly _onDidNewSession: Emitter<IDebugSession>;
 	private readonly _onWillNewSession: Emitter<IDebugSession>;
-	private readonly _onDidEndSession: Emitter<IDebugSession>;
+	private readonly _onDidEndSession: Emitter<{ session: IDebugSession; restart: boolean }>;
+	private readonly restartingSessions = new Set<IDebugSession>();
 	private debugStorage: DebugStorage;
 	private model: DebugModel;
 	private viewModel: ViewModel;
@@ -114,7 +115,7 @@ export class DebugService implements IDebugService {
 		this._onDidChangeState = new Emitter<State>();
 		this._onDidNewSession = new Emitter<IDebugSession>();
 		this._onWillNewSession = new Emitter<IDebugSession>();
-		this._onDidEndSession = new Emitter<IDebugSession>();
+		this._onDidEndSession = new Emitter();
 
 		this.adapterManager = this.instantiationService.createInstance(AdapterManager, { onDidNewSession: this.onDidNewSession });
 		this.disposables.add(this.adapterManager);
@@ -315,7 +316,7 @@ export class DebugService implements IDebugService {
 		return this._onWillNewSession.event;
 	}
 
-	get onDidEndSession(): Event<IDebugSession> {
+	get onDidEndSession(): Event<{ session: IDebugSession; restart: boolean }> {
 		return this._onDidEndSession.event;
 	}
 
@@ -656,22 +657,29 @@ export class DebugService implements IDebugService {
 	}
 
 	private registerSessionListeners(session: DebugSession): void {
-		const sessionRunningScheduler = new RunOnceScheduler(() => {
+		const listenerDisposables = new DisposableStore();
+		this.disposables.add(listenerDisposables);
+
+		const sessionRunningScheduler = listenerDisposables.add(new RunOnceScheduler(() => {
 			// Do not immediatly defocus the stack frame if the session is running
 			if (session.state === State.Running && this.viewModel.focusedSession === session) {
 				this.viewModel.setFocus(undefined, this.viewModel.focusedThread, session, false);
 			}
-		}, 200);
-		this.disposables.add(session.onDidChangeState(() => {
+		}, 200));
+		listenerDisposables.add(session.onDidChangeState(() => {
 			if (session.state === State.Running && this.viewModel.focusedSession === session) {
 				sessionRunningScheduler.schedule();
 			}
 			if (session === this.viewModel.focusedSession) {
 				this.onStateChange();
 			}
 		}));
-
-		this.disposables.add(session.onDidEndAdapter(async adapterExitEvent => {
+		listenerDisposables.add(this.onDidEndSession(e => {
+			if (e.session === session && !e.restart) {
+				this.disposables.delete(listenerDisposables);
+			}
+		}));
+		listenerDisposables.add(session.onDidEndAdapter(async adapterExitEvent => {
 
 			if (adapterExitEvent) {
 				if (adapterExitEvent.error) {
@@ -696,7 +704,7 @@ export class DebugService implements IDebugService {
 			}
 			this.endInitializingState();
 			this.cancelTokens(session.getId());
-			this._onDidEndSession.fire(session);
+			this._onDidEndSession.fire({ session, restart: this.restartingSessions.has(session) });
 
 			const focusedSession = this.viewModel.focusedSession;
 			if (focusedSession && focusedSession.getId() === session.getId()) {
@@ -796,42 +804,66 @@ export class DebugService implements IDebugService {
 		}
 		session.configuration.__restart = restartData;
 
+		const doRestart = async (fn: () => Promise<boolean | undefined>) => {
+			this.restartingSessions.add(session);
+			let didRestart = false;
+			try {
+				didRestart = (await fn()) !== false;
+			} catch (e) {
+				didRestart = false;
+				throw e;
+			} finally {
+				this.restartingSessions.delete(session);
+				// we previously may have issued an onDidEndSession with restart: true,
+				// assuming the adapter exited (in `registerSessionListeners`). But the
+				// restart failed, so emit the final termination now.
+				if (!didRestart) {
+					this._onDidEndSession.fire({ session, restart: false });
+				}
+			}
+		};
+
 		if (session.capabilities.supportsRestartRequest) {
 			const taskResult = await runTasks();
 			if (taskResult === TaskRunResult.Success) {
-				await session.restart();
+				await doRestart(async () => {
+					await session.restart();
+					return true;
+				});
 			}
 
 			return;
 		}
 
 		const shouldFocus = !!this.viewModel.focusedSession && session.getId() === this.viewModel.focusedSession.getId();
-		// If the restart is automatic  -> disconnect, otherwise -> terminate #55064
-		if (isAutoRestart) {
-			await session.disconnect(true);
-		} else {
-			await session.terminate(true);
-		}
+		return doRestart(async () => {
+			// If the restart is automatic  -> disconnect, otherwise -> terminate #55064
+			if (isAutoRestart) {
+				await session.disconnect(true);
+			} else {
+				await session.terminate(true);
+			}
 
-		return new Promise<void>((c, e) => {
-			setTimeout(async () => {
-				const taskResult = await runTasks();
-				if (taskResult !== TaskRunResult.Success) {
-					return;
-				}
+			return new Promise<boolean>((c, e) => {
+				setTimeout(async () => {
+					const taskResult = await runTasks();
+					if (taskResult !== TaskRunResult.Success) {
+						return c(false);
+					}
 
-				if (!resolved) {
-					return c(undefined);
-				}
+					if (!resolved) {
+						return c(false);
+					}
 
-				try {
-					await this.launchOrAttachToSession(session, shouldFocus);
-					this._onDidNewSession.fire(session);
-					c(undefined);
-				} catch (error) {
-					e(error);
-				}
-			}, 300);
+					try {
+						await this.launchOrAttachToSession(session, shouldFocus);
+						this._onDidNewSession.fire(session);
+						c(true);
+					} catch (error) {
+						e(error);
+					}
+				}, 300);
+			});
 		});
 	}
 

src/vs/workbench/contrib/debug/browser/loadedScriptsView.ts

@@ -571,7 +571,7 @@ export class LoadedScriptsView extends ViewPane {
 		this._register(this.debugService.onDidNewSession(registerSessionListeners));
 		this.debugService.getModel().getSessions().forEach(registerSessionListeners);
 
-		this._register(this.debugService.onDidEndSession(session => {
+		this._register(this.debugService.onDidEndSession(({ session }) => {
 			root.remove(session.getId());
 			this.changeScheduler.schedule();
 		}));

src/vs/workbench/contrib/debug/browser/repl.ts

@@ -179,7 +179,7 @@ export class Repl extends FilterViewPane implements IHistoryNavigationWidget {
 			}
 			this.multiSessionRepl.set(this.isMultiSessionView);
 		}));
-		this._register(this.debugService.onDidEndSession(async session => {
+		this._register(this.debugService.onDidEndSession(async () => {
 			// Update view, since orphaned sessions might now be separate
 			await Promise.resolve(); // allow other listeners to go first, so sessions can update parents
 			this.multiSessionRepl.set(this.isMultiSessionView);

src/vs/workbench/contrib/debug/common/debug.ts

@@ -1008,19 +1008,24 @@ export interface IDebugService {
 	onDidChangeState: Event<State>;
 
 	/**
-	 * Allows to register on new session events.
+	 * Allows to register on sessions about to be created (not yet fully initialised).
+	 * This is fired exactly one time for any given session.
 	 */
-	onDidNewSession: Event<IDebugSession>;
+	onWillNewSession: Event<IDebugSession>;
 
 	/**
-	 * Allows to register on sessions about to be created (not yet fully initialised)
+	 * Fired when a new debug session is started. This may fire multiple times
+	 * for a single session due to restarts.
 	 */
-	onWillNewSession: Event<IDebugSession>;
+	onDidNewSession: Event<IDebugSession>;
 
 	/**
 	 * Allows to register on end session events.
+	 *
+	 * Contains a boolean indicating whether the session will restart. If restart
+	 * is true, the session should not considered to be dead yet.
 	 */
-	onDidEndSession: Event<IDebugSession>;
+	onDidEndSession: Event<{ session: IDebugSession; restart: boolean }>;
 
 	/**
 	 * Gets the configuration manager.

src/vs/workbench/contrib/debug/test/common/mockDebug.ts

@@ -33,7 +33,7 @@ export class MockDebugService implements IDebugService {
 		throw new Error('not implemented');
 	}
 
-	get onDidEndSession(): Event<IDebugSession> {
+	get onDidEndSession(): Event<{ session: IDebugSession; restart: boolean }> {
 		throw new Error('not implemented');
 	}
 

src/vs/workbench/contrib/inlineChat/browser/inlineChatDecorations.ts

@@ -79,14 +79,14 @@ export class InlineChatDecorationsContribution extends Disposable implements IEd
 				this._onEnablementOrModelChanged();
 			}
 		}));
-		this._register(this._debugService.onWillNewSession((session) => {
+		this._register(this._debugService.onDidNewSession((session) => {
 			this._debugSessions.add(session);
 			if (!this._hasActiveDebugSession) {
 				this._hasActiveDebugSession = true;
 				this._onEnablementOrModelChanged();
 			}
 		}));
-		this._register(this._debugService.onDidEndSession((session) => {
+		this._register(this._debugService.onDidEndSession(({ session }) => {
 			this._debugSessions.delete(session);
 			if (this._debugSessions.size === 0) {
 				this._hasActiveDebugSession = false;

src/vs/workbench/contrib/remote/browser/urlFinder.ts

@@ -52,7 +52,7 @@ export class UrlFinder extends Disposable {
 				}));
 			}
 		}));
-		this._register(debugService.onDidEndSession(session => {
+		this._register(debugService.onDidEndSession(({ session }) => {
 			if (this.listeners.has(session.getId())) {
 				this.listeners.get(session.getId())?.dispose();
 				this.listeners.delete(session.getId());