@ainetwork/adk-provider-memory-mongodb 0.4.0 → 0.5.0

package/dist/index.cjs

@@ -35,7 +35,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // implements/base.memory.ts
-var import_mongoose10 = __toESM(require("mongoose"), 1);
+var import_mongoose12 = __toESM(require("mongoose"), 1);
 var import_logger2 = require("@ainetwork/adk/utils/logger");
 
 // models/agent.model.ts
@@ -259,14 +259,19 @@ var ThreadObjectSchema = new import_mongoose4.Schema(
       type: String,
       required: false
     },
-    created_at: {
-      type: Number,
-      required: true
+    isPinned: {
+      type: Boolean,
+      required: false,
+      default: false
     },
-    updated_at: {
-      type: Number,
-      required: true
+    workflowId: {
+      type: String,
+      required: false,
+      index: true
     }
+  },
+  {
+    timestamps: true
   }
 );
 var ThreadModel = import_mongoose5.default.model("Thread", ThreadObjectSchema);
@@ -316,8 +321,12 @@ var MessageObjectSchema = new import_mongoose6.Schema(
       type: import_mongoose6.Schema.Types.Mixed,
       default: {}
     }
+  },
+  {
+    timestamps: true
   }
 );
+MessageObjectSchema.index({ threadId: 1, messageId: 1 }, { unique: true });
 var MessageModel = import_mongoose7.default.model("Message", MessageObjectSchema);
 
 // implements/thread.memory.ts
@@ -341,6 +350,8 @@ var MongoDBThread = class {
         userId: thread.userId,
         type: thread.type,
         title: thread.title || "New thread",
+        isPinned: thread.isPinned ?? false,
+        workflowId: thread.workflowId,
         messages: []
       };
       messages.forEach((message) => {
@@ -355,36 +366,39 @@ var MongoDBThread = class {
       return threadObject;
     }, `getThread(${userId}, ${threadId})`);
   }
-  async createThread(type, userId, threadId, title) {
+  async createThread(type, userId, threadId, title, workflowId) {
     return this.executeWithRetry(async () => {
-      const now = Date.now();
       await ThreadModel.create({
         type,
         userId,
         threadId,
         title,
-        updated_at: now,
-        created_at: now
+        workflowId
       });
-      return { type, userId, threadId, title, messages: [] };
+      return { type, userId, threadId, title, workflowId, messages: [] };
     }, `createThread(${userId}, ${threadId})`);
   }
   async addMessagesToThread(userId, threadId, messages) {
     return this.executeWithRetry(async () => {
-      await ThreadModel.updateOne({ threadId, userId }, {
-        updated_at: Date.now()
-      });
-      for (const message of messages) {
-        await MessageModel.create({
-          threadId,
-          messageId: message.messageId,
-          userId,
-          role: message.role,
-          content: message.content,
-          timestamp: message.timestamp,
-          metadata: message.metadata
-        });
+      if (messages.length > 0) {
+        const messageIds = messages.map((m) => m.messageId);
+        await MessageModel.deleteMany({ threadId, userId, messageId: { $in: messageIds } });
+        await MessageModel.insertMany(
+          messages.map((message) => ({
+            threadId,
+            messageId: message.messageId,
+            userId,
+            role: message.role,
+            content: message.content,
+            timestamp: message.timestamp,
+            metadata: message.metadata
+          }))
+        );
       }
+      await ThreadModel.updateOne(
+        { threadId, userId },
+        { $set: { updatedAt: /* @__PURE__ */ new Date() } }
+      );
     }, `addMessagesToThread(${userId}, ${threadId})`);
   }
   async deleteThread(userId, threadId) {
@@ -394,28 +408,43 @@ var MongoDBThread = class {
       await ThreadModel.deleteOne({ userId, threadId }).maxTimeMS(timeout);
     }, `deleteThread(${userId}, ${threadId})`);
   }
-  async listThreads(userId) {
+  async listThreads(userId, filter) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const threads = await ThreadModel.find({ userId }).sort({ updated_at: -1 }).maxTimeMS(timeout);
+      const query = { userId };
+      if (filter?.workflowId) query.workflowId = filter.workflowId;
+      if (filter?.type) query.type = filter.type;
+      const threads = await ThreadModel.find(query).sort({ updatedAt: -1 }).maxTimeMS(timeout);
       const data = threads.map((thread) => {
         return {
           type: thread.type,
           userId,
           threadId: thread.threadId,
           title: thread.title,
-          updatedAt: thread.updated_at
+          isPinned: thread.isPinned ?? false,
+          workflowId: thread.workflowId,
+          createdAt: thread.createdAt?.toISOString(),
+          updatedAt: thread.updatedAt?.toISOString()
         };
       });
       return data;
     }, `listThreads(${userId})`);
   }
