Adding promises and operation timeouts to fix race conditions and infinite loading in OE (#22475)

* Adding promises and operation timeouts to fix race conditions

* cleaning up logic

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Fixing promise type

* Reverting back to old error logic

* Making onsessioncreated async

* Removed polling and converted to event based

* removing connection variable out of promise

* Combining promises

* Update src/sql/workbench/services/objectExplorer/browser/treeUpdateUtils.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Fixing error messages and localizing user facing errors

* Fixing error message

* localizing config

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Update src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

* Fixing comment

---------

Co-authored-by: Charles Gagnon <chgagnon@microsoft.com>

src/sql/workbench/contrib/objectExplorer/common/serverGroup.contribution.ts

@@ -41,9 +41,10 @@ const serverGroupConfig: IConfigurationNode = {
 	}
 };
 
+export const NODE_EXPANSION_CONFIG = 'serverTree.nodeExpansionTimeout';
 const serverTreeConfig: IConfigurationNode = {
 	'id': 'serverTree',
-	'title': 'Server Tree',
+	'title': localize('serverTree.configuration.title', "Server Tree"),
 	'type': 'object',
 	'properties': {
 		'serverTree.useAsyncServerTree': {

src/sql/workbench/services/objectExplorer/browser/asyncServerTreeDataSource.ts

@@ -13,6 +13,7 @@ import Severity from 'vs/base/common/severity';
 import { IErrorMessageService } from 'sql/platform/errorMessage/common/errorMessageService';
 import { IAsyncDataSource } from 'vs/base/browser/ui/tree/tree';
 import { ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 
 /**
  * Implements the DataSource(that returns a parent/children of an element) for the server tree
@@ -22,6 +23,7 @@ export class AsyncServerTreeDataSource implements IAsyncDataSource<ConnectionPro
 	constructor(
 		@IObjectExplorerService private _objectExplorerService: IObjectExplorerService,
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
+		@IConfigurationService private _configurationService: IConfigurationService,
 		@IErrorMessageService private _errorMessageService: IErrorMessageService
 	) {
 	}
@@ -45,9 +47,9 @@ export class AsyncServerTreeDataSource implements IAsyncDataSource<ConnectionPro
 	public async getChildren(element: ServerTreeElement): Promise<ServerTreeElement[]> {
 		try {
 			if (element instanceof ConnectionProfile) {
-				return await TreeUpdateUtils.getAsyncConnectionNodeChildren(element, this._connectionManagementService, this._objectExplorerService);
+				return await TreeUpdateUtils.getAsyncConnectionNodeChildren(element, this._connectionManagementService, this._objectExplorerService, this._configurationService);
 			} else if (element instanceof ConnectionProfileGroup) {
-				return (element as ConnectionProfileGroup).getChildren();
+				return element.getChildren();
 			} else if (element instanceof TreeNode) {
 				if (element.children) {
 					return element.children;

src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

@@ -26,6 +26,7 @@ import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { ObjectExplorerRequestStatus } from 'sql/workbench/services/objectExplorer/browser/treeSelectionHandler';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { INotificationService } from 'vs/platform/notification/common/notification';
+import { NODE_EXPANSION_CONFIG } from 'sql/workbench/contrib/objectExplorer/common/serverGroup.contribution';
 
 export const SERVICE_ID = 'ObjectExplorerService';
 
@@ -174,6 +175,10 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	private _onSelectionOrFocusChange: Emitter<void>;
 	private _onNodeExpandedError: Emitter<NodeExpandInfoWithProviderId> = new Emitter<NodeExpandInfoWithProviderId>();
 
+	private _onCreateNewSession: Emitter<azdata.ObjectExplorerSessionResponse> = new Emitter<azdata.ObjectExplorerSessionResponse>();
+
+	private _connectionsWaitingForSession: Map<string, boolean> = new Map<string, boolean>();
+
 	constructor(
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IAdsTelemetryService private _telemetryService: IAdsTelemetryService,
@@ -264,9 +269,9 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	/**
 	 * Gets called when session is created
 	 */
-	public onSessionCreated(handle: number, session: azdata.ObjectExplorerSession): void {
+	public async onSessionCreated(handle: number, session: azdata.ObjectExplorerSession): Promise<void> {
 		if (session && session.success) {
-			this.handleSessionCreated(session).catch((e) => this.logService.error(e));
+			await this.handleSessionCreated(session).catch((e) => this.logService.error(e));
 		} else {
 			let errorMessage = session && session.errorMessage ? session.errorMessage : errSessionCreateFailed;
 			this.logService.error(errorMessage);
@@ -276,35 +281,85 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	private async handleSessionCreated(session: azdata.ObjectExplorerSession): Promise<void> {
 		let connection: ConnectionProfile | undefined = undefined;
 		let errorMessage: string | undefined = undefined;
-		if (this._sessions[session.sessionId!]) {
+		const sessionId = session.sessionId;
-			connection = this._sessions[session.sessionId!].connection;
 
-			try {
+
-				if (session.success && session.rootNode) {
+		await new Promise<void>((resolve, reject) => {
-					let server = this.toTreeNode(session.rootNode, undefined);
+			const cleanup = () => {
-					server.connection = connection;
+				if (connection) {
-					server.session = session;
+					this._connectionsWaitingForSession.delete(connection.id);
-					this._activeObjectExplorerNodes[connection!.id] = server;
 				}
-				else {
+				createNewSessionListener.dispose();
-					errorMessage = session && session.errorMessage ? session.errorMessage : errSessionCreateFailed;
+				clearTimeout(timeoutHandle);
-					this.logService.error(errorMessage);
-				}
-				// Send on session created about the session to all node providers so they can prepare for node expansion
-				let nodeProviders = this._nodeProviders[connection!.providerName];
-				if (nodeProviders) {
-					const promises = nodeProviders.map(p => p.handleSessionOpen(session));
-					await Promise.all(promises);
-				}
-			} catch (error) {
-				this.logService.warn(`cannot handle the session ${session.sessionId} in all nodeProviders`);
-			} finally {
-				this.sendUpdateNodeEvent(connection!, errorMessage);
 			}
+			const onTimeout = () => {
+				if (!this._sessions[sessionId]) {
+					this.logService.error(`Timed out waiting for session ${sessionId} to be created.
+					This has probably happened because OE service did not recieve a response for createNewSession from the provider.`);
+					reject(new Error(
+						nls.localize('objectExplorerMissingSession',
+							'Timed out waiting for session {0} to be created. This has probably happened because OE service did not recieve a response for createNewSession from the provider.', sessionId)));
+				} else {
+					this.logService.error(`Timeout waiting for session ${sessionId} to be created for connection "${connection.title}".
+					This has probably happened because OE service did not recieve a response for createNewSession from the provider for connection."${connection.title}`);
+					reject(new Error(nls.localize(
+						'objectExplorerMissingConnectionForSession',
+						'Timeout waiting for session {0} to be created for connection "{1}". This has probably happened because OE service did not recieve a response for createNewSession from the provider for connection "{1}"', sessionId, connection.title
+					)));
+				}
+
+				cleanup();
+			}
+			const timeoutHandle = setTimeout(onTimeout, this.getObjectExplorerTimeout() * 1000);
+			const createNewSessionListener = this._onCreateNewSession.event((response) => {
+				checkSessionAndConnection();
+			});
+			const checkSessionAndConnection = () => {
+				/**
+				 * Sometimes ads recieves handleSessionCreated from providers before the createNewSession response is recieved.
+				 * We need to wait while the createNewSession response is recieved and the session map contains the session before we can continue.
+				 */
+				if (this._sessions[sessionId]) {
+					connection = this._sessions[sessionId].connection;
+					/**
+					 * In certain cases, when we try to connect to a previously connected server, we may encounter a situation where the session is present in this._sessions,
+					 * probably becaue the close session request was not completed successfully with the same session id for an older connection. While creating this new session,
+					 * if we recieve the handleSessionCreated event before the createNewSession response is recieved, we will end up using the older connection stored in
+					 * this._sessions[sessionId].connection. To avoid this, we check if the connection id is false in this._connectionsWaitingForSession. If it is not false,
+					 * we know that the createNewSession response has been recieved and we have the correct connection.
+					 */
+					if (connection && this._connectionsWaitingForSession.get(connection.id) === false) {
+						resolve();
+						cleanup();
+					}
+				}
+			}
+			checkSessionAndConnection();
+		});
+
+		try {
+			if (session.success && session.rootNode) {
+				let server = this.toTreeNode(session.rootNode, undefined);
+				server.connection = connection;
+				server.session = session;
+				this._activeObjectExplorerNodes[connection!.id] = server;
+			}
+			else {
+				errorMessage = session && session.errorMessage ? session.errorMessage : errSessionCreateFailed;
+				this.logService.error(errorMessage);
+			}
+			// Send on session created about the session to all node providers so they can prepare for node expansion
+			let nodeProviders = this._nodeProviders[connection!.providerName];
+			if (nodeProviders) {
+				const promises = nodeProviders.map(p => p.handleSessionOpen(session));
+				await Promise.all(promises);
+			}
+		} catch (error) {
+			this.logService.error(`An error occured while handling session ${sessionId} in all nodeProviders.`, error);
+		} finally {
+			this.sendUpdateNodeEvent(connection!, errorMessage);
 		}
-		else {
+
-			this.logService.warn(`cannot find session ${session.sessionId}`);
-		}
 	}
 
 	/**
@@ -360,7 +415,14 @@ export class ObjectExplorerService implements IObjectExplorerService {
 	public async createNewSession(providerId: string, connection: ConnectionProfile): Promise<azdata.ObjectExplorerSessionResponse> {
 		const provider = this._providers[providerId];
 		if (provider) {
+			// set the connection to wait for session
+			this._connectionsWaitingForSession.set(connection.id, true);
 			const result = await provider.createNewSession(connection.toConnectionInfo());
+			// some providers return a malformed create sessions responses which don't have a session id. We should throw an error in this case
+			if (!result?.sessionId) {
+				this.logService.error(`The session ID returned by provider "${providerId}" for connection "${connection.title}" is invalid.`);
+				throw new Error(nls.localize('objectExplorerSessionIdMissing', 'The session ID returned by provider "{0}" for connection "{1}" is invalid.', providerId, connection.title));
+			}
 			if (this._sessions[result.sessionId]) {
 				this.logService.trace(`Overwriting session ${result.sessionId}`);
 			}
@@ -368,8 +430,12 @@ export class ObjectExplorerService implements IObjectExplorerService {
 				connection: connection,
 				nodes: {}
 			};
+			// once the session is created, set the connection to not wait for session
+			this._connectionsWaitingForSession.set(connection.id, false);
+			this._onCreateNewSession.fire(result);
 			return result;
 		} else {
+			this._connectionsWaitingForSession.delete(connection.id);
 			throw new Error(`Provider doesn't exist. id: ${providerId}`);
 		}
 	}
@@ -446,7 +512,6 @@ export class ObjectExplorerService implements IObjectExplorerService {
 						}
 					});
 
-					const expansionTimeoutValueSec = this._configurationService.getValue<number>('serverTree.nodeExpansionTimeout');
 					const expansionTimeout = setTimeout(() => {
 						/**
 						 * If we don't get a response back from all the providers in specified expansion timeout seconds then we assume
@@ -458,7 +523,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 							this._notificationService.error(nls.localize('nodeExpansionTimeout', "Node expansion timed out for node {0} for providers {1}", node.nodePath, missingProviders.map(p => p.providerId).join(', ')));
 						}
 						resolveExpansion();
-					}, expansionTimeoutValueSec * 1000);
+					}, this.getObjectExplorerTimeout() * 1000);
 
 					self._sessions[session.sessionId!].nodes[node.nodePath].expandEmitter.event((expandResult: NodeExpandInfoWithProviderId) => {
 						if (expandResult && expandResult.providerId) {
@@ -932,4 +997,11 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		}
 		return currentNode;
 	}
+
+	/**
+	 * returns object explorer timeout in seconds.
+	 */
+	public getObjectExplorerTimeout(): number {
+		return this._configurationService.getValue<number>(NODE_EXPANSION_CONFIG);
+	}
 }

src/sql/workbench/services/objectExplorer/browser/treeUpdateUtils.ts

@@ -15,6 +15,9 @@ import { Disposable, isDisposable } from 'vs/base/common/lifecycle';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { AsyncServerTree, ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';
 import { ObjectExplorerRequestStatus } from 'sql/workbench/services/objectExplorer/browser/treeSelectionHandler';
+import * as nls from 'vs/nls';
+import { NODE_EXPANSION_CONFIG } from 'sql/workbench/contrib/objectExplorer/common/serverGroup.contribution';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 
 export interface IExpandableTree extends ITree {
 	/**
@@ -222,8 +225,13 @@ export class TreeUpdateUtils {
 	 * @param connectionManagementService Connection management service instance
 	 * @param objectExplorerService Object explorer service instance
 	 */
-	public static async connectAndCreateOeSession(connection: ConnectionProfile, options: IConnectionCompletionOptions,
+	public static async connectAndCreateOeSession(
-		connectionManagementService: IConnectionManagementService, objectExplorerService: IObjectExplorerService, tree: AsyncServerTree | ITree | undefined, requestStatus?: ObjectExplorerRequestStatus | undefined): Promise<boolean> {
+		connection: ConnectionProfile,
+		options: IConnectionCompletionOptions,
+		connectionManagementService: IConnectionManagementService,
+		objectExplorerService: IObjectExplorerService,
+		tree: AsyncServerTree | ITree | undefined,
+		requestStatus?: ObjectExplorerRequestStatus | undefined): Promise<boolean> {
 		const connectedConnection = await TreeUpdateUtils.connectIfNotConnected(connection, options, connectionManagementService, tree);
 		if (connectedConnection) {
 			// append group ID and original display name to build unique OE session ID
@@ -264,7 +272,12 @@ export class TreeUpdateUtils {
 		}
 	}
 
-	public static async getAsyncConnectionNodeChildren(connection: ConnectionProfile, connectionManagementService: IConnectionManagementService, objectExplorerService: IObjectExplorerService): Promise<TreeNode[]> {
+	public static async getAsyncConnectionNodeChildren(
+		connection: ConnectionProfile,
+		connectionManagementService: IConnectionManagementService,
+		objectExplorerService: IObjectExplorerService,
+		configurationService: IConfigurationService
+	): Promise<TreeNode[]> {
 		if (connection.isDisconnecting) {
 			return [];
 		} else {
@@ -281,14 +294,31 @@ export class TreeUpdateUtils {
 					showFirewallRuleOnError: true,
 					showDashboard: false
 				};
+				const expansionTimeoutValueSec = configurationService.getValue<number>(NODE_EXPANSION_CONFIG);
 				// Need to wait for the OE service to update its nodes in order to resolve the children
-				const nodesUpdatedPromise = new Promise((resolve, reject) => {
+				const nodesUpdatedPromise = new Promise<void>((resolve, reject) => {
-					objectExplorerService.onUpdateObjectExplorerNodes(e => {
+					// Clean up timeout and listener
-						if (e.errorMessage) {
+					const cleanup = () => {
-							reject(new Error(e.errorMessage));
+						clearTimeout(nodeUpdateTimer);
-						}
+						nodesUpdatedListener.dispose();
-						if (e.connection.id === connection.id) {
+					}
-							resolve(undefined);
+
+					// If the node update takes too long, reject the promise
+					const nodeUpdateTimeout = () => {
+						reject(new Error(nls.localize('objectExplorerTimeout', "Object Explorer expansion timed out for '{0}'", connection.databaseName)));
+						cleanup();
+					}
+					const nodeUpdateTimer = setTimeout(nodeUpdateTimeout, expansionTimeoutValueSec * 1000);
+
+
+					const nodesUpdatedListener = objectExplorerService.onUpdateObjectExplorerNodes(e => {
+						if (e.connection && e.connection.id === connection.id) {
+							if (e.errorMessage) {
+								reject(new Error(e.errorMessage));
+							} else {
+								resolve();
+							}
+							cleanup();
 						}
 					});
 				});

src/sql/workbench/services/objectExplorer/test/browser/objectExplorerService.test.ts

@@ -327,7 +327,7 @@ suite('SQL Object Explorer Service tests', () => {
 		const session = await objectExplorerService.createNewSession(mssqlProviderName, connection);
 		assert.strictEqual(session !== null || session !== undefined, true);
 		assert.strictEqual(session.sessionId, '1234');
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const node = objectExplorerService.getObjectExplorerNode(connection);
 		assert.notStrictEqual(node, undefined);
 		assert.strictEqual(node.session.success, true);
@@ -338,7 +338,7 @@ suite('SQL Object Explorer Service tests', () => {
 		assert.strictEqual(session !== null || session !== undefined, true);
 		assert.strictEqual(session.sessionId, failedSessionId);
 		const currentNumberOfSuccessfulSessions = numberOfSuccessfulSessions;
-		objectExplorerService.onSessionCreated(1, objectExplorerFailedSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerFailedSession);
 		const node = objectExplorerService.getObjectExplorerNode(connection);
 		assert.strictEqual(node, undefined);
 		assert.strictEqual(currentNumberOfSuccessfulSessions, numberOfSuccessfulSessions);
@@ -354,7 +354,7 @@ suite('SQL Object Explorer Service tests', () => {
 	test('expand node should expand node correctly', async () => {
 		const tablesNode = new TreeNode(NodeType.Folder, '', 'Tables', false, 'testServerName/tables', 'testServerName', '', '', null, null, undefined, undefined);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const expandInfo = await objectExplorerService.expandNode(mssqlProviderName, objectExplorerSession, tablesNode);
 		assert.strictEqual(expandInfo !== null || expandInfo !== undefined, true);
 		assert.strictEqual(expandInfo.sessionId, '1234');
@@ -367,7 +367,7 @@ suite('SQL Object Explorer Service tests', () => {
 	test('refresh node should refresh node correctly', async () => {
 		const tablesNode = new TreeNode(NodeType.Folder, '', 'Tables', false, 'testServerName/tables', 'testServerName', '', '', null, null, undefined, undefined);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const expandInfo = await objectExplorerService.refreshNode(mssqlProviderName, objectExplorerSession, tablesNode);
 		assert.strictEqual(expandInfo !== null || expandInfo !== undefined, true);
 		assert.strictEqual(expandInfo.sessionId, '1234');
@@ -381,7 +381,7 @@ suite('SQL Object Explorer Service tests', () => {
 		const tablesNode = new TreeNode(NodeType.Folder, '', 'Tables', false, 'testServerName/tables', 'testServerName', '', '', null, null, undefined, undefined);
 		tablesNode.connection = connection;
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const children = await objectExplorerService.resolveTreeNodeChildren(objectExplorerSession, tablesNode);
 		assert.strictEqual(children !== null || children !== undefined, true);
 		assert.strictEqual(children[0].label, 'dbo.Table1');
@@ -396,7 +396,7 @@ suite('SQL Object Explorer Service tests', () => {
 		const tablesNode = new TreeNode(NodeType.Folder, '', 'Tables', false, 'testServerName/tables', 'testServerName', '', '', null, null, undefined, undefined);
 		tablesNode.connection = connection;
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const children = await objectExplorerService.refreshTreeNode(objectExplorerSession, tablesNode);
 		assert.strictEqual(children !== null || children !== undefined, true);
 		assert.strictEqual(children[0].label, 'dbo.Table1');
@@ -409,7 +409,7 @@ suite('SQL Object Explorer Service tests', () => {
 
 	test('update object explorer nodes should get active connection, create session, add to the active OE nodes successfully', async () => {
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		await objectExplorerService.updateObjectExplorerNodes(connection);
 		const treeNode = objectExplorerService.getObjectExplorerNode(connection);
 		assert.strictEqual(treeNode !== null || treeNode !== undefined, true);
@@ -421,7 +421,7 @@ suite('SQL Object Explorer Service tests', () => {
 
 	test('delete object explorerNode nodes should delete session, delete the root node to the active OE node', async () => {
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		await objectExplorerService.updateObjectExplorerNodes(connection);
 		let treeNode = objectExplorerService.getObjectExplorerNode(connection);
 		assert.strictEqual(treeNode !== null && treeNode !== undefined, true);
@@ -522,7 +522,7 @@ suite('SQL Object Explorer Service tests', () => {
 		});
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const childNodes = await objectExplorerService.resolveTreeNodeChildren(objectExplorerSession, objectExplorerService.getObjectExplorerNode(connection));
 		sqlOEProvider.setup(x => x.expandNode(TypeMoq.It.isAny())).callback(() => {
 			objectExplorerService.onNodeExpanded(tableExpandInfo);
@@ -540,7 +540,7 @@ suite('SQL Object Explorer Service tests', () => {
 		});
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const childNodes = await objectExplorerService.resolveTreeNodeChildren(objectExplorerSession, objectExplorerService.getObjectExplorerNode(connection));
 		// If I check whether the table is expanded, the answer should be no because only its parent node is expanded
 		const tableNode = childNodes.find(node => node.nodePath === table1NodePath);
@@ -562,7 +562,7 @@ suite('SQL Object Explorer Service tests', () => {
 		});
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const childNodes = await objectExplorerService.resolveTreeNodeChildren(objectExplorerSession, objectExplorerService.getObjectExplorerNode(connection));
 		sqlOEProvider.setup(x => x.expandNode(TypeMoq.It.isAny())).callback(() => {
 			objectExplorerService.onNodeExpanded(tableExpandInfo);
@@ -596,7 +596,7 @@ suite('SQL Object Explorer Service tests', () => {
 		serverTreeView.setup(x => x.reveal(TypeMoq.It.isAny())).returns(() => Promise.resolve());
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		// If I expand the node, then it should get revealed and expanded
 		const tableNode = await objectExplorerService.getTreeNode(connection.id, table1NodePath);
 		await tableNode.setExpandedState(TreeItemCollapsibleState.Expanded);
@@ -611,7 +611,7 @@ suite('SQL Object Explorer Service tests', () => {
 		serverTreeView.setup(x => x.setExpandedState(TypeMoq.It.is(treeNode => treeNode === connection), TypeMoq.It.is(state => state === TreeItemCollapsibleState.Collapsed))).returns(() => Promise.resolve());
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		await objectExplorerService.resolveTreeNodeChildren(objectExplorerSession, objectExplorerService.getObjectExplorerNode(connection));
 		// If I collapse the connection node, then the tree's collapse method should get called
 		const treeNode = await objectExplorerService.getTreeNode(connection.id, undefined);
@@ -625,7 +625,7 @@ suite('SQL Object Explorer Service tests', () => {
 		serverTreeView.setup(x => x.reveal(TypeMoq.It.isAny())).returns(() => Promise.resolve());
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		// If I select the table node, then it should be selected and revealed
 		const tableNode = await objectExplorerService.getTreeNode(connection.id, table1NodePath);
 		await tableNode.setSelected(true);
@@ -636,7 +636,7 @@ suite('SQL Object Explorer Service tests', () => {
 	test('findTreeNode returns the tree node for the relevant node', async () => {
 		const table1NodePath = objectExplorerExpandInfo.nodes[0].nodePath;
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const treeNode = await objectExplorerService.getTreeNode(connection.id, table1NodePath);
 		assert.strictEqual(treeNode.nodePath, objectExplorerExpandInfo.nodes[0].nodePath);
 		assert.strictEqual(treeNode.nodeTypeId, objectExplorerExpandInfo.nodes[0].nodeType);
@@ -646,7 +646,7 @@ suite('SQL Object Explorer Service tests', () => {
 	test('findTreeNode returns undefined if the requested node does not exist', async () => {
 		const invalidNodePath = objectExplorerSession.rootNode.nodePath + '/invalidNode';
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		const nodeInfo = await objectExplorerService.getTreeNode(connection.id, invalidNodePath);
 		assert.strictEqual(nodeInfo, undefined);
 	});
@@ -654,7 +654,7 @@ suite('SQL Object Explorer Service tests', () => {
 	test('refreshInView refreshes the node, expands it, and returns the refreshed node', async () => {
 		// Set up the session and tree view
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 		serverTreeView.setup(x => x.refreshElement(TypeMoq.It.isAny())).returns(() => Promise.resolve());
 		objectExplorerService.registerServerTreeView(serverTreeView.object);
 
@@ -673,7 +673,7 @@ suite('SQL Object Explorer Service tests', () => {
 
 		// Set up the session
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 
 		// Set up the provider to not respond to the second expand request, simulating a request that takes a long time to compconste
 		const nodePath = objectExplorerSession.rootNode.nodePath;
@@ -693,7 +693,7 @@ suite('SQL Object Explorer Service tests', () => {
 
 	test('resolveTreeNodeChildren refreshes a node if it currently has an error', async () => {
 		await objectExplorerService.createNewSession(mssqlProviderName, connection);
-		objectExplorerService.onSessionCreated(1, objectExplorerSession);
+		await objectExplorerService.onSessionCreated(1, objectExplorerSession);
 
 		// If I call resolveTreeNodeChildren once, set an error on the node, and then call it again
 		const tablesNodePath = 'testServerName/tables';