Fix P2 bgs

* Setting enableEndpointDiscovery to true to test client resilience during region outage.
When this is set to false, it can impact the SDK's failover behavior.

src/Common/Constants.ts

@@ -171,6 +171,7 @@ export class Areas {
   public static Tab: string = "Tab";
   public static ShareDialog: string = "Share Access Dialog";
   public static Notebook: string = "Notebook";
+  public static Copilot: string = "Copilot";
 }
 
 export class HttpHeaders {

src/Common/CosmosClient.ts

@@ -148,9 +148,6 @@ export function client(): Cosmos.CosmosClient {
     endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.masterKey,
     tokenProvider,
-    connectionPolicy: {
-      enableEndpointDiscovery: false,
-    },
     userAgentSuffix: "Azure Portal",
     defaultHeaders: _defaultHeaders,
   };

src/Common/dataAccess/readCollections.ts

@@ -35,14 +35,21 @@ export async function readCollections(databaseId: string): Promise<DataModels.Co
       }
     }
 
-    const responses = await Promise.all(promises);
-    responses.forEach((response) => {
-      collections.push(response.resource as DataModels.Collection);
-    });
+    try {
+      const responses = await Promise.all(promises);
+      responses.forEach((response) => {
+        collections.push(response.resource as DataModels.Collection);
+      });
 
-    // Sort collections by id before returning
-    collections.sort((a, b) => a.id.localeCompare(b.id));
-    return collections;
+      // Sort collections by id before returning
+      collections.sort((a, b) => a.id.localeCompare(b.id));
+      return collections;
+    } catch (error) {
+      handleError(error, "ReadCollections", `Error while querying containers for database ${databaseId}`);
+      throw error;
+    } finally {
+      clearMessage();
+    }
   }
 
   try {

src/Common/dataAccess/readDatabases.ts

@@ -22,6 +22,13 @@ export async function readDatabases(): Promise<DataModels.Database[]> {
     for (const collectionResourceId in tokensData.resourceTokens) {
       // Dictionary key looks like this: dbs/SampleDB/colls/Container
       const resourceIdObj = collectionResourceId.split("/");
+
+      if (resourceIdObj.length !== 4) {
+        handleError(`Resource key not recognized: ${resourceIdObj}`, "ReadDatabases", `Error while querying databases`);
+        clearMessage();
+        return [];
+      }
+
       const databaseId = resourceIdObj[1];
 
       databaseIdsSet.add(databaseId);
@@ -37,6 +44,7 @@ export async function readDatabases(): Promise<DataModels.Database[]> {
         id: databaseId,
         collections: [],
       }));
+    clearMessage();
     return databases;
   }
 

src/ConfigContext.ts

