upgrade react to v18 and other packages

* Update to Node 20.

* Update vulnerable dependencies.

.vscode/settings.json

@@ -1,22 +1,6 @@
 // Place your settings in this file to overwrite default and user settings.
 {
-  "files.exclude": {
-    ".vs": true,
-    ".vscode/**": true,
-    "*.trx": true,
-    "**/.DS_Store": true,
-    "**/.git": true,
-    "**/.hg": true,
-    "**/.svn": true,
-    "built/**": true,
-    "coverage/**": true,
-    "libs/**": true,
-    "node_modules/**": true,
-    "package-lock.json": true,
-    "quickstart/**": true,
-    "test/out/**": true,
-    "workers/libs/**": true
-  },
+  
   "typescript.tsdk": "node_modules/typescript/lib",
   "editor.formatOnSave": true,
   "editor.codeActionsOnSave": {
@@ -24,5 +8,8 @@
     "source.organizeImports": "explicit"
   },
   "typescript.preferences.importModuleSpecifier": "non-relative",
-  "editor.defaultFormatter": "esbenp.prettier-vscode"
+  "editor.defaultFormatter": "esbenp.prettier-vscode",
+  "[typescript]": {
+    "editor.defaultFormatter": "vscode.typescript-language-features"
+  }
 }

package.json

