@mastra/dynamodb 1.1.2 → 1.1.3

package/dist/index.js

@@ -1,7 +1,7 @@
 import { DynamoDBClient, DescribeTableCommand } from '@aws-sdk/client-dynamodb';
 import { DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { BackgroundTasksStorage, createStorageErrorId, MemoryStorage, normalizePerPage, calculatePagination, filterByDateRange, ScoresStorage, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TABLE_BACKGROUND_TASKS, TABLE_SCORERS, TABLE_WORKFLOW_SNAPSHOT, TABLE_RESOURCES, TABLE_MESSAGES, TABLE_THREADS } from '@mastra/core/storage';
+import { BackgroundTasksStorage, createStorageErrorId, MemoryStorage, normalizePerPage, calculatePagination, ScoresStorage, SCORERS_SCHEMA, WorkflowsStorage, MastraCompositeStore, TABLE_BACKGROUND_TASKS, TABLE_SCORERS, TABLE_WORKFLOW_SNAPSHOT, TABLE_RESOURCES, TABLE_MESSAGES, TABLE_THREADS } from '@mastra/core/storage';
 import { Entity, Service } from 'electrodb';
 import { MessageList } from '@mastra/core/agent';
 import { saveScorePayloadSchema } from '@mastra/core/evals';
@@ -1798,29 +1798,80 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
           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");
-      if (resourceId) {
-        allThreadMessages = allThreadMessages.filter((msg) => msg.resourceId === resourceId);
-      }
-      allThreadMessages = filterByDateRange(
-        allThreadMessages,
-        (msg) => new Date(msg.createdAt),
-        filter?.dateRange
-      );
-      allThreadMessages.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);
+      const order = direction === "DESC" ? "desc" : "asc";
+      const parseQueryMessages = (data) => data.map((item) => this.parseMessageData(item)).filter((msg) => "content" in msg && typeof msg.content === "object");
+      const toIso = (value) => value instanceof Date ? value.toISOString() : String(value);
+      const applyQueryFilters = (q) => {
+        let query = q;
+        const dateRange = filter?.dateRange;
+        const startIso = dateRange?.start ? toIso(dateRange.start) : void 0;
+        const endIso = dateRange?.end ? toIso(dateRange.end) : void 0;
+        const startExclusive = dateRange?.startExclusive ?? false;
+        const endExclusive = dateRange?.endExclusive ?? false;
+        const startOp = startExclusive ? "gt" : "gte";
+        const endOp = endExclusive ? "lt" : "lte";
+        if (startIso && endIso) {
+          query = query.between({ createdAt: startIso }, { createdAt: endIso });
+          if (startExclusive || endExclusive) {
+            query = query.where(({ createdAt }, { gt, lt }) => {
+              if (startExclusive && endExclusive) {
+                return `${gt(createdAt, startIso)} AND ${lt(createdAt, endIso)}`;
+              }
+              if (startExclusive) {
+                return gt(createdAt, startIso);
+              }
+              return lt(createdAt, endIso);
+            });
+          }
+        } else if (startIso) {
+          query = query[startOp]({ createdAt: startIso });
+        } else if (endIso) {
+          query = query[endOp]({ createdAt: endIso });
         }
-        return direction === "ASC" ? aValue - bValue : bValue - aValue;
-      });
-      const total = allThreadMessages.length;
-      const paginatedMessages = allThreadMessages.slice(offset, offset + perPage);
-      const paginatedCount = paginatedMessages.length;
-      if (total === 0 && paginatedCount === 0 && (!include || include.length === 0)) {
+        if (resourceId) {
+          query = query.where(({ resourceId: rid }, { eq }) => eq(rid, resourceId));
+        }
+        return query;
+      };
+      let paginatedMessages = [];
+      let total = 0;
+      const filteredMessageIds = /* @__PURE__ */ new Set();
+      if (threadIds.length > 1) {
+        const threadResults = await Promise.all(
+          threadIds.map(async (tid) => {
+            const q = applyQueryFilters(
+              this.service.entities.message.query.byThread({ entity: "message", threadId: tid })
+            );
+            const countResult = await q.go({ pages: "all", attributes: ["id"] });
+            const results = await q.go({ pages: "all", order });
+            return { ids: countResult.data.map((item) => item.id), messages: parseQueryMessages(results.data) };
+          })
+        );
+        for (const r of threadResults) {
+          for (const id of r.ids) filteredMessageIds.add(id);
+        }
+        total = threadResults.reduce((sum, r) => sum + r.ids.length, 0);
+        const merged = threadResults.flatMap((r) => r.messages);
+        const sorted = this._sortMessages(merged, field, direction);
+        paginatedMessages = perPageInput === false ? sorted : sorted.slice(offset, offset + perPage);
+      } else {
+        const baseQuery = this.service.entities.message.query.byThread({ entity: "message", threadId: threadIds[0] });
+        const query = applyQueryFilters(baseQuery);
+        const countResult = await query.go({ pages: "all", attributes: ["id"] });
+        total = countResult.data.length;
+        for (const item of countResult.data) filteredMessageIds.add(item.id);
+        if (perPageInput === false) {
+          const results = await query.go({ pages: "all", order });
+          paginatedMessages = parseQueryMessages(results.data);
+        } else {
+          const results = await query.go({ count: offset + perPage, order });
+          paginatedMessages = parseQueryMessages(results.data).slice(offset, offset + perPage);
+        }
+        if (field !== "createdAt") {
+          paginatedMessages = this._sortMessages(paginatedMessages, field, direction);
+        }
+      }
+      if (total === 0 && paginatedMessages.length === 0 && (!include || include.length === 0)) {
         return {
           messages: [],
           total: 0,
@@ -1843,12 +1894,11 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
       const list = new MessageList().add(paginatedMessages, "memory");
       let finalMessages = list.get.all.db();
       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;
-      if (perPageInput !== false && !allThreadMessagesReturned) {
-        hasMore = offset + paginatedCount < total;
-      }
+      const returnedFilteredMessageIds = new Set(
+        finalMessages.filter((m) => filteredMessageIds.has(m.id)).map((m) => m.id)
+      );
+      const allFilteredMessagesReturned = returnedFilteredMessageIds.size >= total;
+      const hasMore = perPageInput !== false && !allFilteredMessagesReturned && offset + perPage < total;
       return {
         messages: finalMessages,
         total,
@@ -2053,43 +2103,29 @@ var MemoryStorageDynamoDB = class extends MemoryStorage {
     const targetMap = /* @__PURE__ */ new Map();
     for (const { id, data } of targetResults) {
       if (data) {
-        targetMap.set(id, { threadId: data.threadId });
+        const createdAt = typeof data.createdAt === "string" ? data.createdAt : new Date(data.createdAt).toISOString();
+        targetMap.set(id, { threadId: data.threadId, createdAt });
       }
     }
     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 parseQueryMessages = (data) => data.map((item) => this.parseMessageData(item)).filter(
+      (msg) => "content" in msg && typeof msg.content === "object"
     );
     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));
+      try {
+        const prevResult = await this.service.entities.message.query.byThread({ entity: "message", threadId: target.threadId }).lte({ createdAt: target.createdAt }).go({ order: "desc", count: withPreviousMessages + 1 });
+        const prevMessages = parseQueryMessages(prevResult.data).reverse();
+        includeMessages.push(...prevMessages);
+        if (withNextMessages > 0) {
+          const nextResult = await this.service.entities.message.query.byThread({ entity: "message", threadId: target.threadId }).gt({ createdAt: target.createdAt }).go({ order: "asc", count: withNextMessages });
+          includeMessages.push(...parseQueryMessages(nextResult.data));
+        }
+      } catch {
+      }
     }
     const seen = /* @__PURE__ */ new Set();
     return includeMessages.filter((msg) => {