@@ -1,3 +1,4 @@
+import { JunoEndpoints } from "Common/Constants";
 import {
   allowedAadEndpoints,
   allowedArcadiaEndpoints,
@@ -78,7 +79,7 @@ let configContext: Readonly<ConfigContext> = {
   ARCADIA_LIVY_ENDPOINT_DNS_ZONE: "dev.azuresynapse.net",
   GITHUB_CLIENT_ID: "6cb2f63cf6f7b5cbdeca", // Registered OAuth app: https://github.com/organizations/AzureCosmosDBNotebooks/settings/applications/1189306
   GITHUB_TEST_ENV_CLIENT_ID: "b63fc8cbf87fd3c6e2eb", // Registered OAuth app: https://github.com/organizations/AzureCosmosDBNotebooks/settings/applications/1777772
-  JUNO_ENDPOINT: "https://tools.cosmos.azure.com",
+  JUNO_ENDPOINT: JunoEndpoints.Prod,
   BACKEND_ENDPOINT: "https://main.documentdb.ext.azure.com",
   isTerminalEnabled: false,
   isPhoenixEnabled: false,

src/Contracts/AzureResourceGraph.ts

@@ -0,0 +1,13 @@
+export interface QueryRequestOptions {
+  $skipToken?: string;
+  $top?: number;
+  subscriptions: string[];
+}
+
+export interface QueryResponse {
+  $skipToken: string;
+  count: number;
+  data: any;
+  resultTruncated: boolean;
+  totalRecords: number;
+}

src/Contracts/DataModels.ts

@@ -88,13 +88,13 @@ export interface GenerateTokenResponse {
 }
 
 export interface Subscription {
-  uniqueDisplayName: string;
+  uniqueDisplayName?: string;
   displayName: string;
   subscriptionId: string;
-  tenantId: string;
+  tenantId?: string;
   state: string;
-  subscriptionPolicies: SubscriptionPolicies;
-  authorizationSource: string;
+  subscriptionPolicies?: SubscriptionPolicies;
+  authorizationSource?: string;
 }
 
 export interface SubscriptionPolicies {
@@ -457,8 +457,11 @@ export interface ContainerInfo {
 }
 
 export interface IProvisionData {
-  cosmosEndpoint: string;
+  cosmosEndpoint?: string;
   poolId: string;
+  databaseId?: string;
+  containerId?: string;
+  mode?: string;
 }
 
 export interface IContainerData {
@@ -601,3 +604,14 @@ export enum PhoenixErrorType {
   PhoenixFlightFallback = "PhoenixFlightFallback",
   UserMissingPermissionsError = "UserMissingPermissionsError",
 }
+
+export interface CopilotEnabledConfiguration {
+  isEnabled: boolean;
+}
+
+export interface FeatureRegistration {
+  name: string;
+  properties: {
+    state: string;
+  };
+}

src/Explorer/Controls/Editor/EditorReact.tsx

@@ -48,7 +48,7 @@ export class EditorReact extends React.Component<EditorReactProps, EditorReactSt
 
   public componentDidUpdate(previous: EditorReactProps) {
     if (this.props.content !== previous.content) {
-      this.editor.setValue(this.props.content);
+      this.editor?.setValue(this.props.content);
     }
   }
 
@@ -111,7 +111,7 @@ export class EditorReact extends React.Component<EditorReactProps, EditorReactSt
 
     this.rootNode.innerHTML = "";
     const monaco = await loadMonaco();
-    createCallback(monaco.editor.create(this.rootNode, options));
+    createCallback(monaco?.editor?.create(this.rootNode, options));
 
     if (this.rootNode.innerHTML) {
       this.setState({

src/Explorer/Explorer.tsx

@@ -3,6 +3,7 @@ import { isPublicInternetAccessAllowed } from "Common/DatabaseAccountUtility";
 import { sendMessage } from "Common/MessageHandler";
 import { Platform, configContext } from "ConfigContext";
 import { MessageTypes } from "Contracts/ExplorerContracts";
+import { getCopilotEnabled, isCopilotFeatureRegistered } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { IGalleryItem } from "Juno/JunoClient";
 import { requestDatabaseResourceTokens } from "Platform/Fabric/FabricUtil";
 import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
@@ -92,7 +93,7 @@ export default class Explorer {
   };
 
   private static readonly MaxNbDatabasesToAutoExpand = 5;
-  private phoenixClient: PhoenixClient;
+  public phoenixClient: PhoenixClient;
   constructor() {
     const startKey: number = TelemetryProcessor.traceStart(Action.InitializeDataExplorer, {
       dataExplorerArea: Constants.Areas.ResourceTree,
@@ -274,6 +275,7 @@ export default class Explorer {
 
     const NINETY_DAYS_IN_MS = 7776000000;
     const ONE_DAY_IN_MS = 86400000;
+    const THREE_DAYS_IN_MS = 259200000;
     const isAccountNewerThanNinetyDays = isAccountNewerThanThresholdInMs(
       userContext.databaseAccount?.systemData?.createdAt || "",
       NINETY_DAYS_IN_MS,
@@ -293,32 +295,32 @@ export default class Explorer {
       }
     }
 
-    // Try Cosmos DB subscription - survey shown to random 25% of users at day 1 in Data Explorer.
+    // Try Cosmos DB subscription - survey shown to 100% of users at day 1 in Data Explorer.
     if (userContext.isTryCosmosDBSubscription) {
-      if (
-        isAccountNewerThanThresholdInMs(userContext.databaseAccount?.systemData?.createdAt || "", ONE_DAY_IN_MS) &&
-        this.getRandomInt(100) < 25
-      ) {
-        sendMessage({ type: MessageTypes.DisplayNPSSurvey });
-        localStorage.setItem("lastSubmitted", Date.now().toString());
+      if (isAccountNewerThanThresholdInMs(userContext.databaseAccount?.systemData?.createdAt || "", ONE_DAY_IN_MS)) {
+        this.sendNPSMessage();
       }
     } else {
-      // An existing account is lesser than 90 days old. For existing account show to random 10 % of users in Data Explorer.
-      if (isAccountNewerThanNinetyDays) {
-        if (this.getRandomInt(100) < 10) {
-          sendMessage({ type: MessageTypes.DisplayNPSSurvey });
-          localStorage.setItem("lastSubmitted", Date.now().toString());
-        }
+      // An existing account is older than 3 days but less than 90 days old. For existing account show to 100% of users in Data Explorer.
+      if (
+        !isAccountNewerThanThresholdInMs(userContext.databaseAccount?.systemData?.createdAt || "", THREE_DAYS_IN_MS) &&
+        isAccountNewerThanNinetyDays
+      ) {
+        this.sendNPSMessage();
       } else {
-        // An existing account is greater than 90 days. For existing account show to random 25 % of users in Data Explorer.
-        if (this.getRandomInt(100) < 25) {
-          sendMessage({ type: MessageTypes.DisplayNPSSurvey });
-          localStorage.setItem("lastSubmitted", Date.now().toString());
+        // An existing account is greater than 90 days. For existing account show to random 33% of users in Data Explorer.
+        if (this.getRandomInt(100) < 33) {
+          this.sendNPSMessage();
         }
       }
     }
   }
 
+  private sendNPSMessage() {
+    sendMessage({ type: MessageTypes.DisplayNPSSurvey });
+    localStorage.setItem("lastSubmitted", Date.now().toString());
+  }
+
   public async refreshDatabaseForResourceToken(): Promise<void> {
     const databaseId = userContext.parsedResourceToken?.databaseId;
     const collectionId = userContext.parsedResourceToken?.collectionId;
@@ -411,7 +413,7 @@ export default class Explorer {
     this._isInitializingNotebooks = false;
   }
 
-  public async allocateContainer(poolId: PoolIdType): Promise<void> {
+  public async allocateContainer(poolId: PoolIdType, mode?: string): Promise<void> {
     const shouldUseNotebookStates = poolId === PoolIdType.DefaultPoolId ? true : false;
     const notebookServerInfo = shouldUseNotebookStates
       ? useNotebook.getState().notebookServerInfo
@@ -425,10 +427,6 @@ export default class Explorer {
       (notebookServerInfo === undefined ||
         (notebookServerInfo && notebookServerInfo.notebookServerEndpoint === undefined))
     ) {
-      const provisionData: IProvisionData = {
-        cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
-        poolId: shouldUseNotebookStates ? undefined : poolId,
-      };
       const connectionStatus: ContainerConnectionInfo = {
         status: ConnectionStatusType.Connecting,
       };
@@ -436,14 +434,26 @@ export default class Explorer {
       shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus);
 
       let connectionInfo;
+      let provisionData: IProvisionData;
       try {
         TelemetryProcessor.traceStart(Action.PhoenixConnection, {
           dataExplorerArea: Areas.Notebook,
         });
-        shouldUseNotebookStates
-          ? useNotebook.getState().setIsAllocating(true)
-          : useQueryCopilot.getState().setIsAllocatingContainer(true);
-
+        if (shouldUseNotebookStates) {
+          useNotebook.getState().setIsAllocating(true);
+          provisionData = {
+            cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
+            poolId: undefined,
+          };
+        } else {
+          useQueryCopilot.getState().setIsAllocatingContainer(true);
+          provisionData = {
+            poolId: poolId,
+            databaseId: useTabs.getState().activeTab.collection.databaseId,
+            containerId: useTabs.getState().activeTab.collection.id(),
+            mode: mode,
+          };
+        }
         connectionInfo = await this.phoenixClient.allocateContainer(provisionData);
         if (!connectionInfo?.data?.phoenixServiceUrl) {
           throw new Error(`PhoenixServiceUrl is invalid!`);
@@ -459,19 +469,21 @@ export default class Explorer {
           error: getErrorMessage(error),
           errorStack: getErrorStack(error),
         });
-        connectionStatus.status = ConnectionStatusType.Failed;
-        shouldUseNotebookStates
-          ? useNotebook.getState().resetContainerConnection(connectionStatus)
-          : useQueryCopilot.getState().resetContainerConnection();
-        if (error?.status === HttpStatusCodes.Forbidden && error.message) {
-          useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`);
-        } else {
-          useDialog
-            .getState()
-            .showOkModalDialog(
-              "Connection Failed",
-              "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.",
-            );
+        if (shouldUseNotebookStates) {
+          connectionStatus.status = ConnectionStatusType.Failed;
+          shouldUseNotebookStates
+            ? useNotebook.getState().resetContainerConnection(connectionStatus)
+            : useQueryCopilot.getState().resetContainerConnection();
+          if (error?.status === HttpStatusCodes.Forbidden && error.message) {
+            useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`);
+          } else {
+            useDialog
+              .getState()
+              .showOkModalDialog(
+                "Connection Failed",
+                "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.",
+              );
+          }
         }
         throw error;
       } finally {
@@ -485,11 +497,11 @@ export default class Explorer {
     }
   }
 
-  private async setNotebookInfo(
+  public async setNotebookInfo(
     shouldUseNotebookStates: boolean,
     connectionInfo: IResponse<IPhoenixServiceInfo>,
     connectionStatus: DataModels.ContainerConnectionInfo,
-  ) {
+  ): Promise<void> {
     const containerData = {
       forwardingId: connectionInfo.data.forwardingId,
       dbAccountName: userContext.databaseAccount.name,
@@ -510,6 +522,7 @@ export default class Explorer {
     shouldUseNotebookStates
       ? useNotebook.getState().setNotebookServerInfo(noteBookServerInfo)
       : useQueryCopilot.getState().setNotebookServerInfo(noteBookServerInfo);
+
     shouldUseNotebookStates &&
       this.notebookManager?.notebookClient
         .getMemoryUsage()
@@ -1372,6 +1385,21 @@ export default class Explorer {
     await this.refreshSampleData();
   }
 
+  public async configureCopilot(): Promise<void> {
+    if (userContext.apiType !== "SQL" || !userContext.subscriptionId) {
+      return;
+    }
+    const copilotEnabledPromise = getCopilotEnabled();
+    const copilotUserDBEnabledPromise = isCopilotFeatureRegistered(userContext.subscriptionId);
+    const [copilotEnabled, copilotUserDBEnabled] = await Promise.all([
+      copilotEnabledPromise,
+      copilotUserDBEnabledPromise,
+    ]);
+    useQueryCopilot.getState().setCopilotEnabled(copilotEnabled);
+    useQueryCopilot.getState().setCopilotUserDBEnabled(copilotUserDBEnabled);
+    useQueryCopilot.getState().setShowWelcomeModal(window.localStorage.getItem("hideWelcomeModal") !== "true");
+  }
+
   public async refreshSampleData(): Promise<void> {
     if (!userContext.sampleDataConnectionInfo) {
       return;

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.test.ts

@@ -345,7 +345,7 @@ describe("CommandBarComponentButtonFactory tests", () => {
 
   describe("Open Postgres and vCore Mongo buttons", () => {
     const openPostgresShellButtonLabel = "Open PSQL shell";
-    const openVCoreMongoShellButtonLabel = "Open MongoDB (vcore) shell";
+    const openVCoreMongoShellButtonLabel = "Open MongoDB (vCore) shell";
 
     beforeAll(() => {
       mockExplorer = {} as Explorer;

src/Explorer/Menus/CommandBar/CommandBarComponentButtonFactory.tsx

@@ -1,6 +1,3 @@
-import { Action } from "Shared/Telemetry/TelemetryConstants";
-import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
-import { ReactTabKind, useTabs } from "hooks/useTabs";
 import * as React from "react";
 import AddCollectionIcon from "../../../../images/AddCollection.svg";
 import AddDatabaseIcon from "../../../../images/AddDatabase.svg";
@@ -337,13 +334,8 @@ function createNewSQLQueryButton(selectedNodeState: SelectedNodeState): CommandB
       iconSrc: AddSqlQueryIcon,
       iconAlt: label,
       onCommandClick: () => {
-        if (useSelectedNode.getState().isQueryCopilotCollectionSelected()) {
-          useTabs.getState().openAndActivateReactTab(ReactTabKind.QueryCopilot);
-          traceOpen(Action.OpenQueryCopilotFromNewQuery, { apiType: userContext.apiType });
-        } else {
-          const selectedCollection: ViewModels.Collection = selectedNodeState.findSelectedCollection();
-          selectedCollection && selectedCollection.onNewQueryClick(selectedCollection);
-        }
+        const selectedCollection: ViewModels.Collection = selectedNodeState.findSelectedCollection();
+        selectedCollection && selectedCollection.onNewQueryClick(selectedCollection);
       },
       commandButtonLabel: label,
       ariaLabel: label,
@@ -523,7 +515,7 @@ function createOpenTerminalButtonByKind(
       case ViewModels.TerminalKind.Postgres:
         return "PSQL";
       case ViewModels.TerminalKind.VCoreMongo:
-        return "MongoDB (vcore)";
+        return "MongoDB (vCore)";
       default:
         return "";
     }

src/Explorer/Menus/CommandBar/CommandBarUtil.tsx

@@ -6,6 +6,7 @@ import {
   IDropdownOption,
   IDropdownStyles,
 } from "@fluentui/react";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as React from "react";
 import _ from "underscore";
 import ChevronDownIcon from "../../../../images/Chevron_down.svg";
@@ -57,7 +58,11 @@ export const convertButton = (btns: CommandButtonComponentProps[], backgroundCol
         },
         onClick: (ev?: React.MouseEvent<HTMLElement, MouseEvent> | React.KeyboardEvent<HTMLElement>) => {
           btn.onCommandClick(ev);
-          TelemetryProcessor.trace(Action.ClickCommandBarButton, ActionModifiers.Mark, { label });
+          let copilotEnabled = false;
+          if (useQueryCopilot.getState().copilotEnabled && useQueryCopilot.getState().copilotUserDBEnabled) {
+            copilotEnabled = useQueryCopilot.getState().copilotEnabledforExecution;
+          }
+          TelemetryProcessor.trace(Action.ClickCommandBarButton, ActionModifiers.Mark, { label, copilotEnabled });
         },
         key: `${btn.commandButtonLabel}${index}`,
         text: label,

src/Explorer/Panes/AddCollectionPanel.tsx

@@ -1431,8 +1431,6 @@ export class AddCollectionPanel extends React.Component<AddCollectionPanelProps,
       this.setState({ isExecuting: false });
       TelemetryProcessor.traceSuccess(Action.CreateCollection, telemetryData, startKey);
       useSidePanel.getState().closeSidePanel();
-      // open NPS Survey Dialog once the collection is created
-      this.props.explorer.openNPSSurveyDialog();
     } catch (error) {
       const errorMessage: string = getErrorMessage(error);
       this.setState({ isExecuting: false, errorMessage, showErrorDetails: true });

src/Explorer/Panes/__snapshots__/AddCollectionPanel.test.tsx.snap

@@ -336,7 +336,8 @@ exports[`AddCollectionPanel should render Default properly 1`] = `
           directionalHint={4}
         >
           <Icon
-            ariaLabel="Enable analytical store capability to perform near real-time analytics on your operational data, without impacting the performance of transactional workloads."
+            ariaLabel="Enable analytical store capability to perform near real-time analytics on your operational data, without
+ impacting the performance of transactional workloads."
             className="panelInfoIcon"
             iconName="Info"
             tabIndex={0}

src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.test.tsx

@@ -1,10 +1,10 @@
-import { Checkbox, ChoiceGroup, DefaultButton, IconButton, PrimaryButton, TextField } from "@fluentui/react";
+import { Checkbox, DefaultButton, IconButton, PrimaryButton, TextField } from "@fluentui/react";
 import Explorer from "Explorer/Explorer";
 import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal";
+import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext";
 import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { getUserEmail } from "Utils/UserUtils";
 import { shallow } from "enzyme";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
 
 jest.mock("Utils/UserUtils");
@@ -13,21 +13,49 @@ jest.mock("Utils/UserUtils");
 jest.mock("Explorer/QueryCopilot/Shared/QueryCopilotClient");
 SubmitFeedback as jest.Mock;
 
+jest.mock("Explorer/QueryCopilot/QueryCopilotContext");
+const mockUseCopilotStore = useCopilotStore as jest.Mock;
+const mockReturnValue = {
+  generatedQuery: "test query",
+  userPrompt: "test prompt",
+  likeQuery: false,
+  showFeedbackModal: false,
+  closeFeedbackModal: jest.fn,
+  setHideFeedbackModalForLikedQueries: jest.fn,
+};
+
 describe("Query Copilot Feedback Modal snapshot test", () => {
   beforeEach(() => {
+    mockUseCopilotStore.mockReturnValue(mockReturnValue);
     jest.clearAllMocks();
   });
   it("shoud render and match snapshot", () => {
-    useQueryCopilot.getState().openFeedbackModal("test query", false, "test prompt");
-
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    mockUseCopilotStore.mockReturnValue({
+      ...mockReturnValue,
+      showFeedbackModal: true,
+    });
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     expect(wrapper.props().isOpen).toBeTruthy();
     expect(wrapper).toMatchSnapshot();
   });
 
   it("should close on cancel click", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const cancelButton = wrapper.find(IconButton);
     cancelButton.simulate("click");
@@ -38,7 +66,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should get user unput", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
     const testUserInput = "test user input";
 
     const userInput = wrapper.find(TextField).first();
@@ -48,30 +83,15 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
     expect(wrapper).toMatchSnapshot();
   });
 
-  it("should record user contact choice no", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
-    const contactAllowed = wrapper.find(ChoiceGroup);
-
-    contactAllowed.simulate("change", {}, { key: "no" });
-
-    expect(getUserEmail).toHaveBeenCalledTimes(3);
-    expect(wrapper.find(ChoiceGroup).props().selectedKey).toEqual("no");
-    expect(wrapper).toMatchSnapshot();
-  });
-
-  it("should record user contact choice yes", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
-    const contactAllowed = wrapper.find(ChoiceGroup);
-
-    contactAllowed.simulate("change", {}, { key: "yes" });
-
-    expect(getUserEmail).toHaveBeenCalledTimes(4);
-    expect(wrapper.find(ChoiceGroup).props().selectedKey).toEqual("yes");
-    expect(wrapper).toMatchSnapshot();
-  });
-
   it("should not render dont show again button", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const dontShowAgain = wrapper.find(Checkbox);
 
@@ -80,8 +100,19 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should render dont show again button and check it ", () => {
-    useQueryCopilot.getState().openFeedbackModal("test query", true, "test prompt");
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    mockUseCopilotStore.mockReturnValue({
+      ...mockReturnValue,
+      showFeedbackModal: true,
+      likeQuery: true,
+    });
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const dontShowAgain = wrapper.find(Checkbox);
     dontShowAgain.simulate("change", {}, true);
@@ -92,7 +123,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should cancel submission", () => {
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={new Explorer()} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={new Explorer()}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const cancelButton = wrapper.find(DefaultButton);
     cancelButton.simulate("click");
@@ -104,7 +142,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
 
   it("should not submit submission if required description field is null", () => {
     const explorer = new Explorer();
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={explorer} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={explorer}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const submitButton = wrapper.find(PrimaryButton);
     submitButton.simulate("click");
@@ -114,9 +159,15 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
   });
 
   it("should submit submission", () => {
-    useQueryCopilot.getState().openFeedbackModal("test query", false, "test prompt");
     const explorer = new Explorer();
-    const wrapper = shallow(<QueryCopilotFeedbackModal explorer={explorer} />);
+    const wrapper = shallow(
+      <QueryCopilotFeedbackModal
+        explorer={explorer}
+        databaseId="CopilotUserDb"
+        containerId="CopilotUserContainer"
+        mode="User"
+      />,
+    );
 
     const submitButton = wrapper.find("form");
     submitButton.simulate("submit");
@@ -124,12 +175,14 @@ describe("Query Copilot Feedback Modal snapshot test", () => {
 
     expect(SubmitFeedback).toHaveBeenCalledTimes(1);
     expect(SubmitFeedback).toHaveBeenCalledWith({
+      containerId: "CopilotUserContainer",
+      databaseId: "CopilotUserDb",
+      mode: "User",
       params: {
         likeQuery: false,
         generatedQuery: "test query",
         userPrompt: "test prompt",
         description: "",
-        contact: getUserEmail(),
       },
       explorer: explorer,
     });

src/Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal.tsx

@@ -1,6 +1,5 @@
 import {
   Checkbox,
-  ChoiceGroup,
   DefaultButton,
   IconButton,
   Link,
@@ -11,12 +10,21 @@ import {
   TextField,
 } from "@fluentui/react";
 import Explorer from "Explorer/Explorer";
+import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext";
 import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
-import { getUserEmail } from "../../../Utils/UserUtils";
 
-export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }): JSX.Element => {
+export const QueryCopilotFeedbackModal = ({
+  explorer,
+  databaseId,
+  containerId,
+  mode,
+}: {
+  explorer: Explorer;
+  databaseId: string;
+  containerId: string;
+  mode: string;
+}): JSX.Element => {
   const {
     generatedQuery,
     userPrompt,
@@ -24,18 +32,19 @@ export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }):
     showFeedbackModal,
     closeFeedbackModal,
     setHideFeedbackModalForLikedQueries,
-  } = useQueryCopilot();
-  const [isContactAllowed, setIsContactAllowed] = React.useState<boolean>(false);
+  } = useCopilotStore();
   const [description, setDescription] = React.useState<string>("");
   const [doNotShowAgainChecked, setDoNotShowAgainChecked] = React.useState<boolean>(false);
-  const [contact, setContact] = React.useState<string>(getUserEmail());
 
   const handleSubmit = () => {
     closeFeedbackModal();
     setHideFeedbackModalForLikedQueries(doNotShowAgainChecked);
     SubmitFeedback({
-      params: { generatedQuery, likeQuery, description, userPrompt, contact },
-      explorer: explorer,
+      params: { generatedQuery, likeQuery, description, userPrompt },
+      explorer,
+      databaseId,
+      containerId,
+      mode: mode,
     });
   };
 
@@ -64,30 +73,6 @@ export const QueryCopilotFeedbackModal = ({ explorer }: { explorer: Explorer }):
             defaultValue={generatedQuery}
             readOnly
           />
-          <ChoiceGroup
-            styles={{
-              root: {
-                marginBottom: 14,
-              },
-              flexContainer: {
-                selectors: {
-                  ".ms-ChoiceField-field::before": { marginTop: 4 },
-                  ".ms-ChoiceField-field::after": { marginTop: 4 },
-                  ".ms-ChoiceFieldLabel": { paddingLeft: 6 },
-                },
-              },
-            }}
-            label="May we contact you about your feedback?"
-            options={[
-              { key: "yes", text: "Yes, you may contact me." },
-              { key: "no", text: "No, do not contact me." },
-            ]}
-            selectedKey={isContactAllowed ? "yes" : "no"}
-            onChange={(_, option) => {
-              setIsContactAllowed(option.key === "yes");
-              setContact(option.key === "yes" ? getUserEmail() : "");
-            }}
-          ></ChoiceGroup>
           <Text style={{ fontSize: 12, marginBottom: 14 }}>
             By pressing submit, your feedback will be used to improve Microsoft products and services. Please see the{" "}
             {

src/Explorer/QueryCopilot/Modal/WelcomeModal.css

@@ -1,7 +1,6 @@
 .modalContentPadding {
   padding-top: 15px;
   width: 513px;
-  height: 638px;
 }
 
 .exitPadding {
@@ -39,7 +38,7 @@
 }
 
 .buttonPadding {
-  padding: 15px 0px 0px 0px;
+  padding: 15px 0px 15px 0px;
 }
 
 .tryButton {

src/Explorer/QueryCopilot/Modal/WelcomeModal.tsx

@@ -1,7 +1,7 @@
 import { IconButton, Image, Link, Modal, PrimaryButton, Stack, StackItem, Text } from "@fluentui/react";
 import { useBoolean } from "@fluentui/react-hooks";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import React from "react";
-import Database from "../../../../images/CopilotDatabase.svg";
 import Flash from "../../../../images/CopilotFlash.svg";
 import Thumb from "../../../../images/CopilotThumb.svg";
 import CoplilotWelcomeIllustration from "../../../../images/CopliotWelcomeIllustration.svg";
@@ -14,7 +14,11 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element =>
     if (visible) {
       window.localStorage.setItem("hideWelcomeModal", "true");
     }
-  });
+  }, []);
+
+  React.useEffect(() => {
+    useQueryCopilot.getState().setShowWelcomeModal(isModalVisible);
+  }, [isModalVisible]);
 
   return (
     <>
@@ -23,8 +27,10 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element =>
         onDismiss={hideModal}
         isBlocking={false}
         styles={{
-          scrollableContent: {
-            minHeight: 680,
+          main: {
+            maxHeight: 600,
+            borderRadius: 10,
+            overflow: "hidden",
           },
         }}
       >
@@ -52,7 +58,7 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element =>
           </Stack>
           <Stack horizontalAlign="center">
             <Stack.Item align="center" style={{ textAlign: "center" }}>
-              <Text className="title bold">Welcome to Copilot in Azure Cosmos DB (Private Preview)</Text>
+              <Text className="title bold">Welcome to Microsoft Copilot for Azure in Cosmos DB (preview)</Text>
             </Stack.Item>
             <Stack.Item align="center" className="text">
               <Stack horizontal>
@@ -69,7 +75,7 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element =>
               <Text>
                 Ask Copilot to generate a query by describing the query in your words.
                 <br />
-                <Link target="_blank" href="https://aka.ms/cdb-copilot-learn-more">
+                <Link target="_blank" href="https://aka.ms/MicrosoftCopilotForAzureInCDBHowTo">
                   Learn more
                 </Link>
               </Text>
@@ -87,31 +93,11 @@ export const WelcomeModal = ({ visible }: { visible: boolean }): JSX.Element =>
                 </StackItem>
               </Stack>
               <Text>
-                AI-generated content can have mistakes. Make sure it’s accurate and appropriate before using it.
+                AI-generated content can have mistakes. Make sure it is accurate and appropriate before executing the
+                query.
                 <br />
                 <Link target="_blank" href="https://aka.ms/cdb-copilot-preview-terms">
-                  Read preview terms
-                </Link>
-              </Text>
-            </Stack.Item>
-            <Stack.Item align="center" className="text">
-              <Stack horizontal>
-                <StackItem align="start" className="imageTextPadding">
-                  <Image src={Database} />
-                </StackItem>
-                <StackItem align="start">
-                  <Text className="bold">
-                    Query Copilot works on a sample database.
-                    <br />
-                  </Text>
-                </StackItem>
-              </Stack>
-              <Text>
-                While in Private Preview, Query Copilot is setup to work on sample database we have configured for you
-                at no cost.
-                <br />
-                <Link target="_blank" href="https://aka.ms/cdb-copilot-learn-more">
-                  Learn more
+                  Read our preview terms here
                 </Link>
               </Text>
             </Stack.Item>

src/Explorer/QueryCopilot/Modal/__snapshots__/QueryCopilotFeedbackModal.test.tsx.snap

@@ -76,43 +76,6 @@ exports[`Query Copilot Feedback Modal snapshot test shoud render and match snaps
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -235,43 +198,6 @@ exports[`Query Copilot Feedback Modal snapshot test should cancel submission 1`]
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -291,21 +217,6 @@ exports[`Query Copilot Feedback Modal snapshot test should cancel submission 1`]
          
         for more information.
       </Text>
-      <StyledCheckboxBase
-        checked={false}
-        label="Don't show me this next time"
-        onChange={[Function]}
-        styles={
-          Object {
-            "label": Object {
-              "paddingLeft": 0,
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Stack
         horizontal={true}
         horizontalAlign="end"
@@ -409,43 +320,6 @@ exports[`Query Copilot Feedback Modal snapshot test should close on cancel click
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -568,43 +442,6 @@ exports[`Query Copilot Feedback Modal snapshot test should get user unput 1`] =
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -727,361 +564,6 @@ exports[`Query Copilot Feedback Modal snapshot test should not render dont show
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
-      <Text
-        style={
-          Object {
-            "fontSize": 12,
-            "marginBottom": 14,
-          }
-        }
-      >
-        By pressing submit, your feedback will be used to improve Microsoft products and services. Please see the
-         
-        <StyledLinkBase
-          href="https://privacy.microsoft.com/privacystatement"
-          target="_blank"
-        >
-          Privacy statement
-        </StyledLinkBase>
-         
-        for more information.
-      </Text>
-      <Stack
-        horizontal={true}
-        horizontalAlign="end"
-      >
-        <CustomizedPrimaryButton
-          styles={
-            Object {
-              "root": Object {
-                "marginRight": 8,
-              },
-            }
-          }
-          type="submit"
-        >
-          Submit
-        </CustomizedPrimaryButton>
-        <CustomizedDefaultButton
-          onClick={[Function]}
-        >
-          Cancel
-        </CustomizedDefaultButton>
-      </Stack>
-    </Stack>
-  </form>
-</Modal>
-`;
-
-exports[`Query Copilot Feedback Modal snapshot test should record user contact choice no 1`] = `
-<Modal
-  isOpen={false}
->
-  <form
-    onSubmit={[Function]}
-  >
-    <Stack
-      style={
-        Object {
-          "padding": 24,
-        }
-      }
-    >
-      <Stack
-        horizontal={true}
-        horizontalAlign="space-between"
-      >
-        <Text
-          style={
-            Object {
-              "fontSize": 20,
-              "fontWeight": 600,
-              "marginBottom": 20,
-            }
-          }
-        >
-          Send feedback to Microsoft
-        </Text>
-        <CustomizedIconButton
-          iconProps={
-            Object {
-              "iconName": "Cancel",
-            }
-          }
-          onClick={[Function]}
-        />
-      </Stack>
-      <Text
-        style={
-          Object {
-            "fontSize": 14,
-            "marginBottom": 14,
-          }
-        }
-      >
-        Your feedback will help improve the experience.
-      </Text>
-      <StyledTextFieldBase
-        label="Description"
-        multiline={true}
-        onChange={[Function]}
-        placeholder="Provide more details"
-        required={true}
-        rows={3}
-        styles={
-          Object {
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-        value=""
-      />
-      <StyledTextFieldBase
-        defaultValue="test query"
-        label="Query generated"
-        readOnly={true}
-        styles={
-          Object {
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
-      <Text
-        style={
-          Object {
-            "fontSize": 12,
-            "marginBottom": 14,
-          }
-        }
-      >
-        By pressing submit, your feedback will be used to improve Microsoft products and services. Please see the
-         
-        <StyledLinkBase
-          href="https://privacy.microsoft.com/privacystatement"
-          target="_blank"
-        >
-          Privacy statement
-        </StyledLinkBase>
-         
-        for more information.
-      </Text>
-      <Stack
-        horizontal={true}
-        horizontalAlign="end"
-      >
-        <CustomizedPrimaryButton
-          styles={
-            Object {
-              "root": Object {
-                "marginRight": 8,
-              },
-            }
-          }
-          type="submit"
-        >
-          Submit
-        </CustomizedPrimaryButton>
-        <CustomizedDefaultButton
-          onClick={[Function]}
-        >
-          Cancel
-        </CustomizedDefaultButton>
-      </Stack>
-    </Stack>
-  </form>
-</Modal>
-`;
-
-exports[`Query Copilot Feedback Modal snapshot test should record user contact choice yes 1`] = `
-<Modal
-  isOpen={false}
->
-  <form
-    onSubmit={[Function]}
-  >
-    <Stack
-      style={
-        Object {
-          "padding": 24,
-        }
-      }
-    >
-      <Stack
-        horizontal={true}
-        horizontalAlign="space-between"
-      >
-        <Text
-          style={
-            Object {
-              "fontSize": 20,
-              "fontWeight": 600,
-              "marginBottom": 20,
-            }
-          }
-        >
-          Send feedback to Microsoft
-        </Text>
-        <CustomizedIconButton
-          iconProps={
-            Object {
-              "iconName": "Cancel",
-            }
-          }
-          onClick={[Function]}
-        />
-      </Stack>
-      <Text
-        style={
-          Object {
-            "fontSize": 14,
-            "marginBottom": 14,
-          }
-        }
-      >
-        Your feedback will help improve the experience.
-      </Text>
-      <StyledTextFieldBase
-        label="Description"
-        multiline={true}
-        onChange={[Function]}
-        placeholder="Provide more details"
-        required={true}
-        rows={3}
-        styles={
-          Object {
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-        value=""
-      />
-      <StyledTextFieldBase
-        defaultValue="test query"
-        label="Query generated"
-        readOnly={true}
-        styles={
-          Object {
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="yes"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -1204,43 +686,6 @@ exports[`Query Copilot Feedback Modal snapshot test should render dont show agai
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {
@@ -1378,43 +823,6 @@ exports[`Query Copilot Feedback Modal snapshot test should submit submission 1`]
           }
         }
       />
-      <StyledChoiceGroup
-        label="May we contact you about your feedback?"
-        onChange={[Function]}
-        options={
-          Array [
-            Object {
-              "key": "yes",
-              "text": "Yes, you may contact me.",
-            },
-            Object {
-              "key": "no",
-              "text": "No, do not contact me.",
-            },
-          ]
-        }
-        selectedKey="no"
-        styles={
-          Object {
-            "flexContainer": Object {
-              "selectors": Object {
-                ".ms-ChoiceField-field::after": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceField-field::before": Object {
-                  "marginTop": 4,
-                },
-                ".ms-ChoiceFieldLabel": Object {
-                  "paddingLeft": 6,
-                },
-              },
-            },
-            "root": Object {
-              "marginBottom": 14,
-            },
-          }
-        }
-      />
       <Text
         style={
           Object {

src/Explorer/QueryCopilot/Modal/__snapshots__/WelcomeModal.test.tsx.snap

@@ -8,8 +8,10 @@ exports[`Query Copilot Welcome Modal snapshot test should render when isOpen is
     onDismiss={[Function]}
     styles={
       Object {
-        "scrollableContent": Object {
-          "minHeight": 680,
+        "main": Object {
+          "borderRadius": 10,
+          "maxHeight": 600,
+          "overflow": "hidden",
         },
       }
     }
@@ -76,7 +78,7 @@ exports[`Query Copilot Welcome Modal snapshot test should render when isOpen is
           <Text
             className="title bold"
           >
-            Welcome to Copilot in Azure Cosmos DB (Private Preview)
+            Welcome to Microsoft Copilot for Azure in Cosmos DB (preview)
           </Text>
         </StackItem>
         <StackItem
@@ -109,7 +111,7 @@ exports[`Query Copilot Welcome Modal snapshot test should render when isOpen is
             Ask Copilot to generate a query by describing the query in your words.
             <br />
             <StyledLinkBase
-              href="https://aka.ms/cdb-copilot-learn-more"
+              href="https://aka.ms/MicrosoftCopilotForAzureInCDBHowTo"
               target="_blank"
             >
               Learn more
@@ -143,50 +145,13 @@ exports[`Query Copilot Welcome Modal snapshot test should render when isOpen is
             </StackItem>
           </Stack>
           <Text>
-            AI-generated content can have mistakes. Make sure it’s accurate and appropriate before using it.
+            AI-generated content can have mistakes. Make sure it is accurate and appropriate before executing the query.
             <br />
             <StyledLinkBase
               href="https://aka.ms/cdb-copilot-preview-terms"
               target="_blank"
             >
-              Read preview terms
-            </StyledLinkBase>
-          </Text>
-        </StackItem>
-        <StackItem
-          align="center"
-          className="text"
-        >
-          <Stack
-            horizontal={true}
-          >
-            <StackItem
-              align="start"
-              className="imageTextPadding"
-            >
-              <Image
-                src={Object {}}
-              />
-            </StackItem>
-            <StackItem
-              align="start"
-            >
-              <Text
-                className="bold"
-              >
-                Query Copilot works on a sample database.
-                <br />
-              </Text>
-            </StackItem>
-          </Stack>
-          <Text>
-            While in Private Preview, Query Copilot is setup to work on sample database we have configured for you at no cost.
-            <br />
-            <StyledLinkBase
-              href="https://aka.ms/cdb-copilot-learn-more"
-              target="_blank"
-            >
-              Learn more
+              Read our preview terms here
             </StyledLinkBase>
           </Text>
         </StackItem>

src/Explorer/QueryCopilot/QueryCopilotContext.tsx

@@ -0,0 +1,135 @@
+import { MinimalQueryIterator } from "Common/IteratorUtilities";
+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";
+const context = createContext(null);
+const useCopilotStore = (): Partial<QueryCopilotState> => useContext(context);
+
+const CopilotProvider = ({ children }: { children: React.ReactNode }): JSX.Element => {
+  const [useStore] = useState(() =>
+    create((set, get) => ({
+      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,
+      errorMessage: "",
+      isSamplePromptsOpen: false,
+      showDeletePopup: false,
+      showFeedbackBar: false,
+      showCopyPopup: false,
+      showErrorMessageBar: false,
+      showInvalidQueryMessageBar: false,
+      generatedQueryComments: "",
+      wasCopilotUsed: false,
+      showWelcomeSidebar: true,
+      showCopilotSidebar: false,
+      chatMessages: [],
+      shouldIncludeInMessages: true,
+      showExplanationBubble: false,
+      isAllocatingContainer: false,
+
+      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 }),
+      setErrorMessage: (errorMessage: string) => set({ errorMessage }),
+      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 }),
+
+      getState: () => {
+        return get();
+      },
+
+      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,
+          errorMessage: "",
+          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,
+          },
+          isAllocatingContainer: false,
+        }));
+      },
+    })),
+  );
+  return <context.Provider value={useStore()}>{children}</context.Provider>;
+};
+
+export { CopilotProvider, useCopilotStore };

src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx

@@ -1,3 +1,5 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/* eslint-disable no-console */
 import {
   Callout,
   CommandBarButton,
@@ -17,50 +19,61 @@ import {
   TextField,
 } from "@fluentui/react";
 import { useBoolean } from "@fluentui/react-hooks";
-import {
-  ContainerStatusType,
-  PoolIdType,
-  QueryCopilotSampleContainerSchema,
-  ShortenedQueryCopilotSampleContainerSchema,
-} from "Common/Constants";
+import { HttpStatusCodes } from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
 import { WelcomeModal } from "Explorer/QueryCopilot/Modal/WelcomeModal";
 import { CopyPopup } from "Explorer/QueryCopilot/Popup/CopyPopup";
 import { DeletePopup } from "Explorer/QueryCopilot/Popup/DeletePopup";
-import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
+import {
+  SuggestedPrompt,
+  getSampleDatabaseSuggestedPrompts,
+  getSuggestedPrompts,
+} from "Explorer/QueryCopilot/QueryCopilotUtilities";
+import { SubmitFeedback, allocatePhoenixContainer } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { GenerateSQLQueryResponse, QueryCopilotProps } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { SamplePrompts, SamplePromptsProps } from "Explorer/QueryCopilot/Shared/SamplePrompts/SamplePrompts";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { userContext } from "UserContext";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
+import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
 import React, { useRef, useState } from "react";
 import HintIcon from "../../../images/Hint.svg";
 import CopilotIcon from "../../../images/QueryCopilotNewLogo.svg";
 import RecentIcon from "../../../images/Recent.svg";
 import errorIcon from "../../../images/close-black.svg";
+import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { useTabs } from "../../hooks/useTabs";
+import { useCopilotStore } from "../QueryCopilot/QueryCopilotContext";
+import { useSelectedNode } from "../useSelectedNode";
 
 type QueryCopilotPromptProps = QueryCopilotProps & {
+  databaseId: string;
+  containerId: string;
   toggleCopilot: (toggle: boolean) => void;
 };
 
-interface SuggestedPrompt {
-  id: number;
-  text: string;
-}
-
 const promptStyles: IButtonStyles = {
   root: { border: 0, selectors: { ":hover": { outline: "1px dashed #605e5c" } } },
-  label: { fontWeight: 400, textAlign: "left", paddingLeft: 8 },
+  label: {
+    fontWeight: 400,
+    textAlign: "left",
+    paddingLeft: 8,
+    overflow: "hidden",
+    whiteSpace: "nowrap",
+    textOverflow: "ellipsis",
+  },
+  textContainer: { overflow: "hidden" },
 };
 
 export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
   explorer,
   toggleCopilot,
+  databaseId,
+  containerId,
 }: QueryCopilotPromptProps): JSX.Element => {
   const [copilotTeachingBubbleVisible, { toggle: toggleCopilotTeachingBubbleVisible }] = useBoolean(false);
-  const inputEdited = useRef(false);
   const {
+    openFeedbackModal,
     hideFeedbackModalForLikedQueries,
     userPrompt,
     setUserPrompt,
@@ -93,7 +106,10 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
     setGeneratedQueryComments,
     setQueryResults,
     setErrorMessage,
-  } = useQueryCopilot();
+    errorMessage,
+  } = useCopilotStore();
+
+  const inputEdited = useRef(!!userPrompt);
 
   const sampleProps: SamplePromptsProps = {
     isSamplePromptsOpen: isSamplePromptsOpen,
@@ -118,14 +134,13 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
     }, 6000);
   };
 
+  const isSampleCopilotActive = useSelectedNode.getState().isQueryCopilotCollectionSelected();
   const cachedHistoriesString = localStorage.getItem(`${userContext.databaseAccount?.id}-queryCopilotHistories`);
   const cachedHistories = cachedHistoriesString?.split("|");
   const [histories, setHistories] = useState<string[]>(cachedHistories || []);
-  const suggestedPrompts: SuggestedPrompt[] = [
-    { id: 1, text: 'Show all products that have the word "ultra" in the name or description' },
-    { id: 2, text: "What are all of the possible categories for the products, and their counts?" },
-    { id: 3, text: 'Show me all products that have been reviewed by someone with a username that contains "bob"' },
-  ];
+  const suggestedPrompts: SuggestedPrompt[] = isSampleCopilotActive
+    ? getSampleDatabaseSuggestedPrompts()
+    : getSuggestedPrompts();
   const [filteredHistories, setFilteredHistories] = useState<string[]>(histories);
   const [filteredSuggestedPrompts, setFilteredSuggestedPrompts] = useState<SuggestedPrompt[]>(suggestedPrompts);
 
@@ -176,28 +191,24 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
       setShowDeletePopup(false);
       useTabs.getState().setIsTabExecuting(true);
       useTabs.getState().setIsQueryErrorThrown(false);
-      if (
-        useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
-        !userContext.features.disableCopilotPhoenixGateaway
-      ) {
-        await explorer.allocateContainer(PoolIdType.QueryCopilot);
-      }
+      const mode: string = isSampleCopilotActive ? "Sample" : "User";
+
+      await allocatePhoenixContainer({ explorer, databaseId, containerId, mode });
+
       const payload = {
-        containerSchema: userContext.features.enableCopilotFullSchema
-          ? QueryCopilotSampleContainerSchema
-          : ShortenedQueryCopilotSampleContainerSchema,
         userPrompt: userPrompt,
       };
       useQueryCopilot.getState().refreshCorrelationId();
       const serverInfo = useQueryCopilot.getState().notebookServerInfo;
       const queryUri = userContext.features.disableCopilotPhoenixGateaway
         ? createUri("https://copilotorchestrater.azurewebsites.net/", "generateSQLQuery")
-        : createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery");
+        : createUri(serverInfo.notebookServerEndpoint, "public/generateSQLQuery");
       const response = await fetch(queryUri, {
         method: "POST",
         headers: {
           "content-type": "application/json",
           "x-ms-correlationid": useQueryCopilot.getState().correlationId,
+          Authorization: `token ${useQueryCopilot.getState().notebookServerInfo.authToken}`,
         },
         body: JSON.stringify(payload),
       });
@@ -215,13 +226,40 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
           setGeneratedQueryComments(generateSQLQueryResponse.explanation);
           setShowFeedbackBar(true);
           resetQueryResults();
+          TelemetryProcessor.traceSuccess(Action.QueryGenerationFromCopilotPrompt, {
+            databaseName: databaseId,
+            collectionId: containerId,
+            copilotLatency:
+              Date.parse(generateSQLQueryResponse?.generateEnd) - Date.parse(generateSQLQueryResponse?.generateStart),
+            responseCode: response.status,
+          });
         } else {
           setShowInvalidQueryMessageBar(true);
+          TelemetryProcessor.traceFailure(Action.QueryGenerationFromCopilotPrompt, {
+            databaseName: databaseId,
+            collectionId: containerId,
+            responseCode: response.status,
+          });
         }
+      } else if (response?.status === HttpStatusCodes.TooManyRequests) {
+        handleError(JSON.stringify(generateSQLQueryResponse), "copilotTooManyRequestError");
+        useTabs.getState().setIsQueryErrorThrown(true);
+        setShowErrorMessageBar(true);
+        setErrorMessage("Ratelimit exceeded 5 per 1 minute. Please try again after sometime");
+        TelemetryProcessor.traceFailure(Action.QueryGenerationFromCopilotPrompt, {
+          databaseName: databaseId,
+          collectionId: containerId,
+          responseCode: response.status,
+        });
       } else {
         handleError(JSON.stringify(generateSQLQueryResponse), "copilotInternalServerError");
         useTabs.getState().setIsQueryErrorThrown(true);
         setShowErrorMessageBar(true);
+        TelemetryProcessor.traceFailure(Action.QueryGenerationFromCopilotPrompt, {
+          databaseName: databaseId,
+          collectionId: containerId,
+          responseCode: response.status,
+        });
       }
     } catch (error) {
       handleError(error, "executeNaturalLanguageQuery");
@@ -237,7 +275,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
   const showTeachingBubble = (): void => {
     if (!inputEdited.current) {
       setTimeout(() => {
-        if (!inputEdited.current && !isWelcomModalVisible()) {
+        if (!useQueryCopilot.getState().showWelcomeModal && !userPrompt && !inputEdited.current) {
           toggleCopilotTeachingBubbleVisible();
           inputEdited.current = true;
         }
@@ -245,10 +283,6 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
     }
   };
 
-  const isWelcomModalVisible = (): boolean => {
-    return localStorage.getItem("hideWelcomeModal") !== "true";
-  };
-
   const clearFeedback = () => {
     resetButtonState();
     resetQueryResults();
@@ -268,6 +302,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
 
   React.useEffect(() => {
     showTeachingBubble();
+    useQueryCopilot.subscribe(showTeachingBubble, (state: QueryCopilotState) => state.showWelcomeModal);
     useTabs.getState().setIsQueryErrorThrown(false);
   }, []);
 
@@ -310,6 +345,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
           styles={{ root: { width: "95%" }, fieldGroup: { borderRadius: 6 } }}
           disabled={isGeneratingQuery}
           autoComplete="off"
+          placeholder="Ask a question in natural language and we’ll generate the query for you."
         />
         {copilotTeachingBubbleVisible && (
           <TeachingBubble
@@ -343,7 +379,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
         {isGeneratingQuery && <Spinner style={{ marginLeft: 8 }} />}
         {showSamplePrompts && (
           <Callout
-            styles={{ root: { minWidth: 400 } }}
+            styles={{ root: { minWidth: 400, maxWidth: "70vw" } }}
             target="#naturalLanguageInput"
             isBeakVisible={false}
             onDismiss={() => setShowSamplePrompts(false)}
@@ -375,7 +411,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
                         setShowSamplePrompts(false);
                         inputEdited.current = true;
                       }}
-                      onRenderIcon={() => <Image src={RecentIcon} />}
+                      onRenderIcon={() => <Image src={RecentIcon} styles={{ root: { overflow: "unset" } }} />}
                       styles={promptStyles}
                     >
                       {history}
@@ -451,7 +487,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
           </Link>
           {showErrorMessageBar && (
             <MessageBar messageBarType={MessageBarType.error}>
-              We ran into an error and were not able to execute query.
+              {errorMessage ? errorMessage : "We ran into an error and were not able to execute query."}
             </MessageBar>
           )}
           {showInvalidQueryMessageBar && (
@@ -489,7 +525,10 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
                     description: "",
                     userPrompt: userPrompt,
                   },
-                  explorer: explorer,
+                  explorer,
+                  databaseId,
+                  containerId,
+                  mode: isSampleCopilotActive ? "Sample" : "User",
                 });
               }}
               directionalHint={DirectionalHint.topCenter}
@@ -499,7 +538,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
                 <Link
                   onClick={() => {
                     setShowCallout(false);
-                    useQueryCopilot.getState().openFeedbackModal(generatedQuery, true, userPrompt);
+                    openFeedbackModal(generatedQuery, true, userPrompt);
                   }}
                 >
                   more feedback?
@@ -524,7 +563,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
             iconProps={{ iconName: dislikeQuery === true ? "DislikeSolid" : "Dislike" }}
             onClick={() => {
               if (!dislikeQuery) {
-                useQueryCopilot.getState().openFeedbackModal(generatedQuery, false, userPrompt);
+                openFeedbackModal(generatedQuery, false, userPrompt);
                 setLikeQuery(false);
               }
               setDislikeQuery(!dislikeQuery);
@@ -550,7 +589,7 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
           </CommandBarButton>
         </Stack>
       )}
-      <WelcomeModal visible={isWelcomModalVisible()} />
+      <WelcomeModal visible={useQueryCopilot.getState().showWelcomeModal} />
       {isSamplePromptsOpen && <SamplePrompts sampleProps={sampleProps} />}
       {query !== "" && query.trim().length !== 0 && (
         <DeletePopup

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -1,5 +1,6 @@
 /* eslint-disable no-console */
 import { Stack } from "@fluentui/react";
+import { QueryCopilotSampleContainerId, QueryCopilotSampleDatabaseId } from "Common/Constants";
 import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
@@ -11,6 +12,7 @@ import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotRe
 import { userContext } from "UserContext";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useSidePanel } from "hooks/useSidePanel";
+import { ReactTabKind, TabsState, useTabs } from "hooks/useTabs";
 import React, { useState } from "react";
 import SplitterLayout from "react-splitter-layout";
 import QueryCommandIcon from "../../../images/CopilotCommand.svg";
@@ -28,13 +30,14 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     ? StringUtility.toBoolean(cachedCopilotToggleStatus)
     : true;
   const [copilotActive, setCopilotActive] = useState<boolean>(copilotInitialActive);
+  const [tabActive, setTabActive] = useState<boolean>(true);
 
   const getCommandbarButtons = (): CommandButtonComponentProps[] => {
     const executeQueryBtnLabel = selectedQuery ? "Execute Selection" : "Execute Query";
     const executeQueryBtn = {
       iconSrc: ExecuteQueryIcon,
       iconAlt: executeQueryBtnLabel,
-      onCommandClick: () => OnExecuteQueryClick(),
+      onCommandClick: () => OnExecuteQueryClick(useQueryCopilot),
       commandButtonLabel: executeQueryBtnLabel,
       ariaLabel: executeQueryBtnLabel,
       hasPopup: false,
@@ -73,10 +76,13 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
 
   React.useEffect(() => {
     return () => {
-      const commandbarButtons = getCommandbarButtons();
-      commandbarButtons.pop();
-      commandbarButtons.map((props: CommandButtonComponentProps) => (props.disabled = true));
-      useCommandBar.getState().setContextButtons(commandbarButtons);
+      useTabs.subscribe((state: TabsState) => {
+        if (state.activeReactTab === ReactTabKind.QueryCopilot) {
+          setTabActive(true);
+        } else {
+          setTabActive(false);
+        }
+      });
     };
   }, []);
 
@@ -88,8 +94,13 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
   return (
     <Stack className="tab-pane" style={{ width: "100%" }}>
       <div style={isGeneratingQuery ? { height: "100%" } : { overflowY: "auto", height: "100%" }}>
-        {copilotActive && (
-          <QueryCopilotPromptbar explorer={explorer} toggleCopilot={toggleCopilot}></QueryCopilotPromptbar>
+        {tabActive && copilotActive && (
+          <QueryCopilotPromptbar
+            explorer={explorer}
+            toggleCopilot={toggleCopilot}
+            databaseId={QueryCopilotSampleDatabaseId}
+            containerId={QueryCopilotSampleContainerId}
+          ></QueryCopilotPromptbar>
         )}
         <Stack className="tabPaneContentContainer">
           <SplitterLayout percentage={true} vertical={true} primaryIndex={0} primaryMinSize={30} secondaryMinSize={70}>

src/Explorer/QueryCopilot/QueryCopilotUtilities.ts

@@ -7,6 +7,11 @@ import { getCommonQueryOptions } from "Common/dataAccess/queryDocuments";
 import DocumentId from "Explorer/Tree/DocumentId";
 import { logConsoleProgress } from "Utils/NotificationConsoleUtils";
 
+export interface SuggestedPrompt {
+  id: number;
+  text: string;
+}
+
 export const querySampleDocuments = (query: string, options: FeedOptions): QueryIterator<ItemDefinition & Resource> => {
   options = getCommonQueryOptions(options);
   return sampleDataClient()
@@ -33,3 +38,19 @@ export const readSampleDocument = async (documentId: DocumentId): Promise<Item>
     clearMessage();
   }
 };
+
+export const getSampleDatabaseSuggestedPrompts = (): SuggestedPrompt[] => {
+  return [
+    { id: 1, text: 'Show all products that have the word "ultra" in the name or description' },
+    { id: 2, text: "What are all of the possible categories for the products, and their counts?" },
+    { id: 3, text: 'Show me all products that have been reviewed by someone with a username that contains "bob"' },
+  ];
+};
+
+export const getSuggestedPrompts = (): SuggestedPrompt[] => {
+  return [
+    { id: 1, text: "Show the first 10 items" },
+    { id: 2, text: 'Count all the items in my data as "numItems"' },
+    { id: 3, text: "Find the oldest item added to my collection" },
+  ];
+};

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.test.ts

@@ -1,4 +1,3 @@
-import { QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainerSchema } from "Common/Constants";
 import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
 import Explorer from "Explorer/Explorer";
@@ -37,9 +36,6 @@ describe("Query Copilot Client", () => {
       userPrompt: "UserPrompt",
       description: "Description",
       contact: "Contact",
-      containerSchema: userContext.features.enableCopilotFullSchema
-        ? QueryCopilotSampleContainerSchema
-        : ShortenedQueryCopilotSampleContainerSchema,
     };
 
     const mockStore = useQueryCopilot.getState();
@@ -52,13 +48,16 @@ describe("Query Copilot Client", () => {
 
     const feedbackUri = userContext.features.disableCopilotPhoenixGateaway
       ? createUri("https://copilotorchestrater.azurewebsites.net/", "feedback")
-      : createUri(useQueryCopilot.getState().notebookServerInfo.notebookServerEndpoint, "feedback");
+      : createUri(useQueryCopilot.getState().notebookServerInfo.notebookServerEndpoint, "public/feedback");
 
     it("should call fetch with the payload with like", async () => {
       const mockFetch = jest.fn().mockResolvedValueOnce({});
 
       globalThis.fetch = mockFetch;
       await SubmitFeedback({
+        databaseId: "test",
+        containerId: "test",
+        mode: "User",
         params: {
           likeQuery: true,
           generatedQuery: "GeneratedQuery",
@@ -91,6 +90,9 @@ describe("Query Copilot Client", () => {
       globalThis.fetch = mockFetch;
 
       await SubmitFeedback({
+        databaseId: "test",
+        containerId: "test",
+        mode: "User",
         params: {
           likeQuery: false,
           generatedQuery: "GeneratedQuery",
@@ -108,6 +110,7 @@ describe("Query Copilot Client", () => {
           headers: {
             "content-type": "application/json",
             "x-ms-correlationid": "mocked-correlation-id",
+            Authorization: "token mocked-token",
           },
         }),
       );
@@ -120,6 +123,9 @@ describe("Query Copilot Client", () => {
       globalThis.fetch = jest.fn().mockRejectedValueOnce(new Error("Mock error"));
 
       await SubmitFeedback({
+        databaseId: "test",
+        containerId: "test",
+        mode: "User",
         params: {
           likeQuery: true,
           generatedQuery: "GeneratedQuery",

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -1,28 +1,192 @@
 import { FeedOptions } from "@azure/cosmos";
 import {
+  Areas,
+  ConnectionStatusType,
   ContainerStatusType,
+  HttpStatusCodes,
   PoolIdType,
   QueryCopilotSampleContainerId,
   QueryCopilotSampleContainerSchema,
   ShortenedQueryCopilotSampleContainerSchema,
 } from "Common/Constants";
-import { getErrorMessage, handleError } from "Common/ErrorHandlingUtils";
+import { getErrorMessage, getErrorStack, handleError } from "Common/ErrorHandlingUtils";
 import { shouldEnableCrossPartitionKey } from "Common/HeadersUtility";
 import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import { createUri } from "Common/UrlUtility";
 import { queryDocumentsPage } from "Common/dataAccess/queryDocumentsPage";
-import { QueryResults } from "Contracts/ViewModels";
+import { configContext } from "ConfigContext";
+import {
+  ContainerConnectionInfo,
+  CopilotEnabledConfiguration,
+  FeatureRegistration,
+  IProvisionData,
+} from "Contracts/DataModels";
+import { AuthorizationTokenHeaderMetadata, QueryResults } from "Contracts/ViewModels";
+import { useDialog } from "Explorer/Controls/Dialog";
 import Explorer from "Explorer/Explorer";
 import { querySampleDocuments } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import { FeedbackParams, GenerateSQLQueryResponse } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceFailure, traceStart, traceSuccess } from "Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "UserContext";
+import { getAuthorizationHeader } from "Utils/AuthorizationUtils";
 import { queryPagesUntilContentPresent } from "Utils/QueryUtils";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
+import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
 import { useTabs } from "hooks/useTabs";
 import * as StringUtility from "../../../Shared/StringUtility";
 
+async function fetchWithTimeout(
+  url: string,
+  headers: {
+    [x: string]: string;
+  },
+) {
+  const timeout = 10000;
+  const options = { timeout };
+
+  const controller = new AbortController();
+  const id = setTimeout(() => controller.abort(), timeout);
+
+  const response = await window.fetch(url, {
+    headers,
+    ...options,
+    signal: controller.signal,
+  });
+  clearTimeout(id);
+
+  return response;
+}
+
+export const isCopilotFeatureRegistered = async (subscriptionId: string): Promise<boolean> => {
+  const api_version = "2021-07-01";
+  const url = `${configContext.ARM_ENDPOINT}/subscriptions/${subscriptionId}/providers/Microsoft.Features/featureProviders/Microsoft.DocumentDB/subscriptionFeatureRegistrations/MicrosoftCopilotForAzureInCDB?api-version=${api_version}`;
+  const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader();
+  const headers = { [authorizationHeader.header]: authorizationHeader.token };
+
+  let response;
+
+  try {
+    response = await fetchWithTimeout(url, headers);
+  } catch (error) {
+    return false;
+  }
+
+  if (!response?.ok) {
+    return false;
+  }
+
+  const featureRegistration = (await response?.json()) as FeatureRegistration;
+  return featureRegistration?.properties?.state === "Registered";
+};
+
+export const getCopilotEnabled = async (): Promise<boolean> => {
+  const url = `${configContext.BACKEND_ENDPOINT}/api/portalsettings/querycopilot`;
+  const authorizationHeader: AuthorizationTokenHeaderMetadata = getAuthorizationHeader();
+  const headers = { [authorizationHeader.header]: authorizationHeader.token };
+
+  let response;
+
+  try {
+    response = await fetchWithTimeout(url, headers);
+  } catch (error) {
+    return false;
+  }
+
+  if (!response?.ok) {
+    return false;
+  }
+
+  const copilotPortalConfiguration = (await response?.json()) as CopilotEnabledConfiguration;
+  return copilotPortalConfiguration?.isEnabled;
+};
+
+export const allocatePhoenixContainer = async ({
+  explorer,
+  databaseId,
+  containerId,
+  mode,
+}: {
+  explorer: Explorer;
+  databaseId: string;
+  containerId: string;
+  mode: string;
+}): Promise<void> => {
+  try {
+    if (
+      useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
+      !userContext.features.disableCopilotPhoenixGateaway
+    ) {
+      await explorer.allocateContainer(PoolIdType.QueryCopilot, mode);
+    } else {
+      const currentAllocatedSchemaInfo = useQueryCopilot.getState().schemaAllocationInfo;
+      if (
+        currentAllocatedSchemaInfo.databaseId !== databaseId ||
+        currentAllocatedSchemaInfo.containerId !== containerId
+      ) {
+        await resetPhoenixContainerSchema({ explorer, databaseId, containerId, mode });
+      }
+    }
+    useQueryCopilot.getState().setSchemaAllocationInfo({
+      databaseId,
+      containerId,
+    });
+  } catch (error) {
+    traceFailure(Action.PhoenixConnection, {
+      dataExplorerArea: Areas.Copilot,
+      status: error.status,
+      error: getErrorMessage(error),
+      errorStack: getErrorStack(error),
+    });
+    useQueryCopilot.getState().resetContainerConnection();
+    if (error?.status === HttpStatusCodes.Forbidden && error.message) {
+      useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`);
+    } else {
+      useDialog
+        .getState()
+        .showOkModalDialog(
+          "Connection Failed",
+          "We are unable to connect to the temporary workspace. Please try again in a few minutes. If the error persists, file a support ticket.",
+        );
+    }
+  } finally {
+    useTabs.getState().setIsTabExecuting(false);
+  }
+};
+
+export const resetPhoenixContainerSchema = async ({
+  explorer,
+  databaseId,
+  containerId,
+  mode,
+}: {
+  explorer: Explorer;
+  databaseId: string;
+  containerId: string;
+  mode: string;
+}): Promise<void> => {
+  try {
+    const provisionData: IProvisionData = {
+      poolId: PoolIdType.QueryCopilot,
+      databaseId: databaseId,
+      containerId: containerId,
+      mode: mode,
+    };
+    const connectionInfo = await explorer.phoenixClient.allocateContainer(provisionData);
+    const connectionStatus: ContainerConnectionInfo = {
+      status: ConnectionStatusType.Connecting,
+    };
+    await explorer.setNotebookInfo(false, connectionInfo, connectionStatus);
+  } catch (error) {
+    traceFailure(Action.PhoenixConnection, {
+      dataExplorerArea: Areas.Copilot,
+      status: error.status,
+      error: getErrorMessage(error),
+      errorStack: getErrorStack(error),
+    });
+    throw error;
+  }
+};
+
 export const SendQueryRequest = async ({
   userPrompt,
   explorer,
@@ -51,7 +215,7 @@ export const SendQueryRequest = async ({
 
       const queryUri = userContext.features.disableCopilotPhoenixGateaway
         ? createUri("https://copilotorchestrater.azurewebsites.net/", "generateSQLQuery")
-        : createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery");
+        : createUri(serverInfo.notebookServerEndpoint, "public/generateSQLQuery");
 
       const payload = {
         containerSchema: userContext.features.enableCopilotFullSchema
@@ -106,16 +270,19 @@ export const SendQueryRequest = async ({
 export const SubmitFeedback = async ({
   params,
   explorer,
+  databaseId,
+  containerId,
+  mode,
 }: {
   params: FeedbackParams;
   explorer: Explorer;
+  databaseId: string;
+  containerId: string;
+  mode: string;
 }): Promise<void> => {
   try {
     const { likeQuery, generatedQuery, userPrompt, description, contact } = params;
     const payload = {
-      containerSchema: userContext.features.enableCopilotFullSchema
-        ? QueryCopilotSampleContainerSchema
-        : ShortenedQueryCopilotSampleContainerSchema,
       like: likeQuery ? "like" : "dislike",
       generatedSql: generatedQuery,
       userPrompt,
@@ -126,17 +293,18 @@ export const SubmitFeedback = async ({
       useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
       !userContext.features.disableCopilotPhoenixGateaway
     ) {
-      await explorer.allocateContainer(PoolIdType.QueryCopilot);
+      await allocatePhoenixContainer({ explorer, databaseId, containerId, mode });
     }
     const serverInfo = useQueryCopilot.getState().notebookServerInfo;
     const feedbackUri = userContext.features.disableCopilotPhoenixGateaway
       ? createUri("https://copilotorchestrater.azurewebsites.net/", "feedback")
-      : createUri(serverInfo.notebookServerEndpoint, "feedback");
+      : createUri(serverInfo.notebookServerEndpoint, "public/feedback");
     await fetch(feedbackUri, {
       method: "POST",
       headers: {
         "content-type": "application/json",
         "x-ms-correlationid": useQueryCopilot.getState().correlationId,
+        Authorization: `token ${useQueryCopilot.getState().notebookServerInfo.authToken}`,
       },
       body: JSON.stringify(payload),
     });
@@ -145,7 +313,7 @@ export const SubmitFeedback = async ({
   }
 };
 
-export const OnExecuteQueryClick = async (): Promise<void> => {
+export const OnExecuteQueryClick = async (useQueryCopilot: Partial<QueryCopilotState>): Promise<void> => {
   traceStart(Action.ExecuteQueryGeneratedFromQueryCopilot, {
     correlationId: useQueryCopilot.getState().correlationId,
     userPrompt: useQueryCopilot.getState().userPrompt,
@@ -160,13 +328,14 @@ export const OnExecuteQueryClick = async (): Promise<void> => {
   useQueryCopilot.getState().setQueryIterator(queryIterator);
 
   setTimeout(async () => {
-    await QueryDocumentsPerPage(0, queryIterator);
+    await QueryDocumentsPerPage(0, queryIterator, useQueryCopilot);
   }, 100);
 };
 
 export const QueryDocumentsPerPage = async (
   firstItemIndex: number,
   queryIterator: MinimalQueryIterator,
+  useQueryCopilot: Partial<QueryCopilotState>,
 ): Promise<void> => {
   try {
     useQueryCopilot.getState().setIsExecuting(true);

src/Explorer/QueryCopilot/Shared/QueryCopilotInterfaces.ts

@@ -32,3 +32,8 @@ export interface FeedbackParams {
 export interface QueryCopilotProps {
   explorer: Explorer;
 }
+
+export interface CopilotSchemaAllocationInfo {
+  databaseId: string;
+  containerId: string;
+}

src/Explorer/QueryCopilot/Shared/QueryCopilotResults.tsx

@@ -12,7 +12,7 @@ export const QueryCopilotResults: React.FC = (): JSX.Element => {
       queryResults={useQueryCopilot.getState().queryResults}
       isExecuting={useQueryCopilot.getState().isExecuting}
       executeQueryDocumentsPage={(firstItemIndex: number) =>
-        QueryDocumentsPerPage(firstItemIndex, useQueryCopilot.getState().queryIterator)
+        QueryDocumentsPerPage(firstItemIndex, useQueryCopilot.getState().queryIterator, useQueryCopilot)
       }
     />
   );

src/Explorer/QueryCopilot/__snapshots__/QueryCopilotTab.test.tsx.snap

@@ -18,6 +18,8 @@ exports[`Query copilot tab snapshot test should render with initial input 1`] =
     }
   >
     <QueryCopilotPromptbar
+      containerId="SampleContainer"
+      databaseId="CopilotSampleDb"
       explorer={
         Explorer {
           "_isInitializingNotebooks": false,

src/Explorer/Quickstart/VCoreMongoQuickstartCommands.ts

@@ -1,5 +1,6 @@
 export const newDbAndCollectionCommand = `use quickstartDB
-db.createCollection('sampleCollection')`;
+db.createCollection('sampleCollection')
+`;
 
 export const newDbAndCollectionCommandForDisplay = `use quickstartDB // Create new database named 'quickstartDB' or switch to it if it already exists
 
@@ -16,19 +17,25 @@ export const loadDataCommand = `db.sampleCollection.insertMany([
     {title: "War and Peace", author: "Leo Tolstoy", pages: 1392},
     {title: "The Odyssey", author: "Homer", pages: 374},
     {title: "Ulysses", author: "James Joyce", pages: 730}
-  ])`;
+  ])
+`;
 
-export const queriesCommand = `db.sampleCollection.find({author: "George Orwell"})
+export const findOrwellCommand = `db.sampleCollection.find({author: "George Orwell"})
+`;
 
+export const findOrwellCommandForDisplay = `// Query to find all books written by "George Orwell"
+db.sampleCollection.find({author: "George Orwell"})`;
+
+export const findByPagesCommand = `db.sampleCollection.find({pages: {$gt: 500}})
+`;
+
+export const findByPagesCommandForDisplay = `// Query to find all books with more than 500 pages
 db.sampleCollection.find({pages: {$gt: 500}})
+`;
 
-db.sampleCollection.find({}).sort({pages: 1})`;
+export const findAndSortCommand = `db.sampleCollection.find({}).sort({pages: 1})
+`;
 
-export const queriesCommandForDisplay = `// Query to find all books written by "George Orwell"
-db.sampleCollection.find({author: "George Orwell"})
-
-// Query to find all books with more than 500 pages
-db.sampleCollection.find({pages: {$gt: 500}})
-
-// Query to find all books and sort them by the number of pages in ascending order
-db.sampleCollection.find({}).sort({pages: 1})`;
+export const findAndSortCommandForDisplay = `// Query to find all books and sort them by the number of pages in ascending order
+db.sampleCollection.find({}).sort({pages: 1})
+`;

src/Explorer/Quickstart/VCoreMongoQuickstartGuide.tsx

@@ -11,11 +11,15 @@ import {
 } from "@fluentui/react";
 import { customPivotHeaderRenderer } from "Explorer/Quickstart/Shared/QuickstartRenderUtilities";
 import {
+  findAndSortCommand,
+  findAndSortCommandForDisplay,
+  findByPagesCommand,
+  findByPagesCommandForDisplay,
+  findOrwellCommand,
+  findOrwellCommandForDisplay,
   loadDataCommand,
   newDbAndCollectionCommand,
   newDbAndCollectionCommandForDisplay,
-  queriesCommand,
-  queriesCommandForDisplay,
 } from "Explorer/Quickstart/VCoreMongoQuickstartCommands";
 import { useTerminal } from "hooks/useTerminal";
 import React, { useState } from "react";
@@ -190,17 +194,17 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                 </Text>
                 <DefaultButton
                   style={{ marginTop: 16, width: 110 }}
-                  onClick={() => useTerminal.getState().sendMessage(queriesCommand)}
+                  onClick={() => useTerminal.getState().sendMessage(findOrwellCommand)}
                 >
                   Try query
                 </DefaultButton>
                 <Stack horizontal style={{ marginTop: 16 }}>
                   <TextField
-                    id="queriesCommand"
+                    id="findOrwellCommand"
                     multiline
-                    rows={5}
+                    rows={2}
                     readOnly
-                    defaultValue={queriesCommandForDisplay}
+                    defaultValue={findOrwellCommandForDisplay}
                     styles={{
                       root: { width: "90%" },
                       field: {
@@ -214,7 +218,65 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                     iconProps={{
                       iconName: "Copy",
                     }}
-                    onClick={() => onCopyBtnClicked("#queriesCommand")}
+                    onClick={() => onCopyBtnClicked("#findOrwellCommand")}
+                  />
+                </Stack>
+                <DefaultButton
+                  style={{ marginTop: 32, width: 110 }}
+                  onClick={() => useTerminal.getState().sendMessage(findByPagesCommand)}
+                >
+                  Try query
+                </DefaultButton>
+                <Stack horizontal style={{ marginTop: 16 }}>
+                  <TextField
+                    id="findByPagesCommand"
+                    multiline
+                    rows={2}
+                    readOnly
+                    defaultValue={findByPagesCommandForDisplay}
+                    styles={{
+                      root: { width: "90%" },
+                      field: {
+                        backgroundColor: "#EEEEEE",
+                        fontFamily:
+                          "Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New",
+                      },
+                    }}
+                  />
+                  <IconButton
+                    iconProps={{
+                      iconName: "Copy",
+                    }}
+                    onClick={() => onCopyBtnClicked("#findByPagesCommand")}
+                  />
+                </Stack>
+                <DefaultButton
+                  style={{ marginTop: 32, width: 110 }}
+                  onClick={() => useTerminal.getState().sendMessage(findAndSortCommand)}
+                >
+                  Try query
+                </DefaultButton>
+                <Stack horizontal style={{ marginTop: 16 }}>
+                  <TextField
+                    id="findAndSortCommand"
+                    multiline
+                    rows={2}
+                    readOnly
+                    defaultValue={findAndSortCommandForDisplay}
+                    styles={{
+                      root: { width: "90%" },
+                      field: {
+                        backgroundColor: "#EEEEEE",
+                        fontFamily:
+                          "Consolas,Monaco,Lucida Console,Liberation Mono,DejaVu Sans Mono,Bitstream Vera Sans Mono,Courier New",
+                      },
+                    }}
+                  />
+                  <IconButton
+                    iconProps={{
+                      iconName: "Copy",
+                    }}
+                    onClick={() => onCopyBtnClicked("#findAndSortCommand")}
                   />
                 </Stack>
               </Stack>
@@ -236,7 +298,7 @@ export const VcoreMongoQuickstartGuide: React.FC = (): JSX.Element => {
                   hosted in the cloud, to Azure Cosmos DB for MongoDB vCore.&nbsp;
                   <Link
                     target="_blank"
-                    href="https://learn.microsoft.com/azure-data-studio/extensions/azure-cosmos-db-mongodb-extension"
+                    href="https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/migration-options"
                   >
                     Learn more
                   </Link>

src/Explorer/SplashScreen/SplashScreen.tsx

@@ -19,6 +19,7 @@ import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
 import { useCarousel } from "hooks/useCarousel";
 import { usePostgres } from "hooks/usePostgres";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
 import * as React from "react";
 import ConnectIcon from "../../../images/Connect_color.svg";
@@ -104,6 +105,12 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
           (state) => state.sampleDataResourceTokenCollection,
         ),
       },
+      {
+        dispose: useQueryCopilot.subscribe(
+          () => this.setState({}),
+          (state) => state.copilotEnabled,
+        ),
+      },
     );
   }
 
@@ -114,9 +121,9 @@ export class SplashScreen extends React.Component<SplashScreenProps> {
 
   private getSplashScreenButtons = (): JSX.Element => {
     if (
-      useDatabases.getState().sampleDataResourceTokenCollection &&
-      userContext.features.enableCopilot &&
-      userContext.apiType === "SQL"
+      userContext.apiType === "SQL" &&
+      useQueryCopilot.getState().copilotEnabled &&
+      useDatabases.getState().sampleDataResourceTokenCollection
     ) {
       return (
         <Stack style={{ width: "66%", cursor: "pointer", margin: "40px auto" }} tokens={{ childrenGap: 16 }}>

src/Explorer/Tabs/DocumentsTab.ts

@@ -1,10 +1,10 @@
 import { ItemDefinition, PartitionKey, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import { querySampleDocuments, readSampleDocument } from "Explorer/QueryCopilot/QueryCopilotUtilities";
+import { QueryConstants } from "Shared/Constants";
+import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import * as ko from "knockout";
 import Q from "q";
 import { format } from "react-string-format";
-import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import DeleteDocumentIcon from "../../../images/DeleteDocument.svg";
 import NewDocumentIcon from "../../../images/NewDocument.svg";
 import UploadIcon from "../../../images/Upload_16x16.svg";
@@ -331,6 +331,7 @@ export default class DocumentsTab extends TabsBase {
 
     this.showPartitionKey = this._shouldShowPartitionKey();
     this._isQueryCopilotSampleContainer =
+      this.collection?.isSampleCollection &&
       this.collection?.databaseId === QueryCopilotSampleDatabaseId &&
       this.collection?.id() === QueryCopilotSampleContainerId;
   }

src/Explorer/Tabs/QueryTab/QueryTab.tsx

@@ -1,11 +1,16 @@
+import { CopilotProvider } from "Explorer/QueryCopilot/QueryCopilotContext";
+import { userContext } from "UserContext";
 import React from "react";
 import * as DataModels from "../../../Contracts/DataModels";
 import type { QueryTabOptions } from "../../../Contracts/ViewModels";
 import { useTabs } from "../../../hooks/useTabs";
 import Explorer from "../../Explorer";
-import { IQueryTabComponentProps, ITabAccessor } from "../../Tabs/QueryTab/QueryTabComponent";
+import QueryTabComponent, {
+  IQueryTabComponentProps,
+  ITabAccessor,
+  QueryTabFunctionComponent,
+} from "../../Tabs/QueryTab/QueryTabComponent";
 import TabsBase from "../TabsBase";
-import QueryTabComponent from "./QueryTabComponent";
 
 export interface IQueryTabProps {
   container: Explorer;
@@ -40,7 +45,13 @@ export class NewQueryTab extends TabsBase {
   }
 
   public render(): JSX.Element {
-    return <QueryTabComponent {...this.iQueryTabComponentProps} />;
+    return userContext.apiType === "SQL" ? (
+      <CopilotProvider>
+        <QueryTabFunctionComponent {...this.iQueryTabComponentProps} />
+      </CopilotProvider>
+    ) : (
+      <QueryTabComponent {...this.iQueryTabComponentProps} />
+    );
   }
 
   public onTabClick(): void {

src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx

@@ -1,6 +1,16 @@
 import { fireEvent, render } from "@testing-library/react";
-import QueryTabComponent, { IQueryTabComponentProps } from "Explorer/Tabs/QueryTab/QueryTabComponent";
+import { CollectionTabKind } from "Contracts/ViewModels";
+import { CopilotProvider } from "Explorer/QueryCopilot/QueryCopilotContext";
+import { QueryCopilotPromptbar } from "Explorer/QueryCopilot/QueryCopilotPromptbar";
+import QueryTabComponent, {
+  IQueryTabComponentProps,
+  QueryTabFunctionComponent,
+} from "Explorer/Tabs/QueryTab/QueryTabComponent";
+import TabsBase from "Explorer/Tabs/TabsBase";
+import { updateUserContext, userContext } from "UserContext";
+import { mount } from "enzyme";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
+import { useTabs } from "hooks/useTabs";
 import React from "react";
 
 jest.mock("Explorer/Controls/Editor/EditorReact");
@@ -11,9 +21,15 @@ describe("QueryTabComponent", () => {
     mockStore.showCopilotSidebar = false;
     mockStore.setShowCopilotSidebar = jest.fn();
   });
-  beforeEach(() => jest.clearAllMocks());
+  afterEach(() => jest.clearAllMocks());
 
-  it("should launch Copilot when ALT+C is pressed", () => {
+  it("should launch conversational Copilot when ALT+C is pressed and when copilot version is 3", () => {
+    updateUserContext({
+      features: {
+        ...userContext.features,
+        copilotVersion: "v3.0",
+      },
+    });
     const propsMock: Readonly<IQueryTabComponentProps> = {
       collection: { databaseId: "CopilotSampleDb" },
       onTabAccessor: () => jest.fn(),
@@ -31,4 +47,32 @@ describe("QueryTabComponent", () => {
 
     expect(mockStore.setShowCopilotSidebar).toHaveBeenCalledWith(true);
   });
+
+  it("copilot should be enabled by default when tab is active", () => {
+    useQueryCopilot.getState().setCopilotEnabled(true);
+    useQueryCopilot.getState().setCopilotUserDBEnabled(true);
+    const activeTab = new TabsBase({
+      tabKind: CollectionTabKind.Query,
+      title: "Query",
+      tabPath: "",
+    });
+    activeTab.tabId = "mockTabId";
+    useTabs.getState().activeTab = activeTab;
+    const propsMock: Readonly<IQueryTabComponentProps> = {
+      collection: { databaseId: "CopilotUserDb", id: () => "CopilotUserContainer" },
+      onTabAccessor: () => jest.fn(),
+      isExecutionError: false,
+      tabId: "mockTabId",
+      tabsBaseInstance: {
+        updateNavbarWithTabsButtons: () => jest.fn(),
+      },
+    } as unknown as IQueryTabComponentProps;
+
+    const container = mount(
+      <CopilotProvider>
+        <QueryTabFunctionComponent {...propsMock} />
+      </CopilotProvider>,
+    );
+    expect(container.find(QueryCopilotPromptbar).exists()).toBe(true);
+  });
 });

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -1,21 +1,29 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+/* eslint-disable no-console */
 import { FeedOptions } from "@azure/cosmos";
 import { useDialog } from "Explorer/Controls/Dialog";
-import { OnExecuteQueryClick } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
-import { QueryCopilotResults } from "Explorer/QueryCopilot/Shared/QueryCopilotResults";
+import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal";
+import { useCopilotStore } from "Explorer/QueryCopilot/QueryCopilotContext";
+import { QueryCopilotPromptbar } from "Explorer/QueryCopilot/QueryCopilotPromptbar";
+import { OnExecuteQueryClick, QueryDocumentsPerPage } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { QueryCopilotSidebar } from "Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
+import { useSelectedNode } from "Explorer/useSelectedNode";
 import { QueryConstants } from "Shared/Constants";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
+import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
+import { TabsState, useTabs } from "hooks/useTabs";
 import React, { Fragment } from "react";
 import SplitterLayout from "react-splitter-layout";
 import "react-splitter-layout/lib/index.css";
 import { format } from "react-string-format";
+import QueryCommandIcon from "../../../../images/CopilotCommand.svg";
 import LaunchCopilot from "../../../../images/CopilotTabIcon.svg";
 import CancelQueryIcon from "../../../../images/Entity_cancel.svg";
 import ExecuteQueryIcon from "../../../../images/ExecuteQuery.svg";
 import SaveQueryIcon from "../../../../images/save-cosmos.svg";
-import { NormalizedEventKey, QueryCopilotSampleDatabaseId } from "../../../Common/Constants";
+import { NormalizedEventKey } from "../../../Common/Constants";
 import { getErrorMessage } from "../../../Common/ErrorHandlingUtils";
 import * as HeadersUtility from "../../../Common/HeadersUtility";
 import { MinimalQueryIterator } from "../../../Common/IteratorUtilities";
@@ -24,6 +32,8 @@ import { queryDocuments } from "../../../Common/dataAccess/queryDocuments";
 import { queryDocumentsPage } from "../../../Common/dataAccess/queryDocumentsPage";
 import * as DataModels from "../../../Contracts/DataModels";
 import * as ViewModels from "../../../Contracts/ViewModels";
+import * as StringUtility from "../../../Shared/StringUtility";
+import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../../UserContext";
 import * as QueryUtils from "../../../Utils/QueryUtils";
 import { useSidePanel } from "../../../hooks/useSidePanel";
@@ -72,6 +82,9 @@ export interface IQueryTabComponentProps {
   isPreferredApiMongoDB?: boolean;
   monacoEditorSetting?: string;
   viewModelcollection?: ViewModels.Collection;
+  copilotEnabled?: boolean;
+  isSampleCopilotActive?: boolean;
+  copilotStore?: Partial<QueryCopilotState>;
 }
 
 interface IQueryTabStates {
@@ -85,8 +98,25 @@ interface IQueryTabStates {
   showCopilotSidebar: boolean;
   queryCopilotGeneratedQuery: string;
   cancelQueryTimeoutID: NodeJS.Timeout;
+  copilotActive: boolean;
+  currentTabActive: boolean;
 }
 
+export const QueryTabFunctionComponent = (props: IQueryTabComponentProps): any => {
+  const copilotStore = useCopilotStore();
+  const copilotGlobalStore = useQueryCopilot();
+  const isSampleCopilotActive = useSelectedNode.getState().isQueryCopilotCollectionSelected();
+  const queryTabProps = {
+    ...props,
+    copilotEnabled:
+      copilotGlobalStore.copilotEnabled &&
+      (copilotGlobalStore.copilotUserDBEnabled || (isSampleCopilotActive && !!userContext.sampleDataConnectionInfo)),
+    isSampleCopilotActive: isSampleCopilotActive,
+    copilotStore: copilotStore,
+  };
+  return <QueryTabComponent {...queryTabProps}></QueryTabComponent>;
+};
+
 export default class QueryTabComponent extends React.Component<IQueryTabComponentProps, IQueryTabStates> {
   public queryEditorId: string;
   public executeQueryButton: Button;
@@ -113,12 +143,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       showCopilotSidebar: useQueryCopilot.getState().showCopilotSidebar,
       queryCopilotGeneratedQuery: useQueryCopilot.getState().query,
       cancelQueryTimeoutID: undefined,
+      copilotActive: this._queryCopilotActive(),
+      currentTabActive: true,
     };
     this.isCloseClicked = false;
     this.splitterId = this.props.tabId + "_splitter";
     this.queryEditorId = `queryeditor${this.props.tabId}`;
     this.isPreferredApiMongoDB = this.props.isPreferredApiMongoDB;
-    this.isCopilotTabActive = QueryCopilotSampleDatabaseId === this.props.collection.databaseId;
+    this.isCopilotTabActive = userContext.features.copilotVersion === "v3.0";
     this.executeQueryButton = {
       enabled: !!this.state.sqlQueryEditorContent && this.state.sqlQueryEditorContent.length > 0,
       visible: true,
@@ -143,6 +175,19 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     });
   }
 
+  private _queryCopilotActive(): boolean {
+    if (this.props.copilotEnabled) {
+      const cachedCopilotToggleStatus: string = localStorage.getItem(
+        `${userContext.databaseAccount?.id}-queryCopilotToggleStatus`,
+      );
+      const copilotInitialActive: boolean = cachedCopilotToggleStatus
+        ? StringUtility.toBoolean(cachedCopilotToggleStatus)
+        : true;
+      return copilotInitialActive;
+    }
+    return false;
+  }
+
   public onCloseClick(isClicked: boolean): void {
     this.isCloseClicked = isClicked;
     if (useQueryCopilot.getState().wasCopilotUsed && this.isCopilotTabActive) {
@@ -167,6 +212,16 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     setTimeout(async () => {
       await this._executeQueryDocumentsPage(0);
     }, 100);
+    if (this.state.copilotActive) {
+      const query = this.state.sqlQueryEditorContent.split("\r\n")?.pop();
+      const isqueryEdited = this.props.copilotStore.generatedQuery && this.props.copilotStore.generatedQuery !== query;
+      if (isqueryEdited) {
+        TelemetryProcessor.traceMark(Action.QueryEdited, {
+          databaseName: this.props.collection.databaseId,
+          collectionId: this.props.collection.id(),
+        });
+      }
+    }
   };
 
   public onSaveQueryClick = (): void => {
@@ -326,7 +381,9 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       buttons.push({
         iconSrc: ExecuteQueryIcon,
         iconAlt: label,
-        onCommandClick: this.isCopilotTabActive ? () => OnExecuteQueryClick() : this.onExecuteQueryClick,
+        onCommandClick: this.props.isSampleCopilotActive
+          ? () => OnExecuteQueryClick(this.props.copilotStore)
+          : this.onExecuteQueryClick,
         commandButtonLabel: label,
         ariaLabel: label,
         hasPopup: false,
@@ -380,6 +437,20 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       buttons.push(launchCopilotButton);
     }
 
+    if (this.props.copilotEnabled) {
+      const toggleCopilotButton = {
+        iconSrc: QueryCommandIcon,
+        iconAlt: "Copilot",
+        onCommandClick: () => {
+          this._toggleCopilot(!this.state.copilotActive);
+        },
+        commandButtonLabel: this.state.copilotActive ? "Disable Copilot" : "Enable Copilot",
+        ariaLabel: "Copilot",
+        hasPopup: false,
+      };
+      buttons.push(toggleCopilotButton);
+    }
+
     if (!this.props.isPreferredApiMongoDB && this.state.isExecuting) {
       const label = "Cancel query";
       buttons.push({
@@ -395,11 +466,31 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     return buttons;
   }
 
+  private _toggleCopilot = (active: boolean) => {
+    this.setState({ copilotActive: active });
+    useQueryCopilot.getState().setCopilotEnabledforExecution(active);
+    localStorage.setItem(`${userContext.databaseAccount?.id}-queryCopilotToggleStatus`, active.toString());
+
+    TelemetryProcessor.traceSuccess(active ? Action.ActivateQueryCopilot : Action.DeactivateQueryCopilot, {
+      databaseName: this.props.collection.databaseId,
+      collectionId: this.props.collection.id(),
+    });
+  };
+
+  componentDidUpdate = (_prevProps: IQueryTabComponentProps, prevState: IQueryTabStates): void => {
+    if (prevState.copilotActive !== this.state.copilotActive) {
+      useCommandBar.getState().setContextButtons(this.getTabsButtons());
+    }
+  };
+
   public onChangeContent(newContent: string): void {
     this.setState({
       sqlQueryEditorContent: newContent,
       queryCopilotGeneratedQuery: "",
     });
+    if (this.state.copilotActive) {
+      this.props.copilotStore?.setQuery(newContent);
+    }
     if (this.isPreferredApiMongoDB) {
       if (newContent.length > 0) {
         this.executeQueryButton = {
@@ -434,6 +525,10 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
         : useQueryCopilot.getState().setSelectedQuery("");
     }
 
+    if (this.state.copilotActive) {
+      this.props.copilotStore?.setSelectedQuery(selectedContent);
+    }
+
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
   }
 
@@ -442,6 +537,10 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       return this.state.queryCopilotGeneratedQuery;
     }
 
+    if (this.state.copilotActive) {
+      return this.props.copilotStore?.query;
+    }
+
     return this.state.sqlQueryEditorContent;
   }
 
@@ -452,12 +551,15 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
   private unsubscribeCopilotSidebar: () => void;
 
   componentDidMount(): void {
-    this.unsubscribeCopilotSidebar = useQueryCopilot.subscribe((state: QueryCopilotState) => {
-      if (this.state.showCopilotSidebar !== state.showCopilotSidebar) {
-        this.setState({ showCopilotSidebar: state.showCopilotSidebar });
-      }
-      if (this.state.queryCopilotGeneratedQuery !== state.query) {
-        this.setState({ queryCopilotGeneratedQuery: state.query });
+    useTabs.subscribe((state: TabsState) => {
+      if (this.state.currentTabActive && state.activeTab?.tabId !== this.props.tabId) {
+        this.setState({
+          currentTabActive: false,
+        });
+      } else if (!this.state.currentTabActive && state.activeTab?.tabId === this.props.tabId) {
+        this.setState({
+          currentTabActive: true,
+        });
       }
     });
 
@@ -466,7 +568,6 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
   }
 
   componentWillUnmount(): void {
-    this.unsubscribeCopilotSidebar();
     document.removeEventListener("keydown", this.handleCopilotKeyDown);
   }
 
@@ -474,6 +575,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     return (
       <Fragment>
         <div className="tab-pane" id={this.props.tabId} role="tabpanel">
+          {this.props.copilotEnabled && this.state.currentTabActive && this.state.copilotActive && (
+            <QueryCopilotPromptbar
+              explorer={this.props.collection.container}
+              toggleCopilot={this._toggleCopilot}
+              databaseId={this.props.collection.databaseId}
+              containerId={this.props.collection.id()}
+            ></QueryCopilotPromptbar>
+          )}
           <div className="tabPaneContentContainer">
             <SplitterLayout vertical={true} primaryIndex={0} primaryMinSize={100} secondaryMinSize={200}>
               <Fragment>
@@ -482,6 +591,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
                     language={"sql"}
                     content={this.setEditorContent()}
                     isReadOnly={false}
+                    wordWrap={"on"}
                     ariaLabel={"Editing Query"}
                     lineNumbers={"on"}
                     onContentChanged={(newContent: string) => this.onChangeContent(newContent)}
@@ -489,8 +599,21 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
                   />
                 </div>
               </Fragment>
-              {this.isCopilotTabActive ? (
-                <QueryCopilotResults />
+              {this.props.isSampleCopilotActive ? (
+                <QueryResultSection
+                  isMongoDB={this.props.isPreferredApiMongoDB}
+                  queryEditorContent={this.state.sqlQueryEditorContent}
+                  error={this.props.copilotStore?.errorMessage}
+                  queryResults={this.props.copilotStore?.queryResults}
+                  isExecuting={this.props.copilotStore?.isExecuting}
+                  executeQueryDocumentsPage={(firstItemIndex: number) =>
+                    QueryDocumentsPerPage(
+                      firstItemIndex,
+                      this.props.copilotStore.queryIterator,
+                      this.props.copilotStore,
+                    )
+                  }
+                />
               ) : (
                 <QueryResultSection
                   isMongoDB={this.props.isPreferredApiMongoDB}
@@ -506,6 +629,14 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
             </SplitterLayout>
           </div>
         </div>
+        {this.props.copilotEnabled && this.props.copilotStore?.showFeedbackModal && (
+          <QueryCopilotFeedbackModal
+            explorer={this.props.collection.container}
+            databaseId={this.props.collection.databaseId}
+            containerId={this.props.collection.id()}
+            mode={this.props.isSampleCopilotActive ? "Sample" : "User"}
+          />
+        )}
       </Fragment>
     );
   }

src/Explorer/Tabs/Tabs.tsx

@@ -11,7 +11,6 @@ import { QuickstartTab } from "Explorer/Tabs/QuickstartTab";
 import { VcoreMongoConnectTab } from "Explorer/Tabs/VCoreMongoConnectTab";
 import { VcoreMongoQuickstartTab } from "Explorer/Tabs/VCoreMongoQuickstartTab";
 import { userContext } from "UserContext";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { useTeachingBubble } from "hooks/useTeachingBubble";
 import ko from "knockout";
 import React, { MutableRefObject, useEffect, useRef, useState } from "react";
@@ -170,7 +169,7 @@ const CloseButton = ({
     onClick={(event: React.MouseEvent<HTMLSpanElement, MouseEvent>) => {
       event.stopPropagation();
       tab ? tab.onCloseTabButtonClick() : useTabs.getState().closeReactTab(tabKind);
-      tabKind === ReactTabKind.QueryCopilot && useQueryCopilot.getState().resetQueryCopilotStates();
+      // tabKind === ReactTabKind.QueryCopilot && useQueryCopilot.getState().resetQueryCopilotStates();
     }}
     tabIndex={active ? 0 : undefined}
     onKeyPress={({ nativeEvent: e }) => tab.onKeyPressClose(undefined, e)}
@@ -256,6 +255,7 @@ const isQueryErrorThrown = (tab?: Tab, tabKind?: ReactTabKind): boolean => {
 };
 
 const getReactTabContent = (activeReactTab: ReactTabKind, explorer: Explorer): JSX.Element => {
+  // eslint-disable-next-line no-console
   switch (activeReactTab) {
     case ReactTabKind.Connect:
       return userContext.apiType === "VCoreMongo" ? (

src/Explorer/Tabs/TabsBase.ts

@@ -3,15 +3,15 @@ import * as Constants from "../../Common/Constants";
 import * as ThemeUtility from "../../Common/ThemeUtility";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
-import { useNotificationConsole } from "../../hooks/useNotificationConsole";
-import { useTabs } from "../../hooks/useTabs";
 import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
+import { useNotificationConsole } from "../../hooks/useNotificationConsole";
+import { useTabs } from "../../hooks/useTabs";
 import { CommandButtonComponentProps } from "../Controls/CommandButton/CommandButtonComponent";
 import Explorer from "../Explorer";
 import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
-import { useSelectedNode } from "../useSelectedNode";
 import { WaitsForTemplateViewModel } from "../WaitsForTemplateViewModel";
+import { useSelectedNode } from "../useSelectedNode";
 // TODO: Use specific actions for logging telemetry data
 export default class TabsBase extends WaitsForTemplateViewModel {
   private static id = 0;

src/Explorer/Tree/Collection.ts

@@ -177,8 +177,8 @@ export default class Collection implements ViewModels.Collection {
     this.children.subscribe(() => {
       // update the database in zustand store
       const database = this.getDatabase();
-      database.collections(
-        database.collections()?.map((collection) => {
+      database?.collections(
+        database?.collections()?.map((collection) => {
           if (collection.id() === this.id()) {
             return this;
           }

src/Explorer/Tree/ResourceTokenCollection.ts

@@ -156,6 +156,6 @@ export default class ResourceTokenCollection implements ViewModels.CollectionBas
   }
 
   public getDatabase(): ViewModels.Database {
-    return useDatabases.getState().findDatabaseWithId(this.databaseId);
+    return useDatabases.getState().findDatabaseWithId(this.databaseId, this.isSampleCollection);
   }
 }

src/Explorer/Tree/ResourceTree.tsx

@@ -1,6 +1,7 @@
 import { Callout, DirectionalHint, ICalloutProps, ILinkProps, Link, Stack, Text } from "@fluentui/react";
 import { SampleDataTree } from "Explorer/Tree/SampleDataTree";
 import { getItemName } from "Utils/APITypeUtils";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as React from "react";
 import shallow from "zustand/shallow";
 import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
@@ -767,7 +768,8 @@ export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: Resourc
   };
 
   const dataRootNode = buildDataTree();
-  const isSampleDataEnabled = userContext.sampleDataConnectionInfo && userContext.apiType === "SQL";
+  const isSampleDataEnabled =
+    useQueryCopilot().copilotEnabled && userContext.sampleDataConnectionInfo && userContext.apiType === "SQL";
   const sampleDataResourceTokenCollection = useDatabases((state) => state.sampleDataResourceTokenCollection);
 
   return (

src/Explorer/useDatabases.ts

@@ -14,7 +14,7 @@ interface DatabasesState {
   deleteDatabase: (database: ViewModels.Database) => void;
   clearDatabases: () => void;
   isSaveQueryEnabled: () => boolean;
-  findDatabaseWithId: (databaseId: string) => ViewModels.Database;
+  findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => ViewModels.Database;
   isLastNonEmptyDatabase: () => boolean;
   findCollection: (databaseId: string, collectionId: string) => ViewModels.Collection;
   isLastCollection: () => boolean;
@@ -33,7 +33,7 @@ export const useDatabases: UseStore<DatabasesState> = create((set, get) => ({
   updateDatabase: (updatedDatabase: ViewModels.Database) =>
     set((state) => {
       const updatedDatabases = state.databases.map((database: ViewModels.Database) => {
-        if (database.id() === updatedDatabase.id()) {
+        if (database?.id() === updatedDatabase?.id()) {
           return updatedDatabase;
         }
 
@@ -67,7 +67,9 @@ export const useDatabases: UseStore<DatabasesState> = create((set, get) => ({
     }
     return true;
   },
-  findDatabaseWithId: (databaseId: string) => get().databases.find((db) => databaseId === db.id()),
+  findDatabaseWithId: (databaseId: string, isSampleDatabase?: boolean) => {
+    return 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());

src/HostedExplorer.tsx

@@ -7,9 +7,6 @@ import "../less/hostedexplorer.less";
 import { AuthType } from "./AuthType";
 import { DatabaseAccount } from "./Contracts/DataModels";
 import "./Explorer/Menus/NavBar/MeControlComponent.less";
-import { useAADAuth } from "./hooks/useAADAuth";
-import { useConfig } from "./hooks/useConfig";
-import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 import { HostedExplorerChildFrame } from "./HostedExplorerChildFrame";
 import { AccountSwitcher } from "./Platform/Hosted/Components/AccountSwitcher";
 import { ConnectExplorer } from "./Platform/Hosted/Components/ConnectExplorer";
@@ -20,6 +17,9 @@ import { SignInButton } from "./Platform/Hosted/Components/SignInButton";
 import "./Platform/Hosted/ConnectScreen.less";
 import { extractMasterKeyfromConnectionString } from "./Platform/Hosted/HostedUtils";
 import "./Shared/appInsights";
+import { useAADAuth } from "./hooks/useAADAuth";
+import { useConfig } from "./hooks/useConfig";
+import { useTokenMetadata } from "./hooks/usePortalAccessToken";
 
 initializeIcons();
 

src/Main.tsx

@@ -1,13 +1,13 @@
 // CSS Dependencies
 import { initializeIcons, loadTheme } from "@fluentui/react";
-import "bootstrap/dist/css/bootstrap.css";
 import { QuickstartCarousel } from "Explorer/Quickstart/QuickstartCarousel";
 import { MongoQuickstartTutorial } from "Explorer/Quickstart/Tutorials/MongoQuickstartTutorial";
 import { SQLQuickstartTutorial } from "Explorer/Quickstart/Tutorials/SQLQuickstartTutorial";
+import { userContext } from "UserContext";
+import "bootstrap/dist/css/bootstrap.css";
 import { useCarousel } from "hooks/useCarousel";
 import React, { useState } from "react";
 import ReactDOM from "react-dom";
-import { userContext } from "UserContext";
 import "../externals/jquery-ui.min.css";
 import "../externals/jquery-ui.min.js";
 import "../externals/jquery-ui.structure.min.css";
@@ -18,21 +18,19 @@ import "../externals/jquery.typeahead.min.js";
 // Image Dependencies
 import { Platform } from "ConfigContext";
 import { QueryCopilotCarousel } from "Explorer/QueryCopilot/CopilotCarousel";
-import { QueryCopilotFeedbackModal } from "Explorer/QueryCopilot/Modal/QueryCopilotFeedbackModal";
-import { useQueryCopilot } from "hooks/useQueryCopilot";
 import "../images/CosmosDB_rgb_ui_lighttheme.ico";
-import "../images/favicon.ico";
 import hdeConnectImage from "../images/HdeConnectCosmosDB.svg";
+import "../images/favicon.ico";
+import "../less/TableStyles/CustomizeColumns.less";
+import "../less/TableStyles/EntityEditor.less";
+import "../less/TableStyles/fulldatatables.less";
+import "../less/TableStyles/queryBuilder.less";
 import "../less/documentDB.less";
 import "../less/forms.less";
 import "../less/infobox.less";
 import "../less/menus.less";
 import "../less/messagebox.less";
 import "../less/resourceTree.less";
-import "../less/TableStyles/CustomizeColumns.less";
-import "../less/TableStyles/EntityEditor.less";
-import "../less/TableStyles/fulldatatables.less";
-import "../less/TableStyles/queryBuilder.less";
 import "../less/tree.less";
 import { CollapsedResourceTree } from "./Common/CollapsedResourceTree";
 import { ResourceTreeContainer } from "./Common/ResourceTreeContainer";
@@ -55,11 +53,11 @@ import "./Explorer/Panes/PanelComponent.less";
 import { SidePanel } from "./Explorer/Panes/PanelContainerComponent";
 import "./Explorer/SplashScreen/SplashScreen.less";
 import { Tabs } from "./Explorer/Tabs/Tabs";
-import { useConfig } from "./hooks/useConfig";
-import { useKnockoutExplorer } from "./hooks/useKnockoutExplorer";
 import "./Libs/jquery";
 import { appThemeFabric } from "./Platform/Fabric/FabricTheme";
 import "./Shared/appInsights";
+import { useConfig } from "./hooks/useConfig";
+import { useKnockoutExplorer } from "./hooks/useKnockoutExplorer";
 
 initializeIcons();
 
@@ -67,7 +65,6 @@ const App: React.FunctionComponent = () => {
   const [isLeftPaneExpanded, setIsLeftPaneExpanded] = useState<boolean>(true);
   const isCarouselOpen = useCarousel((state) => state.shouldOpen);
   const isCopilotCarouselOpen = useCarousel((state) => state.showCopilotCarousel);
-  const shouldShowModal = useQueryCopilot((state) => state.showFeedbackModal);
 
   const config = useConfig();
   if (config?.platform === Platform.Fabric) {
@@ -136,7 +133,6 @@ const App: React.FunctionComponent = () => {
       {<SQLQuickstartTutorial />}
       {<MongoQuickstartTutorial />}
       {<QueryCopilotCarousel isOpen={isCopilotCarouselOpen} explorer={explorer} />}
-      {shouldShowModal && <QueryCopilotFeedbackModal explorer={explorer} />}
     </div>
   );
 };

src/Platform/Hosted/Components/ConnectExplorer.tsx

@@ -69,7 +69,9 @@ export const ConnectExplorer: React.FunctionComponent<Props> = ({
                 setErrorMessage("");
 
                 if (await isAccountRestrictedForConnectionStringLogin(connectionString)) {
-                  setErrorMessage("This account has been blocked from connection-string login.");
+                  setErrorMessage(
+                    "This account has been blocked from connection-string login. Please go to cosmos.azure.com/aad for AAD based login.",
+                  );
                   return;
                 }
 

src/Platform/Hosted/ConnectScreen.less

@@ -66,7 +66,7 @@
 }
 .connectExplorerContainer .connectExplorer .connectExplorerContent .errorDetailsInfoTooltip .errorDetails {
   bottom: 24px;
-  width: 145px;
+  width: 165px;
   visibility: hidden;
   background-color: #393939;
   color: #ffffff;

src/Platform/Hosted/extractFeatures.ts

@@ -14,6 +14,7 @@ export type Features = {
   readonly enableTtl: boolean;
   readonly executeSproc: boolean;
   readonly enableAadDataPlane: boolean;
+  readonly enableResourceGraph: boolean;
   readonly enableKoResourceTree: boolean;
   readonly hostedDataExplorer: boolean;
   readonly junoEndpoint?: string;
@@ -73,6 +74,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     canExceedMaximumValue: "true" === get("canexceedmaximumvalue"),
     cosmosdb: "true" === get("cosmosdb"),
     enableAadDataPlane: "true" === get("enableaaddataplane"),
+    enableResourceGraph: "true" === get("enableresourcegraph"),
     enableChangeFeedPolicy: "true" === get("enablechangefeedpolicy"),
     enableFixedCollectionWithSharedThroughput: "true" === get("enablefixedcollectionwithsharedthroughput"),
     enableKOPanel: "true" === get("enablekopanel"),
@@ -110,7 +112,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     enableLegacyMongoShellV2Debug: "true" === get("enablelegacymongoshellv2debug"),
     loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"),
     enableCopilot: "true" === get("enablecopilot", "true"),
-    copilotVersion: get("copilotversion") ?? "v1.0",
+    copilotVersion: get("copilotversion") ?? "v2.0",
     disableCopilotPhoenixGateaway: "true" === get("disablecopilotphoenixgateaway"),
     enableCopilotFullSchema: "true" === get("enablecopilotfullschema", "true"),
     copilotChatFixedMonacoEditorHeight: "true" === get("copilotchatfixedmonacoeditorheight"),

src/Shared/Telemetry/TelemetryConstants.ts

@@ -133,6 +133,10 @@ export enum Action {
   CompleteUITour,
   OpenQueryCopilotFromSplashScreen,
   OpenQueryCopilotFromNewQuery,
+  ActivateQueryCopilot,
+  DeactivateQueryCopilot,
+  QueryGenerationFromCopilotPrompt,
+  QueryEdited,
   ExecuteQueryGeneratedFromQueryCopilot,
 }
 

src/Terminal/JupyterLabAppFactory.ts

@@ -10,9 +10,13 @@ import { userContext } from "UserContext";
 export class JupyterLabAppFactory {
   private isShellStarted: boolean | undefined;
   private checkShellStarted: ((content: string | undefined) => void) | undefined;
-  private onShellExited: () => void;
+  private onShellExited: (restartShell: boolean) => void;
+  private restartShell: boolean;
 
   private isShellExited(content: string | undefined) {
+    if (userContext.apiType === "VCoreMongo" && content?.includes("MongoServerError: Invalid key")) {
+      this.restartShell = true;
+    }
     return content?.includes("cosmosuser@");
   }
 
@@ -32,10 +36,11 @@ export class JupyterLabAppFactory {
     this.isShellStarted = content?.includes("Enter password");
   }
 
-  constructor(closeTab: () => void) {
+  constructor(closeTab: (restartShell: boolean) => void) {
     this.onShellExited = closeTab;
     this.isShellStarted = false;
     this.checkShellStarted = undefined;
+    this.restartShell = false;
 
     switch (userContext.apiType) {
       case "Mongo":
@@ -69,7 +74,7 @@ export class JupyterLabAppFactory {
         if (!this.isShellStarted) {
           this.checkShellStarted(content);
         } else if (this.isShellExited(content)) {
-          this.onShellExited();
+          this.onShellExited(this.restartShell);
         }
       }
     }, this);

src/Terminal/index.ts

@@ -11,6 +11,8 @@ import { JupyterLabAppFactory } from "./JupyterLabAppFactory";
 import { TerminalProps } from "./TerminalProps";
 import "./index.css";
 
+let session: ITerminalConnection | undefined;
+
 const createServerSettings = (props: TerminalProps): ServerConnection.ISettings => {
   let body: BodyInit | undefined;
   let headers: HeadersInit | undefined;
@@ -49,7 +51,7 @@ const createServerSettings = (props: TerminalProps): ServerConnection.ISettings
   return ServerConnection.makeSettings(options);
 };
 
-const initTerminal = async (props: TerminalProps): Promise<ITerminalConnection | undefined> => {
+const initTerminal = async (props: TerminalProps): Promise<void> => {
   // Initialize userContext (only properties which are needed by TelemetryProcessor)
   updateUserContext({
     subscriptionId: props.subscriptionId,
@@ -59,28 +61,37 @@ const initTerminal = async (props: TerminalProps): Promise<ITerminalConnection |
   });
 
   const serverSettings = createServerSettings(props);
+
+  createTerminalApp(props, serverSettings);
+};
+
+const createTerminalApp = async (props: TerminalProps, serverSettings: ServerConnection.ISettings) => {
   const data = { baseUrl: serverSettings.baseUrl };
   const startTime = TelemetryProcessor.traceStart(Action.OpenTerminal, data);
 
   try {
-    const session = await new JupyterLabAppFactory(() => closeTab(props.tabId)).createTerminalApp(serverSettings);
+    session = await new JupyterLabAppFactory((restartShell: boolean) =>
+      closeTab(props, serverSettings, restartShell),
+    ).createTerminalApp(serverSettings);
     TelemetryProcessor.traceSuccess(Action.OpenTerminal, data, startTime);
-    return session;
   } catch (error) {
     TelemetryProcessor.traceFailure(Action.OpenTerminal, data, startTime);
-    return undefined;
+    session = undefined;
   }
 };
 
-const closeTab = (tabId: string): void => {
-  window.parent.postMessage(
-    { type: MessageTypes.CloseTab, data: { tabId: tabId }, signature: "pcIframe" },
-    window.document.referrer,
-  );
+const closeTab = (props: TerminalProps, serverSettings: ServerConnection.ISettings, restartShell: boolean): void => {
+  if (restartShell) {
+    createTerminalApp(props, serverSettings);
+  } else {
+    window.parent.postMessage(
+      { type: MessageTypes.CloseTab, data: { tabId: props.tabId }, signature: "pcIframe" },
+      window.document.referrer,
+    );
+  }
 };
 
 const main = async (): Promise<void> => {
-  let session: ITerminalConnection | undefined;
   postRobot.on(
     "props",
     {
@@ -91,7 +102,7 @@ const main = async (): Promise<void> => {
       // Typescript definition for event is wrong. So read props by casting to <any>
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       const props = (event as any).data as TerminalProps;
-      session = await initTerminal(props);
+      await initTerminal(props);
     },
   );
 

src/hooks/useDatabaseAccounts.tsx

@@ -1,6 +1,9 @@
+import { HttpHeaders } from "Common/Constants";
+import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { DatabaseAccount } from "../Contracts/DataModels";
+/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface AccountListResult {
   nextLink: string;
@@ -30,10 +33,56 @@ export async function fetchDatabaseAccounts(subscriptionId: string, accessToken:
   return accounts.sort((a, b) => a.name.localeCompare(b.name));
 }
 
+export async function fetchDatabaseAccountsFromGraph(
+  subscriptionId: string,
+  accessToken: string,
+): Promise<DatabaseAccount[]> {
+  const headers = new Headers();
+  const bearer = `Bearer ${accessToken}`;
+
+  headers.append("Authorization", bearer);
+  headers.append(HttpHeaders.contentType, "application/json");
+  const databaseAccountsQuery = "resources | where type =~ 'microsoft.documentdb/databaseaccounts'";
+  const apiVersion = "2021-03-01";
+  const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+
+  const databaseAccounts: DatabaseAccount[] = [];
+  let skipToken: string;
+  do {
+    const body = {
+      query: databaseAccountsQuery,
+      subscriptions: [subscriptionId],
+      ...(skipToken && {
+        options: {
+          $skipToken: skipToken,
+        } as QueryRequestOptions,
+      }),
+    };
+
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(body),
+    });
+
+    if (!response.ok) {
+      throw new Error(await response.text());
+    }
+
+    const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
+    skipToken = queryResponse.$skipToken;
+    queryResponse.data?.map((databaseAccount: any) => {
+      databaseAccounts.push(databaseAccount as DatabaseAccount);
+    });
+  } while (skipToken);
+
+  return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
+}
+
 export function useDatabaseAccounts(subscriptionId: string, armToken: string): DatabaseAccount[] | undefined {
   const { data } = useSWR(
     () => (armToken && subscriptionId ? ["databaseAccounts", subscriptionId, armToken] : undefined),
-    (_, subscriptionId, armToken) => fetchDatabaseAccounts(subscriptionId, armToken),
+    (_, subscriptionId, armToken) => fetchDatabaseAccountsFromGraph(subscriptionId, armToken),
   );
   return data;
 }

src/hooks/useKnockoutExplorer.ts

@@ -7,6 +7,7 @@ import {
   scheduleRefreshDatabaseResourceToken,
 } from "Platform/Fabric/FabricUtil";
 import { getNetworkSettingsWarningMessage } from "Utils/NetworkUtility";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
 import { useEffect, useState } from "react";
 import { AuthType } from "../AuthType";
@@ -57,19 +58,21 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
           userContext.features.phoenixNotebooks = true;
           userContext.features.phoenixFeatures = true;
         }
+        let explorer: Explorer;
         if (platform === Platform.Hosted) {
-          const explorer = await configureHosted();
-          setExplorer(explorer);
+          explorer = await configureHosted();
         } else if (platform === Platform.Emulator) {
-          const explorer = configureEmulator();
-          setExplorer(explorer);
+          explorer = configureEmulator();
         } else if (platform === Platform.Portal) {
-          const explorer = await configurePortal();
-          setExplorer(explorer);
+          explorer = await configurePortal();
         } else if (platform === Platform.Fabric) {
-          const explorer = await configureFabric();
-          setExplorer(explorer);
+          explorer = await configureFabric();
         }
+        if (explorer && userContext.features.enableCopilot) {
+          await updateContextForCopilot(explorer);
+          await updateContextForSampleData(explorer);
+        }
+        setExplorer(explorer);
       }
     };
     effect();
@@ -78,9 +81,6 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
   useEffect(() => {
     if (explorer) {
       applyExplorerBindings(explorer);
-      if (userContext.features.enableCopilot) {
-        updateContextForSampleData(explorer);
-      }
     }
   }, [explorer]);
 
@@ -420,9 +420,11 @@ async function configurePortal(): Promise<Explorer> {
           updateContextsFromPortalMessage(inputs);
           explorer = new Explorer();
           resolve(explorer);
-          if (userContext.apiType === "Postgres") {
-            explorer.openNPSSurveyDialog();
+
+          if (userContext.apiType === "Postgres" || userContext.apiType === "SQL" || userContext.apiType === "Mongo") {
+            setTimeout(() => explorer.openNPSSurveyDialog(), 3000);
           }
+
           if (openAction) {
             handleOpenAction(openAction, useDatabases.getState().databases, explorer);
           }
@@ -554,12 +556,23 @@ interface PortalMessage {
   inputs?: DataExplorerInputsFrame;
 }
 
+async function updateContextForCopilot(explorer: Explorer): Promise<void> {
+  await explorer.configureCopilot();
+}
+
 async function updateContextForSampleData(explorer: Explorer): Promise<void> {
-  if (!userContext.features.enableCopilot) {
+  const copilotEnabled =
+    userContext.apiType === "SQL" && userContext.features.enableCopilot && useQueryCopilot.getState().copilotEnabled;
+
+  if (!copilotEnabled) {
     return;
   }
 
-  const url = createUri(`${configContext.BACKEND_ENDPOINT}`, `/api/tokens/sampledataconnection`);
+  const sampleDatabaseEndpoint = useQueryCopilot.getState().copilotUserDBEnabled
+    ? `/api/tokens/sampledataconnection/v2`
+    : `/api/tokens/sampledataconnection`;
+
+  const url = createUri(`${configContext.BACKEND_ENDPOINT}`, sampleDatabaseEndpoint);
   const authorizationHeader = getAuthorizationHeader();
   const headers = { [authorizationHeader.header]: authorizationHeader.token };
 

src/hooks/useQueryCopilot.ts

@@ -1,6 +1,6 @@
 import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import { QueryResults } from "Contracts/ViewModels";
-import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
+import { CopilotMessage, CopilotSchemaAllocationInfo } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { guid } from "Explorer/Tables/Utilities";
 import { useTabs } from "hooks/useTabs";
 import create, { UseStore } from "zustand";
@@ -8,6 +8,8 @@ import * as DataModels from "../Contracts/DataModels";
 import { ContainerInfo } from "../Contracts/DataModels";
 
 export interface QueryCopilotState {
+  copilotEnabled: boolean;
+  copilotUserDBEnabled: boolean;
   generatedQuery: string;
   likeQuery: boolean;
   userPrompt: string;
@@ -34,14 +36,21 @@ export interface QueryCopilotState {
   generatedQueryComments: string;
   wasCopilotUsed: boolean;
   showWelcomeSidebar: boolean;
+  showWelcomeModal: boolean;
   showCopilotSidebar: boolean;
   chatMessages: CopilotMessage[];
   shouldIncludeInMessages: boolean;
   showExplanationBubble: boolean;
   notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo;
   containerStatus: ContainerInfo;
+  schemaAllocationInfo: CopilotSchemaAllocationInfo;
   isAllocatingContainer: boolean;
+  copilotEnabledforExecution: boolean;
 
+  getState?: () => QueryCopilotState;
+
+  setCopilotEnabled: (copilotEnabled: boolean) => void;
+  setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => void;
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
   setHideFeedbackModalForLikedQueries: (hideFeedbackModalForLikedQueries: boolean) => void;
@@ -69,6 +78,7 @@ export interface QueryCopilotState {
   setGeneratedQueryComments: (generatedQueryComments: string) => void;
   setWasCopilotUsed: (wasCopilotUsed: boolean) => void;
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => void;
+  setShowWelcomeModal: (showWelcomeModal: boolean) => void;
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
   setChatMessages: (chatMessages: CopilotMessage[]) => void;
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
@@ -76,6 +86,8 @@ export interface QueryCopilotState {
   setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
   setContainerStatus: (containerStatus: ContainerInfo) => void;
   setIsAllocatingContainer: (isAllocatingContainer: boolean) => void;
+  setSchemaAllocationInfo: (schemaAllocationInfo: CopilotSchemaAllocationInfo) => void;
+  setCopilotEnabledforExecution: (copilotEnabledforExecution: boolean) => void;
 
   resetContainerConnection: () => void;
   resetQueryCopilotStates: () => void;
@@ -84,6 +96,8 @@ export interface QueryCopilotState {
 type QueryCopilotStore = UseStore<QueryCopilotState>;
 
 export const useQueryCopilot: QueryCopilotStore = create((set) => ({
+  copilotEnabled: false,
+  copilotUserDBEnabled: false,
   generatedQuery: "",
   likeQuery: false,
   userPrompt: "",
@@ -110,6 +124,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   generatedQueryComments: "",
   wasCopilotUsed: false,
   showWelcomeSidebar: true,
+  showWelcomeModal: true,
   showCopilotSidebar: false,
   chatMessages: [],
   shouldIncludeInMessages: true,
@@ -124,8 +139,15 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
     durationLeftInMinutes: undefined,
     phoenixServerInfo: undefined,
   },
+  schemaAllocationInfo: {
+    databaseId: undefined,
+    containerId: undefined,
+  },
   isAllocatingContainer: false,
+  copilotEnabledforExecution: false,
 
+  setCopilotEnabled: (copilotEnabled: boolean) => set({ copilotEnabled }),
+  setCopilotUserDBEnabled: (copilotUserDBEnabled: boolean) => set({ copilotUserDBEnabled }),
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
   closeFeedbackModal: () => set({ showFeedbackModal: false }),
@@ -155,6 +177,7 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setGeneratedQueryComments: (generatedQueryComments: string) => set({ generatedQueryComments }),
   setWasCopilotUsed: (wasCopilotUsed: boolean) => set({ wasCopilotUsed }),
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
+  setShowWelcomeModal: (showWelcomeModal: boolean) => set({ showWelcomeModal }),
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
   setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
@@ -163,6 +186,8 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
     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);
@@ -173,6 +198,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       durationLeftInMinutes: undefined,
       phoenixServerInfo: undefined,
     });
+    useQueryCopilot.getState().setSchemaAllocationInfo({
+      databaseId: undefined,
+      containerId: undefined,
+    });
   },
 
   resetQueryCopilotStates: () => {
@@ -217,6 +246,10 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
         durationLeftInMinutes: undefined,
         phoenixServerInfo: undefined,
       },
+      schemaAllocationInfo: {
+        databaseId: undefined,
+        containerId: undefined,
+      },
       isAllocatingContainer: false,
     }));
   },

src/hooks/useSubscriptions.tsx

@@ -1,6 +1,9 @@
+import { HttpHeaders } from "Common/Constants";
+import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { Subscription } from "../Contracts/DataModels";
+/* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface SubscriptionListResult {
   nextLink: string;
@@ -32,10 +35,58 @@ export async function fetchSubscriptions(accessToken: string): Promise<Subscript
   return subscriptions.sort((a, b) => a.displayName.localeCompare(b.displayName));
 }
 
+export async function fetchSubscriptionsFromGraph(accessToken: string): Promise<Subscription[]> {
+  const headers = new Headers();
+  const bearer = `Bearer ${accessToken}`;
+
+  headers.append("Authorization", bearer);
+  headers.append(HttpHeaders.contentType, "application/json");
+  const subscriptionsQuery =
+    "resources | where type == 'microsoft.documentdb/databaseaccounts' | join kind=inner ( resourcecontainers | where type == 'microsoft.resources/subscriptions' | project subscriptionId, subscriptionName = name, subscriptionState = tostring(parse_json(properties).state) ) on subscriptionId |  summarize by subscriptionId, subscriptionName, subscriptionState";
+  const apiVersion = "2021-03-01";
+  const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
+
+  const subscriptions: Subscription[] = [];
+  let skipToken: string;
+  do {
+    const body = {
+      query: subscriptionsQuery,
+      ...(skipToken && {
+        options: {
+          $skipToken: skipToken,
+        } as QueryRequestOptions,
+      }),
+    };
+
+    const response = await fetch(managementResourceGraphAPIURL, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(body),
+    });
+
+    if (!response.ok) {
+      throw new Error(await response.text());
+    }
+
+    const queryResponse: QueryResponse = (await response.json()) as QueryResponse;
+    skipToken = queryResponse.$skipToken;
+
+    queryResponse.data?.map((subscription: any) => {
+      subscriptions.push({
+        displayName: subscription.subscriptionName,
+        subscriptionId: subscription.subscriptionId,
+        state: subscription.subscriptionState,
+      } as Subscription);
+    });
+  } while (skipToken);
+
+  return subscriptions.sort((a, b) => a.displayName.localeCompare(b.displayName));
+}
+
 export function useSubscriptions(armToken: string): Subscription[] | undefined {
   const { data } = useSWR(
     () => (armToken ? ["subscriptions", armToken] : undefined),
-    (_, armToken) => fetchSubscriptions(armToken),
+    (_, armToken) => fetchSubscriptionsFromGraph(armToken),
   );
   return data;
 }

src/hooks/useTabs.ts

@@ -5,7 +5,7 @@ import NotebookTabV2 from "../Explorer/Tabs/NotebookV2Tab";
 import TabsBase from "../Explorer/Tabs/TabsBase";
 import { Platform, configContext } from "./../ConfigContext";
 
-interface TabsState {
+export interface TabsState {
   openedTabs: TabsBase[];
   openedReactTabs: ReactTabKind[];
   activeTab: TabsBase | undefined;