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

package/src/storage/upstash.test.ts

@@ -1,5 +1,10 @@
 import { randomUUID } from 'crypto';
-import { createSampleMessageV2, createSampleThread, createSampleWorkflowSnapshot } from '@internal/storage-test-utils';
+import {
+  checkWorkflowSnapshot,
+  createSampleMessageV2,
+  createSampleThread,
+  createSampleWorkflowSnapshot,
+} from '@internal/storage-test-utils';
 import type { MastraMessageV2 } from '@mastra/core';
 import type { TABLE_NAMES } from '@mastra/core/storage';
 import {
@@ -17,7 +22,12 @@ import { UpstashStore } from './index';
 // Increase timeout for all tests in this file to 30 seconds
 vi.setConfig({ testTimeout: 200_000, hookTimeout: 200_000 });
 
-const createSampleTrace = (name: string, scope?: string, attributes?: Record<string, string>) => ({
+const createSampleTrace = (
+  name: string,
+  scope?: string,
+  attributes?: Record<string, string>,
+  createdAt: Date = new Date(),
+) => ({
   id: `trace-${randomUUID()}`,
   parentSpanId: `span-${randomUUID()}`,
   traceId: `trace-${randomUUID()}`,
@@ -25,16 +35,16 @@ const createSampleTrace = (name: string, scope?: string, attributes?: Record<str
   scope,
   kind: 'internal',
   status: JSON.stringify({ code: 'success' }),
-  events: JSON.stringify([{ name: 'start', timestamp: Date.now() }]),
+  events: JSON.stringify([{ name: 'start', timestamp: createdAt.getTime() }]),
   links: JSON.stringify([]),
   attributes: attributes ? JSON.stringify(attributes) : undefined,
-  startTime: new Date().toISOString(),
-  endTime: new Date().toISOString(),
+  startTime: createdAt.toISOString(),
+  endTime: new Date(createdAt.getTime() + 1000).toISOString(),
   other: JSON.stringify({ custom: 'data' }),
-  createdAt: new Date().toISOString(),
+  createdAt: createdAt.toISOString(),
 });
 
-const createSampleEval = (agentName: string, isTest = false) => {
+const createSampleEval = (agentName: string, isTest = false, createdAt: Date = new Date()) => {
   const testInfo = isTest ? { testPath: 'test/path.ts', testName: 'Test Name' } : undefined;
 
   return {
@@ -47,17 +57,10 @@ const createSampleEval = (agentName: string, isTest = false) => {
     test_info: testInfo ? JSON.stringify(testInfo) : undefined,
     global_run_id: `global-${randomUUID()}`,
     run_id: `run-${randomUUID()}`,
-    created_at: new Date().toISOString(),
+    created_at: createdAt.toISOString(),
   };
 };
 
-const checkWorkflowSnapshot = (snapshot: WorkflowRunState | string, stepId: string, status: string) => {
-  if (typeof snapshot === 'string') {
-    throw new Error('Expected WorkflowRunState, got string');
-  }
-  expect(snapshot.context?.[stepId]?.status).toBe(status);
-};
-
 describe('UpstashStore', () => {
   let store: UpstashStore;
   const testTableName = 'test_table';
@@ -154,7 +157,7 @@ describe('UpstashStore', () => {
 
     it('should get threads by resource ID', async () => {
       const thread1 = createSampleThread();
-      const thread2 = { ...createSampleThread(), resourceId: thread1.resourceId };
+      const thread2 = createSampleThread({ resourceId: thread1.resourceId });
       const threads = [thread1, thread2];
 
       const resourceId = threads[0].resourceId;
@@ -187,7 +190,7 @@ describe('UpstashStore', () => {
     it('should fetch >100000 threads by resource ID', async () => {
       const resourceId = `resource-${randomUUID()}`;
       const total = 100_000;
-      const threads = Array.from({ length: total }, () => ({ ...createSampleThread(), resourceId }));
+      const threads = Array.from({ length: total }, () => createSampleThread({ resourceId }));
 
       await store.batchInsert({ tableName: TABLE_THREADS, records: threads });
 
@@ -300,13 +303,117 @@ describe('UpstashStore', () => {
         createSampleMessageV2({ threadId, content: 'Third' }),
       ];
 
-      await store.saveMessages({ messages: messages, format: 'v2' });
+      await store.saveMessages({ messages, format: 'v2' });
 
       const retrievedMessages = await store.getMessages({ threadId, format: 'v2' });
       expect(retrievedMessages).toHaveLength(3);
       expect(retrievedMessages.map((m: any) => m.content.parts[0].text)).toEqual(['First', 'Second', 'Third']);
     });
 
+    it('should retrieve messages w/ next/prev messages by message id + resource id', async () => {
+      const thread = createSampleThread({ id: 'thread-one' });
+      await store.saveThread({ thread });
+
+      const thread2 = createSampleThread({ id: 'thread-two' });
+      await store.saveThread({ thread: thread2 });
+
+      const thread3 = createSampleThread({ id: 'thread-three' });
+      await store.saveThread({ thread: thread3 });
+
+      const messages: MastraMessageV2[] = [
+        createSampleMessageV2({ threadId: 'thread-one', content: 'First', resourceId: 'cross-thread-resource' }),
+        createSampleMessageV2({ threadId: 'thread-one', content: 'Second', resourceId: 'cross-thread-resource' }),
+        createSampleMessageV2({ threadId: 'thread-one', content: 'Third', resourceId: 'cross-thread-resource' }),
+
+        createSampleMessageV2({ threadId: 'thread-two', content: 'Fourth', resourceId: 'cross-thread-resource' }),
+        createSampleMessageV2({ threadId: 'thread-two', content: 'Fifth', resourceId: 'cross-thread-resource' }),
+        createSampleMessageV2({ threadId: 'thread-two', content: 'Sixth', resourceId: 'cross-thread-resource' }),
+
+        createSampleMessageV2({ threadId: 'thread-three', content: 'Seventh', resourceId: 'other-resource' }),
+        createSampleMessageV2({ threadId: 'thread-three', content: 'Eighth', resourceId: 'other-resource' }),
+      ];
+
+      await store.saveMessages({ messages: messages, format: 'v2' });
+
+      const retrievedMessages = await store.getMessages({ threadId: 'thread-one', format: 'v2' });
+      expect(retrievedMessages).toHaveLength(3);
+      expect(retrievedMessages.map((m: any) => m.content.parts[0].text)).toEqual(['First', 'Second', 'Third']);
+
+      const retrievedMessages2 = await store.getMessages({ threadId: 'thread-two', format: 'v2' });
+      expect(retrievedMessages2).toHaveLength(3);
+      expect(retrievedMessages2.map((m: any) => m.content.parts[0].text)).toEqual(['Fourth', 'Fifth', 'Sixth']);
+
+      const retrievedMessages3 = await store.getMessages({ threadId: 'thread-three', format: 'v2' });
+      expect(retrievedMessages3).toHaveLength(2);
+      expect(retrievedMessages3.map((m: any) => m.content.parts[0].text)).toEqual(['Seventh', 'Eighth']);
+
+      const crossThreadMessages = await store.getMessages({
+        threadId: 'thread-doesnt-exist',
+        format: 'v2',
+        selectBy: {
+          last: 0,
+          include: [
+            {
+              id: messages[1].id,
+              threadId: 'thread-one',
+              withNextMessages: 2,
+              withPreviousMessages: 2,
+            },
+            {
+              id: messages[4].id,
+              threadId: 'thread-two',
+              withPreviousMessages: 2,
+              withNextMessages: 2,
+            },
+          ],
+        },
+      });
+
+      expect(crossThreadMessages).toHaveLength(6);
+      expect(crossThreadMessages.filter(m => m.threadId === `thread-one`)).toHaveLength(3);
+      expect(crossThreadMessages.filter(m => m.threadId === `thread-two`)).toHaveLength(3);
+
+      const crossThreadMessages2 = await store.getMessages({
+        threadId: 'thread-one',
+        format: 'v2',
+        selectBy: {
+          last: 0,
+          include: [
+            {
+              id: messages[4].id,
+              threadId: 'thread-two',
+              withPreviousMessages: 1,
+              withNextMessages: 1,
+            },
+          ],
+        },
+      });
+
+      expect(crossThreadMessages2).toHaveLength(3);
+      expect(crossThreadMessages2.filter(m => m.threadId === `thread-one`)).toHaveLength(0);
+      expect(crossThreadMessages2.filter(m => m.threadId === `thread-two`)).toHaveLength(3);
+
+      const crossThreadMessages3 = await store.getMessages({
+        threadId: 'thread-two',
+        format: 'v2',
+        selectBy: {
+          last: 0,
+          include: [
+            {
+              id: messages[1].id,
+              threadId: 'thread-one',
+              withNextMessages: 1,
+              withPreviousMessages: 1,
+            },
+          ],
+        },
+      });
+
+      expect(crossThreadMessages3).toHaveLength(3);
+      expect(crossThreadMessages3.filter(m => m.threadId === `thread-one`)).toHaveLength(3);
+      expect(crossThreadMessages3.filter(m => m.threadId === `thread-two`)).toHaveLength(0);
+    });
+
     it('should handle empty message array', async () => {
       const result = await store.saveMessages({ messages: [] });
       expect(result).toEqual([]);
@@ -336,7 +443,7 @@ describe('UpstashStore', () => {
       expect(retrievedMessages[0].content).toEqual(messages[0].content);
     });
 
-    describe('getPaginatedMessages', () => {
+    describe('getMessagesPaginated', () => {
       it('should return paginated messages with total count', async () => {
         const thread = createSampleThread();
         await store.saveThread({ thread });
@@ -347,10 +454,9 @@ describe('UpstashStore', () => {
 
         await store.saveMessages({ messages, format: 'v2' });
 
-        const page1 = await store.getMessages({
+        const page1 = await store.getMessagesPaginated({
           threadId: thread.id,
-          page: 0,
-          perPage: 5,
+          selectBy: { pagination: { page: 0, perPage: 5 } },
           format: 'v2',
         });
         expect(page1.messages).toHaveLength(5);
@@ -359,20 +465,18 @@ describe('UpstashStore', () => {
         expect(page1.perPage).toBe(5);
         expect(page1.hasMore).toBe(true);
 
-        const page3 = await store.getMessages({
+        const page3 = await store.getMessagesPaginated({
           threadId: thread.id,
-          page: 2,
-          perPage: 5,
+          selectBy: { pagination: { page: 2, perPage: 5 } },
           format: 'v2',
         });
         expect(page3.messages).toHaveLength(5);
         expect(page3.total).toBe(15);
         expect(page3.hasMore).toBe(false);
 
-        const page4 = await store.getMessages({
+        const page4 = await store.getMessagesPaginated({
           threadId: thread.id,
-          page: 3,
-          perPage: 5,
+          selectBy: { pagination: { page: 3, perPage: 5 } },
           format: 'v2',
         });
         expect(page4.messages).toHaveLength(0);
@@ -393,10 +497,9 @@ describe('UpstashStore', () => {
 
         await store.saveMessages({ messages, format: 'v2' });
 
-        const page1 = await store.getMessages({
+        const page1 = await store.getMessagesPaginated({
           threadId: thread.id,
-          page: 0,
-          perPage: 3,
+          selectBy: { pagination: { page: 0, perPage: 3 } },
           format: 'v2',
         });
 
@@ -410,32 +513,6 @@ describe('UpstashStore', () => {
         }
       });
 
-      it('should maintain backward compatibility when no pagination params provided', async () => {
-        const thread = createSampleThread();
-        await store.saveThread({ thread });
-
-        const messages = Array.from({ length: 5 }, (_, i) =>
-          createSampleMessageV2({ threadId: thread.id, content: `Message ${i + 1}` }),
-        );
-
-        await store.saveMessages({ messages, format: 'v2' });
-
-        // Test original format without pagination - should return array
-        const messagesV1 = await store.getMessages({
-          threadId: thread.id,
-          format: 'v1',
-        });
-        expect(Array.isArray(messagesV1)).toBe(true);
-        expect(messagesV1).toHaveLength(5);
-
-        const messagesV2 = await store.getMessages({
-          threadId: thread.id,
-          format: 'v2',
-        });
-        expect(Array.isArray(messagesV2)).toBe(true);
-        expect(messagesV2).toHaveLength(5);
-      });
-
       it('should support date filtering with pagination', async () => {
         const thread = createSampleThread();
         await store.saveThread({ thread });
@@ -458,11 +535,15 @@ describe('UpstashStore', () => {
 
         await store.saveMessages({ messages: [...oldMessages, ...newMessages], format: 'v2' });
 
-        const recentMessages = await store.getMessages({
+        const recentMessages = await store.getMessagesPaginated({
           threadId: thread.id,
-          page: 0,
-          perPage: 10,
-          fromDate: now,
+          selectBy: {
+            pagination: {
+              page: 0,
+              perPage: 10,
+              dateRange: { start: now },
+            },
+          },
           format: 'v2',
         });
         expect(recentMessages.messages).toHaveLength(4);
@@ -1098,7 +1179,7 @@ describe('UpstashStore', () => {
         expect(page3.hasMore).toBe(false);
       });
 
-      it('should support limit/offset pagination', async () => {
+      it('should support page/perPage pagination', async () => {
         const agentName = 'test-agent-2';
         const evals = Array.from({ length: 15 }, () => createSampleEval(agentName));
 
@@ -1110,12 +1191,12 @@ describe('UpstashStore', () => {
         }
 
         // Test offset-based pagination
-        const result1 = await store.getEvals({ agentName, limit: 5, offset: 0 });
+        const result1 = await store.getEvals({ agentName, page: 0, perPage: 5 });
         expect(result1.evals).toHaveLength(5);
         expect(result1.total).toBe(15);
         expect(result1.hasMore).toBe(true);
 
-        const result2 = await store.getEvals({ agentName, limit: 5, offset: 10 });
+        const result2 = await store.getEvals({ agentName, page: 2, perPage: 5 });
         expect(result2.evals).toHaveLength(5);
         expect(result2.total).toBe(15);
         expect(result2.hasMore).toBe(false);
@@ -1141,9 +1222,30 @@ describe('UpstashStore', () => {
         expect(liveResults.evals).toHaveLength(5);
         expect(liveResults.total).toBe(8);
       });
+
+      it('should filter by date with pagination', async () => {
+        const agentName = 'test-agent-date';
+        const now = new Date();
+        const yesterday = new Date(now.getTime() - 24 * 60 * 60 * 1000);
+        const evals = [createSampleEval(agentName, false, now), createSampleEval(agentName, false, yesterday)];
+        for (const evalRecord of evals) {
+          await store.insert({
+            tableName: TABLE_EVALS,
+            record: evalRecord,
+          });
+        }
+        const result = await store.getEvals({
+          agentName,
+          page: 0,
+          perPage: 10,
+          dateRange: { start: now },
+        });
+        expect(result.evals).toHaveLength(1);
+        expect(result.total).toBe(1);
+      });
     });
 
-    describe('getTracesPaginated', () => {
+    describe('getTraces with pagination', () => {
       it('should return paginated traces with total count', async () => {
         const traces = Array.from({ length: 18 }, (_, i) => createSampleTrace(`test-trace-${i}`, 'test-scope'));
 
@@ -1154,11 +1256,10 @@ describe('UpstashStore', () => {
           });
         }
 
-        const page1 = await store.getTraces({
+        const page1 = await store.getTracesPaginated({
           scope: 'test-scope',
           page: 0,
           perPage: 8,
-          returnPaginationResults: true,
         });
         expect(page1.traces).toHaveLength(8);
         expect(page1.total).toBe(18);
@@ -1166,75 +1267,62 @@ describe('UpstashStore', () => {
         expect(page1.perPage).toBe(8);
         expect(page1.hasMore).toBe(true);
 
-        const page3 = await store.getTraces({
+        const page3 = await store.getTracesPaginated({
           scope: 'test-scope',
           page: 2,
           perPage: 8,
-          returnPaginationResults: true,
         });
         expect(page3.traces).toHaveLength(2);
         expect(page3.total).toBe(18);
         expect(page3.hasMore).toBe(false);
       });
 
-      it('should filter by attributes with pagination', async () => {
-        const tracesWithAttr = Array.from({ length: 8 }, (_, i) =>
-          createSampleTrace(`trace-${i}`, 'test-scope', { environment: 'prod' }),
-        );
-        const tracesWithoutAttr = Array.from({ length: 5 }, (_, i) =>
-          createSampleTrace(`trace-other-${i}`, 'test-scope', { environment: 'dev' }),
-        );
+      it('should filter by date with pagination', async () => {
+        const scope = 'test-scope-date';
+        const now = new Date();
+        const yesterday = new Date(now.getTime() - 24 * 60 * 60 * 1000);
+
+        const traces = [
+          createSampleTrace(`test-trace-now`, scope, undefined, now),
+          createSampleTrace(`test-trace-yesterday`, scope, undefined, yesterday),
+        ];
 
-        for (const trace of [...tracesWithAttr, ...tracesWithoutAttr]) {
+        for (const trace of traces) {
           await store.insert({
             tableName: TABLE_TRACES,
             record: trace,
           });
         }
 
-        const prodTraces = await store.getTraces({
-          scope: 'test-scope',
-          attributes: { environment: 'prod' },
-          page: 0,
-          perPage: 5,
-          returnPaginationResults: true,
-        });
-        expect(prodTraces.traces).toHaveLength(5);
-        expect(prodTraces.total).toBe(8);
-        expect(prodTraces.hasMore).toBe(true);
-
-        const devTraces = await store.getTraces({
-          scope: 'test-scope',
-          attributes: { environment: 'dev' },
+        const result = await store.getTracesPaginated({
+          scope,
           page: 0,
           perPage: 10,
-          returnPaginationResults: true,
+          dateRange: { start: now },
         });
-        expect(devTraces.traces).toHaveLength(5);
-        expect(devTraces.total).toBe(5);
-        expect(devTraces.hasMore).toBe(false);
+
+        expect(result.traces).toHaveLength(1);
+        expect(result.traces[0].name).toBe('test-trace-now');
+        expect(result.total).toBe(1);
       });
     });
 
     describe('Enhanced existing methods with pagination', () => {
       it('should support pagination in getThreadsByResourceId', async () => {
         const resourceId = 'enhanced-resource';
-        const threads = Array.from({ length: 17 }, () => ({
-          ...createSampleThread(),
-          resourceId,
-        }));
+        const threads = Array.from({ length: 17 }, () => createSampleThread({ resourceId }));
 
         for (const thread of threads) {
           await store.saveThread({ thread });
         }
 
-        const page1 = await store.getThreadsByResourceId({ resourceId, page: 0, perPage: 7 });
+        const page1 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 0, perPage: 7 });
         expect(page1.threads).toHaveLength(7);
 
-        const page3 = await store.getThreadsByResourceId({ resourceId, page: 2, perPage: 7 });
+        const page3 = await store.getThreadsByResourceIdPaginated({ resourceId, page: 2, perPage: 7 });
         expect(page3.threads).toHaveLength(3);
 
-        const limited = await store.getThreadsByResourceId({ resourceId, page: 1, perPage: 5 });
+        const limited = await store.getThreadsByResourceIdPaginated({ resourceId, page: 1, perPage: 5 });
         expect(limited.threads).toHaveLength(5);
       });
     });