totem-llm 0.1.0

package/server/utils/MCP/index.js

@@ -0,0 +1,269 @@
+const MCPHypervisor = require("./hypervisor");
+
+class MCPCompatibilityLayer extends MCPHypervisor {
+  static _instance;
+
+  constructor() {
+    super();
+    if (MCPCompatibilityLayer._instance) return MCPCompatibilityLayer._instance;
+    MCPCompatibilityLayer._instance = this;
+  }
+
+  /**
+   * Get all of the active MCP servers as plugins we can load into agents.
+   * This will also boot all MCP servers if they have not been started yet.
+   * @returns {Promise<string[]>} Array of flow names in @@mcp_{name} format
+   */
+  async activeMCPServers() {
+    await this.bootMCPServers();
+    return Object.keys(this.mcps).flatMap((name) => `@@mcp_${name}`);
+  }
+
+  /**
+   * Convert an MCP server name to an AnythingLLM Agent plugin
+   * @param {string} name - The base name of the MCP server to convert - not the tool name. eg: `docker-mcp` not `docker-mcp:list-containers`
+   * @param {Object} aibitat - The aibitat object to pass to the plugin
+   * @returns {Promise<{name: string, description: string, plugin: Function}[]|null>} Array of plugin configurations or null if not found
+   */
+  async convertServerToolsToPlugins(name, _aibitat = null) {
+    const mcp = this.mcps[name];
+    if (!mcp) return null;
+
+    let tools;
+    try {
+      const response = await mcp.listTools();
+      tools = response.tools;
+    } catch (error) {
+      this.log(`Failed to list tools for MCP server ${name}:`, error);
+      return null;
+    }
+    if (!tools || !tools.length) return null;
+
+    const suppressedTools = this.getSuppressedTools(name);
+    const totalTools = tools.length;
+    tools = tools.filter((tool) => !suppressedTools.includes(tool.name));
+    const suppressedCount = totalTools - tools.length;
+
+    if (suppressedCount > 0) {
+      this.log(
+        `MCP server ${name}: ${suppressedCount} tool(s) suppressed, ${tools.length} tool(s) enabled`
+      );
+    }
+
+    if (!tools.length) {
+      this.log(`MCP server ${name}: All tools are suppressed, skipping`);
+      return null;
+    }
+
+    const plugins = [];
+    for (const tool of tools) {
+      plugins.push({
+        name: `${name}-${tool.name}`,
+        description: tool.description,
+        plugin: function () {
+          return {
+            name: `${name}-${tool.name}`,
+            setup: (aibitat) => {
+              aibitat.function({
+                super: aibitat,
+                name: `${name}-${tool.name}`,
+                controller: new AbortController(),
+                description: tool.description,
+                isMCPTool: true,
+                examples: [],
+                parameters: {
+                  $schema: "http://json-schema.org/draft-07/schema#",
+                  ...tool.inputSchema,
+                },
+                handler: async function (args = {}) {
+                  try {
+                    const mcpLayer = new MCPCompatibilityLayer();
+                    const currentMcp = mcpLayer.mcps[name];
+                    if (!currentMcp)
+                      throw new Error(
+                        `MCP server ${name} is not currently running`
+                      );
+
+                    aibitat.handlerProps.log(
+                      `Executing MCP server: ${name}:${tool.name} with args:`,
+                      args
+                    );
+                    aibitat.introspect(
+                      `Executing MCP server: ${name} with ${JSON.stringify(args, null, 2)}`
+                    );
+                    const result = await currentMcp.callTool({
+                      name: tool.name,
+                      arguments: args,
+                    });
+                    aibitat.handlerProps.log(
+                      `MCP server: ${name}:${tool.name} completed successfully`,
+                      result
+                    );
+                    aibitat.introspect(
+                      `MCP server: ${name}:${tool.name} completed successfully`
+                    );
+                    return MCPCompatibilityLayer.returnMCPResult(result);
+                  } catch (error) {
+                    aibitat.handlerProps.log(
+                      `MCP server: ${name}:${tool.name} failed with error:`,
+                      error
+                    );
+                    aibitat.introspect(
+                      `MCP server: ${name}:${tool.name} failed with error:`,
+                      error
+                    );
+                    return `The tool ${name}:${tool.name} failed with error: ${error?.message || "An unknown error occurred"}`;
+                  }
+                },
+              });
+            },
+          };
+        },
+        toolName: `${name}:${tool.name}`,
+      });
+    }
+
+    return plugins;
+  }
+
+  /**
+   * Returns the MCP servers that were loaded or attempted to be loaded
+   * so that we can display them in the frontend for review or error logging.
+   * @returns {Promise<{
+   *   name: string,
+   *   running: boolean,
+   *   tools: {name: string, description: string, inputSchema: Object}[],
+   *   process: {pid: number, cmd: string}|null,
+   *   error: string|null
+   * }[]>} - The active MCP servers
+   */
+  async servers() {
+    await this.bootMCPServers();
+    const servers = [];
+    for (const [name, result] of Object.entries(this.mcpLoadingResults)) {
+      const config = this.mcpServerConfigs.find((s) => s.name === name);
+
+      if (result.status === "failed") {
+        servers.push({
+          name,
+          config: config?.server || null,
+          running: false,
+          tools: [],
+          error: result.message,
+          process: null,
+        });
+        continue;
+      }
+
+      const mcp = this.mcps[name];
+      if (!mcp) {
+        delete this.mcpLoadingResults[name];
+        delete this.mcps[name];
+        continue;
+      }
+
+      const online = !!(await mcp.ping());
+      const tools = (online ? (await mcp.listTools()).tools : []).filter(
+        (tool) => !tool.name.startsWith("handle_mcp_connection_mcp_")
+      );
+      servers.push({
+        name,
+        config: config?.server || null,
+        running: online,
+        tools,
+        error: null,
+        process: {
+          pid: mcp.transport?.process?.pid || null,
+        },
+      });
+    }
+    return servers;
+  }
+
+  /**
+   * Toggle the MCP server (start or stop)
+   * @param {string} name - The name of the MCP server to toggle
+   * @returns {Promise<{success: boolean, error: string | null}>}
+   */
+  async toggleServerStatus(name) {
+    const server = this.mcpServerConfigs.find((s) => s.name === name);
+    if (!server)
+      return {
+        success: false,
+        error: `MCP server ${name} not found in config file.`,
+      };
+    const mcp = this.mcps[name];
+    const online = !!mcp ? !!(await mcp.ping()) : false; // If the server is not in the mcps object, it is not running
+
+    if (online) {
+      const killed = this.pruneMCPServer(name);
+      return {
+        success: killed,
+        error: killed ? null : `Failed to kill MCP server: ${name}`,
+      };
+    } else {
+      const startupResult = await this.startMCPServer(name);
+      return { success: startupResult.success, error: startupResult.error };
+    }
+  }
+
+  /**
+   * Delete the MCP server - will also remove it from the config file
+   * @param {string} name - The name of the MCP server to delete
+   * @returns {Promise<{success: boolean, error: string | null}>}
+   */
+  async deleteServer(name) {
+    const server = this.mcpServerConfigs.find((s) => s.name === name);
+    if (!server)
+      return {
+        success: false,
+        error: `MCP server ${name} not found in config file.`,
+      };
+
+    const mcp = this.mcps[name];
+    const online = !!mcp ? !!(await mcp.ping()) : false; // If the server is not in the mcps object, it is not running
+    if (online) this.pruneMCPServer(name);
+    this.removeMCPServerFromConfig(name);
+
+    delete this.mcps[name];
+    delete this.mcpLoadingResults[name];
+    this.log(`MCP server was killed and removed from config file: ${name}`);
+    return { success: true, error: null };
+  }
+
+  /**
+   * Return the result of an MCP server call as a string
+   * This will handle circular references and bigints since an MCP server can return any type of data.
+   * @param {Object} result - The result to return
+   * @returns {string} The result as a string
+   */
+  static returnMCPResult(result) {
+    if (typeof result !== "object" || result === null) return String(result);
+
+    const seen = new WeakSet();
+    try {
+      return JSON.stringify(result, (key, value) => {
+        if (typeof value === "bigint") return value.toString();
+        if (typeof value === "object" && value !== null) {
+          if (seen.has(value)) return "[Circular]";
+          seen.add(value);
+        }
+        return value;
+      });
+    } catch (e) {
+      return `[Unserializable: ${e.message}]`;
+    }
+  }
+
+  /**
+   * Toggle tool suppression for an MCP server
+   * @param {string} serverName - The name of the MCP server
+   * @param {string} toolName - The name of the tool to toggle
+   * @param {boolean} enabled - Whether the tool should be enabled (true) or suppressed (false)
+   * @returns {Promise<{success: boolean, error: string | null, suppressedTools: string[]}>}
+   */
+  async toggleToolSuppression(serverName, toolName, enabled) {
+    return this.updateSuppressedTools(serverName, toolName, enabled);
+  }
+}
+module.exports = MCPCompatibilityLayer;

