@mastra/dynamodb 1.0.2 → 1.0.3

package/dist/index.js

@@ -1345,6 +1345,20 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
         field,
         direction
       });
+      if (perPage === 0 && (!include || include.length === 0)) {
+        return { messages: [], total: 0, page, perPage: perPageForResponse, hasMore: false };
+      }
+      if (perPage === 0 && include && include.length > 0) {
+        const includeMessages2 = await this._getIncludedMessages({ include });
+        const list2 = new MessageList().add(includeMessages2, "memory");
+        return {
+          messages: this._sortMessages(list2.get.all.db(), field, direction),
+          total: 0,
+          page,
+          perPage: perPageForResponse,
+          hasMore: false
+        };
+      }
       const query = this.service.entities.message.query.byThread({ entity: "message", threadId });
       const results = await query.go();
       let allThreadMessages = results.data.map((data) => this.parseMessageData(data)).filter((msg) => "content" in msg && typeof msg.content === "object");
@@ -1379,8 +1393,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
       const messageIds = new Set(paginatedMessages.map((m) => m.id));
       let includeMessages = [];
       if (include && include.length > 0) {
-        const selectBy = { include };
-        includeMessages = await this._getIncludedMessages(selectBy);
+        includeMessages = await this._getIncludedMessages({ include });
         for (const includeMsg of includeMessages) {
           if (!messageIds.has(includeMsg.id)) {
             paginatedMessages.push(includeMsg);
@@ -1390,14 +1403,7 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
       }
       const list = new MessageList().add(paginatedMessages, "memory");
       let finalMessages = list.get.all.db();
-      finalMessages = finalMessages.sort((a, b) => {
-        const aValue = field === "createdAt" ? new Date(a.createdAt).getTime() : a[field];
-        const bValue = field === "createdAt" ? new Date(b.createdAt).getTime() : b[field];
-        if (aValue === bValue) {
-          return a.id.localeCompare(b.id);
-        }
-        return direction === "ASC" ? aValue - bValue : bValue - aValue;
-      });
+      finalMessages = this._sortMessages(finalMessages, field, direction);
       const returnedThreadMessageIds = new Set(finalMessages.filter((m) => m.threadId === threadId).map((m) => m.id));
       const allThreadMessagesReturned = returnedThreadMessageIds.size >= total;
       let hasMore = false;
@@ -1584,68 +1590,74 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
     }
   }
   // Helper method to get included messages with context
-  async _getIncludedMessages(selectBy) {
-    if (!selectBy?.include?.length) {
+  _sortMessages(messages, field, direction) {
+    return messages.sort((a, b) => {
+      const aValue = field === "createdAt" ? new Date(a.createdAt).getTime() : a[field];
+      const bValue = field === "createdAt" ? new Date(b.createdAt).getTime() : b[field];
+      if (aValue === bValue) {
+        return a.id.localeCompare(b.id);
+      }
+      return direction === "ASC" ? aValue - bValue : bValue - aValue;
+    });
+  }
+  async _getIncludedMessages({
+    include
+  }) {
+    if (!include?.length) {
       return [];
     }
-    const includeMessages = [];
-    for (const includeItem of selectBy.include) {
-      try {
-        const { id, withPreviousMessages = 0, withNextMessages = 0 } = includeItem;
-        const targetResult = await this.service.entities.message.get({ entity: "message", id }).go();
-        if (!targetResult.data) {
-          this.logger.warn("Target message not found", { id });
-          continue;
-        }
-        const targetMessageData = targetResult.data;
-        const searchThreadId = targetMessageData.threadId;
-        this.logger.debug("Getting included messages for", {
-          id,
-          searchThreadId,
-          withPreviousMessages,
-          withNextMessages
-        });
-        const query = this.service.entities.message.query.byThread({ entity: "message", threadId: searchThreadId });
-        const results = await query.go();
-        const allMessages = results.data.map((data) => this.parseMessageData(data)).filter((msg) => "content" in msg && typeof msg.content === "object");
-        this.logger.debug("Found messages in thread", {
-          threadId: searchThreadId,
-          messageCount: allMessages.length,
-          messageIds: allMessages.map((m) => m.id)
-        });
-        allMessages.sort((a, b) => {
-          const timeA = a.createdAt.getTime();
-          const timeB = b.createdAt.getTime();
-          if (timeA === timeB) {
-            return a.id.localeCompare(b.id);
-          }
-          return timeA - timeB;
-        });
-        const targetIndex = allMessages.findIndex((msg) => msg.id === id);
-        if (targetIndex === -1) {
-          this.logger.warn("Target message not found in thread", { id, threadId: searchThreadId });
-          continue;
-        }
-        this.logger.debug("Found target message at index", { id, targetIndex, totalMessages: allMessages.length });
-        const startIndex = Math.max(0, targetIndex - withPreviousMessages);
-        const endIndex = Math.min(allMessages.length, targetIndex + withNextMessages + 1);
-        const contextMessages = allMessages.slice(startIndex, endIndex);
-        this.logger.debug("Context messages", {
-          startIndex,
-          endIndex,
-          contextCount: contextMessages.length,
-          contextIds: contextMessages.map((m) => m.id)
-        });
-        includeMessages.push(...contextMessages);
-      } catch (error) {
-        this.logger.warn("Failed to get included message", { messageId: includeItem.id, error });
+    const targetResults = await Promise.all(
+      include.map(
+        (inc) => this.service.entities.message.get({ entity: "message", id: inc.id }).go().then((r) => ({ id: inc.id, data: r.data })).catch(() => ({ id: inc.id, data: null }))
+      )
+    );
+    const targetMap = /* @__PURE__ */ new Map();
+    for (const { id, data } of targetResults) {
+      if (data) {
+        targetMap.set(id, { threadId: data.threadId });
       }
     }
-    this.logger.debug("Total included messages", {
-      count: includeMessages.length,
-      ids: includeMessages.map((m) => m.id)
+    if (targetMap.size === 0) return [];
+    const threadCache = /* @__PURE__ */ new Map();
+    const uniqueThreadIds = [...new Set([...targetMap.values()].map((t) => t.threadId))];
+    await Promise.all(
+      uniqueThreadIds.map(async (threadId) => {
+        try {
+          const query = this.service.entities.message.query.byThread({ entity: "message", threadId });
+          const results = await query.go();
+          const messages = results.data.map((data) => this.parseMessageData(data)).filter(
+            (msg) => "content" in msg && typeof msg.content === "object"
+          );
+          messages.sort((a, b) => {
+            const timeA = a.createdAt.getTime();
+            const timeB = b.createdAt.getTime();
+            if (timeA === timeB) return a.id.localeCompare(b.id);
+            return timeA - timeB;
+          });
+          threadCache.set(threadId, messages);
+        } catch {
+        }
+      })
+    );
+    const includeMessages = [];
+    for (const includeItem of include) {
+      const { id, withPreviousMessages = 0, withNextMessages = 0 } = includeItem;
+      const target = targetMap.get(id);
+      if (!target) continue;
+      const allMessages = threadCache.get(target.threadId);
+      if (!allMessages) continue;
+      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, targetIndex + withNextMessages + 1);
+      includeMessages.push(...allMessages.slice(startIndex, endIndex));
+    }
+    const seen = /* @__PURE__ */ new Set();
+    return includeMessages.filter((msg) => {
+      if (seen.has(msg.id)) return false;
+      seen.add(msg.id);
+      return true;
     });
-    return includeMessages;
   }
   async updateMessages(args) {
     const { messages } = args;