Enable RBAC support for MongoDB and Cassandra APIs (#2198)

* enable RBAC support for Mongo & Cassandra API

* fix formatting issue

* Handling AAD integration for Mongo Shell

* remove empty aadToken error

* fix formatting issue

* added environment specific scope endpoints

src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx

@@ -5,6 +5,7 @@
  */
 import { CommandBar as FluentCommandBar, ICommandBarItemProps } from "@fluentui/react";
 import { useNotebook } from "Explorer/Notebook/useNotebook";
+import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 import { KeyboardActionGroup, useKeyboardActionGroup } from "KeyboardShortcuts";
 import { isFabric } from "Platform/Fabric/FabricUtil";
 import { userContext } from "UserContext";
@@ -30,7 +31,7 @@ export interface CommandBarStore {
 }
 
 export const useCommandBar: UseStore<CommandBarStore> = create((set) => ({
-  contextButtons: [],
+  contextButtons: [] as CommandButtonComponentProps[],
   setContextButtons: (contextButtons: CommandButtonComponentProps[]) => set((state) => ({ ...state, contextButtons })),
   isHidden: false,
   setIsHidden: (isHidden: boolean) => set((state) => ({ ...state, isHidden })),
@@ -43,6 +44,15 @@ export const CommandBar: React.FC<Props> = ({ container }: Props) => {
   const backgroundColor = StyleConstants.BaseLight;
   const setKeyboardHandlers = useKeyboardActionGroup(KeyboardActionGroup.COMMAND_BAR);
 
+  // Subscribe to the store changes that affect button creation
+  const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled);
+  const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated);
+
+  // Memoize the expensive button creation
+  const staticButtons = React.useMemo(() => {
+    return CommandBarComponentButtonFactory.createStaticCommandBarButtons(container, selectedNodeState);
+  }, [container, selectedNodeState, dataPlaneRbacEnabled, aadTokenUpdated]);
+
   if (userContext.apiType === "Postgres" || userContext.apiType === "VCoreMongo") {
     const buttons =
       userContext.apiType === "Postgres"
@@ -62,7 +72,6 @@ export const CommandBar: React.FC<Props> = ({ container }: Props) => {
     );
   }
 
-  const staticButtons = CommandBarComponentButtonFactory.createStaticCommandBarButtons(container, selectedNodeState);
   const contextButtons = (buttons || []).concat(
     CommandBarComponentButtonFactory.createContextCommandBarButtons(container, selectedNodeState),
   );

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -1,7 +1,6 @@
 import { KeyboardAction } from "KeyboardShortcuts";
 import { isDataplaneRbacSupported } from "Utils/APITypeUtils";
 import * as React from "react";
-import { useEffect, useState } from "react";
 import AddSqlQueryIcon from "../../../../images/AddSqlQuery_16x16.svg";
 import AddStoredProcedureIcon from "../../../../images/AddStoredProcedure.svg";
 import AddTriggerIcon from "../../../../images/AddTrigger.svg";
@@ -68,15 +67,7 @@ export function createStaticCommandBarButtons(
   }
 
   if (isDataplaneRbacSupported(userContext.apiType)) {
-    const [loginButtonProps, setLoginButtonProps] = useState<CommandButtonComponentProps | undefined>(undefined);
-    const dataPlaneRbacEnabled = useDataPlaneRbac((state) => state.dataPlaneRbacEnabled);
-    const aadTokenUpdated = useDataPlaneRbac((state) => state.aadTokenUpdated);
-
-    useEffect(() => {
-      const buttonProps = createLoginForEntraIDButton(container);
-      setLoginButtonProps(buttonProps);
-    }, [dataPlaneRbacEnabled, aadTokenUpdated, container]);
-
+    const loginButtonProps = createLoginForEntraIDButton(container);
     if (loginButtonProps) {
       addDivider();
       buttons.push(loginButtonProps);

src/Explorer/Tables/TableDataClient.ts

@@ -13,7 +13,7 @@ import { updateDocument } from "../../Common/dataAccess/updateDocument";
 import { configContext } from "../../ConfigContext";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { userContext } from "../../UserContext";
-import { getAuthorizationHeader } from "../../Utils/AuthorizationUtils";
+import { getAuthorizationHeader, isDataplaneRbacEnabledForProxyApi } from "../../Utils/AuthorizationUtils";
 import * as NotificationConsoleUtils from "../../Utils/NotificationConsoleUtils";
 import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import Explorer from "../Explorer";
@@ -551,6 +551,10 @@ export class CassandraAPIDataClient extends TableDataClient {
     const authorizationHeaderMetadata: ViewModels.AuthorizationTokenHeaderMetadata = getAuthorizationHeader();
     xhr.setRequestHeader(authorizationHeaderMetadata.header, authorizationHeaderMetadata.token);
 
+    if (isDataplaneRbacEnabledForProxyApi(userContext)) {
+      xhr.setRequestHeader(Constants.HttpHeaders.entraIdToken, userContext.aadToken);
+    }
+
     return true;
   };
 

src/Explorer/Tabs/CloudShellTab/ShellTypes/AbstractShellHandler.tsx

@@ -24,7 +24,7 @@ export const EXIT_COMMAND_MONGO = ` printf "\\033[1;31mSession ended. Please clo
  * This command runs mongosh in no-database and quiet mode,
  * and evaluates the `disableTelemetry()` function to turn off telemetry collection.
  */
-export const DISABLE_TELEMETRY_COMMAND = `mongosh --nodb --quiet --eval "disableTelemetry()"`;
+export const DISABLE_TELEMETRY_COMMAND = `mongosh --nodb --quiet --eval 'disableTelemetry()'`;
 
 /**
  * Abstract class that defines the interface for shell-specific handlers
@@ -97,7 +97,7 @@ export abstract class AbstractShellHandler {
    * is not already present in the environment.
    */
   protected mongoShellSetupCommands(): string[] {
-    const PACKAGE_VERSION: string = "2.5.5";
+    const PACKAGE_VERSION: string = "2.5.6";
     return [
       "if ! command -v mongosh &> /dev/null; then echo '⚠️ mongosh not found. Installing...'; fi",
       `if ! command -v mongosh &> /dev/null; then curl -LO https://downloads.mongodb.com/compass/mongosh-${PACKAGE_VERSION}-linux-x64.tgz; fi`,

src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.test.tsx

@@ -18,6 +18,12 @@ interface DatabaseAccount {
 
 interface UserContextType {
   databaseAccount: DatabaseAccount;
+  features: {
+    enableAadDataPlane: boolean;
+  };
+  apiType: string;
+  dataPlaneRbacEnabled: boolean;
+  aadToken?: string;
 }
 
 // Mock dependencies
@@ -29,6 +35,8 @@ jest.mock("../../../../UserContext", () => ({
         mongoEndpoint: "https://test-mongo.documents.azure.com:443/",
       },
     },
+    features: { enableAadDataPlane: false },
+    apiType: "Mongo",
   },
 }));
 
@@ -70,7 +78,7 @@ describe("MongoShellHandler", () => {
 
       expect(Array.isArray(commands)).toBe(true);
       expect(commands.length).toBe(7);
-      expect(commands[1]).toContain("mongosh-2.5.5-linux-x64.tgz");
+      expect(commands[1]).toContain("mongosh-2.5.6-linux-x64.tgz");
     });
   });
 
@@ -88,11 +96,12 @@ describe("MongoShellHandler", () => {
         kind: "test-kind",
         properties: { mongoEndpoint: "https://test-mongo.documents.azure.com:443/" },
       };
+      (userContext as UserContextType).dataPlaneRbacEnabled = false;
 
       const command = mongoShellHandler.getConnectionCommand();
 
       expect(command).toBe(
-        'mongosh --nodb --quiet --eval "disableTelemetry()" && mongosh mongodb://test-mongo.documents.azure.com:10255?appName=CosmosExplorerTerminal --username test-account --password test-key --tls --tlsAllowInvalidCertificates',
+        "mongosh --nodb --quiet --eval 'disableTelemetry()'; mongosh mongodb://test-mongo.documents.azure.com:10255?appName=CosmosExplorerTerminal --username test-account --password test-key --tls --tlsAllowInvalidCertificates",
       );
       expect(CommonUtils.getHostFromUrl).toHaveBeenCalledWith("https://test-mongo.documents.azure.com:443/");
 
@@ -115,12 +124,47 @@ describe("MongoShellHandler", () => {
       };
 
       const command = mongoShellHandler.getConnectionCommand();
-
       expect(command).toBe("echo 'Database name not found.'");
 
       // Restore original
       (userContext as UserContextType).databaseAccount = originalDatabaseAccount;
     });
+
+    it("should return echo if endpoint is missing", () => {
+      const testKey = "test-key";
+      (userContext as UserContextType).databaseAccount = {
+        id: "test-id",
+        name: "", // Empty name to simulate missing name
+        location: "test-location",
+        type: "test-type",
+        kind: "test-kind",
+        properties: { mongoEndpoint: "" },
+      };
+      const mongoShellHandler = new MongoShellHandler(testKey);
+      const command = mongoShellHandler.getConnectionCommand();
+      expect(command).toBe("echo 'MongoDB endpoint not found.'");
+    });
+
+    it("should use _getAadConnectionCommand when _isEntraIdEnabled is true", () => {
+      const testKey = "aad-key";
+      (userContext as UserContextType).databaseAccount = {
+        id: "test-id",
+        name: "test-account",
+        location: "test-location",
+        type: "test-type",
+        kind: "test-kind",
+        properties: { mongoEndpoint: "https://test-mongo.documents.azure.com:443/" },
+      };
+      (userContext as UserContextType).dataPlaneRbacEnabled = true;
+
+      const mongoShellHandler = new MongoShellHandler(testKey);
+
+      const command = mongoShellHandler.getConnectionCommand();
+      expect(command).toContain(
+        "mongosh 'mongodb://test-account:aad-key@test-account.mongo.cosmos.azure.com:10255/?ssl=true&replicaSet=globaldb&authMechanism=PLAIN&retryWrites=false' --tls --tlsAllowInvalidCertificates",
+      );
+      expect(command.startsWith("mongosh --nodb")).toBeTruthy();
+    });
   });
 
   describe("getTerminalSuppressedData", () => {

src/Explorer/Tabs/CloudShellTab/ShellTypes/MongoShellHandler.tsx

@@ -1,4 +1,5 @@
 import { userContext } from "../../../../UserContext";
+import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils";
 import { filterAndCleanTerminalOutput, getHostFromUrl, getMongoShellRemoveInfoText } from "../Utils/CommonUtils";
 import { AbstractShellHandler, DISABLE_TELEMETRY_COMMAND, EXIT_COMMAND_MONGO } from "./AbstractShellHandler";
 
@@ -6,12 +7,23 @@ export class MongoShellHandler extends AbstractShellHandler {
   private _key: string;
   private _endpoint: string | undefined;
   private _removeInfoText: string[] = getMongoShellRemoveInfoText();
+  private _isEntraIdEnabled: boolean = isDataplaneRbacEnabledForProxyApi(userContext);
   constructor(private key: string) {
     super();
     this._key = key;
     this._endpoint = userContext?.databaseAccount?.properties?.mongoEndpoint;
   }
 
+  private _getKeyConnectionCommand(dbName: string): string {
+    return `mongosh mongodb://${getHostFromUrl(this._endpoint)}:10255?appName=${
+      this.APP_NAME
+    } --username ${dbName} --password ${this._key} --tls --tlsAllowInvalidCertificates`;
+  }
+
+  private _getAadConnectionCommand(dbName: string): string {
+    return `mongosh 'mongodb://${dbName}:${this._key}@${dbName}.mongo.cosmos.azure.com:10255/?ssl=true&replicaSet=globaldb&authMechanism=PLAIN&retryWrites=false' --tls --tlsAllowInvalidCertificates`;
+  }
+
   public getShellName(): string {
     return "MongoDB";
   }
@@ -29,19 +41,11 @@ export class MongoShellHandler extends AbstractShellHandler {
     if (!dbName) {
       return "echo 'Database name not found.'";
     }
-    return (
-      DISABLE_TELEMETRY_COMMAND +
-      " && " +
-      "mongosh mongodb://" +
-      getHostFromUrl(this._endpoint) +
-      ":10255?appName=" +
-      this.APP_NAME +
-      " --username " +
-      dbName +
-      " --password " +
-      this._key +
-      " --tls --tlsAllowInvalidCertificates"
-    );
+    const connectionCommand = this._isEntraIdEnabled
+      ? this._getAadConnectionCommand(dbName)
+      : this._getKeyConnectionCommand(dbName);
+    const fullCommand = `${DISABLE_TELEMETRY_COMMAND}; ${connectionCommand}`;
+    return fullCommand;
   }
 
   public getTerminalSuppressedData(): string[] {

src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.test.tsx

@@ -7,12 +7,24 @@ import { PostgresShellHandler } from "./PostgresShellHandler";
 import { getHandler, getKey } from "./ShellTypeFactory";
 import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler";
 
+interface UserContextType {
+  databaseAccount: { name: string };
+  subscriptionId: string;
+  resourceGroup: string;
+  features: { enableAadDataPlane: boolean };
+  dataPlaneRbacEnabled: boolean;
+  aadToken?: string;
+  apiType?: string;
+}
+
 // Mock dependencies
 jest.mock("../../../../UserContext", () => ({
   userContext: {
     databaseAccount: { name: "testDbName" },
     subscriptionId: "testSubId",
     resourceGroup: "testResourceGroup",
+    features: { enableAadDataPlane: false },
+    dataPlaneRbacEnabled: false,
   },
 }));
 
@@ -109,5 +121,33 @@ describe("ShellTypeHandlerFactory", () => {
       expect(key).toBe(mockKey);
       expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName");
     });
+
+    it("should return MongoShellHandler with primaryMasterKey for TerminalKind.Mongo when RBAC is disabled", async () => {
+      (listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: "primaryKey123" });
+      (userContext as UserContextType).features.enableAadDataPlane = false;
+      (userContext as UserContextType).dataPlaneRbacEnabled = false;
+      const handler = await getHandler(TerminalKind.Mongo);
+      expect(handler).toBeInstanceOf(MongoShellHandler);
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      expect(handler.key).toBe("primaryKey123");
+    });
+
+    it("should return MongoShellHandler with aadToken for TerminalKind.Mongo when RBAC is enabled", async () => {
+      (userContext as UserContextType).aadToken = "aadToken123";
+      (userContext as UserContextType).features.enableAadDataPlane = true;
+      (userContext as UserContextType).dataPlaneRbacEnabled = true;
+      (userContext as UserContextType).apiType = "Mongo";
+      const handler = await getHandler(TerminalKind.Mongo);
+      expect(handler).toBeInstanceOf(MongoShellHandler);
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      expect(handler.key).toBe("aadToken123");
+    });
+    it("should throw error for unsupported shell type", async () => {
+      await expect(getHandler("UnknownShell" as unknown as TerminalKind)).rejects.toThrow(
+        "Unsupported shell type: UnknownShell",
+      );
+    });
   });
 });

