⌚ fix: Wait for Initial Message Save & Correct Latest Message (#3399)

* chore: assistants, unsupported assistant, better logging

* chore: remove unnecessary logger in validateAssistant middleware

* fix: resolve initial conversation save/promise before saving response

* chore: Import and organize dependencies in Speech component

* fix: conversation statefulness
- Latest Message (at index 0) should not be reset if existing convo
- add debugging context for clearAllLatestMessages
- Added logging concerning latest Message updates (dev mode only)
- update latest message Set logic, also checks for change in conversation Id
- consolidated latest message helpers to client/src/utils/messages.ts

api/app/clients/BaseClient.js

@@ -540,6 +540,9 @@ class BaseClient {
       const completionTokens = this.getTokenCount(completion);
       await this.recordTokenUsage({ promptTokens, completionTokens });
     }
+    if (this.userMessagePromise) {
+      await this.userMessagePromise;
+    }
     this.responsePromise = this.saveMessageToDatabase(responseMessage, saveOptions, user);
     const messageCache = getLogStores(CacheKeys.MESSAGES);
     messageCache.set(
@@ -620,18 +623,23 @@ class BaseClient {
         unfinished: false,
         user,
       },
-      { context: 'api/app/clients/BaseClient.js - saveMessageToDatabase' },
+      { context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveMessage' },
     );
 
     if (this.skipSaveConvo) {
       return { message: savedMessage };
     }
-    const conversation = await saveConvo(user, {
-      conversationId: message.conversationId,
-      endpoint: this.options.endpoint,
-      endpointType: this.options.endpointType,
-      ...endpointOptions,
-    });
+
+    const conversation = await saveConvo(
+      this.options.req,
+      {
+        conversationId: message.conversationId,
+        endpoint: this.options.endpoint,
+        endpointType: this.options.endpointType,
+        ...endpointOptions,
+      },
+      { context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveConvo' },
+    );
 
     return { message: savedMessage, conversation };
   }

api/app/clients/specs/BaseClient.test.js

@@ -1,7 +1,7 @@
 const { Constants } = require('librechat-data-provider');
 const { initializeFakeClient } = require('./FakeClient');
 
-jest.mock('../../../lib/db/connectDb');
+jest.mock('~/lib/db/connectDb');
 jest.mock('~/models', () => ({
   User: jest.fn(),
   Key: jest.fn(),
@@ -631,5 +631,32 @@ describe('BaseClient', () => {
         }),
       );
     });
+
+    test('userMessagePromise is awaited before saving response message', async () => {
+      // Mock the saveMessageToDatabase method
+      TestClient.saveMessageToDatabase = jest.fn().mockImplementation(() => {
+        return new Promise((resolve) => setTimeout(resolve, 100)); // Simulate a delay
+      });
+
+      // Send a message
+      const messagePromise = TestClient.sendMessage('Hello, world!');
+
+      // Wait a short time to ensure the user message save has started
+      await new Promise((resolve) => setTimeout(resolve, 50));
+
+      // Check that saveMessageToDatabase has been called once (for the user message)
+      expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(1);
+
+      // Wait for the message to be fully processed
+      await messagePromise;
+
+      // Check that saveMessageToDatabase has been called twice (once for user message, once for response)
+      expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(2);
+
+      // Check the order of calls
+      const calls = TestClient.saveMessageToDatabase.mock.calls;
+      expect(calls[0][0].isCreatedByUser).toBe(true); // First call should be for user message
+      expect(calls[1][0].isCreatedByUser).toBe(false); // Second call should be for response message
+    });
   });
 });

api/models/Conversation.js

