@mastra/upstash 0.10.3-alpha.2 → 0.10.3-alpha.3

package/src/storage/index.ts

@@ -15,6 +15,9 @@ import type {
   EvalRow,
   WorkflowRuns,
   WorkflowRun,
+  PaginationInfo,
+  PaginationArgs,
+  StorageGetTracesArg,
 } from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import { Redis } from '@upstash/redis';
@@ -272,82 +275,44 @@ export class UpstashStore extends MastraStorage {
     }
   }
 
-  public async getTraces(args: {
-    name?: string;
-    scope?: string;
-    attributes?: Record<string, string>;
-    filters?: Record<string, any>;
-    page: number;
-    perPage?: number;
-    fromDate?: Date;
-    toDate?: Date;
-  }): Promise<any[]>;
-  public async getTraces(args: {
-    name?: string;
-    scope?: string;
-    page: number;
-    perPage?: number;
-    attributes?: Record<string, string>;
-    filters?: Record<string, any>;
-    fromDate?: Date;
-    toDate?: Date;
-    returnPaginationResults: true;
-  }): Promise<{
-    traces: any[];
-    total: number;
-    page: number;
-    perPage: number;
-    hasMore: boolean;
-  }>;
-  public async getTraces(args: {
-    name?: string;
-    scope?: string;
-    page: number;
-    perPage?: number;
-    attributes?: Record<string, string>;
-    filters?: Record<string, any>;
-    fromDate?: Date;
-    toDate?: Date;
-    returnPaginationResults?: boolean;
-  }): Promise<
-    | any[]
-    | {
-        traces: any[];
-        total: number;
-        page: number;
-        perPage: number;
-        hasMore: boolean;
-      }
-  > {
-    const {
-      name,
-      scope,
-      page,
-      perPage: perPageInput,
-      attributes,
-      filters,
-      fromDate,
-      toDate,
-      returnPaginationResults,
-    } = args;
+  /**
+   * @deprecated use getTracesPaginated instead
+   */
+  public async getTraces(args: StorageGetTracesArg): Promise<any[]> {
+    if (args.fromDate || args.toDate) {
+      (args as any).dateRange = {
+        start: args.fromDate,
+        end: args.toDate,
+      };
+    }
+    const { traces } = await this.getTracesPaginated(args);
+    return traces;
+  }
 
-    const perPage = perPageInput !== undefined ? perPageInput : 100;
+  public async getTracesPaginated(
+    args: {
+      name?: string;
+      scope?: string;
+      attributes?: Record<string, string>;
+      filters?: Record<string, any>;
+    } & PaginationArgs,
+  ): Promise<PaginationInfo & { traces: any[] }> {
+    const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
 
     try {
       const pattern = `${TABLE_TRACES}:*`;
       const keys = await this.scanKeys(pattern);
 
       if (keys.length === 0) {
-        if (returnPaginationResults) {
-          return {
-            traces: [],
-            total: 0,
-            page,
-            perPage: perPage || 100,
-            hasMore: false,
-          };
-        }
-        return [];
+        return {
+          traces: [],
+          total: 0,
+          page,
+          perPage: perPage || 100,
+          hasMore: false,
+        };
       }
 
       const pipeline = this.redis.pipeline();
@@ -414,29 +379,23 @@ export class UpstashStore extends MastraStorage {
       const end = start + resolvedPerPage;
       const paginatedTraces = transformedTraces.slice(start, end);
       const hasMore = end < total;
-      if (returnPaginationResults) {
-        return {
-          traces: paginatedTraces,
-          total,
-          page,
-          perPage: resolvedPerPage,
-          hasMore,
-        };
-      } else {
-        return paginatedTraces;
-      }
+
+      return {
+        traces: paginatedTraces,
+        total,
+        page,
+        perPage: resolvedPerPage,
+        hasMore,
+      };
     } catch (error) {
       console.error('Failed to get traces:', error);
-      if (returnPaginationResults) {
-        return {
-          traces: [],
-          total: 0,
-          page,
-          perPage: perPage || 100,
-          hasMore: false,
-        };
-      }
-      return [];
+      return {
+        traces: [],
+        total: 0,
+        page,
+        perPage: perPage || 100,
+        hasMore: false,
+      };
     }
   }
 
@@ -515,44 +474,15 @@ export class UpstashStore extends MastraStorage {
     };
   }
 
-  async getThreadsByResourceId(args: { resourceId: string }): Promise<StorageThreadType[]>;
-  async getThreadsByResourceId(args: { resourceId: string; page: number; perPage?: number }): Promise<{
-    threads: StorageThreadType[];
-    total: number;
-    page: number;
-    perPage: number;
-    hasMore: boolean;
-  }>;
-  async getThreadsByResourceId(args: { resourceId: string; page?: number; perPage?: number }): Promise<
-    | StorageThreadType[]
-    | {
-        threads: StorageThreadType[];
-        total: number;
-        page: number;
-        perPage: number;
-        hasMore: boolean;
-      }
-  > {
-    const resourceId: string = args.resourceId;
-    const page: number | undefined = args.page;
-    // Determine perPage only if page is actually provided. Otherwise, its value is not critical for the non-paginated path.
-    // If page is provided, perPage defaults to 100 if not specified.
-    const perPage: number = page !== undefined ? (args.perPage !== undefined ? args.perPage : 100) : 100;
-
+  /**
+   * @deprecated use getThreadsByResourceIdPaginated instead
+   */
+  async getThreadsByResourceId({ resourceId }: { resourceId: string }): Promise<StorageThreadType[]> {
     try {
       const pattern = `${TABLE_THREADS}:*`;
       const keys = await this.scanKeys(pattern);
 
       if (keys.length === 0) {
-        if (page !== undefined) {
-          return {
-            threads: [],
-            total: 0,
-            page,
-            perPage, // perPage is number here
-            hasMore: false,
-          };
-        }
         return [];
       }
 
@@ -574,40 +504,48 @@ export class UpstashStore extends MastraStorage {
       }
 
       allThreads.sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
-
-      if (page !== undefined) {
-        // If page is defined, perPage is also a number (due to the defaulting logic above)
-        const total = allThreads.length;
-        const start = page * perPage;
-        const end = start + perPage;
-        const paginatedThreads = allThreads.slice(start, end);
-        const hasMore = end < total;
-        return {
-          threads: paginatedThreads,
-          total,
-          page,
-          perPage,
-          hasMore,
-        };
-      } else {
-        // page is undefined, return all threads
-        return allThreads;
-      }
+      return allThreads;
     } catch (error) {
       console.error('Error in getThreadsByResourceId:', error);
-      if (page !== undefined) {
-        return {
-          threads: [],
-          total: 0,
-          page,
-          perPage, // perPage is number here
-          hasMore: false,
-        };
-      }
       return [];
     }
   }
 
+  public async getThreadsByResourceIdPaginated(
+    args: {
+      resourceId: string;
+    } & PaginationArgs,
+  ): Promise<PaginationInfo & { threads: StorageThreadType[] }> {
+    const { resourceId, page = 0, perPage = 100 } = args;
+
+    try {
+      const allThreads = await this.getThreadsByResourceId({ resourceId });
+
+      const total = allThreads.length;
+      const start = page * perPage;
+      const end = start + perPage;
+      const paginatedThreads = allThreads.slice(start, end);
+      const hasMore = end < total;
+
+      return {
+        threads: paginatedThreads,
+        total,
+        page,
+        perPage,
+        hasMore,
+      };
+    } catch (error) {
+      console.error('Error in getThreadsByResourceIdPaginated:', error);
+      return {
+        threads: [],
+        total: 0,
+        page,
+        perPage,
+        hasMore: false,
+      };
+    }
+  }
+
   async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
     await this.insert({
       tableName: TABLE_THREADS,
@@ -717,160 +655,13 @@ export class UpstashStore extends MastraStorage {
     return list.get.all.v1();
   }
 
-  // Function overloads for different return types
-  public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
-  public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
-  public async getMessages(
-    args: StorageGetMessagesArg & {
-      format?: 'v1' | 'v2';
-      page: number;
-      perPage?: number;
-      fromDate?: Date;
-      toDate?: Date;
-    },
-  ): Promise<{
-    messages: MastraMessageV1[] | MastraMessageV2[];
-    total: number;
-    page: number;
-    perPage: number;
-    hasMore: boolean;
-  }>;
-  public async getMessages({
-    threadId,
-    selectBy,
-    format,
-    page,
-    perPage = 40,
-    fromDate,
-    toDate,
-  }: StorageGetMessagesArg & {
-    format?: 'v1' | 'v2';
-    page?: number;
-    perPage?: number;
-    fromDate?: Date;
-    toDate?: Date;
-  }): Promise<
-    | MastraMessageV1[]
-    | MastraMessageV2[]
-    | {
-        messages: MastraMessageV1[] | MastraMessageV2[];
-        total: number;
-        page: number;
-        perPage: number;
-        hasMore: boolean;
-      }
-  > {
-    const threadMessagesKey = this.getThreadMessagesKey(threadId);
-
-    const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
-    // If pagination is requested, use the new pagination logic
-    if (page !== undefined) {
-      try {
-        // Get all message IDs from the sorted set
-
-        if (allMessageIds.length === 0) {
-          return {
-            messages: [],
-            total: 0,
-            page,
-            perPage,
-            hasMore: false,
-          };
-        }
-
-        // Use pipeline to fetch all messages efficiently
-        const pipeline = this.redis.pipeline();
-        allMessageIds.forEach(id => pipeline.get(this.getMessageKey(threadId, id as string)));
-        const results = await pipeline.exec();
-
-        // Process messages and apply filters - handle undefined results from pipeline
-        let messages = results
-          .map((result: any) => result as MastraMessageV2 | null)
-          .filter((msg): msg is MastraMessageV2 => msg !== null) as (MastraMessageV2 & { _index?: number })[];
-
-        // Apply date filters if provided
-        if (fromDate) {
-          messages = messages.filter(msg => msg && new Date(msg.createdAt).getTime() >= fromDate.getTime());
-        }
-
-        if (toDate) {
-          messages = messages.filter(msg => msg && new Date(msg.createdAt).getTime() <= toDate.getTime());
-        }
-
-        // Sort messages by their position in the sorted set
-        messages.sort((a, b) => allMessageIds.indexOf(a!.id) - allMessageIds.indexOf(b!.id));
-
-        const total = messages.length;
-
-        // Apply pagination
-        const start = page * perPage;
-        const end = start + perPage;
-        const hasMore = end < total;
-        const paginatedMessages = messages.slice(start, end);
-
-        // Remove _index before returning and handle format conversion properly
-        const prepared = paginatedMessages
-          .filter(message => message !== null && message !== undefined)
-          .map(message => {
-            const { _index, ...messageWithoutIndex } = message as MastraMessageV2 & { _index?: number };
-            return messageWithoutIndex as unknown as MastraMessageV1;
-          });
-
-        // Return pagination object with correct format
-        if (format === 'v2') {
-          // Convert V1 format back to V2 format
-          const v2Messages = prepared.map(msg => ({
-            ...msg,
-            content: msg.content || { format: 2, parts: [{ type: 'text', text: '' }] },
-          })) as MastraMessageV2[];
-
-          return {
-            messages: v2Messages,
-            total,
-            page,
-            perPage,
-            hasMore,
-          };
-        }
-
-        return {
-          messages: prepared,
-          total,
-          page,
-          perPage,
-          hasMore,
-        };
-      } catch (error) {
-        console.error('Failed to get paginated messages:', error);
-        return {
-          messages: [],
-          total: 0,
-          page,
-          perPage,
-          hasMore: false,
-        };
-      }
-    }
-
-    // Original logic for backward compatibility
-    // When selectBy is undefined or selectBy.last is undefined, get ALL messages (not just 40)
-    let limit: number;
-    if (typeof selectBy?.last === 'number') {
-      limit = Math.max(0, selectBy.last);
-    } else if (selectBy?.last === false) {
-      limit = 0;
-    } else {
-      // No limit specified - get all messages
-      limit = Number.MAX_SAFE_INTEGER;
-    }
-
+  private async _getIncludedMessages(
+    threadId: string,
+    selectBy: StorageGetMessagesArg['selectBy'],
+  ): Promise<MastraMessageV2[] | MastraMessageV1[]> {
     const messageIds = new Set<string>();
     const messageIdToThreadIds: Record<string, string> = {};
 
-    if (limit === 0 && !selectBy?.include) {
-      return [];
-    }
-
     // First, get specifically included messages and their context
     if (selectBy?.include?.length) {
       for (const item of selectBy.include) {
@@ -904,6 +695,47 @@ export class UpstashStore extends MastraStorage {
           });
         }
       }
+
+      const pipeline = this.redis.pipeline();
+      Array.from(messageIds).forEach(id => {
+        const tId = messageIdToThreadIds[id] || threadId;
+        pipeline.get(this.getMessageKey(tId, id as string));
+      });
+      const results = await pipeline.exec();
+      return results.filter(result => result !== null) as MastraMessageV2[] | MastraMessageV1[];
+    }
+
+    return [];
+  }
+
+  /**
+   * @deprecated use getMessagesPaginated instead
+   */
+  public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
+  public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
+  public async getMessages({
+    threadId,
+    selectBy,
+    format,
+  }: StorageGetMessagesArg & { format?: 'v1' | 'v2' }): Promise<MastraMessageV1[] | MastraMessageV2[]> {
+    const threadMessagesKey = this.getThreadMessagesKey(threadId);
+    const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
+    // When selectBy is undefined or selectBy.last is undefined, get ALL messages (not just 40)
+    let limit: number;
+    if (typeof selectBy?.last === 'number') {
+      limit = Math.max(0, selectBy.last);
+    } else if (selectBy?.last === false) {
+      limit = 0;
+    } else {
+      // No limit specified - get all messages
+      limit = Number.MAX_SAFE_INTEGER;
+    }
+
+    const messageIds = new Set<string>();
+    const messageIdToThreadIds: Record<string, string> = {};
+
+    if (limit === 0 && !selectBy?.include) {
+      return [];
     }
 
     // Then get the most recent messages (or all if no limit)
@@ -923,24 +755,36 @@ export class UpstashStore extends MastraStorage {
       });
     }
 
+    const includedMessages = await this._getIncludedMessages(threadId, selectBy);
+
     // Fetch all needed messages in parallel
-    const messages = (
-      await Promise.all(
-        Array.from(messageIds).map(async id => {
-          const tId = messageIdToThreadIds[id] || threadId;
-          const byThreadId = await this.redis.get<MastraMessageV2 & { _index?: number }>(this.getMessageKey(tId, id));
-          if (byThreadId) return byThreadId;
-
-          return null;
-        }),
-      )
-    ).filter(msg => msg !== null) as (MastraMessageV2 & { _index?: number })[];
+    const messages = [
+      ...includedMessages,
+      ...((
+        await Promise.all(
+          Array.from(messageIds).map(async id => {
+            const tId = messageIdToThreadIds[id] || threadId;
+            const byThreadId = await this.redis.get<MastraMessageV2 & { _index?: number }>(this.getMessageKey(tId, id));
+            if (byThreadId) return byThreadId;
+
+            return null;
+          }),
+        )
+      ).filter(msg => msg !== null) as (MastraMessageV2 & { _index?: number })[]),
+    ];
 
     // Sort messages by their position in the sorted set
     messages.sort((a, b) => allMessageIds.indexOf(a!.id) - allMessageIds.indexOf(b!.id));
 
+    const seen = new Set<string>();
+    const dedupedMessages = messages.filter(row => {
+      if (seen.has(row.id)) return false;
+      seen.add(row.id);
+      return true;
+    });
+
     // Remove _index before returning and handle format conversion properly
-    const prepared = messages
+    const prepared = dedupedMessages
       .filter(message => message !== null && message !== undefined)
       .map(message => {
         const { _index, ...messageWithoutIndex } = message as MastraMessageV2 & { _index?: number };
@@ -960,6 +804,89 @@ export class UpstashStore extends MastraStorage {
     return prepared;
   }
 
+  public async getMessagesPaginated(
+    args: StorageGetMessagesArg & {
+      format?: 'v1' | 'v2';
+    },
+  ): Promise<PaginationInfo & { messages: MastraMessageV1[] | MastraMessageV2[] }> {
+    const { threadId, selectBy, format } = args;
+    const { page = 0, perPage = 40, dateRange } = selectBy?.pagination || {};
+    const fromDate = dateRange?.start;
+    const toDate = dateRange?.end;
+    const threadMessagesKey = this.getThreadMessagesKey(threadId);
+    const messages: (MastraMessageV2 | MastraMessageV1)[] = [];
+
+    const includedMessages = await this._getIncludedMessages(threadId, selectBy);
+    messages.push(...includedMessages);
+
+    try {
+      const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
+      if (allMessageIds.length === 0) {
+        return {
+          messages: [],
+          total: 0,
+          page,
+          perPage,
+          hasMore: false,
+        };
+      }
+
+      // Use pipeline to fetch all messages efficiently
+      const pipeline = this.redis.pipeline();
+      allMessageIds.forEach(id => pipeline.get(this.getMessageKey(threadId, id as string)));
+      const results = await pipeline.exec();
+
+      // Process messages and apply filters - handle undefined results from pipeline
+      let messagesData = results.filter((msg): msg is MastraMessageV2 | MastraMessageV1 => msg !== null) as (
+        | MastraMessageV2
+        | MastraMessageV1
+      )[];
+
+      // Apply date filters if provided
+      if (fromDate) {
+        messagesData = messagesData.filter(msg => msg && new Date(msg.createdAt).getTime() >= fromDate.getTime());
+      }
+
+      if (toDate) {
+        messagesData = messagesData.filter(msg => msg && new Date(msg.createdAt).getTime() <= toDate.getTime());
+      }
+
+      // Sort messages by their position in the sorted set
+      messagesData.sort((a, b) => allMessageIds.indexOf(a!.id) - allMessageIds.indexOf(b!.id));
+
+      const total = messagesData.length;
+
+      const start = page * perPage;
+      const end = start + perPage;
+      const hasMore = end < total;
+      const paginatedMessages = messagesData.slice(start, end);
+
+      messages.push(...paginatedMessages);
+
+      const list = new MessageList().add(messages, 'memory');
+      const finalMessages = (format === `v2` ? list.get.all.v2() : list.get.all.v1()) as
+        | MastraMessageV1[]
+        | MastraMessageV2[];
+
+      return {
+        messages: finalMessages,
+        total,
+        page,
+        perPage,
+        hasMore,
+      };
+    } catch (error) {
+      console.error('Failed to get paginated messages:', error);
+      return {
+        messages: [],
+        total: 0,
+        page,
+        perPage,
+        hasMore: false,
+      };
+    }
+  }
+
   async persistWorkflowSnapshot(params: {
     namespace: string;
     workflowName: string;
@@ -1006,28 +933,17 @@ export class UpstashStore extends MastraStorage {
    * @param options Pagination and filtering options
    * @returns Object with evals array and total count
    */
-  async getEvals(options?: {
-    agentName?: string;
-    type?: 'test' | 'live';
-    page?: number;
-    perPage?: number;
-    limit?: number;
-    offset?: number;
-    fromDate?: Date;
-    toDate?: Date;
-  }): Promise<{
-    evals: EvalRow[];
-    total: number;
-    page?: number;
-    perPage?: number;
-    hasMore?: boolean;
-  }> {
+  async getEvals(
+    options?: {
+      agentName?: string;
+      type?: 'test' | 'live';
+    } & PaginationArgs,
+  ): Promise<PaginationInfo & { evals: EvalRow[] }> {
     try {
       // Default pagination parameters
-      const page = options?.page ?? 0;
-      const perPage = options?.perPage ?? 100;
-      const limit = options?.limit;
-      const offset = options?.offset;
+      const { agentName, type, page = 0, perPage = 100, dateRange } = options || {};
+      const fromDate = dateRange?.start;
+      const toDate = dateRange?.end;
 
       // Get all keys that match the evals table pattern using cursor-based scanning
       const pattern = `${TABLE_EVALS}:*`;
@@ -1038,8 +954,8 @@ export class UpstashStore extends MastraStorage {
         return {
           evals: [],
           total: 0,
-          page: options?.page ?? 0,
-          perPage: options?.perPage ?? 100,
+          page,
+          perPage,
           hasMore: false,
         };
       }
@@ -1055,12 +971,12 @@ export class UpstashStore extends MastraStorage {
         .filter((record): record is Record<string, any> => record !== null && typeof record === 'object');
 
       // Apply agent name filter if provided
-      if (options?.agentName) {
-        filteredEvals = filteredEvals.filter(record => record.agent_name === options.agentName);
+      if (agentName) {
+        filteredEvals = filteredEvals.filter(record => record.agent_name === agentName);
       }
 
       // Apply type filter if provided
-      if (options?.type === 'test') {
+      if (type === 'test') {
         filteredEvals = filteredEvals.filter(record => {
           if (!record.test_info) return false;
 
@@ -1074,7 +990,7 @@ export class UpstashStore extends MastraStorage {
             return false;
           }
         });
-      } else if (options?.type === 'live') {
+      } else if (type === 'live') {
         filteredEvals = filteredEvals.filter(record => {
           if (!record.test_info) return true;
 
@@ -1091,17 +1007,17 @@ export class UpstashStore extends MastraStorage {
       }
 
       // Apply date filters if provided
-      if (options?.fromDate) {
+      if (fromDate) {
         filteredEvals = filteredEvals.filter(record => {
           const createdAt = new Date(record.created_at || record.createdAt || 0);
-          return createdAt.getTime() >= options.fromDate!.getTime();
+          return createdAt.getTime() >= fromDate.getTime();
         });
       }
 
-      if (options?.toDate) {
+      if (toDate) {
         filteredEvals = filteredEvals.filter(record => {
           const createdAt = new Date(record.created_at || record.createdAt || 0);
-          return createdAt.getTime() <= options.toDate!.getTime();
+          return createdAt.getTime() <= toDate.getTime();
         });
       }
 
@@ -1114,21 +1030,11 @@ export class UpstashStore extends MastraStorage {
 
       const total = filteredEvals.length;
 
-      // Apply pagination - support both page/perPage and limit/offset patterns
-      let paginatedEvals: Record<string, any>[];
-      let hasMore = false;
-
-      if (limit !== undefined && offset !== undefined) {
-        // Offset-based pagination
-        paginatedEvals = filteredEvals.slice(offset, offset + limit);
-        hasMore = offset + limit < total;
-      } else {
-        // Page-based pagination
-        const start = page * perPage;
-        const end = start + perPage;
-        paginatedEvals = filteredEvals.slice(start, end);
-        hasMore = end < total;
-      }
+      // Apply pagination
+      const start = page * perPage;
+      const end = start + perPage;
+      const paginatedEvals = filteredEvals.slice(start, end);
+      const hasMore = end < total;
 
       // Transform to EvalRow format
       const evals = paginatedEvals.map(record => this.transformEvalRecord(record));
@@ -1136,17 +1042,18 @@ export class UpstashStore extends MastraStorage {
       return {
         evals,
         total,
-        page: limit !== undefined ? undefined : page,
-        perPage: limit !== undefined ? undefined : perPage,
+        page,
+        perPage,
         hasMore,
       };
     } catch (error) {
+      const { page = 0, perPage = 100 } = options || {};
       console.error('Failed to get evals:', error);
       return {
         evals: [],
         total: 0,
-        page: options?.page ?? 0,
-        perPage: options?.perPage ?? 100,
+        page,
+        perPage,
         hasMore: false,
       };
     }