src/Explorer/Tabs/CloudShellTab/ShellTypes/ShellTypeFactory.tsx

@@ -1,6 +1,7 @@
 import { TerminalKind } from "../../../../Contracts/ViewModels";
 import { userContext } from "../../../../UserContext";
 import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
+import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils";
 import { AbstractShellHandler } from "./AbstractShellHandler";
 import { CassandraShellHandler } from "./CassandraShellHandler";
 import { MongoShellHandler } from "./MongoShellHandler";
@@ -30,6 +31,9 @@ export async function getKey(): Promise<string> {
   if (!dbName) {
     return "";
   }
+  if (isDataplaneRbacEnabledForProxyApi(userContext)) {
+    return userContext.aadToken || "";
+  }
 
   const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
   return keys?.primaryMasterKey || "";

src/Explorer/Tabs/CloudShellTab/ShellTypes/VCoreMongoShellHandler.test.tsx

@@ -45,7 +45,7 @@ describe("VCoreMongoShellHandler", () => {
 
       expect(Array.isArray(commands)).toBe(true);
       expect(commands.length).toBe(7);
-      expect(commands[1]).toContain("mongosh-2.5.5-linux-x64.tgz");
+      expect(commands[1]).toContain("mongosh-2.5.6-linux-x64.tgz");
       expect(commands[0]).toContain("mongosh not found");
     });
 

src/Explorer/Tabs/CloudShellTab/Utils/AttachAddOn.tsx

@@ -92,6 +92,18 @@ export class AttachAddon implements ITerminalAddon {
    * @param {Terminal} terminal - The XTerm terminal instance
    */
   public addMessageListener(terminal: Terminal): void {
+    let messageBuffer = "";
+    let bufferTimeout: NodeJS.Timeout | null = null;
+    const BUFFER_TIMEOUT = 50; // ms - short timeout for prompt detection
+
+    const processBuffer = () => {
+      if (messageBuffer.length > 0) {
+        this.handleCompleteTerminalData(terminal, messageBuffer);
+        messageBuffer = "";
+      }
+      bufferTimeout = null;
+    };
+
     this._disposables.push(
       addSocketListener(this._socket, "message", (ev) => {
         let data: ArrayBuffer | string = ev.data;
@@ -103,57 +115,136 @@ export class AttachAddon implements ITerminalAddon {
           data = enc.decode(ev.data as ArrayBuffer);
         }
 
-        // for example of json object look in TerminalHelper in the socket.onMessage
-        if (data.includes(startStatusJson) && data.includes(endStatusJson)) {
-          // process as one line
-          const statusData = data.split(startStatusJson)[1].split(endStatusJson)[0];
-          data = data.replace(statusData, "");
-          data = data.replace(startStatusJson, "");
-          data = data.replace(endStatusJson, "");
-        } else if (data.includes(startStatusJson)) {
-          // check for start
-          const partialStatusData = data.split(startStatusJson)[1];
-          this._socketData += partialStatusData;
-          data = data.replace(partialStatusData, "");
-          data = data.replace(startStatusJson, "");
-        } else if (data.includes(endStatusJson)) {
-          // check for end and process the command
-          const partialStatusData = data.split(endStatusJson)[0];
-          this._socketData += partialStatusData;
-          data = data.replace(partialStatusData, "");
-          data = data.replace(endStatusJson, "");
-          this._socketData = "";
-        } else if (this._socketData.length > 0) {
-          // check if the line is all data then just concatenate
-          this._socketData += data;
-          data = "";
-        }
+        // Handle status messages
+        let processedStatusData = data;
 
-        if (this._allowTerminalWrite && data.includes(this._startMarker)) {
-          this._allowTerminalWrite = false;
-          terminal.write(`Preparing ${this._shellHandler.getShellName()} environment...\r\n`);
-        }
-
-        if (this._allowTerminalWrite) {
-          const updatedData =
-            typeof this._shellHandler?.updateTerminalData === "function"
-              ? this._shellHandler.updateTerminalData(data)
-              : data;
-
-          const suppressedData = this._shellHandler?.getTerminalSuppressedData();
-
-          const shouldNotWrite = suppressedData.filter(Boolean).some((item) => updatedData.includes(item));
-
-          if (!shouldNotWrite) {
-            terminal.write(updatedData);
+        // Process status messages with delimiters
+        // eslint-disable-next-line no-constant-condition
+        while (true) {
+          const startIndex = processedStatusData.indexOf(startStatusJson);
+          if (startIndex === -1) {
+            break;
           }
+
+          const afterStart = processedStatusData.substring(startIndex + startStatusJson.length);
+          const endIndex = afterStart.indexOf(endStatusJson);
+
+          if (endIndex === -1) {
+            // Incomplete status message
+            this._socketData += processedStatusData.substring(startIndex);
+            processedStatusData = processedStatusData.substring(0, startIndex);
+            break;
+          }
+
+          // Remove processed status message
+          processedStatusData =
+            processedStatusData.substring(0, startIndex) + afterStart.substring(endIndex + endStatusJson.length);
         }
 
-        if (data.includes(this._shellHandler.getConnectionCommand())) {
-          this._allowTerminalWrite = true;
+        // Add to message buffer
+        messageBuffer += processedStatusData;
+
+        // Clear existing timeout
+        if (bufferTimeout) {
+          clearTimeout(bufferTimeout);
+          bufferTimeout = null;
+        }
+
+        // Check if this looks like a complete message/command
+        const isComplete = this.isMessageComplete(messageBuffer, processedStatusData);
+
+        if (isComplete) {
+          // Message marked as complete, processing immediately
+          processBuffer();
+        } else {
+          // Set timeout to process buffer after delay
+          bufferTimeout = setTimeout(processBuffer, BUFFER_TIMEOUT);
         }
       }),
     );
+
+    // Clean up timeout on dispose
+    this._disposables.push({
+      dispose: () => {
+        if (bufferTimeout) {
+          clearTimeout(bufferTimeout);
+        }
+      },
+    });
+  }
+
+  private isMessageComplete(fullBuffer: string, currentChunk: string): boolean {
+    // Immediate completion indicators
+    const immediateCompletionPatterns = [
+      /\n$/, // Ends with newline
+      /\r$/, // Ends with carriage return
+      /\r\n$/, // Ends with CRLF
+      /; \} \|\| true;$/, // Your command pattern
+      /disown -a && exit$/, // Exit commands
+      /printf.*?\\033\[0m\\n"$/, // Your printf pattern
+    ];
+
+    // Check current chunk for immediate completion
+    for (const pattern of immediateCompletionPatterns) {
+      if (pattern.test(currentChunk)) {
+        return true;
+      }
+    }
+
+    // ANSI sequence detection - these might be complete prompts
+    const ansiPromptPatterns = [
+      /\[\d+G\[0J.*>\s*\[\d+G$/, // Your specific pattern: [1G[0J...> [26G
+      /\[\d+;\d+H/, // Cursor position sequences
+      /\]\s*\[\d+G$/, // Ends with cursor positioning
+      />\s*\[\d+G$/, // Prompt followed by cursor position
+    ];
+
+    // Check if buffer ends with what looks like a complete prompt
+    for (const pattern of ansiPromptPatterns) {
+      if (pattern.test(fullBuffer)) {
+        return true;
+      }
+    }
+
+    // Check for MongoDB shell prompts specifically
+    const mongoPromptPatterns = [
+      /globaldb \[primary\] \w+>\s*\[\d+G$/, // MongoDB replica set prompt
+      />\s*\[\d+G$/, // General prompt with cursor positioning
+      /\w+>\s*$/, // Simple shell prompt
+    ];
+
+    for (const pattern of mongoPromptPatterns) {
+      if (pattern.test(fullBuffer)) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  private handleCompleteTerminalData(terminal: Terminal, data: string): void {
+    if (this._allowTerminalWrite && data.includes(this._startMarker)) {
+      this._allowTerminalWrite = false;
+      terminal.write(`Preparing ${this._shellHandler.getShellName()} environment...\r\n`);
+    }
+
+    if (this._allowTerminalWrite) {
+      const updatedData =
+        typeof this._shellHandler?.updateTerminalData === "function"
+          ? this._shellHandler.updateTerminalData(data)
+          : data;
+
+      const suppressedData = this._shellHandler?.getTerminalSuppressedData();
+      const shouldNotWrite = suppressedData.filter(Boolean).some((item) => updatedData.includes(item));
+
+      if (!shouldNotWrite) {
+        terminal.write(updatedData);
+      }
+    }
+
+    if (data.includes(this._shellHandler.getConnectionCommand())) {
+      this._allowTerminalWrite = true;
+    }
   }
 
   public dispose(): void {

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2Mongo.test.tsx

@@ -146,10 +146,16 @@ describe("Documents tab (Mongo API)", () => {
       updateConfigContext({ platform: Platform.Hosted });
 
       const props: IDocumentsTabComponentProps = createMockProps();
-
       wrapper = mount(<DocumentsTabComponent {...props} />);
-      wrapper = await waitForComponentToPaint(wrapper);
-    });
+
+      // Wait for all pending promises
+      await act(async () => {
+        await new Promise((resolve) => setTimeout(resolve, 100));
+      });
+
+      // Wait for any async operations to complete
+      wrapper = await waitForComponentToPaint(wrapper, 100);
+    }, 10000);
 
     afterEach(() => {
       wrapper.unmount();