From ff0c1cc8c13f10bad0b33516a6b8c4c1d4f2846c Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Thu, 31 Oct 2024 13:42:09 +0100
Subject: [PATCH 01/29] WIP

---
 .vscode/settings.json                         |   1 -
 src/participant/participant.ts                | 230 ++++++++++++-----
 src/test/suite/index.ts                       |   3 +-
 .../suite/participant/participant.test.ts     | 240 ++++++++++--------
 4 files changed, 303 insertions(+), 171 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 732bb7406..32baf6163 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,6 +1,5 @@
 // Place your settings in this file to overwrite default and user settings.
 {
-  "editor.formatOnSave": false,
   "editor.codeActionsOnSave": {
     "source.fixAll.eslint": "explicit"
   },
diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index ee5655467..44ac41374 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -44,6 +44,7 @@ import formatError from '../utils/formatError';
 import type { ModelInput } from './prompts/promptBase';
 import { processStreamWithIdentifiers } from './streamParsing';
 import type { PromptIntent } from './prompts/intent';
+import type { DataService } from 'mongodb-data-service';
 
 const log = createLogger('participant');
 
@@ -671,64 +672,49 @@ export default class ParticipantController {
     }
   }
 
-  async renderCollectionsTree({
+  renderCollectionsTree({
+    collections,
     command,
     context,
     databaseName,
     stream,
   }: {
+    collections: Awaited<ReturnType<DataService['listCollections']>>;
     command: ParticipantCommand;
     databaseName: string;
     context: vscode.ChatContext;
     stream: vscode.ChatResponseStream;
-  }): Promise<void> {
-    const dataService = this._connectionController.getActiveDataService();
-    if (!dataService) {
-      return;
-    }
-
-    stream.push(
-      new vscode.ChatResponseProgressPart('Fetching collection names...')
+  }): void {
+    collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((coll) =>
+      stream.markdown(
+        createMarkdownLink({
+          commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
+          data: {
+            command,
+            chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+              context.history
+            ),
+            databaseName,
+            collectionName: coll.name,
+          },
+          name: coll.name,
+        })
+      )
     );
-
-    try {
-      const collections = await dataService.listCollections(databaseName);
-      collections.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((coll) =>
-        stream.markdown(
-          createMarkdownLink({
-            commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
-            data: {
-              command,
-              chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-                context.history
-              ),
-              databaseName,
-              collectionName: coll.name,
-            },
-            name: coll.name,
-          })
-        )
+    if (collections.length > MAX_MARKDOWN_LIST_LENGTH) {
+      stream.markdown(
+        createMarkdownLink({
+          commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
+          data: {
+            command,
+            chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+              context.history
+            ),
+            databaseName,
+          },
+          name: 'Show more',
+        })
       );
-      if (collections.length > MAX_MARKDOWN_LIST_LENGTH) {
-        stream.markdown(
-          createMarkdownLink({
-            commandId: EXTENSION_COMMANDS.SELECT_COLLECTION_WITH_PARTICIPANT,
-            data: {
-              command,
-              chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-                context.history
-              ),
-              databaseName,
-            },
-            name: 'Show more',
-          })
-        );
-      }
-    } catch (error) {
-      log.error('Unable to fetch collections:', error);
-
-      // Users can always do this manually when asked to provide a collection name.
-      return;
     }
   }
 
@@ -811,7 +797,67 @@ export default class ParticipantController {
     };
   }
 
