[Query Copilot] Phoenix container implementation (#1594)

* Phoenix implementation

* removing comments

---------

Co-authored-by: Predrag Klepic <v-prklepic@microsoft.com>

src/Explorer/Explorer.tsx

@@ -5,6 +5,7 @@ import { Platform } from "ConfigContext";
 import { MessageTypes } from "Contracts/ExplorerContracts";
 import { IGalleryItem } from "Juno/JunoClient";
 import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointValidation";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import * as ko from "knockout";
 import React from "react";
 import _ from "underscore";
@@ -386,8 +387,14 @@ export default class Explorer {
   }
 
   public async allocateContainer(poolId: PoolIdType): Promise<void> {
-    const notebookServerInfo = useNotebook.getState().notebookServerInfo;
+    const shouldUseNotebookStates = poolId === PoolIdType.DefaultPoolId ? true : false;
-    const isAllocating = useNotebook.getState().isAllocating;
+    const notebookServerInfo = shouldUseNotebookStates
+      ? useNotebook.getState().notebookServerInfo
+      : useQueryCopilot.getState().notebookServerInfo;
+
+    const isAllocating = shouldUseNotebookStates
+      ? useNotebook.getState().isAllocating
+      : useQueryCopilot.getState().isAllocatingContainer;
     if (
       isAllocating === false &&
       (notebookServerInfo === undefined ||
@@ -395,23 +402,28 @@ export default class Explorer {
     ) {
       const provisionData: IProvisionData = {
         cosmosEndpoint: userContext?.databaseAccount?.properties?.documentEndpoint,
-        poolId: poolId === PoolIdType.DefaultPoolId ? undefined : poolId,
+        poolId: shouldUseNotebookStates ? undefined : poolId,
       };
       const connectionStatus: ContainerConnectionInfo = {
         status: ConnectionStatusType.Connecting,
       };
-      useNotebook.getState().setConnectionInfo(connectionStatus);
+
+      shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus);
+
       let connectionInfo;
       try {
         TelemetryProcessor.traceStart(Action.PhoenixConnection, {
           dataExplorerArea: Areas.Notebook,
         });
-        useNotebook.getState().setIsAllocating(true);
+        shouldUseNotebookStates
+          ? useNotebook.getState().setIsAllocating(true)
+          : useQueryCopilot.getState().setIsAllocatingContainer(true);
+
         connectionInfo = await this.phoenixClient.allocateContainer(provisionData);
         if (!connectionInfo?.data?.phoenixServiceUrl) {
           throw new Error(`PhoenixServiceUrl is invalid!`);
         }
-        await this.setNotebookInfo(connectionInfo, connectionStatus);
+        await this.setNotebookInfo(shouldUseNotebookStates, connectionInfo, connectionStatus);
         TelemetryProcessor.traceSuccess(Action.PhoenixConnection, {
           dataExplorerArea: Areas.Notebook,
         });
@@ -423,7 +435,9 @@ export default class Explorer {
           errorStack: getErrorStack(error),
         });
         connectionStatus.status = ConnectionStatusType.Failed;
-        useNotebook.getState().resetContainerConnection(connectionStatus);
+        shouldUseNotebookStates
+          ? useNotebook.getState().resetContainerConnection(connectionStatus)
+          : useQueryCopilot.getState().resetContainerConnection();
         if (error?.status === HttpStatusCodes.Forbidden && error.message) {
           useDialog.getState().showOkModalDialog("Connection Failed", `${error.message}`);
         } else {
@@ -436,7 +450,9 @@ export default class Explorer {
         }
         throw error;
       } finally {
-        useNotebook.getState().setIsAllocating(false);
+        shouldUseNotebookStates
+          ? useNotebook.getState().setIsAllocating(false)
+          : useQueryCopilot.getState().setIsAllocatingContainer(false);
         this.refreshCommandBarButtons();
         this.refreshNotebookList();
         this._isInitializingNotebooks = false;
@@ -445,6 +461,7 @@ export default class Explorer {
   }
 
   private async setNotebookInfo(
+    shouldUseNotebookStates: boolean,
     connectionInfo: IResponse<IPhoenixServiceInfo>,
     connectionStatus: DataModels.ContainerConnectionInfo
   ) {
@@ -452,21 +469,26 @@ export default class Explorer {
       forwardingId: connectionInfo.data.forwardingId,
       dbAccountName: userContext.databaseAccount.name,
     };
-    await this.phoenixClient.initiateContainerHeartBeat(containerData);
+    await this.phoenixClient.initiateContainerHeartBeat(shouldUseNotebookStates, containerData);
 
     connectionStatus.status = ConnectionStatusType.Connected;
-    useNotebook.getState().setConnectionInfo(connectionStatus);
+    shouldUseNotebookStates && useNotebook.getState().setConnectionInfo(connectionStatus);
-    useNotebook.getState().setNotebookServerInfo({
+
+    const noteBookServerInfo = {
       notebookServerEndpoint:
         (validateEndpoint(userContext.features.notebookServerUrl, allowedNotebookServerUrls) &&
           userContext.features.notebookServerUrl) ||
         connectionInfo.data.phoenixServiceUrl,
       authToken: userContext.features.notebookServerToken || connectionInfo.data.authToken,
       forwardingId: connectionInfo.data.forwardingId,
-    });
+    };
-    this.notebookManager?.notebookClient
+    shouldUseNotebookStates
-      .getMemoryUsage()
+      ? useNotebook.getState().setNotebookServerInfo(noteBookServerInfo)
-      .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
+      : useQueryCopilot.getState().setNotebookServerInfo(noteBookServerInfo);
+    shouldUseNotebookStates &&
+      this.notebookManager?.notebookClient
+        .getMemoryUsage()
+        .then((memoryUsageInfo) => useNotebook.getState().setMemoryUsageInfo(memoryUsageInfo));
   }
 
   public resetNotebookWorkspace(): void {
@@ -540,7 +562,7 @@ export default class Explorer {
         throw new Error(`Reset Workspace: PhoenixServiceUrl is invalid!`);
       }
       if (useNotebook.getState().isPhoenixNotebooks) {
-        await this.setNotebookInfo(connectionInfo, connectionStatus);
+        await this.setNotebookInfo(true, connectionInfo, connectionStatus);
         useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
       }
       logConsoleInfo("Successfully reset notebook workspace");

src/Explorer/QueryCopilot/QueryCopilotTab.tsx

@@ -18,6 +18,7 @@ import {
 } from "@fluentui/react";
 import { useBoolean } from "@fluentui/react-hooks";
 import {
+  ContainerStatusType,
   PoolIdType,
   QueryCopilotSampleContainerId,
   QueryCopilotSampleContainerSchema,
@@ -32,7 +33,6 @@ import { QueryResults } from "Contracts/ViewModels";
 import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
-import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { SaveQueryPane } from "Explorer/Panes/SaveQueryPane/SaveQueryPane";
 import { WelcomeModal } from "Explorer/QueryCopilot/Modal/WelcomeModal";
 import { CopyPopup } from "Explorer/QueryCopilot/Popup/CopyPopup";
@@ -111,8 +111,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
     setShowErrorMessageBar,
     generatedQueryComments,
     setGeneratedQueryComments,
-    shouldAllocateContainer,
-    setShouldAllocateContainer,
   } = useQueryCopilot();
 
   const sampleProps: SamplePromptsProps = {
@@ -180,23 +178,24 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
 
   const generateSQLQuery = async (): Promise<void> => {
     try {
-      if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
-        await explorer.allocateContainer(PoolIdType.QueryCopilot);
-        setShouldAllocateContainer(false);
-      }
-
       setIsGeneratingQuery(true);
+      setShowDeletePopup(false);
       useTabs.getState().setIsTabExecuting(true);
       useTabs.getState().setIsQueryErrorThrown(false);
+      if (
+        useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
+        userContext.features.enableCopilotPhoenixGateaway
+      ) {
+        await explorer.allocateContainer(PoolIdType.QueryCopilot);
+      }
       const payload = {
         containerSchema: userContext.features.enableCopilotFullSchema
           ? QueryCopilotSampleContainerSchema
           : ShortenedQueryCopilotSampleContainerSchema,
         userPrompt: userPrompt,
       };
-      setShowDeletePopup(false);
       useQueryCopilot.getState().refreshCorrelationId();
-      const serverInfo = useNotebook.getState().notebookServerInfo;
+      const serverInfo = useQueryCopilot.getState().notebookServerInfo;
       const queryUri = userContext.features.enableCopilotPhoenixGateaway
         ? createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery")
         : createUri("https://copilotorchestrater.azurewebsites.net/", "generateSQLQuery");
@@ -210,9 +209,6 @@ export const QueryCopilotTab: React.FC<QueryCopilotProps> = ({ explorer }: Query
       });
 
       const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
-      if (response.status === 404) {
-        setShouldAllocateContainer(true);
-      }
       if (response.ok) {
         if (generateSQLQueryResponse?.sql) {
           let query = `-- **Prompt:** ${userPrompt}\r\n`;

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

@@ -2,9 +2,9 @@ import { QueryCopilotSampleContainerSchema, ShortenedQueryCopilotSampleContainer
 import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
 import Explorer from "Explorer/Explorer";
-import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { SubmitFeedback } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { userContext } from "UserContext";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 
 jest.mock("@azure/cosmos", () => ({
   Constants: {
@@ -27,22 +27,10 @@ jest.mock("Explorer/Explorer", () => {
   return MockExplorer;
 });
 
-jest.mock("hooks/useQueryCopilot", () => {
-  const mockQueryCopilotStore = {
-    shouldAllocateContainer: true,
-    setShouldAllocateContainer: jest.fn(),
-    correlationId: "mocked-correlation-id",
-  };
-
-  return {
-    useQueryCopilot: jest.fn(() => mockQueryCopilotStore),
-  };
-});
-
 describe("Query Copilot Client", () => {
   beforeEach(() => jest.clearAllMocks());
 
-  describe("submitFeedback", () => {
+  describe("SubmitFeedback", () => {
     const payload = {
       like: "like",
       generatedSql: "GeneratedQuery",
@@ -54,17 +42,16 @@ describe("Query Copilot Client", () => {
         : ShortenedQueryCopilotSampleContainerSchema,
     };
 
-    const mockStore = useNotebook.getState();
+    const mockStore = useQueryCopilot.getState();
-    beforeEach(() => {
+    mockStore.correlationId = "mocked-correlation-id";
-      mockStore.notebookServerInfo = {
+    mockStore.notebookServerInfo = {
-        notebookServerEndpoint: "mocked-endpoint",
+      notebookServerEndpoint: "mocked-endpoint",
-        authToken: "mocked-token",
+      authToken: "mocked-token",
-        forwardingId: "mocked-forwarding-id",
+      forwardingId: "mocked-forwarding-id",
-      };
+    };
-    });
 
     const feedbackUri = userContext.features.enableCopilotPhoenixGateaway
-      ? createUri(useNotebook.getState().notebookServerInfo.notebookServerEndpoint, "feedback")
+      ? createUri(useQueryCopilot.getState().notebookServerInfo.notebookServerEndpoint, "feedback")
       : createUri("https://copilotorchestrater.azurewebsites.net/", "feedback");
 
     it("should call fetch with the payload with like", async () => {

src/Explorer/QueryCopilot/Shared/QueryCopilotClient.ts

@@ -1,4 +1,5 @@
 import {
+  ContainerStatusType,
   PoolIdType,
   QueryCopilotSampleContainerSchema,
   ShortenedQueryCopilotSampleContainerSchema,
@@ -6,7 +7,6 @@ import {
 import { handleError } from "Common/ErrorHandlingUtils";
 import { createUri } from "Common/UrlUtility";
 import Explorer from "Explorer/Explorer";
-import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { FeedbackParams, GenerateSQLQueryResponse } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { userContext } from "UserContext";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
@@ -30,13 +30,15 @@ export const SendQueryRequest = async ({
       .getState()
       .setChatMessages([...useQueryCopilot.getState().chatMessages, { source: 0, message: userPrompt }]);
     try {
-      if (useQueryCopilot.getState().shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
+      if (
-        await explorer.allocateContainer(PoolIdType.DefaultPoolId);
+        useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
-        useQueryCopilot.getState().setShouldAllocateContainer(false);
+        userContext.features.enableCopilotPhoenixGateaway
+      ) {
+        await explorer.allocateContainer(PoolIdType.QueryCopilot);
       }
 
       useQueryCopilot.getState().refreshCorrelationId();
-      const serverInfo = useNotebook.getState().notebookServerInfo;
+      const serverInfo = useQueryCopilot.getState().notebookServerInfo;
 
       const queryUri = userContext.features.enableCopilotPhoenixGateaway
         ? createUri(serverInfo.notebookServerEndpoint, "generateSQLQuery")
@@ -58,9 +60,6 @@ export const SendQueryRequest = async ({
       });
 
       const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
-      if (response.status === 404) {
-        useQueryCopilot.getState().setShouldAllocateContainer(true);
-      }
       if (response.ok) {
         if (generateSQLQueryResponse?.sql) {
           let query = `Here is a query which will help you with provided prompt.\r\n **Prompt:** ${userPrompt}`;
@@ -102,7 +101,6 @@ export const SubmitFeedback = async ({
 }): Promise<void> => {
   try {
     const { likeQuery, generatedQuery, userPrompt, description, contact } = params;
-    const { correlationId, shouldAllocateContainer, setShouldAllocateContainer } = useQueryCopilot();
     const payload = {
       containerSchema: userContext.features.enableCopilotFullSchema
         ? QueryCopilotSampleContainerSchema
@@ -113,25 +111,24 @@ export const SubmitFeedback = async ({
       description: description || "",
       contact: contact || "",
     };
-    if (shouldAllocateContainer && userContext.features.enableCopilotPhoenixGateaway) {
+    if (
-      await explorer.allocateContainer(PoolIdType.DefaultPoolId);
+      useQueryCopilot.getState().containerStatus.status !== ContainerStatusType.Active &&
-      setShouldAllocateContainer(false);
+      userContext.features.enableCopilotPhoenixGateaway
+    ) {
+      await explorer.allocateContainer(PoolIdType.QueryCopilot);
     }
-    const serverInfo = useNotebook.getState().notebookServerInfo;
+    const serverInfo = useQueryCopilot.getState().notebookServerInfo;
     const feedbackUri = userContext.features.enableCopilotPhoenixGateaway
       ? createUri(serverInfo.notebookServerEndpoint, "feedback")
       : createUri("https://copilotorchestrater.azurewebsites.net/", "feedback");
-    const response = await fetch(feedbackUri, {
+    await fetch(feedbackUri, {
       method: "POST",
       headers: {
         "content-type": "application/json",
-        "x-ms-correlationid": correlationId,
+        "x-ms-correlationid": useQueryCopilot.getState().correlationId,
       },
       body: JSON.stringify(payload),
     });
-    if (response.status === 404) {
-      setShouldAllocateContainer(true);
-    }
   } catch (error) {
     handleError(error, "copilotSubmitFeedback");
   }

src/Phoenix/PhoenixClient.ts

@@ -3,6 +3,7 @@ import { useDialog } from "Explorer/Controls/Dialog";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { userContext } from "UserContext";
 import { allowedJunoOrigins, validateEndpoint } from "Utils/EndpointValidation";
+import { useQueryCopilot } from "hooks/useQueryCopilot";
 import promiseRetry, { AbortError } from "p-retry";
 import {
   Areas,
@@ -96,20 +97,27 @@ export class PhoenixClient {
     }
   }
 
-  public async initiateContainerHeartBeat(containerData: IContainerData) {
+  public async initiateContainerHeartBeat(shouldUseNotebookStates: boolean, containerData: IContainerData) {
     if (this.containerHealthHandler) {
       clearTimeout(this.containerHealthHandler);
     }
-    await this.getContainerHealth(Notebook.containerStatusHeartbeatDelayMs, containerData);
+    await this.getContainerHealth(shouldUseNotebookStates, Notebook.containerStatusHeartbeatDelayMs, containerData);
   }
 
-  private scheduleContainerHeartbeat(delayMs: number, containerData: IContainerData): void {
+  private scheduleContainerHeartbeat(
+    shouldUseNotebookStates: boolean,
+    delayMs: number,
+    containerData: IContainerData
+  ): void {
     this.containerHealthHandler = setTimeout(async () => {
-      await this.getContainerHealth(delayMs, containerData);
+      await this.getContainerHealth(shouldUseNotebookStates, delayMs, containerData);
     }, delayMs);
   }
 
-  private async getContainerStatusAsync(containerData: IContainerData): Promise<ContainerInfo> {
+  private async getContainerStatusAsync(
+    shouldUseNotebookStates: boolean,
+    containerData: IContainerData
+  ): Promise<ContainerInfo> {
     try {
       const runContainerStatusAsync = async () => {
         const response = await window.fetch(
@@ -136,14 +144,17 @@ export class PhoenixClient {
             dataExplorerArea: Areas.Notebook,
             message: getErrorMessage(error),
           });
-          useNotebook.getState().resetContainerConnection(connectionStatus);
+          shouldUseNotebookStates
-          useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+            ? useNotebook.getState().resetContainerConnection(connectionStatus)
-          useDialog
+            : useQueryCopilot.getState().resetContainerConnection();
-            .getState()
+          shouldUseNotebookStates && useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
-            .showOkModalDialog(
+          shouldUseNotebookStates &&
-              "Disconnected",
+            useDialog
-              "Disconnected from temporary workspace. Please click on connect button to connect to temporary workspace."
+              .getState()
-            );
+              .showOkModalDialog(
+                "Disconnected",
+                "Disconnected from temporary workspace. Please click on connect button to connect to temporary workspace."
+              );
           throw new AbortError(response.statusText);
         } else if (response?.status === HttpStatusCodes.Forbidden) {
           const validationMessage = this.ConvertToForbiddenErrorString(await response.json());
@@ -163,8 +174,10 @@ export class PhoenixClient {
       const connectionStatus: ContainerConnectionInfo = {
         status: ConnectionStatusType.Failed,
       };
-      useNotebook.getState().resetContainerConnection(connectionStatus);
+      shouldUseNotebookStates
-      useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
+        ? useNotebook.getState().resetContainerConnection(connectionStatus)
+        : useQueryCopilot.getState().resetContainerConnection();
+      shouldUseNotebookStates && useNotebook.getState().setIsRefreshed(!useNotebook.getState().isRefreshed);
       return {
         durationLeftInMinutes: undefined,
         phoenixServerInfo: undefined,
@@ -173,11 +186,17 @@ export class PhoenixClient {
     }
   }
 
-  private async getContainerHealth(delayMs: number, containerData: IContainerData) {
+  private async getContainerHealth(shouldUseNotebookStates: boolean, delayMs: number, containerData: IContainerData) {
-    const containerInfo = await this.getContainerStatusAsync(containerData);
+    const containerInfo = await this.getContainerStatusAsync(shouldUseNotebookStates, containerData);
-    useNotebook.getState().setContainerStatus(containerInfo);
+    shouldUseNotebookStates
-    if (useNotebook.getState().containerStatus?.status === ContainerStatusType.Active) {
+      ? useNotebook.getState().setContainerStatus(containerInfo)
-      this.scheduleContainerHeartbeat(delayMs, containerData);
+      : useQueryCopilot.getState().setContainerStatus(containerInfo);
+
+    const containerStatus = shouldUseNotebookStates
+      ? useNotebook.getState().containerStatus?.status
+      : useQueryCopilot.getState().containerStatus?.status;
+    if (containerStatus === ContainerStatusType.Active) {
+      this.scheduleContainerHeartbeat(shouldUseNotebookStates, delayMs, containerData);
     }
   }
 

src/hooks/useQueryCopilot.ts

@@ -2,7 +2,10 @@ import { MinimalQueryIterator } from "Common/IteratorUtilities";
 import { QueryResults } from "Contracts/ViewModels";
 import { CopilotMessage } from "Explorer/QueryCopilot/Shared/QueryCopilotInterfaces";
 import { guid } from "Explorer/Tables/Utilities";
+import { useTabs } from "hooks/useTabs";
 import create, { UseStore } from "zustand";
+import * as DataModels from "../Contracts/DataModels";
+import { ContainerInfo } from "../Contracts/DataModels";
 
 export interface QueryCopilotState {
   generatedQuery: string;
@@ -32,10 +35,12 @@ export interface QueryCopilotState {
   showWelcomeSidebar: boolean;
   showCopilotSidebar: boolean;
   chatMessages: CopilotMessage[];
-  shouldAllocateContainer: boolean;
   shouldIncludeInMessages: boolean;
   showExplanationBubble: boolean;
   showQueryExplanation: boolean;
+  notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo;
+  containerStatus: ContainerInfo;
+  isAllocatingContainer: boolean;
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) => void;
   closeFeedbackModal: () => void;
@@ -65,11 +70,14 @@ export interface QueryCopilotState {
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => void;
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => void;
   setChatMessages: (chatMessages: CopilotMessage[]) => void;
-  setShouldAllocateContainer: (shouldAllocateContainer: boolean) => void;
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => void;
   setShowExplanationBubble: (showExplanationBubble: boolean) => void;
   setShowQueryExplanation: (showQueryExplanation: boolean) => void;
+  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) => void;
+  setContainerStatus: (containerStatus: ContainerInfo) => void;
+  setIsAllocatingContainer: (isAllocatingContainer: boolean) => void;
 
+  resetContainerConnection: () => void;
   resetQueryCopilotStates: () => void;
 }
 
@@ -103,10 +111,20 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   showWelcomeSidebar: true,
   showCopilotSidebar: false,
   chatMessages: [],
-  shouldAllocateContainer: true,
   shouldIncludeInMessages: true,
   showExplanationBubble: false,
   showQueryExplanation: false,
+  notebookServerInfo: {
+    notebookServerEndpoint: undefined,
+    authToken: undefined,
+    forwardingId: undefined,
+  },
+  containerStatus: {
+    status: undefined,
+    durationLeftInMinutes: undefined,
+    phoenixServerInfo: undefined,
+  },
+  isAllocatingContainer: false,
 
   openFeedbackModal: (generatedQuery: string, likeQuery: boolean, userPrompt: string) =>
     set({ generatedQuery, likeQuery, userPrompt, showFeedbackModal: true }),
@@ -138,10 +156,24 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
   setShowWelcomeSidebar: (showWelcomeSidebar: boolean) => set({ showWelcomeSidebar }),
   setShowCopilotSidebar: (showCopilotSidebar: boolean) => set({ showCopilotSidebar }),
   setChatMessages: (chatMessages: CopilotMessage[]) => set({ chatMessages }),
-  setShouldAllocateContainer: (shouldAllocateContainer: boolean) => set({ shouldAllocateContainer }),
   setShouldIncludeInMessages: (shouldIncludeInMessages: boolean) => set({ shouldIncludeInMessages }),
   setShowExplanationBubble: (showExplanationBubble: boolean) => set({ showExplanationBubble }),
   setShowQueryExplanation: (showQueryExplanation: boolean) => set({ showQueryExplanation }),
+  setNotebookServerInfo: (notebookServerInfo: DataModels.NotebookWorkspaceConnectionInfo) =>
+    set({ notebookServerInfo }),
+  setContainerStatus: (containerStatus: ContainerInfo) => set({ containerStatus }),
+  setIsAllocatingContainer: (isAllocatingContainer: boolean) => set({ isAllocatingContainer }),
+
+  resetContainerConnection: (): void => {
+    useTabs.getState().closeAllNotebookTabs(true);
+    useQueryCopilot.getState().setNotebookServerInfo(undefined);
+    useQueryCopilot.getState().setIsAllocatingContainer(false);
+    useQueryCopilot.getState().setContainerStatus({
+      status: undefined,
+      durationLeftInMinutes: undefined,
+      phoenixServerInfo: undefined,
+    });
+  },
 
   resetQueryCopilotStates: () => {
     set((state) => ({
@@ -172,10 +204,20 @@ export const useQueryCopilot: QueryCopilotStore = create((set) => ({
       wasCopilotUsed: false,
       showCopilotSidebar: false,
       chatMessages: [],
-      shouldAllocateContainer: true,
       shouldIncludeInMessages: true,
       showExplanationBubble: false,
       showQueryExplanation: false,
+      notebookServerInfo: {
+        notebookServerEndpoint: undefined,
+        authToken: undefined,
+        forwardingId: undefined,
+      },
+      containerStatus: {
+        status: undefined,
+        durationLeftInMinutes: undefined,
+        phoenixServerInfo: undefined,
+      },
+      isAllocatingContainer: false,
     }));
   },
 }));