package/server/utils/PasswordRecovery/index.js

@@ -0,0 +1,103 @@
+const bcrypt = require("bcryptjs");
+const { v4, validate } = require("uuid");
+const { User } = require("../../models/user");
+const {
+  RecoveryCode,
+  PasswordResetToken,
+} = require("../../models/passwordRecovery");
+
+async function generateRecoveryCodes(userId) {
+  const newRecoveryCodes = [];
+  const plainTextCodes = [];
+  for (let i = 0; i < 4; i++) {
+    const code = v4();
+    const hashedCode = bcrypt.hashSync(code, 10);
+    newRecoveryCodes.push({
+      user_id: userId,
+      code_hash: hashedCode,
+    });
+    plainTextCodes.push(code);
+  }
+
+  const { error } = await RecoveryCode.createMany(newRecoveryCodes);
+  if (!!error) throw new Error(error);
+
+  const { user: success } = await User._update(userId, {
+    seen_recovery_codes: true,
+  });
+  if (!success) throw new Error("Failed to generate user recovery codes!");
+
+  return plainTextCodes;
+}
+
+async function recoverAccount(username = "", recoveryCodes = []) {
+  const user = await User.get({ username: String(username) });
+  if (!user) return { success: false, error: "Invalid recovery codes." };
+
+  // If hashes do not exist for a user
+  // because this is a user who has not logged out and back in since upgrade.
+  const allUserHashes = await RecoveryCode.hashesForUser(user.id);
+  if (allUserHashes.length < 4)
+    return { success: false, error: "Invalid recovery codes." };
+
+  // If they tried to send more than two unique codes, we only take the first two
+  const uniqueRecoveryCodes = [...new Set(recoveryCodes)]
+    .map((code) => code.trim())
+    .filter((code) => validate(code)) // we know that any provided code must be a uuid v4.
+    .slice(0, 2);
+  if (uniqueRecoveryCodes.length !== 2)
+    return { success: false, error: "Invalid recovery codes." };
+
+  const validCodes = uniqueRecoveryCodes.every((code) => {
+    let valid = false;
+    allUserHashes.forEach((hash) => {
+      if (bcrypt.compareSync(code, hash)) valid = true;
+    });
+    return valid;
+  });
+  if (!validCodes) return { success: false, error: "Invalid recovery codes." };
+
+  const { passwordResetToken, error } = await PasswordResetToken.create(
+    user.id
+  );
+  if (!!error) return { success: false, error };
+  return { success: true, resetToken: passwordResetToken.token };
+}
+
+async function resetPassword(token, _newPassword = "", confirmPassword = "") {
+  const newPassword = String(_newPassword).trim(); // No spaces in passwords
+  if (!newPassword) throw new Error("Invalid password.");
+  if (newPassword !== String(confirmPassword))
+    throw new Error("Passwords do not match");
+
+  const resetToken = await PasswordResetToken.findUnique({
+    token: String(token),
+  });
+  if (!resetToken || resetToken.expiresAt < new Date()) {
+    return { success: false, message: "Invalid reset token" };
+  }
+
+  // JOI password rules will be enforced inside .update.
+  const { error } = await User.update(resetToken.user_id, {
+    password: newPassword,
+  });
+
+  // seen_recovery_codes is not publicly writable
+  // so we have to do direct update here
+  await User._update(resetToken.user_id, {
+    seen_recovery_codes: false,
+  });
+
+  if (error) return { success: false, message: error };
+  await PasswordResetToken.deleteMany({ user_id: resetToken.user_id });
+  await RecoveryCode.deleteMany({ user_id: resetToken.user_id });
+
+  // New codes are provided on first new login.
+  return { success: true, message: "Password reset successful" };
+}
+
+module.exports = {
+  recoverAccount,
+  resetPassword,
+  generateRecoveryCodes,
+};

