feat(playground): add Generate Query with Copilot code lens in playgrounds VSCODE-650 (#881)

Co-authored-by: Alena Khineika <[email protected]>

Author: gagik

src/commands/index.ts

@@ -73,6 +73,8 @@ enum EXTENSION_COMMANDS {
 
   // Chat participant.
   OPEN_PARTICIPANT_CODE_IN_PLAYGROUND = 'mdb.openParticipantCodeInPlayground',
+  SEND_MESSAGE_TO_PARTICIPANT = 'mdb.sendMessageToParticipant',
+  SEND_MESSAGE_TO_PARTICIPANT_FROM_INPUT = 'mdb.sendMessageToParticipantFromInput',
   RUN_PARTICIPANT_CODE = 'mdb.runParticipantCode',
   CONNECT_WITH_PARTICIPANT = 'mdb.connectWithParticipant',
   SELECT_DATABASE_WITH_PARTICIPANT = 'mdb.selectDatabaseWithParticipant',

src/connectionController.ts

@@ -201,7 +201,7 @@ export default class ConnectionController {
           ignoreFocusOut: true,
           placeHolder:
             'e.g. mongodb+srv://username:[email protected]/admin',
-          prompt: 'Enter your connection string (SRV or standard)',
+          prompt: 'Enter your SRV or standard connection string',
           validateInput: (uri: string) => {
             if (
               !uri.startsWith('mongodb://') &&

src/documentSource.ts

@@ -2,4 +2,5 @@ export enum DocumentSource {
   DOCUMENT_SOURCE_TREEVIEW = 'treeview',
   DOCUMENT_SOURCE_PLAYGROUND = 'playground',
   DOCUMENT_SOURCE_COLLECTIONVIEW = 'collectionview',
+  DOCUMENT_SOURCE_QUERY_WITH_COPILOT_CODELENS = 'query with copilot codelens',
 }

src/editors/activeConnectionCodeLensProvider.ts

@@ -24,7 +24,7 @@ export default class ActiveConnectionCodeLensProvider
       this._onDidChangeCodeLenses.fire();
     });
 
-    this._activeConnectionChangedHandler = () => {
+    this._activeConnectionChangedHandler = (): void => {
       this._onDidChangeCodeLenses.fire();
     };
     this._connectionController.addEventListener(
@@ -50,10 +50,10 @@ export default class ActiveConnectionCodeLensProvider
         ? getDBFromConnectionString(connectionString)
         : null;
       message = defaultDB
-        ? `Currently connected to ${this._connectionController.getActiveConnectionName()} with default database ${defaultDB}. Click here to change connection.`
-        : `Currently connected to ${this._connectionController.getActiveConnectionName()}. Click here to change connection.`;
+        ? `Connected to ${this._connectionController.getActiveConnectionName()} with default database ${defaultDB}`
+        : `Connected to ${this._connectionController.getActiveConnectionName()}`;
     } else {
-      message = 'Disconnected. Click here to connect.';
+      message = 'Connect';
     }
 
     codeLens.command = {
@@ -65,7 +65,7 @@ export default class ActiveConnectionCodeLensProvider
     return [codeLens];
   }
 
-  deactivate() {
+  deactivate(): void {
     this._connectionController.removeEventListener(
       DataServiceEventTypes.ACTIVE_CONNECTION_CHANGED,
       this._activeConnectionChangedHandler

src/editors/editorsController.ts

@@ -32,6 +32,7 @@ import type PlaygroundResultProvider from './playgroundResultProvider';
 import { PLAYGROUND_RESULT_SCHEME } from './playgroundResultProvider';
 import { StatusView } from '../views';
 import type TelemetryService from '../telemetry/telemetryService';
+import type { QueryWithCopilotCodeLensProvider } from './queryWithCopilotCodeLensProvider';
 
 const log = createLogger('editors controller');
 
@@ -102,6 +103,7 @@ export default class EditorsController {
   _exportToLanguageCodeLensProvider: ExportToLanguageCodeLensProvider;
   _editDocumentCodeLensProvider: EditDocumentCodeLensProvider;
   _collectionDocumentsCodeLensProvider: CollectionDocumentsCodeLensProvider;
+  _queryWithCopilotCodeLensProvider: QueryWithCopilotCodeLensProvider;
 
   constructor({
     context,
@@ -115,6 +117,7 @@ export default class EditorsController {
     playgroundSelectionCodeActionProvider,
     playgroundDiagnosticsCodeActionProvider,
     editDocumentCodeLensProvider,
+    queryWithCopilotCodeLensProvider,
   }: {
     context: vscode.ExtensionContext;
     connectionController: ConnectionController;
@@ -127,6 +130,7 @@ export default class EditorsController {
     playgroundSelectionCodeActionProvider: PlaygroundSelectionCodeActionProvider;
     playgroundDiagnosticsCodeActionProvider: PlaygroundDiagnosticsCodeActionProvider;
     editDocumentCodeLensProvider: EditDocumentCodeLensProvider;
+    queryWithCopilotCodeLensProvider: QueryWithCopilotCodeLensProvider;
   }) {
     this._connectionController = connectionController;
     this._playgroundController = playgroundController;
@@ -160,6 +164,7 @@ export default class EditorsController {
       playgroundSelectionCodeActionProvider;
     this._playgroundDiagnosticsCodeActionProvider =
       playgroundDiagnosticsCodeActionProvider;
+    this._queryWithCopilotCodeLensProvider = queryWithCopilotCodeLensProvider;
 
     vscode.workspace.onDidCloseTextDocument((e) => {
       const uriParams = new URLSearchParams(e.uri.query);
@@ -410,6 +415,10 @@ export default class EditorsController {
       )
     );
     this._context.subscriptions.push(
+      vscode.languages.registerCodeLensProvider(
+        { language: 'javascript' },
+        this._queryWithCopilotCodeLensProvider
+      ),
       vscode.languages.registerCodeLensProvider(
         { language: 'javascript' },
         this._activeConnectionCodeLensProvider

src/editors/queryWithCopilotCodeLensProvider.ts

@@ -0,0 +1,47 @@
+import * as vscode from 'vscode';
+import EXTENSION_COMMANDS from '../commands';
+import type { SendMessageToParticipantFromInputOptions } from '../participant/participantTypes';
+import { isPlayground } from '../utils/playground';
+import { COPILOT_EXTENSION_ID } from '../participant/constants';
+import { DocumentSource } from '../documentSource';
+
+export class QueryWithCopilotCodeLensProvider
+  implements vscode.CodeLensProvider
+{
+  constructor() {}
+
+  readonly onDidChangeCodeLenses: vscode.Event<void> =
+    vscode.extensions.onDidChange;
+
+  provideCodeLenses(document: vscode.TextDocument): vscode.CodeLens[] {
+    if (!isPlayground(document.uri)) {
+      return [];
+    }
+
+    // We can only detect whether a user has the Copilot extension active
+    // but not whether it has an active subscription.
+    const hasCopilotChatActive =
+      vscode.extensions.getExtension(COPILOT_EXTENSION_ID)?.isActive === true;
+
+    if (!hasCopilotChatActive) {
+      return [];
+    }
+
+    const options: SendMessageToParticipantFromInputOptions = {
+      prompt: 'Describe the query you would like to generate',
+      placeHolder:
+        'e.g. Find the document in sample_mflix.users with the name of Kayden Washington',
+      messagePrefix: '/query',
+      isNewChat: true,
+      source: DocumentSource.DOCUMENT_SOURCE_QUERY_WITH_COPILOT_CODELENS,
+    };
+
+    return [
+      new vscode.CodeLens(new vscode.Range(0, 0, 0, 0), {
+        title: '✨ Generate query with MongoDB Copilot',
+        command: EXTENSION_COMMANDS.SEND_MESSAGE_TO_PARTICIPANT_FROM_INPUT,
+        arguments: [options],
+      }),
+    ];
+  }
+}

src/mdbExtensionController.ts

@@ -46,6 +46,11 @@ import type {
 } from './participant/participant';
 import ParticipantController from './participant/participant';
 import type { OpenSchemaCommandArgs } from './participant/prompts/schema';
+import { QueryWithCopilotCodeLensProvider } from './editors/queryWithCopilotCodeLensProvider';
+import type {
+  SendMessageToParticipantOptions,
+  SendMessageToParticipantFromInputOptions,
+} from './participant/participantTypes';
 
 // This class is the top-level controller for our extension.
 // Commands which the extensions handles are defined in the function `activate`.
@@ -65,6 +70,7 @@ export default class MDBExtensionController implements vscode.Disposable {
   _telemetryService: TelemetryService;
   _languageServerController: LanguageServerController;
   _webviewController: WebviewController;
+  _queryWithCopilotCodeLensProvider: QueryWithCopilotCodeLensProvider;
   _playgroundResultProvider: PlaygroundResultProvider;
   _activeConnectionCodeLensProvider: ActiveConnectionCodeLensProvider;
   _editDocumentCodeLensProvider: EditDocumentCodeLensProvider;
@@ -105,6 +111,8 @@ export default class MDBExtensionController implements vscode.Disposable {
       this._connectionController,
       this._editDocumentCodeLensProvider
     );
+    this._queryWithCopilotCodeLensProvider =
+      new QueryWithCopilotCodeLensProvider();
     this._activeConnectionCodeLensProvider =
       new ActiveConnectionCodeLensProvider(this._connectionController);
     this._exportToLanguageCodeLensProvider =
@@ -143,6 +151,7 @@ export default class MDBExtensionController implements vscode.Disposable {
       playgroundDiagnosticsCodeActionProvider:
         this._playgroundDiagnosticsCodeActionProvider,
       editDocumentCodeLensProvider: this._editDocumentCodeLensProvider,
+      queryWithCopilotCodeLensProvider: this._queryWithCopilotCodeLensProvider,
     });
     this._webviewController = new WebviewController({
       connectionController: this._connectionController,
@@ -305,6 +314,22 @@ export default class MDBExtensionController implements vscode.Disposable {
         });
       }
     );
+    this.registerParticipantCommand(
+      EXTENSION_COMMANDS.SEND_MESSAGE_TO_PARTICIPANT,
+      async (options: SendMessageToParticipantOptions) => {
+        await this._participantController.sendMessageToParticipant(options);
+        return true;
+      }
+    );
+    this.registerParticipantCommand(
+      EXTENSION_COMMANDS.SEND_MESSAGE_TO_PARTICIPANT_FROM_INPUT,
+      async (options: SendMessageToParticipantFromInputOptions) => {
+        await this._participantController.sendMessageToParticipantFromInput(
+          options
+        );
+        return true;
+      }
+    );
     this.registerParticipantCommand(
       EXTENSION_COMMANDS.RUN_PARTICIPANT_CODE,
       ({ runnableContent }: RunParticipantCodeCommandArgs) => {
@@ -949,12 +974,9 @@ export default class MDBExtensionController implements vscode.Disposable {
 
     const copilot = vscode.extensions.getExtension('github.copilot-chat');
     if (result?.title === action) {
-      await vscode.commands.executeCommand('workbench.action.chat.newChat');
-      await vscode.commands.executeCommand(
-        'workbench.action.chat.clearHistory'
-      );
-      await vscode.commands.executeCommand('workbench.action.chat.open', {
-        query: '@MongoDB ',
+      await this._participantController.sendMessageToParticipant({
+        message: '',
+        isNewChat: true,
         isPartialQuery: true,
       });
       this._telemetryService.trackCopilotIntroductionClicked({

src/participant/constants.ts

@@ -3,6 +3,7 @@ import { ChatMetadataStore } from './chatMetadata';
 
 export const CHAT_PARTICIPANT_ID = 'mongodb.participant';
 export const CHAT_PARTICIPANT_MODEL = 'gpt-4o';
+export const COPILOT_EXTENSION_ID = 'GitHub.copilot';
 
 export type ParticipantResponseType =
   | 'query'

src/participant/participant.ts

@@ -49,6 +49,10 @@ import { ParticipantErrorTypes } from './participantErrorTypes';
 import type PlaygroundResultProvider from '../editors/playgroundResultProvider';
 import { isExportToLanguageResult } from '../types/playgroundType';
 import { PromptHistory } from './prompts/promptHistory';
+import type {
+  SendMessageToParticipantOptions,
+  SendMessageToParticipantFromInputOptions,
+} from './participantTypes';
 import { DEFAULT_EXPORT_TO_LANGUAGE_DRIVER_SYNTAX } from '../editors/exportToLanguageCodeLensProvider';
 import { EXPORT_TO_LANGUAGE_ALIASES } from '../editors/playgroundSelectionCodeActionProvider';
 
@@ -130,9 +134,51 @@ export default class ParticipantController {
    * in the chat. To work around this, we can write a message as the user, which will
    * trigger the chat handler and give us access to the model.
    */
-  writeChatMessageAsUser(message: string): Thenable<unknown> {
+  async sendMessageToParticipant(
+    options: SendMessageToParticipantOptions
+  ): Promise<unknown> {
+    const { message, isNewChat = false, isPartialQuery = false } = options;
+
+    if (isNewChat) {
+      await vscode.commands.executeCommand('workbench.action.chat.newChat');
+      await vscode.commands.executeCommand(
+        'workbench.action.chat.clearHistory'
+      );
+    }
+
     return vscode.commands.executeCommand('workbench.action.chat.open', {
       query: `@MongoDB ${message}`,
+      isPartialQuery,
+    });
+  }
+
+  async sendMessageToParticipantFromInput(
+    options: SendMessageToParticipantFromInputOptions
+  ): Promise<unknown> {
+    const {
+      messagePrefix = '',
+      isNewChat = false,
+      isPartialQuery = false,
+      source,
+      ...inputBoxOptions
+    } = options;
+
+    this._telemetryService.trackCopilotParticipantSubmittedFromInputBox({
+      source,
+    });
+
+    const message = await vscode.window.showInputBox({
+      ...inputBoxOptions,
+    });
+
+    if (message === undefined || message.trim() === '') {
+      return Promise.resolve();
+    }
+
+    return this.sendMessageToParticipant({
+      message: `${messagePrefix ? `${messagePrefix} ` : ''}${message}`,
+      isNewChat,
+      isPartialQuery,
     });
   }
 
@@ -460,9 +506,9 @@ export default class ParticipantController {
 
     const connectionName = this._connectionController.getActiveConnectionName();
 
-    return this.writeChatMessageAsUser(
-      `${command ? `${command} ` : ''}${connectionName}`
-    ) as Promise<boolean>;
+    return this.sendMessageToParticipant({
+      message: `${command ? `${command} ` : ''}${connectionName}`,
+    }) as Promise<boolean>;
   }
 
   getConnectionsTree(command: ParticipantCommand): vscode.MarkdownString[] {
@@ -501,7 +547,7 @@ export default class ParticipantController {
     const dataService = this._connectionController.getActiveDataService();
     if (!dataService) {
       // Run a blank command to get the user to connect first.
-      void this.writeChatMessageAsUser(command);
+      void this.sendMessageToParticipant({ message: command });
       return [];
     }
 
@@ -549,9 +595,9 @@ export default class ParticipantController {
       databaseName: databaseName,
     });
 
-    return this.writeChatMessageAsUser(
-      `${command} ${databaseName}`
-    ) as Promise<boolean>;
+    return this.sendMessageToParticipant({
+      message: `${command} ${databaseName}`,
+    }) as Promise<boolean>;
   }
 
   async getCollectionQuickPicks({
@@ -564,7 +610,7 @@ export default class ParticipantController {
     const dataService = this._connectionController.getActiveDataService();
     if (!dataService) {
       // Run a blank command to get the user to connect first.
-      void this.writeChatMessageAsUser(command);
+      void this.sendMessageToParticipant({ message: command });
       return [];
     }
 
@@ -625,9 +671,9 @@ export default class ParticipantController {
       databaseName: databaseName,
       collectionName: collectionName,
     });
-    return this.writeChatMessageAsUser(
-      `${command} ${collectionName}`
-    ) as Promise<boolean>;
+    return this.sendMessageToParticipant({
+      message: `${command} ${collectionName}`,
+    }) as Promise<boolean>;
   }
 
   renderDatabasesTree({

src/participant/participantTypes.ts

@@ -0,0 +1,14 @@
+import type * as vscode from 'vscode';
+import type { DocumentSource } from '../documentSource';
+
+export type SendMessageToParticipantOptions = {
+  message: string;
+  isNewChat?: boolean;
+  isPartialQuery?: boolean;
+};
+
+export type SendMessageToParticipantFromInputOptions = {
+  messagePrefix?: string;
+  source?: DocumentSource;
+} & Omit<SendMessageToParticipantOptions, 'message'> &
+  vscode.InputBoxOptions;