+  async updateThreadPin(userId, threadId, isPinned) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await ThreadModel.updateOne(
+        { threadId, userId },
+        { $set: { isPinned } }
+      ).maxTimeMS(timeout);
+    }, `updateThreadPin(${userId}, ${threadId})`);
+  }
 };
 
-// models/workflow.model.ts
+// models/user-workflow.model.ts
 var import_mongoose8 = require("mongoose");
 var import_mongoose9 = __toESM(require("mongoose"), 1);
-var WorkflowObjectSchema = new import_mongoose8.Schema(
+var UserWorkflowObjectSchema = new import_mongoose8.Schema(
   {
     workflowId: {
       type: String,
@@ -423,81 +452,203 @@ var WorkflowObjectSchema = new import_mongoose8.Schema(
       unique: true
     },
     userId: {
-      type: String
+      type: String,
+      required: true,
+      index: true
     },
     title: {
       type: String,
       required: true
     },
     description: {
-      type: String,
-      required: true
+      type: String
     },
     active: {
       type: Boolean,
       required: true,
       default: false
     },
+    templateId: {
+      type: String
+    },
     content: {
       type: String,
       required: true
+    },
+    variables: {
+      type: import_mongoose8.Schema.Types.Mixed
+    },
+    variableValues: {
+      type: import_mongoose8.Schema.Types.Mixed
+    },
+    schedule: {
+      type: String
+    },
+    timezone: {
+      type: String
+    },
+    lastRunAt: {
+      type: Number
+    },
+    nextRunAt: {
+      type: Number
+    },
+    lastThreadId: {
+      type: String
     }
   },
   {
     timestamps: true
   }
 );
-var WorkflowModel = import_mongoose9.default.model("Workflow", WorkflowObjectSchema);
+var UserWorkflowModel = import_mongoose9.default.model(
+  "UserWorkflow",
+  UserWorkflowObjectSchema
+);
 
-// implements/workflow.memory.ts
-var MongoDBWorkflow = class {
+// implements/user-workflow.memory.ts
+var MongoDBUserWorkflow = class {
   executeWithRetry;
   getOperationTimeout;
   constructor(executeWithRetry, getOperationTimeout) {
     this.executeWithRetry = executeWithRetry;
     this.getOperationTimeout = getOperationTimeout;
   }
-  async createWorkflow(workflow) {
+  async createUserWorkflow(workflow) {
     return this.executeWithRetry(async () => {
-      const timeout = this.getOperationTimeout();
-      const created = await WorkflowModel.create(workflow);
+      const created = await UserWorkflowModel.create(workflow);
       return created.toObject();
-    }, "createWorkflow()");
+    }, "createUserWorkflow()");
   }
-  async getWorkflow(workflowId) {
+  async getUserWorkflow(workflowId) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const workflow = await WorkflowModel.findOne({
-        workflowId
-      }).maxTimeMS(timeout).lean();
+      const workflow = await UserWorkflowModel.findOne({ workflowId }).maxTimeMS(timeout).lean();
       return workflow || void 0;
-    }, "getWorkflow()");
+    }, "getUserWorkflow()");
   }
-  async updateWorkflow(workflowId, updates) {
+  async updateUserWorkflow(workflowId, updates) {
     if (!updates.userId) {
-      throw new Error("userId is required for updateWorkflow");
+      throw new Error("userId is required for updateUserWorkflow");
     }
+    const { userId, workflowId: _workflowId, ...mutableUpdates } = updates;
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      await WorkflowModel.updateOne(
-        { workflowId, userId: updates.userId },
-        { $set: updates }
+      await UserWorkflowModel.updateOne(
+        { workflowId, userId },
+        { $set: mutableUpdates }
       ).maxTimeMS(timeout);
-    }, "updateWorkflow()");
+    }, "updateUserWorkflow()");
+  }
+  async deleteUserWorkflow(workflowId, userId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await UserWorkflowModel.deleteOne({ workflowId, userId }).maxTimeMS(timeout);
+    }, "deleteUserWorkflow()");
   }