-  async _askForNamespace({
+  async _getCollections({
+    stream,
+    databaseName,
+  }: {
+    stream: vscode.ChatResponseStream;
+    databaseName: string;
+  }): Promise<ReturnType<DataService['listCollections']> | undefined> {
+    stream.push(
+      new vscode.ChatResponseProgressPart('Fetching collection names...')
+    );
+
+    const dataService = this._connectionController.getActiveDataService();
+
+    if (!dataService) {
+      return;
+    }
+    try {
+      return await dataService.listCollections(databaseName);
+    } catch (error) {
+      log.error('Unable to fetch collections:', error);
+      return;
+    }
+  }
+
+  /** Gets the collection name if there is only one collection.
+   *  Otherwise returns undefined and asks the user to select the collection. */
+  async _getOrAskForCollectionName({
+    context,
+    databaseName,
+    stream,
+    command,
+  }: {
+    command: ParticipantCommand;
+    context: vscode.ChatContext;
+    databaseName: string;
+    stream: vscode.ChatResponseStream;
+  }): Promise<string | undefined> {
+    const collections = await this._getCollections({ stream, databaseName });
+
+    if (collections !== undefined) {
+      if (collections.length === 1) {
+        return collections[0].name;
+      }
+
+      stream.markdown(
+        `Which collection would you like to use within ${databaseName}?\n\n`
+      );
+
+      this.renderCollectionsTree({
+        collections,
+        command,
+        databaseName,
+        context,
+        stream,
+      });
+    }
+
+    return;
+  }
+
+  async _askForDatabaseName({
     command,
     context,
     databaseName,
@@ -838,16 +884,6 @@ export default class ParticipantController {
         context,
         stream,
       });
-    } else if (!collectionName) {
-      stream.markdown(
-        `Which collection would you like to use within ${databaseName}?\n\n`
-      );
-      await this.renderCollectionsTree({
-        command,
-        databaseName,
-        context,
-        stream,
-      });
     }
 
     return namespaceRequestChatResult({
@@ -1011,12 +1047,33 @@ export default class ParticipantController {
     // we re-ask the question.
     const databaseName = lastMessage.metadata.databaseName;
     if (databaseName) {
+      const collections = await this._getCollections({
+        stream,
+        databaseName,
+      });
+
+      if (!collections) {
+        return namespaceRequestChatResult({
+          databaseName,
+          collectionName: undefined,
+          history: context.history,
+        });
+      }
+
+      // If there is only 1 collection in the database, we can pick it automatically
+      if (collections.length === 1) {
+        stream.markdown(Prompts.generic.getEmptyRequestResponse());
+        return emptyRequestChatResult(context.history);
+      }
+
       stream.markdown(
         vscode.l10n.t(
           'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
         )
       );
-      await this.renderCollectionsTree({
+
+      this.renderCollectionsTree({
+        collections,
         command,
         databaseName,
         context,
@@ -1043,6 +1100,7 @@ export default class ParticipantController {
   }
 
   // @MongoDB /schema
+  // eslint-disable-next-line complexity
   async handleSchemaRequest(
     request: vscode.ChatRequest,
     context: vscode.ChatContext,
@@ -1068,14 +1126,16 @@ export default class ParticipantController {
       });
     }
 
-    const { databaseName, collectionName } = await this._getNamespaceFromChat({
+    const namespace = await this._getNamespaceFromChat({
       request,
       context,
       token,
     });
+    const { databaseName } = namespace;
+    let { collectionName } = namespace;
 
-    if (!databaseName || !collectionName) {
-      return await this._askForNamespace({
+    if (!databaseName) {
+      return await this._askForDatabaseName({
         command: '/schema',
         context,
         databaseName,
@@ -1084,6 +1144,25 @@ export default class ParticipantController {
       });
     }
 
+    if (!collectionName) {
+      collectionName = await this._getOrAskForCollectionName({
+        command: '/schema',
+        context,
+        databaseName,
+        stream,
+      });
+
+      // If the collection name could not get automatically selected,
+      // then the user has been prompted for it.
+      if (!collectionName) {
+        return namespaceRequestChatResult({
+          databaseName,
+          collectionName: undefined,
+          history: context.history,
+        });
+      }
+    }
+
     if (token.isCancellationRequested) {
       return this._handleCancelledRequest({
         context,
@@ -1194,13 +1273,16 @@ export default class ParticipantController {
     // First we ask the model to parse for the database and collection name.
     // If they exist, we can then use them in our final completion.
     // When they don't exist we ask the user for them.
-    const { databaseName, collectionName } = await this._getNamespaceFromChat({
+    const namespace = await this._getNamespaceFromChat({
       request,
       context,
       token,
     });
-    if (!databaseName || !collectionName) {
-      return await this._askForNamespace({
+    const { databaseName } = namespace;
+    let { collectionName } = namespace;
+
+    if (!databaseName) {
+      return await this._askForDatabaseName({
         command: '/query',
         context,
         databaseName,
@@ -1208,6 +1290,24 @@ export default class ParticipantController {
         stream,
       });
     }
+    if (!collectionName) {
+      collectionName = await this._getOrAskForCollectionName({
+        command: '/query',
+        context,
+        databaseName,
+        stream,
+      });
+
+      // If the collection name could not get automatically selected,
+      // then the user has been prompted for it.
+      if (!collectionName) {
+        return namespaceRequestChatResult({
+          databaseName,
+          collectionName: undefined,
+          history: context.history,
+        });
+      }
+    }
 
     if (token.isCancellationRequested) {
       return this._handleCancelledRequest({
diff --git a/src/test/suite/index.ts b/src/test/suite/index.ts
index 9833e8055..cd75fe177 100644
--- a/src/test/suite/index.ts
+++ b/src/test/suite/index.ts
@@ -9,8 +9,8 @@ import { mdbTestExtension } from './stubbableMdbExtension';
 
 export async function run(): Promise<void> {
   const reporterOptions = {
-    spec: '-',
     'mocha-junit-reporter': path.join(__dirname, './test-results.xml'),
+    spec: '-',
   };
 
   // Create the mocha tester.
@@ -19,6 +19,7 @@ export async function run(): Promise<void> {
     reporterOptions,
     ui: 'tdd',
     color: true,
+    grep: 'Participant',
   });
 
   const testsRoot = path.join(__dirname, '..');
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 61eeb5c6f..9b1924465 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -404,6 +404,19 @@ suite('Participant Controller Test Suite', function () {
 
   suite('when connected', function () {
     let sampleStub;
+    let mockedCollectionList = [
+      { name: 'collOne' },
+      { name: 'notifications' },
+      { name: 'products' },
+      { name: 'orders' },
+      { name: 'categories' },
+      { name: 'invoices' },
+      { name: 'transactions' },
+      { name: 'logs' },
+      { name: 'messages' },
+      { name: 'sessions' },
+      { name: 'feedback' },
+    ];
 
     beforeEach(function () {
       sampleStub = sinon.stub();
@@ -426,20 +439,7 @@ suite('Participant Controller Test Suite', function () {
                 { name: 'analytics' },
                 { name: '123' },
               ]),
-            listCollections: () =>
-              Promise.resolve([
-                { name: 'collOne' },
-                { name: 'notifications' },
-                { name: 'products' },
-                { name: 'orders' },
-                { name: 'categories' },
-                { name: 'invoices' },
-                { name: 'transactions' },
-                { name: 'logs' },
-                { name: 'messages' },
-                { name: 'sessions' },
-                { name: 'feedback' },
-              ]),
+            listCollections: () => Promise.resolve(mockedCollectionList),
             getMongoClientConnectionOptions: () => ({
               url: TEST_DATABASE_URI,
               options: {},
@@ -1298,107 +1298,139 @@ suite('Participant Controller Test Suite', function () {
             });
           });
 
-          test('handles empty collection name', async function () {
+          suite('with an empty collection name', function () {
             const chatRequestMock = {
               prompt: '',
               command: 'query',
               references: [],
             };
-            chatContextStub = {
-              history: [
-                Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
-                  prompt: 'find all docs by a name example',
-                  command: 'query',
-                  references: [],
-                  participant: CHAT_PARTICIPANT_ID,
-                }),
-                Object.assign(
-                  Object.create(vscode.ChatResponseTurn.prototype),
-                  {
-                    participant: CHAT_PARTICIPANT_ID,
-                    response: [
-                      {
-                        value: {
-                          value:
-                            'Which database would you like to query within this database?',
-                        } as vscode.MarkdownString,
+
+            beforeEach(function () {
+              chatContextStub = {
+                history: [
+                  Object.assign(
+                    Object.create(vscode.ChatRequestTurn.prototype),
+                    {
+                      prompt: 'find all docs by a name example',
+                      command: 'query',
+                      references: [],
+                      participant: CHAT_PARTICIPANT_ID,
+                    }
+                  ),
+                  Object.assign(
+                    Object.create(vscode.ChatResponseTurn.prototype),
+                    {
+                      participant: CHAT_PARTICIPANT_ID,
+                      response: [
+                        {
+                          value: {
+                            value:
+                              'Which database would you like to query within this database?',
+                          } as vscode.MarkdownString,
+                        },
+                      ],
+                      command: 'query',
+                      result: {
+                        metadata: {
+                          intent: 'askForNamespace',
+                        },
                       },
-                    ],
-                    command: 'query',
-                    result: {
-                      metadata: {
-                        intent: 'askForNamespace',
+                    }
+                  ),
+                  Object.assign(
+                    Object.create(vscode.ChatRequestTurn.prototype),
+                    {
+                      prompt: 'dbOne',
+                      command: 'query',
+                      references: [],
+                      participant: CHAT_PARTICIPANT_ID,
+                    }
+                  ),
+                  Object.assign(
+                    Object.create(vscode.ChatResponseTurn.prototype),
+                    {
+                      participant: CHAT_PARTICIPANT_ID,
+                      response: [
+                        {
+                          value: {
+                            value:
+                              'Which collection would you like to query within dbOne?',
+                          } as vscode.MarkdownString,
+                        },
+                      ],
+                      command: 'query',
+                      result: {
+                        metadata: {
+                          intent: 'askForNamespace',
+                          databaseName: 'dbOne',
+                          collectionName: undefined,
+                          chatId: 'pineapple',
+                        },
                       },
-                    },
+                    }
+                  ),
+                ],
+              };
+            });
+
+            test('prompts to select collection name', async function () {
+              const chatResult = await invokeChatHandler(chatRequestMock);
+
+              const emptyMessage = chatStreamStub.markdown.getCall(0).args[0];
+              expect(emptyMessage).to.include(
+                'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
+              );
+              const listCollsMessage =
+                chatStreamStub.markdown.getCall(1).args[0];
+              expect(listCollsMessage.value).to.include(
+                `- [collOne](command:mdb.selectCollectionWithParticipant?${encodeStringify(
+                  {
+                    command: '/query',
+                    chatId: 'pineapple',
+                    databaseName: 'dbOne',
+                    collectionName: 'collOne',
                   }
-                ),
-                Object.assign(Object.create(vscode.ChatRequestTurn.prototype), {
-                  prompt: 'dbOne',
-                  command: 'query',
-                  references: [],
-                  participant: CHAT_PARTICIPANT_ID,
-                }),
-                Object.assign(
-                  Object.create(vscode.ChatResponseTurn.prototype),
+                )})`
+              );
+              const showMoreCollsMessage =
+                chatStreamStub.markdown.getCall(11).args[0];
+              expect(showMoreCollsMessage.value).to.include(
+                `- [Show more](command:mdb.selectCollectionWithParticipant?${encodeStringify(
                   {
-                    participant: CHAT_PARTICIPANT_ID,
-                    response: [
-                      {
-                        value: {
-                          value:
-                            'Which collection would you like to query within dbOne?',
-                        } as vscode.MarkdownString,
-                      },
-                    ],
-                    command: 'query',
-                    result: {
-                      metadata: {
-                        intent: 'askForNamespace',
-                        databaseName: 'dbOne',
-                        collectionName: undefined,
-                        chatId: 'pineapple',
-                      },
-                    },
+                    command: '/query',
+                    chatId: 'pineapple',
+                    databaseName: 'dbOne',
                   }
-                ),
-              ],
-            };
-            const chatResult = await invokeChatHandler(chatRequestMock);
+                )})`
+              );
+              expect({
+                ...chatResult?.metadata,
+                chatId: undefined,
+              }).to.deep.equal({
+                intent: 'askForNamespace',
+                collectionName: undefined,
+                databaseName: 'dbOne',
+                chatId: undefined,
+              });
+            });
 
-            const emptyMessage = chatStreamStub.markdown.getCall(0).args[0];
-            expect(emptyMessage).to.include(
-              'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
-            );
-            const listCollsMessage = chatStreamStub.markdown.getCall(1).args[0];
-            expect(listCollsMessage.value).to.include(
-              `- [collOne](command:mdb.selectCollectionWithParticipant?${encodeStringify(
-                {
-                  command: '/query',
-                  chatId: 'pineapple',
-                  databaseName: 'dbOne',
-                  collectionName: 'collOne',
-                }
-              )})`
-            );
-            const showMoreCollsMessage =
-              chatStreamStub.markdown.getCall(11).args[0];
-            expect(showMoreCollsMessage.value).to.include(
-              `- [Show more](command:mdb.selectCollectionWithParticipant?${encodeStringify(
-                {
-                  command: '/query',
-                  chatId: 'pineapple',
-                  databaseName: 'dbOne',
-                }
-              )})`
-            );
-            expect({
-              ...chatResult?.metadata,
-              chatId: undefined,
-            }).to.deep.equal({
-              intent: 'askForNamespace',
-              collectionName: undefined,
-              databaseName: 'dbOne',
-              chatId: undefined,
+            test('shows the empty request message if there is only 1 collection in the database', async function () {
+              mockedCollectionList = [{ name: 'onlyColl' }];
+
+              const chatResult = await invokeChatHandler(chatRequestMock);
+
+              const responses = chatStreamStub.markdown
+                .getCalls()
+                .map((call) => call.args[0]);
+              expect(responses.length).equals(1);
+              expect(responses[0]).to.include(
+                'Ask anything about MongoDB, from writing queries to questions about your cluster.'
+              );
+
+              expect(chatResult).deep.equals({
+                intent: 'emptyRequest',
+                chatId: 'pineapple',
+              });
             });
           });
         });

From f9dd536b66e465e7bcde1a9b5d604501de779587 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 12:53:18 +0100
Subject: [PATCH 02/29] Move around dependencies

---
 .../suite/participant/participant.test.ts     | 53 ++++++++++++++++++-
 1 file changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 9b1924465..0046614e4 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1427,7 +1427,7 @@ suite('Participant Controller Test Suite', function () {
                 'Ask anything about MongoDB, from writing queries to questions about your cluster.'
               );
 
-              expect(chatResult).deep.equals({
+              expect(chatResult?.metadata).deep.equals({
                 intent: 'emptyRequest',
                 chatId: 'pineapple',
               });
@@ -1539,6 +1539,57 @@ suite('Participant Controller Test Suite', function () {
               'see previous messages'
             );
           });
+
+          suite('with an empty collection name', function () {
+            beforeEach(function () {
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: 'dbOne',
+                  collectionName: undefined,
+                })
+              );
+            });
+
+            test('if collection name is not provided but there is only one collection in the database, it automatically picks it', async function () {
+              mockedCollectionList = [{ name: 'onlyOneColl' }];
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(chatResult?.metadata).deep.equals({
+                chatId: 'test-chat-id',
+                intent: 'schema',
+              });
+
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.args[0]
+              ).to.include({
+                collectionName: 'onlyOneColl',
+              });
+
+              const responseContents = chatStreamStub.markdown
+                .getCalls()
+                .map((call) => call.args[0]);
+              expect(responseContents.length).equals(1);
+              expect(responseContents[0]).equals(
+                'Unable to generate a schema from the collection, no documents found.'
+              );
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+            });
+          });
         });
 
         suite(

From 09c74149c637e9170ba301646beff07681b7e8a7 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 12:53:45 +0100
Subject: [PATCH 03/29] Remove grep

---
 src/test/suite/index.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/test/suite/index.ts b/src/test/suite/index.ts
index cd75fe177..ac5d52aa2 100644
--- a/src/test/suite/index.ts
+++ b/src/test/suite/index.ts
@@ -19,7 +19,6 @@ export async function run(): Promise<void> {
     reporterOptions,
     ui: 'tdd',
     color: true,
-    grep: 'Participant',
   });
 
   const testsRoot = path.join(__dirname, '..');

From aada8e1a5687836cb4e40b2837cabf248a43fbad Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 13:15:03 +0100
Subject: [PATCH 04/29] Use firstCall

---
 src/test/suite/participant/participant.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 0046614e4..cf32ec6b4 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1575,7 +1575,7 @@ suite('Participant Controller Test Suite', function () {
               });
 
               expect(
-                fetchCollectionSchemaAndSampleDocumentsSpy.args[0]
+                fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
               ).to.include({
                 collectionName: 'onlyOneColl',
               });

From eb2c7dca166c978e95ef47b935f88a25b3623284 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 13:37:56 +0100
Subject: [PATCH 05/29] Add test filtering

---
 .vscode/launch.json     | 10 ++++++++++
 CONTRIBUTING.md         | 16 ++++++++++++++++
 package.json            |  4 ++--
 src/test/suite/index.ts |  1 +
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index 9d596e5fd..533cbba63 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -52,10 +52,20 @@
         "--extensionDevelopmentPath=${workspaceFolder}",
         "--extensionTestsPath=${workspaceFolder}/out/test/suite"
       ],
+      "env": {
+        "MOCHA_GREP": "${input:mochaGrep}"
+      },
       "outFiles": ["${workspaceFolder}/out/**/*.js"],
       "preLaunchTask": "npm: compile:extension",
     }
   ],
+  "inputs": [
+    {
+      "id": "mochaGrep",
+      "type": "promptString",
+      "description": "Enter an optional grep filter to run specific tests. Leave blank for all.",
+    }
+  ],
   "compounds": [
     {
       "name": "Extension + Server Inspector",
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 0e3168a2c..5b2af7c22 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -35,6 +35,22 @@ npm run watch
 
 2. Inside of [VS Code Insiders](https://code.visualstudio.com/insiders/) open this directory and press `F5` to begin debugging the extension. This should launch a new VSCode window which is running the extension.
 
+### Running Tests
+
+#### Using the VSCode debugger
+
+You can launch a debugging task for tests inside VSCode with the **"Run Tests"** task. There you also can specify an optional test filter.
+
+#### Using command line
+
+You can run tests using command line along with an optional `MOCHA_GREP` environment variable to apply a grep filter on tests to run.
+
+```shell
+MOCHA_GREP="Participant .* prompt builders" npm test
+```
+
+It may be quicker to be more specific and use `npm run test-extension` or `npm run test-webview`.
+
 ### Using Proposed API
 
 The vscode extension will occasionally need to use [proposed API](https://code.visualstudio.com/api/advanced-topics/using-proposed-api) that haven't been promoted to stable yet. To enable an API proposal, add it to the `enabledApiProposals` section in `package.json`, then run `cd src/vscode-dts && npx @vscode/dts dev` to install the type definitions for the API you want to enable.
diff --git a/package.json b/package.json
index 2a35a4935..09198e5d5 100644
--- a/package.json
+++ b/package.json
@@ -53,8 +53,8 @@
     "pretest": "npm run compile",
     "test": "npm run test-webview && npm run test-extension",
     "test-extension": "cross-env NODE_OPTIONS=--no-force-async-hooks-checks xvfb-maybe node ./out/test/runTest.js",
-    "test-webview": "mocha -r ts-node/register --file ./src/test/setup-webview.ts src/test/suite/views/webview-app/**/*.test.tsx",
-    "ai-accuracy-tests": "env TS_NODE_FILES=true mocha -r ts-node/register --file ./src/test/ai-accuracy-tests/test-setup.ts ./src/test/ai-accuracy-tests/ai-accuracy-tests.ts",
+    "test-webview": "mocha -r ts-node/register --grep=${MOCHA_GREP} --file ./src/test/setup-webview.ts src/test/suite/views/webview-app/**/*.test.tsx",
+    "ai-accuracy-tests": "env TS_NODE_FILES=true mocha -r ts-node/register --grep=${MOCHA_GREP} --file ./src/test/ai-accuracy-tests/test-setup.ts ./src/test/ai-accuracy-tests/ai-accuracy-tests.ts",
     "analyze-bundle": "webpack --mode production --analyze",
     "vscode:prepublish": "npm run clean && npm run compile:constants && npm run compile:resources && webpack --mode production",
     "check": "npm run lint && npm run depcheck",
diff --git a/src/test/suite/index.ts b/src/test/suite/index.ts
index 9833e8055..b9c7f2f42 100644
--- a/src/test/suite/index.ts
+++ b/src/test/suite/index.ts
@@ -19,6 +19,7 @@ export async function run(): Promise<void> {
     reporterOptions,
     ui: 'tdd',
     color: true,
+    grep: process.env.MOCHA_GREP,
   });
 
   const testsRoot = path.join(__dirname, '..');

From bcd260bbae9fe8f607daa93832c94020179884f0 Mon Sep 17 00:00:00 2001
From: Gagik Amaryan <me@gagik.co>
Date: Fri, 1 Nov 2024 13:42:43 +0100
Subject: [PATCH 06/29] Update CONTRIBUTING.md

---
 CONTRIBUTING.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 5b2af7c22..114b53f98 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -49,7 +49,7 @@ You can run tests using command line along with an optional `MOCHA_GREP` environ
 MOCHA_GREP="Participant .* prompt builders" npm test
 ```
 
-It may be quicker to be more specific and use `npm run test-extension` or `npm run test-webview`.
+It may be quicker to be more specific and use `npm run test-extension` or `npm run test-webview` after compiling.
 
 ### Using Proposed API
 

From a9beef11920316e7e44e521c0553a2b3104b26f3 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 13:45:09 +0100
Subject: [PATCH 07/29] Escape the environment variable

---
 package.json | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/package.json b/package.json
index 09198e5d5..a8a10c25e 100644
--- a/package.json
+++ b/package.json
@@ -53,8 +53,8 @@
     "pretest": "npm run compile",
     "test": "npm run test-webview && npm run test-extension",
     "test-extension": "cross-env NODE_OPTIONS=--no-force-async-hooks-checks xvfb-maybe node ./out/test/runTest.js",
-    "test-webview": "mocha -r ts-node/register --grep=${MOCHA_GREP} --file ./src/test/setup-webview.ts src/test/suite/views/webview-app/**/*.test.tsx",
-    "ai-accuracy-tests": "env TS_NODE_FILES=true mocha -r ts-node/register --grep=${MOCHA_GREP} --file ./src/test/ai-accuracy-tests/test-setup.ts ./src/test/ai-accuracy-tests/ai-accuracy-tests.ts",
+    "test-webview": "mocha -r ts-node/register --grep=\"${MOCHA_GREP}\" --file ./src/test/setup-webview.ts src/test/suite/views/webview-app/**/*.test.tsx",
+    "ai-accuracy-tests": "env TS_NODE_FILES=true mocha -r ts-node/register --grep=\"${MOCHA_GREP}\" --file ./src/test/ai-accuracy-tests/test-setup.ts ./src/test/ai-accuracy-tests/ai-accuracy-tests.ts",
     "analyze-bundle": "webpack --mode production --analyze",
     "vscode:prepublish": "npm run clean && npm run compile:constants && npm run compile:resources && webpack --mode production",
     "check": "npm run lint && npm run depcheck",

From 33c995e3ad9ea6870d32c8cb75f8df576442929e Mon Sep 17 00:00:00 2001
From: Gagik Amaryan <me@gagik.co>
Date: Fri, 1 Nov 2024 14:18:53 +0100
Subject: [PATCH 08/29] Fix wording

---
 CONTRIBUTING.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 114b53f98..87aa3e686 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -39,7 +39,7 @@ npm run watch
 
 #### Using the VSCode debugger
 
-You can launch a debugging task for tests inside VSCode with the **"Run Tests"** task. There you also can specify an optional test filter.
+You can launch a debugging task for tests inside VSCode with the **"Run Tests"** task. There you can also specify an optional test filter.
 
 #### Using command line
 

From 858219385e241d6ebf25c44c5ef789965c4fb3b1 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 14:25:08 +0100
Subject: [PATCH 09/29] Add schema tests

---
 .../suite/participant/participant.test.ts     | 45 ++++++++++++++-----
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index cf32ec6b4..7dc7dc326 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1552,7 +1552,7 @@ suite('Participant Controller Test Suite', function () {
               );
             });
 
-            test('if collection name is not provided but there is only one collection in the database, it automatically picks it', async function () {
+            test('collection name gets picked automatically if there is only 1', async function () {
               mockedCollectionList = [{ name: 'onlyOneColl' }];
               const renderCollectionsTreeSpy = sinon.spy(
                 testParticipantController,
@@ -1569,10 +1569,7 @@ suite('Participant Controller Test Suite', function () {
                 references: [],
               });
 
-              expect(chatResult?.metadata).deep.equals({
-                chatId: 'test-chat-id',
-                intent: 'schema',
-              });
+              expect(renderCollectionsTreeSpy.called).to.be.false;
 
               expect(
                 fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
@@ -1580,14 +1577,38 @@ suite('Participant Controller Test Suite', function () {
                 collectionName: 'onlyOneColl',
               });
 
-              const responseContents = chatStreamStub.markdown
-                .getCalls()
-                .map((call) => call.args[0]);
-              expect(responseContents.length).equals(1);
-              expect(responseContents[0]).equals(
-                'Unable to generate a schema from the collection, no documents found.'
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'schema',
+              });
+            });
+
+            test('prompts for collection name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
               );
-              expect(renderCollectionsTreeSpy.called).to.be.false;
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.called
+              ).to.be.false;
+
+              expect(chatResult?.metadata).deep.equals({
+                intent: 'askForNamespace',
+                chatId: testChatId,
+                databaseName: 'dbOne',
+              });
             });
           });
         });

From a8bc30d997905f677efadda751f79029da049940 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 15:01:49 +0100
Subject: [PATCH 10/29] align tests and use a stub

---
 .vscode/settings.json                         |  1 +
 .../suite/participant/participant.test.ts     | 79 +++++++++++--------
 2 files changed, 45 insertions(+), 35 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 32baf6163..732bb7406 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,5 +1,6 @@
 // Place your settings in this file to overwrite default and user settings.
 {
+  "editor.formatOnSave": false,
   "editor.codeActionsOnSave": {
     "source.fixAll.eslint": "explicit"
   },
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 7dc7dc326..5d7960115 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -404,22 +404,26 @@ suite('Participant Controller Test Suite', function () {
 
   suite('when connected', function () {
     let sampleStub;
-    let mockedCollectionList = [
-      { name: 'collOne' },
-      { name: 'notifications' },
-      { name: 'products' },
-      { name: 'orders' },
-      { name: 'categories' },
-      { name: 'invoices' },
-      { name: 'transactions' },
-      { name: 'logs' },
-      { name: 'messages' },
-      { name: 'sessions' },
-      { name: 'feedback' },
-    ];
+    let listCollectionsStub;
 
     beforeEach(function () {
       sampleStub = sinon.stub();
+      listCollectionsStub = sinon
+        .stub()
+        .resolves([
+          { name: 'collOne' },
+          { name: 'notifications' },
+          { name: 'products' },
+          { name: 'orders' },
+          { name: 'categories' },
+          { name: 'invoices' },
+          { name: 'transactions' },
+          { name: 'logs' },
+          { name: 'messages' },
+          { name: 'sessions' },
+          { name: 'feedback' },
+        ]);
+
       sinon.replace(
         testParticipantController._connectionController,
         'getActiveDataService',
@@ -439,7 +443,7 @@ suite('Participant Controller Test Suite', function () {
                 { name: 'analytics' },
                 { name: '123' },
               ]),
-            listCollections: () => Promise.resolve(mockedCollectionList),
+            listCollections: listCollectionsStub,
             getMongoClientConnectionOptions: () => ({
               url: TEST_DATABASE_URI,
               options: {},
@@ -1373,7 +1377,30 @@ suite('Participant Controller Test Suite', function () {
               };
             });
 
-            test('prompts to select collection name', async function () {
+            afterEach(function () {
+              sinon.restore();
+            });
+
+            test('collection name gets picked automatically if there is only', async function () {
+              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
+
+              const chatResult = await invokeChatHandler(chatRequestMock);
+
+              const responses = chatStreamStub.markdown
+                .getCalls()
+                .map((call) => call.args[0]);
+              expect(responses.length).equals(1);
+              expect(responses[0]).to.include(
+                'Ask anything about MongoDB, from writing queries to questions about your cluster.'
+              );
+
+              expect(chatResult?.metadata).deep.equals({
+                intent: 'emptyRequest',
+                chatId: 'pineapple',
+              });
+            });
+
+            test('prompts for collection name if there are multiple available', async function () {
               const chatResult = await invokeChatHandler(chatRequestMock);
 
               const emptyMessage = chatStreamStub.markdown.getCall(0).args[0];
@@ -1413,25 +1440,6 @@ suite('Participant Controller Test Suite', function () {
                 chatId: undefined,
               });
             });
-
-            test('shows the empty request message if there is only 1 collection in the database', async function () {
-              mockedCollectionList = [{ name: 'onlyColl' }];
-
-              const chatResult = await invokeChatHandler(chatRequestMock);
-
-              const responses = chatStreamStub.markdown
-                .getCalls()
-                .map((call) => call.args[0]);
-              expect(responses.length).equals(1);
-              expect(responses[0]).to.include(
-                'Ask anything about MongoDB, from writing queries to questions about your cluster.'
-              );
-
-              expect(chatResult?.metadata).deep.equals({
-                intent: 'emptyRequest',
-                chatId: 'pineapple',
-              });
-            });
           });
         });
       });
@@ -1553,7 +1561,7 @@ suite('Participant Controller Test Suite', function () {
             });
 
             test('collection name gets picked automatically if there is only 1', async function () {
-              mockedCollectionList = [{ name: 'onlyOneColl' }];
+              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
               const renderCollectionsTreeSpy = sinon.spy(
                 testParticipantController,
                 'renderCollectionsTree'
@@ -1608,6 +1616,7 @@ suite('Participant Controller Test Suite', function () {
                 intent: 'askForNamespace',
                 chatId: testChatId,
                 databaseName: 'dbOne',
+                collectionName: undefined,
               });
             });
           });

From 5ddc7fb0abdb513a1f6aa2cc6ea54a7dda1dd202 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 1 Nov 2024 15:30:14 +0100
Subject: [PATCH 11/29] Add saving to metadata

---
 src/participant/participant.ts | 18 ++++++++++++++++++
 src/test/suite/index.ts        |  2 +-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 44ac41374..c84126f9b 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -1161,6 +1161,15 @@ export default class ParticipantController {
           history: context.history,
         });
       }
+
+      // Save the collection name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        collectionName,
+      });
     }
 
     if (token.isCancellationRequested) {
@@ -1307,6 +1316,15 @@ export default class ParticipantController {
           history: context.history,
         });
       }
+
+      // Save the collection name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        collectionName,
+      });
     }
 
     if (token.isCancellationRequested) {
diff --git a/src/test/suite/index.ts b/src/test/suite/index.ts
index ac5d52aa2..9833e8055 100644
--- a/src/test/suite/index.ts
+++ b/src/test/suite/index.ts
@@ -9,8 +9,8 @@ import { mdbTestExtension } from './stubbableMdbExtension';
 
 export async function run(): Promise<void> {
   const reporterOptions = {
-    'mocha-junit-reporter': path.join(__dirname, './test-results.xml'),
     spec: '-',
+    'mocha-junit-reporter': path.join(__dirname, './test-results.xml'),
   };
 
   // Create the mocha tester.

From 06435e8b8ef92a369d1e8abc47fa6805d3324ac0 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Sun, 3 Nov 2024 22:42:45 +0100
Subject: [PATCH 12/29] Move things

---
 src/participant/participant.ts | 182 +++++++++++++++++++--------------
 1 file changed, 106 insertions(+), 76 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index c84126f9b..6a709e47a 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -613,62 +613,49 @@ export default class ParticipantController {
     ) as Promise<boolean>;
   }
 
-  async renderDatabasesTree({
+  renderDatabasesTree({
     command,
     context,
     stream,
+    databases,
   }: {
     command: ParticipantCommand;
     context: vscode.ChatContext;
     stream: vscode.ChatResponseStream;
-  }): Promise<void> {
-    const dataService = this._connectionController.getActiveDataService();
-    if (!dataService) {
-      return;
-    }
-
-    stream.push(
-      new vscode.ChatResponseProgressPart('Fetching database names...')
+    databases: {
+      _id: string;
+      name: string;
+    }[];
+  }): void {
+    databases.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((db) =>
+      stream.markdown(
+        createMarkdownLink({
+          commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT,
+          data: {
+            command,
+            chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+              context.history
+            ),
+            databaseName: db.name,
+          },
+          name: db.name,
+        })
+      )
     );
 
-    try {
-      const databases = await dataService.listDatabases({
-        nameOnly: true,
-      });
-      databases.slice(0, MAX_MARKDOWN_LIST_LENGTH).forEach((db) =>
-        stream.markdown(
-          createMarkdownLink({
-            commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT,
-            data: {
-              command,
-              chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-                context.history
-              ),
-              databaseName: db.name,
-            },
-            name: db.name,
-          })
-        )
+    if (databases.length > MAX_MARKDOWN_LIST_LENGTH) {
+      stream.markdown(
+        createMarkdownLink({
+          data: {
+            command,
+            chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+              context.history
+            ),
+          },
+          commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT,
+          name: 'Show more',
+        })
       );
-      if (databases.length > MAX_MARKDOWN_LIST_LENGTH) {
-        stream.markdown(
-          createMarkdownLink({
-            data: {
-              command,
-              chatId: ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-                context.history
-              ),
-            },
-            commandId: EXTENSION_COMMANDS.SELECT_DATABASE_WITH_PARTICIPANT,
-            name: 'Show more',
-          })
-        );
-      }
-    } catch (error) {
-      log.error('Unable to fetch databases:', error);
-
-      // Users can always do this manually when asked to provide a database name.
-      return;
     }
   }
 
@@ -797,6 +784,37 @@ export default class ParticipantController {
     };
   }
 
+  async _getDatabases({
+    stream,
+  }: {
+    stream: vscode.ChatResponseStream;
+  }): Promise<
+    | {
+        _id: string;
+        name: string;
+      }[]
+    | undefined
+  > {
+    stream.push(
+      new vscode.ChatResponseProgressPart('Fetching database names...')
+    );
+    const dataService = this._connectionController.getActiveDataService();
+    if (!dataService) {
+      return;
+    }
+
+    try {
+      const databases = await dataService.listDatabases({
+        nameOnly: true,
+      });
+      return databases;
+    } catch (error) {
+      log.error('Unable to fetch databases:', error);
+
+      return;
+    }
+  }
+
   async _getCollections({
     stream,
     databaseName,
@@ -857,40 +875,44 @@ export default class ParticipantController {
     return;
   }
 
-  async _askForDatabaseName({
+  /** Gets the database name if there is only one collection.
+   *  Otherwise returns undefined and asks the user to select the database. */
+  async _getOrAskForDatabaseName({
     command,
     context,
-    databaseName,
-    collectionName,
     stream,
   }: {
     command: ParticipantCommand;
     context: vscode.ChatContext;
-    databaseName: string | undefined;
-    collectionName: string | undefined;
     stream: vscode.ChatResponseStream;
-  }): Promise<ChatResult> {
+  }): Promise<string | undefined> {
     // If no database or collection name is found in the user prompt,
     // we retrieve the available namespaces from the current connection.
     // Users can then select a value by clicking on an item in the list.
-    if (!databaseName) {
-      stream.markdown(
-        `What is the name of the database you would like${
-          command === '/query' ? ' this query' : ''
-        } to run against?\n\n`
-      );
-      await this.renderDatabasesTree({
-        command,
-        context,
-        stream,
-      });
+    stream.markdown(
+      `What is the name of the database you would like${
+        command === '/query' ? ' this query' : ''
+      } to run against?\n\n`
+    );
+    const databases = await this._getDatabases({ stream });
+
+    if (databases === undefined || databases.length === 0) {
+      log.error('No databases found');
+      return;
     }
 
-    return namespaceRequestChatResult({
-      databaseName,
-      collectionName,
-      history: context.history,
+    if (databases.length === 1) {
+      return databases[0].name;
+    }
+
+    this.renderDatabasesTree({
+      databases,
+      command,
+      context,
+      stream,
     });
+
+    return;
   }
 
   _doesLastMessageAskForNamespace(
@@ -1045,7 +1067,7 @@ export default class ParticipantController {
 
     // When the last message was asking for a database or collection name,
     // we re-ask the question.
-    const databaseName = lastMessage.metadata.databaseName;
+    let databaseName = lastMessage.metadata.databaseName;
     if (databaseName) {
       const collections = await this._getCollections({
         stream,
@@ -1085,10 +1107,11 @@ export default class ParticipantController {
           'Please select a database by either clicking on an item in the list or typing the name manually in the chat.'
         )
       );
-      await this.renderDatabasesTree({
+
+      databaseName = await this._getOrAskForDatabaseName({
         command,
-        context,
         stream,
+        context,
       });
     }
 
@@ -1131,17 +1154,24 @@ export default class ParticipantController {
       context,
       token,
     });
-    const { databaseName } = namespace;
-    let { collectionName } = namespace;
+    let { databaseName, collectionName } = namespace;
 
     if (!databaseName) {
-      return await this._askForDatabaseName({
+      databaseName = await this._getOrAskForDatabaseName({
         command: '/schema',
         context,
-        databaseName,
-        collectionName,
         stream,
       });
+
+      // If the database name could not get automatically selected,
+      // then the user has been prompted for it.
+      if (!databaseName) {
+        return namespaceRequestChatResult({
+          databaseName,
+          collectionName,
+          history: context.history,
+        });
+      }
     }
 
     if (!collectionName) {
@@ -1157,7 +1187,7 @@ export default class ParticipantController {
       if (!collectionName) {
         return namespaceRequestChatResult({
           databaseName,
-          collectionName: undefined,
+          collectionName,
           history: context.history,
         });
       }

From 5a5817198956214112b6c25b912681537cd85a4e Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Sun, 3 Nov 2024 22:56:05 +0100
Subject: [PATCH 13/29] Better org

---
 src/participant/participant.ts | 160 ++++++++++++++++++---------------
 1 file changed, 90 insertions(+), 70 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 6a709e47a..19d4d69c5 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -784,6 +784,78 @@ export default class ParticipantController {
     };
   }
 
+  /** Helper which either automatically picks and returns missing parts of the namespace (if any)
+   *  or prompts the user to pick the missing namespace.
+   */
+  async _getOrAskForMissingNamespace({
+    databaseName,
+    collectionName,
+    context,
+    stream,
+    command,
+  }: {
+    databaseName: string | undefined;
+    collectionName: string | undefined;
+    context: vscode.ChatContext;
+    stream: vscode.ChatResponseStream;
+    command: ParticipantCommand;
+  }): Promise<{
+    databaseName: string | undefined;
+    collectionName: string | undefined;
+  }> {
+    if (!databaseName) {
+      databaseName = await this._getOrAskForDatabaseName({
+        command,
+        context,
+        stream,
+      });
+
+      // If the database name could not get automatically selected,
+      // then the user has been prompted for it instead.
+      if (!databaseName) {
+        return { databaseName, collectionName };
+      }
+
+      // Save the database name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        databaseName,
+      });
+    }
+
+    if (!collectionName) {
+      collectionName = await this._getOrAskForCollectionName({
+        command: '/schema',
+        context,
+        databaseName,
+        stream,
+      });
+
+      // If the collection name could not get automatically selected,
+      // then the user has been prompted for it instead.
+      if (!collectionName) {
+        return {
+          databaseName,
+          collectionName,
+        };
+      }
+
+      // Save the collection name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        collectionName,
+      });
+    }
+
+    return { collectionName, databaseName };
+  }
+
   async _getDatabases({
     stream,
   }: {
@@ -1154,51 +1226,21 @@ export default class ParticipantController {
       context,
       token,
     });
-    let { databaseName, collectionName } = namespace;
-
-    if (!databaseName) {
-      databaseName = await this._getOrAskForDatabaseName({
-        command: '/schema',
+    const { databaseName, collectionName } =
+      await this._getOrAskForMissingNamespace({
+        ...namespace,
         context,
         stream,
-      });
-
-      // If the database name could not get automatically selected,
-      // then the user has been prompted for it.
-      if (!databaseName) {
-        return namespaceRequestChatResult({
-          databaseName,
-          collectionName,
-          history: context.history,
-        });
-      }
-    }
-
-    if (!collectionName) {
-      collectionName = await this._getOrAskForCollectionName({
         command: '/schema',
-        context,
-        databaseName,
-        stream,
       });
 
-      // If the collection name could not get automatically selected,
-      // then the user has been prompted for it.
-      if (!collectionName) {
-        return namespaceRequestChatResult({
-          databaseName,
-          collectionName,
-          history: context.history,
-        });
-      }
-
-      // Save the collection name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
+    // If either the database or collection name could not be automatically picked
+    // then the user has be prompted to select one manually.
+    if (databaseName === undefined || collectionName === undefined) {
+      return namespaceRequestChatResult({
+        databaseName,
         collectionName,
+        history: context.history,
       });
     }
 
@@ -1317,43 +1359,21 @@ export default class ParticipantController {
       context,
       token,
     });
-    const { databaseName } = namespace;
-    let { collectionName } = namespace;
-
-    if (!databaseName) {
-      return await this._askForDatabaseName({
-        command: '/query',
+    const { databaseName, collectionName } =
+      await this._getOrAskForMissingNamespace({
+        ...namespace,
         context,
-        databaseName,
-        collectionName,
         stream,
-      });
-    }
-    if (!collectionName) {
-      collectionName = await this._getOrAskForCollectionName({
         command: '/query',
-        context,
-        databaseName,
-        stream,
       });
 
-      // If the collection name could not get automatically selected,
-      // then the user has been prompted for it.
-      if (!collectionName) {
-        return namespaceRequestChatResult({
-          databaseName,
-          collectionName: undefined,
-          history: context.history,
-        });
-      }
-
-      // Save the collection name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
+    // If either the database or collection name could not be automatically picked
+    // then the user has be prompted to select one manually.
+    if (databaseName === undefined || collectionName === undefined) {
+      return namespaceRequestChatResult({
+        databaseName,
         collectionName,
+        history: context.history,
       });
     }
 

From f23e19ff7c17cd5f7d1b384dad179e9aa91d56a5 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Mon, 4 Nov 2024 00:25:10 +0100
Subject: [PATCH 14/29] simplify tests and picking logic

---
 src/participant/participant.ts                | 167 ++++-----
 .../suite/participant/participant.test.ts     | 348 +++++++++++-------
 2 files changed, 298 insertions(+), 217 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 19d4d69c5..5701e0d64 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -784,78 +784,6 @@ export default class ParticipantController {
     };
   }
 
-  /** Helper which either automatically picks and returns missing parts of the namespace (if any)
-   *  or prompts the user to pick the missing namespace.
-   */
-  async _getOrAskForMissingNamespace({
-    databaseName,
-    collectionName,
-    context,
-    stream,
-    command,
-  }: {
-    databaseName: string | undefined;
-    collectionName: string | undefined;
-    context: vscode.ChatContext;
-    stream: vscode.ChatResponseStream;
-    command: ParticipantCommand;
-  }): Promise<{
-    databaseName: string | undefined;
-    collectionName: string | undefined;
-  }> {
-    if (!databaseName) {
-      databaseName = await this._getOrAskForDatabaseName({
-        command,
-        context,
-        stream,
-      });
-
-      // If the database name could not get automatically selected,
-      // then the user has been prompted for it instead.
-      if (!databaseName) {
-        return { databaseName, collectionName };
-      }
-
-      // Save the database name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
-        databaseName,
-      });
-    }
-
-    if (!collectionName) {
-      collectionName = await this._getOrAskForCollectionName({
-        command: '/schema',
-        context,
-        databaseName,
-        stream,
-      });
-
-      // If the collection name could not get automatically selected,
-      // then the user has been prompted for it instead.
-      if (!collectionName) {
-        return {
-          databaseName,
-          collectionName,
-        };
-      }
-
-      // Save the collection name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
-        collectionName,
-      });
-    }
-
-    return { collectionName, databaseName };
-  }
-
   async _getDatabases({
     stream,
   }: {
@@ -958,14 +886,6 @@ export default class ParticipantController {
     context: vscode.ChatContext;
     stream: vscode.ChatResponseStream;
   }): Promise<string | undefined> {
-    // If no database or collection name is found in the user prompt,
-    // we retrieve the available namespaces from the current connection.
-    // Users can then select a value by clicking on an item in the list.
-    stream.markdown(
-      `What is the name of the database you would like${
-        command === '/query' ? ' this query' : ''
-      } to run against?\n\n`
-    );
     const databases = await this._getDatabases({ stream });
 
     if (databases === undefined || databases.length === 0) {
@@ -977,6 +897,15 @@ export default class ParticipantController {
       return databases[0].name;
     }
 
+    // If no database or collection name is found in the user prompt,
+    // we retrieve the available namespaces from the current connection.
+    // Users can then select a value by clicking on an item in the list.
+    stream.markdown(
+      `What is the name of the database you would like${
+        command === '/query' ? ' this query' : ''
+      } to run against?\n\n`
+    );
+
     this.renderDatabasesTree({
       databases,
       command,
@@ -987,6 +916,78 @@ export default class ParticipantController {
     return;
   }
 
+  /** Helper which either automatically picks and returns missing parts of the namespace (if any)
+   *  or prompts the user to pick the missing namespace.
+   */
+  async _getOrAskForMissingNamespace({
+    databaseName,
+    collectionName,
+    context,
+    stream,
+    command,
+  }: {
+    databaseName: string | undefined;
+    collectionName: string | undefined;
+    context: vscode.ChatContext;
+    stream: vscode.ChatResponseStream;
+    command: ParticipantCommand;
+  }): Promise<{
+    databaseName: string | undefined;
+    collectionName: string | undefined;
+  }> {
+    if (!databaseName) {
+      databaseName = await this._getOrAskForDatabaseName({
+        command,
+        context,
+        stream,
+      });
+
+      // If the database name could not get automatically selected,
+      // then the user has been prompted for it instead.
+      if (!databaseName) {
+        return { databaseName, collectionName };
+      }
+
+      // Save the database name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        databaseName,
+      });
+    }
+
+    if (!collectionName) {
+      collectionName = await this._getOrAskForCollectionName({
+        command,
+        context,
+        databaseName,
+        stream,
+      });
+
+      // If the collection name could not get automatically selected,
+      // then the user has been prompted for it instead.
+      if (!collectionName) {
+        return {
+          databaseName,
+          collectionName,
+        };
+      }
+
+      // Save the collection name in the metadata.
+      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+        context.history
+      );
+      this._chatMetadataStore.setChatMetadata(chatId, {
+        ...this._chatMetadataStore.getChatMetadata(chatId),
+        collectionName,
+      });
+    }
+
+    return { collectionName, databaseName };
+  }
+
   _doesLastMessageAskForNamespace(
     history: ReadonlyArray<vscode.ChatRequestTurn | vscode.ChatResponseTurn>
   ): boolean {
@@ -1154,12 +1155,6 @@ export default class ParticipantController {
         });
       }
 
-      // If there is only 1 collection in the database, we can pick it automatically
-      if (collections.length === 1) {
-        stream.markdown(Prompts.generic.getEmptyRequestResponse());
-        return emptyRequestChatResult(context.history);
-      }
-
       stream.markdown(
         vscode.l10n.t(
           'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 5d7960115..a418da8cf 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -405,9 +405,25 @@ suite('Participant Controller Test Suite', function () {
   suite('when connected', function () {
     let sampleStub;
     let listCollectionsStub;
+    let listDatabasesStub;
 
     beforeEach(function () {
       sampleStub = sinon.stub();
+      listDatabasesStub = sinon
+        .stub()
+        .resolves([
+          { name: 'dbOne' },
+          { name: 'customer' },
+          { name: 'inventory' },
+          { name: 'sales' },
+          { name: 'employee' },
+          { name: 'financialReports' },
+          { name: 'productCatalog' },
+          { name: 'projectTracker' },
+          { name: 'user' },
+          { name: 'analytics' },
+          { name: '123' },
+        ]);
       listCollectionsStub = sinon
         .stub()
         .resolves([
@@ -429,20 +445,7 @@ suite('Participant Controller Test Suite', function () {
         'getActiveDataService',
         () =>
           ({
-            listDatabases: () =>
-              Promise.resolve([
-                { name: 'dbOne' },
-                { name: 'customer' },
-                { name: 'inventory' },
-                { name: 'sales' },
-                { name: 'employee' },
-                { name: 'financialReports' },
-                { name: 'productCatalog' },
-                { name: 'projectTracker' },
-                { name: 'user' },
-                { name: 'analytics' },
-                { name: '123' },
-              ]),
+            listDatabases: listDatabasesStub,
             listCollections: listCollectionsStub,
             getMongoClientConnectionOptions: () => ({
               url: TEST_DATABASE_URI,
@@ -1019,7 +1022,7 @@ suite('Participant Controller Test Suite', function () {
           });
         });
 
-        suite('unknown namespace', function () {
+        suite('no namespace provided', function () {
           test('asks for a namespace and generates a query', async function () {
             const chatRequestMock = {
               prompt: 'find all docs by a name example',
@@ -1228,7 +1231,7 @@ suite('Participant Controller Test Suite', function () {
             });
           });
 
-          test('handles empty database name', async function () {
+          test('asks for the empty database name again if the last prompt was doing so', async function () {
             const chatRequestMock = {
               prompt: '',
               command: 'query',
@@ -1302,142 +1305,152 @@ suite('Participant Controller Test Suite', function () {
             });
           });
 
-          suite('with an empty collection name', function () {
-            const chatRequestMock = {
-              prompt: '',
-              command: 'query',
-              references: [],
-            };
-
+          suite('with an empty database name', function () {
             beforeEach(function () {
-              chatContextStub = {
-                history: [
-                  Object.assign(
-                    Object.create(vscode.ChatRequestTurn.prototype),
-                    {
-                      prompt: 'find all docs by a name example',
-                      command: 'query',
-                      references: [],
-                      participant: CHAT_PARTICIPANT_ID,
-                    }
-                  ),
-                  Object.assign(
-                    Object.create(vscode.ChatResponseTurn.prototype),
-                    {
-                      participant: CHAT_PARTICIPANT_ID,
-                      response: [
-                        {
-                          value: {
-                            value:
-                              'Which database would you like to query within this database?',
-                          } as vscode.MarkdownString,
-                        },
-                      ],
-                      command: 'query',
-                      result: {
-                        metadata: {
-                          intent: 'askForNamespace',
-                        },
-                      },
-                    }
-                  ),
-                  Object.assign(
-                    Object.create(vscode.ChatRequestTurn.prototype),
-                    {
-                      prompt: 'dbOne',
-                      command: 'query',
-                      references: [],
-                      participant: CHAT_PARTICIPANT_ID,
-                    }
-                  ),
-                  Object.assign(
-                    Object.create(vscode.ChatResponseTurn.prototype),
-                    {
-                      participant: CHAT_PARTICIPANT_ID,
-                      response: [
-                        {
-                          value: {
-                            value:
-                              'Which collection would you like to query within dbOne?',
-                          } as vscode.MarkdownString,
-                        },
-                      ],
-                      command: 'query',
-                      result: {
-                        metadata: {
-                          intent: 'askForNamespace',
-                          databaseName: 'dbOne',
-                          collectionName: undefined,
-                          chatId: 'pineapple',
-                        },
-                      },
-                    }
-                  ),
-                ],
-              };
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: undefined,
+                  collectionName: undefined,
+                })
+              );
             });
 
             afterEach(function () {
               sinon.restore();
             });
 
-            test('collection name gets picked automatically if there is only', async function () {
-              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
+            test('database name gets picked automatically if there is only 1', async function () {
+              listDatabasesStub.resolves([{ name: 'onlyOneDb' }]);
+
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
+              );
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'what is this',
+                command: 'query',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.called).to.be.false;
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
 
-              const chatResult = await invokeChatHandler(chatRequestMock);
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'askForNamespace',
+                databaseName: 'onlyOneDb',
+                collectionName: undefined,
+              });
+            });
 
-              const responses = chatStreamStub.markdown
-                .getCalls()
-                .map((call) => call.args[0]);
-              expect(responses.length).equals(1);
-              expect(responses[0]).to.include(
-                'Ask anything about MongoDB, from writing queries to questions about your cluster.'
+            test('prompts for database name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
               );
 
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'query',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.calledOnce).to.be.true;
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+
               expect(chatResult?.metadata).deep.equals({
-                intent: 'emptyRequest',
-                chatId: 'pineapple',
+                intent: 'askForNamespace',
+                chatId: testChatId,
+                databaseName: undefined,
+                collectionName: undefined,
               });
             });
+          });
 
-            test('prompts for collection name if there are multiple available', async function () {
-              const chatResult = await invokeChatHandler(chatRequestMock);
+          suite('with an empty collection name', function () {
+            beforeEach(function () {
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: 'dbOne',
+                  collectionName: undefined,
+                })
+              );
+            });
 
-              const emptyMessage = chatStreamStub.markdown.getCall(0).args[0];
-              expect(emptyMessage).to.include(
-                'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
+            afterEach(function () {
+              sinon.restore();
+            });
+
+            test('collection name gets picked automatically if there is only 1', async function () {
+              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
               );
-              const listCollsMessage =
-                chatStreamStub.markdown.getCall(1).args[0];
-              expect(listCollsMessage.value).to.include(
-                `- [collOne](command:mdb.selectCollectionWithParticipant?${encodeStringify(
-                  {
-                    command: '/query',
-                    chatId: 'pineapple',
-                    databaseName: 'dbOne',
-                    collectionName: 'collOne',
-                  }
-                )})`
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
               );
-              const showMoreCollsMessage =
-                chatStreamStub.markdown.getCall(11).args[0];
-              expect(showMoreCollsMessage.value).to.include(
-                `- [Show more](command:mdb.selectCollectionWithParticipant?${encodeStringify(
-                  {
-                    command: '/query',
-                    chatId: 'pineapple',
-                    databaseName: 'dbOne',
-                  }
-                )})`
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'what is this',
+                command: 'query',
+                references: [],
+              });
+
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
+              ).to.include({
+                collectionName: 'onlyOneColl',
+              });
+
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'query',
+              });
+            });
+
+            test('prompts for collection name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
               );
-              expect({
-                ...chatResult?.metadata,
-                chatId: undefined,
-              }).to.deep.equal({
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'query',
+                references: [],
+              });
+
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.called
+              ).to.be.false;
+
+              expect(chatResult?.metadata).deep.equals({
                 intent: 'askForNamespace',
-                collectionName: undefined,
+                chatId: testChatId,
                 databaseName: 'dbOne',
-                chatId: undefined,
+                collectionName: undefined,
               });
             });
           });
@@ -1548,6 +1561,79 @@ suite('Participant Controller Test Suite', function () {
             );
           });
 
+          suite('with an empty database name', function () {
+            beforeEach(function () {
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: undefined,
+                  collectionName: undefined,
+                })
+              );
+            });
+
+            afterEach(function () {
+              sinon.restore();
+            });
+
+            test('database name gets picked automatically if there is only 1', async function () {
+              listDatabasesStub.resolves([{ name: 'onlyOneDb' }]);
+
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
+              );
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'what is this',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.called).to.be.false;
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
+
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'askForNamespace',
+                databaseName: 'onlyOneDb',
+                collectionName: undefined,
+              });
+            });
+
+            test('prompts for database name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.calledOnce).to.be.true;
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+
+              expect(chatResult?.metadata).deep.equals({
+                intent: 'askForNamespace',
+                chatId: testChatId,
+                databaseName: undefined,
+                collectionName: undefined,
+              });
+            });
+          });
+
           suite('with an empty collection name', function () {
             beforeEach(function () {
               sinon.replace(

From 274fe173788d9405dfc7f0ad8e15731a24327de0 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Mon, 4 Nov 2024 00:30:20 +0100
Subject: [PATCH 15/29] typos

---
 src/participant/participant.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 5701e0d64..7336520da 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -1230,7 +1230,7 @@ export default class ParticipantController {
       });
 
     // If either the database or collection name could not be automatically picked
-    // then the user has be prompted to select one manually.
+    // then the user has been prompted to select one manually.
     if (databaseName === undefined || collectionName === undefined) {
       return namespaceRequestChatResult({
         databaseName,
@@ -1363,7 +1363,7 @@ export default class ParticipantController {
       });
 
     // If either the database or collection name could not be automatically picked
-    // then the user has be prompted to select one manually.
+    // then the user has been prompted to select one manually.
     if (databaseName === undefined || collectionName === undefined) {
       return namespaceRequestChatResult({
         databaseName,

From 54885d66883aebb38cdec259366a68f0f341680b Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Mon, 4 Nov 2024 10:50:41 +0100
Subject: [PATCH 16/29] Align with broken test

---
 src/participant/participant.ts | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 7336520da..cc73c5e3c 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -1140,7 +1140,7 @@ export default class ParticipantController {
 
     // When the last message was asking for a database or collection name,
     // we re-ask the question.
-    let databaseName = lastMessage.metadata.databaseName;
+    const databaseName = lastMessage.metadata.databaseName;
     if (databaseName) {
       const collections = await this._getCollections({
         stream,
@@ -1169,16 +1169,29 @@ export default class ParticipantController {
         stream,
       });
     } else {
+      const databases = await this._getDatabases({
+        stream,
+      });
+
+      if (!databases) {
+        return namespaceRequestChatResult({
+          databaseName,
+          collectionName: undefined,
+          history: context.history,
+        });
+      }
+
       stream.markdown(
         vscode.l10n.t(
           'Please select a database by either clicking on an item in the list or typing the name manually in the chat.'
         )
       );
 
-      databaseName = await this._getOrAskForDatabaseName({
+      this.renderDatabasesTree({
+        databases,
         command,
-        stream,
         context,
+        stream,
       });
     }
 

From f08001665add17e4471a3eb6dcacdf979cb0eca4 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Mon, 4 Nov 2024 15:49:54 +0100
Subject: [PATCH 17/29] wrap l10n

---
 src/participant/participant.ts | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index cc73c5e3c..c11f3c7db 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -860,7 +860,9 @@ export default class ParticipantController {
       }
 
       stream.markdown(
-        `Which collection would you like to use within ${databaseName}?\n\n`
+        vscode.l10n.t(
+          `Which collection would you like to use within ${databaseName}?\n\n`
+        )
       );
 
       this.renderCollectionsTree({

From 62c07ef161e87c50d99a4d089e4a8b6c26d2604f Mon Sep 17 00:00:00 2001
From: Gagik Amaryan <me@gagik.co>
Date: Wed, 6 Nov 2024 10:52:26 +0100
Subject: [PATCH 18/29] Apply suggestions from code review

Co-authored-by: Rhys <Anemy@users.noreply.github.com>
---
 src/participant/participant.ts | 43 +++++++++++++++-------------------
 1 file changed, 19 insertions(+), 24 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index c11f3c7db..57163d92b 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -810,8 +810,6 @@ export default class ParticipantController {
       return databases;
     } catch (error) {
       log.error('Unable to fetch databases:', error);
-
-      return;
     }
   }
 
@@ -835,7 +833,6 @@ export default class ParticipantController {
       return await dataService.listCollections(databaseName);
     } catch (error) {
       log.error('Unable to fetch collections:', error);
-      return;
     }
   }
 
@@ -854,27 +851,27 @@ export default class ParticipantController {
   }): Promise<string | undefined> {
     const collections = await this._getCollections({ stream, databaseName });
 
-    if (collections !== undefined) {
-      if (collections.length === 1) {
-        return collections[0].name;
-      }
-
-      stream.markdown(
-        vscode.l10n.t(
-          `Which collection would you like to use within ${databaseName}?\n\n`
-        )
-      );
-
-      this.renderCollectionsTree({
-        collections,
-        command,
-        databaseName,
-        context,
-        stream,
-      });
+    if (collections === undefined) {
+      return;
+    }
+    
+    if (collections.length === 1) {
+      return collections[0].name;
     }
 
-    return;
+    stream.markdown(
+      vscode.l10n.t(
+        `Which collection would you like to use within ${databaseName}?\n\n`
+      )
+    );
+
+    this.renderCollectionsTree({
+      collections,
+      command,
+      databaseName,
+      context,
+      stream,
+    });
   }
 
   /** Gets the database name if there is only one collection.
@@ -914,8 +911,6 @@ export default class ParticipantController {
       context,
       stream,
     });
-
-    return;
   }
 
   /** Helper which either automatically picks and returns missing parts of the namespace (if any)

From f2b54ea47603744815b61a5255d1a6b7998430f1 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Wed, 6 Nov 2024 14:43:46 +0100
Subject: [PATCH 19/29] changes from review

---
 src/participant/chatMetadata.ts | 14 ++++++++++
 src/participant/participant.ts  | 45 ++++++++++++++-------------------
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/src/participant/chatMetadata.ts b/src/participant/chatMetadata.ts
index 1462c4ae4..1daade9bb 100644
--- a/src/participant/chatMetadata.ts
+++ b/src/participant/chatMetadata.ts
@@ -20,6 +20,20 @@ export class ChatMetadataStore {
     return this._chats[chatId];
   }
 
+  patchChatMetadata(
+    context: vscode.ChatContext,
+    patchedMetadata: Partial<ChatMetadata>
+  ): void {
+    const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
+      context.history
+    );
+
+    this.setChatMetadata(chatId, {
+      ...this.getChatMetadata(chatId),
+      ...patchedMetadata,
+    });
+  }
+
   // Exposed for stubbing in tests.
   static createNewChatId(): string {
     return uuidv4();
diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 57163d92b..7caf5366b 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -795,21 +795,23 @@ export default class ParticipantController {
       }[]
     | undefined
   > {
-    stream.push(
-      new vscode.ChatResponseProgressPart('Fetching database names...')
-    );
     const dataService = this._connectionController.getActiveDataService();
     if (!dataService) {
-      return;
+      return undefined;
     }
 
+    stream.push(
+      new vscode.ChatResponseProgressPart('Fetching database names...')
+    );
+
     try {
-      const databases = await dataService.listDatabases({
+      return await dataService.listDatabases({
         nameOnly: true,
       });
-      return databases;
     } catch (error) {
       log.error('Unable to fetch databases:', error);
+
+      return undefined;
     }
   }
 
@@ -820,19 +822,22 @@ export default class ParticipantController {
     stream: vscode.ChatResponseStream;
     databaseName: string;
   }): Promise<ReturnType<DataService['listCollections']> | undefined> {
-    stream.push(
-      new vscode.ChatResponseProgressPart('Fetching collection names...')
-    );
-
     const dataService = this._connectionController.getActiveDataService();
 
     if (!dataService) {
-      return;
+      return undefined;
     }
+
+    stream.push(
+      new vscode.ChatResponseProgressPart('Fetching collection names...')
+    );
+
     try {
       return await dataService.listCollections(databaseName);
     } catch (error) {
       log.error('Unable to fetch collections:', error);
+
+      return undefined;
     }
   }
 
@@ -854,7 +859,6 @@ export default class ParticipantController {
     if (collections === undefined) {
       return;
     }
-    
     if (collections.length === 1) {
       return collections[0].name;
     }
@@ -939,18 +943,12 @@ export default class ParticipantController {
         stream,
       });
 
-      // If the database name could not get automatically selected,
-      // then the user has been prompted for it instead.
+      // databaseName will be undefined if some error occurs.
       if (!databaseName) {
         return { databaseName, collectionName };
       }
 
-      // Save the database name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
+      this._chatMetadataStore.patchChatMetadata(context, {
         databaseName,
       });
     }
@@ -972,12 +970,7 @@ export default class ParticipantController {
         };
       }
 
-      // Save the collection name in the metadata.
-      const chatId = ChatMetadataStore.getChatIdFromHistoryOrNewChatId(
-        context.history
-      );
-      this._chatMetadataStore.setChatMetadata(chatId, {
-        ...this._chatMetadataStore.getChatMetadata(chatId),
+      this._chatMetadataStore.patchChatMetadata(context, {
         collectionName,
       });
     }

From bf2a0c7e212b1c117b84cffd4e67870f643675ee Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Wed, 6 Nov 2024 17:09:10 +0100
Subject: [PATCH 20/29] switch to parametrized tests

---
 .vscode/settings.json                         |   1 -
 .../suite/participant/participant.test.ts     | 298 +++++++++---------
 2 files changed, 152 insertions(+), 147 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 732bb7406..32baf6163 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,6 +1,5 @@
 // Place your settings in this file to overwrite default and user settings.
 {
-  "editor.formatOnSave": false,
   "editor.codeActionsOnSave": {
     "source.fixAll.eslint": "explicit"
   },
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index a418da8cf..dd50fb824 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1560,152 +1560,6 @@ suite('Participant Controller Test Suite', function () {
               'see previous messages'
             );
           });
-
-          suite('with an empty database name', function () {
-            beforeEach(function () {
-              sinon.replace(
-                testParticipantController._chatMetadataStore,
-                'getChatMetadata',
-                () => ({
-                  databaseName: undefined,
-                  collectionName: undefined,
-                })
-              );
-            });
-
-            afterEach(function () {
-              sinon.restore();
-            });
-
-            test('database name gets picked automatically if there is only 1', async function () {
-              listDatabasesStub.resolves([{ name: 'onlyOneDb' }]);
-
-              const renderDatabasesTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderDatabasesTree'
-              );
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'what is this',
-                command: 'schema',
-                references: [],
-              });
-
-              expect(renderDatabasesTreeSpy.called).to.be.false;
-              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
-
-              expect(chatResult?.metadata).deep.equals({
-                chatId: testChatId,
-                intent: 'askForNamespace',
-                databaseName: 'onlyOneDb',
-                collectionName: undefined,
-              });
-            });
-
-            test('prompts for database name if there are multiple available', async function () {
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const renderDatabasesTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderDatabasesTree'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'dbOne',
-                command: 'schema',
-                references: [],
-              });
-
-              expect(renderDatabasesTreeSpy.calledOnce).to.be.true;
-              expect(renderCollectionsTreeSpy.called).to.be.false;
-
-              expect(chatResult?.metadata).deep.equals({
-                intent: 'askForNamespace',
-                chatId: testChatId,
-                databaseName: undefined,
-                collectionName: undefined,
-              });
-            });
-          });
-
-          suite('with an empty collection name', function () {
-            beforeEach(function () {
-              sinon.replace(
-                testParticipantController._chatMetadataStore,
-                'getChatMetadata',
-                () => ({
-                  databaseName: 'dbOne',
-                  collectionName: undefined,
-                })
-              );
-            });
-
-            test('collection name gets picked automatically if there is only 1', async function () {
-              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
-                testParticipantController,
-                '_fetchCollectionSchemaAndSampleDocuments'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'dbOne',
-                command: 'schema',
-                references: [],
-              });
-
-              expect(renderCollectionsTreeSpy.called).to.be.false;
-
-              expect(
-                fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
-              ).to.include({
-                collectionName: 'onlyOneColl',
-              });
-
-              expect(chatResult?.metadata).deep.equals({
-                chatId: testChatId,
-                intent: 'schema',
-              });
-            });
-
-            test('prompts for collection name if there are multiple available', async function () {
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
-                testParticipantController,
-                '_fetchCollectionSchemaAndSampleDocuments'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'dbOne',
-                command: 'schema',
-                references: [],
-              });
-
-              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
-              expect(
-                fetchCollectionSchemaAndSampleDocumentsSpy.called
-              ).to.be.false;
-
-              expect(chatResult?.metadata).deep.equals({
-                intent: 'askForNamespace',
-                chatId: testChatId,
-                databaseName: 'dbOne',
-                collectionName: undefined,
-              });
-            });
-          });
         });
 
         suite(
@@ -2059,6 +1913,158 @@ Schema:
         });
       });
     });
+
+    suite('determining the namespace', function () {
+      ['query', 'schema'].forEach(function (command) {
+        suite(`${command} command`, function () {
+          suite('with an empty database name', function () {
+            beforeEach(function () {
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: undefined,
+                  collectionName: undefined,
+                })
+              );
+            });
+
+            afterEach(function () {
+              sinon.restore();
+            });
+
+            test('database name gets picked automatically if there is only 1', async function () {
+              listDatabasesStub.resolves([{ name: 'onlyOneDb' }]);
+
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
+              );
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'what is this',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.called).to.be.false;
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
+
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'askForNamespace',
+                databaseName: 'onlyOneDb',
+                collectionName: undefined,
+              });
+            });
+
+            test('prompts for database name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const renderDatabasesTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderDatabasesTree'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderDatabasesTreeSpy.calledOnce).to.be.true;
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+
+              expect(chatResult?.metadata).deep.equals({
+                intent: 'askForNamespace',
+                chatId: testChatId,
+                databaseName: undefined,
+                collectionName: undefined,
+              });
+            });
+          });
+
+          suite('with an empty collection name', function () {
+            beforeEach(function () {
+              sinon.replace(
+                testParticipantController._chatMetadataStore,
+                'getChatMetadata',
+                () => ({
+                  databaseName: 'dbOne',
+                  collectionName: undefined,
+                })
+              );
+            });
+
+            test('collection name gets picked automatically if there is only 1', async function () {
+              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderCollectionsTreeSpy.called).to.be.false;
+
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
+              ).to.include({
+                collectionName: 'onlyOneColl',
+              });
+
+              expect(chatResult?.metadata).deep.equals({
+                chatId: testChatId,
+                intent: 'schema',
+              });
+            });
+
+            test('prompts for collection name if there are multiple available', async function () {
+              const renderCollectionsTreeSpy = sinon.spy(
+                testParticipantController,
+                'renderCollectionsTree'
+              );
+              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
+                testParticipantController,
+                '_fetchCollectionSchemaAndSampleDocuments'
+              );
+
+              const chatResult = await invokeChatHandler({
+                prompt: 'dbOne',
+                command: 'schema',
+                references: [],
+              });
+
+              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
+              expect(
+                fetchCollectionSchemaAndSampleDocumentsSpy.called
+              ).to.be.false;
+
+              expect(chatResult?.metadata).deep.equals({
+                intent: 'askForNamespace',
+                chatId: testChatId,
+                databaseName: 'dbOne',
+                collectionName: undefined,
+              });
+            });
+          });
+        });
+      });
+    });
   });
 
   suite('prompt builders', function () {

From ccacb8a36fe129fc250f9a843bce43f92a168681 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Wed, 6 Nov 2024 17:24:49 +0100
Subject: [PATCH 21/29] remove vscode

---
 .vscode/settings.json | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 32baf6163..732bb7406 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,5 +1,6 @@
 // Place your settings in this file to overwrite default and user settings.
 {
+  "editor.formatOnSave": false,
   "editor.codeActionsOnSave": {
     "source.fixAll.eslint": "explicit"
   },

From 886aa451d2efd3b47a956cd7cc935c5659b6359b Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Thu, 7 Nov 2024 00:59:36 +0100
Subject: [PATCH 22/29] better comments

---
 src/participant/participant.ts | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 7caf5366b..dc64ed4d6 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -943,7 +943,9 @@ export default class ParticipantController {
         stream,
       });
 
-      // databaseName will be undefined if some error occurs.
+      // databaseName will be undefined if it cannot be found from
+      // the metadata or history, in which case the user will be prompted
+      // to select it or if some error occurs.
       if (!databaseName) {
         return { databaseName, collectionName };
       }
@@ -1233,7 +1235,7 @@ export default class ParticipantController {
       });
 
     // If either the database or collection name could not be automatically picked
-    // then the user has been prompted to select one manually.
+    // then the user has been prompted to select one manually or been presented with an error.
     if (databaseName === undefined || collectionName === undefined) {
       return namespaceRequestChatResult({
         databaseName,

From 52dbdf600a897da55d9aa13bf9a4fd52e065c67a Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Thu, 7 Nov 2024 10:09:51 +0100
Subject: [PATCH 23/29] fix potential CI discrepancy

---
 src/test/suite/participant/participant.test.ts | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index dd50fb824..b4b6fee8c 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1917,6 +1917,12 @@ Schema:
     suite('determining the namespace', function () {
       ['query', 'schema'].forEach(function (command) {
         suite(`${command} command`, function () {
+          beforeEach(function () {
+            sendRequestStub.onCall(0).resolves({
+              text: ['determining the namespace'],
+            });
+          });
+
           suite('with an empty database name', function () {
             beforeEach(function () {
               sinon.replace(

From fee668e20f20a3bc870cd1169419bd1a029ffc1c Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Thu, 7 Nov 2024 13:32:54 +0100
Subject: [PATCH 24/29] combine message text

---
 src/participant/participant.ts                | 69 ++++---------------
 .../suite/participant/participant.test.ts     | 22 +++---
 2 files changed, 25 insertions(+), 66 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index 803237b31..cf3f70e5a 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -865,7 +865,7 @@ export default class ParticipantController {
 
     stream.markdown(
       vscode.l10n.t(
-        `Which collection would you like to use within ${databaseName}?\n\n`
+        `Which collection would you like to use within ${databaseName}? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n`
       )
     );
 
@@ -902,11 +902,12 @@ export default class ParticipantController {
 
     // If no database or collection name is found in the user prompt,
     // we retrieve the available namespaces from the current connection.
-    // Users can then select a value by clicking on an item in the list.
+    // Users can then select a value by clicking on an item in the list
+    // or typing the name manually.
     stream.markdown(
-      `What is the name of the database you would like${
-        command === '/query' ? ' this query' : ''
-      } to run against?\n\n`
+      `Which database would you like ${
+        command === '/query' ? 'this query to run against' : 'to use'
+      }? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n`
     );
 
     this.renderDatabasesTree({
@@ -1132,64 +1133,22 @@ export default class ParticipantController {
 
     // When the last message was asking for a database or collection name,
     // we re-ask the question.
-    const databaseName = lastMessage.metadata.databaseName;
-    if (databaseName) {
-      const collections = await this._getCollections({
-        stream,
-        databaseName,
-      });
-
-      if (!collections) {
-        return namespaceRequestChatResult({
-          databaseName,
-          collectionName: undefined,
-          history: context.history,
-        });
-      }
-
-      stream.markdown(
-        vscode.l10n.t(
-          'Please select a collection by either clicking on an item in the list or typing the name manually in the chat.'
-        )
-      );
+    const metadataDatabaseName = lastMessage.metadata.databaseName;
 
-      this.renderCollectionsTree({
-        collections,
-        command,
-        databaseName,
-        context,
-        stream,
-      });
-    } else {
-      const databases = await this._getDatabases({
-        stream,
-      });
-
-      if (!databases) {
-        return namespaceRequestChatResult({
-          databaseName,
-          collectionName: undefined,
-          history: context.history,
-        });
-      }
-
-      stream.markdown(
-        vscode.l10n.t(
-          'Please select a database by either clicking on an item in the list or typing the name manually in the chat.'
-        )
-      );
-
-      this.renderDatabasesTree({
-        databases,
+    // This will prompt the user for the missing databaseName or the collectionName.
+    // If anything in the namespace can be automatically picked, it will be returned.
+    const { databaseName, collectionName } =
+      await this._getOrAskForMissingNamespace({
         command,
         context,
         stream,
+        databaseName: metadataDatabaseName,
+        collectionName: undefined,
       });
-    }
 
     return namespaceRequestChatResult({
       databaseName,
-      collectionName: undefined,
+      collectionName,
       history: context.history,
     });
   }
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index aeecbe65e..29f57d00c 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1033,7 +1033,7 @@ suite('Participant Controller Test Suite', function () {
             const chatResult = await invokeChatHandler(chatRequestMock);
             const askForDBMessage = chatStreamStub.markdown.getCall(0).args[0];
             expect(askForDBMessage).to.include(
-              'What is the name of the database you would like this query to run against?'
+              'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n'
             );
             const listDBsMessage = chatStreamStub.markdown.getCall(1).args[0];
             const expectedContent = encodeStringify({
@@ -1087,7 +1087,7 @@ suite('Participant Controller Test Suite', function () {
                       {
                         value: {
                           value:
-                            'What is the name of the database you would like this query to run against?',
+                            'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n',
                         } as vscode.MarkdownString,
                       },
                     ],
@@ -1108,7 +1108,7 @@ suite('Participant Controller Test Suite', function () {
             const askForCollMessage =
               chatStreamStub.markdown.getCall(12).args[0];
             expect(askForCollMessage).to.include(
-              'Which collection would you like to use within dbOne?'
+              'Which collection would you like to use within dbOne? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n'
             );
             const listCollsMessage =
               chatStreamStub.markdown.getCall(13).args[0];
@@ -1164,7 +1164,7 @@ suite('Participant Controller Test Suite', function () {
                       {
                         value: {
                           value:
-                            'Which database would you like to query within this database?',
+                            'Which database would you like to this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n',
                         } as vscode.MarkdownString,
                       },
                     ],
@@ -1190,7 +1190,7 @@ suite('Participant Controller Test Suite', function () {
                       {
                         value: {
                           value:
-                            'Which collection would you like to query within dbOne?',
+                            'Which collection would you like to query within dbOne? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n',
                         } as vscode.MarkdownString,
                       },
                     ],
@@ -1254,7 +1254,7 @@ suite('Participant Controller Test Suite', function () {
                       {
                         value: {
                           value:
-                            'What is the name of the database you would like this query to run against?',
+                            'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n',
                         } as vscode.MarkdownString,
                       },
                     ],
@@ -1272,8 +1272,8 @@ suite('Participant Controller Test Suite', function () {
             const chatResult = await invokeChatHandler(chatRequestMock);
 
             const emptyMessage = chatStreamStub.markdown.getCall(0).args[0];
-            expect(emptyMessage).to.include(
-              'Please select a database by either clicking on an item in the list or typing the name manually in the chat.'
+            expect(emptyMessage).to.equal(
+              'Which database would you like this query to run against? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n'
             );
             const listDBsMessage = chatStreamStub.markdown.getCall(1).args[0];
             expect(listDBsMessage.value).to.include(
@@ -1488,7 +1488,7 @@ suite('Participant Controller Test Suite', function () {
             expect(sendRequestStub.called).to.be.false;
             const askForDBMessage = chatStreamStub.markdown.getCall(0).args[0];
             expect(askForDBMessage).to.include(
-              'What is the name of the database you would like to run against?'
+              'Which database would you like to use? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n'
             );
           });
 
@@ -1508,8 +1508,8 @@ suite('Participant Controller Test Suite', function () {
             );
 
             const askForDBMessage = chatStreamStub.markdown.getCall(0).args[0];
-            expect(askForDBMessage).to.include(
-              'What is the name of the database you would like to run against?'
+            expect(askForDBMessage).to.equals(
+              'Which database would you like to use? Select one by either clicking on an item in the list or typing the name manually in the chat.\n\n'
             );
           });
 

From e1d12bf6c39684c28da40add3c68e717eb71adea Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Thu, 7 Nov 2024 13:42:13 +0100
Subject: [PATCH 25/29] add explicit undefined returns

---
 src/participant/participant.ts | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index cf3f70e5a..d4301d61d 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -857,7 +857,8 @@ export default class ParticipantController {
     const collections = await this._getCollections({ stream, databaseName });
 
     if (collections === undefined) {
-      return;
+      log.error('No collections found');
+      return undefined;
     }
     if (collections.length === 1) {
       return collections[0].name;
@@ -876,6 +877,8 @@ export default class ParticipantController {
       context,
       stream,
     });
+
+    return undefined;
   }
 
   /** Gets the database name if there is only one collection.
@@ -893,7 +896,7 @@ export default class ParticipantController {
 
     if (databases === undefined || databases.length === 0) {
       log.error('No databases found');
-      return;
+      return undefined;
     }
 
     if (databases.length === 1) {
@@ -916,6 +919,8 @@ export default class ParticipantController {
       context,
       stream,
     });
+
+    return undefined;
   }
 
   /** Helper which either automatically picks and returns missing parts of the namespace (if any)

From 9292c5f3d9d0935052b062f97977052604a55803 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 8 Nov 2024 12:50:13 +0100
Subject: [PATCH 26/29] Use command from test

---
 src/test/suite/participant/participant.test.ts | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 29f57d00c..0a965e442 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1954,7 +1954,7 @@ Schema:
 
               const chatResult = await invokeChatHandler({
                 prompt: 'what is this',
-                command: 'schema',
+                command,
                 references: [],
               });
 
@@ -1981,7 +1981,7 @@ Schema:
 
               const chatResult = await invokeChatHandler({
                 prompt: 'dbOne',
-                command: 'schema',
+                command,
                 references: [],
               });
 
@@ -2022,7 +2022,7 @@ Schema:
 
               const chatResult = await invokeChatHandler({
                 prompt: 'dbOne',
-                command: 'schema',
+                command,
                 references: [],
               });
 
@@ -2036,7 +2036,7 @@ Schema:
 
               expect(chatResult?.metadata).deep.equals({
                 chatId: testChatId,
-                intent: 'schema',
+                intent: command,
               });
             });
 
@@ -2052,7 +2052,7 @@ Schema:
 
               const chatResult = await invokeChatHandler({
                 prompt: 'dbOne',
-                command: 'schema',
+                command,
                 references: [],
               });
 

From 7391721ace9f3c1652d1ea72e92d799f3d01c499 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 8 Nov 2024 14:07:29 +0100
Subject: [PATCH 27/29] resolve correctly

---
 src/test/suite/participant/participant.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 0a965e442..f760527bc 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1919,7 +1919,7 @@ Schema:
       ['query', 'schema'].forEach(function (command) {
         suite(`${command} command`, function () {
           beforeEach(function () {
-            sendRequestStub.onCall(0).resolves({
+            sendRequestStub.resolves({
               text: ['determining the namespace'],
             });
           });

From affdf746638bcc3fa61d5f20a3165a7fbd22e458 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 8 Nov 2024 14:20:43 +0100
Subject: [PATCH 28/29] move participant error types

---
 src/participant/participant.ts                            | 2 +-
 src/{test/suite => }/participant/participantErrorTypes.ts | 0
 src/participant/prompts/promptBase.ts                     | 2 +-
 src/telemetry/telemetryService.ts                         | 2 +-
 src/test/suite/participant/participant.test.ts            | 2 +-
 5 files changed, 4 insertions(+), 4 deletions(-)
 rename src/{test/suite => }/participant/participantErrorTypes.ts (100%)

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
index d4301d61d..b0254c16f 100644
--- a/src/participant/participant.ts
+++ b/src/participant/participant.ts
@@ -44,7 +44,7 @@ import type { ModelInput } from './prompts/promptBase';
 import { processStreamWithIdentifiers } from './streamParsing';
 import type { PromptIntent } from './prompts/intent';
 import type { DataService } from 'mongodb-data-service';
-import { ParticipantErrorTypes } from '../test/suite/participant/participantErrorTypes';
+import { ParticipantErrorTypes } from './participantErrorTypes';
 
 const log = createLogger('participant');
 
diff --git a/src/test/suite/participant/participantErrorTypes.ts b/src/participant/participantErrorTypes.ts
similarity index 100%
rename from src/test/suite/participant/participantErrorTypes.ts
rename to src/participant/participantErrorTypes.ts
diff --git a/src/participant/prompts/promptBase.ts b/src/participant/prompts/promptBase.ts
index 40c430eec..364284240 100644
--- a/src/participant/prompts/promptBase.ts
+++ b/src/participant/prompts/promptBase.ts
@@ -4,7 +4,7 @@ import type {
   InternalPromptPurpose,
   ParticipantPromptProperties,
 } from '../../telemetry/telemetryService';
-import { ParticipantErrorTypes } from '../../test/suite/participant/participantErrorTypes';
+import { ParticipantErrorTypes } from '../participantErrorTypes';
 
 export interface PromptArgsBase {
   request: {
diff --git a/src/telemetry/telemetryService.ts b/src/telemetry/telemetryService.ts
index 8e0482606..8a58258dc 100644
--- a/src/telemetry/telemetryService.ts
+++ b/src/telemetry/telemetryService.ts
@@ -13,7 +13,7 @@ import type { NewConnectionTelemetryEventProperties } from './connectionTelemetr
 import type { ShellEvaluateResult } from '../types/playgroundType';
 import type { StorageController } from '../storage';
 import type { ParticipantResponseType } from '../participant/constants';
-import { ParticipantErrorTypes } from '../test/suite/participant/participantErrorTypes';
+import { ParticipantErrorTypes } from '../participant/participantErrorTypes';
 
 const log = createLogger('telemetry');
 // eslint-disable-next-line @typescript-eslint/no-var-requires
diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index f760527bc..3085a4c88 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -34,7 +34,7 @@ import { Prompts } from '../../../participant/prompts';
 import { createMarkdownLink } from '../../../participant/markdown';
 import EXTENSION_COMMANDS from '../../../commands';
 import { getContentLength } from '../../../participant/prompts/promptBase';
-import { ParticipantErrorTypes } from './participantErrorTypes';
+import { ParticipantErrorTypes } from '../../../participant/participantErrorTypes';
 
 // The Copilot's model in not available in tests,
 // therefore we need to mock its methods and returning values.

From 4dea018a87d0d9c5e3d7c410becc4c80ca5fd261 Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 8 Nov 2024 14:33:21 +0100
Subject: [PATCH 29/29] remove redundant test code

---
 .../suite/participant/participant.test.ts     | 150 ------------------
 1 file changed, 150 deletions(-)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index 3085a4c88..8b5916ecf 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -1305,156 +1305,6 @@ suite('Participant Controller Test Suite', function () {
               chatId: undefined,
             });
           });
-
-          suite('with an empty database name', function () {
-            beforeEach(function () {
-              sinon.replace(
-                testParticipantController._chatMetadataStore,
-                'getChatMetadata',
-                () => ({
-                  databaseName: undefined,
-                  collectionName: undefined,
-                })
-              );
-            });
-
-            afterEach(function () {
-              sinon.restore();
-            });
-
-            test('database name gets picked automatically if there is only 1', async function () {
-              listDatabasesStub.resolves([{ name: 'onlyOneDb' }]);
-
-              const renderDatabasesTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderDatabasesTree'
-              );
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'what is this',
-                command: 'query',
-                references: [],
-              });
-
-              expect(renderDatabasesTreeSpy.called).to.be.false;
-              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
-
-              expect(chatResult?.metadata).deep.equals({
-                chatId: testChatId,
-                intent: 'askForNamespace',
-                databaseName: 'onlyOneDb',
-                collectionName: undefined,
-              });
-            });
-
-            test('prompts for database name if there are multiple available', async function () {
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const renderDatabasesTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderDatabasesTree'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'dbOne',
-                command: 'query',
-                references: [],
-              });
-
-              expect(renderDatabasesTreeSpy.calledOnce).to.be.true;
-              expect(renderCollectionsTreeSpy.called).to.be.false;
-
-              expect(chatResult?.metadata).deep.equals({
-                intent: 'askForNamespace',
-                chatId: testChatId,
-                databaseName: undefined,
-                collectionName: undefined,
-              });
-            });
-          });
-
-          suite('with an empty collection name', function () {
-            beforeEach(function () {
-              sinon.replace(
-                testParticipantController._chatMetadataStore,
-                'getChatMetadata',
-                () => ({
-                  databaseName: 'dbOne',
-                  collectionName: undefined,
-                })
-              );
-            });
-
-            afterEach(function () {
-              sinon.restore();
-            });
-
-            test('collection name gets picked automatically if there is only 1', async function () {
-              listCollectionsStub.resolves([{ name: 'onlyOneColl' }]);
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
-                testParticipantController,
-                '_fetchCollectionSchemaAndSampleDocuments'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'what is this',
-                command: 'query',
-                references: [],
-              });
-
-              expect(renderCollectionsTreeSpy.called).to.be.false;
-
-              expect(
-                fetchCollectionSchemaAndSampleDocumentsSpy.firstCall.args[0]
-              ).to.include({
-                collectionName: 'onlyOneColl',
-              });
-
-              expect(chatResult?.metadata).deep.equals({
-                chatId: testChatId,
-                intent: 'query',
-              });
-            });
-
-            test('prompts for collection name if there are multiple available', async function () {
-              const renderCollectionsTreeSpy = sinon.spy(
-                testParticipantController,
-                'renderCollectionsTree'
-              );
-              const fetchCollectionSchemaAndSampleDocumentsSpy = sinon.spy(
-                testParticipantController,
-                '_fetchCollectionSchemaAndSampleDocuments'
-              );
-
-              const chatResult = await invokeChatHandler({
-                prompt: 'dbOne',
-                command: 'query',
-                references: [],
-              });
-
-              expect(renderCollectionsTreeSpy.calledOnce).to.be.true;
-              expect(
-                fetchCollectionSchemaAndSampleDocumentsSpy.called
-              ).to.be.false;
-
-              expect(chatResult?.metadata).deep.equals({
-                intent: 'askForNamespace',
-                chatId: testChatId,
-                databaseName: 'dbOne',
-                collectionName: undefined,
-              });
-            });
-          });
         });
       });