@mastra/mongodb 1.5.4 → 1.5.5

package/dist/index.js

@@ -12,7 +12,7 @@ import { saveScorePayloadSchema } from '@mastra/core/evals';
 
 // package.json
 var package_default = {
-  version: "1.5.4"};
+  version: "1.5.5"};
 var MongoDBFilterTranslator = class extends BaseFilterTranslator {
   getSupportedOperators() {
     return {
@@ -2905,22 +2905,37 @@ var MemoryStorageMongoDB = class _MemoryStorageMongoDB extends MemoryStorage {
     if (row.type && row.type !== "v2") result.type = row.type;
     return result;
   }
+  _sortMessages(messages, field, direction) {
+    return messages.sort((a, b) => {
+      const isDateField = field === "createdAt" || field === "updatedAt";
+      const aValue = isDateField ? new Date(a[field]).getTime() : a[field];
+      const bValue = isDateField ? new Date(b[field]).getTime() : b[field];
+      if (typeof aValue === "number" && typeof bValue === "number") {
+        return direction === "ASC" ? aValue - bValue : bValue - aValue;
+      }
+      return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
+    });
+  }
   async _getIncludedMessages({ include }) {
     if (!include || include.length === 0) return null;
     const collection = await this.getCollection(TABLE_MESSAGES);
+    const targetIds = include.map((inc) => inc.id).filter(Boolean);
+    if (targetIds.length === 0) return null;
+    const targetDocs = await collection.find({ id: { $in: targetIds } }, { projection: { id: 1, thread_id: 1, createdAt: 1 } }).toArray();
+    if (targetDocs.length === 0) return null;
+    const targetMap = new Map(
+      targetDocs.map((doc) => [doc.id, { threadId: doc.thread_id, createdAt: doc.createdAt }])
+    );
     const includedMessages = [];
     for (const inc of include) {
       const { id, withPreviousMessages = 0, withNextMessages = 0 } = inc;
-      const targetMessage = await collection.findOne({ id });
-      if (!targetMessage) continue;
-      const messageThreadId = targetMessage.thread_id;
-      const allMessages = await collection.find({ thread_id: messageThreadId }).sort({ createdAt: 1 }).toArray();
-      const targetIndex = allMessages.findIndex((msg) => msg.id === id);
-      if (targetIndex === -1) continue;
-      const startIndex = Math.max(0, targetIndex - withPreviousMessages);
-      const endIndex = Math.min(allMessages.length - 1, targetIndex + withNextMessages);
-      for (let i = startIndex; i <= endIndex; i++) {
-        includedMessages.push(allMessages[i]);
+      const target = targetMap.get(id);
+      if (!target) continue;
+      const prevMessages = await collection.find({ thread_id: target.threadId, createdAt: { $lte: target.createdAt } }).sort({ createdAt: -1, id: -1 }).limit(withPreviousMessages + 1).toArray();
+      includedMessages.push(...prevMessages);
+      if (withNextMessages > 0) {
+        const nextMessages = await collection.find({ thread_id: target.threadId, createdAt: { $gt: target.createdAt } }).sort({ createdAt: 1, id: 1 }).limit(withNextMessages).toArray();
+        includedMessages.push(...nextMessages);
       }
     }
     const seen = /* @__PURE__ */ new Set();
@@ -2996,6 +3011,20 @@ var MemoryStorageMongoDB = class _MemoryStorageMongoDB extends MemoryStorage {
         const endOp = filter.dateRange.endExclusive ? "$lt" : "$lte";
         query.createdAt = { ...query.createdAt, [endOp]: formatDateForMongoDB(filter.dateRange.end) };
       }
+      if (perPage === 0 && (!include || include.length === 0)) {
+        return { messages: [], total: 0, page, perPage: perPageForResponse, hasMore: false };
+      }
+      if (perPage === 0 && include && include.length > 0) {
+        const includeMessages = await this._getIncludedMessages({ include });
+        const list2 = new MessageList().add(includeMessages ?? [], "memory");
+        return {
+          messages: this._sortMessages(list2.get.all.db(), field, direction),
+          total: 0,
+          page,
+          perPage: perPageForResponse,
+          hasMore: false
+        };
+      }
       const total = await collection.countDocuments(query);
       const messages = [];
       if (perPage !== 0) {
@@ -3029,16 +3058,7 @@ var MemoryStorageMongoDB = class _MemoryStorageMongoDB extends MemoryStorage {
         }
       }
       const list = new MessageList().add(messages, "memory");
-      let finalMessages = list.get.all.db();
-      finalMessages = finalMessages.sort((a, b) => {
-        const isDateField = field === "createdAt" || field === "updatedAt";
-        const aValue = isDateField ? new Date(a[field]).getTime() : a[field];
-        const bValue = isDateField ? new Date(b[field]).getTime() : b[field];
-        if (typeof aValue === "number" && typeof bValue === "number") {
-          return direction === "ASC" ? aValue - bValue : bValue - aValue;
-        }
-        return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
-      });
+      const finalMessages = this._sortMessages(list.get.all.db(), field, direction);
       const threadIdSet = new Set(threadIds);
       const returnedThreadMessageIds = new Set(
         finalMessages.filter((m) => m.threadId && threadIdSet.has(m.threadId)).map((m) => m.id)
@@ -3116,6 +3136,23 @@ var MemoryStorageMongoDB = class _MemoryStorageMongoDB extends MemoryStorage {
         const endOp = filter.dateRange.endExclusive ? "$lt" : "$lte";
         query.createdAt = { ...query.createdAt, [endOp]: formatDateForMongoDB(filter.dateRange.end) };
       }
+      if (perPage === 0 && (!include || include.length === 0)) {
+        return { messages: [], total: 0, page, perPage: perPageForResponse, hasMore: false };
+      }
+      if (perPage === 0 && include && include.length > 0) {
+        const includeMessages = await this._getIncludedMessages({ include });
+        if (!includeMessages || includeMessages.length === 0) {
+          return { messages: [], total: 0, page, perPage: perPageForResponse, hasMore: false };
+        }
+        const list2 = new MessageList().add(includeMessages, "memory");
+        return {
+          messages: this._sortMessages(list2.get.all.db(), field, direction),
+          total: 0,
+          page,
+          perPage: perPageForResponse,
+          hasMore: false
+        };
+      }
       const total = await collection.countDocuments(query);
       const messages = [];
       if (perPage !== 0) {
@@ -3149,16 +3186,7 @@ var MemoryStorageMongoDB = class _MemoryStorageMongoDB extends MemoryStorage {
         }
       }
       const list = new MessageList().add(messages, "memory");
-      let finalMessages = list.get.all.db();
-      finalMessages = finalMessages.sort((a, b) => {
-        const isDateField = field === "createdAt" || field === "updatedAt";
-        const aValue = isDateField ? new Date(a[field]).getTime() : a[field];
-        const bValue = isDateField ? new Date(b[field]).getTime() : b[field];
-        if (typeof aValue === "number" && typeof bValue === "number") {
-          return direction === "ASC" ? aValue - bValue : bValue - aValue;
-        }
-        return direction === "ASC" ? String(aValue).localeCompare(String(bValue)) : String(bValue).localeCompare(String(aValue));
-      });
+      const finalMessages = this._sortMessages(list.get.all.db(), field, direction);
       const hasMore = perPageInput !== false && offset + perPage < total;
       return {
         messages: finalMessages,
@@ -4993,7 +5021,8 @@ Note: This migration may take some time for large collections.
   }
   async listTraces(args) {
     const { filters, pagination, orderBy } = listTracesArgsSchema.parse(args);
-    const { page, perPage } = pagination;
+    const page = pagination?.page ?? 0;
+    const perPage = pagination?.perPage ?? 10;
     try {
       const collection = await this.getCollection(TABLE_SPANS);
       const mongoFilter = {
@@ -5098,8 +5127,8 @@ Note: This migration may take some time for large collections.
       if (andConditions.length) {
         mongoFilter.$and = andConditions;
       }
-      const sortField = orderBy.field;
-      const sortDirection = orderBy.direction === "ASC" ? 1 : -1;
+      const sortField = orderBy?.field ?? "startedAt";
+      const sortDirection = (orderBy?.direction ?? "DESC") === "ASC" ? 1 : -1;
       if (filters?.hasChildError !== void 0) {
         const pipeline = [
           { $match: mongoFilter },