-  async deleteWorkflow(workflowId, userId) {
+  async listUserWorkflows(userId) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      await WorkflowModel.deleteOne({ workflowId, userId }).maxTimeMS(timeout);
-    }, "deleteWorkflow()");
+      const query = userId ? { userId } : {};
+      const workflows = await UserWorkflowModel.find(query).maxTimeMS(timeout).lean();
+      return workflows;
+    }, "listUserWorkflows()");
   }
-  async listWorkflows(userId) {
+  async listActiveScheduledWorkflows() {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const query = userId ? { $or: [{ userId }, { userId: { $exists: false } }, { userId: null }] } : { $or: [{ userId: { $exists: false } }, { userId: null }] };
-      const workflows = await WorkflowModel.find(query).maxTimeMS(timeout).lean();
+      const workflows = await UserWorkflowModel.find({
+        active: true,
+        schedule: { $exists: true, $ne: null }
+      }).maxTimeMS(timeout).lean();
       return workflows;
-    }, "listWorkflows()");
+    }, "listActiveScheduledWorkflows()");
+  }
+};
+
+// models/workflow-template.model.ts
+var import_mongoose10 = require("mongoose");
+var import_mongoose11 = __toESM(require("mongoose"), 1);
+var WorkflowTemplateObjectSchema = new import_mongoose10.Schema(
+  {
+    templateId: {
+      type: String,
+      required: true,
+      unique: true
+    },
+    title: {
+      type: String,
+      required: true
+    },
+    description: {
+      type: String,
+      required: true
+    },
+    active: {
+      type: Boolean,
+      required: true,
+      default: false
+    },
+    content: {
+      type: String,
+      required: true
+    },
+    variables: {
+      type: import_mongoose10.Schema.Types.Mixed
+    }
+  },
+  {
+    timestamps: true
+  }
+);
+var WorkflowTemplateModel = import_mongoose11.default.model(
+  "WorkflowTemplate",
+  WorkflowTemplateObjectSchema
+);
+
+// implements/workflow-template.memory.ts
+var MongoDBWorkflowTemplate = class {
+  executeWithRetry;
+  getOperationTimeout;
+  constructor(executeWithRetry, getOperationTimeout) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+  async createTemplate(template) {
+    return this.executeWithRetry(async () => {
+      const created = await WorkflowTemplateModel.create(template);
+      return created.toObject();
+    }, "createTemplate()");
+  }
+  async getTemplate(templateId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const template = await WorkflowTemplateModel.findOne({ templateId }).maxTimeMS(timeout).lean();
+      return template || void 0;
+    }, "getTemplate()");
+  }
+  async updateTemplate(templateId, updates) {
+    const { templateId: _templateId, ...mutableUpdates } = updates;
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await WorkflowTemplateModel.updateOne(
+        { templateId },
+        { $set: mutableUpdates }
+      ).maxTimeMS(timeout);
+    }, "updateTemplate()");
+  }
+  async deleteTemplate(templateId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await WorkflowTemplateModel.deleteOne({ templateId }).maxTimeMS(timeout);
+    }, "deleteTemplate()");
+  }
+  async listTemplates() {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const templates = await WorkflowTemplateModel.find().maxTimeMS(timeout).lean();
+      return templates;
+    }, "listTemplates()");
   }
 };
 