@@ -19,18 +19,32 @@ const getConvo = async (user, conversationId) => {
 
 module.exports = {
   Conversation,
-  saveConvo: async (user, { conversationId, newConversationId, ...convo }) => {
+  /**
+   * Saves a conversation to the database.
+   * @param {Object} req - The request object.
+   * @param {string} conversationId - The conversation's ID.
+   * @param {Object} metadata - Additional metadata to log for operation.
+   * @returns {Promise<TConversation>} The conversation object.
+   */
+  saveConvo: async (req, { conversationId, newConversationId, ...convo }, metadata) => {
     try {
+      if (metadata && metadata?.context) {
+        logger.info(`[saveConvo] ${metadata.context}`);
+      }
       const messages = await getMessages({ conversationId }, '_id');
-      const update = { ...convo, messages, user };
+      const update = { ...convo, messages, user: req.user.id };
       if (newConversationId) {
         update.conversationId = newConversationId;
       }
 
-      const conversation = await Conversation.findOneAndUpdate({ conversationId, user }, update, {
-        new: true,
-        upsert: true,
-      });
+      const conversation = await Conversation.findOneAndUpdate(
+        { conversationId, user: req.user.id },
+        update,
+        {
+          new: true,
+          upsert: true,
+        },
+      );
 
       return conversation.toObject();
     } catch (error) {

api/server/controllers/assistants/chatV1.js

@@ -383,6 +383,9 @@ const chatV1 = async (req, res) => {
       return files;
     };
 
+    /** @type {Promise<Run>|undefined} */
+    let userMessagePromise;
+
     const initializeThread = async () => {
       /** @type {[ undefined | MongoFile[]]}*/
       const [processedFiles] = await Promise.all([addVisionPrompt(), getRequestFileIds()]);
@@ -439,7 +442,7 @@ const chatV1 = async (req, res) => {
       previousMessages.push(requestMessage);
 
       /* asynchronous */
-      saveUserMessage({ ...requestMessage, model });
+      userMessagePromise = saveUserMessage(req, { ...requestMessage, model });
 
       conversation = {
         conversationId,
@@ -583,7 +586,10 @@ const chatV1 = async (req, res) => {
     });
     res.end();
 
-    await saveAssistantMessage({ ...responseMessage, model });
+    if (userMessagePromise) {
+      await userMessagePromise;
+    }
+    await saveAssistantMessage(req, { ...responseMessage, model });
 
     if (parentMessageId === Constants.NO_PARENT && !_thread_id) {
       addTitle(req, {

api/server/controllers/assistants/chatV2.js

@@ -246,6 +246,9 @@ const chatV2 = async (req, res) => {
       }
     };
 
+    /** @type {Promise<Run>|undefined} */
+    let userMessagePromise;
+
     const initializeThread = async () => {
       await getRequestFileIds();
 
@@ -288,7 +291,7 @@ const chatV2 = async (req, res) => {
       previousMessages.push(requestMessage);
 
       /* asynchronous */
-      saveUserMessage({ ...requestMessage, model });
+      userMessagePromise = saveUserMessage(req, { ...requestMessage, model });
 
       conversation = {
         conversationId,
@@ -449,7 +452,10 @@ const chatV2 = async (req, res) => {
     });
     res.end();
 
-    await saveAssistantMessage({ ...responseMessage, model });
+    if (userMessagePromise) {
+      await userMessagePromise;
+    }
+    await saveAssistantMessage(req, { ...responseMessage, model });
 
     if (parentMessageId === Constants.NO_PARENT && !_thread_id) {
       addTitle(req, {

api/server/middleware/assistants/validate.js

@@ -19,7 +19,8 @@ const validateAssistant = async (req, res, next) => {
   }
 
   const { supportedIds, excludedIds } = assistantsConfig;
-  const error = { message: 'Assistant not supported' };
+  const error = { message: 'validateAssistant: Assistant not supported' };
+
   if (supportedIds?.length && !supportedIds.includes(assistant_id)) {
     return await handleAbortError(res, req, error, {
       sender: 'System',

api/server/routes/ask/askChatGPTBrowser.js

@@ -52,7 +52,7 @@ router.post('/', setHeaders, async (req, res) => {
 
   if (!overrideParentMessageId) {
     await saveMessage(req, { ...userMessage, user: req.user.id });
-    await saveConvo(req.user.id, {
+    await saveConvo(req, {
       ...userMessage,
       ...endpointOption,
       conversationId,
@@ -183,7 +183,7 @@ const ask = async ({
       }
     }
 
-    await saveConvo(user, conversationUpdate);
+    await saveConvo(req, conversationUpdate);
     conversationId = newConversationId;
 
     // STEP3 update the user message
@@ -213,7 +213,7 @@ const ask = async ({
     if (userParentMessageId == Constants.NO_PARENT) {
       // const title = await titleConvo({ endpoint: endpointOption?.endpoint, text, response: responseMessage });
       const title = await response.details.title;
-      await saveConvo(user, {
+      await saveConvo(req, {
         conversationId: conversationId,
         title,
       });

api/server/routes/ask/bingAI.js

@@ -71,7 +71,7 @@ router.post('/', setHeaders, async (req, res) => {
 
   if (!overrideParentMessageId) {
     await saveMessage(req, { ...userMessage, user: req.user.id });
-    await saveConvo(req.user.id, {
+    await saveConvo(req, {
       ...userMessage,
       ...endpointOption,
       conversationId,
@@ -216,7 +216,7 @@ const ask = async ({
       conversationUpdate.invocationId = response.invocationId;
     }
 
-    await saveConvo(user, conversationUpdate);
+    await saveConvo(req, conversationUpdate);
     userMessage.messageId = newUserMessageId;
 
     // If response has parentMessageId, the fake userMessage.messageId should be updated to the real one.
@@ -245,7 +245,7 @@ const ask = async ({
         response: responseMessage,
       });
 
-      await saveConvo(user, {
+      await saveConvo(req, {
         conversationId: conversationId,
         title,
       });

api/server/routes/convos.js

@@ -104,7 +104,7 @@ router.post('/update', async (req, res) => {
   const update = req.body.arg;
 
   try {
-    const dbResponse = await saveConvo(req.user.id, update);
+    const dbResponse = await saveConvo(req, update, { context: 'POST /api/convos/update' });
     res.status(201).json(dbResponse);
   } catch (error) {
     logger.error('Error updating conversation', error);

api/server/routes/messages.js

@@ -30,7 +30,7 @@ router.post('/:conversationId', validateMessageReq, async (req, res) => {
     if (!savedMessage) {
       return res.status(400).json({ error: 'Message not saved' });
     }
-    await saveConvo(req.user.id, savedMessage);
+    await saveConvo(req, savedMessage, { context: 'POST /api/messages/:conversationId' });
     res.status(201).json(savedMessage);
   } catch (error) {
     logger.error('Error saving message:', error);

api/server/services/Endpoints/anthropic/addTitle.js

@@ -23,10 +23,14 @@ const addTitle = async (req, { text, response, client }) => {
 
   const title = await client.titleConvo({ text, responseText: response?.text });
   await titleCache.set(key, title, 120000);
-  await saveConvo(req.user.id, {
-    conversationId: response.conversationId,
-    title,
-  });
+  await saveConvo(
+    req,
+    {
+      conversationId: response.conversationId,
+      title,
+    },
+    { context: 'api/server/services/Endpoints/anthropic/addTitle.js' },
+  );
 };
 
 module.exports = addTitle;

api/server/services/Endpoints/assistants/addTitle.js

@@ -19,10 +19,14 @@ const addTitle = async (req, { text, responseText, conversationId, client }) =>
   const title = await client.titleConvo({ text, conversationId, responseText });
   await titleCache.set(key, title, 120000);
 
-  await saveConvo(req.user.id, {
-    conversationId,
-    title,
-  });
+  await saveConvo(
+    req,
+    {
+      conversationId,
+      title,
+    },
+    { context: 'api/server/services/Endpoints/assistants/addTitle.js' },
+  );
 };
 
 module.exports = addTitle;

api/server/services/Endpoints/google/addTitle.js

@@ -49,10 +49,14 @@ const addTitle = async (req, { text, response, client }) => {
 
   const title = await titleClient.titleConvo({ text, responseText: response?.text });
   await titleCache.set(key, title, 120000);
-  await saveConvo(req.user.id, {
-    conversationId: response.conversationId,
-    title,
-  });
+  await saveConvo(
+    req,
+    {
+      conversationId: response.conversationId,
+      title,
+    },
+    { context: 'api/server/services/Endpoints/google/addTitle.js' },
+  );
 };
 
 module.exports = addTitle;

api/server/services/Endpoints/openAI/addTitle.js

@@ -23,10 +23,14 @@ const addTitle = async (req, { text, response, client }) => {
 
   const title = await client.titleConvo({ text, responseText: response?.text });
   await titleCache.set(key, title, 120000);
-  await saveConvo(req.user.id, {
-    conversationId: response.conversationId,
-    title,
-  });
+  await saveConvo(
+    req,
+    {
+      conversationId: response.conversationId,
+      title,
+    },
+    { context: 'api/server/services/Endpoints/openAI/addTitle.js' },
+  );
 };
 
 module.exports = addTitle;

api/server/services/Threads/manage.js

@@ -41,6 +41,7 @@ async function initThread({ openai, body, thread_id: _thread_id }) {
 /**
  * Saves a user message to the DB in the Assistants endpoint format.
  *
+ * @param {Object} req - The request object.
  * @param {Object} params - The parameters of the user message
  * @param {string} params.user - The user's ID.
  * @param {string} params.text - The user's prompt.
@@ -59,7 +60,7 @@ async function initThread({ openai, body, thread_id: _thread_id }) {
  * @param {string[]} [params.file_ids] - Optional. List of File IDs attached to the userMessage.
  * @return {Promise<Run>} A promise that resolves to the created run object.
  */
-async function saveUserMessage(params) {
+async function saveUserMessage(req, params) {
   const tokenCount = await countTokens(params.text);
 
   // todo: do this on the frontend
@@ -110,14 +111,16 @@ async function saveUserMessage(params) {
   }
 
   const message = await recordMessage(userMessage);
-  await saveConvo(params.user, convo);
-
+  await saveConvo(req, convo, {
+    context: 'api/server/services/Threads/manage.js #saveUserMessage',
+  });
   return message;
 }
 
 /**
  * Saves an Assistant message to the DB in the Assistants endpoint format.
  *
+ * @param {Object} req - The request object.
  * @param {Object} params - The parameters of the Assistant message
  * @param {string} params.user - The user's ID.
  * @param {string} params.messageId - The message Id.
@@ -134,7 +137,7 @@ async function saveUserMessage(params) {
  * @param {string} [params.promptPrefix] - Optional: from preset for `additional_instructions` field.
  * @return {Promise<Run>} A promise that resolves to the created run object.
  */
-async function saveAssistantMessage(params) {
+async function saveAssistantMessage(req, params) {
   // const tokenCount = // TODO: need to count each content part
 
   const message = await recordMessage({
@@ -154,14 +157,18 @@ async function saveAssistantMessage(params) {
     // tokenCount,
   });
 
-  await saveConvo(params.user, {
-    endpoint: params.endpoint,
-    conversationId: params.conversationId,
-    promptPrefix: params.promptPrefix,
-    instructions: params.instructions,
-    assistant_id: params.assistant_id,
-    model: params.model,
-  });
+  await saveConvo(
+    req,
+    {
+      endpoint: params.endpoint,
+      conversationId: params.conversationId,
+      promptPrefix: params.promptPrefix,
+      instructions: params.instructions,
+      assistant_id: params.assistant_id,
+      model: params.model,
+    },
+    { context: 'api/server/services/Threads/manage.js #saveAssistantMessage' },
+  );
 
   return message;
 }
@@ -338,10 +345,14 @@ async function syncMessages({
   await Promise.all(modifyPromises);
   await Promise.all(recordPromises);
 
-  await saveConvo(openai.req.user.id, {
-    conversationId,
-    file_ids: attached_file_ids,
-  });
+  await saveConvo(
+    openai.req,
+    {
+      conversationId,
+      file_ids: attached_file_ids,
+    },
+    { context: 'api/server/services/Threads/manage.js #syncMessages' },
+  );
 
   return result;
 }