🛠️ fix(Azure/Assistants): Handle Long Domain Names & Other Minor chores (#2475)

* chore: replace violation cache accessors with enum

* chore: fix test

* chore(fileSchema): index timestamps

* fix(ActionService): use encoding/caching strategy for handling assistant function character length limit

* refactor(actions): async `domainParser` also resolve retrieved model (which is deployment name) to user-defined model

* style(AssistantAction): add `whitespace-nowrap` for ellipsis

* refactor(ActionService): if domain is less than or equal to encoded domain fixed length, return domain with replacement of separator

* refactor(actions): use sessions/transactions for updating Assistant Action database records

* chore: remove TTL from ENCODED_DOMAINS cache

* refactor(domainParser): minor optimization and add tests

* fix(spendTokens): use txData.user for token usage logging

* refactor(actions): add helper function `withSession` for database operations with sessions/transactions

* fix(PluginsClient): logger debug `message` field edge case

api/server/middleware/checkBan.js

@@ -1,14 +1,15 @@
 const Keyv = require('keyv');
 const uap = require('ua-parser-js');
-const denyRequest = require('./denyRequest');
-const { getLogStores } = require('../../cache');
+const { ViolationTypes } = require('librechat-data-provider');
 const { isEnabled, removePorts } = require('../utils');
-const keyvRedis = require('../../cache/keyvRedis');
-const User = require('../../models/User');
+const keyvRedis = require('~/cache/keyvRedis');
+const denyRequest = require('./denyRequest');
+const { getLogStores } = require('~/cache');
+const User = require('~/models/User');
 
 const banCache = isEnabled(process.env.USE_REDIS)
   ? new Keyv({ store: keyvRedis })
-  : new Keyv({ namespace: 'bans', ttl: 0 });
+  : new Keyv({ namespace: ViolationTypes.BAN, ttl: 0 });
 const message = 'Your account has been temporarily banned due to violations of our service.';
 
 /**
@@ -28,7 +29,7 @@ const banResponse = async (req, res) => {
   if (!ua.browser.name) {
     return res.status(403).json({ message });
   } else if (baseUrl === '/api/ask' || baseUrl === '/api/edit') {
-    return await denyRequest(req, res, { type: 'ban' });
+    return await denyRequest(req, res, { type: ViolationTypes.BAN });
   }
 
   return res.status(403).json({ message });
@@ -87,7 +88,7 @@ const checkBan = async (req, res, next = () => {}) => {
     return await banResponse(req, res);
   }
 
-  const banLogs = getLogStores('ban');
+  const banLogs = getLogStores(ViolationTypes.BAN);
   const duration = banLogs.opts.ttl;
 
   if (duration <= 0) {

api/server/routes/assistants/actions.js

@@ -1,10 +1,11 @@
 const { v4 } = require('uuid');
 const express = require('express');
-const { actionDelimiter } = require('librechat-data-provider');
-const { initializeClient } = require('~/server/services/Endpoints/assistants');
 const { encryptMetadata, domainParser } = require('~/server/services/ActionService');
+const { actionDelimiter, EModelEndpoint } = require('librechat-data-provider');
+const { initializeClient } = require('~/server/services/Endpoints/assistants');
 const { updateAction, getActions, deleteAction } = require('~/models/Action');
 const { updateAssistant, getAssistant } = require('~/models/Assistant');
+const { withSession } = require('~/server/utils');
 const { logger } = require('~/config');
 
 const router = express.Router();
@@ -46,7 +47,7 @@ router.post('/:assistant_id', async (req, res) => {
 
     let { domain } = metadata;
     /* Azure doesn't support periods in function names */
-    domain = domainParser(req, domain, true);
+    domain = await domainParser(req, domain, true);
 
     if (!domain) {
       return res.status(400).json({ message: 'No domain provided' });
@@ -110,7 +111,8 @@ router.post('/:assistant_id', async (req, res) => {
 
     const promises = [];
     promises.push(
-      updateAssistant(
+      withSession(
+        updateAssistant,
         { assistant_id },
         {
           actions,
@@ -119,7 +121,9 @@ router.post('/:assistant_id', async (req, res) => {
       ),
     );
     promises.push(openai.beta.assistants.update(assistant_id, { tools }));
-    promises.push(updateAction({ action_id }, { metadata, assistant_id, user: req.user.id }));
+    promises.push(
+      withSession(updateAction, { action_id }, { metadata, assistant_id, user: req.user.id }),
+    );
 
     /** @type {[AssistantDocument, Assistant, Action]} */
     const resolved = await Promise.all(promises);
@@ -129,6 +133,15 @@ router.post('/:assistant_id', async (req, res) => {
         delete resolved[2].metadata[field];
       }
     }
+
+    /* Map Azure OpenAI model to the assistant as defined by config */
+    if (req.app.locals[EModelEndpoint.azureOpenAI]?.assistants) {
+      resolved[1] = {
+        ...resolved[1],
+        model: req.body.model,
+      };
+    }
+
     res.json(resolved);
   } catch (error) {
     const message = 'Trouble updating the Assistant Action';
@@ -171,7 +184,7 @@ router.delete('/:assistant_id/:action_id/:model', async (req, res) => {
       return true;
     });
 
-    domain = domainParser(req, domain, true);
+    domain = await domainParser(req, domain, true);
 
     const updatedTools = tools.filter(
       (tool) => !(tool.function && tool.function.name.includes(domain)),
@@ -179,7 +192,8 @@ router.delete('/:assistant_id/:action_id/:model', async (req, res) => {
 
     const promises = [];
     promises.push(
-      updateAssistant(
+      withSession(
+        updateAssistant,
         { assistant_id },
         {
           actions: updatedActions,
@@ -188,7 +202,7 @@ router.delete('/:assistant_id/:action_id/:model', async (req, res) => {
       ),
     );
     promises.push(openai.beta.assistants.update(assistant_id, { tools: updatedTools }));
-    promises.push(deleteAction({ action_id }));
+    promises.push(withSession(deleteAction, { action_id }));
 
     await Promise.all(promises);
     res.status(200).json({ message: 'Action deleted successfully' });

api/server/services/ActionService.js

@@ -1,20 +1,27 @@
-const { AuthTypeEnum, EModelEndpoint, actionDomainSeparator } = require('librechat-data-provider');
+const {
+  AuthTypeEnum,
+  EModelEndpoint,
+  actionDomainSeparator,
+  CacheKeys,
+  Constants,
+} = require('librechat-data-provider');
 const { encryptV2, decryptV2 } = require('~/server/utils/crypto');
 const { getActions } = require('~/models/Action');
+const { getLogStores } = require('~/cache');
 const { logger } = require('~/config');
 
 /**
- * Parses the domain for an action.
+ * Encodes or decodes a domain name to/from base64, or replacing periods with a custom separator.
  *
- * Azure OpenAI Assistants API doesn't support periods in function
- * names due to `[a-zA-Z0-9_-]*` Regex Validation.
+ * Necessary because Azure OpenAI Assistants API doesn't support periods in function
+ * names due to `[a-zA-Z0-9_-]*` Regex Validation, limited to a 64-character maximum.
  *
- * @param {Express.Request} req - Express Request object
- * @param {string} domain - The domain for the actoin
- * @param {boolean} inverse - If true, replaces periods with `actionDomainSeparator`
- * @returns {string} The parsed domain
+ * @param {Express.Request} req - The Express Request object.
+ * @param {string} domain - The domain name to encode/decode.
+ * @param {boolean} inverse - False to decode from base64, true to encode to base64.
+ * @returns {Promise<string>} Encoded or decoded domain string.
  */
-function domainParser(req, domain, inverse = false) {
+async function domainParser(req, domain, inverse = false) {
   if (!domain) {
     return;
   }
@@ -23,11 +30,35 @@ function domainParser(req, domain, inverse = false) {
     return domain;
   }
 
-  if (inverse) {
+  const domainsCache = getLogStores(CacheKeys.ENCODED_DOMAINS);
+  const cachedDomain = await domainsCache.get(domain);
+  if (inverse && cachedDomain) {
+    return domain;
+  }
+
+  if (inverse && domain.length <= Constants.ENCODED_DOMAIN_LENGTH) {
     return domain.replace(/\./g, actionDomainSeparator);
   }
 
-  return domain.replace(actionDomainSeparator, '.');
+  if (inverse) {
+    const modifiedDomain = Buffer.from(domain).toString('base64');
+    const key = modifiedDomain.substring(0, Constants.ENCODED_DOMAIN_LENGTH);
+    await domainsCache.set(key, modifiedDomain);
+    return key;
+  }
+
+  const replaceSeparatorRegex = new RegExp(actionDomainSeparator, 'g');
+
+  if (!cachedDomain) {
+    return domain.replace(replaceSeparatorRegex, '.');
+  }
+
+  try {
+    return Buffer.from(cachedDomain, 'base64').toString('utf-8');
+  } catch (error) {
+    logger.error(`Failed to parse domain (possibly not base64): ${domain}`, error);
+    return domain;
+  }
 }
 
 /**

api/server/services/ActionService.spec.js

@@ -0,0 +1,196 @@
+const { Constants, EModelEndpoint, actionDomainSeparator } = require('librechat-data-provider');
+const { domainParser } = require('./ActionService');
+
+jest.mock('keyv');
+
+const globalCache = {};
+jest.mock('~/cache/getLogStores', () => {
+  return jest.fn().mockImplementation(() => {
+    const EventEmitter = require('events');
+    const { CacheKeys } = require('librechat-data-provider');
+
+    class KeyvMongo extends EventEmitter {
+      constructor(url = 'mongodb://127.0.0.1:27017', options) {
+        super();
+        this.ttlSupport = false;
+        url = url ?? {};
+        if (typeof url === 'string') {
+          url = { url };
+        }
+        if (url.uri) {
+          url = { url: url.uri, ...url };
+        }
+        this.opts = {
+          url,
+          collection: 'keyv',
+          ...url,
+          ...options,
+        };
+      }
+
+      get = async (key) => {
+        return new Promise((resolve) => {
+          resolve(globalCache[key] || null);
+        });
+      };
+
+      set = async (key, value) => {
+        return new Promise((resolve) => {
+          globalCache[key] = value;
+          resolve(true);
+        });
+      };
+    }
+
+    return new KeyvMongo('', {
+      namespace: CacheKeys.ENCODED_DOMAINS,
+      ttl: 0,
+    });
+  });
+});
+
+describe('domainParser', () => {
+  const req = {
+    app: {
+      locals: {
+        [EModelEndpoint.azureOpenAI]: {
+          assistants: true,
+        },
+      },
+    },
+  };
+
+  const reqNoAzure = {
+    app: {
+      locals: {
+        [EModelEndpoint.azureOpenAI]: {
+          assistants: false,
+        },
+      },
+    },
+  };
+
+  const TLD = '.com';
+
+  // Non-azure request
+  it('returns domain as is if not azure', async () => {
+    const domain = `example.com${actionDomainSeparator}test${actionDomainSeparator}`;
+    const result1 = await domainParser(reqNoAzure, domain, false);
+    const result2 = await domainParser(reqNoAzure, domain, true);
+    expect(result1).toEqual(domain);
+    expect(result2).toEqual(domain);
+  });
+
+  // Test for Empty or Null Inputs
+  it('returns undefined for null domain input', async () => {
+    const result = await domainParser(req, null, true);
+    expect(result).toBeUndefined();
+  });
+
+  it('returns undefined for empty domain input', async () => {
+    const result = await domainParser(req, '', true);
+    expect(result).toBeUndefined();
+  });
+
+  // Verify Correct Caching Behavior
+  it('caches encoded domain correctly', async () => {
+    const domain = 'longdomainname.com';
+    const encodedDomain = Buffer.from(domain)
+      .toString('base64')
+      .substring(0, Constants.ENCODED_DOMAIN_LENGTH);
+
+    await domainParser(req, domain, true);
+
+    const cachedValue = await globalCache[encodedDomain];
+    expect(cachedValue).toEqual(Buffer.from(domain).toString('base64'));
+  });
+
+  // Test for Edge Cases Around Length Threshold
+  it('encodes domain exactly at threshold without modification', async () => {
+    const domain = 'a'.repeat(Constants.ENCODED_DOMAIN_LENGTH - TLD.length) + TLD;
+    const expected = domain.replace(/\./g, actionDomainSeparator);
+    const result = await domainParser(req, domain, true);
+    expect(result).toEqual(expected);
+  });
+
+  it('encodes domain just below threshold without modification', async () => {
+    const domain = 'a'.repeat(Constants.ENCODED_DOMAIN_LENGTH - 1 - TLD.length) + TLD;
+    const expected = domain.replace(/\./g, actionDomainSeparator);
+    const result = await domainParser(req, domain, true);
+    expect(result).toEqual(expected);
+  });
+
+  // Test for Unicode Domain Names
+  it('handles unicode characters in domain names correctly when encoding', async () => {
+    const unicodeDomain = 'täst.example.com';
+    const encodedDomain = Buffer.from(unicodeDomain)
+      .toString('base64')
+      .substring(0, Constants.ENCODED_DOMAIN_LENGTH);
+    const result = await domainParser(req, unicodeDomain, true);
+    expect(result).toEqual(encodedDomain);
+  });
+
+  it('decodes unicode domain names correctly', async () => {
+    const unicodeDomain = 'täst.example.com';
+    const encodedDomain = Buffer.from(unicodeDomain).toString('base64');
+    globalCache[encodedDomain.substring(0, Constants.ENCODED_DOMAIN_LENGTH)] = encodedDomain; // Simulate caching
+
+    const result = await domainParser(
+      req,
+      encodedDomain.substring(0, Constants.ENCODED_DOMAIN_LENGTH),
+      false,
+    );
+    expect(result).toEqual(unicodeDomain);
+  });
+
+  // Core Functionality Tests
+  it('returns domain with replaced separators if no cached domain exists', async () => {
+    const domain = 'example.com';
+    const withSeparator = domain.replace(/\./g, actionDomainSeparator);
+    const result = await domainParser(req, withSeparator, false);
+    expect(result).toEqual(domain);
+  });
+
+  it('returns domain with replaced separators when inverse is false and under encoding length', async () => {
+    const domain = 'examp.com';
+    const withSeparator = domain.replace(/\./g, actionDomainSeparator);
+    const result = await domainParser(req, withSeparator, false);
+    expect(result).toEqual(domain);
+  });
+
+  it('replaces periods with actionDomainSeparator when inverse is true and under encoding length', async () => {
+    const domain = 'examp.com';
+    const expected = domain.replace(/\./g, actionDomainSeparator);
+    const result = await domainParser(req, domain, true);
+    expect(result).toEqual(expected);
+  });
+
+  it('encodes domain when length is above threshold and inverse is true', async () => {
+    const domain = 'a'.repeat(Constants.ENCODED_DOMAIN_LENGTH + 1).concat('.com');
+    const result = await domainParser(req, domain, true);
+    expect(result).not.toEqual(domain);
+    expect(result.length).toBeLessThanOrEqual(Constants.ENCODED_DOMAIN_LENGTH);
+  });
+
+  it('returns encoded value if no encoded value is cached, and inverse is false', async () => {
+    const originalDomain = 'example.com';
+    const encodedDomain = Buffer.from(
+      originalDomain.replace(/\./g, actionDomainSeparator),
+    ).toString('base64');
+    const result = await domainParser(req, encodedDomain, false);
+    expect(result).toEqual(encodedDomain);
+  });
+
+  it('decodes encoded value if cached and encoded value is provided, and inverse is false', async () => {
+    const originalDomain = 'example.com';
+    const encodedDomain = await domainParser(req, originalDomain, true);
+    const result = await domainParser(req, encodedDomain, false);
+    expect(result).toEqual(originalDomain);
+  });
+
+  it('handles invalid base64 encoded values gracefully', async () => {
+    const invalidBase64Domain = 'not_base64_encoded';
+    const result = await domainParser(req, invalidBase64Domain, false);
+    expect(result).toEqual(invalidBase64Domain);
+  });
+});

api/server/services/ToolService.js

@@ -274,9 +274,16 @@ async function processRequiredActions(client, requiredActions) {
           })) ?? [];
       }
 
-      const actionSet = actionSets.find((action) =>
-        currentAction.tool.includes(domainParser(client.req, action.metadata.domain, true)),
-      );
+      let actionSet = null;
+      let currentDomain = '';
+      for (let action of actionSets) {
+        const domain = await domainParser(client.req, action.metadata.domain, true);
+        if (currentAction.tool.includes(domain)) {
+          currentDomain = domain;
+          actionSet = action;
+          break;
+        }
+      }
 
       if (!actionSet) {
         // TODO: try `function` if no action set is found
@@ -298,10 +305,8 @@ async function processRequiredActions(client, requiredActions) {
         builders = requestBuilders;
       }
 
-      const functionName = currentAction.tool.replace(
-        `${actionDelimiter}${domainParser(client.req, actionSet.metadata.domain, true)}`,
-        '',
-      );
+      const functionName = currentAction.tool.replace(`${actionDelimiter}${currentDomain}`, '');
+
       const requestBuilder = builders[functionName];
 
       if (!requestBuilder) {

api/server/utils/index.js

@@ -5,6 +5,7 @@ const handleText = require('./handleText');
 const cryptoUtils = require('./crypto');
 const citations = require('./citations');
 const sendEmail = require('./sendEmail');
+const mongoose = require('./mongoose');
 const queue = require('./queue');
 const files = require('./files');
 const math = require('./math');
@@ -14,6 +15,7 @@ module.exports = {
   ...cryptoUtils,
   ...handleText,
   ...citations,
+  ...mongoose,
   countTokens,
   removePorts,
   sendEmail,

api/server/utils/mongoose.js

@@ -0,0 +1,25 @@
+const mongoose = require('mongoose');
+/**
+ * Executes a database operation within a session.
+ * @param {() => Promise<any>} method - The method to execute. This method must accept a session as its first argument.
+ * @param {...any} args - Additional arguments to pass to the method.
+ * @returns {Promise<any>} - The result of the executed method.
+ */
+async function withSession(method, ...args) {
+  const session = await mongoose.startSession();
+  session.startTransaction();
+  try {
+    const result = await method(...args, session);
+    await session.commitTransaction();
+    return result;
+  } catch (error) {
+    if (session.inTransaction()) {
+      await session.abortTransaction();
+    }
+    throw error;
+  } finally {
+    await session.endSession();
+  }
+}
+
+module.exports = { withSession };