chore(participant): refactor chat participant state VSCODE-583

src/connectionController.ts

@@ -77,7 +77,7 @@
 }
 
 // Exported for testing.
 export function getNotifyDeviceFlowForConnectionAttempt(
   connectionOptions: ConnectionOptions
 ) {
   const isOIDCConnectionAttempt = isOIDCAuth(
@@ -91,7 +91,7 @@
     | undefined;
 
   if (isOIDCConnectionAttempt) {
     notifyDeviceFlow = ({
       verificationUrl,
       userCode,
     }: {
@@ -381,7 +381,7 @@
           ...cloneDeep(connectionOptions.oidc),
           openBrowser: browserAuthCommand
             ? { command: browserAuthCommand }
             : async ({ signal, url }) => {
                 try {
                   await openLink(url);
                 } catch (err) {
@@ -469,7 +469,7 @@
   }
 
   // Used to re-authenticate with OIDC.
   async _reauthenticationHandler() {
     const removeConfirmationResponse =
       await vscode.window.showInformationMessage(
         'You need to re-authenticate to the database in order to continue.',
@@ -482,7 +482,7 @@
     }
   }
 
   private async onConnectSuccess({
     connectionInfo,
     dataService,
   }: {
@@ -1054,15 +1054,15 @@
     );
 
     if (!selectedQuickPickItem) {
-      return true;
+      return false;
     }
 
     if (selectedQuickPickItem.data.type === NewConnectionType.NEW_CONNECTION) {
       return this.connectWithURI();
     }
 
     if (!selectedQuickPickItem.data.connectionId) {
-      return true;
+      return false;
     }
 
     const { successfullyConnected } = await this.connectWithConnectionId(

src/mdbExtensionController.ts

@@ -42,6 +42,7 @@ import WebviewController from './views/webviewController';
 import { createIdFactory, generateId } from './utils/objectIdHelper';
 import { ConnectionStorage } from './storage/connectionStorage';
 import type StreamProcessorTreeItem from './explorer/streamProcessorTreeItem';
+import type { RunParticipantQueryCommandArgs } from './participant/participant';
 import ParticipantController from './participant/participant';
 
 // This class is the top-level controller for our extension.
@@ -291,36 +292,47 @@ export default class MDBExtensionController implements vscode.Disposable {
     // ------ CHAT PARTICIPANT ------ //
     this.registerParticipantCommand(
       EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND,
-      () => {
+      ({ runnableContent }: RunParticipantQueryCommandArgs) => {
         return this._playgroundController.createPlaygroundFromParticipantQuery({
-          text:
-            this._participantController._chatResult?.metadata
-              ?.responseContent || '',
+          text: runnableContent,
         });
       }
     );
     this.registerParticipantCommand(
       EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY,
-      () => {
+      ({ runnableContent }: RunParticipantQueryCommandArgs) => {
         return this._playgroundController.evaluateParticipantQuery(
-          this._participantController._chatResult?.metadata?.responseContent ||
-            ''
+          runnableContent
         );
       }
     );
     this.registerCommand(
       EXTENSION_COMMANDS.CONNECT_WITH_PARTICIPANT,
-      (id: string) => this._participantController.connectWithParticipant(id)
+      (id?: string) =>
+        this._participantController.connectWithParticipant(
+          id ? decodeURIComponent(id) : id
+        )
     );
     this.registerCommand(
       EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT,
-      (name: string) =>
-        this._participantController.selectDatabaseWithParticipant(name)
+      (_data: string) => {
+        const data = JSON.parse(decodeURIComponent(_data));
+        return this._participantController.selectDatabaseWithParticipant({
+          chatId: data.chatId,
+          databaseName: data.databaseName,
+        });
+      }
     );
     this.registerCommand(
       EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
-      (name: string) =>
-        this._participantController.selectCollectionWithParticipant(name)
+      (_data: string) => {
+        const data = JSON.parse(decodeURIComponent(_data));
+        return this._participantController.selectCollectionWithParticipant({
+          chatId: data.chatId,
+          databaseName: data.databaseName,
+          collectionName: data.collectionName,
+        });
+      }
     );
   };
 

src/participant/chatMetadata.ts

@@ -0,0 +1,36 @@
+import * as vscode from 'vscode';
+import { v4 as uuidv4 } from 'uuid';
+
+export type ChatMetadata = {
+  databaseName?: string;
+  collectionName?: string;
+};
+
+export class ChatMetadataStore {
+  _chats: { [chatId: string]: ChatMetadata } = {};
+
+  constructor() {}
+
+  setChatMetadata(chatId: string, metadata: ChatMetadata): void {
+    this._chats[chatId] = metadata;
+  }
+
+  getChatMetadata(chatId: string): ChatMetadata | undefined {
+    return this._chats[chatId];
+  }
+
+  static getChatIdFromHistoryOrNewChatId(
+    history: ReadonlyArray<vscode.ChatRequestTurn | vscode.ChatResponseTurn>
+  ): string {
+    for (const historyItem of history) {
+      if (
+        historyItem instanceof vscode.ChatResponseTurn &&
+        historyItem.result?.metadata?.chatId
+      ) {
+        return historyItem.result.metadata.chatId;
+      }
+    }
+
+    return uuidv4();
+  }
+}

src/participant/constants.ts

@@ -1,2 +1,63 @@
+import type * as vscode from 'vscode';
+import { ChatMetadataStore } from './chatMetadata';
+
 export const CHAT_PARTICIPANT_ID = 'mongodb.participant';
 export const CHAT_PARTICIPANT_MODEL = 'gpt-4o';
+
+export class NamespaceRequestChatResult implements vscode.ChatResult {
+  readonly metadata: {
+    chatId: string;
+    intent: 'askForNamespace';
+    databaseName?: string | undefined;
+    collectionName?: string | undefined;
+  };
+
+  constructor({
+    databaseName,
+    collectionName,
+    history,
+  }: {
+    history: ReadonlyArray<vscode.ChatRequestTurn | vscode.ChatResponseTurn>;
+    databaseName: string | undefined;
+    collectionName: string | undefined;
+  }) {
+    this.metadata = {
+      chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(history),
+      intent: 'askForNamespace',
+      databaseName,
+      collectionName,
+    };
+  }
+}
+
+export class EmptyRequestChatResult implements vscode.ChatResult {
+  readonly metadata: {
+    chatId: string;
+    intent: 'emptyRequest';
+  };
+
+  constructor(
+    history: ReadonlyArray<vscode.ChatRequestTurn | vscode.ChatResponseTurn>
+  ) {
+    this.metadata = {
+      chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(history),
+      intent: 'emptyRequest',
+    };
+  }
+}
+
+export class AskToConnectChatResult implements vscode.ChatResult {
+  readonly metadata: {
+    chatId: string;
+    intent: 'askToConnect';
+  };
+
+  constructor(
+    history: ReadonlyArray<vscode.ChatRequestTurn | vscode.ChatResponseTurn>
+  ) {
+    this.metadata = {
+      chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(history),
+      intent: 'askToConnect',
+    };
+  }
+}

src/participant/markdown.ts

@@ -0,0 +1,32 @@
+import * as vscode from 'vscode';
+
+export function createMarkdownLink({
+  commandId,
+  data,
+  name,
+}: {
+  commandId: string;
+  data?:
+    | {
+        [field: string]: any;
+      }
+    | string;
+  name: string;
+}): vscode.MarkdownString {
+  const encodedData = data
+    ? encodeURIComponent(
+        `["${
+          typeof data === 'string'
+            ? data
+            : encodeURIComponent(JSON.stringify(data))
+        }"]`
+      )
+    : undefined;
+  const commandQueryString = data ? `?${encodedData}` : '';
+  const connName = new vscode.MarkdownString(
+    `- <a href="command:${commandId}${commandQueryString}">${name}</a>\n`
+  );
+  connName.supportHtml = true;
+  connName.isTrusted = { enabledCommands: [commandId] };
+  return connName;
+}