@@ -4,200 +4,270 @@
   "description": "Cosmos Explorer",
   "main": "index.js",
   "dependencies": {
-    "@azure/arm-cosmosdb": "9.1.0",
+    "@azure/arm-cosmosdb": "16.3.0",
     "@azure/cosmos": "4.5.0",
     "@azure/cosmos-language-service": "0.0.5",
-    "@azure/identity": "4.5.0",
-    "@azure/msal-browser": "2.14.2",
-    "@babel/plugin-proposal-class-properties": "7.12.1",
-    "@babel/plugin-proposal-decorators": "7.12.12",
-    "@fluentui/react": "8.119.0",
-    "@fluentui/react-components": "9.54.2",
-    "@jupyterlab/services": "6.0.2",
-    "@jupyterlab/terminal": "3.0.3",
-    "@microsoft/applicationinsights-web": "2.6.1",
+    "@azure/identity": "4.10.1",
+    "@azure/msal-browser": "4.24.0",
+    "@babel/plugin-transform-class-properties": "^7.24.7",
+    "@babel/plugin-proposal-decorators": "7.28.0",
+    "@fluentui/react": "8.123.6",
+    "@fluentui/react-components": "9.70.0",
+    "@jupyterlab/services": "7.4.9",
+    "@jupyterlab/terminal": "4.4.9",
+    "@microsoft/applicationinsights-web": "3.3.10",
     "@nteract/commutable": "7.5.1",
-    "@nteract/connected-components": "6.8.2",
+    "@nteract/connected-components": "6.9.0",
     "@nteract/core": "15.1.9",
-    "@nteract/data-explorer": "8.0.3",
-    "@nteract/directory-listing": "2.0.6",
-    "@nteract/dropdown-menu": "1.0.1",
+    "@nteract/data-explorer": "8.2.12",
+    "@nteract/directory-listing": "2.1.0",
+    "@nteract/dropdown-menu": "1.1.9",
     "@nteract/editor": "10.1.12",
-    "@nteract/fixtures": "2.3.0",
+    "@nteract/fixtures": "2.3.19",
     "@nteract/iron-icons": "1.0.0",
-    "@nteract/jupyter-widgets": "2.0.0",
+    "@nteract/jupyter-widgets": "4.1.19",
     "@nteract/logos": "1.0.0",
-    "@nteract/markdown": "4.6.0",
+    "@nteract/markdown": "4.6.2",
     "@nteract/monaco-editor": "3.2.2",
     "@nteract/octicons": "2.0.0",
-    "@nteract/outputs": "3.0.9",
-    "@nteract/presentational-components": "3.0.7",
-    "@nteract/stateful-components": "1.7.0",
-    "@nteract/styles": "2.0.2",
-    "@nteract/transform-geojson": "5.1.8",
+    "@nteract/outputs": "5.1.14",
+    "@nteract/presentational-components": "3.4.12",
+    "@nteract/stateful-components": "1.7.15",
+    "@nteract/styles": "2.2.11",
+    "@nteract/transform-geojson": "5.1.13",
     "@nteract/transform-model-debug": "5.0.1",
-    "@nteract/transform-plotly": "6.1.6",
-    "@nteract/transform-vdom": "4.0.11",
+    "@nteract/transform-plotly": "7.0.1",
+    "@nteract/transform-vdom": "4.0.15",
     "@nteract/transform-vega": "7.0.6",
-    "@octokit/rest": "17.9.2",
+    "@octokit/rest": "21.1.1",
     "@phosphor/widgets": "1.9.3",
-    "@testing-library/jest-dom": "6.4.6",
-    "@types/lodash": "4.14.171",
-    "@types/mkdirp": "1.0.1",
-    "@types/node-fetch": "2.5.7",
-    "@xmldom/xmldom": "0.7.13",
+    "@testing-library/jest-dom": "6.8.0",
+    "@types/lodash": "4.17.20",
+    "@types/mkdirp": "1.0.2",
+    "@types/node-fetch": "2.6.13",
+    "@xmldom/xmldom": "0.9.8",
     "@xterm/xterm": "5.5.0",
     "@xterm/addon-fit": "0.10.0",
-    "allotment": "1.20.2",
-    "applicationinsights": "1.8.0",
+    "allotment": "1.20.4",
+    "applicationinsights": "3.12.0",
     "bootstrap": "3.4.1",
-    "canvas": "2.11.2",
+    "canvas": "3.2.0",
     "clean-webpack-plugin": "4.0.0",
     "clipboard-copy": "4.0.1",
-    "copy-webpack-plugin": "11.0.0",
+    "copy-webpack-plugin": "13.0.1",
     "crossroads": "0.12.2",
-    "css-element-queries": "1.1.1",
-    "d3": "7.8.5",
-    "datatables.net-colreorder-dt": "1.7.0",
-    "datatables.net-dt": "1.13.8",
-    "date-fns": "1.29.0",
-    "dayjs": "1.8.19",
+    "css-element-queries": "1.2.3",
+    "d3": "7.9.0",
+    "datatables.net-colreorder-dt": "2.1.1",
+    "datatables.net-dt": "2.3.4",
+    "date-fns": "4.1.0",
+    "dayjs": "1.11.18",
     "dom-to-image": "2.6.0",
-    "dotenv": "8.2.0",
-    "eslint-plugin-jest": "27.4.2",
-    "eslint-plugin-react": "7.33.2",
+    "dotenv": "17.2.3",
+    "eslint-plugin-jest": "28.14.0",
+    "eslint-plugin-react": "7.37.5",
     "hasher": "1.2.0",
-    "html2canvas": "1.0.0-rc.5",
-    "i18next": "23.11.5",
-    "i18next-browser-languagedetector": "6.0.1",
-    "i18next-http-backend": "1.0.23",
-    "iframe-resizer-react": "1.1.0",
-    "immer": "9.0.6",
+    "html2canvas": "1.4.1",
+    "i18next": "25.5.2",
+    "i18next-browser-languagedetector": "8.2.0",
+    "i18next-http-backend": "3.0.2",
+    "iframe-resizer-react": "5.1.5",
+    "immer": "10.1.3",
     "immutable": "4.0.0-rc.12",
-    "is-ci": "2.0.0",
+    "is-ci": "4.1.0",
     "jquery": "3.7.1",
     "jquery-typeahead": "2.11.1",
-    "jquery-ui-dist": "1.13.2",
+    "jquery-ui-dist": "1.13.3",
     "knockout": "3.5.1",
-    "loader-utils": "2.0.3",
-    "mkdirp": "1.0.4",
-    "monaco-editor": "0.44.0",
+    "loader-utils": "3.3.1",
+    "mkdirp": "3.0.1",
+    "monaco-editor": "0.53.0",
     "ms": "2.1.3",
     "p-retry": "6.2.1",
-    "patch-package": "8.0.0",
-    "plotly.js-cartesian-dist-min": "1.52.3",
+    "patch-package": "8.0.1",
+    "plotly.js-cartesian-dist-min": "3.1.1",
     "post-robot": "10.0.42",
-    "q": "1.5.1",
-    "react": "16.14.0",
-    "react-animate-height": "2.0.8",
-    "react-dnd": "14.0.2",
-    "react-dnd-html5-backend": "14.0.0",
-    "react-dom": "16.14.0",
+    "q": "2.0.3",
+    "react": "18.2.0",
+    "react-animate-height": "3.2.3",
+    "react-dnd": "16.0.1",
+    "react-dnd-html5-backend": "16.0.1",
+    "react-dom": "18.2.0",
     "react-hotkeys": "2.0.0",
-    "react-i18next": "14.1.2",
+    "react-i18next": "16.0.0",
     "react-notification-system": "0.2.17",
-    "react-redux": "7.1.3",
+    "react-redux": "7.2.9",
     "react-splitter-layout": "4.0.0",
-    "react-string-format": "1.0.1",
+    "react-string-format": "1.2.0",
     "react-window": "1.8.10",
-    "react-youtube": "9.0.1",
-    "reflect-metadata": "0.1.13",
-    "rx-jupyter": "5.5.12",
-    "sanitize-html": "2.3.3",
-    "shell-quote": "1.7.3",
-    "styled-components": "5.0.1",
-    "swr": "0.4.0",
-    "terser-webpack-plugin": "5.3.9",
-    "tinykeys": "2.1.0",
-    "underscore": "1.12.1",
-    "utility-types": "3.10.0",
-    "zustand": "3.5.0"
+    "react-youtube": "10.1.0",
+    "reflect-metadata": "0.2.2",
+    "rx-jupyter": "5.5.21",
+    "sanitize-html": "2.17.0",
+    "shell-quote": "1.8.3",
+    "styled-components": "6.1.19",
+    "swr": "2.3.6",
+    "terser-webpack-plugin": "5.3.14",
+    "tinykeys": "3.0.0",
+    "underscore": "1.13.7",
+    "utility-types": "3.11.0",
+    "zustand": "5.0.8"
   },
   "devDependencies": {
-    "@babel/core": "7.24.7",
-    "@babel/preset-env": "7.24.7",
-    "@babel/preset-react": "7.24.7",
-    "@babel/preset-typescript": "7.24.7",
-    "@playwright/test": "1.49.1",
-    "@testing-library/react": "11.2.3",
-    "@types/applicationinsights-js": "1.0.7",
-    "@types/codemirror": "0.0.56",
-    "@types/crossroads": "0.0.30",
-    "@types/d3": "5.9.2",
+    "@babel/core": "7.28.4",
+    "@babel/preset-env": "7.28.3",
+    "@babel/preset-react": "7.27.1",
+    "@babel/preset-typescript": "7.27.1",
+    "@playwright/test": "1.55.1",
+    "@testing-library/react": "16.3.0",
+    "@types/applicationinsights-js": "1.0.9",
+    "@types/codemirror": "5.60.16",
+    "@types/crossroads": "0.0.33",
+    "@types/d3": "7.4.3",
     "@types/datatables.net": "1.10.28",
     "@types/datatables.net-colreorder": "1.4.5",
-    "@types/dom-to-image": "2.6.2",
-    "@types/enzyme": "3.10.12",
+    "@types/dom-to-image": "2.6.7",
+    "@types/enzyme": "3.10.19",
     "@types/enzyme-adapter-react-16": "1.0.9",
-    "@types/hasher": "0.0.31",
-    "@types/jest": "29.5.12",
-    "@types/jquery": "3.5.29",
-    "@types/node": "12.11.1",
-    "@types/post-robot": "10.0.1",
-    "@types/q": "1.5.1",
-    "@types/react": "17.0.44",
-    "@types/react-dom": "17.0.15",
-    "@types/react-notification-system": "0.2.39",
-    "@types/react-redux": "7.1.7",
-    "@types/react-splitter-layout": "3.0.1",
+    "@types/hasher": "0.0.35",
+    "@types/jest": "30.0.0",
+    "@types/jquery": "3.5.33",
+    "@types/node": "24.6.0",
+    "@types/post-robot": "10.0.6",
+    "@types/q": "1.5.8",
+    "@types/react": "18.3.7",
+    "@types/react-dom": "18.3.7",
+    "@types/react-notification-system": "0.2.46",
+    "@types/react-redux": "7.1.34",
+    "@types/react-splitter-layout": "4.0.0",
     "@types/react-window": "1.8.8",
-    "@types/sanitize-html": "1.27.2",
-    "@types/sinon": "2.3.3",
-    "@types/styled-components": "5.1.1",
-    "@types/underscore": "1.7.36",
-    "@types/youtube-player": "5.5.6",
-    "@typescript-eslint/eslint-plugin": "6.7.4",
-    "@typescript-eslint/parser": "6.7.4",
-    "@webpack-cli/serve": "2.0.5",
-    "babel-jest": "29.7.0",
-    "babel-loader": "8.1.0",
-    "buffer": "5.1.0",
+    "@types/sanitize-html": "2.16.0",
+    "@types/sinon": "17.0.4",
+    "@types/styled-components": "5.1.34",
+    "@types/underscore": "1.13.0",
+    "@types/youtube-player": "5.5.11",
+    "@typescript-eslint/eslint-plugin": "8.45.0",
+    "@typescript-eslint/parser": "8.45.0",
+    "@webpack-cli/serve": "3.0.1",
+    "babel-jest": "30.2.0",
+    "babel-loader": "10.0.0",
+    "buffer": "6.0.3",
     "case-sensitive-paths-webpack-plugin": "2.4.0",
     "create-file-webpack": "1.0.2",
-    "css-loader": "6.8.1",
+    "css-loader": "7.1.2",
     "enzyme": "3.11.0",
     "enzyme-adapter-react-16": "1.15.8",
     "enzyme-to-json": "3.6.2",
-    "eslint": "8.50.0",
+    "eslint": "9.36.0",
     "eslint-cli": "1.1.1",
     "eslint-plugin-no-null": "1.0.2",
     "eslint-plugin-prefer-arrow": "1.2.3",
-    "eslint-plugin-react-hooks": "4.6.0",
-    "fast-glob": "3.2.5",
-    "fs-extra": "7.0.0",
+    "eslint-plugin-react-hooks": "5.2.0",
+    "fast-glob": "3.3.3",
+    "fs-extra": "11.3.2",
     "html-inline-css-webpack-plugin": "1.11.2",
-    "html-loader": "5.0.0",
-    "html-webpack-plugin": "5.5.3",
-    "jest": "29.7.0",
+    "html-loader": "5.1.0",
+    "html-webpack-plugin": "5.6.4",
+    "jest": "30.2.0",
     "jest-canvas-mock": "2.5.2",
-    "jest-circus": "29.7.0",
-    "jest-environment-jsdom": "29.7.0",
+    "jest-circus": "30.2.0",
+    "jest-environment-jsdom": "30.2.0",
     "jest-html-loader": "1.0.0",
     "jest-react-hooks-shallow": "1.5.1",
     "jest-trx-results-processor": "3.0.2",
-    "less": "3.8.1",
-    "less-loader": "11.1.3",
+    "less": "4.4.1",
+    "less-loader": "12.3.0",
     "less-vars-loader": "1.1.0",
-    "mini-css-extract-plugin": "2.1.0",
+    "mini-css-extract-plugin": "2.9.4",
     "monaco-editor-webpack-plugin": "7.1.0",
-    "node-fetch": "2.6.7",
-    "prettier": "3.0.3",
+    "node-fetch": "3.3.2",
+    "prettier": "3.6.2",
     "process": "0.11.10",
     "querystring-es3": "0.2.1",
-    "raw-loader": "0.5.1",
+    "raw-loader": "4.0.2",
     "react-dev-utils": "12.0.1",
-    "rimraf": "3.0.0",
-    "sinon": "3.2.1",
-    "style-loader": "0.23.0",
-    "ts-loader": "9.2.4",
-    "typedoc": "0.26.2",
-    "typescript": "4.9.5",
+    "rimraf": "5.0.10",
+    "sinon": "21.0.0",
+    "style-loader": "4.0.0",
+    "ts-loader": "9.5.4",
+    "typedoc": "0.28.13",
+    "typescript": "5.9.2",
     "url-loader": "4.1.1",
-    "wait-on": "4.0.2",
-    "webpack": "5.88.2",
-    "webpack-bundle-analyzer": "4.9.1",
-    "webpack-cli": "5.1.4",
-    "webpack-dev-server": "4.15.2"
+    "wait-on": "8.0.5",
+    "webpack": "5.102.0",
+    "webpack-bundle-analyzer": "4.10.2",
+    "webpack-cli": "6.0.1",
+    "webpack-dev-server": "5.2.2"
+  },
+  "overrides": {
+    "@nteract/connected-components": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/core": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/data-explorer": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/dropdown-menu": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/editor": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/iron-icons": {
+      "react": "18.2.0"
+    },
+    "@nteract/jupyter-widgets": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/logos": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/markdown": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/monaco-editor": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/octicons": {
+      "react": "18.2.0"
+    },
+    "@nteract/outputs": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/presentational-components": {
+      "react": "18.2.0"
+    },
+    "@nteract/stateful-components": {
+      "react": "18.2.0",
+      "react-dom": "18.2.0"
+    },
+    "@nteract/transform-geojson": {
+      "react": "18.2.0"
+    },
+    "@nteract/transform-model-debug": {
+      "react": "18.2.0"
+    },
+    "@nteract/transform-plotly": {
+      "react": "18.2.0"
+    },
+    "@nteract/transform-vdom": {
+      "react": "18.2.0"
+    },
+    "@nteract/transform-vega": {
+      "react": "18.2.0"
+    }
   },
   "scripts": {
     "postinstall": "patch-package",

patches/datatables.net-colreorder+2.1.1.patch

@@ -6,9 +6,9 @@ index e5dc283..1930c2b 100644
  
  /// <reference types="jquery" />
  
--import DataTables, {Api} from 'datatables.net';
-+import DataTables, { Api } from 'datatables.net';
- 
+-import DataTables, {Api, ColumnSelector} from 'datatables.net';
++import DataTables, { Api, ColumnSelector } from 'datatables.net';
+
  export default DataTables;
  
 @@ -40,6 +40,8 @@ declare module 'datatables.net' {
@@ -17,6 +17,6 @@ index e5dc283..1930c2b 100644
  			 */
 +      // Ignore this error: error TS7013: Construct signature, which lacks return-type annotation, implicitly has an 'any' return type.
 +      // @ts-ignore
- 			new (dt: Api<any>, settings: boolean | ConfigColReorder);
+ 			new (dt: Api<any>, settings: boolean | ConfigColReorder): DataTablesStatic['ColReorder'];
  
  			/**

preview/.azure/config

@@ -1,6 +1,6 @@
 [defaults]
 group = dataexplorer-preview
-sku = P1V2
+sku = P1v2
 appserviceplan = dataexplorer-preview
 location = westus2
 web = dataexplorer-preview

preview/package.json

@@ -4,16 +4,18 @@
   "description": "",
   "main": "index.js",
   "scripts": {
-    "deploy": "az webapp up --name \"dataexplorer-preview\" --subscription \"cosmosdb-portalteam-runners\" --resource-group \"dataexplorer-preview\" --runtime \"NODE:18-lts\" --sku P1V2",
+    "deploy": "az webapp up --name \"dataexplorer-preview\" --subscription \"cosmosdb-portalteam-runners\" --resource-group \"dataexplorer-preview\" --runtime \"NODE:20-lts\" --sku P1V2",
     "start": "node index.js",
     "test": "echo \"Error: no test specified\" && exit 1"
   },
   "keywords": [],
   "author": "Microsoft Corporation",
   "dependencies": {
-    "express": "^4.17.1",
+    "body-parser": "^1.20.3",
+    "express": "^4.21.2",
     "http-proxy-middleware": "^3.0.3",
-    "node": "^18.20.6",
-    "node-fetch": "^2.6.1"
+    "node": "^20.19.5",
+    "node-fetch": "^2.6.1",
+    "path-to-regexp": "^0.1.12"
   }
 }

src/Common/Constants.ts

@@ -138,6 +138,14 @@ export enum MongoBackendEndpointType {
   remote,
 }
 
+export class AadScopeEndpoints {
+  public static readonly Development: string = "https://cosmos.azure.com";
+  public static readonly MPAC: string = "https://cosmos.azure.com";
+  public static readonly Prod: string = "https://cosmos.azure.com";
+  public static readonly Fairfax: string = "https://cosmos.azure.us";
+  public static readonly Mooncake: string = "https://cosmos.azure.cn";
+}
+
 export class PortalBackendEndpoints {
   public static readonly Development: string = "https://localhost:7235";
   public static readonly Mpac: string = "https://cdb-ms-mpac-pbe.cosmos.azure.com";
@@ -255,6 +263,7 @@ export class HttpHeaders {
   public static activityId: string = "x-ms-activity-id";
   public static apiType: string = "x-ms-cosmos-apitype";
   public static authorization: string = "authorization";
+  public static entraIdToken: string = "x-ms-entraid-token";
   public static collectionIndexTransformationProgress: string =
     "x-ms-documentdb-collection-index-transformation-progress";
   public static continuation: string = "x-ms-continuation";

src/Common/EnvironmentUtility.test.ts

@@ -28,3 +28,39 @@ describe("Environment Utility Test", () => {
     expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Development);
   });
 });
+describe("normalizeArmEndpoint", () => {
+  it("should append '/' if not present", () => {
+    expect(EnvironmentUtility.normalizeArmEndpoint("https://example.com")).toBe("https://example.com/");
+  });
+
+  it("should return the same uri if '/' is present at the end", () => {
+    expect(EnvironmentUtility.normalizeArmEndpoint("https://example.com/")).toBe("https://example.com/");
+  });
+
+  it("should handle empty string", () => {
+    expect(EnvironmentUtility.normalizeArmEndpoint("")).toBe("");
+  });
+});
+
+describe("getEnvironment", () => {
+  it("should return Prod environment", () => {
+    updateConfigContext({
+      PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Prod,
+    });
+    expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Prod);
+  });
+
+  it("should return Fairfax environment", () => {
+    updateConfigContext({
+      PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Fairfax,
+    });
+    expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Fairfax);
+  });
+
+  it("should return Mooncake environment", () => {
+    updateConfigContext({
+      PORTAL_BACKEND_ENDPOINT: PortalBackendEndpoints.Mooncake,
+    });
+    expect(EnvironmentUtility.getEnvironment()).toBe(EnvironmentUtility.Environment.Mooncake);
+  });
+});

src/Common/EnvironmentUtility.ts

@@ -1,4 +1,5 @@
-import { PortalBackendEndpoints } from "Common/Constants";
+import { AadScopeEndpoints, PortalBackendEndpoints } from "Common/Constants";
+import * as Logger from "Common/Logger";
 import { configContext } from "ConfigContext";
 
 export function normalizeArmEndpoint(uri: string): string {
@@ -27,3 +28,17 @@ export const getEnvironment = (): Environment => {
 
   return environmentMap[configContext.PORTAL_BACKEND_ENDPOINT];
 };
+
+export const getEnvironmentScopeEndpoint = (): string => {
+  const environment = getEnvironment();
+  const endpoint = AadScopeEndpoints[environment];
+  if (!endpoint) {
+    throw new Error("Cannot determine AAD scope endpoint");
+  }
+  const hrefEndpoint = new URL(endpoint).href.replace(/\/+$/, "/.default");
+  Logger.logInfo(
+    `Using AAD scope endpoint: ${hrefEndpoint}, Environment: ${environment}`,
+    "EnvironmentUtility/getEnvironmentScopeEndpoint",
+  );
+  return hrefEndpoint;
+};

src/Common/MongoProxyClient.ts

@@ -7,6 +7,7 @@ import { MessageTypes } from "../Contracts/ExplorerContracts";
 import { Collection } from "../Contracts/ViewModels";
 import DocumentId from "../Explorer/Tree/DocumentId";
 import { userContext } from "../UserContext";
+import { isDataplaneRbacEnabledForProxyApi } from "../Utils/AuthorizationUtils";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import { ApiType, ContentType, HttpHeaders, HttpStatusCodes } from "./Constants";
 import { MinimalQueryIterator } from "./IteratorUtilities";
@@ -22,7 +23,13 @@ function authHeaders() {
   if (userContext.authType === AuthType.EncryptedToken) {
     return { [HttpHeaders.guestAccessToken]: userContext.accessToken };
   } else {
-    return { [HttpHeaders.authorization]: userContext.authorizationToken };
+    const headers: { [key: string]: string } = {
+      [HttpHeaders.authorization]: userContext.authorizationToken,
+    };
+    if (isDataplaneRbacEnabledForProxyApi(userContext)) {
+      headers[HttpHeaders.entraIdToken] = userContext.aadToken;
+    }
+    return headers;
   }
 }
 

src/Explorer/Controls/Dialog.tsx

@@ -16,7 +16,7 @@ import {
   TextField,
 } from "@fluentui/react";
 import React, { FC, useEffect } from "react";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 export interface DialogState {
   visible: boolean;
@@ -38,7 +38,7 @@ export interface DialogState {
   showOkModalDialog: (title: string, subText: string, linkProps?: LinkProps) => void;
 }
 
-export const useDialog: UseStore<DialogState> = create((set, get) => ({
+export const useDialog = create<DialogState>((set, get) => ({
   visible: false,
   openDialog: (props: DialogProps) => set(() => ({ visible: true, dialogProps: props })),
   closeDialog: () =>

src/Explorer/Explorer.tsx

@@ -23,7 +23,7 @@ import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
-import shallow from "zustand/shallow";
+import { shallow } from "zustand/shallow";
 import { AuthType } from "../AuthType";
 import { BindingHandlersRegisterer } from "../Bindings/BindingHandlersRegisterer";
 import * as Constants from "../Common/Constants";
@@ -112,8 +112,8 @@ export default class Explorer {
 
     this.phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
     useNotebook.subscribe(
-      () => this.refreshCommandBarButtons(),
       (state) => state.isNotebooksEnabledForAccount,
+      () => this.refreshCommandBarButtons(),
     );
 
     this.queriesClient = new QueriesClient(this);
@@ -136,13 +136,13 @@ export default class Explorer {
     });
 
     useTabs.subscribe(
+      (state) => state.openedTabs,
       (openedTabs: TabsBase[]) => {
         if (openedTabs.length === 0) {
           useSelectedNode.getState().setSelectedNode(undefined);
           useCommandBar.getState().setContextButtons([]);
         }
-      },
-      (state) => state.openedTabs,
+      }
     );
 
     this.isTabsContentExpanded = ko.observable(false);
@@ -170,9 +170,9 @@ export default class Explorer {
     );
 
     useNotebook.subscribe(
-      async () => this.initiateAndRefreshNotebookList(),
       (state) => [state.isNotebookEnabled, state.isRefreshed],
-      shallow,
+      async () => this.initiateAndRefreshNotebookList(),
+      { equalityFn: shallow },
     );
 
     this.resourceTree = new ResourceTreeAdapter(this);

src/Explorer/Menus/CommandBar/CommandBarComponentAdapter.tsx

@@ -5,11 +5,12 @@
  */
 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";
 import * as React from "react";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 import { ConnectionStatusType, PoolIdType } from "../../../Common/Constants";
 import { StyleConstants } from "../../../Common/StyleConstants";
 import { CommandButtonComponentProps } from "../../Controls/CommandButton/CommandButtonComponent";
@@ -29,8 +30,8 @@ export interface CommandBarStore {
   setIsHidden: (isHidden: boolean) => void;
 }
 
-export const useCommandBar: UseStore<CommandBarStore> = create((set) => ({
-  contextButtons: [],
+export const useCommandBar = create<CommandBarStore>((set) => ({
+  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/Notebook/NotebookRenderer/StatusBar.tsx

@@ -1,5 +1,5 @@
 import { AppState, ContentRef, selectors } from "@nteract/core";
-import distanceInWordsToNow from "date-fns/distance_in_words_to_now";
+import { formatDistanceToNow } from "date-fns";
 import React from "react";
 import { connect } from "react-redux";
 import styled from "styled-components";
@@ -59,7 +59,7 @@ export class StatusBar extends React.Component<Props> {
         <Bar data-test="notebookStatusBar">
           <RightStatus>
             {this.props.lastSaved ? (
-              <p data-test="saveStatus"> Last saved {distanceInWordsToNow(this.props.lastSaved)} </p>
+              <p data-test="saveStatus"> Last saved {formatDistanceToNow(this.props.lastSaved)} ago </p>
             ) : (
               <p> Not saved yet </p>
             )}

src/Explorer/Notebook/useNotebook.ts

@@ -1,7 +1,8 @@
 import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import { PhoenixClient } from "Phoenix/PhoenixClient";
 import { cloneDeep } from "lodash";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from 'zustand/middleware';
 import { AuthType } from "../../AuthType";
 import * as Constants from "../../Common/Constants";
 import { ConnectionStatusType, HttpStatusCodes } from "../../Common/Constants";
@@ -66,270 +67,274 @@ interface NotebookState {
   setIsPhoenixFeatures: (isPhoenixFeatures: boolean) => void;
 }
 
-export const useNotebook: UseStore<NotebookState> = create((set, get) => ({
-  isNotebookEnabled: false,
-  isNotebooksEnabledForAccount: false,
-  notebookServerInfo: {
-    notebookServerEndpoint: undefined,
-    authToken: undefined,
-    forwardingId: undefined,
-  },
-  sparkClusterConnectionInfo: {
-    userName: undefined,
-    password: undefined,
-    endpoints: [],
-  },
-  isSynapseLinkUpdating: false,
-  memoryUsageInfo: undefined,
-  isShellEnabled: false,
-  notebookBasePath: Constants.Notebook.defaultBasePath,
-  isInitializingNotebooks: false,
-  myNotebooksContentRoot: undefined,
-  gitHubNotebooksContentRoot: undefined,
-  galleryContentRoot: undefined,
-  connectionInfo: {
-    status: ConnectionStatusType.Connect,
-  },
-  notebookFolderName: undefined,
-  isAllocating: false,
-  isRefreshed: false,
-  containerStatus: {
-    status: undefined,
-    durationLeftInMinutes: undefined,
-    phoenixServerInfo: undefined,
-  },
-  isPhoenixNotebooks: undefined,
-  isPhoenixFeatures: undefined,
-  setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
-  setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => set({ isNotebooksEnabledForAccount }),
-  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
-    set({ notebookServerInfo }),
-  setSparkClusterConnectionInfo: (sparkClusterConnectionInfo: DataModels.SparkClusterConnectionInfo) =>
-    set({ sparkClusterConnectionInfo }),
-  setIsSynapseLinkUpdating: (isSynapseLinkUpdating: boolean) => set({ isSynapseLinkUpdating }),
-  setMemoryUsageInfo: (memoryUsageInfo: DataModels.MemoryUsageInfo) => set({ memoryUsageInfo }),
-  setIsShellEnabled: (isShellEnabled: boolean) => set({ isShellEnabled }),
-  setNotebookBasePath: (notebookBasePath: string) => set({ notebookBasePath }),
-  setNotebookFolderName: (notebookFolderName: string) => set({ notebookFolderName }),
-  refreshNotebooksEnabledStateForAccount: async (): Promise<void> => {
-    await get().getPhoenixStatus();
-    const { databaseAccount, authType } = userContext;
-    if (
-      authType === AuthType.EncryptedToken ||
-      authType === AuthType.ResourceToken ||
-      authType === AuthType.MasterKey
-    ) {
-      set({ isNotebooksEnabledForAccount: false });
-      return;
-    }
-
-    const firstWriteLocation =
-      userContext.apiType === "Postgres" || userContext.apiType === "VCoreMongo"
-        ? databaseAccount?.location
-        : databaseAccount?.properties?.writeLocations?.[0]?.locationName.toLowerCase();
-    const disallowedLocationsUri: string = `${configContext.PORTAL_BACKEND_ENDPOINT}/api/disallowedlocations`;
-    const authorizationHeader = getAuthorizationHeader();
-    try {
-      const response = await fetch(disallowedLocationsUri, {
-        method: "POST",
-        body: JSON.stringify({
-          resourceTypes: [Constants.ArmResourceTypes.notebookWorkspaces],
-        }),
-        headers: {
-          [authorizationHeader.header]: authorizationHeader.token,
-          [Constants.HttpHeaders.contentType]: "application/json",
-        },
-      });
-
-      if (!response.ok) {
-        throw new Error("Failed to fetch disallowed locations");
-      }
-
-      const disallowedLocations: string[] = await response.json();
-      if (!disallowedLocations) {
-        Logger.logInfo("No disallowed locations found", "Explorer/isNotebooksEnabledForAccount");
-        set({ isNotebooksEnabledForAccount: true });
-        return;
-      }
-
-      // firstWriteLocation should not be disallowed
-      const isAccountInAllowedLocation = firstWriteLocation && disallowedLocations.indexOf(firstWriteLocation) === -1;
-      set({ isNotebooksEnabledForAccount: isAccountInAllowedLocation });
-    } catch (error) {
-      Logger.logError(getErrorMessage(error), "Explorer/isNotebooksEnabledForAccount");
-      set({ isNotebooksEnabledForAccount: false });
-    }
-  },
-  findItem: (root: NotebookContentItem, item: NotebookContentItem): NotebookContentItem => {
-    const currentItem = root || get().myNotebooksContentRoot;
-
-    if (currentItem) {
-      if (currentItem.path === item.path && currentItem.name === item.name) {
-        return currentItem;
-      }
-
-      if (currentItem.children) {
-        for (const childItem of currentItem.children) {
-          const result = get().findItem(childItem, item);
-          if (result) {
-            return result;
-          }
+export const useNotebook = create<NotebookState>()(
+  subscribeWithSelector(
+    (set, get) => ({
+      isNotebookEnabled: false,
+      isNotebooksEnabledForAccount: false,
+      notebookServerInfo: {
+        notebookServerEndpoint: "",
+        authToken: "",
+        forwardingId: "",
+      },
+      sparkClusterConnectionInfo: {
+        userName: "",
+        password: "",
+        endpoints: [] as DataModels.SparkClusterEndpoint[],
+      },
+      isSynapseLinkUpdating: false,
+      memoryUsageInfo: undefined as DataModels.MemoryUsageInfo,
+      isShellEnabled: false,
+      notebookBasePath: Constants.Notebook.defaultBasePath,
+      isInitializingNotebooks: false,
+      myNotebooksContentRoot: undefined as NotebookContentItem,
+      gitHubNotebooksContentRoot: undefined as NotebookContentItem,
+      galleryContentRoot: undefined as NotebookContentItem,
+      connectionInfo: {
+        status: ConnectionStatusType.Connect,
+      },
+      notebookFolderName: "",
+      isAllocating: false,
+      isRefreshed: false,
+      containerStatus: {
+        status: undefined,
+        durationLeftInMinutes: undefined,
+        phoenixServerInfo: undefined,
+      } as ContainerInfo,
+      isPhoenixNotebooks: undefined as boolean,
+      isPhoenixFeatures: undefined as boolean,
+      setIsNotebookEnabled: (isNotebookEnabled: boolean) => set({ isNotebookEnabled }),
+      setIsNotebooksEnabledForAccount: (isNotebooksEnabledForAccount: boolean) => set({ isNotebooksEnabledForAccount }),
+      setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
+        set({ notebookServerInfo }),
+      setSparkClusterConnectionInfo: (sparkClusterConnectionInfo: DataModels.SparkClusterConnectionInfo) =>
+        set({ sparkClusterConnectionInfo }),
+      setIsSynapseLinkUpdating: (isSynapseLinkUpdating: boolean) => set({ isSynapseLinkUpdating }),
+      setMemoryUsageInfo: (memoryUsageInfo: DataModels.MemoryUsageInfo) => set({ memoryUsageInfo }),
+      setIsShellEnabled: (isShellEnabled: boolean) => set({ isShellEnabled }),
+      setNotebookBasePath: (notebookBasePath: string) => set({ notebookBasePath }),
+      setNotebookFolderName: (notebookFolderName: string) => set({ notebookFolderName }),
+      refreshNotebooksEnabledStateForAccount: async (): Promise<void> => {
+        await get().getPhoenixStatus();
+        const { databaseAccount, authType } = userContext;
+        if (
+          authType === AuthType.EncryptedToken ||
+          authType === AuthType.ResourceToken ||
+          authType === AuthType.MasterKey
+        ) {
+          set({ isNotebooksEnabledForAccount: false });
+          return;
         }
-      }
-    }
 
-    return undefined;
-  },
-  insertNotebookItem: (parent: NotebookContentItem, item: NotebookContentItem, isGithubTree?: boolean): void => {
-    const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
-    const parentItem = get().findItem(root, parent);
-    item.parent = parentItem;
-    if (parentItem.children) {
-      parentItem.children.push(item);
-    } else {
-      parentItem.children = [item];
-    }
-    isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
-  },
-  updateNotebookItem: (item: NotebookContentItem, isGithubTree?: boolean): void => {
-    const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
-    const parentItem = get().findItem(root, item.parent);
-    parentItem.children = parentItem.children.filter((child) => child.path !== item.path);
-    parentItem.children.push(item);
-    item.parent = parentItem;
-    isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
-  },
-  deleteNotebookItem: (item: NotebookContentItem, isGithubTree?: boolean): void => {
-    const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
-    const parentItem = get().findItem(root, item.parent);
-    parentItem.children = parentItem.children.filter((child) => child.path !== item.path);
-    isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
-  },
-  initializeNotebooksTree: async (notebookManager: NotebookManager): Promise<void> => {
-    const notebookFolderName = get().isPhoenixNotebooks ? "Temporary Notebooks" : "My Notebooks";
-    set({ notebookFolderName });
-    const myNotebooksContentRoot = {
-      name: get().notebookFolderName,
-      path: get().notebookBasePath,
-      type: NotebookContentItemType.Directory,
-    };
-    const galleryContentRoot = {
-      name: "Gallery",
-      path: "Gallery",
-      type: NotebookContentItemType.File,
-    };
-    const gitHubNotebooksContentRoot = notebookManager?.gitHubOAuthService?.isLoggedIn()
-      ? {
-          name: "GitHub repos",
-          path: "PsuedoDir",
-          type: NotebookContentItemType.Directory,
-        }
-      : undefined;
-
-    set({
-      myNotebooksContentRoot,
-      galleryContentRoot,
-      gitHubNotebooksContentRoot,
-    });
-
-    if (get().notebookServerInfo?.notebookServerEndpoint) {
-      const updatedRoot = await notebookManager?.notebookContentClient?.updateItemChildren(myNotebooksContentRoot);
-      set({ myNotebooksContentRoot: updatedRoot });
-
-      if (updatedRoot?.children) {
-        // Count 1st generation children (tree is lazy-loaded)
-        const nodeCounts = { files: 0, notebooks: 0, directories: 0 };
-        updatedRoot.children.forEach((notebookItem) => {
-          switch (notebookItem.type) {
-            case NotebookContentItemType.File:
-              nodeCounts.files++;
-              break;
-            case NotebookContentItemType.Directory:
-              nodeCounts.directories++;
-              break;
-            case NotebookContentItemType.Notebook:
-              nodeCounts.notebooks++;
-              break;
-            default:
-              break;
-          }
-        });
-        TelemetryProcessor.trace(Action.RefreshResourceTreeMyNotebooks, ActionModifiers.Mark, { ...nodeCounts });
-      }
-    }
-  },
-  initializeGitHubRepos: (pinnedRepos: IPinnedRepo[]): void => {
-    const gitHubNotebooksContentRoot = cloneDeep(get().gitHubNotebooksContentRoot);
-    if (gitHubNotebooksContentRoot) {
-      gitHubNotebooksContentRoot.children = [];
-      pinnedRepos?.forEach((pinnedRepo) => {
-        const repoFullName = GitHubUtils.toRepoFullName(pinnedRepo.owner, pinnedRepo.name);
-        const repoTreeItem: NotebookContentItem = {
-          name: repoFullName,
-          path: "PsuedoDir",
-          type: NotebookContentItemType.Directory,
-          children: [],
-          parent: gitHubNotebooksContentRoot,
-        };
-
-        pinnedRepo.branches.forEach((branch) => {
-          repoTreeItem.children.push({
-            name: branch.name,
-            path: GitHubUtils.toContentUri(pinnedRepo.owner, pinnedRepo.name, branch.name, ""),
-            type: NotebookContentItemType.Directory,
-            parent: repoTreeItem,
+        const firstWriteLocation =
+          userContext.apiType === "Postgres" || userContext.apiType === "VCoreMongo"
+            ? databaseAccount?.location
+            : databaseAccount?.properties?.writeLocations?.[0]?.locationName.toLowerCase();
+        const disallowedLocationsUri: string = `${configContext.PORTAL_BACKEND_ENDPOINT}/api/disallowedlocations`;
+        const authorizationHeader = getAuthorizationHeader();
+        try {
+          const response = await fetch(disallowedLocationsUri, {
+            method: "POST",
+            body: JSON.stringify({
+              resourceTypes: [Constants.ArmResourceTypes.notebookWorkspaces],
+            }),
+            headers: {
+              [authorizationHeader.header]: authorizationHeader.token,
+              [Constants.HttpHeaders.contentType]: "application/json",
+            },
           });
+
+          if (!response.ok) {
+            throw new Error("Failed to fetch disallowed locations");
+          }
+
+          const disallowedLocations: string[] = await response.json();
+          if (!disallowedLocations) {
+            Logger.logInfo("No disallowed locations found", "Explorer/isNotebooksEnabledForAccount");
+            set({ isNotebooksEnabledForAccount: true });
+            return;
+          }
+
+          // firstWriteLocation should not be disallowed
+          const isAccountInAllowedLocation = firstWriteLocation && disallowedLocations.indexOf(firstWriteLocation) === -1;
+          set({ isNotebooksEnabledForAccount: isAccountInAllowedLocation });
+        } catch (error) {
+          Logger.logError(getErrorMessage(error), "Explorer/isNotebooksEnabledForAccount");
+          set({ isNotebooksEnabledForAccount: false });
+        }
+      },
+      findItem: (root: NotebookContentItem, item: NotebookContentItem): NotebookContentItem => {
+        const currentItem = root || get().myNotebooksContentRoot;
+
+        if (currentItem) {
+          if (currentItem.path === item.path && currentItem.name === item.name) {
+            return currentItem;
+          }
+
+          if (currentItem.children) {
+            for (const childItem of currentItem.children) {
+              const result = get().findItem(childItem, item);
+              if (result) {
+                return result;
+              }
+            }
+          }
+        }
+
+        return undefined;
+      },
+      insertNotebookItem: (parent: NotebookContentItem, item: NotebookContentItem, isGithubTree?: boolean): void => {
+        const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
+        const parentItem = get().findItem(root, parent);
+        item.parent = parentItem;
+        if (parentItem.children) {
+          parentItem.children.push(item);
+        } else {
+          parentItem.children = [item];
+        }
+        isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
+      },
+      updateNotebookItem: (item: NotebookContentItem, isGithubTree?: boolean): void => {
+        const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
+        const parentItem = get().findItem(root, item.parent);
+        parentItem.children = parentItem.children.filter((child) => child.path !== item.path);
+        parentItem.children.push(item);
+        item.parent = parentItem;
+        isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
+      },
+      deleteNotebookItem: (item: NotebookContentItem, isGithubTree?: boolean): void => {
+        const root = isGithubTree ? cloneDeep(get().gitHubNotebooksContentRoot) : cloneDeep(get().myNotebooksContentRoot);
+        const parentItem = get().findItem(root, item.parent);
+        parentItem.children = parentItem.children.filter((child) => child.path !== item.path);
+        isGithubTree ? set({ gitHubNotebooksContentRoot: root }) : set({ myNotebooksContentRoot: root });
+      },
+      initializeNotebooksTree: async (notebookManager: NotebookManager): Promise<void> => {
+        const notebookFolderName = get().isPhoenixNotebooks ? "Temporary Notebooks" : "My Notebooks";
+        set({ notebookFolderName });
+        const myNotebooksContentRoot = {
+          name: get().notebookFolderName,
+          path: get().notebookBasePath,
+          type: NotebookContentItemType.Directory,
+        };
+        const galleryContentRoot = {
+          name: "Gallery",
+          path: "Gallery",
+          type: NotebookContentItemType.File,
+        };
+        const gitHubNotebooksContentRoot = notebookManager?.gitHubOAuthService?.isLoggedIn()
+          ? {
+            name: "GitHub repos",
+            path: "PsuedoDir",
+            type: NotebookContentItemType.Directory,
+          }
+          : undefined;
+
+        set({
+          myNotebooksContentRoot,
+          galleryContentRoot,
+          gitHubNotebooksContentRoot,
         });
 
-        gitHubNotebooksContentRoot.children.push(repoTreeItem);
-      });
+        if (get().notebookServerInfo?.notebookServerEndpoint) {
+          const updatedRoot = await notebookManager?.notebookContentClient?.updateItemChildren(myNotebooksContentRoot);
+          set({ myNotebooksContentRoot: updatedRoot });
 
-      set({ gitHubNotebooksContentRoot });
-    }
-  },
-  setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => set({ connectionInfo }),
-  setIsAllocating: (isAllocating: boolean) => set({ isAllocating }),
-  resetContainerConnection: (connectionStatus: ContainerConnectionInfo): void => {
-    useTabs.getState().closeAllNotebookTabs(true);
-    useNotebook.getState().setConnectionInfo(connectionStatus);
-    useNotebook.getState().setNotebookServerInfo(undefined);
-    useNotebook.getState().setIsAllocating(false);
-    useNotebook.getState().setContainerStatus({
-      status: undefined,
-      durationLeftInMinutes: undefined,
-      phoenixServerInfo: undefined,
-    });
-  },
-  setIsRefreshed: (isRefreshed: boolean) => set({ isRefreshed }),
-  setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
-  getPhoenixStatus: async () => {
-    if (get().isPhoenixNotebooks === undefined || get().isPhoenixFeatures === undefined) {
-      let isPhoenixNotebooks = false;
-      let isPhoenixFeatures = false;
-
-      const isPublicInternetAllowed = isPublicInternetAccessAllowed();
-      const phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
-      const dbAccountAllowedInfo = await phoenixClient.getDbAccountAllowedStatus();
-
-      if (dbAccountAllowedInfo.status === HttpStatusCodes.OK) {
-        if (dbAccountAllowedInfo?.type === PhoenixErrorType.PhoenixFlightFallback) {
-          isPhoenixNotebooks = isPublicInternetAllowed && userContext.features.phoenixNotebooks === true;
-          isPhoenixFeatures =
-            isPublicInternetAllowed &&
-            // phoenix needs to be enabled for Postgres and VCoreMongo accounts since the PSQL and mongo shell requires phoenix containers
-            (userContext.features.phoenixFeatures === true ||
-              userContext.apiType === "Postgres" ||
-              userContext.apiType === "VCoreMongo");
-        } else {
-          isPhoenixNotebooks = isPhoenixFeatures = isPublicInternetAllowed;
+          if (updatedRoot?.children) {
+            // Count 1st generation children (tree is lazy-loaded)
+            const nodeCounts = { files: 0, notebooks: 0, directories: 0 };
+            updatedRoot.children.forEach((notebookItem) => {
+              switch (notebookItem.type) {
+                case NotebookContentItemType.File:
+                  nodeCounts.files++;
+                  break;
+                case NotebookContentItemType.Directory:
+                  nodeCounts.directories++;
+                  break;
+                case NotebookContentItemType.Notebook:
+                  nodeCounts.notebooks++;
+                  break;
+                default:
+                  break;
+              }
+            });
+            TelemetryProcessor.trace(Action.RefreshResourceTreeMyNotebooks, ActionModifiers.Mark, { ...nodeCounts });
+          }
         }
-      } else {
-        isPhoenixNotebooks = isPhoenixFeatures = false;
-      }
-      set({ isPhoenixNotebooks: isPhoenixNotebooks });
-      set({ isPhoenixFeatures: isPhoenixFeatures });
-    }
-  },
-  setIsPhoenixNotebooks: (isPhoenixNotebooks: boolean) => set({ isPhoenixNotebooks: isPhoenixNotebooks }),
-  setIsPhoenixFeatures: (isPhoenixFeatures: boolean) => set({ isPhoenixFeatures: isPhoenixFeatures }),
-}));
+      },
+      initializeGitHubRepos: (pinnedRepos: IPinnedRepo[]): void => {
+        const gitHubNotebooksContentRoot = cloneDeep(get().gitHubNotebooksContentRoot);
+        if (gitHubNotebooksContentRoot) {
+          gitHubNotebooksContentRoot.children = [];
+          pinnedRepos?.forEach((pinnedRepo) => {
+            const repoFullName = GitHubUtils.toRepoFullName(pinnedRepo.owner, pinnedRepo.name);
+            const repoTreeItem: NotebookContentItem = {
+              name: repoFullName,
+              path: "PsuedoDir",
+              type: NotebookContentItemType.Directory,
+              children: [],
+              parent: gitHubNotebooksContentRoot,
+            };
+
+            pinnedRepo.branches.forEach((branch) => {
+              repoTreeItem.children.push({
+                name: branch.name,
+                path: GitHubUtils.toContentUri(pinnedRepo.owner, pinnedRepo.name, branch.name, ""),
+                type: NotebookContentItemType.Directory,
+                parent: repoTreeItem,
+              });
+            });
+
+            gitHubNotebooksContentRoot.children.push(repoTreeItem);
+          });
+
+          set({ gitHubNotebooksContentRoot });
+        }
+      },
+      setConnectionInfo: (connectionInfo: ContainerConnectionInfo) => set({ connectionInfo }),
+      setIsAllocating: (isAllocating: boolean) => set({ isAllocating }),
+      resetContainerConnection: (connectionStatus: ContainerConnectionInfo): void => {
+        useTabs.getState().closeAllNotebookTabs(true);
+        useNotebook.getState().setConnectionInfo(connectionStatus);
+        useNotebook.getState().setNotebookServerInfo(undefined);
+        useNotebook.getState().setIsAllocating(false);
+        useNotebook.getState().setContainerStatus({
+          status: undefined,
+          durationLeftInMinutes: undefined,
+          phoenixServerInfo: undefined,
+        });
+      },
+      setIsRefreshed: (isRefreshed: boolean) => set({ isRefreshed }),
+      setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
+      getPhoenixStatus: async () => {
+        if (get().isPhoenixNotebooks === undefined || get().isPhoenixFeatures === undefined) {
+          let isPhoenixNotebooks = false;
+          let isPhoenixFeatures = false;
+
+          const isPublicInternetAllowed = isPublicInternetAccessAllowed();
+          const phoenixClient = new PhoenixClient(userContext?.databaseAccount?.id);
+          const dbAccountAllowedInfo = await phoenixClient.getDbAccountAllowedStatus();
+
+          if (dbAccountAllowedInfo.status === HttpStatusCodes.OK) {
+            if (dbAccountAllowedInfo?.type === PhoenixErrorType.PhoenixFlightFallback) {
+              isPhoenixNotebooks = isPublicInternetAllowed && userContext.features.phoenixNotebooks === true;
+              isPhoenixFeatures =
+                isPublicInternetAllowed &&
+                // phoenix needs to be enabled for Postgres and VCoreMongo accounts since the PSQL and mongo shell requires phoenix containers
+                (userContext.features.phoenixFeatures === true ||
+                  userContext.apiType === "Postgres" ||
+                  userContext.apiType === "VCoreMongo");
+            } else {
+              isPhoenixNotebooks = isPhoenixFeatures = isPublicInternetAllowed;
+            }
+          } else {
+            isPhoenixNotebooks = isPhoenixFeatures = false;
+          }
+          set({ isPhoenixNotebooks: isPhoenixNotebooks });
+          set({ isPhoenixFeatures: isPhoenixFeatures });
+        }
+      },
+      setIsPhoenixNotebooks: (isPhoenixNotebooks: boolean) => set({ isPhoenixNotebooks: isPhoenixNotebooks }),
+      setIsPhoenixFeatures: (isPhoenixFeatures: boolean) => set({ isPhoenixFeatures: isPhoenixFeatures }),
+    })
+  )
+);

src/Explorer/Panes/SettingsPane/SettingsPane.tsx

@@ -51,7 +51,7 @@ import { useClientWriteEnabled } from "hooks/useClientWriteEnabled";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useSidePanel } from "hooks/useSidePanel";
 import React, { FunctionComponent, useState } from "react";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 import Explorer from "../../Explorer";
 import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm";
 
@@ -65,8 +65,6 @@ export interface DataPlaneRbacState {
   setAadDataPlaneUpdated: (aadTokenUpdated: boolean) => void;
 }
 
-type DataPlaneRbacStore = UseStore<Partial<DataPlaneRbacState>>;
-
 const useStyles = makeStyles({
   bulletList: {
     listStyleType: "disc",
@@ -100,7 +98,7 @@ const useStyles = makeStyles({
   },
 });
 
-export const useDataPlaneRbac: DataPlaneRbacStore = create(() => ({
+export const useDataPlaneRbac = create<Partial<DataPlaneRbacState>>(() => ({
   dataPlaneRbacEnabled: false,
 }));
 

src/Explorer/QueryCopilot/QueryCopilotContext.tsx

@@ -1,12 +1,14 @@
 import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import QueryError from "Common/QueryError";
+import * as DataModels from "Contracts/DataModels";
 import { QueryResults } from "Contracts/ViewModels";
 import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { guid } from "Explorer/Tables/Utilities";
 import { QueryCopilotState } from "hooks/useQueryCopilot";
 
 import React, { createContext, useContext, useState } from "react";
-import create from "zustand";
+import { create } from "zustand";
+
 const context = createContext(null);
 const useCopilotStore = (): Partial<QueryCopilotState> => useContext(context);
 
@@ -24,12 +26,12 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
       isGeneratingQuery: false,
       isGeneratingExplanation: false,
       isExecuting: false,
-      dislikeQuery: undefined,
+      dislikeQuery: undefined as boolean,
       showCallout: false,
       showSamplePrompts: false,
-      queryIterator: undefined,
-      queryResults: undefined,
-      errors: [],
+      queryIterator: undefined as MinimalQueryIterator,
+      queryResults: undefined as QueryResults,
+      errors: [] as QueryError[],
       isSamplePromptsOpen: false,
       showPromptTeachingBubble: true,
       showDeletePopup: false,
@@ -41,7 +43,7 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
       wasCopilotUsed: false,
       showWelcomeSidebar: true,
       showCopilotSidebar: false,
-      chatMessages: [],
+      chatMessages: [] as CopilotMessage[],
       shouldIncludeInMessages: true,
       showExplanationBubble: false,
       isAllocatingContainer: false,
@@ -86,7 +88,7 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
       },
 
       resetQueryCopilotStates: () => {
-        set((state) => ({
+        set((state: QueryCopilotState) => ({
           ...state,
           generatedQuery: "",
           likeQuery: false,
@@ -99,11 +101,11 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
           isGeneratingQuery: false,
           isGeneratingExplanation: false,
           isExecuting: false,
-          dislikeQuery: undefined,
+          dislikeQuery: undefined as boolean,
           showCallout: false,
           showSamplePrompts: false,
-          queryIterator: undefined,
-          queryResults: undefined,
+          queryIterator: undefined as MinimalQueryIterator,
+          queryResults: undefined as QueryResults,
           errorMessage: "",
           isSamplePromptsOpen: false,
           showPromptTeachingBubble: true,
@@ -115,19 +117,19 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
           generatedQueryComments: "",
           wasCopilotUsed: false,
           showCopilotSidebar: false,
-          chatMessages: [],
+          chatMessages: [] as CopilotMessage[],
           shouldIncludeInMessages: true,
           showExplanationBubble: false,
           notebookServerInfo: {
             notebookServerEndpoint: undefined,
             authToken: undefined,
             forwardingId: undefined,
-          },
+          } as DataModels.NotebookWorkspaceConnectionInfo,
           containerStatus: {
             status: undefined,
             durationLeftInMinutes: undefined,
             phoenixServerInfo: undefined,
-          },
+          } as DataModels.ContainerInfo,
           isAllocatingContainer: false,
         }));
       },
@@ -137,3 +139,4 @@ const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Eleme
 };
 
 export { CopilotProvider, useCopilotStore };
+

src/Explorer/SplashScreen/SplashScreen.tsx

@@ -77,39 +77,39 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
     this.subscriptions.push(
       {
         dispose: useNotebook.subscribe(
-          () => this.setState({}),
           (state) => state.isNotebookEnabled,
+          () => this.setState({}),
         ),
       },
       { dispose: useSelectedNode.subscribe(() => this.setState({})) },
       {
         dispose: useCarousel.subscribe(
-          () => this.setState({}),
           (state) => state.showCoachMark,
+          () => this.setState({}),
         ),
       },
       {
         dispose: usePostgres.subscribe(
-          () => this.setState({}),
           (state) => state.showPostgreTeachingBubble,
+          () => this.setState({}),
         ),
       },
       {
         dispose: usePostgres.subscribe(
-          () => this.setState({}),
           (state) => state.showResetPasswordBubble,
+          () => this.setState({}),
         ),
       },
       {
         dispose: useDatabases.subscribe(
-          () => this.setState({}),
           (state) => state.sampleDataResourceTokenCollection,
+          () => this.setState({}),
         ),
       },
       {
         dispose: useQueryCopilot.subscribe(
-          () => this.setState({}),
           (state) => state.copilotEnabled,
+          () => this.setState({}),
         ),
       },
     );

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/CloudShellTerminalComponent.tsx

@@ -1,7 +1,7 @@
 import { FitAddon } from "@xterm/addon-fit";
 import { Terminal } from "@xterm/xterm";
+import "@xterm/xterm/css/xterm.css";
 import React, { useEffect, useRef } from "react";
-import "xterm/css/xterm.css";
 import { DatabaseAccount } from "../../../Contracts/DataModels";
 import { TerminalKind } from "../../../Contracts/ViewModels";
 import { startCloudShellTerminal } from "./CloudShellTerminalCore";

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();

src/Explorer/Tabs/NotebookV2Tab.ts

@@ -44,8 +44,8 @@ export default class NotebookTabV2 extends NotebookTabBase {
     this.container = options.container;
     this.notebookPath = ko.observable(options.notebookContentItem.path);
     useNotebook.subscribe(
-      () => logConsoleInfo("New notebook server info received."),
       (state) => state.notebookServerInfo,
+      () => logConsoleInfo("New notebook server info received."),
     );
     this.notebookComponentAdapter = new NotebookComponentAdapter({
       contentItem: options.notebookContentItem,
@@ -165,7 +165,7 @@ export default class NotebookTabV2 extends NotebookTabBase {
       {
         iconSrc: null,
         iconAlt: kernelLabel,
-        onCommandClick: () => {},
+        onCommandClick: () => { },
         commandButtonLabel: null,
         hasPopup: false,
         disabled: availableKernels.length < 1,
@@ -276,7 +276,7 @@ export default class NotebookTabV2 extends NotebookTabBase {
       {
         iconSrc: null,
         iconAlt: null,
-        onCommandClick: () => {},
+        onCommandClick: () => { },
         commandButtonLabel: null,
         ariaLabel: cellTypeLabel,
         hasPopup: false,

src/Explorer/Tree/ResourceTree.tsx

@@ -16,7 +16,6 @@ import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
 import * as React from "react";
 import { useEffect, useMemo } from "react";
-import shallow from "zustand/shallow";
 import Explorer from "../Explorer";
 import { useNotebook } from "../Notebook/useNotebook";
 
@@ -38,13 +37,8 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ explorer }: Resource
   const [openItems, setOpenItems] = React.useState<TreeItemValue[]>([]);
   const treeStyles = useTreeStyles();
 
-  const { isNotebookEnabled } = useNotebook(
-    (state) => ({
-      isNotebookEnabled: state.isNotebookEnabled,
-    }),
-    shallow,
-  );
-
+  const isNotebookEnabled = useNotebook((state) => state.isNotebookEnabled)
+  
   // We intentionally avoid using a state selector here because we want to re-render the tree if the active tab changes.
   const { refreshActiveTab } = useTabs();
 

src/Explorer/Tree/ResourceTreeAdapter.tsx

@@ -14,7 +14,6 @@ import PublishIcon from "../../../images/notebook/publish_content.svg";
 import RefreshIcon from "../../../images/refresh-cosmos.svg";
 import CollectionIcon from "../../../images/tree-collection.svg";
 import { ReactAdapter } from "../../Bindings/ReactBindingHandler";
-import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { IPinnedRepo } from "../../Juno/JunoClient";
@@ -58,12 +57,12 @@ export class ResourceTreeAdapter implements ReactAdapter {
 
     useSelectedNode.subscribe(() => this.triggerRender());
     useTabs.subscribe(
-      () => this.triggerRender(),
       (state) => state.activeTab,
+      () => this.triggerRender(),
     );
     useNotebook.subscribe(
-      () => this.triggerRender(),
       (state) => state.isNotebookEnabled,
+      () => this.triggerRender(),
     );
 
     useDatabases.subscribe(() => this.triggerRender());

src/Explorer/useDatabases.ts

@@ -1,5 +1,6 @@
 import _ from "underscore";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 import * as Constants from "../Common/Constants";
 import * as ViewModels from "../Contracts/ViewModels";
 import { userContext } from "../UserContext";
@@ -26,143 +27,147 @@ interface DatabasesState {
   validateCollectionId: (databaseId: string, collectionId: string) => Promise<boolean>;
 }
 
-export const useDatabases: UseStore<DatabasesState> = create((set, get) => ({
-  databases: [],
-  resourceTokenCollection: undefined,
-  sampleDataResourceTokenCollection: undefined,
-  updateDatabase: (updatedDatabase: ViewModels.Database) =>
-    set((state) => {
-      const updatedDatabases = state.databases.map((database: ViewModels.Database) => {
-        if (database?.id() === updatedDatabase?.id()) {
-          return updatedDatabase;
+export const useDatabases = create<DatabasesState>()(
+  subscribeWithSelector(
+    (set, get) => ({
+      databases: [] as ViewModels.Database[],
+      resourceTokenCollection: undefined as ViewModels.CollectionBase,
+      sampleDataResourceTokenCollection: undefined as ViewModels.CollectionBase,
+      updateDatabase: (updatedDatabase: ViewModels.Database) =>
+        set((state) => {
+          const updatedDatabases = state.databases.map((database: ViewModels.Database) => {
+            if (database?.id() === updatedDatabase?.id()) {
+              return updatedDatabase;
+            }
+
+            return database;
+          });
+          return { databases: updatedDatabases };
+        }),
+      addDatabases: (databases: ViewModels.Database[]) =>
+        set((state) => ({
+          databases: [...state.databases, ...databases].sort((db1, db2) => db1.id().localeCompare(db2.id())),
+        })),
+      deleteDatabase: (database: ViewModels.Database) =>
+        set((state) => ({ databases: state.databases.filter((db) => database.id() !== db.id()) })),
+      clearDatabases: () => set(() => ({ databases: [] })),
+      isSaveQueryEnabled: () => {
+        const savedQueriesDatabase: ViewModels.Database = _.find(
+          get().databases,
+          (database: ViewModels.Database) => database.id() === Constants.SavedQueries.DatabaseName,
+        );
+        if (!savedQueriesDatabase) {
+          return false;
+        }
+        const savedQueriesCollection: ViewModels.Collection =
+          savedQueriesDatabase &&
+          _.find(
+            savedQueriesDatabase.collections(),
+            (collection: ViewModels.Collection) => collection.id() === Constants.SavedQueries.CollectionName,
+          );
+        if (!savedQueriesCollection) {
+          return false;
+        }
+        return true;
+      },
+      findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => {
+        return isSampleDatabase === undefined
+          ? get().databases.find((db) => databaseId === db.id())
+          : get().databases.find((db) => databaseId === db.id() && db.isSampleDB === isSampleDatabase);
+      },
+      isLastNonEmptyDatabase: () => {
+        const databases = get().databases;
+        return databases.length === 1 && (databases[0].collections()?.length > 0 || !!databases[0].offer());
+      },
+      findCollection: (databaseId: string, collectionId: string) => {
+        const database = get().findDatabaseWithId(databaseId);
+        return database?.collections()?.find((collection) => collection.id() === collectionId);
+      },
+      isLastCollection: () => {
+        const databases = get().databases;
+        if (databases.length === 0) {
+          return false;
         }
 
-        return database;
-      });
-      return { databases: updatedDatabases };
-    }),
-  addDatabases: (databases: ViewModels.Database[]) =>
-    set((state) => ({
-      databases: [...state.databases, ...databases].sort((db1, db2) => db1.id().localeCompare(db2.id())),
-    })),
-  deleteDatabase: (database: ViewModels.Database) =>
-    set((state) => ({ databases: state.databases.filter((db) => database.id() !== db.id()) })),
-  clearDatabases: () => set(() => ({ databases: [] })),
-  isSaveQueryEnabled: () => {
-    const savedQueriesDatabase: ViewModels.Database = _.find(
-      get().databases,
-      (database: ViewModels.Database) => database.id() === Constants.SavedQueries.DatabaseName,
-    );
-    if (!savedQueriesDatabase) {
-      return false;
-    }
-    const savedQueriesCollection: ViewModels.Collection =
-      savedQueriesDatabase &&
-      _.find(
-        savedQueriesDatabase.collections(),
-        (collection: ViewModels.Collection) => collection.id() === Constants.SavedQueries.CollectionName,
-      );
-    if (!savedQueriesCollection) {
-      return false;
-    }
-    return true;
-  },
-  findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => {
-    return isSampleDatabase === undefined
-      ? get().databases.find((db) => databaseId === db.id())
-      : get().databases.find((db) => databaseId === db.id() && db.isSampleDB === isSampleDatabase);
-  },
-  isLastNonEmptyDatabase: () => {
-    const databases = get().databases;
-    return databases.length === 1 && (databases[0].collections()?.length > 0 || !!databases[0].offer());
-  },
-  findCollection: (databaseId: string, collectionId: string) => {
-    const database = get().findDatabaseWithId(databaseId);
-    return database?.collections()?.find((collection) => collection.id() === collectionId);
-  },
-  isLastCollection: () => {
-    const databases = get().databases;
-    if (databases.length === 0) {
-      return false;
-    }
+        let collectionCount = 0;
+        for (let i = 0; i < databases.length; i++) {
+          const database = databases[i];
+          collectionCount += database.collections().length;
+          if (collectionCount > 1) {
+            return false;
+          }
+        }
 
-    let collectionCount = 0;
-    for (let i = 0; i < databases.length; i++) {
-      const database = databases[i];
-      collectionCount += database.collections().length;
-      if (collectionCount > 1) {
-        return false;
-      }
-    }
-
-    return true;
-  },
-  loadDatabaseOffers: async () => {
-    await Promise.all(
-      get().databases?.map(async (database: ViewModels.Database) => {
-        await database.loadOffer();
-      }),
-    );
-  },
-  loadAllOffers: async () => {
-    await Promise.all(
-      get().databases?.map(async (database: ViewModels.Database) => {
-        await database.loadOffer();
-        await database.loadCollections();
+        return true;
+      },
+      loadDatabaseOffers: async () => {
         await Promise.all(
-          (database.collections() || []).map(async (collection: ViewModels.Collection) => {
-            await collection.loadOffer();
+          get().databases?.map(async (database: ViewModels.Database) => {
+            await database.loadOffer();
           }),
         );
-      }),
-    );
-  },
-  isFirstResourceCreated: () => {
-    const databases = get().databases;
+      },
+      loadAllOffers: async () => {
+        await Promise.all(
+          get().databases?.map(async (database: ViewModels.Database) => {
+            await database.loadOffer();
+            await database.loadCollections();
+            await Promise.all(
+              (database.collections() || []).map(async (collection: ViewModels.Collection) => {
+                await collection.loadOffer();
+              }),
+            );
+          }),
+        );
+      },
+      isFirstResourceCreated: () => {
+        const databases = get().databases;
 
-    if (!databases || databases.length === 0) {
-      return false;
-    }
+        if (!databases || databases.length === 0) {
+          return false;
+        }
 
-    return databases.some((database) => {
-      // user has created at least one collection
-      if (database.collections()?.length > 0) {
-        return true;
-      }
-      // user has created a database with shared throughput
-      if (database.offer()) {
-        return true;
-      }
-      // use has created an empty database without shared throughput
-      return false;
-    });
-  },
-  findSelectedDatabase: (): ViewModels.Database => {
-    const selectedNode = useSelectedNode.getState().selectedNode;
-    if (!selectedNode) {
-      return undefined;
-    }
-    if (selectedNode.nodeKind === "Database") {
-      return _.find(get().databases, (database: ViewModels.Database) => database.id() === selectedNode.id());
-    }
+        return databases.some((database) => {
+          // user has created at least one collection
+          if (database.collections()?.length > 0) {
+            return true;
+          }
+          // user has created a database with shared throughput
+          if (database.offer()) {
+            return true;
+          }
+          // use has created an empty database without shared throughput
+          return false;
+        });
+      },
+      findSelectedDatabase: (): ViewModels.Database => {
+        const selectedNode = useSelectedNode.getState().selectedNode;
+        if (!selectedNode) {
+          return undefined;
+        }
+        if (selectedNode.nodeKind === "Database") {
+          return _.find(get().databases, (database: ViewModels.Database) => database.id() === selectedNode.id());
+        }
 
-    if (selectedNode.nodeKind === "Collection") {
-      return selectedNode.database;
-    }
+        if (selectedNode.nodeKind === "Collection") {
+          return selectedNode.database;
+        }
 
-    return selectedNode.collection?.database;
-  },
-  validateDatabaseId: (id: string): boolean => {
-    return !get().databases.some((database) => database.id() === id);
-  },
-  validateCollectionId: async (databaseId: string, collectionId: string): Promise<boolean> => {
-    const database = get().databases.find((db) => db.id() === databaseId);
-    // For a new tables account, database is undefined when creating the first table
-    if (!database && userContext.apiType === "Tables") {
-      return true;
-    }
+        return selectedNode.collection?.database;
+      },
+      validateDatabaseId: (id: string): boolean => {
+        return !get().databases.some((database) => database.id() === id);
+      },
+      validateCollectionId: async (databaseId: string, collectionId: string): Promise<boolean> => {
+        const database = get().databases.find((db) => db.id() === databaseId);
+        // For a new tables account, database is undefined when creating the first table
+        if (!database && userContext.apiType === "Tables") {
+          return true;
+        }
 
-    await database.loadCollections();
-    return !database.collections().some((collection) => collection.id() === collectionId);
-  },
-}));
+        await database.loadCollections();
+        return !database.collections().some((collection) => collection.id() === collectionId);
+      },
+    })
+  )
+);

src/Explorer/useSelectedNode.ts

@@ -1,6 +1,6 @@
 import { ConnectionStatusType, QueryCopilotSampleContainerId, QueryCopilotSampleDatabaseId } from "Common/Constants";
 import { useNotebook } from "Explorer/Notebook/useNotebook";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 import * as ViewModels from "../Contracts/ViewModels";
 import { useTabs } from "../hooks/useTabs";
 export interface SelectedNodeState {
@@ -17,7 +17,7 @@ export interface SelectedNodeState {
   isQueryCopilotCollectionSelected: () => boolean;
 }
 
-export const useSelectedNode: UseStore<SelectedNodeState> = create((set, get) => ({
+export const useSelectedNode = create<SelectedNodeState>((set, get) => ({
   selectedNode: undefined,
   setSelectedNode: (node: ViewModels.TreeNode) => set({ selectedNode: node }),
   isDatabaseNodeOrNoneSelected: (): boolean => {

src/KeyboardShortcuts.tsx

@@ -1,7 +1,7 @@
 import * as React from "react";
 import { PropsWithChildren, useEffect } from "react";
 import { KeyBindingMap, tinykeys } from "tinykeys";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 /**
  * Represents a keyboard shortcut handler.
@@ -126,7 +126,7 @@ export const clearKeyboardActionGroup = (group: KeyboardActionGroup) => {
   useKeyboardActionHandlers.getState().setHandlers(group, {});
 };
 
-const useKeyboardActionHandlers: UseStore<KeyboardShortcutState> = create((set, get) => ({
+const useKeyboardActionHandlers = create<KeyboardShortcutState>((set, get) => ({
   allHandlers: {},
   groups: {},
   setHandlers: (group: KeyboardActionGroup, handlers: KeyboardHandlerMap) => {

src/Utils/APITypeUtils.ts

@@ -91,5 +91,11 @@ export const getItemName = (): string => {
 };
 
 export const isDataplaneRbacSupported = (apiType: string): boolean => {
-  return apiType === "SQL" || apiType === "Tables" || apiType === "Gremlin";
+  return (
+    apiType === "SQL" || apiType === "Tables" || apiType === "Gremlin" || apiType === "Mongo" || apiType === "Cassandra"
+  );
+};
+
+export const hasProxyServer = (apiType: string): boolean => {
+  return apiType === "Mongo" || apiType === "Cassandra";
 };

src/Utils/AuthorizationUtils.test.ts

@@ -104,7 +104,7 @@ describe("AuthorizationUtils", () => {
 
     it("should return true if dataPlaneRbacEnabled is set to true and API supports RBAC", () => {
       setAadDataPlane(false);
-      ["SQL", "Tables", "Gremlin"].forEach((type) => {
+      ["SQL", "Tables", "Gremlin", "Mongo", "Cassandra"].forEach((type) => {
         updateUserContext({
           dataPlaneRbacEnabled: true,
           apiType: type as ApiType,
@@ -115,7 +115,7 @@ describe("AuthorizationUtils", () => {
 
     it("should return false if dataPlaneRbacEnabled is set to true and API does not support RBAC", () => {
       setAadDataPlane(false);
-      ["Mongo", "Cassandra", "Postgres", "VCoreMongo"].forEach((type) => {
+      ["Postgres", "VCoreMongo"].forEach((type) => {
         updateUserContext({
           dataPlaneRbacEnabled: true,
           apiType: type as ApiType,

src/Utils/AuthorizationUtils.ts

@@ -1,6 +1,7 @@
 import * as msal from "@azure/msal-browser";
+import { getEnvironmentScopeEndpoint } from "Common/EnvironmentUtility";
 import { Action, ActionModifiers } from "Shared/Telemetry/TelemetryConstants";
-import { isDataplaneRbacSupported } from "Utils/APITypeUtils";
+import { hasProxyServer, isDataplaneRbacSupported } from "Utils/APITypeUtils";
 import { AuthType } from "../AuthType";
 import * as Constants from "../Common/Constants";
 import * as Logger from "../Common/Logger";
@@ -74,10 +75,12 @@ export async function acquireMsalTokenForAccount(
   if (userContext.databaseAccount.properties?.documentEndpoint === undefined) {
     throw new Error("Database account has no document endpoint defined");
   }
-  const hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace(
-    /\/+$/,
-    "/.default",
-  );
+  let hrefEndpoint = "";
+  if (isDataplaneRbacEnabledForProxyApi(userContext)) {
+    hrefEndpoint = getEnvironmentScopeEndpoint();
+  } else {
+    hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace(/\/+$/, "/.default");
+  }
   const msalInstance = await getMsalInstance();
   const knownAccounts = msalInstance.getAllAccounts();
   // If user_hint is provided, we will try to use it to find the account.
@@ -183,7 +186,11 @@ export async function acquireTokenWithMsal(
 
 export function useDataplaneRbacAuthorization(userContext: UserContext): boolean {
   return (
-    userContext.features.enableAadDataPlane ||
+    userContext.features?.enableAadDataPlane ||
     (userContext.dataPlaneRbacEnabled && isDataplaneRbacSupported(userContext.apiType))
   );
 }
+
+export function isDataplaneRbacEnabledForProxyApi(userContext: UserContext): boolean {
+  return useDataplaneRbacAuthorization(userContext) && hasProxyServer(userContext.apiType);
+}

src/hooks/useCarousel.ts

@@ -1,4 +1,5 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 
 interface CarouselState {
   shouldOpen: boolean;
@@ -9,11 +10,15 @@ interface CarouselState {
   setShowCopilotCarousel: (showCopilotCarousel: boolean) => void;
 }
 
-export const useCarousel: UseStore<CarouselState> = create((set) => ({
-  shouldOpen: false,
-  showCoachMark: false,
-  showCopilotCarousel: false,
-  setShouldOpen: (shouldOpen: boolean) => set({ shouldOpen }),
-  setShowCoachMark: (showCoachMark: boolean) => set({ showCoachMark }),
-  setShowCopilotCarousel: (showCopilotCarousel: boolean) => set({ showCopilotCarousel }),
-}));
+export const useCarousel = create<CarouselState>()(
+  subscribeWithSelector(
+    (set) => ({
+      shouldOpen: false,
+      showCoachMark: false,
+      showCopilotCarousel: false,
+      setShouldOpen: (shouldOpen: boolean) => set({ shouldOpen }),
+      setShowCoachMark: (showCoachMark: boolean) => set({ showCoachMark }),
+      setShowCopilotCarousel: (showCopilotCarousel: boolean) => set({ showCopilotCarousel }),
+    })
+  )
+);

src/hooks/useClientWriteEnabled.ts

@@ -1,10 +1,10 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 interface ClientWriteEnabledState {
   clientWriteEnabled: boolean;
   setClientWriteEnabled: (writeEnabled: boolean) => void;
 }
 
-export const useClientWriteEnabled: UseStore<ClientWriteEnabledState> = create((set) => ({
+export const useClientWriteEnabled = create<ClientWriteEnabledState>((set) => ({
   clientWriteEnabled: true,
   setClientWriteEnabled: (clientWriteEnabled: boolean) => set({ clientWriteEnabled }),
 }));

src/hooks/useDataTransferJobs.tsx

@@ -1,6 +1,6 @@
 import { getDataTransferJobs } from "Common/dataAccess/dataTransfers";
 import { DataTransferJobGetResults } from "Utils/arm/generatedClients/dataTransferService/types";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 export interface DataTransferJobsState {
   dataTransferJobs: DataTransferJobGetResults[];
@@ -9,9 +9,7 @@ export interface DataTransferJobsState {
   setPollingDataTransferJobs: (pollingDataTransferJobs: Set<string>) => void;
 }
 
-type DataTransferJobStore = UseStore<DataTransferJobsState>;
-
-export const useDataTransferJobs: DataTransferJobStore = create((set) => ({
+export const useDataTransferJobs = create<DataTransferJobsState>((set) => ({
   dataTransferJobs: [],
   pollingDataTransferJobs: new Set<string>(),
   setDataTransferJobs: (dataTransferJobs: DataTransferJobGetResults[]) => set({ dataTransferJobs }),

src/hooks/useKnockoutExplorer.ts

@@ -1,4 +1,5 @@
 import * as Constants from "Common/Constants";
+import { getEnvironmentScopeEndpoint } from "Common/EnvironmentUtility";
 import { createUri } from "Common/UrlUtility";
 import { DATA_EXPLORER_RPC_VERSION } from "Contracts/DataExplorerMessagesContract";
 import { FabricMessageTypes } from "Contracts/FabricMessageTypes";
@@ -62,6 +63,7 @@ import {
   acquireTokenWithMsal,
   getAuthorizationHeader,
   getMsalInstance,
+  isDataplaneRbacEnabledForProxyApi,
 } from "../Utils/AuthorizationUtils";
 import { isInvalidParentFrameOrigin, shouldProcessMessage } from "../Utils/MessageValidation";
 import { get, getReadOnlyKeys, listKeys } from "../Utils/arm/generatedClients/cosmos/databaseAccounts";
@@ -331,7 +333,12 @@ async function configureHostedWithAAD(config: AAD): Promise<Explorer> {
   const resourceGroup = accountResourceId && accountResourceId.split("resourceGroups/")[1].split("/")[0];
   let aadToken;
   if (account.properties?.documentEndpoint) {
-    const hrefEndpoint = new URL(account.properties.documentEndpoint).href.replace(/\/$/, "/.default");
+    let hrefEndpoint = "";
+    if (isDataplaneRbacEnabledForProxyApi(userContext)) {
+      hrefEndpoint = getEnvironmentScopeEndpoint();
+    } else {
+      hrefEndpoint = new URL(account.properties.documentEndpoint).href.replace(/\/$/, "/.default");
+    }
     const msalInstance = await getMsalInstance();
     const cachedAccount = msalInstance.getAllAccounts()?.[0];
     msalInstance.setActiveAccount(cachedAccount);

src/hooks/useNotebookSnapshotStore.ts

@@ -1,4 +1,4 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 export interface NotebookSnapshotHooks {
   snapshot?: string;
@@ -7,7 +7,7 @@ export interface NotebookSnapshotHooks {
   setError: (error: string) => void;
 }
 
-export const useNotebookSnapshotStore: UseStore<NotebookSnapshotHooks> = create((set) => ({
+export const useNotebookSnapshotStore = create<NotebookSnapshotHooks>((set) => ({
   snapshot: undefined,
   error: undefined,
   setSnapshot: (imageSrc: string) => set((state) => ({ ...state, snapshot: imageSrc })),

src/hooks/useNotificationConsole.ts

@@ -1,4 +1,4 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 import { ConsoleData } from "../Explorer/Menus/NotificationConsole/ConsoleData";
 
 export interface NotificationConsoleState {
@@ -15,7 +15,7 @@ export interface NotificationConsoleState {
   setConsoleAnimationFinished: (consoleAnimationFinished: boolean) => void;
 }
 
-export const useNotificationConsole: UseStore<NotificationConsoleState> = create((set) => ({
+export const useNotificationConsole = create<NotificationConsoleState>((set) => ({
   isExpanded: false,
   consoleData: undefined,
   inProgressConsoleDataIdToBeDeleted: "",

src/hooks/usePostgres.ts

@@ -1,4 +1,5 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 
 interface TeachingBubbleState {
   showPostgreTeachingBubble: boolean;
@@ -7,9 +8,13 @@ interface TeachingBubbleState {
   setShowResetPasswordBubble: (showResetPasswordBubble: boolean) => void;
 }
 
-export const usePostgres: UseStore<TeachingBubbleState> = create((set) => ({
-  showPostgreTeachingBubble: false,
-  showResetPasswordBubble: false,
-  setShowPostgreTeachingBubble: (showPostgreTeachingBubble: boolean) => set({ showPostgreTeachingBubble }),
-  setShowResetPasswordBubble: (showResetPasswordBubble: boolean) => set({ showResetPasswordBubble }),
-}));
+export const usePostgres = create<TeachingBubbleState>()(
+  subscribeWithSelector(
+    (set) => ({
+      showPostgreTeachingBubble: false,
+      showResetPasswordBubble: false,
+      setShowPostgreTeachingBubble: (showPostgreTeachingBubble: boolean) => set({ showPostgreTeachingBubble }),
+      setShowResetPasswordBubble: (showResetPasswordBubble: boolean) => set({ showResetPasswordBubble }),
+    })
+  )
+);

src/hooks/useQueryCopilot.ts

@@ -4,7 +4,8 @@ import { QueryResults } from "Contracts/ViewModels";
 import { CopilotMessage, CopilotSchemaAllocationInfo } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { guid } from "Explorer/Tables/Utilities";
 import { useTabs } from "hooks/useTabs";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 import * as DataModels from "../Contracts/DataModels";
 import { ContainerInfo } from "../Contracts/DataModels";
 
@@ -96,120 +97,12 @@ export interface QueryCopilotState {
   resetQueryCopilotStates: () => void;
 }
 
-type QueryCopilotStore = UseStore<Partial<QueryCopilotState>>;
-
-export const useQueryCopilot: QueryCopilotStore = create((set) => ({
-  copilotEnabled: false,
-  copilotUserDBEnabled: false,
-  copilotSampleDBEnabled: false,
-  generatedQuery: "",
-  likeQuery: false,
-  userPrompt: "",
-  showFeedbackModal: false,
-  hideFeedbackModalForLikedQueries: false,
-  correlationId: "",
-  query: "SELECT * FROM c",
-  selectedQuery: "",
-  isGeneratingQuery: null,
-  isGeneratingExplanation: false,
-  isExecuting: false,
-  dislikeQuery: undefined,
-  showCallout: false,
-  showSamplePrompts: false,
-  queryIterator: undefined,
-  queryResults: undefined,
-  errors: [],
-  isSamplePromptsOpen: false,
-  showDeletePopup: false,
-  showFeedbackBar: false,
-  showCopyPopup: false,
-  showErrorMessageBar: false,
-  showInvalidQueryMessageBar: false,
-  generatedQueryComments: "",
-  wasCopilotUsed: false,
-  showWelcomeSidebar: true,
-  showCopilotSidebar: false,
-  chatMessages: [],
-  shouldIncludeInMessages: true,
-  showExplanationBubble: false,
-  notebookServerInfo: {
-    notebookServerEndpoint: undefined,
-    authToken: undefined,
-    forwardingId: undefined,
-  },
-  containerStatus: {
-    status: undefined,
-    durationLeftInMinutes: undefined,
-    phoenixServerInfo: undefined,
-  },
-  schemaAllocationInfo: {
-    databaseId: undefined,
-    containerId: undefined,
-  },
-  isAllocatingContainer: false,
-  copilotEnabledforExecution: false,
-
-  setCopilotEnabled: (copilotEnabled: boolean) => set({ copilotEnabled }),
-  setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => set({ copilotUserDBEnabled }),
-  setCopilotSampleDBEnabled: (copilotSampleDBEnabled: boolean) => set({ copilotSampleDBEnabled }),
-  openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
-    set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
-  closeFeedbackModal: () => set({ showFeedbackModal: false }),
-  setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) =>
-    set({ hideFeedbackModalForLikedQueries }),
-  refreshCorrelationId: () => set({ correlationId: guid() }),
-  setUserPrompt: (userPrompt: string) => set({ userPrompt }),
-  setQuery: (query: string) => set({ query }),
-  setGeneratedQuery: (generatedQuery: string) => set({ generatedQuery }),
-  setSelectedQuery: (selectedQuery: string) => set({ selectedQuery }),
-  setIsGeneratingQuery: (isGeneratingQuery: boolean) => set({ isGeneratingQuery }),
-  setIsGeneratingExplanation: (isGeneratingExplanation: boolean) => set({ isGeneratingExplanation }),
-  setIsExecuting: (isExecuting: boolean) => set({ isExecuting }),
-  setLikeQuery: (likeQuery: boolean) => set({ likeQuery }),
-  setDislikeQuery: (dislikeQuery: boolean | undefined) => set({ dislikeQuery }),
-  setShowCallout: (showCallout: boolean) => set({ showCallout }),
-  setShowSamplePrompts: (showSamplePrompts: boolean) => set({ showSamplePrompts }),
-  setQueryIterator: (queryIterator: MinimalQueryIterator | undefined) => set({ queryIterator }),
-  setQueryResults: (queryResults: QueryResults | undefined) => set({ queryResults }),
-  setErrors: (errors: QueryError[]) => set({ errors }),
-  setIsSamplePromptsOpen: (isSamplePromptsOpen: boolean) => set({ isSamplePromptsOpen }),
-  setShowDeletePopup: (showDeletePopup: boolean) => set({ showDeletePopup }),
-  setShowFeedbackBar: (showFeedbackBar: boolean) => set({ showFeedbackBar }),
-  setshowCopyPopup: (showCopyPopup: boolean) => set({ showCopyPopup }),
-  setShowErrorMessageBar: (showErrorMessageBar: boolean) => set({ showErrorMessageBar }),
-  setShowInvalidQueryMessageBar: (showInvalidQueryMessageBar: boolean) => set({ showInvalidQueryMessageBar }),
-  setGeneratedQueryComments: (generatedQueryComments: string) => set({ generatedQueryComments }),
-  setWasCopilotUsed: (wasCopilotUsed: boolean) => set({ wasCopilotUsed }),
-  setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
-  setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
-  setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
-  setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
-  setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
-  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
-    set({ notebookServerInfo }),
-  setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
-  setIsAllocatingContainer: (isAllocatingContainer: boolean) => set({ isAllocatingContainer }),
-  setSchemaAllocationInfo: (schemaAllocationInfo: CopilotSchemaAllocationInfo) => set({ schemaAllocationInfo }),
-  setCopilotEnabledforExecution: (copilotEnabledforExecution: boolean) => set({ copilotEnabledforExecution }),
-
-  resetContainerConnection: (): void => {
-    useTabs.getState().closeAllNotebookTabs(true);
-    useQueryCopilot.getState().setNotebookServerInfo(undefined);
-    useQueryCopilot.getState().setIsAllocatingContainer(false);
-    useQueryCopilot.getState().setContainerStatus({
-      status: undefined,
-      durationLeftInMinutes: undefined,
-      phoenixServerInfo: undefined,
-    });
-    useQueryCopilot.getState().setSchemaAllocationInfo({
-      databaseId: undefined,
-      containerId: undefined,
-    });
-  },
-
-  resetQueryCopilotStates: () => {
-    set((state) => ({
-      ...state,
+export const useQueryCopilot = create<Partial<QueryCopilotState>>()(
+  subscribeWithSelector(
+    (set) => ({
+      copilotEnabled: false,
+      copilotUserDBEnabled: false,
+      copilotSampleDBEnabled: false,
       generatedQuery: "",
       likeQuery: false,
       userPrompt: "",
@@ -218,15 +111,15 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       correlationId: "",
       query: "SELECT * FROM c",
       selectedQuery: "",
-      isGeneratingQuery: false,
+      isGeneratingQuery: null as boolean,
       isGeneratingExplanation: false,
       isExecuting: false,
-      dislikeQuery: undefined,
+      dislikeQuery: undefined as (boolean | undefined),
       showCallout: false,
       showSamplePrompts: false,
-      queryIterator: undefined,
-      queryResults: undefined,
-      errors: [],
+      queryIterator: undefined as MinimalQueryIterator | undefined,
+      queryResults: undefined as QueryResults | undefined,
+      errors: [] as QueryError[],
       isSamplePromptsOpen: false,
       showDeletePopup: false,
       showFeedbackBar: false,
@@ -235,25 +128,135 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       showInvalidQueryMessageBar: false,
       generatedQueryComments: "",
       wasCopilotUsed: false,
+      showWelcomeSidebar: true,
       showCopilotSidebar: false,
-      chatMessages: [],
+      chatMessages: [] as CopilotMessage[],
       shouldIncludeInMessages: true,
       showExplanationBubble: false,
       notebookServerInfo: {
         notebookServerEndpoint: undefined,
         authToken: undefined,
         forwardingId: undefined,
-      },
+      } as DataModels.NotebookWorkspaceConnectionInfo,
       containerStatus: {
         status: undefined,
         durationLeftInMinutes: undefined,
         phoenixServerInfo: undefined,
-      },
+      } as ContainerInfo,
       schemaAllocationInfo: {
         databaseId: undefined,
         containerId: undefined,
-      },
+      } as CopilotSchemaAllocationInfo,
       isAllocatingContainer: false,
-    }));
-  },
-}));
+      copilotEnabledforExecution: false,
+
+      setCopilotEnabled: (copilotEnabled: boolean) => set({ copilotEnabled }),
+      setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => set({ copilotUserDBEnabled }),
+      setCopilotSampleDBEnabled: (copilotSampleDBEnabled: boolean) => set({ copilotSampleDBEnabled }),
+      openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
+        set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
+      closeFeedbackModal: () => set({ showFeedbackModal: false }),
+      setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) =>
+        set({ hideFeedbackModalForLikedQueries }),
+      refreshCorrelationId: () => set({ correlationId: guid() }),
+      setUserPrompt: (userPrompt: string) => set({ userPrompt }),
+      setQuery: (query: string) => set({ query }),
+      setGeneratedQuery: (generatedQuery: string) => set({ generatedQuery }),
+      setSelectedQuery: (selectedQuery: string) => set({ selectedQuery }),
+      setIsGeneratingQuery: (isGeneratingQuery: boolean) => set({ isGeneratingQuery }),
+      setIsGeneratingExplanation: (isGeneratingExplanation: boolean) => set({ isGeneratingExplanation }),
+      setIsExecuting: (isExecuting: boolean) => set({ isExecuting }),
+      setLikeQuery: (likeQuery: boolean) => set({ likeQuery }),
+      setDislikeQuery: (dislikeQuery: boolean | undefined) => set({ dislikeQuery }),
+      setShowCallout: (showCallout: boolean) => set({ showCallout }),
+      setShowSamplePrompts: (showSamplePrompts: boolean) => set({ showSamplePrompts }),
+      setQueryIterator: (queryIterator: MinimalQueryIterator | undefined) => set({ queryIterator }),
+      setQueryResults: (queryResults: QueryResults | undefined) => set({ queryResults }),
+      setErrors: (errors: QueryError[]) => set({ errors }),
+      setIsSamplePromptsOpen: (isSamplePromptsOpen: boolean) => set({ isSamplePromptsOpen }),
+      setShowDeletePopup: (showDeletePopup: boolean) => set({ showDeletePopup }),
+      setShowFeedbackBar: (showFeedbackBar: boolean) => set({ showFeedbackBar }),
+      setshowCopyPopup: (showCopyPopup: boolean) => set({ showCopyPopup }),
+      setShowErrorMessageBar: (showErrorMessageBar: boolean) => set({ showErrorMessageBar }),
+      setShowInvalidQueryMessageBar: (showInvalidQueryMessageBar: boolean) => set({ showInvalidQueryMessageBar }),
+      setGeneratedQueryComments: (generatedQueryComments: string) => set({ generatedQueryComments }),
+      setWasCopilotUsed: (wasCopilotUsed: boolean) => set({ wasCopilotUsed }),
+      setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
+      setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
+      setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
+      setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
+      setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
+      setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
+        set({ notebookServerInfo }),
+      setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
+      setIsAllocatingContainer: (isAllocatingContainer: boolean) => set({ isAllocatingContainer }),
+      setSchemaAllocationInfo: (schemaAllocationInfo: CopilotSchemaAllocationInfo) => set({ schemaAllocationInfo }),
+      setCopilotEnabledforExecution: (copilotEnabledforExecution: boolean) => set({ copilotEnabledforExecution }),
+
+      resetContainerConnection: (): void => {
+        useTabs.getState().closeAllNotebookTabs(true);
+        useQueryCopilot.getState().setNotebookServerInfo(undefined);
+        useQueryCopilot.getState().setIsAllocatingContainer(false);
+        useQueryCopilot.getState().setContainerStatus({
+          status: undefined,
+          durationLeftInMinutes: undefined,
+          phoenixServerInfo: undefined,
+        });
+        useQueryCopilot.getState().setSchemaAllocationInfo({
+          databaseId: undefined,
+          containerId: undefined,
+        });
+      },
+
+      resetQueryCopilotStates: () => {
+        set((state) => ({
+          ...state,
+          generatedQuery: "",
+          likeQuery: false,
+          userPrompt: "",
+          showFeedbackModal: false,
+          hideFeedbackModalForLikedQueries: false,
+          correlationId: "",
+          query: "SELECT * FROM c",
+          selectedQuery: "",
+          isGeneratingQuery: false,
+          isGeneratingExplanation: false,
+          isExecuting: false,
+          dislikeQuery: undefined,
+          showCallout: false,
+          showSamplePrompts: false,
+          queryIterator: undefined,
+          queryResults: undefined,
+          errors: [],
+          isSamplePromptsOpen: false,
+          showDeletePopup: false,
+          showFeedbackBar: false,
+          showCopyPopup: false,
+          showErrorMessageBar: false,
+          showInvalidQueryMessageBar: false,
+          generatedQueryComments: "",
+          wasCopilotUsed: false,
+          showCopilotSidebar: false,
+          chatMessages: [],
+          shouldIncludeInMessages: true,
+          showExplanationBubble: false,
+          notebookServerInfo: {
+            notebookServerEndpoint: undefined,
+            authToken: undefined,
+            forwardingId: undefined,
+          },
+          containerStatus: {
+            status: undefined,
+            durationLeftInMinutes: undefined,
+            phoenixServerInfo: undefined,
+          },
+          schemaAllocationInfo: {
+            databaseId: undefined,
+            containerId: undefined,
+          },
+          isAllocatingContainer: false,
+        }));
+      },
+    })
+  )
+);

src/hooks/useSidePanel.ts

@@ -1,4 +1,4 @@
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 export interface SidePanelState {
   isOpen: boolean;
@@ -9,7 +9,7 @@ export interface SidePanelState {
   closeSidePanel: () => void;
   getRef?: React.RefObject<HTMLElement>; // Optional ref for focusing the last element.
 }
-export const useSidePanel: UseStore<SidePanelState> = create((set) => ({
+export const useSidePanel = create<SidePanelState>((set) => ({
   isOpen: false,
   panelWidth: "440px",
   openSidePanel: (headerText, panelContent, panelWidth = "440px") =>

src/hooks/useTabs.ts

@@ -7,7 +7,8 @@ import {
   OPEN_TABS_SUBCOMPONENT_NAME,
   saveSubComponentState,
 } from "Shared/AppStatePersistenceUtility";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
+import { subscribeWithSelector } from "zustand/middleware";
 import * as ViewModels from "../Contracts/ViewModels";
 import { CollectionTabKind } from "../Contracts/ViewModels";
 import NotebookTabV2 from "../Explorer/Tabs/NotebookV2Tab";
@@ -51,194 +52,198 @@ export enum ReactTabKind {
   QueryCopilot,
 }
 
-export const useTabs: UseStore<TabsState> = create((set, get) => ({
-  openedTabs: [] as TabsBase[],
-  openedReactTabs: [ReactTabKind.Home],
-  activeTab: undefined as TabsBase,
-  activeReactTab: ReactTabKind.Home,
-  queryCopilotTabInitialInput: "",
-  isTabExecuting: false,
-  isQueryErrorThrown: false,
-  activateTab: (tab: TabsBase): void => {
-    if (get().openedTabs.some((openedTab) => openedTab.tabId === tab.tabId)) {
-      set({ activeTab: tab, activeReactTab: undefined });
-      tab.onActivate();
-    }
-  },
-  activateNewTab: (tab: TabsBase): void => {
-    set((state) => ({ openedTabs: [...state.openedTabs, tab], activeTab: tab, activeReactTab: undefined }));
-    tab.triggerPersistState = get().persistTabsState;
-    tab.onActivate();
-    get().persistTabsState();
-  },
-  activateReactTab: (tabKind: ReactTabKind): void => {
-    // Clear the selected node when switching to a react tab.
-    useSelectedNode.getState().setSelectedNode(undefined);
-    set({ activeTab: undefined, activeReactTab: tabKind });
-  },
-  updateTab: (tab: TabsBase) => {
-    if (get().activeTab?.tabId === tab.tabId) {
-      set({ activeTab: tab });
-    }
-
-    set((state) => ({
-      openedTabs: state.openedTabs.map((openedTab) => {
-        if (openedTab.tabId === tab.tabId) {
-          return tab;
+export const useTabs = create<TabsState>()(
+  subscribeWithSelector(
+    (set, get) => ({
+      openedTabs: [] as TabsBase[],
+      openedReactTabs: [ReactTabKind.Home],
+      activeTab: undefined as TabsBase,
+      activeReactTab: ReactTabKind.Home,
+      queryCopilotTabInitialInput: "",
+      isTabExecuting: false,
+      isQueryErrorThrown: false,
+      activateTab: (tab: TabsBase): void => {
+        if (get().openedTabs.some((openedTab) => openedTab.tabId === tab.tabId)) {
+          set({ activeTab: tab, activeReactTab: undefined });
+          tab.onActivate();
         }
-        return openedTab;
-      }),
-    }));
-  },
-  getTabs: (tabKind: ViewModels.CollectionTabKind, comparator?: (tab: TabsBase) => boolean): TabsBase[] =>
-    get().openedTabs.filter((tab) => tab.tabKind === tabKind && (!comparator || comparator(tab))),
-  refreshActiveTab: (comparator: (tab: TabsBase) => boolean): void => {
-    // ensures that the tab selects/highlights the right node based on resource tree expand/collapse state
-    const activeTab = get().activeTab;
-    activeTab && comparator(activeTab) && activeTab.onActivate();
-  },
-  closeTabsByComparator: (comparator: (tab: TabsBase) => boolean): void =>
-    get()
-      .openedTabs.filter(comparator)
-      .forEach((tab) => tab.onCloseTabButtonClick()),
-  closeTab: (tab: TabsBase): void => {
-    let tabIndex: number;
-    const { activeTab, openedTabs, openedReactTabs } = get();
-    const updatedTabs = openedTabs.filter((openedTab, index) => {
-      if (tab.tabId === openedTab.tabId) {
-        tabIndex = index;
-        return false;
-      }
-      return true;
-    });
-    if (updatedTabs.length === 0 && !isFabricMirrored()) {
-      set({ activeTab: undefined, activeReactTab: undefined });
-    }
-
-    if (tab.tabId === activeTab?.tabId && tabIndex !== -1) {
-      const tabToTheRight = updatedTabs[tabIndex];
-      const lastOpenTab = updatedTabs[updatedTabs.length - 1];
-      const newActiveTab = tabToTheRight ?? lastOpenTab;
-      set({ activeTab: newActiveTab });
-      if (newActiveTab) {
-        newActiveTab.onActivate();
-      }
-    }
-
-    set({ openedTabs: updatedTabs });
-
-    if (updatedTabs.length === 0 && openedReactTabs.length > 0) {
-      set({ activeTab: undefined, activeReactTab: openedReactTabs[openedReactTabs.length - 1] });
-    }
-
-    get().persistTabsState();
-  },
-  closeAllNotebookTabs: (hardClose): void => {
-    const isNotebook = (tabKind: CollectionTabKind): boolean => {
-      if (
-        tabKind === CollectionTabKind.Notebook ||
-        tabKind === CollectionTabKind.NotebookV2 ||
-        tabKind === CollectionTabKind.SchemaAnalyzer ||
-        tabKind === CollectionTabKind.Terminal
-      ) {
-        return true;
-      }
-      return false;
-    };
-
-    const tabList = get().openedTabs;
-    if (tabList && tabList.length > 0) {
-      tabList.forEach((tab: NotebookTabV2) => {
-        const tabKind: CollectionTabKind = tab.tabKind;
-        if (tabKind && isNotebook(tabKind)) {
-          tab.onCloseTabButtonClick(hardClose);
+      },
+      activateNewTab: (tab: TabsBase): void => {
+        set((state) => ({ openedTabs: [...state.openedTabs, tab], activeTab: tab, activeReactTab: undefined }));
+        tab.triggerPersistState = get().persistTabsState;
+        tab.onActivate();
+        get().persistTabsState();
+      },
+      activateReactTab: (tabKind: ReactTabKind): void => {
+        // Clear the selected node when switching to a react tab.
+        useSelectedNode.getState().setSelectedNode(undefined);
+        set({ activeTab: undefined, activeReactTab: tabKind });
+      },
+      updateTab: (tab: TabsBase) => {
+        if (get().activeTab?.tabId === tab.tabId) {
+          set({ activeTab: tab });
         }
-      });
 
-      if (get().openedTabs.length === 0 && !isFabricMirrored()) {
-        set({ activeTab: undefined, activeReactTab: undefined });
-      }
-    }
-  },
-  openAndActivateReactTab: (tabKind: ReactTabKind) => {
-    if (get().openedReactTabs.indexOf(tabKind) === -1) {
-      set((state) => ({
-        openedReactTabs: [...state.openedReactTabs, tabKind],
-      }));
-    }
+        set((state) => ({
+          openedTabs: state.openedTabs.map((openedTab) => {
+            if (openedTab.tabId === tab.tabId) {
+              return tab;
+            }
+            return openedTab;
+          }),
+        }));
+      },
+      getTabs: (tabKind: ViewModels.CollectionTabKind, comparator?: (tab: TabsBase) => boolean): TabsBase[] =>
+        get().openedTabs.filter((tab) => tab.tabKind === tabKind && (!comparator || comparator(tab))),
+      refreshActiveTab: (comparator: (tab: TabsBase) => boolean): void => {
+        // ensures that the tab selects/highlights the right node based on resource tree expand/collapse state
+        const activeTab = get().activeTab;
+        activeTab && comparator(activeTab) && activeTab.onActivate();
+      },
+      closeTabsByComparator: (comparator: (tab: TabsBase) => boolean): void =>
+        get()
+          .openedTabs.filter(comparator)
+          .forEach((tab) => tab.onCloseTabButtonClick()),
+      closeTab: (tab: TabsBase): void => {
+        let tabIndex: number;
+        const { activeTab, openedTabs, openedReactTabs } = get();
+        const updatedTabs = openedTabs.filter((openedTab, index) => {
+          if (tab.tabId === openedTab.tabId) {
+            tabIndex = index;
+            return false;
+          }
+          return true;
+        });
+        if (updatedTabs.length === 0 && !isFabricMirrored()) {
+          set({ activeTab: undefined, activeReactTab: undefined });
+        }
 
-    set({ activeTab: undefined, activeReactTab: tabKind });
-  },
-  closeReactTab: (tabKind: ReactTabKind) => {
-    const { activeReactTab, openedTabs, openedReactTabs } = get();
-    const updatedOpenedReactTabs = openedReactTabs.filter((tab: ReactTabKind) => tabKind !== tab);
-    if (activeReactTab === tabKind) {
-      openedTabs?.length > 0
-        ? set({ activeTab: openedTabs[0], activeReactTab: undefined })
-        : set({ activeTab: undefined, activeReactTab: updatedOpenedReactTabs[0] });
-    }
+        if (tab.tabId === activeTab?.tabId && tabIndex !== -1) {
+          const tabToTheRight = updatedTabs[tabIndex];
+          const lastOpenTab = updatedTabs[updatedTabs.length - 1];
+          const newActiveTab = tabToTheRight ?? lastOpenTab;
+          set({ activeTab: newActiveTab });
+          if (newActiveTab) {
+            newActiveTab.onActivate();
+          }
+        }
 
-    set({ openedReactTabs: updatedOpenedReactTabs });
-  },
-  setQueryCopilotTabInitialInput: (input: string) => set({ queryCopilotTabInitialInput: input }),
-  setIsTabExecuting: (state: boolean) => {
-    set({ isTabExecuting: state });
-  },
-  setIsQueryErrorThrown: (state: boolean) => {
-    set({ isQueryErrorThrown: state });
-  },
-  getCurrentTabIndex: () => {
-    const state = get();
-    if (state.activeReactTab !== undefined) {
-      return state.openedReactTabs.indexOf(state.activeReactTab);
-    } else if (state.activeTab !== undefined) {
-      const nonReactTabIndex = state.openedTabs.indexOf(state.activeTab);
-      if (nonReactTabIndex !== -1) {
-        return state.openedReactTabs.length + nonReactTabIndex;
-      }
-    }
+        set({ openedTabs: updatedTabs });
 
-    return -1;
-  },
-  selectTabByIndex: (index: number) => {
-    const state = get();
-    const totalTabCount = state.openedReactTabs.length + state.openedTabs.length;
-    const clampedIndex = clamp(index, totalTabCount - 1, 0);
+        if (updatedTabs.length === 0 && openedReactTabs.length > 0) {
+          set({ activeTab: undefined, activeReactTab: openedReactTabs[openedReactTabs.length - 1] });
+        }
 
-    if (clampedIndex < state.openedReactTabs.length) {
-      set({ activeTab: undefined, activeReactTab: state.openedReactTabs[clampedIndex] });
-    } else {
-      set({ activeTab: state.openedTabs[clampedIndex - state.openedReactTabs.length], activeReactTab: undefined });
-    }
-  },
-  selectLeftTab: () => {
-    const state = get();
-    state.selectTabByIndex(state.getCurrentTabIndex() - 1);
-  },
-  selectRightTab: () => {
-    const state = get();
-    state.selectTabByIndex(state.getCurrentTabIndex() + 1);
-  },
-  closeActiveTab: () => {
-    const state = get();
-    if (state.activeReactTab !== undefined) {
-      state.closeReactTab(state.activeReactTab);
-    } else if (state.activeTab !== undefined) {
-      state.closeTab(state.activeTab);
-    }
-  },
-  closeAllTabs: () => {
-    set({ openedTabs: [], openedReactTabs: [], activeTab: undefined, activeReactTab: undefined });
-  },
-  persistTabsState: () => {
-    const state = get();
-    const openTabsStates = state.openedTabs.map((tab) => tab.getPersistedState());
+        get().persistTabsState();
+      },
+      closeAllNotebookTabs: (hardClose): void => {
+        const isNotebook = (tabKind: CollectionTabKind): boolean => {
+          if (
+            tabKind === CollectionTabKind.Notebook ||
+            tabKind === CollectionTabKind.NotebookV2 ||
+            tabKind === CollectionTabKind.SchemaAnalyzer ||
+            tabKind === CollectionTabKind.Terminal
+          ) {
+            return true;
+          }
+          return false;
+        };
 
-    saveSubComponentState<OpenTab[]>(
-      AppStateComponentNames.DataExplorerAction,
-      OPEN_TABS_SUBCOMPONENT_NAME,
-      undefined,
-      openTabsStates,
-    );
-  },
-}));
+        const tabList = get().openedTabs;
+        if (tabList && tabList.length > 0) {
+          tabList.forEach((tab: NotebookTabV2) => {
+            const tabKind: CollectionTabKind = tab.tabKind;
+            if (tabKind && isNotebook(tabKind)) {
+              tab.onCloseTabButtonClick(hardClose);
+            }
+          });
+
+          if (get().openedTabs.length === 0 && !isFabricMirrored()) {
+            set({ activeTab: undefined, activeReactTab: undefined });
+          }
+        }
+      },
+      openAndActivateReactTab: (tabKind: ReactTabKind) => {
+        if (get().openedReactTabs.indexOf(tabKind) === -1) {
+          set((state) => ({
+            openedReactTabs: [...state.openedReactTabs, tabKind],
+          }));
+        }
+
+        set({ activeTab: undefined, activeReactTab: tabKind });
+      },
+      closeReactTab: (tabKind: ReactTabKind) => {
+        const { activeReactTab, openedTabs, openedReactTabs } = get();
+        const updatedOpenedReactTabs = openedReactTabs.filter((tab: ReactTabKind) => tabKind !== tab);
+        if (activeReactTab === tabKind) {
+          openedTabs?.length > 0
+            ? set({ activeTab: openedTabs[0], activeReactTab: undefined })
+            : set({ activeTab: undefined, activeReactTab: updatedOpenedReactTabs[0] });
+        }
+
+        set({ openedReactTabs: updatedOpenedReactTabs });
+      },
+      setQueryCopilotTabInitialInput: (input: string) => set({ queryCopilotTabInitialInput: input }),
+      setIsTabExecuting: (state: boolean) => {
+        set({ isTabExecuting: state });
+      },
+      setIsQueryErrorThrown: (state: boolean) => {
+        set({ isQueryErrorThrown: state });
+      },
+      getCurrentTabIndex: () => {
+        const state = get();
+        if (state.activeReactTab !== undefined) {
+          return state.openedReactTabs.indexOf(state.activeReactTab);
+        } else if (state.activeTab !== undefined) {
+          const nonReactTabIndex = state.openedTabs.indexOf(state.activeTab);
+          if (nonReactTabIndex !== -1) {
+            return state.openedReactTabs.length + nonReactTabIndex;
+          }
+        }
+
+        return -1;
+      },
+      selectTabByIndex: (index: number) => {
+        const state = get();
+        const totalTabCount = state.openedReactTabs.length + state.openedTabs.length;
+        const clampedIndex = clamp(index, totalTabCount - 1, 0);
+
+        if (clampedIndex < state.openedReactTabs.length) {
+          set({ activeTab: undefined, activeReactTab: state.openedReactTabs[clampedIndex] });
+        } else {
+          set({ activeTab: state.openedTabs[clampedIndex - state.openedReactTabs.length], activeReactTab: undefined });
+        }
+      },
+      selectLeftTab: () => {
+        const state = get();
+        state.selectTabByIndex(state.getCurrentTabIndex() - 1);
+      },
+      selectRightTab: () => {
+        const state = get();
+        state.selectTabByIndex(state.getCurrentTabIndex() + 1);
+      },
+      closeActiveTab: () => {
+        const state = get();
+        if (state.activeReactTab !== undefined) {
+          state.closeReactTab(state.activeReactTab);
+        } else if (state.activeTab !== undefined) {
+          state.closeTab(state.activeTab);
+        }
+      },
+      closeAllTabs: () => {
+        set({ openedTabs: [], openedReactTabs: [], activeTab: undefined, activeReactTab: undefined });
+      },
+      persistTabsState: () => {
+        const state = get();
+        const openTabsStates = state.openedTabs.map((tab) => tab.getPersistedState());
+
+        saveSubComponentState<OpenTab[]>(
+          AppStateComponentNames.DataExplorerAction,
+          OPEN_TABS_SUBCOMPONENT_NAME,
+          undefined,
+          openTabsStates,
+        );
+      },
+    })
+  )
+);

src/hooks/useTeachingBubble.ts

@@ -1,5 +1,5 @@
 import { Collection } from "Contracts/ViewModels";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 interface TeachingBubbleState {
   step: number;
@@ -12,7 +12,7 @@ interface TeachingBubbleState {
   setSampleCollection: (sampleCollection: Collection) => void;
 }
 
-export const useTeachingBubble: UseStore<TeachingBubbleState> = create((set) => ({
+export const useTeachingBubble = create<TeachingBubbleState>((set) => ({
   step: 1,
   isSampleDBExpanded: false,
   isDocumentsTabOpened: false,

src/hooks/useTerminal.ts

@@ -1,5 +1,5 @@
 import postRobot from "post-robot";
-import create, { UseStore } from "zustand";
+import { create } from "zustand";
 
 interface TerminalState {
   terminalWindow: Window;
@@ -7,7 +7,7 @@ interface TerminalState {
   sendMessage: (message: string) => void;
 }
 
-export const useTerminal: UseStore<TerminalState> = create((set, get) => ({
+export const useTerminal = create<TerminalState>((set, get) => ({
   terminalWindow: undefined,
   setTerminal: (terminalWindow: Window) => {
     set({ terminalWindow });