package/server/utils/PushNotifications/index.js

@@ -0,0 +1,228 @@
+const webpush = require("web-push");
+const fs = require("fs");
+const path = require("path");
+const { User } = require("../../models/user");
+const { SystemSettings } = require("../../models/systemSettings");
+const { safeJsonParse } = require("../http");
+
+/**
+ * For more options, see:
+ * https://developer.mozilla.org/en-US/docs/Web/API/ServiceWorkerRegistration/showNotification#options
+ * @typedef {Object} PushNotificationPayload
+ * @property {string} title - The title of the notification.
+ * @property {string} body - The message of the notification.
+ * @property {Object} data - Unstructured data for the notification. Use this for anything non-standard.
+ * @property {string} [data.onClickUrl] - The URL to open when the notification is clicked. Note: Can be relative or absolute.
+ * @property {Object[]} actions - The actions for the notification.
+ * @property {string} [actions[].action] - The action to perform when the notification is clicked. Handled in the service worker.
+ * @property {string} [actions[].title] - The title of the action to show in the Options dropdown
+ * @property {string} image - A string containing the URL of an image to be displayed in the notification.
+ */
+
+class PushNotifications {
+  static mailTo = "anythingllm@localhost";
+  /**
+   * @type {PushNotifications}
+   */
+  static instance = null;
+
+  /**
+   * The VAPID keys for the push notification service.
+   * @type {{publicKey: string | null, privateKey: string | null}}
+   */
+  #vapidKeys = {
+    publicKey: null,
+    privateKey: null,
+  };
+
+  /**
+   * The subscriptions for the push notification service.
+   * @type {Map<string, Object>}
+   */
+  #subscriptions = new Map();
+
+  constructor() {
+    if (PushNotifications.instance) return PushNotifications.instance;
+    PushNotifications.instance = this;
+  }
+
+  #log(text, ...args) {
+    console.log(`\x1b[36m[PushNotifications]\x1b[0m ${text}`, ...args);
+  }
+
+  get pushService() {
+    try {
+      const vapidKeys = this.existingVapidKeys;
+      if (!vapidKeys.publicKey || !vapidKeys.privateKey)
+        throw new Error(
+          "VAPID keys not found. Make sure they are generated in the main process first."
+        );
+      webpush.setVapidDetails(
+        `mailto:${this.mailTo}`,
+        vapidKeys.publicKey,
+        vapidKeys.privateKey
+      );
+      return webpush;
+    } catch (e) {
+      console.error("Failed to set VAPID details", e);
+      return null;
+    }
+  }
+
+  get storagePath() {
+    return process.env.NODE_ENV === "development"
+      ? path.resolve(__dirname, `../../storage`, "push-notifications")
+      : path.resolve(process.env.STORAGE_DIR, "push-notifications");
+  }
+
+  get primarySubscriptionPath() {
+    return path.resolve(this.storagePath, `primary-subscription.json`);
+  }
+
+  get existingVapidKeys() {
+    // Already loaded and binded to the instance
+    if (this.#vapidKeys.publicKey && this.#vapidKeys.privateKey)
+      return this.#vapidKeys;
+
+    const vapidKeysPath = path.resolve(this.storagePath, `vapid-keys.json`);
+    if (!fs.existsSync(vapidKeysPath))
+      return { publicKey: null, privateKey: null };
+
+    const existingVapidKeys = JSON.parse(
+      fs.readFileSync(vapidKeysPath, "utf8")
+    );
+    this.#log(`Loaded existing VAPID keys!`);
+    this.#vapidKeys.publicKey = existingVapidKeys.publicKey;
+    this.#vapidKeys.privateKey = existingVapidKeys.privateKey;
+    return this.#vapidKeys;
+  }
+
+  get publicVapidKey() {
+    return this.existingVapidKeys.publicKey;
+  }
+
+  /**
+   * Load the subscriptions for the push notification service.
+   * In single user mode, the subscription is stored in the primary-subscription.json file.
+   * In multi user mode, the subscriptions are stored in the database so we grab them from there
+   * and store them in the #subscriptions map for reference later.
+   * @returns {Promise<void>}
+   */
+  async loadSubscriptions() {
+    const isMultiUserMode = await SystemSettings.isMultiUserMode();
+    if (isMultiUserMode) {
+      const users = await User._where({
+        web_push_subscription_config: { not: null },
+      });
+      for (const user of users) {
+        const subscription = safeJsonParse(
+          user.web_push_subscription_config,
+          null
+        );
+        if (subscription) this.#subscriptions.set(user.id, subscription);
+      }
+      this.#log(`Loaded ${this.#subscriptions.size} existing subscriptions.`);
+      return;
+    }
+
+    this.#log("Loading single user mode subscriptions...");
+    if (!fs.existsSync(this.primarySubscriptionPath)) return;
+    const subscription = JSON.parse(
+      fs.readFileSync(this.primarySubscriptionPath, "utf8")
+    );
+    if (subscription) this.#subscriptions.set("primary", subscription);
+    this.#log(`Loaded primary user's existing subscription.`);
+  }
+
+  /**
+   * Register a new subscription for a user.
+   * In single user mode, the userId is mapped to "primary"
+   * In multi user mode, the userId is the user's id in the database
+   *
+   * @param {Object|null} user - The user to register the subscription for.
+   * @param {Object} subscription - The subscription to register.
+   * @returns {Promise<PushNotifications>}
+   */
+  async registerSubscription(user = null, subscription) {
+    let userId = user?.id || "primary";
+    this.#subscriptions.set(userId, subscription);
+
+    // If this was a real user, write the subscription to the database
+    if (!!user) {
+      await User._update(user.id, {
+        web_push_subscription_config: JSON.stringify(subscription),
+      });
+      this.#log(`Registered or updated subscription for user - ${user.id}`);
+    } else {
+      if (!fs.existsSync(this.storagePath))
+        fs.mkdirSync(this.storagePath, { recursive: true });
+      fs.writeFileSync(
+        this.primarySubscriptionPath,
+        JSON.stringify(subscription, null, 2)
+      );
+      this.#log(`Registered or updated primary user's subscription.`);
+    }
+    return this;
+  }
+
+  /**
+   * Send a push notification to all subscribed clients.
+   * @param {Object} options - The options for the notification.
+   * @param {"primary"|number} [options.to] - The subscription to send the notification to. "all" sends to all subscriptions, "primary" sends to the primary user (single user mode only), a number sends subscription to specific user
+   * @param {PushNotificationPayload} [options.payload] - The payload to send to the clients.
+   * @returns {Promise<void>}
+   */
+  sendNotification({ to = "primary", payload = {} } = {}) {
+    if (this.#subscriptions.size === 0)
+      return this.#log(".sendNotification() - No subscriptions found");
+    if (!this.#subscriptions.has(to))
+      return this.#log(
+        `.sendNotification() - Subscription for user ${to} not found`
+      );
+    this.#log(`.sendNotification() - Sending notification to user ${to}`);
+    return this.pushService
+      .sendNotification(this.#subscriptions.get(to), JSON.stringify(payload))
+      .then((res) => {
+        this.#log(
+          `.sendNotification() - Delivered (status: ${res.statusCode})`
+        );
+      })
+      .catch((err) => {
+        this.#log(`.sendNotification() - Failed: ${err.message}`);
+      });
+  }
+
+  /**
+   * Setup the push notification service.
+   * This will generate new VAPID keys if they don't exist and save them to the storage path.
+   * It will also load the subscriptions from the database or the primary-subscription.json file.
+   * @returns {Promise<void>}
+   */
+  static async setupPushNotificationService() {
+    const instance = PushNotifications.instance;
+    const existingVapidKeys = instance.existingVapidKeys;
+
+    if (!existingVapidKeys.publicKey || !existingVapidKeys.privateKey) {
+      instance.#log("Generating new VAPID keys...");
+      const vapidKeys = webpush.generateVAPIDKeys();
+      instance.#vapidKeys.publicKey = vapidKeys.publicKey;
+      instance.#vapidKeys.privateKey = vapidKeys.privateKey;
+      instance.#log(`New VAPID keys generated!`);
+      if (!fs.existsSync(instance.storagePath))
+        fs.mkdirSync(instance.storagePath, { recursive: true });
+      fs.writeFileSync(
+        path.resolve(instance.storagePath, `vapid-keys.json`),
+        JSON.stringify(vapidKeys, null, 2)
+      );
+    }
+
+    await instance.loadSubscriptions();
+    instance.pushService;
+    return;
+  }
+}
+
+module.exports = {
+  pushNotificationService: new PushNotifications(),
+  PushNotifications,
+};