@@ -513,18 +664,26 @@ var MongoDBMemory = class _MongoDBMemory {
   connectionConfig;
   eventListenersSetup = false;
   operationTimeoutMS;
+  threadTTLSeconds;
+  orphanCleanupTimer;
   agentMemory;
   intentMemory;
   threadMemory;
-  workflowMemory;
+  workflowTemplateMemory;
+  userWorkflowMemory;
   constructor(config) {
     const cfg = typeof config === "string" ? { uri: config } : config;
     this.uri = cfg.uri;
     this.maxReconnectAttempts = cfg.maxReconnectAttempts ?? 5;
     this.reconnectInterval = cfg.reconnectInterval ?? 5e3;
     this.operationTimeoutMS = cfg.operationTimeoutMS ?? 1e4;
+    if (cfg.threadTTLSeconds !== void 0 && cfg.threadTTLSeconds > 0) {
+      this.threadTTLSeconds = cfg.threadTTLSeconds;
+    }
     this.connectionConfig = {
-      maxPoolSize: cfg.maxPoolSize ?? 1,
+      maxPoolSize: cfg.maxPoolSize ?? 10,
+      minPoolSize: 0,
+      maxIdleTimeMS: 3e4,
       serverSelectionTimeoutMS: cfg.serverSelectionTimeoutMS ?? 3e4,
       socketTimeoutMS: cfg.socketTimeoutMS ?? 45e3,
       connectTimeoutMS: cfg.connectTimeoutMS ?? 3e4,
@@ -549,7 +708,11 @@ var MongoDBMemory = class _MongoDBMemory {
       this.executeWithRetry.bind(this),
       this.getOperationTimeout.bind(this)
     );
-    this.workflowMemory = new MongoDBWorkflow(
+    this.workflowTemplateMemory = new MongoDBWorkflowTemplate(
+      this.executeWithRetry.bind(this),
+      this.getOperationTimeout.bind(this)
+    );
+    this.userWorkflowMemory = new MongoDBUserWorkflow(
       this.executeWithRetry.bind(this),
       this.getOperationTimeout.bind(this)
     );
@@ -563,29 +726,32 @@ var MongoDBMemory = class _MongoDBMemory {
   getIntentMemory() {
     return this.intentMemory;
   }
-  getWorkflowMemory() {
-    return this.workflowMemory;
+  getWorkflowTemplateMemory() {
+    return this.workflowTemplateMemory;
+  }
+  getUserWorkflowMemory() {
+    return this.userWorkflowMemory;
   }
   setupMongooseEventListeners() {
     if (this.eventListenersSetup) return;
     this.eventListenersSetup = true;
-    import_mongoose10.default.connection.on("connected", () => {
+    import_mongoose12.default.connection.on("connected", () => {
       this.connected = true;
       this.reconnectAttempts = 0;
       this.reconnecting = false;
       import_logger2.loggers.agent.info("MongoDB connected successfully");
     });
-    import_mongoose10.default.connection.on("disconnected", () => {
+    import_mongoose12.default.connection.on("disconnected", () => {
       this.connected = false;
       import_logger2.loggers.agent.warn("MongoDB disconnected");
       this.handleDisconnection();
     });
-    import_mongoose10.default.connection.on("error", (error) => {
+    import_mongoose12.default.connection.on("error", (error) => {
       this.connected = false;
       import_logger2.loggers.agent.error("MongoDB connection error:", error);
       this.handleDisconnection();
     });
-    import_mongoose10.default.connection.on("reconnected", () => {
+    import_mongoose12.default.connection.on("reconnected", () => {
       this.connected = true;
       this.reconnectAttempts = 0;
       this.reconnecting = false;
@@ -597,13 +763,13 @@ var MongoDBMemory = class _MongoDBMemory {
       return;
     }
     this.reconnecting = true;
-    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.isConnected) {
+    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.connected) {
       this.reconnectAttempts++;
       import_logger2.loggers.agent.info(
         `Attempting to reconnect to MongoDB (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`
       );
       try {
-        await import_mongoose10.default.connect(this.uri, this.connectionConfig);
+        await import_mongoose12.default.connect(this.uri, this.connectionConfig);
         this.connected = true;
         this.reconnectAttempts = 0;
         this.reconnecting = false;
@@ -622,7 +788,7 @@ var MongoDBMemory = class _MongoDBMemory {
       }
     }
     this.reconnecting = false;
-    if (!this.isConnected) {
+    if (!this.connected) {
       import_logger2.loggers.agent.error(
         `Failed to reconnect to MongoDB after ${this.maxReconnectAttempts} attempts`
       );
@@ -633,20 +799,26 @@ var MongoDBMemory = class _MongoDBMemory {
       return;
     }
     try {
-      await import_mongoose10.default.connect(this.uri, this.connectionConfig);
+      await import_mongoose12.default.connect(this.uri, this.connectionConfig);
       this.connected = true;
       this.reconnectAttempts = 0;
+      await this.setupTTLIndex();
+      this.startOrphanCleanup();
     } catch (error) {
       import_logger2.loggers.agent.error("Failed to connect to MongoDB:", error);
       throw error;
     }
   }
   async disconnect() {
-    if (!this.isConnected) {
+    if (!this.connected) {
       return;
     }
     try {
-      await import_mongoose10.default.disconnect();
+      if (this.orphanCleanupTimer) {
+        clearInterval(this.orphanCleanupTimer);
+        this.orphanCleanupTimer = void 0;
+      }
+      await import_mongoose12.default.disconnect();
       this.connected = false;
     } catch (error) {
       import_logger2.loggers.agent.error("Failed to disconnect from MongoDB:", error);
@@ -657,7 +829,7 @@ var MongoDBMemory = class _MongoDBMemory {
     return this.connected;
   }
   async ensureConnection() {
-    if (!this.isConnected && !this.reconnecting) {
+    if (!this.connected && !this.reconnecting) {
       await this.connect();
     }
     const maxWaitTime = 3e4;
@@ -665,10 +837,66 @@ var MongoDBMemory = class _MongoDBMemory {
     while (this.reconnecting && Date.now() - startTime < maxWaitTime) {
       await new Promise((resolve) => setTimeout(resolve, 100));
     }
-    if (!this.isConnected) {
+    if (!this.connected) {
       throw new Error("MongoDB is not connected and reconnection failed");
     }
   }
+  async setupTTLIndex() {
+    if (this.threadTTLSeconds === void 0) return;
+    try {
+      const db = import_mongoose12.default.connection.db;
+      if (!db) return;
+      const collection = db.collection("threads");
+      const indexes = await collection.indexes();
+      const existingTTL = indexes.find(
+        (idx) => idx.key?.updatedAt !== void 0 && idx.expireAfterSeconds !== void 0
+      );
+      if (existingTTL) {
+        if (existingTTL.expireAfterSeconds !== this.threadTTLSeconds) {
+          await db.command({
+            collMod: "threads",
+            index: { keyPattern: { updatedAt: 1 }, expireAfterSeconds: this.threadTTLSeconds }
+          });
+          import_logger2.loggers.agent.info(`Thread TTL index updated to ${this.threadTTLSeconds} seconds`);
+        }
+      } else {
+        await collection.createIndex(
+          { updatedAt: 1 },
+          { expireAfterSeconds: this.threadTTLSeconds }
+        );
+        import_logger2.loggers.agent.info(`Thread TTL index created with ${this.threadTTLSeconds} seconds`);
+      }
+    } catch (error) {
+      import_logger2.loggers.agent.error("Failed to setup TTL index:", error);
+    }
+  }
+  startOrphanCleanup() {
+    if (this.threadTTLSeconds === void 0) return;
+    if (this.orphanCleanupTimer) return;
+    const intervalMs = Math.max(6e4, Math.min(this.threadTTLSeconds * 500, 36e5));
+    this.orphanCleanupTimer = setInterval(() => {
+      this.cleanupOrphanedMessages().catch((error) => {
+        import_logger2.loggers.agent.error("Orphaned message cleanup failed:", error);
+      });
+    }, intervalMs);
+    import_logger2.loggers.agent.info(`Orphaned message cleanup scheduled every ${Math.round(intervalMs / 1e3)}s`);
+  }
+  async cleanupOrphanedMessages() {
+    if (!this.connected) return;
+    try {
+      const db = import_mongoose12.default.connection.db;
+      if (!db) return;
+      const existingThreadIds = await db.collection("threads").distinct("threadId");
+      const result = await MessageModel.deleteMany({
+        threadId: { $nin: existingThreadIds }
+      });
+      if (result.deletedCount > 0) {
+        import_logger2.loggers.agent.info(`Cleaned up ${result.deletedCount} orphaned messages`);
+      }
+    } catch (error) {
+      import_logger2.loggers.agent.error("Failed to cleanup orphaned messages:", error);
+    }
+  }
   /**
    * Get the operation timeout in milliseconds
    */
@@ -688,6 +916,25 @@ var MongoDBMemory = class _MongoDBMemory {
         import_logger2.loggers.agent.error(`${operationName} exceeded time limit`);
         throw error;
       }
+      if (error.code === 261 || error.codeName === "TooManyLogicalSessions") {
+        import_logger2.loggers.agent.warn(
+          `${operationName} failed due to too many sessions, disconnecting to release sessions...`
+        );
+        try {
+          await import_mongoose12.default.disconnect();
+          this.connected = false;
+        } catch (disconnectError) {
+          import_logger2.loggers.agent.error("Failed to disconnect during session cleanup:", disconnectError);
+        }
+        await new Promise((resolve) => setTimeout(resolve, 5e3));
+        await this.ensureConnection();
+        try {
+          return await operation();
+        } catch (retryError) {
+          import_logger2.loggers.agent.error(`${operationName} failed after session cleanup retry:`, retryError);
+          throw retryError;
+        }
+      }
       if (error.name === "MongoNetworkError" || error.name === "MongoServerError" || error.message?.includes("connection") || error.message?.includes("disconnect")) {
         import_logger2.loggers.agent.warn(
           `${operationName} failed due to connection issue, attempting reconnection...`